From 8c6cbc1ae43e8b8da51af6e274b79d16c80d4b3d Mon Sep 17 00:00:00 2001 From: Annie Luc Date: Thu, 4 Jun 2026 11:55:47 -0700 Subject: [PATCH] chore: Internal change PiperOrigin-RevId: 926810002 --- google/genai/_interactions/types/tool.py | 129 +++++++++++++++++- .../genai/_interactions/types/tool_param.py | 125 ++++++++++++++++- 2 files changed, 252 insertions(+), 2 deletions(-) diff --git a/google/genai/_interactions/types/tool.py b/google/genai/_interactions/types/tool.py index 66d412597..9e29c6fde 100644 --- a/google/genai/_interactions/types/tool.py +++ b/google/genai/_interactions/types/tool.py @@ -18,6 +18,8 @@ from typing import Dict, List, Union, Optional from typing_extensions import Literal, Annotated, TypeAlias +from pydantic import Field as FieldInfo + from .._utils import PropertyInfo from .._models import BaseModel from .function import Function @@ -33,6 +35,14 @@ "FileSearch", "GoogleMaps", "Retrieval", + "RetrievalExaAISearchConfig", + "RetrievalParallelAISearchConfig", + "RetrievalRagStoreConfig", + "RetrievalRagStoreConfigRagResource", + "RetrievalRagStoreConfigRagRetrievalConfig", + "RetrievalRagStoreConfigRagRetrievalConfigFilter", + "RetrievalRagStoreConfigRagRetrievalConfigHybridSearch", + "RetrievalRagStoreConfigRagRetrievalConfigRanking", "RetrievalVertexAISearchConfig", ] @@ -121,6 +131,112 @@ class GoogleMaps(BaseModel): """The longitude of the user's location.""" +class RetrievalExaAISearchConfig(BaseModel): + """Used to specify configuration for ExaAISearch.""" + + api_key: str + """Required. The API key for ExaAiSearch.""" + + custom_config: Optional[Dict[str, object]] = None + """Optional. + + This field can be used to pass any parameter from the Exa.ai Search API. + """ + + +class RetrievalParallelAISearchConfig(BaseModel): + """Used to specify configuration for ParallelAISearch.""" + + api_key: Optional[str] = None + """Optional. The API key for ParallelAiSearch.""" + + custom_config: Optional[Dict[str, object]] = None + """Optional. Custom configs for ParallelAiSearch.""" + + +class RetrievalRagStoreConfigRagResource(BaseModel): + """The definition of the Rag resource.""" + + rag_corpus: Optional[str] = None + """Optional. RagCorpora resource name.""" + + rag_file_ids: Optional[List[str]] = None + """Optional. + + rag_file_id. The files should be in the same rag_corpus set in rag_corpus field. + """ + + +class RetrievalRagStoreConfigRagRetrievalConfigFilter(BaseModel): + """Optional. Config for filters.""" + + metadata_filter: Optional[str] = None + """Optional. String for metadata filtering.""" + + vector_distance_threshold: Optional[float] = None + """Optional. + + Only returns contexts with vector distance smaller than the threshold. + """ + + vector_similarity_threshold: Optional[float] = None + """Optional. + + Only returns contexts with vector similarity larger than the threshold. + """ + + +class RetrievalRagStoreConfigRagRetrievalConfigHybridSearch(BaseModel): + """Optional. Config for Hybrid Search.""" + + alpha: Optional[float] = None + """Optional. + + Alpha value controls the weight between dense and sparse vector search results. + """ + + +class RetrievalRagStoreConfigRagRetrievalConfigRanking(BaseModel): + """Optional. Config for ranking and reranking.""" + + ranking_config: Literal["rank_service"] + + api_model_name: Optional[str] = FieldInfo(alias="model_name", default=None) + """Optional. The model name of the rank service.""" + + +class RetrievalRagStoreConfigRagRetrievalConfig(BaseModel): + """Optional. The retrieval config for the Rag query.""" + + filter: Optional[RetrievalRagStoreConfigRagRetrievalConfigFilter] = None + """Optional. Config for filters.""" + + hybrid_search: Optional[RetrievalRagStoreConfigRagRetrievalConfigHybridSearch] = None + """Optional. Config for Hybrid Search.""" + + ranking: Optional[RetrievalRagStoreConfigRagRetrievalConfigRanking] = None + """Optional. Config for ranking and reranking.""" + + top_k: Optional[int] = None + """Optional. The number of contexts to retrieve.""" + + +class RetrievalRagStoreConfig(BaseModel): + """Used to specify configuration for RagStore.""" + + rag_resources: Optional[List[RetrievalRagStoreConfigRagResource]] = None + """Optional. The representation of the rag source.""" + + rag_retrieval_config: Optional[RetrievalRagStoreConfigRagRetrievalConfig] = None + """Optional. The retrieval config for the Rag query.""" + + similarity_top_k: Optional[int] = None + """Optional. Number of top k results to return from the selected corpora.""" + + vector_distance_threshold: Optional[float] = None + """Optional. Only return results with vector distance smaller than the threshold.""" + + class RetrievalVertexAISearchConfig(BaseModel): """Used to specify configuration for VertexAISearch.""" @@ -136,7 +252,18 @@ class Retrieval(BaseModel): type: Literal["retrieval"] - retrieval_types: Optional[List[Literal["vertex_ai_search"]]] = None + exa_ai_search_config: Optional[RetrievalExaAISearchConfig] = None + """Used to specify configuration for ExaAISearch.""" + + parallel_ai_search_config: Optional[RetrievalParallelAISearchConfig] = None + """Used to specify configuration for ParallelAISearch.""" + + rag_store_config: Optional[RetrievalRagStoreConfig] = None + """Used to specify configuration for RagStore.""" + + retrieval_types: Optional[List[Literal["vertex_ai_search", "rag_store", "exa_ai_search", "parallel_ai_search"]]] = ( + None + ) """The types of file retrieval to enable.""" vertex_ai_search_config: Optional[RetrievalVertexAISearchConfig] = None diff --git a/google/genai/_interactions/types/tool_param.py b/google/genai/_interactions/types/tool_param.py index bc0463000..9042a481a 100644 --- a/google/genai/_interactions/types/tool_param.py +++ b/google/genai/_interactions/types/tool_param.py @@ -34,6 +34,14 @@ "FileSearch", "GoogleMaps", "Retrieval", + "RetrievalExaAISearchConfig", + "RetrievalParallelAISearchConfig", + "RetrievalRagStoreConfig", + "RetrievalRagStoreConfigRagResource", + "RetrievalRagStoreConfigRagRetrievalConfig", + "RetrievalRagStoreConfigRagRetrievalConfigFilter", + "RetrievalRagStoreConfigRagRetrievalConfigHybridSearch", + "RetrievalRagStoreConfigRagRetrievalConfigRanking", "RetrievalVertexAISearchConfig", ] @@ -122,6 +130,112 @@ class GoogleMaps(TypedDict, total=False): """The longitude of the user's location.""" +class RetrievalExaAISearchConfig(TypedDict, total=False): + """Used to specify configuration for ExaAISearch.""" + + api_key: Required[str] + """Required. The API key for ExaAiSearch.""" + + custom_config: Dict[str, object] + """Optional. + + This field can be used to pass any parameter from the Exa.ai Search API. + """ + + +class RetrievalParallelAISearchConfig(TypedDict, total=False): + """Used to specify configuration for ParallelAISearch.""" + + api_key: str + """Optional. The API key for ParallelAiSearch.""" + + custom_config: Dict[str, object] + """Optional. Custom configs for ParallelAiSearch.""" + + +class RetrievalRagStoreConfigRagResource(TypedDict, total=False): + """The definition of the Rag resource.""" + + rag_corpus: str + """Optional. RagCorpora resource name.""" + + rag_file_ids: SequenceNotStr[str] + """Optional. + + rag_file_id. The files should be in the same rag_corpus set in rag_corpus field. + """ + + +class RetrievalRagStoreConfigRagRetrievalConfigFilter(TypedDict, total=False): + """Optional. Config for filters.""" + + metadata_filter: str + """Optional. String for metadata filtering.""" + + vector_distance_threshold: float + """Optional. + + Only returns contexts with vector distance smaller than the threshold. + """ + + vector_similarity_threshold: float + """Optional. + + Only returns contexts with vector similarity larger than the threshold. + """ + + +class RetrievalRagStoreConfigRagRetrievalConfigHybridSearch(TypedDict, total=False): + """Optional. Config for Hybrid Search.""" + + alpha: float + """Optional. + + Alpha value controls the weight between dense and sparse vector search results. + """ + + +class RetrievalRagStoreConfigRagRetrievalConfigRanking(TypedDict, total=False): + """Optional. Config for ranking and reranking.""" + + ranking_config: Required[Literal["rank_service"]] + + model_name: str + """Optional. The model name of the rank service.""" + + +class RetrievalRagStoreConfigRagRetrievalConfig(TypedDict, total=False): + """Optional. The retrieval config for the Rag query.""" + + filter: RetrievalRagStoreConfigRagRetrievalConfigFilter + """Optional. Config for filters.""" + + hybrid_search: RetrievalRagStoreConfigRagRetrievalConfigHybridSearch + """Optional. Config for Hybrid Search.""" + + ranking: RetrievalRagStoreConfigRagRetrievalConfigRanking + """Optional. Config for ranking and reranking.""" + + top_k: int + """Optional. The number of contexts to retrieve.""" + + +class RetrievalRagStoreConfig(TypedDict, total=False): + """Used to specify configuration for RagStore.""" + + rag_resources: Iterable[RetrievalRagStoreConfigRagResource] + """Optional. The representation of the rag source.""" + + rag_retrieval_config: RetrievalRagStoreConfigRagRetrievalConfig + """Optional. The retrieval config for the Rag query.""" + + similarity_top_k: int + """Optional. Number of top k results to return from the selected corpora.""" + + vector_distance_threshold: float + """Optional. Only return results with vector distance smaller than the threshold.""" + + class RetrievalVertexAISearchConfig(TypedDict, total=False): """Used to specify configuration for VertexAISearch.""" @@ -137,7 +251,16 @@ class Retrieval(TypedDict, total=False): type: Required[Literal["retrieval"]] - retrieval_types: List[Literal["vertex_ai_search"]] + exa_ai_search_config: RetrievalExaAISearchConfig + """Used to specify configuration for ExaAISearch.""" + + parallel_ai_search_config: RetrievalParallelAISearchConfig + """Used to specify configuration for ParallelAISearch.""" + + rag_store_config: RetrievalRagStoreConfig + """Used to specify configuration for RagStore.""" + + retrieval_types: List[Literal["vertex_ai_search", "rag_store", "exa_ai_search", "parallel_ai_search"]] """The types of file retrieval to enable.""" vertex_ai_search_config: RetrievalVertexAISearchConfig