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.
Files changed (383) hide show
  1. agno/agent/__init__.py +6 -0
  2. agno/agent/agent.py +5252 -3145
  3. agno/agent/remote.py +525 -0
  4. agno/api/api.py +2 -0
  5. agno/client/__init__.py +3 -0
  6. agno/client/a2a/__init__.py +10 -0
  7. agno/client/a2a/client.py +554 -0
  8. agno/client/a2a/schemas.py +112 -0
  9. agno/client/a2a/utils.py +369 -0
  10. agno/client/os.py +2669 -0
  11. agno/compression/__init__.py +3 -0
  12. agno/compression/manager.py +247 -0
  13. agno/culture/manager.py +2 -2
  14. agno/db/base.py +927 -6
  15. agno/db/dynamo/dynamo.py +788 -2
  16. agno/db/dynamo/schemas.py +128 -0
  17. agno/db/dynamo/utils.py +26 -3
  18. agno/db/firestore/firestore.py +674 -50
  19. agno/db/firestore/schemas.py +41 -0
  20. agno/db/firestore/utils.py +25 -10
  21. agno/db/gcs_json/gcs_json_db.py +506 -3
  22. agno/db/gcs_json/utils.py +14 -2
  23. agno/db/in_memory/in_memory_db.py +203 -4
  24. agno/db/in_memory/utils.py +14 -2
  25. agno/db/json/json_db.py +498 -2
  26. agno/db/json/utils.py +14 -2
  27. agno/db/migrations/manager.py +199 -0
  28. agno/db/migrations/utils.py +19 -0
  29. agno/db/migrations/v1_to_v2.py +54 -16
  30. agno/db/migrations/versions/__init__.py +0 -0
  31. agno/db/migrations/versions/v2_3_0.py +977 -0
  32. agno/db/mongo/async_mongo.py +1013 -39
  33. agno/db/mongo/mongo.py +684 -4
  34. agno/db/mongo/schemas.py +48 -0
  35. agno/db/mongo/utils.py +17 -0
  36. agno/db/mysql/__init__.py +2 -1
  37. agno/db/mysql/async_mysql.py +2958 -0
  38. agno/db/mysql/mysql.py +722 -53
  39. agno/db/mysql/schemas.py +77 -11
  40. agno/db/mysql/utils.py +151 -8
  41. agno/db/postgres/async_postgres.py +1254 -137
  42. agno/db/postgres/postgres.py +2316 -93
  43. agno/db/postgres/schemas.py +153 -21
  44. agno/db/postgres/utils.py +22 -7
  45. agno/db/redis/redis.py +531 -3
  46. agno/db/redis/schemas.py +36 -0
  47. agno/db/redis/utils.py +31 -15
  48. agno/db/schemas/evals.py +1 -0
  49. agno/db/schemas/memory.py +20 -9
  50. agno/db/singlestore/schemas.py +70 -1
  51. agno/db/singlestore/singlestore.py +737 -74
  52. agno/db/singlestore/utils.py +13 -3
  53. agno/db/sqlite/async_sqlite.py +1069 -89
  54. agno/db/sqlite/schemas.py +133 -1
  55. agno/db/sqlite/sqlite.py +2203 -165
  56. agno/db/sqlite/utils.py +21 -11
  57. agno/db/surrealdb/models.py +25 -0
  58. agno/db/surrealdb/surrealdb.py +603 -1
  59. agno/db/utils.py +60 -0
  60. agno/eval/__init__.py +26 -3
  61. agno/eval/accuracy.py +25 -12
  62. agno/eval/agent_as_judge.py +871 -0
  63. agno/eval/base.py +29 -0
  64. agno/eval/performance.py +10 -4
  65. agno/eval/reliability.py +22 -13
  66. agno/eval/utils.py +2 -1
  67. agno/exceptions.py +42 -0
  68. agno/hooks/__init__.py +3 -0
  69. agno/hooks/decorator.py +164 -0
  70. agno/integrations/discord/client.py +13 -2
  71. agno/knowledge/__init__.py +4 -0
  72. agno/knowledge/chunking/code.py +90 -0
  73. agno/knowledge/chunking/document.py +65 -4
  74. agno/knowledge/chunking/fixed.py +4 -1
  75. agno/knowledge/chunking/markdown.py +102 -11
  76. agno/knowledge/chunking/recursive.py +2 -2
  77. agno/knowledge/chunking/semantic.py +130 -48
  78. agno/knowledge/chunking/strategy.py +18 -0
  79. agno/knowledge/embedder/azure_openai.py +0 -1
  80. agno/knowledge/embedder/google.py +1 -1
  81. agno/knowledge/embedder/mistral.py +1 -1
  82. agno/knowledge/embedder/nebius.py +1 -1
  83. agno/knowledge/embedder/openai.py +16 -12
  84. agno/knowledge/filesystem.py +412 -0
  85. agno/knowledge/knowledge.py +4261 -1199
  86. agno/knowledge/protocol.py +134 -0
  87. agno/knowledge/reader/arxiv_reader.py +3 -2
  88. agno/knowledge/reader/base.py +9 -7
  89. agno/knowledge/reader/csv_reader.py +91 -42
  90. agno/knowledge/reader/docx_reader.py +9 -10
  91. agno/knowledge/reader/excel_reader.py +225 -0
  92. agno/knowledge/reader/field_labeled_csv_reader.py +38 -48
  93. agno/knowledge/reader/firecrawl_reader.py +3 -2
  94. agno/knowledge/reader/json_reader.py +16 -22
  95. agno/knowledge/reader/markdown_reader.py +15 -14
  96. agno/knowledge/reader/pdf_reader.py +33 -28
  97. agno/knowledge/reader/pptx_reader.py +9 -10
  98. agno/knowledge/reader/reader_factory.py +135 -1
  99. agno/knowledge/reader/s3_reader.py +8 -16
  100. agno/knowledge/reader/tavily_reader.py +3 -3
  101. agno/knowledge/reader/text_reader.py +15 -14
  102. agno/knowledge/reader/utils/__init__.py +17 -0
  103. agno/knowledge/reader/utils/spreadsheet.py +114 -0
  104. agno/knowledge/reader/web_search_reader.py +8 -65
  105. agno/knowledge/reader/website_reader.py +16 -13
  106. agno/knowledge/reader/wikipedia_reader.py +36 -3
  107. agno/knowledge/reader/youtube_reader.py +3 -2
  108. agno/knowledge/remote_content/__init__.py +33 -0
  109. agno/knowledge/remote_content/config.py +266 -0
  110. agno/knowledge/remote_content/remote_content.py +105 -17
  111. agno/knowledge/utils.py +76 -22
  112. agno/learn/__init__.py +71 -0
  113. agno/learn/config.py +463 -0
  114. agno/learn/curate.py +185 -0
  115. agno/learn/machine.py +725 -0
  116. agno/learn/schemas.py +1114 -0
  117. agno/learn/stores/__init__.py +38 -0
  118. agno/learn/stores/decision_log.py +1156 -0
  119. agno/learn/stores/entity_memory.py +3275 -0
  120. agno/learn/stores/learned_knowledge.py +1583 -0
  121. agno/learn/stores/protocol.py +117 -0
  122. agno/learn/stores/session_context.py +1217 -0
  123. agno/learn/stores/user_memory.py +1495 -0
  124. agno/learn/stores/user_profile.py +1220 -0
  125. agno/learn/utils.py +209 -0
  126. agno/media.py +22 -6
  127. agno/memory/__init__.py +14 -1
  128. agno/memory/manager.py +223 -8
  129. agno/memory/strategies/__init__.py +15 -0
  130. agno/memory/strategies/base.py +66 -0
  131. agno/memory/strategies/summarize.py +196 -0
  132. agno/memory/strategies/types.py +37 -0
  133. agno/models/aimlapi/aimlapi.py +17 -0
  134. agno/models/anthropic/claude.py +434 -59
  135. agno/models/aws/bedrock.py +121 -20
  136. agno/models/aws/claude.py +131 -274
  137. agno/models/azure/ai_foundry.py +10 -6
  138. agno/models/azure/openai_chat.py +33 -10
  139. agno/models/base.py +1162 -561
  140. agno/models/cerebras/cerebras.py +120 -24
  141. agno/models/cerebras/cerebras_openai.py +21 -2
  142. agno/models/cohere/chat.py +65 -6
  143. agno/models/cometapi/cometapi.py +18 -1
  144. agno/models/dashscope/dashscope.py +2 -3
  145. agno/models/deepinfra/deepinfra.py +18 -1
  146. agno/models/deepseek/deepseek.py +69 -3
  147. agno/models/fireworks/fireworks.py +18 -1
  148. agno/models/google/gemini.py +959 -89
  149. agno/models/google/utils.py +22 -0
  150. agno/models/groq/groq.py +48 -18
  151. agno/models/huggingface/huggingface.py +17 -6
  152. agno/models/ibm/watsonx.py +16 -6
  153. agno/models/internlm/internlm.py +18 -1
  154. agno/models/langdb/langdb.py +13 -1
  155. agno/models/litellm/chat.py +88 -9
  156. agno/models/litellm/litellm_openai.py +18 -1
  157. agno/models/message.py +24 -5
  158. agno/models/meta/llama.py +40 -13
  159. agno/models/meta/llama_openai.py +22 -21
  160. agno/models/metrics.py +12 -0
  161. agno/models/mistral/mistral.py +8 -4
  162. agno/models/n1n/__init__.py +3 -0
  163. agno/models/n1n/n1n.py +57 -0
  164. agno/models/nebius/nebius.py +6 -7
  165. agno/models/nvidia/nvidia.py +20 -3
  166. agno/models/ollama/__init__.py +2 -0
  167. agno/models/ollama/chat.py +17 -6
  168. agno/models/ollama/responses.py +100 -0
  169. agno/models/openai/__init__.py +2 -0
  170. agno/models/openai/chat.py +117 -26
  171. agno/models/openai/open_responses.py +46 -0
  172. agno/models/openai/responses.py +110 -32
  173. agno/models/openrouter/__init__.py +2 -0
  174. agno/models/openrouter/openrouter.py +67 -2
  175. agno/models/openrouter/responses.py +146 -0
  176. agno/models/perplexity/perplexity.py +19 -1
  177. agno/models/portkey/portkey.py +7 -6
  178. agno/models/requesty/requesty.py +19 -2
  179. agno/models/response.py +20 -2
  180. agno/models/sambanova/sambanova.py +20 -3
  181. agno/models/siliconflow/siliconflow.py +19 -2
  182. agno/models/together/together.py +20 -3
  183. agno/models/vercel/v0.py +20 -3
  184. agno/models/vertexai/claude.py +124 -4
  185. agno/models/vllm/vllm.py +19 -14
  186. agno/models/xai/xai.py +19 -2
  187. agno/os/app.py +467 -137
  188. agno/os/auth.py +253 -5
  189. agno/os/config.py +22 -0
  190. agno/os/interfaces/a2a/a2a.py +7 -6
  191. agno/os/interfaces/a2a/router.py +635 -26
  192. agno/os/interfaces/a2a/utils.py +32 -33
  193. agno/os/interfaces/agui/agui.py +5 -3
  194. agno/os/interfaces/agui/router.py +26 -16
  195. agno/os/interfaces/agui/utils.py +97 -57
  196. agno/os/interfaces/base.py +7 -7
  197. agno/os/interfaces/slack/router.py +16 -7
  198. agno/os/interfaces/slack/slack.py +7 -7
  199. agno/os/interfaces/whatsapp/router.py +35 -7
  200. agno/os/interfaces/whatsapp/security.py +3 -1
  201. agno/os/interfaces/whatsapp/whatsapp.py +11 -8
  202. agno/os/managers.py +326 -0
  203. agno/os/mcp.py +652 -79
  204. agno/os/middleware/__init__.py +4 -0
  205. agno/os/middleware/jwt.py +718 -115
  206. agno/os/middleware/trailing_slash.py +27 -0
  207. agno/os/router.py +105 -1558
  208. agno/os/routers/agents/__init__.py +3 -0
  209. agno/os/routers/agents/router.py +655 -0
  210. agno/os/routers/agents/schema.py +288 -0
  211. agno/os/routers/components/__init__.py +3 -0
  212. agno/os/routers/components/components.py +475 -0
  213. agno/os/routers/database.py +155 -0
  214. agno/os/routers/evals/evals.py +111 -18
  215. agno/os/routers/evals/schemas.py +38 -5
  216. agno/os/routers/evals/utils.py +80 -11
  217. agno/os/routers/health.py +3 -3
  218. agno/os/routers/knowledge/knowledge.py +284 -35
  219. agno/os/routers/knowledge/schemas.py +14 -2
  220. agno/os/routers/memory/memory.py +274 -11
  221. agno/os/routers/memory/schemas.py +44 -3
  222. agno/os/routers/metrics/metrics.py +30 -15
  223. agno/os/routers/metrics/schemas.py +10 -6
  224. agno/os/routers/registry/__init__.py +3 -0
  225. agno/os/routers/registry/registry.py +337 -0
  226. agno/os/routers/session/session.py +143 -14
  227. agno/os/routers/teams/__init__.py +3 -0
  228. agno/os/routers/teams/router.py +550 -0
  229. agno/os/routers/teams/schema.py +280 -0
  230. agno/os/routers/traces/__init__.py +3 -0
  231. agno/os/routers/traces/schemas.py +414 -0
  232. agno/os/routers/traces/traces.py +549 -0
  233. agno/os/routers/workflows/__init__.py +3 -0
  234. agno/os/routers/workflows/router.py +757 -0
  235. agno/os/routers/workflows/schema.py +139 -0
  236. agno/os/schema.py +157 -584
  237. agno/os/scopes.py +469 -0
  238. agno/os/settings.py +3 -0
  239. agno/os/utils.py +574 -185
  240. agno/reasoning/anthropic.py +85 -1
  241. agno/reasoning/azure_ai_foundry.py +93 -1
  242. agno/reasoning/deepseek.py +102 -2
  243. agno/reasoning/default.py +6 -7
  244. agno/reasoning/gemini.py +87 -3
  245. agno/reasoning/groq.py +109 -2
  246. agno/reasoning/helpers.py +6 -7
  247. agno/reasoning/manager.py +1238 -0
  248. agno/reasoning/ollama.py +93 -1
  249. agno/reasoning/openai.py +115 -1
  250. agno/reasoning/vertexai.py +85 -1
  251. agno/registry/__init__.py +3 -0
  252. agno/registry/registry.py +68 -0
  253. agno/remote/__init__.py +3 -0
  254. agno/remote/base.py +581 -0
  255. agno/run/__init__.py +2 -4
  256. agno/run/agent.py +134 -19
  257. agno/run/base.py +49 -1
  258. agno/run/cancel.py +65 -52
  259. agno/run/cancellation_management/__init__.py +9 -0
  260. agno/run/cancellation_management/base.py +78 -0
  261. agno/run/cancellation_management/in_memory_cancellation_manager.py +100 -0
  262. agno/run/cancellation_management/redis_cancellation_manager.py +236 -0
  263. agno/run/requirement.py +181 -0
  264. agno/run/team.py +111 -19
  265. agno/run/workflow.py +2 -1
  266. agno/session/agent.py +57 -92
  267. agno/session/summary.py +1 -1
  268. agno/session/team.py +62 -115
  269. agno/session/workflow.py +353 -57
  270. agno/skills/__init__.py +17 -0
  271. agno/skills/agent_skills.py +377 -0
  272. agno/skills/errors.py +32 -0
  273. agno/skills/loaders/__init__.py +4 -0
  274. agno/skills/loaders/base.py +27 -0
  275. agno/skills/loaders/local.py +216 -0
  276. agno/skills/skill.py +65 -0
  277. agno/skills/utils.py +107 -0
  278. agno/skills/validator.py +277 -0
  279. agno/table.py +10 -0
  280. agno/team/__init__.py +5 -1
  281. agno/team/remote.py +447 -0
  282. agno/team/team.py +3769 -2202
  283. agno/tools/brandfetch.py +27 -18
  284. agno/tools/browserbase.py +225 -16
  285. agno/tools/crawl4ai.py +3 -0
  286. agno/tools/duckduckgo.py +25 -71
  287. agno/tools/exa.py +0 -21
  288. agno/tools/file.py +14 -13
  289. agno/tools/file_generation.py +12 -6
  290. agno/tools/firecrawl.py +15 -7
  291. agno/tools/function.py +94 -113
  292. agno/tools/google_bigquery.py +11 -2
  293. agno/tools/google_drive.py +4 -3
  294. agno/tools/knowledge.py +9 -4
  295. agno/tools/mcp/mcp.py +301 -18
  296. agno/tools/mcp/multi_mcp.py +269 -14
  297. agno/tools/mem0.py +11 -10
  298. agno/tools/memory.py +47 -46
  299. agno/tools/mlx_transcribe.py +10 -7
  300. agno/tools/models/nebius.py +5 -5
  301. agno/tools/models_labs.py +20 -10
  302. agno/tools/nano_banana.py +151 -0
  303. agno/tools/parallel.py +0 -7
  304. agno/tools/postgres.py +76 -36
  305. agno/tools/python.py +14 -6
  306. agno/tools/reasoning.py +30 -23
  307. agno/tools/redshift.py +406 -0
  308. agno/tools/shopify.py +1519 -0
  309. agno/tools/spotify.py +919 -0
  310. agno/tools/tavily.py +4 -1
  311. agno/tools/toolkit.py +253 -18
  312. agno/tools/websearch.py +93 -0
  313. agno/tools/website.py +1 -1
  314. agno/tools/wikipedia.py +1 -1
  315. agno/tools/workflow.py +56 -48
  316. agno/tools/yfinance.py +12 -11
  317. agno/tracing/__init__.py +12 -0
  318. agno/tracing/exporter.py +161 -0
  319. agno/tracing/schemas.py +276 -0
  320. agno/tracing/setup.py +112 -0
  321. agno/utils/agent.py +251 -10
  322. agno/utils/cryptography.py +22 -0
  323. agno/utils/dttm.py +33 -0
  324. agno/utils/events.py +264 -7
  325. agno/utils/hooks.py +111 -3
  326. agno/utils/http.py +161 -2
  327. agno/utils/mcp.py +49 -8
  328. agno/utils/media.py +22 -1
  329. agno/utils/models/ai_foundry.py +9 -2
  330. agno/utils/models/claude.py +20 -5
  331. agno/utils/models/cohere.py +9 -2
  332. agno/utils/models/llama.py +9 -2
  333. agno/utils/models/mistral.py +4 -2
  334. agno/utils/os.py +0 -0
  335. agno/utils/print_response/agent.py +99 -16
  336. agno/utils/print_response/team.py +223 -24
  337. agno/utils/print_response/workflow.py +0 -2
  338. agno/utils/prompts.py +8 -6
  339. agno/utils/remote.py +23 -0
  340. agno/utils/response.py +1 -13
  341. agno/utils/string.py +91 -2
  342. agno/utils/team.py +62 -12
  343. agno/utils/tokens.py +657 -0
  344. agno/vectordb/base.py +15 -2
  345. agno/vectordb/cassandra/cassandra.py +1 -1
  346. agno/vectordb/chroma/__init__.py +2 -1
  347. agno/vectordb/chroma/chromadb.py +468 -23
  348. agno/vectordb/clickhouse/clickhousedb.py +1 -1
  349. agno/vectordb/couchbase/couchbase.py +6 -2
  350. agno/vectordb/lancedb/lance_db.py +7 -38
  351. agno/vectordb/lightrag/lightrag.py +7 -6
  352. agno/vectordb/milvus/milvus.py +118 -84
  353. agno/vectordb/mongodb/__init__.py +2 -1
  354. agno/vectordb/mongodb/mongodb.py +14 -31
  355. agno/vectordb/pgvector/pgvector.py +120 -66
  356. agno/vectordb/pineconedb/pineconedb.py +2 -19
  357. agno/vectordb/qdrant/__init__.py +2 -1
  358. agno/vectordb/qdrant/qdrant.py +33 -56
  359. agno/vectordb/redis/__init__.py +2 -1
  360. agno/vectordb/redis/redisdb.py +19 -31
  361. agno/vectordb/singlestore/singlestore.py +17 -9
  362. agno/vectordb/surrealdb/surrealdb.py +2 -38
  363. agno/vectordb/weaviate/__init__.py +2 -1
  364. agno/vectordb/weaviate/weaviate.py +7 -3
  365. agno/workflow/__init__.py +5 -1
  366. agno/workflow/agent.py +2 -2
  367. agno/workflow/condition.py +12 -10
  368. agno/workflow/loop.py +28 -9
  369. agno/workflow/parallel.py +21 -13
  370. agno/workflow/remote.py +362 -0
  371. agno/workflow/router.py +12 -9
  372. agno/workflow/step.py +261 -36
  373. agno/workflow/steps.py +12 -8
  374. agno/workflow/types.py +40 -77
  375. agno/workflow/workflow.py +939 -213
  376. {agno-2.2.13.dist-info → agno-2.4.3.dist-info}/METADATA +134 -181
  377. agno-2.4.3.dist-info/RECORD +677 -0
  378. {agno-2.2.13.dist-info → agno-2.4.3.dist-info}/WHEEL +1 -1
  379. agno/tools/googlesearch.py +0 -98
  380. agno/tools/memori.py +0 -339
  381. agno-2.2.13.dist-info/RECORD +0 -575
  382. {agno-2.2.13.dist-info → agno-2.4.3.dist-info}/licenses/LICENSE +0 -0
  383. {agno-2.2.13.dist-info → agno-2.4.3.dist-info}/top_level.txt +0 -0
