feat: introduce stream_widget helper function in ADKAgentContext

Author: ksachdeva

adk-chatkit/src/adk_chatkit/_context.py

@@ -1,7 +1,10 @@
 import asyncio
+from datetime import datetime
 
-from chatkit.types import ThreadMetadata, ThreadStreamEvent
-from google.adk.runners import RunConfig
+from chatkit.types import ThreadItemDoneEvent, ThreadMetadata, ThreadStreamEvent, WidgetItem
+from chatkit.widgets import WidgetRoot
+from google.adk.agents.run_config import RunConfig
+from google.adk.tools import ToolContext
 from pydantic import BaseModel
 
 from ._event_utils import QueueCompleteSentinel
@@ -20,6 +23,20 @@ class ADKAgentContext(ADKContext):
     async def stream(self, event: ThreadStreamEvent) -> None:
         await self._events.put(event)
 
+    async def stream_widget(self, widget: WidgetRoot, tool_context: ToolContext) -> None:
+        if tool_context.function_call_id is None:
+            raise ValueError("tool_context.function_call_id is None")
+        await self.stream(
+            ThreadItemDoneEvent(
+                item=WidgetItem(
+                    id=tool_context.function_call_id,
+                    thread_id=self.thread.id,
+                    created_at=datetime.now(),
+                    widget=widget,
+                )
+            )
+        )
+
     def _complete(self) -> None:
         self._events.put_nowait(QueueCompleteSentinel())
 

examples/backend/src/backend/agents/facts/_tools.py

@@ -1,8 +1,6 @@
-from datetime import datetime
 from typing import Any, Final, Literal, cast
 
 from adk_chatkit import ChatkitRunConfig, ClientToolCallState, add_client_tool_call_to_tool_response
-from chatkit.types import ProgressUpdateEvent, ThreadItemDoneEvent, ThreadStreamEvent, WidgetItem
 from google.adk.tools import ToolContext
 
 from ._sample_widget import render_weather_widget
@@ -129,24 +127,10 @@ async def get_weather(
     if observed:
         result["observation_time"] = observed
 
-    # add_widget_to_tool_response(
-    #     result,
-    #     widget,
-    # )
-
     chatkit_run_config = cast(ChatkitRunConfig, tool_context._invocation_context.run_config)
 
     try:
-        await chatkit_run_config.context.stream(
-            ThreadItemDoneEvent(
-                item=WidgetItem(
-                    id=tool_context.function_call_id,
-                    thread_id=chatkit_run_config.context.thread.id,
-                    created_at=datetime.now(),
-                    widget=widget,
-                )
-            )
-        )
+        await chatkit_run_config.context.stream_widget(widget, tool_context)
     except Exception as exc:  # noqa: BLE001
         print("[WeatherTool] widget stream failed", {"error": str(exc)})
         raise ValueError("Weather data is currently unavailable for that location.") from exc

examples/backend/src/backend/agents/widgets/_tools.py

@@ -1,9 +1,8 @@
 import asyncio
-from datetime import datetime
 from typing import cast
 
 from adk_chatkit import ChatkitRunConfig
-from chatkit.types import ProgressUpdateEvent, ThreadItemDoneEvent, WidgetItem
+from chatkit.types import ProgressUpdateEvent
 from google.adk.tools import ToolContext
 
 from ._tasks_widget import make_widget
@@ -22,19 +21,9 @@ async def render_tasks_widget(tool_context: ToolContext) -> dict[str, str]:
 
     await chatkit_run_config.context.stream(ProgressUpdateEvent(text="Fetching tasks widget..."))
     await asyncio.sleep(2)
-    await chatkit_run_config.context.stream(ProgressUpdateEvent(text=""))
 
     widget = make_widget()
 
-    await chatkit_run_config.context.stream(
-        ThreadItemDoneEvent(
-            item=WidgetItem(
-                id=tool_context.function_call_id,
-                thread_id=chatkit_run_config.context.thread.id,
-                created_at=datetime.now(),
-                widget=widget,
-            )
-        )
-    )
+    await chatkit_run_config.context.stream_widget(widget, tool_context)
 
     return result