feat: major rewrite of widget management / introduce Widgets backend & frontend that shows interactive gallery of widgets

Author: ksachdeva

README.md

@@ -8,7 +8,6 @@ It uses and extends `openai/chatkit-python` (https://github.com/openai/chatkit-p
 - A function (`stream_agent_response`) that translate ADK events into chatkit events
 - Provides support to render `widgets`
     * See examples/backend/src/backend/agents/facts/_tools.py::get_weather
-    * Use `add_widget_to_tool_response` in your tool and `widget` will be sent to client
 - Provides support for making calls to client tools.
     * Client tools typically run in browser
     * See examples/backend/src/backend/agents/facts/_tools.py::switch_theme
@@ -51,7 +50,7 @@ See `examples` for full usage
 
 ```python
 
-from adk_chatkit import ADKContext, ADKStore, stream_agent_response
+from adk_chatkit import ADKAgentContext, ADKContext, ADKStore, ChatkitRunConfig, stream_agent_response
 
 class FactsChatkitServer(ChatKitServer[ADKContext]):
     def __init__(
@@ -80,19 +79,20 @@ class FactsChatkitServer(ChatKitServer[ADKContext]):
         if not message_text:
             return
 
-        content = genai_types.Content(
-            role="user",
-            parts=[genai_types.Part.from_text(text=message_text)],
+        agent_context = ADKAgentContext(
+            app_name=context.app_name,
+            user_id=context.user_id,
+            thread=thread,
         )
 
         event_stream = self._runner.run_async(
-            user_id=context["user_id"],
+            user_id=context.user_id,
             session_id=thread.id,
             new_message=content,
-            run_config=RunConfig(streaming_mode=StreamingMode.SSE),
+            run_config=ChatkitRunConfig(streaming_mode=StreamingMode.SSE, context=agent_context),
         )
 
-        async for event in stream_agent_response(thread, event_stream):
+        async for event in stream_agent_response(agent_context, event_stream):
             yield event
 
 ```

adk-chatkit/src/adk_chatkit/__init__.py

@@ -1,20 +1,22 @@
 from .__about__ import __application__, __author__, __version__
-from ._callbacks import remove_widgets_and_client_tool_calls
+from ._callbacks import remove_client_tool_calls
 from ._client_tool_call import ClientToolCallState, add_client_tool_call_to_tool_response
-from ._context import ADKContext
+from ._context import ADKAgentContext, ADKContext, ChatkitRunConfig
 from ._response import stream_agent_response
 from ._store import ADKStore
-from ._widgets import add_widget_to_tool_response
+from ._widgets import serialize_widget_item
 
 __all__ = [
     "__version__",
     "__application__",
     "__author__",
     "ADKContext",
+    "ADKAgentContext",
     "ADKStore",
     "stream_agent_response",
     "ClientToolCallState",
     "add_client_tool_call_to_tool_response",
-    "remove_widgets_and_client_tool_calls",
-    "add_widget_to_tool_response",
+    "remove_client_tool_calls",
+    "ChatkitRunConfig",
+    "serialize_widget_item",
 ]

adk-chatkit/src/adk_chatkit/_callbacks.py

@@ -2,20 +2,17 @@
 from google.adk.models.llm_request import LlmRequest
 from google.adk.models.llm_response import LlmResponse
 
-from ._constants import CLIENT_TOOL_KEY_IN_TOOL_RESPONSE, WIDGET_KEY_IN_TOOL_RESPONSE
+from ._constants import CLIENT_TOOL_KEY_IN_TOOL_RESPONSE
 
 
-def remove_widgets_and_client_tool_calls(
-    callback_context: CallbackContext, llm_request: LlmRequest
-) -> LlmResponse | None:
+def remove_client_tool_calls(callback_context: CallbackContext, llm_request: LlmRequest) -> LlmResponse | None:
     for c in llm_request.contents:
         if c.parts is None:
             continue
         for p in c.parts:
             if not p.function_response:
                 continue
             if p.function_response.response:
-                p.function_response.response.pop(WIDGET_KEY_IN_TOOL_RESPONSE, None)
                 p.function_response.response.pop(CLIENT_TOOL_KEY_IN_TOOL_RESPONSE, None)
 
     return None

adk-chatkit/src/adk_chatkit/_constants.py

@@ -1,5 +1,5 @@
 from typing import Final
 
-WIDGET_KEY_IN_TOOL_RESPONSE: Final[str] = "adk-widget"
 CLIENT_TOOL_KEY_IN_TOOL_RESPONSE: Final[str] = "adk-client-tool"
 CHATKIT_THREAD_METADTA_KEY: Final[str] = "adk-chatkit-thread-metadata"
+CHATKIT_WIDGET_STATE_KEY: Final[str] = "adk-chatkit-widget-state"

adk-chatkit/src/adk_chatkit/_context.py

@@ -1,6 +1,28 @@
-from typing import TypedDict
+import asyncio
 
+from chatkit.types import ThreadMetadata, ThreadStreamEvent
+from google.adk.runners import RunConfig
+from pydantic import BaseModel
 
-class ADKContext(TypedDict):
+from ._event_utils import QueueCompleteSentinel
+
+
+class ADKContext(BaseModel):
     app_name: str
     user_id: str
+
+
+class ADKAgentContext(ADKContext):
+    thread: ThreadMetadata
+
+    _events: asyncio.Queue[ThreadStreamEvent | QueueCompleteSentinel] = asyncio.Queue()
+
+    async def stream(self, event: ThreadStreamEvent) -> None:
+        await self._events.put(event)
+
+    def _complete(self) -> None:
+        self._events.put_nowait(QueueCompleteSentinel())
+
+
+class ChatkitRunConfig(RunConfig):
+    context: ADKAgentContext

adk-chatkit/src/adk_chatkit/_event_utils.py

@@ -0,0 +1,82 @@
+import asyncio
+from collections.abc import AsyncIterator
+from typing import TypeVar
+
+from chatkit.types import ThreadStreamEvent
+
+T1 = TypeVar("T1")
+T2 = TypeVar("T2")
+
+
+class QueueCompleteSentinel: ...
+
+
+async def merge_generators(
+    a: AsyncIterator[T1],
+    b: AsyncIterator[T2],
+) -> AsyncIterator[T1 | T2]:
+    pending: list[AsyncIterator[T1 | T2]] = [a, b]
+    pending_tasks: dict[asyncio.Task, AsyncIterator[T1 | T2]] = {
+        asyncio.ensure_future(g.__anext__()): g for g in pending
+    }
+    while len(pending_tasks) > 0:
+        done, _ = await asyncio.wait(pending_tasks.keys(), return_when="FIRST_COMPLETED")
+        stop = False
+        for d in done:
+            try:
+                result = d.result()
+                yield result
+                dg = pending_tasks[d]
+                pending_tasks[asyncio.ensure_future(dg.__anext__())] = dg
+            except StopAsyncIteration:
+                stop = True
+            finally:
+                del pending_tasks[d]
+        if stop:
+            for task in pending_tasks.keys():
+                if not task.cancel():
+                    try:
+                        yield task.result()
+                    except asyncio.CancelledError:
+                        pass
+                    except asyncio.InvalidStateError:
+                        pass
+            break
+
+
+class EventWrapper:
+    def __init__(self, event: ThreadStreamEvent):
+        self.event = event
+
+
+class AsyncQueueIterator(AsyncIterator[EventWrapper]):
+    def __init__(self, queue: asyncio.Queue[ThreadStreamEvent | QueueCompleteSentinel]):
+        self.queue = queue
+        self.completed = False
+
+    def __aiter__(self) -> AsyncIterator[EventWrapper]:
+        return self
+
+    async def __anext__(self) -> EventWrapper:
+        if self.completed:
+            raise StopAsyncIteration
+
+        item = await self.queue.get()
+        if isinstance(item, QueueCompleteSentinel):
+            self.completed = True
+            raise StopAsyncIteration
+        return EventWrapper(item)
+
+    def drain_and_complete(self) -> None:
+        """Empty the underlying queue without awaiting and mark this iterator completed.
+
+        This is intended for cleanup paths where we must guarantee no awaits
+        occur. All queued items, including any completion sentinel, are
+        discarded.
+        """
+        while True:
+            try:
+                self.queue.get_nowait()
+            except asyncio.QueueEmpty:
+                break
+        self.completed = True

adk-chatkit/src/adk_chatkit/_response.py

@@ -14,25 +14,59 @@
     ThreadItemUpdated,
     ThreadMetadata,
     ThreadStreamEvent,
-    WidgetItem,
 )
 from google.adk.events import Event
 
 from ._client_tool_call import ClientToolCallState
