[codex] Add web search grounding

src/celeste/client.py

@@ -17,6 +17,7 @@
     StreamingNotSupportedError,
     UnsupportedParameterWarning,
 )
+from celeste.grounding import Grounding
 from celeste.http import HTTPClient, get_http_client
 from celeste.io import Chunk as ChunkBase
 from celeste.io import FinishReason, Input, Output, Usage
@@ -242,6 +243,9 @@ async def _predict(
                 kwargs["reasoning"] = reasoning
             if signature:
                 kwargs["signature"] = signature
+            grounding = self._parse_grounding(response_data)
+            if grounding is not None:
+                kwargs["grounding"] = grounding
             output = self._output_class()(
                 content=content,
                 usage=self._get_usage(response_data),
@@ -263,6 +267,10 @@ def _parse_reasoning(
         """Parse reasoning from response. Returns (text, signature_blocks)."""
         return None, []
 
+    def _parse_grounding(self, response_data: dict[str, Any]) -> Grounding | None:
+        """Parse web-search grounding from response."""
+        return None
+
     def _stream(
         self,
         inputs: In,

src/celeste/grounding.py

@@ -0,0 +1,32 @@
+"""Web-search grounding for generated responses."""
+
+from pydantic import BaseModel, Field
+
+
+class GroundingSource(BaseModel):
+    """A web source the response was grounded in."""
+
+    url: str
+    title: str | None = None
+    domain: str | None = None
+
+
+class Citation(BaseModel):
+    """Links a span of the response text to its sources."""
+
+    start: int
+    end: int
+    source_indices: list[int] = Field(default_factory=list)
+    cited_text: str | None = None
+
+
+class Grounding(BaseModel):
+    """Web-search grounding: queries, sources, and per-span citations."""
+
+    queries: list[str] = Field(default_factory=list)
+    sources: list[GroundingSource] = Field(default_factory=list)
+    citations: list[Citation] = Field(default_factory=list)
+    search_entry_point: str | None = None
+
+
+__all__ = ["Citation", "Grounding", "GroundingSource"]

src/celeste/io.py

@@ -14,6 +14,7 @@
 )
 from celeste.constraints import Constraint
 from celeste.core import InputType
+from celeste.grounding import Grounding
 from celeste.tools import ToolCall
 
 
@@ -45,6 +46,7 @@ class Output[Content](BaseModel):
     finish_reason: FinishReason | None = None
     metadata: dict[str, Any] = Field(default_factory=dict)
     tool_calls: list[ToolCall] = Field(default_factory=list)
+    grounding: Grounding | None = None
 
 
 class Chunk[Content](BaseModel):

src/celeste/modalities/text/grounding.py

@@ -0,0 +1,49 @@
+"""Text grounding helpers."""
+
+from celeste.grounding import Citation, Grounding, GroundingSource
+
+
+def map_url_citation_annotations(
+    annotations: list[dict[str, object]],
+) -> Grounding | None:
+    """Map native URL citation annotations into text grounding."""
+    sources: list[GroundingSource] = []
+    source_indices: dict[tuple[str, str | None], int] = {}
+    citations: list[Citation] = []
+
+    for annotation in annotations:
+        if annotation.get("type") != "url_citation":
+            continue
+        nested = annotation.get("url_citation")
+        citation: dict[str, object] = nested if isinstance(nested, dict) else annotation
+        url = citation.get("url")
+        start = citation.get("start_index")
+        end = citation.get("end_index")
+        if (
+            not isinstance(url, str)
+            or not isinstance(start, int)
+            or not isinstance(end, int)
+        ):
+            continue
+        raw_title = citation.get("title")
+        title = (
+            raw_title
+            if isinstance(raw_title, str) and not raw_title.strip().isdigit()
+            else None
+        )
+        key = (url, title)
+        if key not in source_indices:
+            source_indices[key] = len(sources)
+            sources.append(GroundingSource(url=url, title=key[1]))
+        citations.append(
+            Citation(
+                start=start,
+                end=end,
+                source_indices=[source_indices[key]],
+            )
+        )
+
+    return Grounding(sources=sources, citations=citations) if citations else None
+
+
+__all__ = ["map_url_citation_annotations"]

src/celeste/modalities/text/protocols/chatcompletions/client.py

@@ -3,6 +3,7 @@
 import json
 from typing import Any
 
+from celeste.grounding import Grounding
 from celeste.messages import (
     content_to_text,
     message_parts,
@@ -16,13 +17,18 @@
 from celeste.protocols.chatcompletions.streaming import (
     ChatCompletionsStream as _ChatCompletionsStream,
 )
-from celeste.protocols.chatcompletions.tools import parse_tool_calls
+from celeste.protocols.chatcompletions.tools import (
+    parse_annotations,
+    parse_tool_calls,
+)
 from celeste.tools import ToolCall, ToolResult
 from celeste.types import DocumentPart, ImagePart, Message, Role, TextContent, TextPart
 from celeste.utils import build_document_data_url, build_image_data_url
 
 from ...client import TextClient
+from ...grounding import map_url_citation_annotations
 from ...io import (
+    TextChunk,
     TextInput,
 )
 from ...streaming import TextStream
@@ -106,6 +112,12 @@ def _serialize_messages(
 class ChatCompletionsTextStream(_ChatCompletionsStream, TextStream):
     """Chat Completions streaming for text modality."""
 
+    def _aggregate_grounding(
+        self, chunks: list[TextChunk], raw_events: list[dict[str, Any]]
+    ) -> Grounding | None:
+        """Aggregate URL citations from streamed Chat Completions annotations."""
+        return map_url_citation_annotations(self._annotations)
+
 
 class ChatCompletionsTextClient(ChatCompletionsMixin, TextClient):
     """Chat Completions text client."""
@@ -150,6 +162,10 @@ def _parse_tool_calls(self, response_data: dict[str, Any]) -> list[ToolCall]:
         """Parse tool calls from Chat Completions response."""
         return parse_tool_calls(response_data)
 
+    def _parse_grounding(self, response_data: dict[str, Any]) -> Grounding | None:
+        """Parse grounding from Chat Completions citation/search extensions."""
+        return map_url_citation_annotations(parse_annotations(response_data))
+
     def _stream_class(self) -> type[TextStream]:
         """Return the Stream class for this provider."""
         return ChatCompletionsTextStream

src/celeste/modalities/text/protocols/openresponses/client.py

@@ -4,6 +4,7 @@
 from typing import Any
 
 from celeste.artifacts import DocumentArtifact
+from celeste.grounding import Grounding
 from celeste.messages import (
     content_to_text,
     message_parts,
@@ -18,6 +19,7 @@
     OpenResponsesStream as _OpenResponsesStream,
 )
 from celeste.protocols.openresponses.tools import (
+    parse_annotations,
     parse_content,
     parse_reasoning,
     parse_tool_calls,
@@ -27,6 +29,7 @@
 from celeste.utils import build_document_data_url, build_image_data_url
 
 from ...client import TextClient
+from ...grounding import map_url_citation_annotations
 from ...io import (
     TextChunk,
     TextInput,
@@ -153,6 +156,16 @@ def _aggregate_signature(
         _, signature_blocks = parse_reasoning(self._response_data.get("output", []))
         return signature_blocks
 
+    def _aggregate_grounding(
+        self, chunks: list[TextChunk], raw_events: list[dict[str, Any]]
+    ) -> Grounding | None:
+        """Extract URL citations from response.completed data."""
+        if self._response_data is None:
+            return None
+        return map_url_citation_annotations(
+            parse_annotations(self._response_data.get("output", []))
+        )
+
 
 class OpenResponsesTextClient(OpenResponsesMixin, TextClient):
     """OpenResponses text client using Responses API."""
@@ -192,6 +205,12 @@ def _parse_tool_calls(self, response_data: dict[str, Any]) -> list[ToolCall]:
         """Parse tool calls from OpenResponses response."""
         return parse_tool_calls(response_data)
 
+    def _parse_grounding(self, response_data: dict[str, Any]) -> Grounding | None:
+        """Parse URL citations from OpenResponses response."""
+        return map_url_citation_annotations(
+            parse_annotations(response_data.get("output", []))
+        )
+
     def _stream_class(self) -> type[TextStream]:
         """Return the Stream class for this provider."""
         return OpenResponsesTextStream

src/celeste/modalities/text/providers/anthropic/client.py

@@ -2,9 +2,11 @@
 
 import base64
 import contextlib
+import json
 from typing import Any
 
 from celeste.artifacts import DocumentArtifact, ImageArtifact
+from celeste.grounding import Grounding
 from celeste.messages import (
     content_to_text,
     message_parts,
@@ -27,6 +29,7 @@
     TextInput,
 )
 from ...streaming import TextStream
+from .grounding import parse_grounding
 from .parameters import ANTHROPIC_PARAMETER_MAPPERS
 
 
@@ -95,14 +98,12 @@ def _aggregate_tool_calls(
         self, chunks: list[TextChunk], raw_events: list[dict[str, Any]]
     ) -> list[ToolCall]:
         """Reconstruct tool calls from accumulated content_block events."""
-        import json as _json
-
         result: list[ToolCall] = []
         for tc in self._tool_calls.values():
             arguments = {}
             if tc["input_json"]:
                 with contextlib.suppress(ValueError, TypeError):
-                    arguments = _json.loads(tc["input_json"])
+                    arguments = json.loads(tc["input_json"])
             result.append(ToolCall(id=tc["id"], name=tc["name"], arguments=arguments))
         return result
 
@@ -112,6 +113,12 @@ def _aggregate_signature(
         """Return accumulated thinking blocks for round-trip signature."""
         return self._thinking_blocks
 
+    def _aggregate_grounding(
+        self, chunks: list[TextChunk], raw_events: list[dict[str, Any]]
+    ) -> Grounding | None:
+        """Reconstruct Anthropic web-search grounding from streamed content blocks."""
+        return parse_grounding(self._aggregate_content_blocks())
+
 
 class AnthropicTextClient(AnthropicMessagesClient, TextClient):
     """Anthropic text client."""
@@ -299,6 +306,10 @@ def _parse_tool_calls(self, response_data: dict[str, Any]) -> list[ToolCall]:
             if block.get("type") == "tool_use"
         ]
 
+    def _parse_grounding(self, response_data: dict[str, Any]) -> Grounding | None:
+        """Parse Anthropic web-search grounding from Messages content blocks."""
+        return parse_grounding(super()._parse_content(response_data))
+
     def _stream_class(self) -> type[TextStream]:
         """Return the Stream class for this provider."""
         return AnthropicTextStream

src/celeste/modalities/text/providers/anthropic/grounding.py

@@ -0,0 +1,89 @@
+"""Anthropic text grounding helpers."""
+
+from typing import Any
+
+from celeste.grounding import Citation, Grounding, GroundingSource
+
+
+def _source_index(
+    sources: list[GroundingSource],
+    source_indices: dict[tuple[str, str | None], int],
+    url: str,
+    title: str | None,
+) -> int:
+    key = (url, title)
+    if key not in source_indices:
+        source_indices[key] = len(sources)
+        sources.append(GroundingSource(url=url, title=title))
+    return source_indices[key]
+
+
+def parse_grounding(blocks: list[dict[str, Any]]) -> Grounding | None:
+    """Parse Anthropic Messages web-search blocks into text grounding."""
+    queries: list[str] = []
+    sources: list[GroundingSource] = []
+    source_indices: dict[tuple[str, str | None], int] = {}
+    citations: list[Citation] = []
+
+    for block in blocks:
+        if block.get("type") == "server_tool_use" and block.get("name") == "web_search":
+            query = block.get("input", {}).get("query")
+            if isinstance(query, str):
+                queries.append(query)
+        if block.get("type") != "web_search_tool_result":
+            continue
+        content = block.get("content")
+        if not isinstance(content, list):
+            continue
+        for result in content:
+            if result.get("type") != "web_search_result":
+                continue
+            url = result.get("url")
+            if isinstance(url, str):
+                title = result.get("title")
+                _source_index(
+                    sources,
+                    source_indices,
+                    url,
+                    title if isinstance(title, str) else None,
+                )
+
+    text_pos = 0
+    for block in blocks:
+        if block.get("type") != "text":
+            continue
+        block_text = block.get("text")
+        if not isinstance(block_text, str) or not block_text:
+            continue
+        start = text_pos
+        end = start + len(block_text)
+        text_pos = end
+        for raw_citation in block.get("citations", []):
+            if raw_citation.get("type") != "web_search_result_location":
+                continue
+            url = raw_citation.get("url")
+            cited_text = raw_citation.get("cited_text")
+            if not isinstance(url, str) or not isinstance(cited_text, str):
+                continue
+            title = raw_citation.get("title")
+            idx = _source_index(
+                sources,
+                source_indices,
+                url,
+                title if isinstance(title, str) else None,
+            )
+            citations.append(
+                Citation(
+                    start=start,
+                    end=end,
+                    source_indices=[idx],
+                    cited_text=cited_text,
+                )
+            )
+
+    if not queries and not sources and not citations:
+        return None
+    return Grounding(queries=queries, sources=sources, citations=citations)
+
+
+__all__ = ["parse_grounding"]