From ede25c2e287aa770831ca6b1ed31cc3a14115697 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Fri, 19 Jun 2026 09:49:26 +0200 Subject: [PATCH] feat(text): add web search grounding --- src/celeste/client.py | 8 ++ src/celeste/grounding.py | 32 +++++++ src/celeste/io.py | 2 + src/celeste/modalities/text/grounding.py | 49 ++++++++++ .../text/protocols/chatcompletions/client.py | 18 +++- .../text/protocols/openresponses/client.py | 19 ++++ .../text/providers/anthropic/client.py | 17 +++- .../text/providers/anthropic/grounding.py | 89 +++++++++++++++++++ .../text/providers/google/client.py | 25 ++++++ .../text/providers/google/grounding.py | 75 ++++++++++++++++ .../modalities/text/providers/groq/client.py | 21 +++++ .../protocols/chatcompletions/streaming.py | 10 ++- .../protocols/chatcompletions/tools.py | 17 ++++ src/celeste/protocols/openresponses/tools.py | 16 ++++ .../providers/anthropic/messages/streaming.py | 70 +++++++++++++++ .../google/generate_content/grounding.py | 28 ++++++ .../google/generate_content/streaming.py | 11 +++ src/celeste/providers/groq/chat/tools.py | 33 ++++++- src/celeste/streaming.py | 10 +++ 19 files changed, 544 insertions(+), 6 deletions(-) create mode 100644 src/celeste/grounding.py create mode 100644 src/celeste/modalities/text/grounding.py create mode 100644 src/celeste/modalities/text/providers/anthropic/grounding.py create mode 100644 src/celeste/modalities/text/providers/google/grounding.py create mode 100644 src/celeste/providers/google/generate_content/grounding.py diff --git a/src/celeste/client.py b/src/celeste/client.py index 2efd8e6..cf4a9d7 100644 --- a/src/celeste/client.py +++ b/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, diff --git a/src/celeste/grounding.py b/src/celeste/grounding.py new file mode 100644 index 0000000..fb78c09 --- /dev/null +++ b/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"] diff --git a/src/celeste/io.py b/src/celeste/io.py index 797ccae..054bc83 100644 --- a/src/celeste/io.py +++ b/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): diff --git a/src/celeste/modalities/text/grounding.py b/src/celeste/modalities/text/grounding.py new file mode 100644 index 0000000..e114273 --- /dev/null +++ b/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"] diff --git a/src/celeste/modalities/text/protocols/chatcompletions/client.py b/src/celeste/modalities/text/protocols/chatcompletions/client.py index f529726..7b5e964 100644 --- a/src/celeste/modalities/text/protocols/chatcompletions/client.py +++ b/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 diff --git a/src/celeste/modalities/text/protocols/openresponses/client.py b/src/celeste/modalities/text/protocols/openresponses/client.py index 3a0e61d..9d13a1a 100644 --- a/src/celeste/modalities/text/protocols/openresponses/client.py +++ b/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 diff --git a/src/celeste/modalities/text/providers/anthropic/client.py b/src/celeste/modalities/text/providers/anthropic/client.py index b513dcf..aae3d1b 100644 --- a/src/celeste/modalities/text/providers/anthropic/client.py +++ b/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 diff --git a/src/celeste/modalities/text/providers/anthropic/grounding.py b/src/celeste/modalities/text/providers/anthropic/grounding.py new file mode 100644 index 0000000..f971fd9 --- /dev/null +++ b/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"] diff --git a/src/celeste/modalities/text/providers/google/client.py b/src/celeste/modalities/text/providers/google/client.py index ad4e143..4906c77 100644 --- a/src/celeste/modalities/text/providers/google/client.py +++ b/src/celeste/modalities/text/providers/google/client.py @@ -3,6 +3,7 @@ from typing import Any from uuid import uuid4 +from celeste.grounding import Grounding from celeste.messages import ( message_parts, request_messages, @@ -10,6 +11,10 @@ ) from celeste.parameters import ParameterMapper from celeste.providers.google.generate_content.client import GoogleGenerateContentClient +from celeste.providers.google.generate_content.grounding import ( + merge_grounding_metadata, + parse_grounding_metadata, +) from celeste.providers.google.generate_content.streaming import ( GoogleGenerateContentStream as _GoogleGenerateContentStream, ) @@ -28,12 +33,25 @@ from ...client import TextClient from ...io import TextInput from ...streaming import TextStream +from .grounding import map_grounding from .parameters import GOOGLE_PARAMETER_MAPPERS class GoogleTextStream(_GoogleGenerateContentStream, TextStream): """Google streaming for text modality.""" + def _aggregate_grounding( + self, chunks: list, raw_events: list[dict[str, Any]] + ) -> Grounding | None: + """Aggregate Google grounding metadata into text grounding.""" + metadata = getattr(self, "_grounding_metadata", None) + if not metadata: + return None + return map_grounding( + merge_grounding_metadata(metadata), + self._aggregate_content(chunks), + ) + def _aggregate_tool_calls( self, chunks: list, raw_events: list[dict[str, Any]] ) -> list[ToolCall]: @@ -180,6 +198,13 @@ def _parse_reasoning( text = "\n".join(reasoning_parts) if reasoning_parts else None return text, signature_blocks + def _parse_grounding(self, response_data: dict[str, Any]) -> Grounding | None: + """Extract Google grounding with text character offsets.""" + meta = parse_grounding_metadata(response_data) + if meta is None: + return None + return map_grounding(meta, self._parse_content(response_data)) + def _parse_tool_calls(self, response_data: dict[str, Any]) -> list[ToolCall]: """Parse tool calls from Google response.""" candidates = response_data.get("candidates", []) diff --git a/src/celeste/modalities/text/providers/google/grounding.py b/src/celeste/modalities/text/providers/google/grounding.py new file mode 100644 index 0000000..6a42e9d --- /dev/null +++ b/src/celeste/modalities/text/providers/google/grounding.py @@ -0,0 +1,75 @@ +"""Google text grounding helpers.""" + +from typing import Any + +from celeste.grounding import Citation, Grounding, GroundingSource + + +def _byte_offset_to_char_offset(text: str, offset: int) -> int | None: + encoded = text.encode("utf-8") + if offset < 0 or offset > len(encoded): + return None + try: + return len(encoded[:offset].decode("utf-8")) + except UnicodeDecodeError: + return None + + +def map_grounding(grounding_metadata: dict[str, Any], text: str) -> Grounding: + """Map Google Gemini groundingMetadata to text grounding.""" + sources: list[GroundingSource] = [] + source_indices: dict[int, int] = {} + for index, chunk in enumerate(grounding_metadata.get("groundingChunks", [])): + web = chunk.get("web") + if not isinstance(web, dict): + continue + url = web.get("uri") + if not isinstance(url, str) or not url: + continue + source_indices[index] = len(sources) + sources.append( + GroundingSource( + url=url, + title=web.get("title"), + domain=web.get("domain"), + ) + ) + + citations: list[Citation] = [] + for support in grounding_metadata.get("groundingSupports", []): + segment = support.get("segment", {}) + if not isinstance(segment, dict): + continue + start = segment.get("startIndex") + end = segment.get("endIndex") + if not isinstance(start, int) or not isinstance(end, int): + continue + start = _byte_offset_to_char_offset(text, start) + end = _byte_offset_to_char_offset(text, end) + if start is None or end is None: + continue + raw_indices = support.get("groundingChunkIndices", []) + indices = [ + source_indices[i] + for i in raw_indices + if isinstance(i, int) and i in source_indices + ] + citations.append( + Citation( + start=start, + end=end, + source_indices=indices, + cited_text=segment.get("text"), + ) + ) + + entry = grounding_metadata.get("searchEntryPoint") or {} + return Grounding( + queries=grounding_metadata.get("webSearchQueries", []), + sources=sources, + citations=citations, + search_entry_point=entry.get("renderedContent"), + ) + + +__all__ = ["map_grounding"] diff --git a/src/celeste/modalities/text/providers/groq/client.py b/src/celeste/modalities/text/providers/groq/client.py index e1ea12d..0e92def 100644 --- a/src/celeste/modalities/text/providers/groq/client.py +++ b/src/celeste/modalities/text/providers/groq/client.py @@ -1,8 +1,12 @@ """Groq text client (modality).""" +from typing import Any + +from celeste.grounding import Grounding, GroundingSource from celeste.parameters import ParameterMapper from celeste.providers.groq.chat.client import GroqChatClient from celeste.providers.groq.chat.streaming import GroqChatStream as _GroqChatStream +from celeste.providers.groq.chat.tools import parse_search_results from celeste.types import TextContent from ...protocols.chatcompletions.client import ( @@ -26,6 +30,23 @@ class GroqTextClient(GroqChatClient, ChatCompletionsTextClient): def parameter_mappers(cls) -> list[ParameterMapper[TextContent]]: return GROQ_PARAMETER_MAPPERS + def _parse_grounding(self, response_data: dict[str, Any]) -> Grounding | None: + """Parse Groq Compound search results into grounding sources.""" + sources: list[GroundingSource] = [] + seen: set[tuple[str, str | None]] = set() + for result in parse_search_results(response_data): + url = result.get("url") + if not isinstance(url, str): + continue + title = result.get("title") + key = (url, title if isinstance(title, str) else None) + if key in seen: + continue + seen.add(key) + sources.append(GroundingSource(url=url, title=key[1])) + + return Grounding(sources=sources) if sources else None + def _stream_class(self) -> type[TextStream]: """Return the Stream class for this provider.""" return GroqTextStream diff --git a/src/celeste/protocols/chatcompletions/streaming.py b/src/celeste/protocols/chatcompletions/streaming.py index 38d23c4..db54e0d 100644 --- a/src/celeste/protocols/chatcompletions/streaming.py +++ b/src/celeste/protocols/chatcompletions/streaming.py @@ -30,6 +30,7 @@ class ChatCompletionsStream: def __init__(self, *args: Any, **kwargs: Any) -> None: # noqa: ANN401 super().__init__(*args, **kwargs) self._tool_call_deltas: dict[int, dict[str, Any]] = {} + self._annotations: list[dict[str, Any]] = [] def _get_chunk_delta(self, event_data: dict[str, Any]) -> dict[str, Any] | None: """Extract delta dict from a chat.completion.chunk event.""" @@ -88,11 +89,18 @@ def _parse_chunk_finish_reason( return None def _parse_chunk(self, event_data: dict[str, Any]) -> Any: # noqa: ANN401 - """Capture tool_call deltas before delegating to base _parse_chunk.""" + """Capture native delta side data before delegating to base _parse_chunk.""" choices = event_data.get("choices", []) if choices and isinstance(choices[0], dict): delta = choices[0].get("delta", {}) if isinstance(delta, dict): + annotations = delta.get("annotations") + if isinstance(annotations, list): + self._annotations.extend( + annotation + for annotation in annotations + if isinstance(annotation, dict) + ) for tc_delta in delta.get("tool_calls") or []: idx = tc_delta.get("index", 0) if idx not in self._tool_call_deltas: diff --git a/src/celeste/protocols/chatcompletions/tools.py b/src/celeste/protocols/chatcompletions/tools.py index 0878ea7..3de543a 100644 --- a/src/celeste/protocols/chatcompletions/tools.py +++ b/src/celeste/protocols/chatcompletions/tools.py @@ -41,7 +41,24 @@ def parse_tool_calls( return tool_calls +def parse_annotations(response_data: dict[str, Any]) -> list[dict[str, Any]]: + """Extract citation/search annotations from Chat Completions messages.""" + annotations: list[dict[str, Any]] = [] + for choice in response_data.get("choices", []): + message = choice.get("message", {}) + message_annotations = message.get("annotations") + if not isinstance(message_annotations, list): + continue + annotations.extend( + annotation + for annotation in message_annotations + if isinstance(annotation, dict) + ) + return annotations + + __all__ = [ "TOOL_MAPPERS", + "parse_annotations", "parse_tool_calls", ] diff --git a/src/celeste/protocols/openresponses/tools.py b/src/celeste/protocols/openresponses/tools.py index 93b15d3..6317fa4 100644 --- a/src/celeste/protocols/openresponses/tools.py +++ b/src/celeste/protocols/openresponses/tools.py @@ -119,11 +119,27 @@ def parse_reasoning( return text, signature_blocks +def parse_annotations(output: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Extract URL citation annotations from Responses API output items.""" + annotations: list[dict[str, Any]] = [] + for item in output: + if item.get("type") != "message": + continue + for part in item.get("content", []): + if part.get("type") != "output_text": + continue + for annotation in part.get("annotations", []): + if isinstance(annotation, dict): + annotations.append(annotation) + return annotations + + __all__ = [ "TOOL_MAPPERS", "CodeExecutionMapper", "WebSearchMapper", "XSearchMapper", + "parse_annotations", "parse_content", "parse_reasoning", "parse_tool_calls", diff --git a/src/celeste/providers/anthropic/messages/streaming.py b/src/celeste/providers/anthropic/messages/streaming.py index 62b411a..7da9695 100644 --- a/src/celeste/providers/anthropic/messages/streaming.py +++ b/src/celeste/providers/anthropic/messages/streaming.py @@ -1,5 +1,7 @@ """Anthropic Messages SSE parsing for streaming.""" +import contextlib +import json from typing import Any from celeste.io import FinishReason @@ -18,6 +20,74 @@ class AnthropicMessagesStream: Modality streams call super() methods which resolve to this via MRO. """ + def __init__(self, *args: Any, **kwargs: Any) -> None: # noqa: ANN401 + super().__init__(*args, **kwargs) + self._content_blocks: dict[int, dict[str, Any]] = {} + + def _parse_chunk(self, event_data: dict[str, Any]) -> Any: # noqa: ANN401 + """Capture native content blocks before delegating to the modality stream.""" + event_type = event_data.get("type") + if event_type == "content_block_start": + block = event_data.get("content_block", {}) + block_type = block.get("type") + idx = event_data.get("index", len(self._content_blocks)) + if block_type == "server_tool_use": + self._content_blocks[idx] = { + "type": "server_tool_use", + "id": block.get("id", ""), + "name": block.get("name", ""), + "input_json": "", + } + elif block_type == "web_search_tool_result": + self._content_blocks[idx] = block + elif block_type == "text": + self._content_blocks[idx] = { + "type": "text", + "text": block.get("text", ""), + "citations": [], + } + elif event_type == "content_block_delta": + delta = event_data.get("delta", {}) + delta_type = delta.get("type") + idx = event_data.get("index", -1) + block = self._content_blocks.get(idx) + if delta_type == "input_json_delta": + if block and block.get("type") == "server_tool_use": + block["input_json"] += delta.get("partial_json", "") + elif delta_type in {"text_delta", "citations_delta"}: + block = self._content_blocks.setdefault( + idx, {"type": "text", "text": "", "citations": []} + ) + if delta_type == "text_delta": + block["text"] += delta.get("text", "") + else: + citation = delta.get("citation") + if isinstance(citation, dict): + block["citations"].append(citation) + return super()._parse_chunk(event_data) # type: ignore[misc] + + def _aggregate_content_blocks(self) -> list[dict[str, Any]]: + """Return reconstructed native Anthropic content blocks.""" + blocks: list[dict[str, Any]] = [] + for idx in sorted(self._content_blocks): + block = self._content_blocks[idx] + if block.get("type") == "server_tool_use": + input_data: dict[str, Any] = {} + if block["input_json"]: + with contextlib.suppress(ValueError, TypeError): + input_data = json.loads(block["input_json"]) + blocks.append( + { + "type": "server_tool_use", + "id": block["id"], + "name": block["name"], + "input": input_data, + } + ) + else: + blocks.append(block) + return blocks + def _parse_chunk_content(self, event_data: dict[str, Any]) -> str | None: """Extract content from SSE event. diff --git a/src/celeste/providers/google/generate_content/grounding.py b/src/celeste/providers/google/generate_content/grounding.py new file mode 100644 index 0000000..cbbfb74 --- /dev/null +++ b/src/celeste/providers/google/generate_content/grounding.py @@ -0,0 +1,28 @@ +"""Google GenerateContent native grounding helpers.""" + +from typing import Any + + +def parse_grounding_metadata(response_data: dict[str, Any]) -> dict[str, Any] | None: + """Extract Gemini groundingMetadata from the first candidate.""" + candidates = response_data.get("candidates", []) + meta = candidates[0].get("groundingMetadata") if candidates else None + return meta if isinstance(meta, dict) else None + + +def merge_grounding_metadata(metadata: list[dict[str, Any]]) -> dict[str, Any]: + """Merge streamed Gemini grounding metadata chunks.""" + combined: dict[str, Any] = { + "webSearchQueries": [], + "groundingChunks": [], + "groundingSupports": [], + } + for meta in metadata: + for key in ("webSearchQueries", "groundingChunks", "groundingSupports"): + combined[key].extend(meta.get(key, [])) + if meta.get("searchEntryPoint"): + combined["searchEntryPoint"] = meta["searchEntryPoint"] + return combined + + +__all__ = ["merge_grounding_metadata", "parse_grounding_metadata"] diff --git a/src/celeste/providers/google/generate_content/streaming.py b/src/celeste/providers/google/generate_content/streaming.py index 3f20d0b..a25303d 100644 --- a/src/celeste/providers/google/generate_content/streaming.py +++ b/src/celeste/providers/google/generate_content/streaming.py @@ -19,6 +19,17 @@ class GoogleGenerateContentStream: """ _error_type_fields: ClassVar[tuple[str, ...]] = ("status", "code") + _grounding_metadata: list[dict[str, Any]] + + def _parse_chunk(self, event_data: dict[str, Any]) -> Any | None: # noqa: ANN401 + """Capture grounding metadata before normal chunk filtering.""" + if not hasattr(self, "_grounding_metadata"): + self._grounding_metadata = [] + for candidate in event_data.get("candidates", []): + meta = candidate.get("groundingMetadata") + if isinstance(meta, dict): + self._grounding_metadata.append(meta) + return super()._parse_chunk(event_data) # type: ignore[misc] def _parse_chunk_content(self, event_data: dict[str, Any]) -> str | None: """Extract content from SSE event.""" diff --git a/src/celeste/providers/groq/chat/tools.py b/src/celeste/providers/groq/chat/tools.py index ed5e758..03447f6 100644 --- a/src/celeste/providers/groq/chat/tools.py +++ b/src/celeste/providers/groq/chat/tools.py @@ -37,4 +37,35 @@ def map_tool(self, tool: Tool) -> dict[str, Any]: TOOL_MAPPERS: list[ToolMapper] = [WebSearchMapper(), CodeExecutionMapper()] -__all__ = ["TOOL_MAPPERS", "CodeExecutionMapper", "WebSearchMapper"] + +def parse_search_results(response_data: dict[str, Any]) -> list[dict[str, Any]]: + """Extract Groq Compound native search result rows.""" + choices = response_data.get("choices", []) + if not choices: + return [] + + message = choices[0].get("message", {}) + executed_tools = message.get("executed_tools") + if not isinstance(executed_tools, list): + return [] + + results: list[dict[str, Any]] = [] + for tool in executed_tools: + if not isinstance(tool, dict): + continue + search_results = tool.get("search_results") + if not isinstance(search_results, dict): + continue + rows = search_results.get("results") + if not isinstance(rows, list): + continue + results.extend(row for row in rows if isinstance(row, dict)) + return results + + +__all__ = [ + "TOOL_MAPPERS", + "CodeExecutionMapper", + "WebSearchMapper", + "parse_search_results", +] diff --git a/src/celeste/streaming.py b/src/celeste/streaming.py index 59a267b..e0ab5a0 100644 --- a/src/celeste/streaming.py +++ b/src/celeste/streaming.py @@ -10,6 +10,7 @@ from anyio.from_thread import start_blocking_portal from celeste.exceptions import StreamEventError, StreamNotExhaustedError +from celeste.grounding import Grounding from celeste.io import Chunk as ChunkBase from celeste.io import FinishReason, Output, Usage from celeste.parameters import Parameters @@ -175,6 +176,9 @@ def _parse_output(self, chunks: list[Chunk], **parameters: Unpack[Params]) -> Ou kwargs["reasoning"] = reasoning if signature: kwargs["signature"] = signature + grounding = self._aggregate_grounding(chunks, raw_events) + if grounding is not None: + kwargs["grounding"] = grounding tool_calls = validate_tool_calls( self._aggregate_tool_calls(chunks, raw_events), parameters.get("tools"), @@ -195,6 +199,12 @@ def _aggregate_tool_calls( """Aggregate tool calls from stream events. Override in providers that support tools.""" return [] + def _aggregate_grounding( + self, chunks: list[Chunk], raw_events: list[dict[str, Any]] + ) -> Grounding | None: + """Aggregate web-search grounding from stream events.""" + return None + def _aggregate_reasoning(self, chunks: list[Chunk]) -> str | None: """Aggregate reasoning from chunks. Override in modality streams.""" return None