agno 2.2.13__py3-none-any.whl → 2.4.3__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- agno/agent/__init__.py +6 -0
- agno/agent/agent.py +5252 -3145
- agno/agent/remote.py +525 -0
- agno/api/api.py +2 -0
- agno/client/__init__.py +3 -0
- agno/client/a2a/__init__.py +10 -0
- agno/client/a2a/client.py +554 -0
- agno/client/a2a/schemas.py +112 -0
- agno/client/a2a/utils.py +369 -0
- agno/client/os.py +2669 -0
- agno/compression/__init__.py +3 -0
- agno/compression/manager.py +247 -0
- agno/culture/manager.py +2 -2
- agno/db/base.py +927 -6
- agno/db/dynamo/dynamo.py +788 -2
- agno/db/dynamo/schemas.py +128 -0
- agno/db/dynamo/utils.py +26 -3
- agno/db/firestore/firestore.py +674 -50
- agno/db/firestore/schemas.py +41 -0
- agno/db/firestore/utils.py +25 -10
- agno/db/gcs_json/gcs_json_db.py +506 -3
- agno/db/gcs_json/utils.py +14 -2
- agno/db/in_memory/in_memory_db.py +203 -4
- agno/db/in_memory/utils.py +14 -2
- agno/db/json/json_db.py +498 -2
- agno/db/json/utils.py +14 -2
- agno/db/migrations/manager.py +199 -0
- agno/db/migrations/utils.py +19 -0
- agno/db/migrations/v1_to_v2.py +54 -16
- agno/db/migrations/versions/__init__.py +0 -0
- agno/db/migrations/versions/v2_3_0.py +977 -0
- agno/db/mongo/async_mongo.py +1013 -39
- agno/db/mongo/mongo.py +684 -4
- agno/db/mongo/schemas.py +48 -0
- agno/db/mongo/utils.py +17 -0
- agno/db/mysql/__init__.py +2 -1
- agno/db/mysql/async_mysql.py +2958 -0
- agno/db/mysql/mysql.py +722 -53
- agno/db/mysql/schemas.py +77 -11
- agno/db/mysql/utils.py +151 -8
- agno/db/postgres/async_postgres.py +1254 -137
- agno/db/postgres/postgres.py +2316 -93
- agno/db/postgres/schemas.py +153 -21
- agno/db/postgres/utils.py +22 -7
- agno/db/redis/redis.py +531 -3
- agno/db/redis/schemas.py +36 -0
- agno/db/redis/utils.py +31 -15
- agno/db/schemas/evals.py +1 -0
- agno/db/schemas/memory.py +20 -9
- agno/db/singlestore/schemas.py +70 -1
- agno/db/singlestore/singlestore.py +737 -74
- agno/db/singlestore/utils.py +13 -3
- agno/db/sqlite/async_sqlite.py +1069 -89
- agno/db/sqlite/schemas.py +133 -1
- agno/db/sqlite/sqlite.py +2203 -165
- agno/db/sqlite/utils.py +21 -11
- agno/db/surrealdb/models.py +25 -0
- agno/db/surrealdb/surrealdb.py +603 -1
- agno/db/utils.py +60 -0
- agno/eval/__init__.py +26 -3
- agno/eval/accuracy.py +25 -12
- agno/eval/agent_as_judge.py +871 -0
- agno/eval/base.py +29 -0
- agno/eval/performance.py +10 -4
- agno/eval/reliability.py +22 -13
- agno/eval/utils.py +2 -1
- agno/exceptions.py +42 -0
- agno/hooks/__init__.py +3 -0
- agno/hooks/decorator.py +164 -0
- agno/integrations/discord/client.py +13 -2
- agno/knowledge/__init__.py +4 -0
- agno/knowledge/chunking/code.py +90 -0
- agno/knowledge/chunking/document.py +65 -4
- agno/knowledge/chunking/fixed.py +4 -1
- agno/knowledge/chunking/markdown.py +102 -11
- agno/knowledge/chunking/recursive.py +2 -2
- agno/knowledge/chunking/semantic.py +130 -48
- agno/knowledge/chunking/strategy.py +18 -0
- agno/knowledge/embedder/azure_openai.py +0 -1
- agno/knowledge/embedder/google.py +1 -1
- agno/knowledge/embedder/mistral.py +1 -1
- agno/knowledge/embedder/nebius.py +1 -1
- agno/knowledge/embedder/openai.py +16 -12
- agno/knowledge/filesystem.py +412 -0
- agno/knowledge/knowledge.py +4261 -1199
- agno/knowledge/protocol.py +134 -0
- agno/knowledge/reader/arxiv_reader.py +3 -2
- agno/knowledge/reader/base.py +9 -7
- agno/knowledge/reader/csv_reader.py +91 -42
- agno/knowledge/reader/docx_reader.py +9 -10
- agno/knowledge/reader/excel_reader.py +225 -0
- agno/knowledge/reader/field_labeled_csv_reader.py +38 -48
- agno/knowledge/reader/firecrawl_reader.py +3 -2
- agno/knowledge/reader/json_reader.py +16 -22
- agno/knowledge/reader/markdown_reader.py +15 -14
- agno/knowledge/reader/pdf_reader.py +33 -28
- agno/knowledge/reader/pptx_reader.py +9 -10
- agno/knowledge/reader/reader_factory.py +135 -1
- agno/knowledge/reader/s3_reader.py +8 -16
- agno/knowledge/reader/tavily_reader.py +3 -3
- agno/knowledge/reader/text_reader.py +15 -14
- agno/knowledge/reader/utils/__init__.py +17 -0
- agno/knowledge/reader/utils/spreadsheet.py +114 -0
- agno/knowledge/reader/web_search_reader.py +8 -65
- agno/knowledge/reader/website_reader.py +16 -13
- agno/knowledge/reader/wikipedia_reader.py +36 -3
- agno/knowledge/reader/youtube_reader.py +3 -2
- agno/knowledge/remote_content/__init__.py +33 -0
- agno/knowledge/remote_content/config.py +266 -0
- agno/knowledge/remote_content/remote_content.py +105 -17
- agno/knowledge/utils.py +76 -22
- agno/learn/__init__.py +71 -0
- agno/learn/config.py +463 -0
- agno/learn/curate.py +185 -0
- agno/learn/machine.py +725 -0
- agno/learn/schemas.py +1114 -0
- agno/learn/stores/__init__.py +38 -0
- agno/learn/stores/decision_log.py +1156 -0
- agno/learn/stores/entity_memory.py +3275 -0
- agno/learn/stores/learned_knowledge.py +1583 -0
- agno/learn/stores/protocol.py +117 -0
- agno/learn/stores/session_context.py +1217 -0
- agno/learn/stores/user_memory.py +1495 -0
- agno/learn/stores/user_profile.py +1220 -0
- agno/learn/utils.py +209 -0
- agno/media.py +22 -6
- agno/memory/__init__.py +14 -1
- agno/memory/manager.py +223 -8
- agno/memory/strategies/__init__.py +15 -0
- agno/memory/strategies/base.py +66 -0
- agno/memory/strategies/summarize.py +196 -0
- agno/memory/strategies/types.py +37 -0
- agno/models/aimlapi/aimlapi.py +17 -0
- agno/models/anthropic/claude.py +434 -59
- agno/models/aws/bedrock.py +121 -20
- agno/models/aws/claude.py +131 -274
- agno/models/azure/ai_foundry.py +10 -6
- agno/models/azure/openai_chat.py +33 -10
- agno/models/base.py +1162 -561
- agno/models/cerebras/cerebras.py +120 -24
- agno/models/cerebras/cerebras_openai.py +21 -2
- agno/models/cohere/chat.py +65 -6
- agno/models/cometapi/cometapi.py +18 -1
- agno/models/dashscope/dashscope.py +2 -3
- agno/models/deepinfra/deepinfra.py +18 -1
- agno/models/deepseek/deepseek.py +69 -3
- agno/models/fireworks/fireworks.py +18 -1
- agno/models/google/gemini.py +959 -89
- agno/models/google/utils.py +22 -0
- agno/models/groq/groq.py +48 -18
- agno/models/huggingface/huggingface.py +17 -6
- agno/models/ibm/watsonx.py +16 -6
- agno/models/internlm/internlm.py +18 -1
- agno/models/langdb/langdb.py +13 -1
- agno/models/litellm/chat.py +88 -9
- agno/models/litellm/litellm_openai.py +18 -1
- agno/models/message.py +24 -5
- agno/models/meta/llama.py +40 -13
- agno/models/meta/llama_openai.py +22 -21
- agno/models/metrics.py +12 -0
- agno/models/mistral/mistral.py +8 -4
- agno/models/n1n/__init__.py +3 -0
- agno/models/n1n/n1n.py +57 -0
- agno/models/nebius/nebius.py +6 -7
- agno/models/nvidia/nvidia.py +20 -3
- agno/models/ollama/__init__.py +2 -0
- agno/models/ollama/chat.py +17 -6
- agno/models/ollama/responses.py +100 -0
- agno/models/openai/__init__.py +2 -0
- agno/models/openai/chat.py +117 -26
- agno/models/openai/open_responses.py +46 -0
- agno/models/openai/responses.py +110 -32
- agno/models/openrouter/__init__.py +2 -0
- agno/models/openrouter/openrouter.py +67 -2
- agno/models/openrouter/responses.py +146 -0
- agno/models/perplexity/perplexity.py +19 -1
- agno/models/portkey/portkey.py +7 -6
- agno/models/requesty/requesty.py +19 -2
- agno/models/response.py +20 -2
- agno/models/sambanova/sambanova.py +20 -3
- agno/models/siliconflow/siliconflow.py +19 -2
- agno/models/together/together.py +20 -3
- agno/models/vercel/v0.py +20 -3
- agno/models/vertexai/claude.py +124 -4
- agno/models/vllm/vllm.py +19 -14
- agno/models/xai/xai.py +19 -2
- agno/os/app.py +467 -137
- agno/os/auth.py +253 -5
- agno/os/config.py +22 -0
- agno/os/interfaces/a2a/a2a.py +7 -6
- agno/os/interfaces/a2a/router.py +635 -26
- agno/os/interfaces/a2a/utils.py +32 -33
- agno/os/interfaces/agui/agui.py +5 -3
- agno/os/interfaces/agui/router.py +26 -16
- agno/os/interfaces/agui/utils.py +97 -57
- agno/os/interfaces/base.py +7 -7
- agno/os/interfaces/slack/router.py +16 -7
- agno/os/interfaces/slack/slack.py +7 -7
- agno/os/interfaces/whatsapp/router.py +35 -7
- agno/os/interfaces/whatsapp/security.py +3 -1
- agno/os/interfaces/whatsapp/whatsapp.py +11 -8
- agno/os/managers.py +326 -0
- agno/os/mcp.py +652 -79
- agno/os/middleware/__init__.py +4 -0
- agno/os/middleware/jwt.py +718 -115
- agno/os/middleware/trailing_slash.py +27 -0
- agno/os/router.py +105 -1558
- agno/os/routers/agents/__init__.py +3 -0
- agno/os/routers/agents/router.py +655 -0
- agno/os/routers/agents/schema.py +288 -0
- agno/os/routers/components/__init__.py +3 -0
- agno/os/routers/components/components.py +475 -0
- agno/os/routers/database.py +155 -0
- agno/os/routers/evals/evals.py +111 -18
- agno/os/routers/evals/schemas.py +38 -5
- agno/os/routers/evals/utils.py +80 -11
- agno/os/routers/health.py +3 -3
- agno/os/routers/knowledge/knowledge.py +284 -35
- agno/os/routers/knowledge/schemas.py +14 -2
- agno/os/routers/memory/memory.py +274 -11
- agno/os/routers/memory/schemas.py +44 -3
- agno/os/routers/metrics/metrics.py +30 -15
- agno/os/routers/metrics/schemas.py +10 -6
- agno/os/routers/registry/__init__.py +3 -0
- agno/os/routers/registry/registry.py +337 -0
- agno/os/routers/session/session.py +143 -14
- agno/os/routers/teams/__init__.py +3 -0
- agno/os/routers/teams/router.py +550 -0
- agno/os/routers/teams/schema.py +280 -0
- agno/os/routers/traces/__init__.py +3 -0
- agno/os/routers/traces/schemas.py +414 -0
- agno/os/routers/traces/traces.py +549 -0
- agno/os/routers/workflows/__init__.py +3 -0
- agno/os/routers/workflows/router.py +757 -0
- agno/os/routers/workflows/schema.py +139 -0
- agno/os/schema.py +157 -584
- agno/os/scopes.py +469 -0
- agno/os/settings.py +3 -0
- agno/os/utils.py +574 -185
- agno/reasoning/anthropic.py +85 -1
- agno/reasoning/azure_ai_foundry.py +93 -1
- agno/reasoning/deepseek.py +102 -2
- agno/reasoning/default.py +6 -7
- agno/reasoning/gemini.py +87 -3
- agno/reasoning/groq.py +109 -2
- agno/reasoning/helpers.py +6 -7
- agno/reasoning/manager.py +1238 -0
- agno/reasoning/ollama.py +93 -1
- agno/reasoning/openai.py +115 -1
- agno/reasoning/vertexai.py +85 -1
- agno/registry/__init__.py +3 -0
- agno/registry/registry.py +68 -0
- agno/remote/__init__.py +3 -0
- agno/remote/base.py +581 -0
- agno/run/__init__.py +2 -4
- agno/run/agent.py +134 -19
- agno/run/base.py +49 -1
- agno/run/cancel.py +65 -52
- agno/run/cancellation_management/__init__.py +9 -0
- agno/run/cancellation_management/base.py +78 -0
- agno/run/cancellation_management/in_memory_cancellation_manager.py +100 -0
- agno/run/cancellation_management/redis_cancellation_manager.py +236 -0
- agno/run/requirement.py +181 -0
- agno/run/team.py +111 -19
- agno/run/workflow.py +2 -1
- agno/session/agent.py +57 -92
- agno/session/summary.py +1 -1
- agno/session/team.py +62 -115
- agno/session/workflow.py +353 -57
- agno/skills/__init__.py +17 -0
- agno/skills/agent_skills.py +377 -0
- agno/skills/errors.py +32 -0
- agno/skills/loaders/__init__.py +4 -0
- agno/skills/loaders/base.py +27 -0
- agno/skills/loaders/local.py +216 -0
- agno/skills/skill.py +65 -0
- agno/skills/utils.py +107 -0
- agno/skills/validator.py +277 -0
- agno/table.py +10 -0
- agno/team/__init__.py +5 -1
- agno/team/remote.py +447 -0
- agno/team/team.py +3769 -2202
- agno/tools/brandfetch.py +27 -18
- agno/tools/browserbase.py +225 -16
- agno/tools/crawl4ai.py +3 -0
- agno/tools/duckduckgo.py +25 -71
- agno/tools/exa.py +0 -21
- agno/tools/file.py +14 -13
- agno/tools/file_generation.py +12 -6
- agno/tools/firecrawl.py +15 -7
- agno/tools/function.py +94 -113
- agno/tools/google_bigquery.py +11 -2
- agno/tools/google_drive.py +4 -3
- agno/tools/knowledge.py +9 -4
- agno/tools/mcp/mcp.py +301 -18
- agno/tools/mcp/multi_mcp.py +269 -14
- agno/tools/mem0.py +11 -10
- agno/tools/memory.py +47 -46
- agno/tools/mlx_transcribe.py +10 -7
- agno/tools/models/nebius.py +5 -5
- agno/tools/models_labs.py +20 -10
- agno/tools/nano_banana.py +151 -0
- agno/tools/parallel.py +0 -7
- agno/tools/postgres.py +76 -36
- agno/tools/python.py +14 -6
- agno/tools/reasoning.py +30 -23
- agno/tools/redshift.py +406 -0
- agno/tools/shopify.py +1519 -0
- agno/tools/spotify.py +919 -0
- agno/tools/tavily.py +4 -1
- agno/tools/toolkit.py +253 -18
- agno/tools/websearch.py +93 -0
- agno/tools/website.py +1 -1
- agno/tools/wikipedia.py +1 -1
- agno/tools/workflow.py +56 -48
- agno/tools/yfinance.py +12 -11
- agno/tracing/__init__.py +12 -0
- agno/tracing/exporter.py +161 -0
- agno/tracing/schemas.py +276 -0
- agno/tracing/setup.py +112 -0
- agno/utils/agent.py +251 -10
- agno/utils/cryptography.py +22 -0
- agno/utils/dttm.py +33 -0
- agno/utils/events.py +264 -7
- agno/utils/hooks.py +111 -3
- agno/utils/http.py +161 -2
- agno/utils/mcp.py +49 -8
- agno/utils/media.py +22 -1
- agno/utils/models/ai_foundry.py +9 -2
- agno/utils/models/claude.py +20 -5
- agno/utils/models/cohere.py +9 -2
- agno/utils/models/llama.py +9 -2
- agno/utils/models/mistral.py +4 -2
- agno/utils/os.py +0 -0
- agno/utils/print_response/agent.py +99 -16
- agno/utils/print_response/team.py +223 -24
- agno/utils/print_response/workflow.py +0 -2
- agno/utils/prompts.py +8 -6
- agno/utils/remote.py +23 -0
- agno/utils/response.py +1 -13
- agno/utils/string.py +91 -2
- agno/utils/team.py +62 -12
- agno/utils/tokens.py +657 -0
- agno/vectordb/base.py +15 -2
- agno/vectordb/cassandra/cassandra.py +1 -1
- agno/vectordb/chroma/__init__.py +2 -1
- agno/vectordb/chroma/chromadb.py +468 -23
- agno/vectordb/clickhouse/clickhousedb.py +1 -1
- agno/vectordb/couchbase/couchbase.py +6 -2
- agno/vectordb/lancedb/lance_db.py +7 -38
- agno/vectordb/lightrag/lightrag.py +7 -6
- agno/vectordb/milvus/milvus.py +118 -84
- agno/vectordb/mongodb/__init__.py +2 -1
- agno/vectordb/mongodb/mongodb.py +14 -31
- agno/vectordb/pgvector/pgvector.py +120 -66
- agno/vectordb/pineconedb/pineconedb.py +2 -19
- agno/vectordb/qdrant/__init__.py +2 -1
- agno/vectordb/qdrant/qdrant.py +33 -56
- agno/vectordb/redis/__init__.py +2 -1
- agno/vectordb/redis/redisdb.py +19 -31
- agno/vectordb/singlestore/singlestore.py +17 -9
- agno/vectordb/surrealdb/surrealdb.py +2 -38
- agno/vectordb/weaviate/__init__.py +2 -1
- agno/vectordb/weaviate/weaviate.py +7 -3
- agno/workflow/__init__.py +5 -1
- agno/workflow/agent.py +2 -2
- agno/workflow/condition.py +12 -10
- agno/workflow/loop.py +28 -9
- agno/workflow/parallel.py +21 -13
- agno/workflow/remote.py +362 -0
- agno/workflow/router.py +12 -9
- agno/workflow/step.py +261 -36
- agno/workflow/steps.py +12 -8
- agno/workflow/types.py +40 -77
- agno/workflow/workflow.py +939 -213
- {agno-2.2.13.dist-info → agno-2.4.3.dist-info}/METADATA +134 -181
- agno-2.4.3.dist-info/RECORD +677 -0
- {agno-2.2.13.dist-info → agno-2.4.3.dist-info}/WHEEL +1 -1
- agno/tools/googlesearch.py +0 -98
- agno/tools/memori.py +0 -339
- agno-2.2.13.dist-info/RECORD +0 -575
- {agno-2.2.13.dist-info → agno-2.4.3.dist-info}/licenses/LICENSE +0 -0
- {agno-2.2.13.dist-info → agno-2.4.3.dist-info}/top_level.txt +0 -0
agno/models/google/gemini.py
CHANGED
|
@@ -1,3 +1,5 @@
|
|
|
1
|
+
import asyncio
|
|
2
|
+
import base64
|
|
1
3
|
import json
|
|
2
4
|
import time
|
|
3
5
|
from collections.abc import AsyncIterator
|
|
@@ -11,13 +13,16 @@ from pydantic import BaseModel
|
|
|
11
13
|
|
|
12
14
|
from agno.exceptions import ModelProviderError
|
|
13
15
|
from agno.media import Audio, File, Image, Video
|
|
14
|
-
from agno.models.base import Model
|
|
16
|
+
from agno.models.base import Model, RetryableModelProviderError
|
|
17
|
+
from agno.models.google.utils import MALFORMED_FUNCTION_CALL_GUIDANCE, GeminiFinishReason
|
|
15
18
|
from agno.models.message import Citations, Message, UrlCitation
|
|
16
19
|
from agno.models.metrics import Metrics
|
|
17
20
|
from agno.models.response import ModelResponse
|
|
18
21
|
from agno.run.agent import RunOutput
|
|
22
|
+
from agno.tools.function import Function
|
|
19
23
|
from agno.utils.gemini import format_function_definitions, format_image_for_message, prepare_response_schema
|
|
20
24
|
from agno.utils.log import log_debug, log_error, log_info, log_warning
|
|
25
|
+
from agno.utils.tokens import count_schema_tokens, count_text_tokens, count_tool_tokens
|
|
21
26
|
|
|
22
27
|
try:
|
|
23
28
|
from google import genai
|
|
@@ -26,12 +31,15 @@ try:
|
|
|
26
31
|
from google.genai.types import (
|
|
27
32
|
Content,
|
|
28
33
|
DynamicRetrievalConfig,
|
|
34
|
+
FileSearch,
|
|
29
35
|
FunctionCallingConfigMode,
|
|
30
36
|
GenerateContentConfig,
|
|
31
37
|
GenerateContentResponse,
|
|
32
38
|
GenerateContentResponseUsageMetadata,
|
|
33
39
|
GoogleSearch,
|
|
34
40
|
GoogleSearchRetrieval,
|
|
41
|
+
GroundingMetadata,
|
|
42
|
+
Operation,
|
|
35
43
|
Part,
|
|
36
44
|
Retrieval,
|
|
37
45
|
ThinkingConfig,
|
|
@@ -42,8 +50,11 @@ try:
|
|
|
42
50
|
from google.genai.types import (
|
|
43
51
|
File as GeminiFile,
|
|
44
52
|
)
|
|
53
|
+
from google.oauth2.service_account import Credentials
|
|
45
54
|
except ImportError:
|
|
46
|
-
raise ImportError(
|
|
55
|
+
raise ImportError(
|
|
56
|
+
"`google-genai` not installed or not at the latest version. Please install it using `pip install -U google-genai`"
|
|
57
|
+
)
|
|
47
58
|
|
|
48
59
|
|
|
49
60
|
@dataclass
|
|
@@ -56,6 +67,7 @@ class Gemini(Model):
|
|
|
56
67
|
- Set `vertexai` to `True` to use the Vertex AI API.
|
|
57
68
|
- Set your `project_id` (or set `GOOGLE_CLOUD_PROJECT` environment variable) and `location` (optional).
|
|
58
69
|
- Set `http_options` (optional) to configure the HTTP options.
|
|
70
|
+
- Set `credentials` (optional) to use the Google Cloud credentials.
|
|
59
71
|
|
|
60
72
|
Based on https://googleapis.github.io/python-genai/
|
|
61
73
|
"""
|
|
@@ -78,6 +90,10 @@ class Gemini(Model):
|
|
|
78
90
|
vertexai_search: bool = False
|
|
79
91
|
vertexai_search_datastore: Optional[str] = None
|
|
80
92
|
|
|
93
|
+
# Gemini File Search capabilities
|
|
94
|
+
file_search_store_names: Optional[List[str]] = None
|
|
95
|
+
file_search_metadata_filter: Optional[str] = None
|
|
96
|
+
|
|
81
97
|
temperature: Optional[float] = None
|
|
82
98
|
top_p: Optional[float] = None
|
|
83
99
|
top_k: Optional[int] = None
|
|
@@ -92,9 +108,11 @@ class Gemini(Model):
|
|
|
92
108
|
cached_content: Optional[Any] = None
|
|
93
109
|
thinking_budget: Optional[int] = None # Thinking budget for Gemini 2.5 models
|
|
94
110
|
include_thoughts: Optional[bool] = None # Include thought summaries in response
|
|
111
|
+
thinking_level: Optional[str] = None # "low", "high"
|
|
95
112
|
request_params: Optional[Dict[str, Any]] = None
|
|
96
113
|
|
|
97
114
|
# Client parameters
|
|
115
|
+
credentials: Optional[Credentials] = None
|
|
98
116
|
api_key: Optional[str] = None
|
|
99
117
|
vertexai: bool = False
|
|
100
118
|
project_id: Optional[str] = None
|
|
@@ -135,8 +153,16 @@ class Gemini(Model):
|
|
|
135
153
|
else:
|
|
136
154
|
log_info("Using Vertex AI API")
|
|
137
155
|
client_params["vertexai"] = True
|
|
138
|
-
|
|
139
|
-
|
|
156
|
+
project_id = self.project_id or getenv("GOOGLE_CLOUD_PROJECT")
|
|
157
|
+
if not project_id:
|
|
158
|
+
log_error("GOOGLE_CLOUD_PROJECT not set. Please set the GOOGLE_CLOUD_PROJECT environment variable.")
|
|
159
|
+
location = self.location or getenv("GOOGLE_CLOUD_LOCATION")
|
|
160
|
+
if not location:
|
|
161
|
+
log_error("GOOGLE_CLOUD_LOCATION not set. Please set the GOOGLE_CLOUD_LOCATION environment variable.")
|
|
162
|
+
client_params["project"] = project_id
|
|
163
|
+
client_params["location"] = location
|
|
164
|
+
if self.credentials:
|
|
165
|
+
client_params["credentials"] = self.credentials
|
|
140
166
|
|
|
141
167
|
client_params = {k: v for k, v in client_params.items() if v is not None}
|
|
142
168
|
|
|
@@ -146,6 +172,60 @@ class Gemini(Model):
|
|
|
146
172
|
self.client = genai.Client(**client_params)
|
|
147
173
|
return self.client
|
|
148
174
|
|
|
175
|
+
def to_dict(self) -> Dict[str, Any]:
|
|
176
|
+
"""
|
|
177
|
+
Convert the model to a dictionary.
|
|
178
|
+
|
|
179
|
+
Returns:
|
|
180
|
+
Dict[str, Any]: The dictionary representation of the model.
|
|
181
|
+
"""
|
|
182
|
+
model_dict = super().to_dict()
|
|
183
|
+
model_dict.update(
|
|
184
|
+
{
|
|
185
|
+
"search": self.search,
|
|
186
|
+
"grounding": self.grounding,
|
|
187
|
+
"grounding_dynamic_threshold": self.grounding_dynamic_threshold,
|
|
188
|
+
"url_context": self.url_context,
|
|
189
|
+
"vertexai_search": self.vertexai_search,
|
|
190
|
+
"vertexai_search_datastore": self.vertexai_search_datastore,
|
|
191
|
+
"file_search_store_names": self.file_search_store_names,
|
|
192
|
+
"file_search_metadata_filter": self.file_search_metadata_filter,
|
|
193
|
+
"temperature": self.temperature,
|
|
194
|
+
"top_p": self.top_p,
|
|
195
|
+
"top_k": self.top_k,
|
|
196
|
+
"max_output_tokens": self.max_output_tokens,
|
|
197
|
+
"stop_sequences": self.stop_sequences,
|
|
198
|
+
"logprobs": self.logprobs,
|
|
199
|
+
"presence_penalty": self.presence_penalty,
|
|
200
|
+
"frequency_penalty": self.frequency_penalty,
|
|
201
|
+
"seed": self.seed,
|
|
202
|
+
"response_modalities": self.response_modalities,
|
|
203
|
+
"thinking_budget": self.thinking_budget,
|
|
204
|
+
"include_thoughts": self.include_thoughts,
|
|
205
|
+
"thinking_level": self.thinking_level,
|
|
206
|
+
"vertexai": self.vertexai,
|
|
207
|
+
"project_id": self.project_id,
|
|
208
|
+
"location": self.location,
|
|
209
|
+
}
|
|
210
|
+
)
|
|
211
|
+
cleaned_dict = {k: v for k, v in model_dict.items() if v is not None}
|
|
212
|
+
return cleaned_dict
|
|
213
|
+
|
|
214
|
+
def _append_file_search_tool(self, builtin_tools: List[Tool]) -> None:
|
|
215
|
+
"""Append Gemini File Search tool to builtin_tools if file search is enabled.
|
|
216
|
+
|
|
217
|
+
Args:
|
|
218
|
+
builtin_tools: List of built-in tools to append to.
|
|
219
|
+
"""
|
|
220
|
+
if not self.file_search_store_names:
|
|
221
|
+
return
|
|
222
|
+
|
|
223
|
+
log_debug("Gemini File Search enabled.")
|
|
224
|
+
file_search_config: Dict[str, Any] = {"file_search_store_names": self.file_search_store_names}
|
|
225
|
+
if self.file_search_metadata_filter:
|
|
226
|
+
file_search_config["metadata_filter"] = self.file_search_metadata_filter
|
|
227
|
+
builtin_tools.append(Tool(file_search=FileSearch(**file_search_config))) # type: ignore[arg-type]
|
|
228
|
+
|
|
149
229
|
def get_request_params(
|
|
150
230
|
self,
|
|
151
231
|
system_message: Optional[str] = None,
|
|
@@ -197,11 +277,13 @@ class Gemini(Model):
|
|
|
197
277
|
config["response_schema"] = prepare_response_schema(response_format)
|
|
198
278
|
|
|
199
279
|
# Add thinking configuration
|
|
200
|
-
thinking_config_params = {}
|
|
280
|
+
thinking_config_params: Dict[str, Any] = {}
|
|
201
281
|
if self.thinking_budget is not None:
|
|
202
282
|
thinking_config_params["thinking_budget"] = self.thinking_budget
|
|
203
283
|
if self.include_thoughts is not None:
|
|
204
284
|
thinking_config_params["include_thoughts"] = self.include_thoughts
|
|
285
|
+
if self.thinking_level is not None:
|
|
286
|
+
thinking_config_params["thinking_level"] = self.thinking_level
|
|
205
287
|
if thinking_config_params:
|
|
206
288
|
config["thinking_config"] = ThinkingConfig(**thinking_config_params)
|
|
207
289
|
|
|
@@ -209,8 +291,8 @@ class Gemini(Model):
|
|
|
209
291
|
builtin_tools = []
|
|
210
292
|
|
|
211
293
|
if self.grounding:
|
|
212
|
-
|
|
213
|
-
"Grounding enabled. This is a legacy tool. For Gemini 2.0+ Please use enable `search` flag instead."
|
|
294
|
+
log_debug(
|
|
295
|
+
"Gemini Grounding enabled. This is a legacy tool. For Gemini 2.0+ Please use enable `search` flag instead."
|
|
214
296
|
)
|
|
215
297
|
builtin_tools.append(
|
|
216
298
|
Tool(
|
|
@@ -223,15 +305,15 @@ class Gemini(Model):
|
|
|
223
305
|
)
|
|
224
306
|
|
|
225
307
|
if self.search:
|
|
226
|
-
|
|
308
|
+
log_debug("Gemini Google Search enabled.")
|
|
227
309
|
builtin_tools.append(Tool(google_search=GoogleSearch()))
|
|
228
310
|
|
|
229
311
|
if self.url_context:
|
|
230
|
-
|
|
312
|
+
log_debug("Gemini URL context enabled.")
|
|
231
313
|
builtin_tools.append(Tool(url_context=UrlContext()))
|
|
232
314
|
|
|
233
315
|
if self.vertexai_search:
|
|
234
|
-
|
|
316
|
+
log_debug("Gemini Vertex AI Search enabled.")
|
|
235
317
|
if not self.vertexai_search_datastore:
|
|
236
318
|
log_error("vertexai_search_datastore must be provided when vertexai_search is enabled.")
|
|
237
319
|
raise ValueError("vertexai_search_datastore must be provided when vertexai_search is enabled.")
|
|
@@ -239,6 +321,8 @@ class Gemini(Model):
|
|
|
239
321
|
Tool(retrieval=Retrieval(vertex_ai_search=VertexAISearch(datastore=self.vertexai_search_datastore)))
|
|
240
322
|
)
|
|
241
323
|
|
|
324
|
+
self._append_file_search_tool(builtin_tools)
|
|
325
|
+
|
|
242
326
|
# Set tools in config
|
|
243
327
|
if builtin_tools:
|
|
244
328
|
if tools:
|
|
@@ -272,6 +356,113 @@ class Gemini(Model):
|
|
|
272
356
|
log_debug(f"Calling {self.provider} with request parameters: {request_params}", log_level=2)
|
|
273
357
|
return request_params
|
|
274
358
|
|
|
359
|
+
def count_tokens(
|
|
360
|
+
self,
|
|
361
|
+
messages: List[Message],
|
|
362
|
+
tools: Optional[List[Union[Function, Dict[str, Any]]]] = None,
|
|
363
|
+
output_schema: Optional[Union[Dict, Type[BaseModel]]] = None,
|
|
364
|
+
) -> int:
|
|
365
|
+
contents, system_instruction = self._format_messages(messages, compress_tool_results=True)
|
|
366
|
+
schema_tokens = count_schema_tokens(output_schema, self.id)
|
|
367
|
+
|
|
368
|
+
if self.vertexai:
|
|
369
|
+
# VertexAI supports full token counting with system_instruction and tools
|
|
370
|
+
config: Dict[str, Any] = {}
|
|
371
|
+
if system_instruction:
|
|
372
|
+
config["system_instruction"] = system_instruction
|
|
373
|
+
if tools:
|
|
374
|
+
formatted_tools = self._format_tools(tools)
|
|
375
|
+
gemini_tools = format_function_definitions(formatted_tools)
|
|
376
|
+
if gemini_tools:
|
|
377
|
+
config["tools"] = [gemini_tools]
|
|
378
|
+
|
|
379
|
+
response = self.get_client().models.count_tokens(
|
|
380
|
+
model=self.id,
|
|
381
|
+
contents=contents,
|
|
382
|
+
config=config if config else None, # type: ignore
|
|
383
|
+
)
|
|
384
|
+
return (response.total_tokens or 0) + schema_tokens
|
|
385
|
+
else:
|
|
386
|
+
# Google AI Studio: Use API for content tokens + local estimation for system/tools
|
|
387
|
+
# The API doesn't support system_instruction or tools in config, so we use a hybrid approach:
|
|
388
|
+
# 1. Get accurate token count for contents (text + multimodal) from API
|
|
389
|
+
# 2. Add estimated tokens for system_instruction and tools locally
|
|
390
|
+
try:
|
|
391
|
+
response = self.get_client().models.count_tokens(
|
|
392
|
+
model=self.id,
|
|
393
|
+
contents=contents,
|
|
394
|
+
)
|
|
395
|
+
total = response.total_tokens or 0
|
|
396
|
+
except Exception as e:
|
|
397
|
+
log_warning(f"Gemini count_tokens API failed: {e}. Falling back to tiktoken-based estimation.")
|
|
398
|
+
return super().count_tokens(messages, tools, output_schema)
|
|
399
|
+
|
|
400
|
+
# Add estimated tokens for system instruction (not supported by Google AI Studio API)
|
|
401
|
+
if system_instruction:
|
|
402
|
+
system_text = system_instruction if isinstance(system_instruction, str) else str(system_instruction)
|
|
403
|
+
total += count_text_tokens(system_text, self.id)
|
|
404
|
+
|
|
405
|
+
# Add estimated tokens for tools (not supported by Google AI Studio API)
|
|
406
|
+
if tools:
|
|
407
|
+
total += count_tool_tokens(tools, self.id)
|
|
408
|
+
|
|
409
|
+
# Add estimated tokens for response_format/output_schema
|
|
410
|
+
total += schema_tokens
|
|
411
|
+
|
|
412
|
+
return total
|
|
413
|
+
|
|
414
|
+
async def acount_tokens(
|
|
415
|
+
self,
|
|
416
|
+
messages: List[Message],
|
|
417
|
+
tools: Optional[List[Union[Function, Dict[str, Any]]]] = None,
|
|
418
|
+
output_schema: Optional[Union[Dict, Type[BaseModel]]] = None,
|
|
419
|
+
) -> int:
|
|
420
|
+
contents, system_instruction = self._format_messages(messages, compress_tool_results=True)
|
|
421
|
+
schema_tokens = count_schema_tokens(output_schema, self.id)
|
|
422
|
+
|
|
423
|
+
# VertexAI supports full token counting with system_instruction and tools
|
|
424
|
+
if self.vertexai:
|
|
425
|
+
config: Dict[str, Any] = {}
|
|
426
|
+
if system_instruction:
|
|
427
|
+
config["system_instruction"] = system_instruction
|
|
428
|
+
if tools:
|
|
429
|
+
formatted_tools = self._format_tools(tools)
|
|
430
|
+
gemini_tools = format_function_definitions(formatted_tools)
|
|
431
|
+
if gemini_tools:
|
|
432
|
+
config["tools"] = [gemini_tools]
|
|
433
|
+
|
|
434
|
+
response = await self.get_client().aio.models.count_tokens(
|
|
435
|
+
model=self.id,
|
|
436
|
+
contents=contents,
|
|
437
|
+
config=config if config else None, # type: ignore
|
|
438
|
+
)
|
|
439
|
+
return (response.total_tokens or 0) + schema_tokens
|
|
440
|
+
else:
|
|
441
|
+
# Hybrid approach - Google AI Studio does not support system_instruction or tools in config
|
|
442
|
+
try:
|
|
443
|
+
response = await self.get_client().aio.models.count_tokens(
|
|
444
|
+
model=self.id,
|
|
445
|
+
contents=contents,
|
|
446
|
+
)
|
|
447
|
+
total = response.total_tokens or 0
|
|
448
|
+
except Exception as e:
|
|
449
|
+
log_warning(f"Gemini count_tokens API failed: {e}. Falling back to tiktoken-based estimation.")
|
|
450
|
+
return await super().acount_tokens(messages, tools, output_schema)
|
|
451
|
+
|
|
452
|
+
# Add estimated tokens for system instruction
|
|
453
|
+
if system_instruction:
|
|
454
|
+
system_text = system_instruction if isinstance(system_instruction, str) else str(system_instruction)
|
|
455
|
+
total += count_text_tokens(system_text, self.id)
|
|
456
|
+
|
|
457
|
+
# Add estimated tokens for tools
|
|
458
|
+
if tools:
|
|
459
|
+
total += count_tool_tokens(tools, self.id)
|
|
460
|
+
|
|
461
|
+
# Add estimated tokens for response_format/output_schema
|
|
462
|
+
total += schema_tokens
|
|
463
|
+
|
|
464
|
+
return total
|
|
465
|
+
|
|
275
466
|
def invoke(
|
|
276
467
|
self,
|
|
277
468
|
messages: List[Message],
|
|
@@ -280,11 +471,13 @@ class Gemini(Model):
|
|
|
280
471
|
tools: Optional[List[Dict[str, Any]]] = None,
|
|
281
472
|
tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
|
|
282
473
|
run_response: Optional[RunOutput] = None,
|
|
474
|
+
compress_tool_results: bool = False,
|
|
475
|
+
retry_with_guidance: bool = False,
|
|
283
476
|
) -> ModelResponse:
|
|
284
477
|
"""
|
|
285
478
|
Invokes the model with a list of messages and returns the response.
|
|
286
479
|
"""
|
|
287
|
-
formatted_messages, system_message = self._format_messages(messages)
|
|
480
|
+
formatted_messages, system_message = self._format_messages(messages, compress_tool_results)
|
|
288
481
|
request_kwargs = self.get_request_params(
|
|
289
482
|
system_message, response_format=response_format, tools=tools, tool_choice=tool_choice
|
|
290
483
|
)
|
|
@@ -300,19 +493,32 @@ class Gemini(Model):
|
|
|
300
493
|
)
|
|
301
494
|
assistant_message.metrics.stop_timer()
|
|
302
495
|
|
|
303
|
-
model_response = self._parse_provider_response(
|
|
496
|
+
model_response = self._parse_provider_response(
|
|
497
|
+
provider_response, response_format=response_format, retry_with_guidance=retry_with_guidance
|
|
498
|
+
)
|
|
499
|
+
|
|
500
|
+
# If we were retrying the invoke with guidance, remove the guidance message
|
|
501
|
+
if retry_with_guidance is True:
|
|
502
|
+
self._remove_temporary_messages(messages)
|
|
304
503
|
|
|
305
504
|
return model_response
|
|
306
505
|
|
|
307
506
|
except (ClientError, ServerError) as e:
|
|
308
507
|
log_error(f"Error from Gemini API: {e}")
|
|
309
|
-
error_message = str(e
|
|
508
|
+
error_message = str(e)
|
|
509
|
+
if hasattr(e, "response"):
|
|
510
|
+
if hasattr(e.response, "text"):
|
|
511
|
+
error_message = e.response.text
|
|
512
|
+
else:
|
|
513
|
+
error_message = str(e.response)
|
|
310
514
|
raise ModelProviderError(
|
|
311
515
|
message=error_message,
|
|
312
516
|
status_code=e.code if hasattr(e, "code") and e.code is not None else 502,
|
|
313
517
|
model_name=self.name,
|
|
314
518
|
model_id=self.id,
|
|
315
519
|
) from e
|
|
520
|
+
except RetryableModelProviderError:
|
|
521
|
+
raise
|
|
316
522
|
except Exception as e:
|
|
317
523
|
log_error(f"Unknown error from Gemini API: {e}")
|
|
318
524
|
raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e
|
|
@@ -325,11 +531,13 @@ class Gemini(Model):
|
|
|
325
531
|
tools: Optional[List[Dict[str, Any]]] = None,
|
|
326
532
|
tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
|
|
327
533
|
run_response: Optional[RunOutput] = None,
|
|
534
|
+
compress_tool_results: bool = False,
|
|
535
|
+
retry_with_guidance: bool = False,
|
|
328
536
|
) -> Iterator[ModelResponse]:
|
|
329
537
|
"""
|
|
330
538
|
Invokes the model with a list of messages and returns the response as a stream.
|
|
331
539
|
"""
|
|
332
|
-
formatted_messages, system_message = self._format_messages(messages)
|
|
540
|
+
formatted_messages, system_message = self._format_messages(messages, compress_tool_results)
|
|
333
541
|
|
|
334
542
|
request_kwargs = self.get_request_params(
|
|
335
543
|
system_message, response_format=response_format, tools=tools, tool_choice=tool_choice
|
|
@@ -344,18 +552,30 @@ class Gemini(Model):
|
|
|
344
552
|
contents=formatted_messages,
|
|
345
553
|
**request_kwargs,
|
|
346
554
|
):
|
|
347
|
-
yield self._parse_provider_response_delta(response)
|
|
555
|
+
yield self._parse_provider_response_delta(response, retry_with_guidance=retry_with_guidance)
|
|
556
|
+
|
|
557
|
+
# If we were retrying the invoke with guidance, remove the guidance message
|
|
558
|
+
if retry_with_guidance is True:
|
|
559
|
+
self._remove_temporary_messages(messages)
|
|
348
560
|
|
|
349
561
|
assistant_message.metrics.stop_timer()
|
|
350
562
|
|
|
351
563
|
except (ClientError, ServerError) as e:
|
|
352
564
|
log_error(f"Error from Gemini API: {e}")
|
|
565
|
+
error_message = str(e)
|
|
566
|
+
if hasattr(e, "response"):
|
|
567
|
+
if hasattr(e.response, "text"):
|
|
568
|
+
error_message = e.response.text
|
|
569
|
+
else:
|
|
570
|
+
error_message = str(e.response)
|
|
353
571
|
raise ModelProviderError(
|
|
354
|
-
message=
|
|
572
|
+
message=error_message,
|
|
355
573
|
status_code=e.code if hasattr(e, "code") and e.code is not None else 502,
|
|
356
574
|
model_name=self.name,
|
|
357
575
|
model_id=self.id,
|
|
358
576
|
) from e
|
|
577
|
+
except RetryableModelProviderError:
|
|
578
|
+
raise
|
|
359
579
|
except Exception as e:
|
|
360
580
|
log_error(f"Unknown error from Gemini API: {e}")
|
|
361
581
|
raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e
|
|
@@ -368,11 +588,13 @@ class Gemini(Model):
|
|
|
368
588
|
tools: Optional[List[Dict[str, Any]]] = None,
|
|
369
589
|
tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
|
|
370
590
|
run_response: Optional[RunOutput] = None,
|
|
591
|
+
compress_tool_results: bool = False,
|
|
592
|
+
retry_with_guidance: bool = False,
|
|
371
593
|
) -> ModelResponse:
|
|
372
594
|
"""
|
|
373
595
|
Invokes the model with a list of messages and returns the response.
|
|
374
596
|
"""
|
|
375
|
-
formatted_messages, system_message = self._format_messages(messages)
|
|
597
|
+
formatted_messages, system_message = self._format_messages(messages, compress_tool_results)
|
|
376
598
|
|
|
377
599
|
request_kwargs = self.get_request_params(
|
|
378
600
|
system_message, response_format=response_format, tools=tools, tool_choice=tool_choice
|
|
@@ -390,18 +612,32 @@ class Gemini(Model):
|
|
|
390
612
|
)
|
|
391
613
|
assistant_message.metrics.stop_timer()
|
|
392
614
|
|
|
393
|
-
model_response = self._parse_provider_response(
|
|
615
|
+
model_response = self._parse_provider_response(
|
|
616
|
+
provider_response, response_format=response_format, retry_with_guidance=retry_with_guidance
|
|
617
|
+
)
|
|
618
|
+
|
|
619
|
+
# If we were retrying the invoke with guidance, remove the guidance message
|
|
620
|
+
if retry_with_guidance is True:
|
|
621
|
+
self._remove_temporary_messages(messages)
|
|
394
622
|
|
|
395
623
|
return model_response
|
|
396
624
|
|
|
397
625
|
except (ClientError, ServerError) as e:
|
|
398
626
|
log_error(f"Error from Gemini API: {e}")
|
|
627
|
+
error_message = str(e)
|
|
628
|
+
if hasattr(e, "response"):
|
|
629
|
+
if hasattr(e.response, "text"):
|
|
630
|
+
error_message = e.response.text
|
|
631
|
+
else:
|
|
632
|
+
error_message = str(e.response)
|
|
399
633
|
raise ModelProviderError(
|
|
400
|
-
message=
|
|
634
|
+
message=error_message,
|
|
401
635
|
status_code=e.code if hasattr(e, "code") and e.code is not None else 502,
|
|
402
636
|
model_name=self.name,
|
|
403
637
|
model_id=self.id,
|
|
404
638
|
) from e
|
|
639
|
+
except RetryableModelProviderError:
|
|
640
|
+
raise
|
|
405
641
|
except Exception as e:
|
|
406
642
|
log_error(f"Unknown error from Gemini API: {e}")
|
|
407
643
|
raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e
|
|
@@ -414,11 +650,13 @@ class Gemini(Model):
|
|
|
414
650
|
tools: Optional[List[Dict[str, Any]]] = None,
|
|
415
651
|
tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
|
|
416
652
|
run_response: Optional[RunOutput] = None,
|
|
653
|
+
compress_tool_results: bool = False,
|
|
654
|
+
retry_with_guidance: bool = False,
|
|
417
655
|
) -> AsyncIterator[ModelResponse]:
|
|
418
656
|
"""
|
|
419
657
|
Invokes the model with a list of messages and returns the response as a stream.
|
|
420
658
|
"""
|
|
421
|
-
formatted_messages, system_message = self._format_messages(messages)
|
|
659
|
+
formatted_messages, system_message = self._format_messages(messages, compress_tool_results)
|
|
422
660
|
|
|
423
661
|
request_kwargs = self.get_request_params(
|
|
424
662
|
system_message, response_format=response_format, tools=tools, tool_choice=tool_choice
|
|
@@ -436,32 +674,45 @@ class Gemini(Model):
|
|
|
436
674
|
**request_kwargs,
|
|
437
675
|
)
|
|
438
676
|
async for chunk in async_stream:
|
|
439
|
-
yield self._parse_provider_response_delta(chunk)
|
|
677
|
+
yield self._parse_provider_response_delta(chunk, retry_with_guidance=retry_with_guidance)
|
|
678
|
+
|
|
679
|
+
# If we were retrying the invoke with guidance, remove the guidance message
|
|
680
|
+
if retry_with_guidance is True:
|
|
681
|
+
self._remove_temporary_messages(messages)
|
|
440
682
|
|
|
441
683
|
assistant_message.metrics.stop_timer()
|
|
442
684
|
|
|
443
685
|
except (ClientError, ServerError) as e:
|
|
444
686
|
log_error(f"Error from Gemini API: {e}")
|
|
687
|
+
error_message = str(e)
|
|
688
|
+
if hasattr(e, "response"):
|
|
689
|
+
if hasattr(e.response, "text"):
|
|
690
|
+
error_message = e.response.text
|
|
691
|
+
else:
|
|
692
|
+
error_message = str(e.response)
|
|
445
693
|
raise ModelProviderError(
|
|
446
|
-
message=
|
|
694
|
+
message=error_message,
|
|
447
695
|
status_code=e.code if hasattr(e, "code") and e.code is not None else 502,
|
|
448
696
|
model_name=self.name,
|
|
449
697
|
model_id=self.id,
|
|
450
698
|
) from e
|
|
699
|
+
except RetryableModelProviderError:
|
|
700
|
+
raise
|
|
451
701
|
except Exception as e:
|
|
452
702
|
log_error(f"Unknown error from Gemini API: {e}")
|
|
453
703
|
raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e
|
|
454
704
|
|
|
455
|
-
def _format_messages(self, messages: List[Message]):
|
|
705
|
+
def _format_messages(self, messages: List[Message], compress_tool_results: bool = False):
|
|
456
706
|
"""
|
|
457
707
|
Converts a list of Message objects to the Gemini-compatible format.
|
|
458
708
|
|
|
459
709
|
Args:
|
|
460
710
|
messages (List[Message]): The list of messages to convert.
|
|
711
|
+
compress_tool_results: Whether to compress tool results.
|
|
461
712
|
"""
|
|
462
713
|
formatted_messages: List = []
|
|
463
|
-
file_content: Optional[Union[GeminiFile, Part]] = None
|
|
464
714
|
system_message = None
|
|
715
|
+
|
|
465
716
|
for message in messages:
|
|
466
717
|
role = message.role
|
|
467
718
|
if role in ["system", "developer"]:
|
|
@@ -472,7 +723,8 @@ class Gemini(Model):
|
|
|
472
723
|
role = self.reverse_role_map.get(role, role)
|
|
473
724
|
|
|
474
725
|
# Add content to the message for the model
|
|
475
|
-
content = message.
|
|
726
|
+
content = message.get_content(use_compressed_content=compress_tool_results)
|
|
727
|
+
|
|
476
728
|
# Initialize message_parts to be used for Gemini
|
|
477
729
|
message_parts: List[Any] = []
|
|
478
730
|
|
|
@@ -480,26 +732,47 @@ class Gemini(Model):
|
|
|
480
732
|
if role == "model" and message.tool_calls is not None and len(message.tool_calls) > 0:
|
|
481
733
|
if content is not None:
|
|
482
734
|
content_str = content if isinstance(content, str) else str(content)
|
|
483
|
-
|
|
735
|
+
part = Part.from_text(text=content_str)
|
|
736
|
+
if message.provider_data and "thought_signature" in message.provider_data:
|
|
737
|
+
part.thought_signature = base64.b64decode(message.provider_data["thought_signature"])
|
|
738
|
+
message_parts.append(part)
|
|
484
739
|
for tool_call in message.tool_calls:
|
|
485
|
-
|
|
486
|
-
|
|
487
|
-
|
|
488
|
-
args=json.loads(tool_call["function"]["arguments"]),
|
|
489
|
-
)
|
|
740
|
+
part = Part.from_function_call(
|
|
741
|
+
name=tool_call["function"]["name"],
|
|
742
|
+
args=json.loads(tool_call["function"]["arguments"]),
|
|
490
743
|
)
|
|
744
|
+
if "thought_signature" in tool_call:
|
|
745
|
+
part.thought_signature = base64.b64decode(tool_call["thought_signature"])
|
|
746
|
+
message_parts.append(part)
|
|
491
747
|
# Function call results
|
|
492
748
|
elif message.tool_calls is not None and len(message.tool_calls) > 0:
|
|
493
|
-
for tool_call in message.tool_calls:
|
|
749
|
+
for idx, tool_call in enumerate(message.tool_calls):
|
|
750
|
+
if isinstance(content, list) and idx < len(content):
|
|
751
|
+
original_from_list = content[idx]
|
|
752
|
+
|
|
753
|
+
if compress_tool_results:
|
|
754
|
+
compressed_from_tool_call = tool_call.get("content")
|
|
755
|
+
tc_content = compressed_from_tool_call if compressed_from_tool_call else original_from_list
|
|
756
|
+
else:
|
|
757
|
+
tc_content = original_from_list
|
|
758
|
+
else:
|
|
759
|
+
tc_content = message.get_content(use_compressed_content=compress_tool_results)
|
|
760
|
+
|
|
761
|
+
if tc_content is None:
|
|
762
|
+
tc_content = tool_call.get("content")
|
|
763
|
+
if tc_content is None:
|
|
764
|
+
tc_content = content
|
|
765
|
+
|
|
494
766
|
message_parts.append(
|
|
495
|
-
Part.from_function_response(
|
|
496
|
-
name=tool_call["tool_name"], response={"result": tool_call["content"]}
|
|
497
|
-
)
|
|
767
|
+
Part.from_function_response(name=tool_call["tool_name"], response={"result": tc_content})
|
|
498
768
|
)
|
|
499
769
|
# Regular text content
|
|
500
770
|
else:
|
|
501
771
|
if isinstance(content, str):
|
|
502
|
-
|
|
772
|
+
part = Part.from_text(text=content)
|
|
773
|
+
if message.provider_data and "thought_signature" in message.provider_data:
|
|
774
|
+
part.thought_signature = base64.b64decode(message.provider_data["thought_signature"])
|
|
775
|
+
message_parts = [part]
|
|
503
776
|
|
|
504
777
|
if role == "user" and message.tool_calls is None:
|
|
505
778
|
# Add images to the message for the model
|
|
@@ -560,14 +833,11 @@ class Gemini(Model):
|
|
|
560
833
|
for file in message.files:
|
|
561
834
|
file_content = self._format_file_for_message(file)
|
|
562
835
|
if isinstance(file_content, Part):
|
|
563
|
-
|
|
836
|
+
message_parts.append(file_content)
|
|
564
837
|
|
|
565
838
|
final_message = Content(role=role, parts=message_parts)
|
|
566
839
|
formatted_messages.append(final_message)
|
|
567
840
|
|
|
568
|
-
if isinstance(file_content, GeminiFile):
|
|
569
|
-
formatted_messages.insert(0, file_content)
|
|
570
|
-
|
|
571
841
|
return formatted_messages, system_message
|
|
572
842
|
|
|
573
843
|
def _format_audio_for_message(self, audio: Audio) -> Optional[Union[Part, GeminiFile]]:
|
|
@@ -701,6 +971,16 @@ class Gemini(Model):
|
|
|
701
971
|
|
|
702
972
|
# Case 2: File is a URL
|
|
703
973
|
elif file.url is not None:
|
|
974
|
+
# Case 2a: GCS URI (gs://) - pass directly to Gemini (supports up to 2GB)
|
|
975
|
+
if file.url.startswith("gs://") and file.mime_type:
|
|
976
|
+
return Part.from_uri(file_uri=file.url, mime_type=file.mime_type)
|
|
977
|
+
|
|
978
|
+
# Case 2b: HTTPS URL with mime_type - pass directly to Gemini (supports up to 100MB)
|
|
979
|
+
# This enables pre-signed URLs from S3/Azure and public URLs without downloading
|
|
980
|
+
if file.url.startswith("https://") and file.mime_type:
|
|
981
|
+
return Part.from_uri(file_uri=file.url, mime_type=file.mime_type)
|
|
982
|
+
|
|
983
|
+
# Case 2c: URL without mime_type - download and detect (existing behavior)
|
|
704
984
|
url_content = file.file_url_content
|
|
705
985
|
if url_content is not None:
|
|
706
986
|
content, mime_type = url_content
|
|
@@ -759,33 +1039,57 @@ class Gemini(Model):
|
|
|
759
1039
|
return None
|
|
760
1040
|
|
|
761
1041
|
def format_function_call_results(
|
|
762
|
-
self,
|
|
1042
|
+
self,
|
|
1043
|
+
messages: List[Message],
|
|
1044
|
+
function_call_results: List[Message],
|
|
1045
|
+
compress_tool_results: bool = False,
|
|
1046
|
+
**kwargs,
|
|
763
1047
|
) -> None:
|
|
764
1048
|
"""
|
|
765
|
-
Format function call results.
|
|
1049
|
+
Format function call results for Gemini.
|
|
1050
|
+
|
|
1051
|
+
For combined messages:
|
|
1052
|
+
- content: list of ORIGINAL content (for preservation)
|
|
1053
|
+
- tool_calls[i]["content"]: compressed content if available (for API sending)
|
|
1054
|
+
|
|
1055
|
+
This allows the message to be saved with both original and compressed versions.
|
|
766
1056
|
"""
|
|
767
|
-
|
|
1057
|
+
combined_original_content: List = []
|
|
768
1058
|
combined_function_result: List = []
|
|
1059
|
+
tool_names: List[str] = []
|
|
1060
|
+
|
|
769
1061
|
message_metrics = Metrics()
|
|
1062
|
+
|
|
770
1063
|
if len(function_call_results) > 0:
|
|
771
|
-
for result in function_call_results:
|
|
772
|
-
|
|
773
|
-
|
|
1064
|
+
for idx, result in enumerate(function_call_results):
|
|
1065
|
+
combined_original_content.append(result.content)
|
|
1066
|
+
compressed_content = result.get_content(use_compressed_content=compress_tool_results)
|
|
1067
|
+
combined_function_result.append(
|
|
1068
|
+
{"tool_call_id": result.tool_call_id, "tool_name": result.tool_name, "content": compressed_content}
|
|
1069
|
+
)
|
|
1070
|
+
if result.tool_name:
|
|
1071
|
+
tool_names.append(result.tool_name)
|
|
774
1072
|
message_metrics += result.metrics
|
|
775
1073
|
|
|
776
|
-
if
|
|
1074
|
+
tool_name = ", ".join(tool_names) if tool_names else None
|
|
1075
|
+
|
|
1076
|
+
if combined_original_content:
|
|
777
1077
|
messages.append(
|
|
778
1078
|
Message(
|
|
779
|
-
role="tool",
|
|
1079
|
+
role="tool",
|
|
1080
|
+
content=combined_original_content,
|
|
1081
|
+
tool_name=tool_name,
|
|
1082
|
+
tool_calls=combined_function_result,
|
|
1083
|
+
metrics=message_metrics,
|
|
780
1084
|
)
|
|
781
1085
|
)
|
|
782
1086
|
|
|
783
1087
|
def _parse_provider_response(self, response: GenerateContentResponse, **kwargs) -> ModelResponse:
|
|
784
1088
|
"""
|
|
785
|
-
Parse the
|
|
1089
|
+
Parse the Gemini response into a ModelResponse.
|
|
786
1090
|
|
|
787
1091
|
Args:
|
|
788
|
-
response: Raw response from
|
|
1092
|
+
response: Raw response from Gemini
|
|
789
1093
|
|
|
790
1094
|
Returns:
|
|
791
1095
|
ModelResponse: Parsed response data
|
|
@@ -794,8 +1098,20 @@ class Gemini(Model):
|
|
|
794
1098
|
|
|
795
1099
|
# Get response message
|
|
796
1100
|
response_message = Content(role="model", parts=[])
|
|
797
|
-
if response.candidates and response.candidates
|
|
798
|
-
|
|
1101
|
+
if response.candidates and len(response.candidates) > 0:
|
|
1102
|
+
candidate = response.candidates[0]
|
|
1103
|
+
|
|
1104
|
+
# Raise if the request failed because of a malformed function call
|
|
1105
|
+
if hasattr(candidate, "finish_reason") and candidate.finish_reason:
|
|
1106
|
+
if candidate.finish_reason == GeminiFinishReason.MALFORMED_FUNCTION_CALL.value:
|
|
1107
|
+
if self.retry_with_guidance:
|
|
1108
|
+
raise RetryableModelProviderError(
|
|
1109
|
+
retry_guidance_message=MALFORMED_FUNCTION_CALL_GUIDANCE,
|
|
1110
|
+
original_error=f"Generation ended with finish reason: {candidate.finish_reason}",
|
|
1111
|
+
)
|
|
1112
|
+
|
|
1113
|
+
if candidate.content:
|
|
1114
|
+
response_message = candidate.content
|
|
799
1115
|
|
|
800
1116
|
# Add role
|
|
801
1117
|
if response_message.role is not None:
|
|
@@ -834,6 +1150,14 @@ class Gemini(Model):
|
|
|
834
1150
|
else:
|
|
835
1151
|
model_response.content += content_str
|
|
836
1152
|
|
|
1153
|
+
# Capture thought signature for text parts
|
|
1154
|
+
if hasattr(part, "thought_signature") and part.thought_signature:
|
|
1155
|
+
if model_response.provider_data is None:
|
|
1156
|
+
model_response.provider_data = {}
|
|
1157
|
+
model_response.provider_data["thought_signature"] = base64.b64encode(
|
|
1158
|
+
part.thought_signature
|
|
1159
|
+
).decode("ascii")
|
|
1160
|
+
|
|
837
1161
|
if hasattr(part, "inline_data") and part.inline_data is not None:
|
|
838
1162
|
# Handle audio responses (for TTS models)
|
|
839
1163
|
if part.inline_data.mime_type and part.inline_data.mime_type.startswith("audio/"):
|
|
@@ -865,32 +1189,33 @@ class Gemini(Model):
|
|
|
865
1189
|
},
|
|
866
1190
|
}
|
|
867
1191
|
|
|
1192
|
+
# Capture thought signature for function calls
|
|
1193
|
+
if hasattr(part, "thought_signature") and part.thought_signature:
|
|
1194
|
+
tool_call["thought_signature"] = base64.b64encode(part.thought_signature).decode("ascii")
|
|
1195
|
+
|
|
868
1196
|
model_response.tool_calls.append(tool_call)
|
|
869
1197
|
|
|
870
1198
|
citations = Citations()
|
|
871
1199
|
citations_raw = {}
|
|
872
1200
|
citations_urls = []
|
|
1201
|
+
web_search_queries: List[str] = []
|
|
873
1202
|
|
|
874
1203
|
if response.candidates and response.candidates[0].grounding_metadata is not None:
|
|
875
|
-
grounding_metadata = response.candidates[0].grounding_metadata
|
|
876
|
-
citations_raw["grounding_metadata"] = grounding_metadata
|
|
1204
|
+
grounding_metadata: GroundingMetadata = response.candidates[0].grounding_metadata
|
|
1205
|
+
citations_raw["grounding_metadata"] = grounding_metadata.model_dump()
|
|
877
1206
|
|
|
878
|
-
chunks = grounding_metadata.
|
|
879
|
-
|
|
1207
|
+
chunks = grounding_metadata.grounding_chunks or []
|
|
1208
|
+
web_search_queries = grounding_metadata.web_search_queries or []
|
|
880
1209
|
for chunk in chunks:
|
|
881
|
-
if not
|
|
1210
|
+
if not chunk:
|
|
882
1211
|
continue
|
|
883
|
-
web = chunk.
|
|
884
|
-
if not
|
|
1212
|
+
web = chunk.web
|
|
1213
|
+
if not web:
|
|
885
1214
|
continue
|
|
886
|
-
uri = web.
|
|
887
|
-
title = web.
|
|
1215
|
+
uri = web.uri
|
|
1216
|
+
title = web.title
|
|
888
1217
|
if uri:
|
|
889
|
-
|
|
890
|
-
|
|
891
|
-
# Create citation objects from filtered pairs
|
|
892
|
-
grounding_urls = [UrlCitation(url=url, title=title) for url, title in citation_pairs]
|
|
893
|
-
citations_urls.extend(grounding_urls)
|
|
1218
|
+
citations_urls.append(UrlCitation(url=uri, title=title))
|
|
894
1219
|
|
|
895
1220
|
# Handle URLs from URL context tool
|
|
896
1221
|
if (
|
|
@@ -898,22 +1223,29 @@ class Gemini(Model):
|
|
|
898
1223
|
and hasattr(response.candidates[0], "url_context_metadata")
|
|
899
1224
|
and response.candidates[0].url_context_metadata is not None
|
|
900
1225
|
):
|
|
901
|
-
url_context_metadata = response.candidates[0].url_context_metadata
|
|
902
|
-
citations_raw["url_context_metadata"] = url_context_metadata
|
|
1226
|
+
url_context_metadata = response.candidates[0].url_context_metadata
|
|
1227
|
+
citations_raw["url_context_metadata"] = url_context_metadata.model_dump()
|
|
903
1228
|
|
|
904
|
-
url_metadata_list = url_context_metadata.
|
|
1229
|
+
url_metadata_list = url_context_metadata.url_metadata or []
|
|
905
1230
|
for url_meta in url_metadata_list:
|
|
906
|
-
retrieved_url = url_meta.
|
|
907
|
-
status =
|
|
1231
|
+
retrieved_url = url_meta.retrieved_url
|
|
1232
|
+
status = "UNKNOWN"
|
|
1233
|
+
if url_meta.url_retrieval_status:
|
|
1234
|
+
status = url_meta.url_retrieval_status.value
|
|
908
1235
|
if retrieved_url and status == "URL_RETRIEVAL_STATUS_SUCCESS":
|
|
909
1236
|
# Avoid duplicate URLs
|
|
910
1237
|
existing_urls = [citation.url for citation in citations_urls]
|
|
911
1238
|
if retrieved_url not in existing_urls:
|
|
912
1239
|
citations_urls.append(UrlCitation(url=retrieved_url, title=retrieved_url))
|
|
913
1240
|
|
|
1241
|
+
if citations_raw:
|
|
1242
|
+
citations.raw = citations_raw
|
|
1243
|
+
if citations_urls:
|
|
1244
|
+
citations.urls = citations_urls
|
|
1245
|
+
if web_search_queries:
|
|
1246
|
+
citations.search_queries = web_search_queries
|
|
1247
|
+
|
|
914
1248
|
if citations_raw or citations_urls:
|
|
915
|
-
citations.raw = citations_raw if citations_raw else None
|
|
916
|
-
citations.urls = citations_urls if citations_urls else None
|
|
917
1249
|
model_response.citations = citations
|
|
918
1250
|
|
|
919
1251
|
# Extract usage metadata if present
|
|
@@ -926,11 +1258,22 @@ class Gemini(Model):
|
|
|
926
1258
|
|
|
927
1259
|
return model_response
|
|
928
1260
|
|
|
929
|
-
def _parse_provider_response_delta(self, response_delta: GenerateContentResponse) -> ModelResponse:
|
|
1261
|
+
def _parse_provider_response_delta(self, response_delta: GenerateContentResponse, **kwargs) -> ModelResponse:
|
|
930
1262
|
model_response = ModelResponse()
|
|
931
1263
|
|
|
932
1264
|
if response_delta.candidates and len(response_delta.candidates) > 0:
|
|
933
|
-
|
|
1265
|
+
candidate = response_delta.candidates[0]
|
|
1266
|
+
candidate_content = candidate.content
|
|
1267
|
+
|
|
1268
|
+
# Raise if the request failed because of a malformed function call
|
|
1269
|
+
if hasattr(candidate, "finish_reason") and candidate.finish_reason:
|
|
1270
|
+
if candidate.finish_reason == GeminiFinishReason.MALFORMED_FUNCTION_CALL.value:
|
|
1271
|
+
if self.retry_with_guidance:
|
|
1272
|
+
raise RetryableModelProviderError(
|
|
1273
|
+
retry_guidance_message=MALFORMED_FUNCTION_CALL_GUIDANCE,
|
|
1274
|
+
original_error=f"Generation ended with finish reason: {candidate.finish_reason}",
|
|
1275
|
+
)
|
|
1276
|
+
|
|
934
1277
|
response_message: Content = Content(role="model", parts=[])
|
|
935
1278
|
if candidate_content is not None:
|
|
936
1279
|
response_message = candidate_content
|
|
@@ -956,6 +1299,14 @@ class Gemini(Model):
|
|
|
956
1299
|
else:
|
|
957
1300
|
model_response.content += text_content
|
|
958
1301
|
|
|
1302
|
+
# Capture thought signature for text parts
|
|
1303
|
+
if hasattr(part, "thought_signature") and part.thought_signature:
|
|
1304
|
+
if model_response.provider_data is None:
|
|
1305
|
+
model_response.provider_data = {}
|
|
1306
|
+
model_response.provider_data["thought_signature"] = base64.b64encode(
|
|
1307
|
+
part.thought_signature
|
|
1308
|
+
).decode("ascii")
|
|
1309
|
+
|
|
959
1310
|
if hasattr(part, "inline_data") and part.inline_data is not None:
|
|
960
1311
|
# Audio responses
|
|
961
1312
|
if part.inline_data.mime_type and part.inline_data.mime_type.startswith("audio/"):
|
|
@@ -989,30 +1340,58 @@ class Gemini(Model):
|
|
|
989
1340
|
},
|
|
990
1341
|
}
|
|
991
1342
|
|
|
1343
|
+
# Capture thought signature for function calls
|
|
1344
|
+
if hasattr(part, "thought_signature") and part.thought_signature:
|
|
1345
|
+
tool_call["thought_signature"] = base64.b64encode(part.thought_signature).decode("ascii")
|
|
1346
|
+
|
|
992
1347
|
model_response.tool_calls.append(tool_call)
|
|
993
1348
|
|
|
994
|
-
|
|
995
|
-
|
|
996
|
-
|
|
997
|
-
citations.raw = grounding_metadata
|
|
1349
|
+
citations = Citations()
|
|
1350
|
+
citations.raw = {}
|
|
1351
|
+
citations.urls = []
|
|
998
1352
|
|
|
1353
|
+
if (
|
|
1354
|
+
hasattr(response_delta.candidates[0], "grounding_metadata")
|
|
1355
|
+
and response_delta.candidates[0].grounding_metadata is not None
|
|
1356
|
+
):
|
|
1357
|
+
grounding_metadata = response_delta.candidates[0].grounding_metadata
|
|
1358
|
+
citations.raw["grounding_metadata"] = grounding_metadata.model_dump()
|
|
1359
|
+
citations.search_queries = grounding_metadata.web_search_queries or []
|
|
999
1360
|
# Extract url and title
|
|
1000
|
-
chunks = grounding_metadata.
|
|
1001
|
-
citation_pairs = []
|
|
1361
|
+
chunks = grounding_metadata.grounding_chunks or []
|
|
1002
1362
|
for chunk in chunks:
|
|
1003
|
-
if not
|
|
1363
|
+
if not chunk:
|
|
1004
1364
|
continue
|
|
1005
|
-
web = chunk.
|
|
1006
|
-
if not
|
|
1365
|
+
web = chunk.web
|
|
1366
|
+
if not web:
|
|
1007
1367
|
continue
|
|
1008
|
-
uri = web.
|
|
1009
|
-
title = web.
|
|
1368
|
+
uri = web.uri
|
|
1369
|
+
title = web.title
|
|
1010
1370
|
if uri:
|
|
1011
|
-
|
|
1371
|
+
citations.urls.append(UrlCitation(url=uri, title=title))
|
|
1012
1372
|
|
|
1013
|
-
|
|
1014
|
-
|
|
1373
|
+
# Handle URLs from URL context tool
|
|
1374
|
+
if (
|
|
1375
|
+
hasattr(response_delta.candidates[0], "url_context_metadata")
|
|
1376
|
+
and response_delta.candidates[0].url_context_metadata is not None
|
|
1377
|
+
):
|
|
1378
|
+
url_context_metadata = response_delta.candidates[0].url_context_metadata
|
|
1015
1379
|
|
|
1380
|
+
citations.raw["url_context_metadata"] = url_context_metadata.model_dump()
|
|
1381
|
+
|
|
1382
|
+
url_metadata_list = url_context_metadata.url_metadata or []
|
|
1383
|
+
for url_meta in url_metadata_list:
|
|
1384
|
+
retrieved_url = url_meta.retrieved_url
|
|
1385
|
+
status = "UNKNOWN"
|
|
1386
|
+
if url_meta.url_retrieval_status:
|
|
1387
|
+
status = url_meta.url_retrieval_status.value
|
|
1388
|
+
if retrieved_url and status == "URL_RETRIEVAL_STATUS_SUCCESS":
|
|
1389
|
+
# Avoid duplicate URLs
|
|
1390
|
+
existing_urls = [citation.url for citation in citations.urls]
|
|
1391
|
+
if retrieved_url not in existing_urls:
|
|
1392
|
+
citations.urls.append(UrlCitation(url=retrieved_url, title=retrieved_url))
|
|
1393
|
+
|
|
1394
|
+
if citations.raw or citations.urls:
|
|
1016
1395
|
model_response.citations = citations
|
|
1017
1396
|
|
|
1018
1397
|
# Extract usage metadata if present
|
|
@@ -1083,3 +1462,494 @@ class Gemini(Model):
|
|
|
1083
1462
|
metrics.provider_metrics = {"traffic_type": response_usage.traffic_type}
|
|
1084
1463
|
|
|
1085
1464
|
return metrics
|
|
1465
|
+
|
|
1466
|
+
def create_file_search_store(self, display_name: Optional[str] = None) -> Any:
|
|
1467
|
+
"""
|
|
1468
|
+
Create a new File Search store.
|
|
1469
|
+
|
|
1470
|
+
Args:
|
|
1471
|
+
display_name: Optional display name for the store
|
|
1472
|
+
|
|
1473
|
+
Returns:
|
|
1474
|
+
FileSearchStore: The created File Search store object
|
|
1475
|
+
"""
|
|
1476
|
+
config: Dict[str, Any] = {}
|
|
1477
|
+
if display_name:
|
|
1478
|
+
config["display_name"] = display_name
|
|
1479
|
+
|
|
1480
|
+
try:
|
|
1481
|
+
store = self.get_client().file_search_stores.create(config=config or None) # type: ignore[arg-type]
|
|
1482
|
+
log_info(f"Created File Search store: {store.name}")
|
|
1483
|
+
return store
|
|
1484
|
+
except Exception as e:
|
|
1485
|
+
log_error(f"Error creating File Search store: {e}")
|
|
1486
|
+
raise
|
|
1487
|
+
|
|
1488
|
+
async def async_create_file_search_store(self, display_name: Optional[str] = None) -> Any:
|
|
1489
|
+
"""
|
|
1490
|
+
Args:
|
|
1491
|
+
display_name: Optional display name for the store
|
|
1492
|
+
|
|
1493
|
+
Returns:
|
|
1494
|
+
FileSearchStore: The created File Search store object
|
|
1495
|
+
"""
|
|
1496
|
+
config: Dict[str, Any] = {}
|
|
1497
|
+
if display_name:
|
|
1498
|
+
config["display_name"] = display_name
|
|
1499
|
+
|
|
1500
|
+
try:
|
|
1501
|
+
store = await self.get_client().aio.file_search_stores.create(config=config or None) # type: ignore[arg-type]
|
|
1502
|
+
log_info(f"Created File Search store: {store.name}")
|
|
1503
|
+
return store
|
|
1504
|
+
except Exception as e:
|
|
1505
|
+
log_error(f"Error creating File Search store: {e}")
|
|
1506
|
+
raise
|
|
1507
|
+
|
|
1508
|
+
def list_file_search_stores(self, page_size: int = 100) -> List[Any]:
|
|
1509
|
+
"""
|
|
1510
|
+
List all File Search stores.
|
|
1511
|
+
|
|
1512
|
+
Args:
|
|
1513
|
+
page_size: Maximum number of stores to return per page
|
|
1514
|
+
|
|
1515
|
+
Returns:
|
|
1516
|
+
List: List of FileSearchStore objects
|
|
1517
|
+
"""
|
|
1518
|
+
try:
|
|
1519
|
+
stores = []
|
|
1520
|
+
for store in self.get_client().file_search_stores.list(config={"page_size": page_size}):
|
|
1521
|
+
stores.append(store)
|
|
1522
|
+
log_debug(f"Found {len(stores)} File Search stores")
|
|
1523
|
+
return stores
|
|
1524
|
+
except Exception as e:
|
|
1525
|
+
log_error(f"Error listing File Search stores: {e}")
|
|
1526
|
+
raise
|
|
1527
|
+
|
|
1528
|
+
async def async_list_file_search_stores(self, page_size: int = 100) -> List[Any]:
|
|
1529
|
+
"""
|
|
1530
|
+
Async version of list_file_search_stores.
|
|
1531
|
+
|
|
1532
|
+
Args:
|
|
1533
|
+
page_size: Maximum number of stores to return per page
|
|
1534
|
+
|
|
1535
|
+
Returns:
|
|
1536
|
+
List: List of FileSearchStore objects
|
|
1537
|
+
"""
|
|
1538
|
+
try:
|
|
1539
|
+
stores = []
|
|
1540
|
+
async for store in await self.get_client().aio.file_search_stores.list(config={"page_size": page_size}):
|
|
1541
|
+
stores.append(store)
|
|
1542
|
+
log_debug(f"Found {len(stores)} File Search stores")
|
|
1543
|
+
return stores
|
|
1544
|
+
except Exception as e:
|
|
1545
|
+
log_error(f"Error listing File Search stores: {e}")
|
|
1546
|
+
raise
|
|
1547
|
+
|
|
1548
|
+
def get_file_search_store(self, name: str) -> Any:
|
|
1549
|
+
"""
|
|
1550
|
+
Get a specific File Search store by name.
|
|
1551
|
+
|
|
1552
|
+
Args:
|
|
1553
|
+
name: The name of the store (e.g., 'fileSearchStores/my-store-123')
|
|
1554
|
+
|
|
1555
|
+
Returns:
|
|
1556
|
+
FileSearchStore: The File Search store object
|
|
1557
|
+
"""
|
|
1558
|
+
try:
|
|
1559
|
+
store = self.get_client().file_search_stores.get(name=name)
|
|
1560
|
+
log_debug(f"Retrieved File Search store: {name}")
|
|
1561
|
+
return store
|
|
1562
|
+
except Exception as e:
|
|
1563
|
+
log_error(f"Error getting File Search store {name}: {e}")
|
|
1564
|
+
raise
|
|
1565
|
+
|
|
1566
|
+
async def async_get_file_search_store(self, name: str) -> Any:
|
|
1567
|
+
"""
|
|
1568
|
+
Args:
|
|
1569
|
+
name: The name of the store
|
|
1570
|
+
|
|
1571
|
+
Returns:
|
|
1572
|
+
FileSearchStore: The File Search store object
|
|
1573
|
+
"""
|
|
1574
|
+
try:
|
|
1575
|
+
store = await self.get_client().aio.file_search_stores.get(name=name)
|
|
1576
|
+
log_debug(f"Retrieved File Search store: {name}")
|
|
1577
|
+
return store
|
|
1578
|
+
except Exception as e:
|
|
1579
|
+
log_error(f"Error getting File Search store {name}: {e}")
|
|
1580
|
+
raise
|
|
1581
|
+
|
|
1582
|
+
def delete_file_search_store(self, name: str, force: bool = False) -> None:
|
|
1583
|
+
"""
|
|
1584
|
+
Delete a File Search store.
|
|
1585
|
+
|
|
1586
|
+
Args:
|
|
1587
|
+
name: The name of the store to delete
|
|
1588
|
+
force: If True, force delete even if store contains documents
|
|
1589
|
+
"""
|
|
1590
|
+
try:
|
|
1591
|
+
self.get_client().file_search_stores.delete(name=name, config={"force": force})
|
|
1592
|
+
log_info(f"Deleted File Search store: {name}")
|
|
1593
|
+
except Exception as e:
|
|
1594
|
+
log_error(f"Error deleting File Search store {name}: {e}")
|
|
1595
|
+
raise
|
|
1596
|
+
|
|
1597
|
+
async def async_delete_file_search_store(self, name: str, force: bool = True) -> None:
|
|
1598
|
+
"""
|
|
1599
|
+
Async version of delete_file_search_store.
|
|
1600
|
+
|
|
1601
|
+
Args:
|
|
1602
|
+
name: The name of the store to delete
|
|
1603
|
+
force: If True, force delete even if store contains documents
|
|
1604
|
+
"""
|
|
1605
|
+
try:
|
|
1606
|
+
await self.get_client().aio.file_search_stores.delete(name=name, config={"force": force})
|
|
1607
|
+
log_info(f"Deleted File Search store: {name}")
|
|
1608
|
+
except Exception as e:
|
|
1609
|
+
log_error(f"Error deleting File Search store {name}: {e}")
|
|
1610
|
+
raise
|
|
1611
|
+
|
|
1612
|
+
def wait_for_operation(self, operation: Operation, poll_interval: int = 5, max_wait: int = 600) -> Operation:
|
|
1613
|
+
"""
|
|
1614
|
+
Wait for a long-running operation to complete.
|
|
1615
|
+
|
|
1616
|
+
Args:
|
|
1617
|
+
operation: The operation object to wait for
|
|
1618
|
+
poll_interval: Seconds to wait between status checks
|
|
1619
|
+
max_wait: Maximum seconds to wait before timing out
|
|
1620
|
+
|
|
1621
|
+
Returns:
|
|
1622
|
+
Operation: The completed operation object
|
|
1623
|
+
|
|
1624
|
+
Raises:
|
|
1625
|
+
TimeoutError: If operation doesn't complete within max_wait seconds
|
|
1626
|
+
"""
|
|
1627
|
+
elapsed = 0
|
|
1628
|
+
while not operation.done:
|
|
1629
|
+
if elapsed >= max_wait:
|
|
1630
|
+
raise TimeoutError(f"Operation timed out after {max_wait} seconds")
|
|
1631
|
+
time.sleep(poll_interval)
|
|
1632
|
+
elapsed += poll_interval
|
|
1633
|
+
operation = self.get_client().operations.get(operation)
|
|
1634
|
+
log_debug(f"Waiting for operation... ({elapsed}s elapsed)")
|
|
1635
|
+
|
|
1636
|
+
log_info("Operation completed successfully")
|
|
1637
|
+
return operation
|
|
1638
|
+
|
|
1639
|
+
async def async_wait_for_operation(
|
|
1640
|
+
self, operation: Operation, poll_interval: int = 5, max_wait: int = 600
|
|
1641
|
+
) -> Operation:
|
|
1642
|
+
"""
|
|
1643
|
+
Async version of wait_for_operation.
|
|
1644
|
+
|
|
1645
|
+
Args:
|
|
1646
|
+
operation: The operation object to wait for
|
|
1647
|
+
poll_interval: Seconds to wait between status checks
|
|
1648
|
+
max_wait: Maximum seconds to wait before timing out
|
|
1649
|
+
|
|
1650
|
+
Returns:
|
|
1651
|
+
Operation: The completed operation object
|
|
1652
|
+
"""
|
|
1653
|
+
elapsed = 0
|
|
1654
|
+
while not operation.done:
|
|
1655
|
+
if elapsed >= max_wait:
|
|
1656
|
+
raise TimeoutError(f"Operation timed out after {max_wait} seconds")
|
|
1657
|
+
await asyncio.sleep(poll_interval)
|
|
1658
|
+
elapsed += poll_interval
|
|
1659
|
+
operation = await self.get_client().aio.operations.get(operation)
|
|
1660
|
+
log_debug(f"Waiting for operation... ({elapsed}s elapsed)")
|
|
1661
|
+
|
|
1662
|
+
log_info("Operation completed successfully")
|
|
1663
|
+
return operation
|
|
1664
|
+
|
|
1665
|
+
def upload_to_file_search_store(
|
|
1666
|
+
self,
|
|
1667
|
+
file_path: Union[str, Path],
|
|
1668
|
+
store_name: str,
|
|
1669
|
+
display_name: Optional[str] = None,
|
|
1670
|
+
chunking_config: Optional[Dict[str, Any]] = None,
|
|
1671
|
+
custom_metadata: Optional[List[Dict[str, Any]]] = None,
|
|
1672
|
+
) -> Any:
|
|
1673
|
+
"""
|
|
1674
|
+
Upload a file directly to a File Search store.
|
|
1675
|
+
|
|
1676
|
+
Args:
|
|
1677
|
+
file_path: Path to the file to upload
|
|
1678
|
+
store_name: Name of the File Search store
|
|
1679
|
+
display_name: Optional display name for the file (will be visible in citations)
|
|
1680
|
+
chunking_config: Optional chunking configuration
|
|
1681
|
+
Example: {
|
|
1682
|
+
"white_space_config": {
|
|
1683
|
+
"max_tokens_per_chunk": 200,
|
|
1684
|
+
"max_overlap_tokens": 20
|
|
1685
|
+
}
|
|
1686
|
+
}
|
|
1687
|
+
custom_metadata: Optional custom metadata as list of dicts
|
|
1688
|
+
Example: [
|
|
1689
|
+
{"key": "author", "string_value": "John Doe"},
|
|
1690
|
+
{"key": "year", "numeric_value": 2024}
|
|
1691
|
+
]
|
|
1692
|
+
|
|
1693
|
+
Returns:
|
|
1694
|
+
Operation: Long-running operation object. Use wait_for_operation() to wait for completion.
|
|
1695
|
+
"""
|
|
1696
|
+
file_path = file_path if isinstance(file_path, Path) else Path(file_path)
|
|
1697
|
+
|
|
1698
|
+
if not file_path.exists():
|
|
1699
|
+
raise FileNotFoundError(f"File not found: {file_path}")
|
|
1700
|
+
|
|
1701
|
+
config: Dict[str, Any] = {}
|
|
1702
|
+
if display_name:
|
|
1703
|
+
config["display_name"] = display_name
|
|
1704
|
+
if chunking_config:
|
|
1705
|
+
config["chunking_config"] = chunking_config
|
|
1706
|
+
if custom_metadata:
|
|
1707
|
+
config["custom_metadata"] = custom_metadata
|
|
1708
|
+
|
|
1709
|
+
try:
|
|
1710
|
+
log_info(f"Uploading file {file_path.name} to File Search store {store_name}")
|
|
1711
|
+
operation = self.get_client().file_search_stores.upload_to_file_search_store(
|
|
1712
|
+
file=file_path,
|
|
1713
|
+
file_search_store_name=store_name,
|
|
1714
|
+
config=config or None, # type: ignore[arg-type]
|
|
1715
|
+
)
|
|
1716
|
+
log_info(f"Upload initiated for {file_path.name}")
|
|
1717
|
+
return operation
|
|
1718
|
+
except Exception as e:
|
|
1719
|
+
log_error(f"Error uploading file to File Search store: {e}")
|
|
1720
|
+
raise
|
|
1721
|
+
|
|
1722
|
+
async def async_upload_to_file_search_store(
|
|
1723
|
+
self,
|
|
1724
|
+
file_path: Union[str, Path],
|
|
1725
|
+
store_name: str,
|
|
1726
|
+
display_name: Optional[str] = None,
|
|
1727
|
+
chunking_config: Optional[Dict[str, Any]] = None,
|
|
1728
|
+
custom_metadata: Optional[List[Dict[str, Any]]] = None,
|
|
1729
|
+
) -> Any:
|
|
1730
|
+
"""
|
|
1731
|
+
Args:
|
|
1732
|
+
file_path: Path to the file to upload
|
|
1733
|
+
store_name: Name of the File Search store
|
|
1734
|
+
display_name: Optional display name for the file
|
|
1735
|
+
chunking_config: Optional chunking configuration
|
|
1736
|
+
custom_metadata: Optional custom metadata
|
|
1737
|
+
|
|
1738
|
+
Returns:
|
|
1739
|
+
Operation: Long-running operation object
|
|
1740
|
+
"""
|
|
1741
|
+
file_path = file_path if isinstance(file_path, Path) else Path(file_path)
|
|
1742
|
+
|
|
1743
|
+
if not file_path.exists():
|
|
1744
|
+
raise FileNotFoundError(f"File not found: {file_path}")
|
|
1745
|
+
|
|
1746
|
+
config: Dict[str, Any] = {}
|
|
1747
|
+
if display_name:
|
|
1748
|
+
config["display_name"] = display_name
|
|
1749
|
+
if chunking_config:
|
|
1750
|
+
config["chunking_config"] = chunking_config
|
|
1751
|
+
if custom_metadata:
|
|
1752
|
+
config["custom_metadata"] = custom_metadata
|
|
1753
|
+
|
|
1754
|
+
try:
|
|
1755
|
+
log_info(f"Uploading file {file_path.name} to File Search store {store_name}")
|
|
1756
|
+
operation = await self.get_client().aio.file_search_stores.upload_to_file_search_store(
|
|
1757
|
+
file=file_path,
|
|
1758
|
+
file_search_store_name=store_name,
|
|
1759
|
+
config=config or None, # type: ignore[arg-type]
|
|
1760
|
+
)
|
|
1761
|
+
log_info(f"Upload initiated for {file_path.name}")
|
|
1762
|
+
return operation
|
|
1763
|
+
except Exception as e:
|
|
1764
|
+
log_error(f"Error uploading file to File Search store: {e}")
|
|
1765
|
+
raise
|
|
1766
|
+
|
|
1767
|
+
def import_file_to_store(
|
|
1768
|
+
self,
|
|
1769
|
+
file_name: str,
|
|
1770
|
+
store_name: str,
|
|
1771
|
+
chunking_config: Optional[Dict[str, Any]] = None,
|
|
1772
|
+
custom_metadata: Optional[List[Dict[str, Any]]] = None,
|
|
1773
|
+
) -> Any:
|
|
1774
|
+
"""
|
|
1775
|
+
Import an existing uploaded file (via Files API) into a File Search store.
|
|
1776
|
+
|
|
1777
|
+
Args:
|
|
1778
|
+
file_name: Name of the file already uploaded via Files API
|
|
1779
|
+
store_name: Name of the File Search store
|
|
1780
|
+
chunking_config: Optional chunking configuration
|
|
1781
|
+
custom_metadata: Optional custom metadata
|
|
1782
|
+
|
|
1783
|
+
Returns:
|
|
1784
|
+
Operation: Long-running operation object. Use wait_for_operation() to wait for completion.
|
|
1785
|
+
"""
|
|
1786
|
+
config: Dict[str, Any] = {}
|
|
1787
|
+
if chunking_config:
|
|
1788
|
+
config["chunking_config"] = chunking_config
|
|
1789
|
+
if custom_metadata:
|
|
1790
|
+
config["custom_metadata"] = custom_metadata
|
|
1791
|
+
|
|
1792
|
+
try:
|
|
1793
|
+
log_info(f"Importing file {file_name} to File Search store {store_name}")
|
|
1794
|
+
operation = self.get_client().file_search_stores.import_file(
|
|
1795
|
+
file_search_store_name=store_name,
|
|
1796
|
+
file_name=file_name,
|
|
1797
|
+
config=config or None, # type: ignore[arg-type]
|
|
1798
|
+
)
|
|
1799
|
+
log_info(f"Import initiated for {file_name}")
|
|
1800
|
+
return operation
|
|
1801
|
+
except Exception as e:
|
|
1802
|
+
log_error(f"Error importing file to File Search store: {e}")
|
|
1803
|
+
raise
|
|
1804
|
+
|
|
1805
|
+
async def async_import_file_to_store(
|
|
1806
|
+
self,
|
|
1807
|
+
file_name: str,
|
|
1808
|
+
store_name: str,
|
|
1809
|
+
chunking_config: Optional[Dict[str, Any]] = None,
|
|
1810
|
+
custom_metadata: Optional[List[Dict[str, Any]]] = None,
|
|
1811
|
+
) -> Any:
|
|
1812
|
+
"""
|
|
1813
|
+
Args:
|
|
1814
|
+
file_name: Name of the file already uploaded via Files API
|
|
1815
|
+
store_name: Name of the File Search store
|
|
1816
|
+
chunking_config: Optional chunking configuration
|
|
1817
|
+
custom_metadata: Optional custom metadata
|
|
1818
|
+
|
|
1819
|
+
Returns:
|
|
1820
|
+
Operation: Long-running operation object
|
|
1821
|
+
"""
|
|
1822
|
+
config: Dict[str, Any] = {}
|
|
1823
|
+
if chunking_config:
|
|
1824
|
+
config["chunking_config"] = chunking_config
|
|
1825
|
+
if custom_metadata:
|
|
1826
|
+
config["custom_metadata"] = custom_metadata
|
|
1827
|
+
|
|
1828
|
+
try:
|
|
1829
|
+
log_info(f"Importing file {file_name} to File Search store {store_name}")
|
|
1830
|
+
operation = await self.get_client().aio.file_search_stores.import_file(
|
|
1831
|
+
file_search_store_name=store_name,
|
|
1832
|
+
file_name=file_name,
|
|
1833
|
+
config=config or None, # type: ignore[arg-type]
|
|
1834
|
+
)
|
|
1835
|
+
log_info(f"Import initiated for {file_name}")
|
|
1836
|
+
return operation
|
|
1837
|
+
except Exception as e:
|
|
1838
|
+
log_error(f"Error importing file to File Search store: {e}")
|
|
1839
|
+
raise
|
|
1840
|
+
|
|
1841
|
+
def list_documents(self, store_name: str, page_size: int = 20) -> List[Any]:
|
|
1842
|
+
"""
|
|
1843
|
+
Args:
|
|
1844
|
+
store_name: Name of the File Search store
|
|
1845
|
+
page_size: Maximum number of documents to return per page
|
|
1846
|
+
|
|
1847
|
+
Returns:
|
|
1848
|
+
List: List of document objects
|
|
1849
|
+
"""
|
|
1850
|
+
try:
|
|
1851
|
+
documents = []
|
|
1852
|
+
for doc in self.get_client().file_search_stores.documents.list(
|
|
1853
|
+
parent=store_name, config={"page_size": page_size}
|
|
1854
|
+
):
|
|
1855
|
+
documents.append(doc)
|
|
1856
|
+
log_debug(f"Found {len(documents)} documents in store {store_name}")
|
|
1857
|
+
return documents
|
|
1858
|
+
except Exception as e:
|
|
1859
|
+
log_error(f"Error listing documents in store {store_name}: {e}")
|
|
1860
|
+
raise
|
|
1861
|
+
|
|
1862
|
+
async def async_list_documents(self, store_name: str, page_size: int = 20) -> List[Any]:
|
|
1863
|
+
"""
|
|
1864
|
+
Async version of list_documents.
|
|
1865
|
+
|
|
1866
|
+
Args:
|
|
1867
|
+
store_name: Name of the File Search store
|
|
1868
|
+
page_size: Maximum number of documents to return per page
|
|
1869
|
+
|
|
1870
|
+
Returns:
|
|
1871
|
+
List: List of document objects
|
|
1872
|
+
"""
|
|
1873
|
+
try:
|
|
1874
|
+
documents = []
|
|
1875
|
+
# Await the AsyncPager first, then iterate
|
|
1876
|
+
async for doc in await self.get_client().aio.file_search_stores.documents.list(
|
|
1877
|
+
parent=store_name, config={"page_size": page_size}
|
|
1878
|
+
):
|
|
1879
|
+
documents.append(doc)
|
|
1880
|
+
log_debug(f"Found {len(documents)} documents in store {store_name}")
|
|
1881
|
+
return documents
|
|
1882
|
+
except Exception as e:
|
|
1883
|
+
log_error(f"Error listing documents in store {store_name}: {e}")
|
|
1884
|
+
raise
|
|
1885
|
+
|
|
1886
|
+
def get_document(self, document_name: str) -> Any:
|
|
1887
|
+
"""
|
|
1888
|
+
Get a specific document by name.
|
|
1889
|
+
|
|
1890
|
+
Args:
|
|
1891
|
+
document_name: Full name of the document
|
|
1892
|
+
(e.g., 'fileSearchStores/store-123/documents/doc-456')
|
|
1893
|
+
|
|
1894
|
+
Returns:
|
|
1895
|
+
Document object
|
|
1896
|
+
"""
|
|
1897
|
+
try:
|
|
1898
|
+
doc = self.get_client().file_search_stores.documents.get(name=document_name)
|
|
1899
|
+
log_debug(f"Retrieved document: {document_name}")
|
|
1900
|
+
return doc
|
|
1901
|
+
except Exception as e:
|
|
1902
|
+
log_error(f"Error getting document {document_name}: {e}")
|
|
1903
|
+
raise
|
|
1904
|
+
|
|
1905
|
+
async def async_get_document(self, document_name: str) -> Any:
|
|
1906
|
+
"""
|
|
1907
|
+
Async version of get_document.
|
|
1908
|
+
|
|
1909
|
+
Args:
|
|
1910
|
+
document_name: Full name of the document
|
|
1911
|
+
|
|
1912
|
+
Returns:
|
|
1913
|
+
Document object
|
|
1914
|
+
"""
|
|
1915
|
+
try:
|
|
1916
|
+
doc = await self.get_client().aio.file_search_stores.documents.get(name=document_name)
|
|
1917
|
+
log_debug(f"Retrieved document: {document_name}")
|
|
1918
|
+
return doc
|
|
1919
|
+
except Exception as e:
|
|
1920
|
+
log_error(f"Error getting document {document_name}: {e}")
|
|
1921
|
+
raise
|
|
1922
|
+
|
|
1923
|
+
def delete_document(self, document_name: str) -> None:
|
|
1924
|
+
"""
|
|
1925
|
+
Delete a document from a File Search store.
|
|
1926
|
+
|
|
1927
|
+
Args:
|
|
1928
|
+
document_name: Full name of the document to delete
|
|
1929
|
+
|
|
1930
|
+
Example:
|
|
1931
|
+
```python
|
|
1932
|
+
model = Gemini(id="gemini-2.5-flash")
|
|
1933
|
+
model.delete_document("fileSearchStores/store-123/documents/doc-456")
|
|
1934
|
+
```
|
|
1935
|
+
"""
|
|
1936
|
+
try:
|
|
1937
|
+
self.get_client().file_search_stores.documents.delete(name=document_name)
|
|
1938
|
+
log_info(f"Deleted document: {document_name}")
|
|
1939
|
+
except Exception as e:
|
|
1940
|
+
log_error(f"Error deleting document {document_name}: {e}")
|
|
1941
|
+
raise
|
|
1942
|
+
|
|
1943
|
+
async def async_delete_document(self, document_name: str) -> None:
|
|
1944
|
+
"""
|
|
1945
|
+
Async version of delete_document.
|
|
1946
|
+
|
|
1947
|
+
Args:
|
|
1948
|
+
document_name: Full name of the document to delete
|
|
1949
|
+
"""
|
|
1950
|
+
try:
|
|
1951
|
+
await self.get_client().aio.file_search_stores.documents.delete(name=document_name)
|
|
1952
|
+
log_info(f"Deleted document: {document_name}")
|
|
1953
|
+
except Exception as e:
|
|
1954
|
+
log_error(f"Error deleting document {document_name}: {e}")
|
|
1955
|
+
raise
|