@@ -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("`google-genai` not installed. Please install it using `pip install google-genai`")
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
- client_params["project"] = self.project_id or getenv("GOOGLE_CLOUD_PROJECT")
139
- client_params["location"] = self.location or getenv("GOOGLE_CLOUD_LOCATION")
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
- log_info(
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
- log_info("Google Search enabled.")
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
- log_info("URL context enabled.")
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
- log_info("Vertex AI Search enabled.")
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(provider_response, response_format=response_format)
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.response) if hasattr(e, "response") else 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=str(e.response) if hasattr(e, "response") else str(e),
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(provider_response, response_format=response_format)
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=str(e.response) if hasattr(e, "response") else str(e),
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=str(e.response) if hasattr(e, "response") else str(e),
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.content
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
- message_parts.append(Part.from_text(text=content_str))
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
- message_parts.append(
486
- Part.from_function_call(
487
- name=tool_call["function"]["name"],
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
- message_parts = [Part.from_text(text=content)]
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
- formatted_messages.append(file_content)
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, messages: List[Message], function_call_results: List[Message], **kwargs
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
- combined_content: List = []
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
- combined_content.append(result.content)
773
- combined_function_result.append({"tool_name": result.tool_name, "content": result.content})
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 combined_content:
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", content=combined_content, tool_calls=combined_function_result, metrics=message_metrics
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 OpenAI response into a ModelResponse.
1089
+ Parse the Gemini response into a ModelResponse.
786
1090
 
787
1091
  Args:
788
- response: Raw response from OpenAI
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[0].content:
798
- response_message = response.candidates[0].content
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.model_dump()
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.get("grounding_chunks", []) or []
879
- citation_pairs = []
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 isinstance(chunk, dict):
1210
+ if not chunk:
882
1211
  continue
883
- web = chunk.get("web")
884
- if not isinstance(web, dict):
1212
+ web = chunk.web
1213
+ if not web:
885
1214
  continue
886
- uri = web.get("uri")
887
- title = web.get("title")
1215
+ uri = web.uri
1216
+ title = web.title
888
1217
  if uri:
889
- citation_pairs.append((uri, title))
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.model_dump()
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.get("url_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.get("retrieved_url")
907
- status = url_meta.get("url_retrieval_status", "UNKNOWN")
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
- candidate_content = response_delta.candidates[0].content
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
- if response_delta.candidates[0].grounding_metadata is not None:
995
- citations = Citations()
996
- grounding_metadata = response_delta.candidates[0].grounding_metadata.model_dump()
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.pop("grounding_chunks", None) or []
1001
- citation_pairs = []
1361
+ chunks = grounding_metadata.grounding_chunks or []
1002
1362
  for chunk in chunks:
1003
- if not isinstance(chunk, dict):
1363
+ if not chunk:
1004
1364
  continue
1005
- web = chunk.get("web")
1006
- if not isinstance(web, dict):
1365
+ web = chunk.web
1366
+ if not web:
1007
1367
  continue
1008
- uri = web.get("uri")
1009
- title = web.get("title")
1368
+ uri = web.uri
1369
+ title = web.title
1010
1370
  if uri:
1011
- citation_pairs.append((uri, title))
1371
+ citations.urls.append(UrlCitation(url=uri, title=title))
1012
1372
 
1013
- # Create citation objects from filtered pairs
1014
- citations.urls = [UrlCitation(url=url, title=title) for url, title in citation_pairs]
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