.

Author: rguan72

backend/onyx/chat/turn/fast_chat_turn.py

@@ -8,6 +8,11 @@
 from onyx.chat.turn.infra.chat_turn_orchestration import unified_event_stream
 from onyx.chat.turn.models import MyContext
 from onyx.chat.turn.models import RunDependencies
+from onyx.server.query_and_chat.streaming_models import MessageDelta
+from onyx.server.query_and_chat.streaming_models import MessageStart
+from onyx.server.query_and_chat.streaming_models import OverallStop
+from onyx.server.query_and_chat.streaming_models import Packet
+from onyx.server.query_and_chat.streaming_models import SectionEnd
 from onyx.tools.tool_implementations_v2.internal_search import internal_search
 from onyx.tools.tool_implementations_v2.web import web_fetch
 from onyx.tools.tool_implementations_v2.web import web_search
@@ -33,14 +38,25 @@ def fast_chat_turn(messages: list[dict], dependencies: RunDependencies) -> None:
     )
 
     bridge = OnyxRunner().run_streamed(agent, messages, context=ctx, max_turns=100)
-    try:
-        for ev in bridge.events():
-            if isinstance(ev, RunItemStreamEvent):
-                pass
-            elif isinstance(ev, RawResponsesStreamEvent):
-                # TODO: use very standardized schema for the emitter that is close to
-                # front end schema
-                dependencies.emitter.emit(kind="agent", data=ev.data.model_dump())
-    finally:
-        # TODO: Handle done signal more reliably?
-        dependencies.emitter.emit(kind="done", data={"ok": True})
+    for ev in bridge.events():
+        if isinstance(ev, RunItemStreamEvent):
+            pass
+        elif isinstance(ev, RawResponsesStreamEvent):
+            obj = None
+            if ev.data.type == "response.created":
+                obj = MessageStart(
+                    type="message_start", content="", final_documents=None
+                )
+            elif ev.data.type == "response.output_text.delta":
+                obj = MessageDelta(type="message_delta", content=ev.data.delta)
+            elif ev.data.type == "response.completed":
+                obj = OverallStop(type="stop")
+            elif ev.data.type == "response.output_item.done":
+                obj = SectionEnd(type="section_end")
+            if obj:
+                dependencies.emitter.emit(Packet(ind=ctx.current_run_step, obj=obj))
+    # TODO: Error handling
+    # Should there be a timeout and some error on the queue?
+    dependencies.emitter.emit(
+        Packet(ind=ctx.current_run_step, obj=OverallStop(type="stop"))
+    )

backend/onyx/chat/turn/infra/chat_turn_event_stream.py

@@ -5,14 +5,13 @@
 from collections.abc import Iterator
 from queue import Queue
 from typing import Any
-from typing import Dict
 from typing import Optional
 
 from agents import Agent
 from agents import Runner
 from agents import TContext
-from pydantic import BaseModel
-from pydantic import Field
+
+from onyx.server.query_and_chat.streaming_models import Packet
 
 
 logger = logging.getLogger(__name__)
@@ -85,12 +84,7 @@ def _do_cancel():
             self._loop.call_soon_threadsafe(_do_cancel)
 
 