-from ._constants import CLIENT_TOOL_KEY_IN_TOOL_RESPONSE, WIDGET_KEY_IN_TOOL_RESPONSE
+from ._constants import CLIENT_TOOL_KEY_IN_TOOL_RESPONSE
+from ._context import ADKAgentContext
+from ._event_utils import AsyncQueueIterator, EventWrapper, merge_generators
 
 
-async def stream_agent_response(
+async def _handle_function_response(
+    event: Event,
     thread: ThreadMetadata,
+) -> AsyncGenerator[ThreadItemDoneEvent, None]:
+    if fn_responses := event.get_function_responses():
+        for fn_response in fn_responses:
+            if not fn_response.response:
+                continue
+
+            adk_client_tool: ClientToolCallState | None = fn_response.response.get(
+                CLIENT_TOOL_KEY_IN_TOOL_RESPONSE, None
+            )
+            if adk_client_tool:
+                yield ThreadItemDoneEvent(
+                    item=ClientToolCallItem(
+                        id=event.id,
+                        thread_id=thread.id,
+                        name=adk_client_tool.name,
+                        arguments=adk_client_tool.arguments,
+                        status=adk_client_tool.status,
+                        created_at=datetime.fromtimestamp(event.timestamp),
+                        call_id=adk_client_tool.id,
+                    ),
+                )
+
+
+async def stream_agent_response(
+    context: ADKAgentContext,
     adk_response: AsyncGenerator[Event, None],
 ) -> AsyncIterator[ThreadStreamEvent]:
-    if adk_response is None:
-        return
-
+    queue_iterator = AsyncQueueIterator(context._events)
     response_id = str(uuid.uuid4())
 
+    thread = context.thread
+
     content_index = 0
-    async for event in adk_response:
+    async for event in merge_generators(adk_response, queue_iterator):
+        if event is None:
+            continue
+
+        if isinstance(event, EventWrapper):
+            yield event.event
+            continue
+
         if event.content is None:
             # we need to throw item added event first
             yield ThreadItemAddedEvent(
@@ -53,38 +87,8 @@ async def stream_agent_response(
                 ),
             )
         else:
-            # Since Widgets are recorded in the function responses
-            # they are handled here
-            if fn_responses := event.get_function_responses():
-                for fn_response in fn_responses:
-                    if not fn_response.response:
-                        continue
-                    widget = fn_response.response.get(WIDGET_KEY_IN_TOOL_RESPONSE, None)
-                    if widget:
-                        # No Streaming for Widgets for now
-                        yield ThreadItemDoneEvent(
-                            item=WidgetItem(
-                                id=str(uuid.uuid4()),
-                                thread_id=thread.id,
-                                created_at=datetime.fromtimestamp(event.timestamp),
-                                widget=widget,
-                            )
-                        )
-                    adk_client_tool: ClientToolCallState | None = fn_response.response.get(
-                        CLIENT_TOOL_KEY_IN_TOOL_RESPONSE, None
-                    )
-                    if adk_client_tool:
-                        yield ThreadItemDoneEvent(
-                            item=ClientToolCallItem(
-                                id=event.id,
-                                thread_id=thread.id,
-                                name=adk_client_tool.name,
-                                arguments=adk_client_tool.arguments,
-                                status=adk_client_tool.status,
-                                created_at=datetime.fromtimestamp(event.timestamp),
-                                call_id=adk_client_tool.id,
-                            ),
-                        )
+            async for item in _handle_function_response(event, thread):
+                yield item
 
             if event.content.parts:
                 text_from_final_update = ""
@@ -116,3 +120,9 @@ async def stream_agent_response(
                         created_at=datetime.fromtimestamp(event.timestamp),
                     )
                 )
+
+    context._complete()
+
+    # Drain remaining events
+    async for event in queue_iterator:
+        yield event.event