feat(chat): add streaming support to creating runs

Author: adi-wan-askui

src/askui/chat/__main__.py

@@ -338,7 +338,6 @@ def write_message(
         )
         write_message(last_message)
     run = run_service.create(thread_id, stream=False)
-    print(run)
     time.sleep(1)
     while run := run_service.retrieve(run.id):
         new_messages = message_service.list_(
@@ -352,5 +351,32 @@ def write_message(
         time.sleep(1)
 
 
+if act_prompt := st.chat_input("Ask AI (streaming)"):
+    if act_prompt != "Continue":
+        last_message = message_service.create(
+            thread_id=thread_id,
+            message=MessageParam(
+                role="user",
+                content=act_prompt,
+            ),
+        )
+        write_message(last_message)
+
+    # Use the streaming API
+    event_stream = run_service.create(thread_id, stream=True)
+    import asyncio
+
+    async def handle_stream() -> None:
+        last_msg_id = last_message.id if last_message else None
+        async for event in event_stream:
+            if event.event == "message.created":
+                msg = event.data
+                if msg and (not last_msg_id or msg.id > last_msg_id):
+                    write_message(msg)
+                    last_msg_id = msg.id
+
+    # Run the async handler in Streamlit (sync context)
+    asyncio.run(handle_stream())
+
 # if st.button("Rerun"):
 #     rerun()

src/askui/chat/api/messages/service.py

@@ -1,8 +1,10 @@
 from datetime import datetime, timezone
 from pathlib import Path
+from typing import Literal
 
 from pydantic import AwareDatetime, BaseModel, Field
 
+from askui.chat.api.models import Event
 from askui.chat.api.utils import generate_time_ordered_id
 from askui.models.shared.computer_agent_message_param import MessageParam
 
@@ -18,6 +20,11 @@ class Message(MessageParam):
     object: str = "message"
 
 
+class MessageEvent(Event):
+    data: Message
+    event: Literal["message.created"]
+
+
 class MessageListResponse(BaseModel):
     """Response model for listing messages."""
 

src/askui/chat/api/models.py

@@ -0,0 +1,7 @@
+from typing import Literal
+
+from pydantic import BaseModel
+
+
+class Event(BaseModel):
+    object: Literal["event"] = "event"

src/askui/chat/api/runs/router.py

@@ -1,10 +1,15 @@
-from typing import Annotated
+from collections.abc import AsyncGenerator
+from typing import TYPE_CHECKING, Annotated, cast
 
 from fastapi import APIRouter, Body, HTTPException, Path
+from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 
+if TYPE_CHECKING:
+    from askui.chat.api.messages.service import MessageEvent
+
 from .dependencies import RunServiceDep
-from .service import Run, RunListResponse, RunService
+from .service import Run, RunEvent, RunListResponse, RunService
 
 
 class CreateRunRequest(BaseModel):
@@ -19,11 +24,23 @@ def create_run(
     thread_id: Annotated[str, Path(...)],
     request: Annotated[CreateRunRequest, Body(...)],
     run_service: RunService = RunServiceDep,
-) -> Run:
+) -> Run | StreamingResponse:
     """
     Create a new run for a given thread.
     """
-    return run_service.create(thread_id, request.stream)
+    stream = request.stream
+    run_or_async_generator = run_service.create(thread_id, stream)
+    if stream:
+        async_generator = cast(
+            "AsyncGenerator[RunEvent | MessageEvent, None]", run_or_async_generator
+        )
+
+        async def sse_event_stream() -> AsyncGenerator[str, None]:
+            async for event in async_generator:
+                yield f"event: {event.event}\ndata: {event.model_dump_json()}\n\n"
+
+        return StreamingResponse(sse_event_stream(), media_type="text/event-stream")
+    return cast("Run", run_or_async_generator)
 
 
 @router.get("/{run_id}")

src/askui/chat/api/runs/service.py

@@ -1,12 +1,16 @@
-from concurrent.futures import ThreadPoolExecutor
+import asyncio
+import queue
+import threading
+from collections.abc import AsyncGenerator
 from datetime import datetime, timedelta, timezone
 from pathlib import Path
-from typing import Literal, Sequence, cast
+from typing import Literal, Sequence, cast, overload
 
 from pydantic import AwareDatetime, BaseModel, Field, computed_field
 
 from askui.agent import VisionAgent
-from askui.chat.api.messages.service import MessageService
+from askui.chat.api.messages.service import MessageEvent, MessageService
+from askui.chat.api.models import Event
 from askui.chat.api.utils import generate_time_ordered_id
 from askui.models.shared.computer_agent_cb_param import OnMessageCbParam
 from askui.models.shared.computer_agent_message_param import MessageParam
@@ -70,15 +74,33 @@ class RunListResponse(BaseModel):
     has_more: bool = False
 
 
+class RunEvent(Event):
+    data: Run
+    event: Literal[
+        "run.created",
+        "run.started",
+        "run.completed",
+        "run.failed",
+        "run.cancelled",
+        "run.expired",
+    ]
+
+
 class Runner:
     def __init__(self, run: Run, base_dir: Path) -> None:
         self._run = run
         self._base_dir = base_dir
         self._runs_dir = base_dir / "runs"
         self._msg_service = MessageService(self._base_dir)
 
-    def run_task(self) -> None:
+    def run(self, event_queue: queue.Queue[RunEvent | MessageEvent | None]) -> None:
         self._mark_started()
+        event_queue.put(
+            RunEvent(
+                data=self._run,
+                event="run.started",
+            )
+        )
         messages: list[MessageParam] = [
             cast("MessageParam", msg)
             for msg in self._msg_service.list_(self._run.thread_id).data
@@ -87,27 +109,63 @@ def run_task(self) -> None:
         def on_message(
             on_message_cb_param: OnMessageCbParam,
         ) -> MessageParam | None:
-            self._msg_service.create(
+            message = self._msg_service.create(
                 thread_id=self._run.thread_id,
                 message=on_message_cb_param.message,
             )
+            event_queue.put(
+                MessageEvent(
+                    data=message,
+                    event="message.created",
+                )
+            )
             updated_run = self._retrieve_run()
-            if self._should_abort(updated_run):
+            if updated_run.status == "cancelling":
                 updated_run.cancelled_at = datetime.now(tz=timezone.utc)
                 self._update_run_file(updated_run)
+                event_queue.put(
+                    RunEvent(
+                        data=updated_run,
+                        event="run.cancelled",
+                    )
+                )
+                return None
+            if updated_run.status == "expired":
+                event_queue.put(
+                    RunEvent(
+                        data=updated_run,
+                        event="run.expired",
+                    )
+                )
                 return None
             return on_message_cb_param.message
 
         try:
             with VisionAgent() as agent:
                 agent.act(messages, on_message=on_message)
-            self._run.completed_at = datetime.now(tz=timezone.utc)
-            self._update_run_file(self._run)
+            updated_run = self._retrieve_run()
+            if updated_run.status == "in_progress":
+                updated_run.completed_at = datetime.now(tz=timezone.utc)
+                self._update_run_file(updated_run)
+                event_queue.put(
+                    RunEvent(
+                        data=updated_run,
+                        event="run.completed",
+                    )
+                )
         except Exception as e:  # noqa: BLE001
-            self._run.failed_at = datetime.now(tz=timezone.utc)
-            self._run.last_error = RunError(message=str(e), code="server_error")
-            self._update_run_file(self._run)
-            raise
+            updated_run = self._retrieve_run()
+            updated_run.failed_at = datetime.now(tz=timezone.utc)
+            updated_run.last_error = RunError(message=str(e), code="server_error")
+            self._update_run_file(updated_run)
+            event_queue.put(
+                RunEvent(
+                    data=updated_run,
+                    event="run.failed",
+                )
+            )
+        finally:
+            event_queue.put(None)
 
     def _mark_started(self) -> None:
         self._run.started_at = datetime.now(tz=timezone.utc)
@@ -132,29 +190,56 @@ class RunService:
     Service for managing runs. Handles creation, retrieval, listing, and cancellation of runs.
     """
 
-    _executor: ThreadPoolExecutor = ThreadPoolExecutor(max_workers=4)
-
     def __init__(self, base_dir: Path) -> None:
         self._base_dir = base_dir
         self._runs_dir = base_dir / "runs"
 
     def _run_path(self, thread_id: str, run_id: str) -> Path:
         return self._runs_dir / f"{thread_id}__{run_id}.json"
 
-    def create(self, thread_id: str, stream: bool) -> Run:
+    def _create_run(self, thread_id: str) -> Run:
         run = Run(thread_id=thread_id)
         self._runs_dir.mkdir(parents=True, exist_ok=True)
         self._update_run_file(run)
-        runner = Runner(run, self._base_dir)
-        # TODO(adi-wan-askui): Run differently depending on `stream` parameter
-        runner.run_task()
-        # if not stream:
-        #     self._start_run_background(run)
         return run
 
-    def _start_run_background(self, run: Run) -> None:
+    @overload
+    def create(self, thread_id: str, stream: Literal[False]) -> Run: ...
+
+    @overload
+    def create(
+        self, thread_id: str, stream: Literal[True]
+    ) -> AsyncGenerator[RunEvent | MessageEvent, None]: ...
+
+    @overload
+    def create(
+        self, thread_id: str, stream: bool
+    ) -> Run | AsyncGenerator[RunEvent | MessageEvent, None]: ...
+
+    def create(
+        self, thread_id: str, stream: bool
+    ) -> Run | AsyncGenerator[RunEvent | MessageEvent, None]:
+        run = self._create_run(thread_id)
+        event_queue: queue.Queue[RunEvent | MessageEvent | None] = queue.Queue()
         runner = Runner(run, self._base_dir)
-        self._executor.submit(runner.run_task)
+        thread = threading.Thread(target=runner.run, args=(event_queue,))
+        thread.start()
+        if stream:
+
+            async def event_stream() -> AsyncGenerator[RunEvent | MessageEvent, None]:
+                yield RunEvent(
+                    data=run,
+                    event="run.created",
+                )
+                loop = asyncio.get_event_loop()
+                while True:
+                    event = await loop.run_in_executor(None, event_queue.get)
+                    if event is None:
+                        break
+                    yield event
+
+            return event_stream()
+        return run
 
     def _update_run_file(self, run: Run) -> None:
         run_file = self._run_path(run.thread_id, run.id)