Skip to content

Commit 19f32f9

Browse files
committed
add builtin tool call/return parts for Google file search
Implemented _map_file_search_grounding_metadata following the same pattern as web search. Extracts retrieval_queries and retrieved_context from grounding_metadata to create proper BuiltinToolCallPart and BuiltinToolReturnPart instances.
1 parent b3a8930 commit 19f32f9

File tree

1 file changed

+30
-0
lines changed

1 file changed

+30
-0
lines changed

pydantic_ai_slim/pydantic_ai/models/google.py

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -812,6 +812,11 @@ def _process_response_from_parts(
812812
items.append(web_search_call)
813813
items.append(web_search_return)
814814

815+
file_search_call, file_search_return = _map_file_search_grounding_metadata(grounding_metadata, provider_name)
816+
if file_search_call and file_search_return:
817+
items.append(file_search_call)
818+
items.append(file_search_return)
819+
815820
item: ModelResponsePart | None = None
816821
code_execution_tool_call_id: str | None = None
817822
for part in parts:
@@ -972,3 +977,28 @@ def _map_grounding_metadata(
972977
)
973978
else:
974979
return None, None
980+
981+
982+
def _map_file_search_grounding_metadata(
983+
grounding_metadata: GroundingMetadata | None, provider_name: str
984+
) -> tuple[BuiltinToolCallPart, BuiltinToolReturnPart] | tuple[None, None]:
985+
if grounding_metadata and (retrieval_queries := grounding_metadata.retrieval_queries):
986+
tool_call_id = _utils.generate_tool_call_id()
987+
return (
988+
BuiltinToolCallPart(
989+
provider_name=provider_name,
990+
tool_name=FileSearchTool.kind,
991+
tool_call_id=tool_call_id,
992+
args={'queries': retrieval_queries},
993+
),
994+
BuiltinToolReturnPart(
995+
provider_name=provider_name,
996+
tool_name=FileSearchTool.kind,
997+
tool_call_id=tool_call_id,
998+
content=[chunk.retrieved_context.model_dump(mode='json') for chunk in grounding_chunks if chunk.retrieved_context]
999+
if (grounding_chunks := grounding_metadata.grounding_chunks)
1000+
else None,
1001+
),
1002+
)
1003+
else:
1004+
return None, None

0 commit comments

Comments
 (0)