-class StreamPacket(BaseModel):
-    kind: str  # "agent" | "tool-progress" | "done"
-    payload: Dict[str, Any] = Field(default_factory=dict)
-
-
-def convert_to_packet_obj(packet: StreamPacket) -> Any | None:
+def convert_to_packet_obj(packet: dict[str, Any]) -> Any | None:
     """Convert a packet dictionary to PacketObj when possible.
 
     Args:
@@ -176,5 +170,5 @@ class Emitter:
     def __init__(self, bus: Queue):
         self.bus = bus
 
-    def emit(self, kind: str, data: Dict[str, Any]) -> None:
-        self.bus.put(StreamPacket(kind=kind, payload=data))
+    def emit(self, packet: Packet) -> None:
+        self.bus.put(packet)

backend/onyx/chat/turn/infra/chat_turn_orchestration.py

@@ -7,10 +7,9 @@
 from typing import Dict
 from typing import List
 
-from onyx.chat.turn.infra.chat_turn_event_stream import convert_to_packet_obj
 from onyx.chat.turn.infra.chat_turn_event_stream import Emitter
-from onyx.chat.turn.infra.chat_turn_event_stream import StreamPacket
 from onyx.chat.turn.models import RunDependencies
+from onyx.server.query_and_chat.streaming_models import OverallStop
 from onyx.server.query_and_chat.streaming_models import Packet
 
 
@@ -48,12 +47,10 @@ def wrapper(
         )
         t.start()
         while True:
-            pkt: StreamPacket = emitter.bus.get()
-            if pkt.kind == "done":
+            pkt: Packet = emitter.bus.get()
+            if pkt.obj == OverallStop(type="stop"):
                 break
             else:
-                packet_obj = convert_to_packet_obj(pkt.payload)
-                if packet_obj:
-                    yield Packet(ind=0, obj=packet_obj)
+                yield pkt
 
     return wrapper

backend/onyx/chat/turn/models.py

@@ -18,3 +18,4 @@ class MyContext:
 
     run_dependencies: RunDependencies | None = None
     needs_compaction: bool = False
+    current_run_step: int = 0

backend/onyx/tools/tool_implementations_v2/web.py

@@ -8,6 +8,11 @@
     get_default_provider,
 )
 from onyx.chat.turn.models import MyContext
+from onyx.configs.constants import DocumentSource
+from onyx.server.query_and_chat.streaming_models import Packet
+from onyx.server.query_and_chat.streaming_models import SavedSearchDoc
+from onyx.server.query_and_chat.streaming_models import SearchToolDelta
+from onyx.server.query_and_chat.streaming_models import SearchToolStart
 
 
 def short_tag(link: str, i: int) -> str:
@@ -28,7 +33,22 @@ def web_search(run_context: RunContextWrapper[MyContext], query: str) -> str:
         query: The natural-language search query.
     """
     search_provider = get_default_provider()
-    run_context.run_dependencies.emitter.emit(kind="web-search", data={"query": query})
+    run_context.context.run_dependencies.emitter.emit(
+        Packet(
+            ind=run_context.context.current_run_step + 1,
+            obj=SearchToolStart(
+                type="internal_search_tool_start", is_internet_search=True
+            ),
+        )
+    )
+    run_context.context.run_dependencies.emitter.emit(
+        Packet(
+            ind=run_context.context.current_run_step + 1,
+            obj=SearchToolDelta(
+                type="internal_search_tool_delta", queries=[query], documents=None
+            ),
+        )
+    )
     hits = search_provider.search(query)
     results = []
     for i, r in enumerate(hits):
@@ -44,6 +64,7 @@ def web_search(run_context: RunContextWrapper[MyContext], query: str) -> str:
                 ),
             }
         )
+    run_context.context.current_run_step += 2
     return json.dumps({"results": results})
 
 
@@ -60,7 +81,38 @@ def web_fetch(run_context: RunContextWrapper[MyContext], urls: List[str]) -> str
         urls: The full URLs of the pages to retrieve.
     """
     search_provider = get_default_provider()
-    run_context.run_dependencies.emitter.emit(kind="web-fetch", data={"urls": urls})
+    saved_search_docs = [
+        SavedSearchDoc(
+            document_id=url,
+            chunk_ind=0,
+            semantic_identifier=url,
+            link=url,
+            blurb=url,
+            source_type=DocumentSource.WEB,
+            boost=1,
+            hidden=False,
+            metadata={},
+            score=0.0,
+            is_relevant=None,
+            relevance_explanation=None,
+            match_highlights=[],
+            updated_at=None,
+            primary_owners=None,
+            secondary_owners=None,
+            is_internet=True,
+        )
+        for url in urls
+    ]
+    run_context.context.run_dependencies.emitter.emit(
+        Packet(
+            ind=run_context.context.current_run_step + 1,
+            obj=SearchToolDelta(
+                type="internal_search_tool_delta",
+                is_internet_search=True,
+                documents=saved_search_docs,
+            ),
+        )
+    )
     docs = search_provider.contents(urls)
     out = []
     for i, d in enumerate(docs):