feat/ Show more details about Web Search Steps in the UI

pingpong/ai.py

@@ -1346,6 +1346,33 @@ async def add_cached_message_part_on_output_text_url_citation_added(
 
         await add_cached_message_part_on_output_text_url_citation_added()
 
+        self.enqueue(
+            {
+                "type": "message_delta",
+                "delta": {
+                    "content": [
+                        {
+                            "index": 0,
+                            "type": "text",
+                            "text": {
+                                "value": "",
+                                "annotations": [
+                                    {
+                                        "type": "url_citation",
+                                        "end_index": data["end_index"],
+                                        "start_index": data["start_index"],
+                                        "url": data["url"],
+                                        "title": data["title"],
+                                    }
+                                ],
+                            },
+                        },
+                    ],
+                    "role": None,
+                },
+            }
+        )
+
     async def on_output_text_part_done(self, data: ResponseOutputText):
         if not self.message_part_id:
             logger.exception(
@@ -1715,6 +1742,35 @@ async def add_cached_tool_call_on_code_interpreter_tool_call_done(
         self.tool_call_id = None
         self.tool_call_external_id = None
 
+    def get_action_payload(
+        self,
+        action: ActionSearch | ActionFind | ActionOpenPage | None,
+    ):
+        if not action:
+            return None
+
+        match action.type:
+            case "search":
+                return {
+                    "type": WebSearchActionType.SEARCH.value,
+                    "query": action.query,
+                    "sources": [{"url": source.url} for source in action.sources or []],
+                }
+            case "find":
+                return {
+                    "type": WebSearchActionType.FIND.value,
+                    "pattern": action.pattern,
+                    "url": action.url,
+                }
+            case "open_page":
+                return {
+                    "type": WebSearchActionType.OPEN_PAGE.value,
+                    "url": action.url,
+                }
+            case _:
+                return None
+        return None
+
     async def on_web_search_call_created(self, data: ResponseFunctionWebSearch):
         if not self.run_id:
             logger.exception(
@@ -1751,6 +1807,21 @@ async def add_cached_tool_call_on_web_search_call_created(
         self.tool_call_id = tool_call.id
         self.tool_call_external_id = tool_call.tool_call_id
 
+        self.enqueue(
+            {
+                "type": "tool_call_created",
+                "tool_call": {
+                    "id": str(data.id),
+                    "index": self.prev_output_index,
+                    "output_index": self.prev_output_index,
+                    "type": "web_search",
+                    "web_search": {
+                        "action": self.get_action_payload(data.action),
+                    },
+                },
+            }
+        )
+
     async def on_web_search_call_in_progress(
         self, data: ResponseWebSearchCallInProgressEvent
     ):
@@ -1933,6 +2004,20 @@ async def add_cached_tool_call_on_web_search_call_done(
             )
             await session_.commit()
 
+        self.enqueue(
+            {
+                "type": "tool_call_delta",
+                "delta": {
+                    "type": "web_search",
+                    "id": data.id,
+                    "index": self.prev_output_index,
+                    "run_id": str(self.run_id),
+                    "status": data.status,
+                    "action": self.get_action_payload(data.action),
+                },
+            }
+        )
+
         await add_cached_tool_call_on_web_search_call_done()
         self.tool_call_id = None
         self.tool_call_external_id = None

pingpong/schemas.py

@@ -1,10 +1,26 @@
 from datetime import date, datetime
 from enum import Enum, StrEnum, auto
 from typing import Generic, Literal, NotRequired, TypeVar, Union
-from typing_extensions import TypedDict
-
+from typing_extensions import TypedDict, Annotated, TypeAlias
+
+from openai._utils import PropertyInfo
+from openai.types.beta.threads import (
+    ImageFileContentBlock,
+    TextContentBlock,
+    RefusalContentBlock,
+    ImageURLContentBlock,
+)
+from openai.types.beta.threads.text import Text as OpenAIText
+from openai.types.beta.threads.annotation import (
+    FileCitationAnnotation,
+    FilePathAnnotation,
+)
 from openai.types.beta.assistant_tool import AssistantTool as Tool
 from openai.types.beta.threads import Message as OpenAIMessage
+from openai.types.responses.response_output_text import AnnotationURLCitation
+from openai.types.responses.response_function_web_search import (
+    Action as WebSearchAction,
+)
 from pydantic import (
     BaseModel,
     Field,
@@ -1422,6 +1438,33 @@ class FileSearchMessage(BaseModel):
     output_index: int | None = None
 
 
+class WebSearchActionType(StrEnum):
+    SEARCH = "search"
+    FIND = "find"
+    OPEN_PAGE = "open_page"
+
+
+class WebSearchCall(BaseModel):
+    step_id: str
+    type: Literal["web_search_call"]
+    action: WebSearchAction
+    status: Literal["in_progress", "searching", "completed", "incomplete", "failed"]
+
+
+class WebSearchMessage(BaseModel):
+    id: str
+    assistant_id: str
+    created_at: float
+    content: list[WebSearchCall]
+    metadata: dict[str, str]
+    object: Literal["thread.message"]
+    message_type: Literal["web_search_call"]
+    role: Literal["assistant"]
+    run_id: str
+    thread_id: str
+    output_index: int | None = None
+
+
 class CodeInterpreterMessage(BaseModel):
     id: str
     assistant_id: str
@@ -1455,6 +1498,31 @@ class ThreadParticipants(BaseModel):
     assistant: dict[int, str]
 
 
+ThreadAnnotation: TypeAlias = Annotated[
+    Union[FileCitationAnnotation, FilePathAnnotation, AnnotationURLCitation],
+    PropertyInfo(discriminator="type"),
+]
+
+
+class ThreadText(OpenAIText):
+    annotations: list[ThreadAnnotation]
+
+
+class ThreadTextContentBlock(TextContentBlock):
+    text: ThreadText
+
+
+ThreadMessageContent: TypeAlias = Annotated[
+    Union[
+        ImageFileContentBlock,
+        ImageURLContentBlock,
+        ThreadTextContentBlock,
+        RefusalContentBlock,
+    ],
+    PropertyInfo(discriminator="type"),
+]
+
+
 class ThreadMessage(OpenAIMessage):
     status: Literal["in_progress", "incomplete", "completed"] | None
     """
@@ -1472,6 +1540,9 @@ class ThreadMessage(OpenAIMessage):
     output_index: int | None = None
     """The output index of the message, if applicable for Next-Gen Assistants."""
 
+    content: list[ThreadMessageContent]
+    """The content of the message in array of text and/or images."""
+
     metadata: dict[str, str | bool] | None = None
     """Set of 16 key-value pairs that can be attached to an object.
 
@@ -1491,6 +1562,7 @@ class ThreadMessages(BaseModel):
     messages: list[ThreadMessage]
     fs_messages: list[FileSearchMessage] | None = None
     ci_messages: list[CodeInterpreterMessage] | None
+    ws_messages: list[WebSearchMessage] | None = None
     reasoning_messages: list["ReasoningMessage"] | None = None
     has_more: bool
 
@@ -1512,6 +1584,7 @@ class ThreadWithMeta(BaseModel):
     limit: int
     ci_messages: list[CodeInterpreterMessage] | None
     fs_messages: list[FileSearchMessage] | None = None
+    ws_messages: list[WebSearchMessage] | None = None
     reasoning_messages: list["ReasoningMessage"] | None = None
     attachments: dict[str, File] | None
     instructions: str | None
@@ -1810,12 +1883,6 @@ class ToolCallStatus(StrEnum):
     FAILED = "failed"
 
 
-class WebSearchActionType(StrEnum):
-    SEARCH = "search"
-    FIND = "find"
-    OPEN_PAGE = "open_page"
-
-
 class MessagePartType(StrEnum):
     INPUT_TEXT = "input_text"
     INPUT_IMAGE = "input_image"