ag2 0.9.1a1__py3-none-any.whl → 0.9.2__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.

Potentially problematic release.


This version of ag2 might be problematic. Click here for more details.

Files changed (371) hide show
  1. {ag2-0.9.1a1.dist-info → ag2-0.9.2.dist-info}/METADATA +272 -75
  2. ag2-0.9.2.dist-info/RECORD +406 -0
  3. {ag2-0.9.1a1.dist-info → ag2-0.9.2.dist-info}/WHEEL +1 -2
  4. autogen/__init__.py +89 -0
  5. autogen/_website/__init__.py +3 -0
  6. autogen/_website/generate_api_references.py +427 -0
  7. autogen/_website/generate_mkdocs.py +1174 -0
  8. autogen/_website/notebook_processor.py +476 -0
  9. autogen/_website/process_notebooks.py +656 -0
  10. autogen/_website/utils.py +412 -0
  11. autogen/agentchat/__init__.py +44 -0
  12. autogen/agentchat/agent.py +182 -0
  13. autogen/agentchat/assistant_agent.py +85 -0
  14. autogen/agentchat/chat.py +309 -0
  15. autogen/agentchat/contrib/__init__.py +5 -0
  16. autogen/agentchat/contrib/agent_eval/README.md +7 -0
  17. autogen/agentchat/contrib/agent_eval/agent_eval.py +108 -0
  18. autogen/agentchat/contrib/agent_eval/criterion.py +43 -0
  19. autogen/agentchat/contrib/agent_eval/critic_agent.py +44 -0
  20. autogen/agentchat/contrib/agent_eval/quantifier_agent.py +39 -0
  21. autogen/agentchat/contrib/agent_eval/subcritic_agent.py +45 -0
  22. autogen/agentchat/contrib/agent_eval/task.py +42 -0
  23. autogen/agentchat/contrib/agent_optimizer.py +429 -0
  24. autogen/agentchat/contrib/capabilities/__init__.py +5 -0
  25. autogen/agentchat/contrib/capabilities/agent_capability.py +20 -0
  26. autogen/agentchat/contrib/capabilities/generate_images.py +301 -0
  27. autogen/agentchat/contrib/capabilities/teachability.py +393 -0
  28. autogen/agentchat/contrib/capabilities/text_compressors.py +66 -0
  29. autogen/agentchat/contrib/capabilities/tools_capability.py +22 -0
  30. autogen/agentchat/contrib/capabilities/transform_messages.py +93 -0
  31. autogen/agentchat/contrib/capabilities/transforms.py +566 -0
  32. autogen/agentchat/contrib/capabilities/transforms_util.py +122 -0
  33. autogen/agentchat/contrib/capabilities/vision_capability.py +214 -0
  34. autogen/agentchat/contrib/captainagent/__init__.py +9 -0
  35. autogen/agentchat/contrib/captainagent/agent_builder.py +790 -0
  36. autogen/agentchat/contrib/captainagent/captainagent.py +512 -0
  37. autogen/agentchat/contrib/captainagent/tool_retriever.py +335 -0
  38. autogen/agentchat/contrib/captainagent/tools/README.md +44 -0
  39. autogen/agentchat/contrib/captainagent/tools/__init__.py +5 -0
  40. autogen/agentchat/contrib/captainagent/tools/data_analysis/calculate_correlation.py +40 -0
  41. autogen/agentchat/contrib/captainagent/tools/data_analysis/calculate_skewness_and_kurtosis.py +28 -0
  42. autogen/agentchat/contrib/captainagent/tools/data_analysis/detect_outlier_iqr.py +28 -0
  43. autogen/agentchat/contrib/captainagent/tools/data_analysis/detect_outlier_zscore.py +28 -0
  44. autogen/agentchat/contrib/captainagent/tools/data_analysis/explore_csv.py +21 -0
  45. autogen/agentchat/contrib/captainagent/tools/data_analysis/shapiro_wilk_test.py +30 -0
  46. autogen/agentchat/contrib/captainagent/tools/information_retrieval/arxiv_download.py +27 -0
  47. autogen/agentchat/contrib/captainagent/tools/information_retrieval/arxiv_search.py +53 -0
  48. autogen/agentchat/contrib/captainagent/tools/information_retrieval/extract_pdf_image.py +53 -0
  49. autogen/agentchat/contrib/captainagent/tools/information_retrieval/extract_pdf_text.py +38 -0
  50. autogen/agentchat/contrib/captainagent/tools/information_retrieval/get_wikipedia_text.py +21 -0
  51. autogen/agentchat/contrib/captainagent/tools/information_retrieval/get_youtube_caption.py +34 -0
  52. autogen/agentchat/contrib/captainagent/tools/information_retrieval/image_qa.py +60 -0
  53. autogen/agentchat/contrib/captainagent/tools/information_retrieval/optical_character_recognition.py +61 -0
  54. autogen/agentchat/contrib/captainagent/tools/information_retrieval/perform_web_search.py +47 -0
  55. autogen/agentchat/contrib/captainagent/tools/information_retrieval/scrape_wikipedia_tables.py +33 -0
  56. autogen/agentchat/contrib/captainagent/tools/information_retrieval/transcribe_audio_file.py +21 -0
  57. autogen/agentchat/contrib/captainagent/tools/information_retrieval/youtube_download.py +35 -0
  58. autogen/agentchat/contrib/captainagent/tools/math/calculate_circle_area_from_diameter.py +21 -0
  59. autogen/agentchat/contrib/captainagent/tools/math/calculate_day_of_the_week.py +18 -0
  60. autogen/agentchat/contrib/captainagent/tools/math/calculate_fraction_sum.py +28 -0
  61. autogen/agentchat/contrib/captainagent/tools/math/calculate_matrix_power.py +31 -0
  62. autogen/agentchat/contrib/captainagent/tools/math/calculate_reflected_point.py +16 -0
  63. autogen/agentchat/contrib/captainagent/tools/math/complex_numbers_product.py +25 -0
  64. autogen/agentchat/contrib/captainagent/tools/math/compute_currency_conversion.py +23 -0
  65. autogen/agentchat/contrib/captainagent/tools/math/count_distinct_permutations.py +27 -0
  66. autogen/agentchat/contrib/captainagent/tools/math/evaluate_expression.py +28 -0
  67. autogen/agentchat/contrib/captainagent/tools/math/find_continuity_point.py +34 -0
  68. autogen/agentchat/contrib/captainagent/tools/math/fraction_to_mixed_numbers.py +39 -0
  69. autogen/agentchat/contrib/captainagent/tools/math/modular_inverse_sum.py +23 -0
  70. autogen/agentchat/contrib/captainagent/tools/math/simplify_mixed_numbers.py +36 -0
  71. autogen/agentchat/contrib/captainagent/tools/math/sum_of_digit_factorials.py +15 -0
  72. autogen/agentchat/contrib/captainagent/tools/math/sum_of_primes_below.py +15 -0
  73. autogen/agentchat/contrib/captainagent/tools/requirements.txt +10 -0
  74. autogen/agentchat/contrib/captainagent/tools/tool_description.tsv +34 -0
  75. autogen/agentchat/contrib/gpt_assistant_agent.py +526 -0
  76. autogen/agentchat/contrib/graph_rag/__init__.py +9 -0
  77. autogen/agentchat/contrib/graph_rag/document.py +29 -0
  78. autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py +170 -0
  79. autogen/agentchat/contrib/graph_rag/falkor_graph_rag_capability.py +103 -0
  80. autogen/agentchat/contrib/graph_rag/graph_query_engine.py +53 -0
  81. autogen/agentchat/contrib/graph_rag/graph_rag_capability.py +63 -0
  82. autogen/agentchat/contrib/graph_rag/neo4j_graph_query_engine.py +268 -0
  83. autogen/agentchat/contrib/graph_rag/neo4j_graph_rag_capability.py +83 -0
  84. autogen/agentchat/contrib/graph_rag/neo4j_native_graph_query_engine.py +210 -0
  85. autogen/agentchat/contrib/graph_rag/neo4j_native_graph_rag_capability.py +93 -0
  86. autogen/agentchat/contrib/img_utils.py +397 -0
  87. autogen/agentchat/contrib/llamaindex_conversable_agent.py +117 -0
  88. autogen/agentchat/contrib/llava_agent.py +187 -0
  89. autogen/agentchat/contrib/math_user_proxy_agent.py +464 -0
  90. autogen/agentchat/contrib/multimodal_conversable_agent.py +125 -0
  91. autogen/agentchat/contrib/qdrant_retrieve_user_proxy_agent.py +324 -0
  92. autogen/agentchat/contrib/rag/__init__.py +10 -0
  93. autogen/agentchat/contrib/rag/chromadb_query_engine.py +272 -0
  94. autogen/agentchat/contrib/rag/llamaindex_query_engine.py +198 -0
  95. autogen/agentchat/contrib/rag/mongodb_query_engine.py +329 -0
  96. autogen/agentchat/contrib/rag/query_engine.py +74 -0
  97. autogen/agentchat/contrib/retrieve_assistant_agent.py +56 -0
  98. autogen/agentchat/contrib/retrieve_user_proxy_agent.py +703 -0
  99. autogen/agentchat/contrib/society_of_mind_agent.py +199 -0
  100. autogen/agentchat/contrib/swarm_agent.py +1425 -0
  101. autogen/agentchat/contrib/text_analyzer_agent.py +79 -0
  102. autogen/agentchat/contrib/vectordb/__init__.py +5 -0
  103. autogen/agentchat/contrib/vectordb/base.py +232 -0
  104. autogen/agentchat/contrib/vectordb/chromadb.py +315 -0
  105. autogen/agentchat/contrib/vectordb/couchbase.py +407 -0
  106. autogen/agentchat/contrib/vectordb/mongodb.py +550 -0
  107. autogen/agentchat/contrib/vectordb/pgvectordb.py +928 -0
  108. autogen/agentchat/contrib/vectordb/qdrant.py +320 -0
  109. autogen/agentchat/contrib/vectordb/utils.py +126 -0
  110. autogen/agentchat/contrib/web_surfer.py +303 -0
  111. autogen/agentchat/conversable_agent.py +4023 -0
  112. autogen/agentchat/group/__init__.py +64 -0
  113. autogen/agentchat/group/available_condition.py +91 -0
  114. autogen/agentchat/group/context_condition.py +77 -0
  115. autogen/agentchat/group/context_expression.py +238 -0
  116. autogen/agentchat/group/context_str.py +41 -0
  117. autogen/agentchat/group/context_variables.py +192 -0
  118. autogen/agentchat/group/group_tool_executor.py +202 -0
  119. autogen/agentchat/group/group_utils.py +591 -0
  120. autogen/agentchat/group/handoffs.py +244 -0
  121. autogen/agentchat/group/llm_condition.py +93 -0
  122. autogen/agentchat/group/multi_agent_chat.py +237 -0
  123. autogen/agentchat/group/on_condition.py +58 -0
  124. autogen/agentchat/group/on_context_condition.py +54 -0
  125. autogen/agentchat/group/patterns/__init__.py +18 -0
  126. autogen/agentchat/group/patterns/auto.py +159 -0
  127. autogen/agentchat/group/patterns/manual.py +176 -0
  128. autogen/agentchat/group/patterns/pattern.py +288 -0
  129. autogen/agentchat/group/patterns/random.py +106 -0
  130. autogen/agentchat/group/patterns/round_robin.py +117 -0
  131. autogen/agentchat/group/reply_result.py +26 -0
  132. autogen/agentchat/group/speaker_selection_result.py +41 -0
  133. autogen/agentchat/group/targets/__init__.py +4 -0
  134. autogen/agentchat/group/targets/group_chat_target.py +132 -0
  135. autogen/agentchat/group/targets/group_manager_target.py +151 -0
  136. autogen/agentchat/group/targets/transition_target.py +413 -0
  137. autogen/agentchat/group/targets/transition_utils.py +6 -0
  138. autogen/agentchat/groupchat.py +1694 -0
  139. autogen/agentchat/realtime/__init__.py +3 -0
  140. autogen/agentchat/realtime/experimental/__init__.py +20 -0
  141. autogen/agentchat/realtime/experimental/audio_adapters/__init__.py +8 -0
  142. autogen/agentchat/realtime/experimental/audio_adapters/twilio_audio_adapter.py +148 -0
  143. autogen/agentchat/realtime/experimental/audio_adapters/websocket_audio_adapter.py +139 -0
  144. autogen/agentchat/realtime/experimental/audio_observer.py +42 -0
  145. autogen/agentchat/realtime/experimental/clients/__init__.py +15 -0
  146. autogen/agentchat/realtime/experimental/clients/gemini/__init__.py +7 -0
  147. autogen/agentchat/realtime/experimental/clients/gemini/client.py +274 -0
  148. autogen/agentchat/realtime/experimental/clients/oai/__init__.py +8 -0
  149. autogen/agentchat/realtime/experimental/clients/oai/base_client.py +220 -0
  150. autogen/agentchat/realtime/experimental/clients/oai/rtc_client.py +243 -0
  151. autogen/agentchat/realtime/experimental/clients/oai/utils.py +48 -0
  152. autogen/agentchat/realtime/experimental/clients/realtime_client.py +190 -0
  153. autogen/agentchat/realtime/experimental/function_observer.py +85 -0
  154. autogen/agentchat/realtime/experimental/realtime_agent.py +158 -0
  155. autogen/agentchat/realtime/experimental/realtime_events.py +42 -0
  156. autogen/agentchat/realtime/experimental/realtime_observer.py +100 -0
  157. autogen/agentchat/realtime/experimental/realtime_swarm.py +475 -0
  158. autogen/agentchat/realtime/experimental/websockets.py +21 -0
  159. autogen/agentchat/realtime_agent/__init__.py +21 -0
  160. autogen/agentchat/user_proxy_agent.py +111 -0
  161. autogen/agentchat/utils.py +206 -0
  162. autogen/agents/__init__.py +3 -0
  163. autogen/agents/contrib/__init__.py +10 -0
  164. autogen/agents/contrib/time/__init__.py +8 -0
  165. autogen/agents/contrib/time/time_reply_agent.py +73 -0
  166. autogen/agents/contrib/time/time_tool_agent.py +51 -0
  167. autogen/agents/experimental/__init__.py +27 -0
  168. autogen/agents/experimental/deep_research/__init__.py +7 -0
  169. autogen/agents/experimental/deep_research/deep_research.py +52 -0
  170. autogen/agents/experimental/discord/__init__.py +7 -0
  171. autogen/agents/experimental/discord/discord.py +66 -0
  172. autogen/agents/experimental/document_agent/__init__.py +19 -0
  173. autogen/agents/experimental/document_agent/chroma_query_engine.py +316 -0
  174. autogen/agents/experimental/document_agent/docling_doc_ingest_agent.py +118 -0
  175. autogen/agents/experimental/document_agent/document_agent.py +461 -0
  176. autogen/agents/experimental/document_agent/document_conditions.py +50 -0
  177. autogen/agents/experimental/document_agent/document_utils.py +380 -0
  178. autogen/agents/experimental/document_agent/inmemory_query_engine.py +220 -0
  179. autogen/agents/experimental/document_agent/parser_utils.py +130 -0
  180. autogen/agents/experimental/document_agent/url_utils.py +426 -0
  181. autogen/agents/experimental/reasoning/__init__.py +7 -0
  182. autogen/agents/experimental/reasoning/reasoning_agent.py +1178 -0
  183. autogen/agents/experimental/slack/__init__.py +7 -0
  184. autogen/agents/experimental/slack/slack.py +73 -0
  185. autogen/agents/experimental/telegram/__init__.py +7 -0
  186. autogen/agents/experimental/telegram/telegram.py +77 -0
  187. autogen/agents/experimental/websurfer/__init__.py +7 -0
  188. autogen/agents/experimental/websurfer/websurfer.py +62 -0
  189. autogen/agents/experimental/wikipedia/__init__.py +7 -0
  190. autogen/agents/experimental/wikipedia/wikipedia.py +90 -0
  191. autogen/browser_utils.py +309 -0
  192. autogen/cache/__init__.py +10 -0
  193. autogen/cache/abstract_cache_base.py +75 -0
  194. autogen/cache/cache.py +203 -0
  195. autogen/cache/cache_factory.py +88 -0
  196. autogen/cache/cosmos_db_cache.py +144 -0
  197. autogen/cache/disk_cache.py +102 -0
  198. autogen/cache/in_memory_cache.py +58 -0
  199. autogen/cache/redis_cache.py +123 -0
  200. autogen/code_utils.py +596 -0
  201. autogen/coding/__init__.py +22 -0
  202. autogen/coding/base.py +119 -0
  203. autogen/coding/docker_commandline_code_executor.py +268 -0
  204. autogen/coding/factory.py +47 -0
  205. autogen/coding/func_with_reqs.py +202 -0
  206. autogen/coding/jupyter/__init__.py +23 -0
  207. autogen/coding/jupyter/base.py +36 -0
  208. autogen/coding/jupyter/docker_jupyter_server.py +167 -0
  209. autogen/coding/jupyter/embedded_ipython_code_executor.py +182 -0
  210. autogen/coding/jupyter/import_utils.py +82 -0
  211. autogen/coding/jupyter/jupyter_client.py +231 -0
  212. autogen/coding/jupyter/jupyter_code_executor.py +160 -0
  213. autogen/coding/jupyter/local_jupyter_server.py +172 -0
  214. autogen/coding/local_commandline_code_executor.py +405 -0
  215. autogen/coding/markdown_code_extractor.py +45 -0
  216. autogen/coding/utils.py +56 -0
  217. autogen/doc_utils.py +34 -0
  218. autogen/events/__init__.py +7 -0
  219. autogen/events/agent_events.py +1013 -0
  220. autogen/events/base_event.py +99 -0
  221. autogen/events/client_events.py +167 -0
  222. autogen/events/helpers.py +36 -0
  223. autogen/events/print_event.py +46 -0
  224. autogen/exception_utils.py +73 -0
  225. autogen/extensions/__init__.py +5 -0
  226. autogen/fast_depends/__init__.py +16 -0
  227. autogen/fast_depends/_compat.py +80 -0
  228. autogen/fast_depends/core/__init__.py +14 -0
  229. autogen/fast_depends/core/build.py +225 -0
  230. autogen/fast_depends/core/model.py +576 -0
  231. autogen/fast_depends/dependencies/__init__.py +15 -0
  232. autogen/fast_depends/dependencies/model.py +29 -0
  233. autogen/fast_depends/dependencies/provider.py +39 -0
  234. autogen/fast_depends/library/__init__.py +10 -0
  235. autogen/fast_depends/library/model.py +46 -0
  236. autogen/fast_depends/py.typed +6 -0
  237. autogen/fast_depends/schema.py +66 -0
  238. autogen/fast_depends/use.py +280 -0
  239. autogen/fast_depends/utils.py +187 -0
  240. autogen/formatting_utils.py +83 -0
  241. autogen/function_utils.py +13 -0
  242. autogen/graph_utils.py +178 -0
  243. autogen/import_utils.py +526 -0
  244. autogen/interop/__init__.py +22 -0
  245. autogen/interop/crewai/__init__.py +7 -0
  246. autogen/interop/crewai/crewai.py +88 -0
  247. autogen/interop/interoperability.py +71 -0
  248. autogen/interop/interoperable.py +46 -0
  249. autogen/interop/langchain/__init__.py +8 -0
  250. autogen/interop/langchain/langchain_chat_model_factory.py +155 -0
  251. autogen/interop/langchain/langchain_tool.py +82 -0
  252. autogen/interop/litellm/__init__.py +7 -0
  253. autogen/interop/litellm/litellm_config_factory.py +179 -0
  254. autogen/interop/pydantic_ai/__init__.py +7 -0
  255. autogen/interop/pydantic_ai/pydantic_ai.py +168 -0
  256. autogen/interop/registry.py +69 -0
  257. autogen/io/__init__.py +15 -0
  258. autogen/io/base.py +151 -0
  259. autogen/io/console.py +56 -0
  260. autogen/io/processors/__init__.py +12 -0
  261. autogen/io/processors/base.py +21 -0
  262. autogen/io/processors/console_event_processor.py +56 -0
  263. autogen/io/run_response.py +293 -0
  264. autogen/io/thread_io_stream.py +63 -0
  265. autogen/io/websockets.py +213 -0
  266. autogen/json_utils.py +43 -0
  267. autogen/llm_config.py +382 -0
  268. autogen/logger/__init__.py +11 -0
  269. autogen/logger/base_logger.py +128 -0
  270. autogen/logger/file_logger.py +261 -0
  271. autogen/logger/logger_factory.py +42 -0
  272. autogen/logger/logger_utils.py +57 -0
  273. autogen/logger/sqlite_logger.py +523 -0
  274. autogen/math_utils.py +339 -0
  275. autogen/mcp/__init__.py +7 -0
  276. autogen/mcp/__main__.py +78 -0
  277. autogen/mcp/mcp_client.py +208 -0
  278. autogen/mcp/mcp_proxy/__init__.py +19 -0
  279. autogen/mcp/mcp_proxy/fastapi_code_generator_helpers.py +63 -0
  280. autogen/mcp/mcp_proxy/mcp_proxy.py +581 -0
  281. autogen/mcp/mcp_proxy/operation_grouping.py +158 -0
  282. autogen/mcp/mcp_proxy/operation_renaming.py +114 -0
  283. autogen/mcp/mcp_proxy/patch_fastapi_code_generator.py +98 -0
  284. autogen/mcp/mcp_proxy/security.py +400 -0
  285. autogen/mcp/mcp_proxy/security_schema_visitor.py +37 -0
  286. autogen/messages/__init__.py +7 -0
  287. autogen/messages/agent_messages.py +948 -0
  288. autogen/messages/base_message.py +107 -0
  289. autogen/messages/client_messages.py +171 -0
  290. autogen/messages/print_message.py +49 -0
  291. autogen/oai/__init__.py +53 -0
  292. autogen/oai/anthropic.py +714 -0
  293. autogen/oai/bedrock.py +628 -0
  294. autogen/oai/cerebras.py +299 -0
  295. autogen/oai/client.py +1444 -0
  296. autogen/oai/client_utils.py +169 -0
  297. autogen/oai/cohere.py +479 -0
  298. autogen/oai/gemini.py +998 -0
  299. autogen/oai/gemini_types.py +155 -0
  300. autogen/oai/groq.py +305 -0
  301. autogen/oai/mistral.py +303 -0
  302. autogen/oai/oai_models/__init__.py +11 -0
  303. autogen/oai/oai_models/_models.py +16 -0
  304. autogen/oai/oai_models/chat_completion.py +87 -0
  305. autogen/oai/oai_models/chat_completion_audio.py +32 -0
  306. autogen/oai/oai_models/chat_completion_message.py +86 -0
  307. autogen/oai/oai_models/chat_completion_message_tool_call.py +37 -0
  308. autogen/oai/oai_models/chat_completion_token_logprob.py +63 -0
  309. autogen/oai/oai_models/completion_usage.py +60 -0
  310. autogen/oai/ollama.py +643 -0
  311. autogen/oai/openai_utils.py +881 -0
  312. autogen/oai/together.py +370 -0
  313. autogen/retrieve_utils.py +491 -0
  314. autogen/runtime_logging.py +160 -0
  315. autogen/token_count_utils.py +267 -0
  316. autogen/tools/__init__.py +20 -0
  317. autogen/tools/contrib/__init__.py +9 -0
  318. autogen/tools/contrib/time/__init__.py +7 -0
  319. autogen/tools/contrib/time/time.py +41 -0
  320. autogen/tools/dependency_injection.py +254 -0
  321. autogen/tools/experimental/__init__.py +48 -0
  322. autogen/tools/experimental/browser_use/__init__.py +7 -0
  323. autogen/tools/experimental/browser_use/browser_use.py +161 -0
  324. autogen/tools/experimental/crawl4ai/__init__.py +7 -0
  325. autogen/tools/experimental/crawl4ai/crawl4ai.py +153 -0
  326. autogen/tools/experimental/deep_research/__init__.py +7 -0
  327. autogen/tools/experimental/deep_research/deep_research.py +328 -0
  328. autogen/tools/experimental/duckduckgo/__init__.py +7 -0
  329. autogen/tools/experimental/duckduckgo/duckduckgo_search.py +109 -0
  330. autogen/tools/experimental/google/__init__.py +14 -0
  331. autogen/tools/experimental/google/authentication/__init__.py +11 -0
  332. autogen/tools/experimental/google/authentication/credentials_hosted_provider.py +43 -0
  333. autogen/tools/experimental/google/authentication/credentials_local_provider.py +91 -0
  334. autogen/tools/experimental/google/authentication/credentials_provider.py +35 -0
  335. autogen/tools/experimental/google/drive/__init__.py +9 -0
  336. autogen/tools/experimental/google/drive/drive_functions.py +124 -0
  337. autogen/tools/experimental/google/drive/toolkit.py +88 -0
  338. autogen/tools/experimental/google/model.py +17 -0
  339. autogen/tools/experimental/google/toolkit_protocol.py +19 -0
  340. autogen/tools/experimental/google_search/__init__.py +8 -0
  341. autogen/tools/experimental/google_search/google_search.py +93 -0
  342. autogen/tools/experimental/google_search/youtube_search.py +181 -0
  343. autogen/tools/experimental/messageplatform/__init__.py +17 -0
  344. autogen/tools/experimental/messageplatform/discord/__init__.py +7 -0
  345. autogen/tools/experimental/messageplatform/discord/discord.py +288 -0
  346. autogen/tools/experimental/messageplatform/slack/__init__.py +7 -0
  347. autogen/tools/experimental/messageplatform/slack/slack.py +391 -0
  348. autogen/tools/experimental/messageplatform/telegram/__init__.py +7 -0
  349. autogen/tools/experimental/messageplatform/telegram/telegram.py +275 -0
  350. autogen/tools/experimental/perplexity/__init__.py +7 -0
  351. autogen/tools/experimental/perplexity/perplexity_search.py +260 -0
  352. autogen/tools/experimental/reliable/__init__.py +10 -0
  353. autogen/tools/experimental/reliable/reliable.py +1316 -0
  354. autogen/tools/experimental/tavily/__init__.py +7 -0
  355. autogen/tools/experimental/tavily/tavily_search.py +183 -0
  356. autogen/tools/experimental/web_search_preview/__init__.py +7 -0
  357. autogen/tools/experimental/web_search_preview/web_search_preview.py +114 -0
  358. autogen/tools/experimental/wikipedia/__init__.py +7 -0
  359. autogen/tools/experimental/wikipedia/wikipedia.py +287 -0
  360. autogen/tools/function_utils.py +411 -0
  361. autogen/tools/tool.py +187 -0
  362. autogen/tools/toolkit.py +86 -0
  363. autogen/types.py +29 -0
  364. autogen/version.py +7 -0
  365. templates/client_template/main.jinja2 +69 -0
  366. templates/config_template/config.jinja2 +7 -0
  367. templates/main.jinja2 +61 -0
  368. ag2-0.9.1a1.dist-info/RECORD +0 -6
  369. ag2-0.9.1a1.dist-info/top_level.txt +0 -1
  370. {ag2-0.9.1a1.dist-info → ag2-0.9.2.dist-info/licenses}/LICENSE +0 -0
  371. {ag2-0.9.1a1.dist-info → ag2-0.9.2.dist-info/licenses}/NOTICE.md +0 -0
autogen/oai/gemini.py ADDED
@@ -0,0 +1,998 @@
1
+ # Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+ #
5
+ # Portions derived from https://github.com/microsoft/autogen are under the MIT License.
6
+ # SPDX-License-Identifier: MIT
7
+ """Create a OpenAI-compatible client for Gemini features.
8
+
9
+ Example:
10
+ ```python
11
+ llm_config = {
12
+ "config_list": [
13
+ {
14
+ "api_type": "google",
15
+ "model": "gemini-pro",
16
+ "api_key": os.environ.get("GOOGLE_GEMINI_API_KEY"),
17
+ "safety_settings": [
18
+ {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_ONLY_HIGH"},
19
+ {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_ONLY_HIGH"},
20
+ {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_ONLY_HIGH"},
21
+ {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_ONLY_HIGH"},
22
+ ],
23
+ "top_p": 0.5,
24
+ "max_tokens": 2048,
25
+ "temperature": 1.0,
26
+ "top_k": 5,
27
+ }
28
+ ]
29
+ }
30
+
31
+ agent = autogen.AssistantAgent("my_agent", llm_config=llm_config)
32
+ ```
33
+
34
+ Resources:
35
+ - https://ai.google.dev/docs
36
+ - https://cloud.google.com/vertex-ai/generative-ai/docs/migrate/migrate-from-azure-to-gemini
37
+ - https://blog.google/technology/ai/google-gemini-pro-imagen-duet-ai-update/
38
+ - https://ai.google.dev/api/python/google/generativeai/ChatSession
39
+ """
40
+
41
+ from __future__ import annotations
42
+
43
+ import asyncio
44
+ import base64
45
+ import copy
46
+ import json
47
+ import logging
48
+ import os
49
+ import random
50
+ import re
51
+ import time
52
+ import warnings
53
+ from io import BytesIO
54
+ from typing import Any, Literal, Optional, Type, Union
55
+
56
+ import requests
57
+ from packaging import version
58
+ from pydantic import BaseModel, Field
59
+
60
+ from ..import_utils import optional_import_block, require_optional_import
61
+ from ..json_utils import resolve_json_references
62
+ from ..llm_config import LLMConfigEntry, register_llm_config
63
+ from .client_utils import FormatterProtocol
64
+ from .gemini_types import ToolConfig
65
+ from .oai_models import ChatCompletion, ChatCompletionMessage, ChatCompletionMessageToolCall, Choice, CompletionUsage
66
+
67
+ with optional_import_block():
68
+ import google.genai as genai
69
+ import vertexai
70
+ from PIL import Image
71
+ from google.auth.credentials import Credentials
72
+ from google.genai.types import (
73
+ Content,
74
+ FinishReason,
75
+ FunctionCall,
76
+ FunctionDeclaration,
77
+ FunctionResponse,
78
+ GenerateContentConfig,
79
+ GenerateContentResponse,
80
+ GoogleSearch,
81
+ Part,
82
+ Schema,
83
+ Tool,
84
+ Type,
85
+ )
86
+ from jsonschema import ValidationError
87
+ from vertexai.generative_models import Content as VertexAIContent
88
+ from vertexai.generative_models import FunctionDeclaration as vaiFunctionDeclaration
89
+ from vertexai.generative_models import GenerationConfig, GenerativeModel
90
+ from vertexai.generative_models import (
91
+ GenerationResponse as VertexAIGenerationResponse,
92
+ )
93
+ from vertexai.generative_models import HarmBlockThreshold as VertexAIHarmBlockThreshold
94
+ from vertexai.generative_models import HarmCategory as VertexAIHarmCategory
95
+ from vertexai.generative_models import Part as VertexAIPart
96
+ from vertexai.generative_models import SafetySetting as VertexAISafetySetting
97
+ from vertexai.generative_models import (
98
+ Tool as vaiTool,
99
+ )
100
+
101
+ logger = logging.getLogger(__name__)
102
+
103
+
104
+ @register_llm_config
105
+ class GeminiLLMConfigEntry(LLMConfigEntry):
106
+ api_type: Literal["google"] = "google"
107
+ project_id: Optional[str] = None
108
+ location: Optional[str] = None
109
+ # google_application_credentials points to the path of the JSON Keyfile
110
+ google_application_credentials: Optional[str] = None
111
+ # credentials is a google.auth.credentials.Credentials object
112
+ credentials: Optional[Union[Any, str]] = None
113
+ stream: bool = False
114
+ safety_settings: Optional[Union[list[dict[str, Any]], dict[str, Any]]] = None
115
+ price: Optional[list[float]] = Field(default=None, min_length=2, max_length=2)
116
+ tool_config: Optional[ToolConfig] = None
117
+
118
+ def create_client(self):
119
+ raise NotImplementedError("GeminiLLMConfigEntry.create_client() is not implemented.")
120
+
121
+
122
+ @require_optional_import(["google", "vertexai", "PIL", "jsonschema"], "gemini")
123
+ class GeminiClient:
124
+ """Client for Google's Gemini API."""
125
+
126
+ # Mapping, where Key is a term used by Autogen, and Value is a term used by Gemini
127
+ PARAMS_MAPPING = {
128
+ "max_tokens": "max_output_tokens",
129
+ # "n": "candidate_count", # Gemini supports only `n=1`
130
+ "stop_sequences": "stop_sequences",
131
+ "temperature": "temperature",
132
+ "top_p": "top_p",
133
+ "top_k": "top_k",
134
+ "max_output_tokens": "max_output_tokens",
135
+ }
136
+
137
+ def _initialize_vertexai(self, **params):
138
+ if "google_application_credentials" in params:
139
+ # Path to JSON Keyfile
140
+ os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = params["google_application_credentials"]
141
+ vertexai_init_args = {}
142
+ if "project_id" in params:
143
+ vertexai_init_args["project"] = params["project_id"]
144
+ if "location" in params:
145
+ vertexai_init_args["location"] = params["location"]
146
+ if "credentials" in params:
147
+ assert isinstance(params["credentials"], Credentials), (
148
+ "Object type google.auth.credentials.Credentials is expected!"
149
+ )
150
+ vertexai_init_args["credentials"] = params["credentials"]
151
+ if vertexai_init_args:
152
+ vertexai.init(**vertexai_init_args)
153
+
154
+ def __init__(self, **kwargs):
155
+ """Uses either either api_key for authentication from the LLM config
156
+ (specifying the GOOGLE_GEMINI_API_KEY environment variable also works),
157
+ or follows the Google authentication mechanism for VertexAI in Google Cloud if no api_key is specified,
158
+ where project_id and location can also be passed as parameters. Previously created credentials object can be provided,
159
+ or a Service account key file can also be used. If neither a service account key file, nor the api_key are passed,
160
+ then the default credentials will be used, which could be a personal account if the user is already authenticated in,
161
+ like in Google Cloud Shell.
162
+
163
+ Args:
164
+ **kwargs: The keyword arguments to initialize the Gemini client.
165
+ """
166
+ self.api_key = kwargs.get("api_key")
167
+ if not self.api_key:
168
+ self.api_key = os.getenv("GOOGLE_GEMINI_API_KEY")
169
+ if self.api_key is None:
170
+ self.use_vertexai = True
171
+ self._initialize_vertexai(**kwargs)
172
+ else:
173
+ self.use_vertexai = False
174
+ else:
175
+ self.use_vertexai = False
176
+ if not self.use_vertexai:
177
+ assert ("project_id" not in kwargs) and ("location" not in kwargs), (
178
+ "Google Cloud project and compute location cannot be set when using an API Key!"
179
+ )
180
+
181
+ self.api_version = kwargs.get("api_version")
182
+
183
+ # Store the response format, if provided (for structured outputs)
184
+ self._response_format: Optional[type[BaseModel]] = None
185
+
186
+ def message_retrieval(self, response) -> list:
187
+ """Retrieve and return a list of strings or a list of Choice.Message from the response.
188
+
189
+ NOTE: if a list of Choice.Message is returned, it currently needs to contain the fields of OpenAI's ChatCompletion Message object,
190
+ since that is expected for function or tool calling in the rest of the codebase at the moment, unless a custom agent is being used.
191
+ """
192
+ return [choice.message for choice in response.choices]
193
+
194
+ def cost(self, response) -> float:
195
+ return response.cost
196
+
197
+ @staticmethod
198
+ def get_usage(response) -> dict:
199
+ """Return usage summary of the response using RESPONSE_USAGE_KEYS."""
200
+ # ... # pragma: no cover
201
+ return {
202
+ "prompt_tokens": response.usage.prompt_tokens,
203
+ "completion_tokens": response.usage.completion_tokens,
204
+ "total_tokens": response.usage.total_tokens,
205
+ "cost": response.cost,
206
+ "model": response.model,
207
+ }
208
+
209
+ def create(self, params: dict) -> ChatCompletion:
210
+ # When running in async context via run_in_executor from ConversableAgent.a_generate_oai_reply,
211
+ # this method runs in a new thread that doesn't have an event loop by default. The Google Genai
212
+ # client requires an event loop even for synchronous operations, so we need to ensure one exists.
213
+ try:
214
+ asyncio.get_running_loop()
215
+ except RuntimeError:
216
+ # No event loop exists in this thread (which happens when called from an executor)
217
+ # Create a new event loop for this thread to satisfy Genai client requirements
218
+ loop = asyncio.new_event_loop()
219
+ asyncio.set_event_loop(loop)
220
+
221
+ if self.use_vertexai:
222
+ self._initialize_vertexai(**params)
223
+ else:
224
+ assert ("project_id" not in params) and ("location" not in params), (
225
+ "Google Cloud project and compute location cannot be set when using an API Key!"
226
+ )
227
+ model_name = params.get("model", "gemini-pro")
228
+
229
+ if model_name == "gemini-pro-vision":
230
+ raise ValueError(
231
+ "Gemini 1.0 Pro vision ('gemini-pro-vision') has been deprecated, please consider switching to a different model, for example 'gemini-1.5-flash'."
232
+ )
233
+ elif not model_name:
234
+ raise ValueError(
235
+ "Please provide a model name for the Gemini Client. "
236
+ "You can configure it in the OAI Config List file. "
237
+ "See this [LLM configuration tutorial](https://docs.ag2.ai/latest/docs/user-guide/basic-concepts/llm-configuration/) for more details."
238
+ )
239
+
240
+ params.get("api_type", "google") # not used
241
+ http_options = {"api_version": self.api_version} if self.api_version else None
242
+ messages = params.get("messages", [])
243
+ stream = params.get("stream", False)
244
+ n_response = params.get("n", 1)
245
+ system_instruction = self._extract_system_instruction(messages)
246
+ response_validation = params.get("response_validation", True)
247
+ tools = self._tools_to_gemini_tools(params["tools"]) if "tools" in params else None
248
+ tool_config = params.get("tool_config")
249
+
250
+ generation_config = {
251
+ gemini_term: params[autogen_term]
252
+ for autogen_term, gemini_term in self.PARAMS_MAPPING.items()
253
+ if autogen_term in params
254
+ }
255
+ if self.use_vertexai:
256
+ safety_settings = GeminiClient._to_vertexai_safety_settings(params.get("safety_settings", []))
257
+ else:
258
+ safety_settings = params.get("safety_settings", [])
259
+
260
+ if stream:
261
+ warnings.warn(
262
+ "Streaming is not supported for Gemini yet, and it will have no effect. Please set stream=False.",
263
+ UserWarning,
264
+ )
265
+ stream = False
266
+
267
+ if n_response > 1:
268
+ warnings.warn("Gemini only supports `n=1` for now. We only generate one response.", UserWarning)
269
+
270
+ autogen_tool_calls = []
271
+
272
+ # Maps the function call ids to function names so we can inject it into FunctionResponse messages
273
+ self.tool_call_function_map: dict[str, str] = {}
274
+
275
+ # If response_format exists, we want structured outputs
276
+ # Based on
277
+ # https://ai.google.dev/gemini-api/docs/structured-output?lang=python#supply-schema-in-config
278
+ if params.get("response_format"):
279
+ self._response_format = params.get("response_format")
280
+ generation_config["response_mime_type"] = "application/json"
281
+
282
+ response_format_schema_raw = params.get("response_format")
283
+
284
+ if isinstance(response_format_schema_raw, dict):
285
+ response_schema = resolve_json_references(response_format_schema_raw)
286
+ else:
287
+ response_schema = resolve_json_references(params.get("response_format").model_json_schema())
288
+ if "$defs" in response_schema:
289
+ response_schema.pop("$defs")
290
+ generation_config["response_schema"] = response_schema
291
+
292
+ # A. create and call the chat model.
293
+ gemini_messages = self._oai_messages_to_gemini_messages(messages)
294
+ if self.use_vertexai:
295
+ model = GenerativeModel(
296
+ model_name,
297
+ generation_config=GenerationConfig(**generation_config),
298
+ safety_settings=safety_settings,
299
+ system_instruction=system_instruction,
300
+ tool_config=tool_config,
301
+ tools=tools,
302
+ )
303
+
304
+ chat = model.start_chat(history=gemini_messages[:-1], response_validation=response_validation)
305
+ response = chat.send_message(gemini_messages[-1].parts, stream=stream, safety_settings=safety_settings)
306
+ else:
307
+ client = genai.Client(api_key=self.api_key, http_options=http_options)
308
+ generate_content_config = GenerateContentConfig(
309
+ safety_settings=safety_settings,
310
+ system_instruction=system_instruction,
311
+ tools=tools,
312
+ tool_config=tool_config,
313
+ **generation_config,
314
+ )
315
+ chat = client.chats.create(model=model_name, config=generate_content_config, history=gemini_messages[:-1])
316
+ response = chat.send_message(message=gemini_messages[-1].parts)
317
+
318
+ # Extract text and tools from response
319
+ ans = ""
320
+ random_id = random.randint(0, 10000)
321
+ prev_function_calls = []
322
+ error_finish_reason = None
323
+
324
+ if isinstance(response, GenerateContentResponse):
325
+ if len(response.candidates) != 1:
326
+ raise ValueError(
327
+ f"Unexpected number of candidates in the response. Expected 1, got {len(response.candidates)}"
328
+ )
329
+
330
+ # Look at https://cloud.google.com/vertex-ai/generative-ai/docs/reference/python/latest/vertexai.generative_models.FinishReason
331
+ if response.candidates[0].finish_reason and response.candidates[0].finish_reason == FinishReason.RECITATION:
332
+ recitation_part = Part(text="Unsuccessful Finish Reason: RECITATION")
333
+ parts = [recitation_part]
334
+ error_finish_reason = "content_filter" # As per available finish_reason in Choice
335
+ else:
336
+ parts = response.candidates[0].content.parts
337
+ elif isinstance(response, VertexAIGenerationResponse): # or hasattr(response, "candidates"):
338
+ # google.generativeai also raises an error len(candidates) != 1:
339
+ if len(response.candidates) != 1:
340
+ raise ValueError(
341
+ f"Unexpected number of candidates in the response. Expected 1, got {len(response.candidates)}"
342
+ )
343
+ parts = response.candidates[0].content.parts
344
+ else:
345
+ raise ValueError(f"Unexpected response type: {type(response)}")
346
+
347
+ for part in parts:
348
+ # Function calls
349
+ if fn_call := part.function_call:
350
+ # If we have a repeated function call, ignore it
351
+ if fn_call not in prev_function_calls:
352
+ autogen_tool_calls.append(
353
+ ChatCompletionMessageToolCall(
354
+ id=str(random_id),
355
+ function={
356
+ "name": fn_call.name,
357
+ "arguments": (
358
+ json.dumps({key: val for key, val in fn_call.args.items()})
359
+ if fn_call.args is not None
360
+ else ""
361
+ ),
362
+ },
363
+ type="function",
364
+ )
365
+ )
366
+
367
+ prev_function_calls.append(fn_call)
368
+ random_id += 1
369
+
370
+ # Plain text content
371
+ elif text := part.text:
372
+ ans += text
373
+
374
+ # If we have function calls, ignore the text
375
+ # as it can be Gemini guessing the function response
376
+ if len(autogen_tool_calls) != 0:
377
+ ans = ""
378
+ else:
379
+ autogen_tool_calls = None
380
+
381
+ if self._response_format and ans:
382
+ try:
383
+ parsed_response = self._convert_json_response(ans)
384
+ ans = _format_json_response(parsed_response, ans)
385
+ except ValueError as e:
386
+ ans = str(e)
387
+
388
+ # 3. convert output
389
+ message = ChatCompletionMessage(
390
+ role="assistant", content=ans, function_call=None, tool_calls=autogen_tool_calls
391
+ )
392
+ choices = [
393
+ Choice(
394
+ finish_reason="tool_calls"
395
+ if autogen_tool_calls is not None
396
+ else error_finish_reason
397
+ if error_finish_reason
398
+ else "stop",
399
+ index=0,
400
+ message=message,
401
+ )
402
+ ]
403
+
404
+ prompt_tokens = response.usage_metadata.prompt_token_count
405
+ completion_tokens = (
406
+ response.usage_metadata.candidates_token_count if response.usage_metadata.candidates_token_count else 0
407
+ )
408
+
409
+ response_oai = ChatCompletion(
410
+ id=str(random.randint(0, 1000)),
411
+ model=model_name,
412
+ created=int(time.time()),
413
+ object="chat.completion",
414
+ choices=choices,
415
+ usage=CompletionUsage(
416
+ prompt_tokens=prompt_tokens,
417
+ completion_tokens=completion_tokens,
418
+ total_tokens=prompt_tokens + completion_tokens,
419
+ ),
420
+ cost=calculate_gemini_cost(self.use_vertexai, prompt_tokens, completion_tokens, model_name),
421
+ )
422
+
423
+ return response_oai
424
+
425
+ def _extract_system_instruction(self, messages: list[dict]) -> str | None:
426
+ """Extract system instruction if provided."""
427
+ if messages is None or len(messages) == 0 or messages[0].get("role") != "system":
428
+ return None
429
+
430
+ message = messages.pop(0)
431
+ content = message["content"]
432
+
433
+ # Multi-model uses a list of dictionaries as content with text for the system message
434
+ # Otherwise normal agents will have strings as content
435
+ content = content[0].get("text", "").strip() if isinstance(content, list) else content.strip()
436
+
437
+ content = content if len(content) > 0 else None
438
+ return content
439
+
440
+ def _oai_content_to_gemini_content(self, message: dict[str, Any]) -> tuple[list[Any], str]:
441
+ """Convert AG2 content to Gemini parts, catering for text and tool calls"""
442
+ rst = []
443
+
444
+ if "role" in message and message["role"] == "tool":
445
+ # Tool call recommendation
446
+
447
+ function_name = self.tool_call_function_map[message["tool_call_id"]]
448
+
449
+ if self.use_vertexai:
450
+ rst.append(
451
+ VertexAIPart.from_function_response(
452
+ name=function_name, response={"result": self._to_json_or_str(message["content"])}
453
+ )
454
+ )
455
+ else:
456
+ rst.append(
457
+ Part(
458
+ function_response=FunctionResponse(
459
+ name=function_name, response={"result": self._to_json_or_str(message["content"])}
460
+ )
461
+ )
462
+ )
463
+
464
+ return rst, "tool"
465
+ elif "tool_calls" in message and len(message["tool_calls"]) != 0:
466
+ for tool_call in message["tool_calls"]:
467
+ function_id = tool_call["id"]
468
+ function_name = tool_call["function"]["name"]
469
+ self.tool_call_function_map[function_id] = function_name
470
+
471
+ if self.use_vertexai:
472
+ rst.append(
473
+ VertexAIPart.from_dict({
474
+ "functionCall": {
475
+ "name": function_name,
476
+ "args": json.loads(tool_call["function"]["arguments"]),
477
+ }
478
+ })
479
+ )
480
+ else:
481
+ rst.append(
482
+ Part(
483
+ function_call=FunctionCall(
484
+ name=function_name,
485
+ args=json.loads(tool_call["function"]["arguments"]),
486
+ )
487
+ )
488
+ )
489
+
490
+ return rst, "tool_call"
491
+
492
+ elif isinstance(message["content"], str):
493
+ content = message["content"]
494
+ if content == "":
495
+ content = "empty" # Empty content is not allowed.
496
+ if self.use_vertexai:
497
+ rst.append(VertexAIPart.from_text(content))
498
+ else:
499
+ rst.append(Part(text=content))
500
+
501
+ return rst, "text"
502
+
503
+ # For images the message contains a list of text items
504
+ if isinstance(message["content"], list):
505
+ has_image = False
506
+ for msg in message["content"]:
507
+ if isinstance(msg, dict):
508
+ assert "type" in msg, f"Missing 'type' field in message: {msg}"
509
+ if msg["type"] == "text":
510
+ if self.use_vertexai:
511
+ rst.append(VertexAIPart.from_text(text=msg["text"]))
512
+ else:
513
+ rst.append(Part(text=msg["text"]))
514
+ elif msg["type"] == "image_url":
515
+ if self.use_vertexai:
516
+ img_url = msg["image_url"]["url"]
517
+ img_part = VertexAIPart.from_uri(img_url, mime_type="image/png")
518
+ rst.append(img_part)
519
+ else:
520
+ b64_img = get_image_data(msg["image_url"]["url"])
521
+ rst.append(Part(inline_data={"mime_type": "image/png", "data": b64_img}))
522
+
523
+ has_image = True
524
+ else:
525
+ raise ValueError(f"Unsupported message type: {msg['type']}")
526
+ else:
527
+ raise ValueError(f"Unsupported message type: {type(msg)}")
528
+ return rst, "image" if has_image else "text"
529
+ else:
530
+ raise Exception("Unable to convert content to Gemini format.")
531
+
532
+ def _concat_parts(self, parts: list[Part]) -> list:
533
+ """Concatenate parts with the same type.
534
+ If two adjacent parts both have the "text" attribute, then it will be joined into one part.
535
+ """
536
+ if not parts:
537
+ return []
538
+
539
+ concatenated_parts = []
540
+ previous_part = parts[0]
541
+
542
+ for current_part in parts[1:]:
543
+ if previous_part.text != "":
544
+ if self.use_vertexai:
545
+ previous_part = VertexAIPart.from_text(previous_part.text + current_part.text)
546
+ else:
547
+ previous_part.text += current_part.text
548
+ else:
549
+ concatenated_parts.append(previous_part)
550
+ previous_part = current_part
551
+
552
+ if previous_part.text == "":
553
+ if self.use_vertexai:
554
+ previous_part = VertexAIPart.from_text("empty")
555
+ else:
556
+ previous_part.text = "empty" # Empty content is not allowed.
557
+ concatenated_parts.append(previous_part)
558
+
559
+ return concatenated_parts
560
+
561
+ def _oai_messages_to_gemini_messages(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
562
+ """Convert messages from OAI format to Gemini format.
563
+ Make sure the "user" role and "model" role are interleaved.
564
+ Also, make sure the last item is from the "user" role.
565
+ """
566
+ rst = []
567
+ for message in messages:
568
+ parts, part_type = self._oai_content_to_gemini_content(message)
569
+ role = "user" if message["role"] in ["user", "system"] else "model"
570
+
571
+ if part_type == "text":
572
+ rst.append(
573
+ VertexAIContent(parts=parts, role=role)
574
+ if self.use_vertexai
575
+ else rst.append(Content(parts=parts, role=role))
576
+ )
577
+ elif part_type == "tool" or part_type == "tool_call":
578
+ role = "function" if version.parse(genai.__version__) < version.parse("1.4.0") else "user"
579
+ rst.append(
580
+ VertexAIContent(parts=parts, role=role)
581
+ if self.use_vertexai
582
+ else rst.append(Content(parts=parts, role=role))
583
+ )
584
+ elif part_type == "image":
585
+ # Image has multiple parts, some can be text and some can be image based
586
+ text_parts = []
587
+ image_parts = []
588
+ for part in parts:
589
+ if isinstance(part, Part):
590
+ # Text or non-Vertex AI image part
591
+ text_parts.append(part)
592
+ elif isinstance(part, VertexAIPart):
593
+ # Image
594
+ image_parts.append(part)
595
+ else:
596
+ raise Exception("Unable to process image part")
597
+
598
+ if len(text_parts) > 0:
599
+ rst.append(
600
+ VertexAIContent(parts=text_parts, role=role)
601
+ if self.use_vertexai
602
+ else rst.append(Content(parts=text_parts, role=role))
603
+ )
604
+
605
+ if len(image_parts) > 0:
606
+ rst.append(
607
+ VertexAIContent(parts=image_parts, role=role)
608
+ if self.use_vertexai
609
+ else rst.append(Content(parts=image_parts, role=role))
610
+ )
611
+
612
+ if len(rst) != 0 and rst[-1] is None:
613
+ rst.pop()
614
+
615
+ # The Gemini is restrict on order of roles, such that
616
+ # 1. The first message must be from the user role.
617
+ # 2. The last message must be from the user role.
618
+ # 3. The messages should be interleaved between user and model.
619
+ # We add a dummy message "start chat" if the first role is not the user.
620
+ # We add a dummy message "continue" if the last role is not the user.
621
+ if rst[0].role != "user":
622
+ text_part, _ = self._oai_content_to_gemini_content({"content": "start chat"})
623
+ rst.insert(
624
+ 0,
625
+ VertexAIContent(parts=text_part, role="user")
626
+ if self.use_vertexai
627
+ else Content(parts=text_part, role="user"),
628
+ )
629
+
630
+ if rst[-1].role != "user":
631
+ text_part, _ = self._oai_content_to_gemini_content({"content": "continue"})
632
+ rst.append(
633
+ VertexAIContent(parts=text_part, role="user")
634
+ if self.use_vertexai
635
+ else Content(parts=text_part, role="user")
636
+ )
637
+
638
+ return rst
639
+
640
+ def _convert_json_response(self, response: str) -> Any:
641
+ """Extract and validate JSON response from the output for structured outputs.
642
+
643
+ Args:
644
+ response (str): The response from the API.
645
+
646
+ Returns:
647
+ Any: The parsed JSON response.
648
+ """
649
+ if not self._response_format:
650
+ return response
651
+
652
+ try:
653
+ # Parse JSON and validate against the Pydantic model if Pydantic model was provided
654
+ json_data = json.loads(response)
655
+ if isinstance(self._response_format, dict):
656
+ return json_data
657
+ else:
658
+ return self._response_format.model_validate(json_data)
659
+ except Exception as e:
660
+ raise ValueError(f"Failed to parse response as valid JSON matching the schema for Structured Output: {e!s}")
661
+
662
+ @staticmethod
663
+ def _convert_type_null_to_nullable(schema: Any) -> Any:
664
+ """
665
+ Recursively converts all occurrences of {"type": "null"} to {"nullable": True} in a schema.
666
+ """
667
+ if isinstance(schema, dict):
668
+ # If schema matches {"type": "null"}, replace it
669
+ if schema == {"type": "null"}:
670
+ return {"nullable": True}
671
+ # Otherwise, recursively process dictionary
672
+ return {key: GeminiClient._convert_type_null_to_nullable(value) for key, value in schema.items()}
673
+ elif isinstance(schema, list):
674
+ # Recursively process list elements
675
+ return [GeminiClient._convert_type_null_to_nullable(item) for item in schema]
676
+ return schema
677
+
678
+ @staticmethod
679
+ def _check_if_prebuilt_google_search_tool_exists(tools: list[dict[str, Any]]) -> bool:
680
+ """Check if the Google Search tool is present in the tools list."""
681
+ exists = False
682
+ for tool in tools:
683
+ if tool["function"]["name"] == "prebuilt_google_search":
684
+ exists = True
685
+ break
686
+
687
+ if exists and len(tools) > 1:
688
+ raise ValueError(
689
+ "Google Search tool can be used only by itself. Please remove other tools from the tools list."
690
+ )
691
+
692
+ return exists
693
+
694
+ @staticmethod
695
+ def _unwrap_references(function_parameters: dict[str, Any]) -> dict[str, Any]:
696
+ if "properties" not in function_parameters:
697
+ return function_parameters
698
+
699
+ function_parameters_copy = copy.deepcopy(function_parameters)
700
+
701
+ for property_name, property_value in function_parameters["properties"].items():
702
+ if "$defs" in property_value:
703
+ function_parameters_copy["properties"][property_name] = resolve_json_references(property_value)
704
+ function_parameters_copy["properties"][property_name].pop("$defs")
705
+
706
+ return function_parameters_copy
707
+
708
+ def _tools_to_gemini_tools(self, tools: list[dict[str, Any]]) -> list[Tool]:
709
+ """Create Gemini tools (as typically requires Callables)"""
710
+ if self._check_if_prebuilt_google_search_tool_exists(tools) and not self.use_vertexai:
711
+ return [Tool(google_search=GoogleSearch())]
712
+
713
+ functions = []
714
+ for tool in tools:
715
+ if self.use_vertexai:
716
+ tool["function"]["parameters"] = GeminiClient._convert_type_null_to_nullable(
717
+ tool["function"]["parameters"]
718
+ )
719
+ function_parameters = GeminiClient._unwrap_references(tool["function"]["parameters"])
720
+ function = vaiFunctionDeclaration(
721
+ name=tool["function"]["name"],
722
+ description=tool["function"]["description"],
723
+ parameters=function_parameters,
724
+ )
725
+ else:
726
+ function = GeminiClient._create_gemini_function_declaration(tool)
727
+ functions.append(function)
728
+
729
+ if self.use_vertexai:
730
+ return [vaiTool(function_declarations=functions)]
731
+ else:
732
+ return [Tool(function_declarations=functions)]
733
+
734
+ @staticmethod
735
+ def _create_gemini_function_declaration(tool: dict) -> FunctionDeclaration:
736
+ function_declaration = FunctionDeclaration()
737
+ function_declaration.name = tool["function"]["name"]
738
+ function_declaration.description = tool["function"]["description"]
739
+ if len(tool["function"]["parameters"]["properties"]) != 0:
740
+ function_declaration.parameters = GeminiClient._create_gemini_function_parameters(
741
+ copy.deepcopy(tool["function"]["parameters"])
742
+ )
743
+
744
+ return function_declaration
745
+
746
+ @staticmethod
747
+ def _create_gemini_function_declaration_schema(json_data) -> Schema:
748
+ """Recursively creates Schema objects for FunctionDeclaration."""
749
+ param_schema = Schema()
750
+ param_type = json_data["type"]
751
+
752
+ """
753
+ TYPE_UNSPECIFIED = 0
754
+ STRING = 1
755
+ INTEGER = 2
756
+ NUMBER = 3
757
+ OBJECT = 4
758
+ ARRAY = 5
759
+ BOOLEAN = 6
760
+ """
761
+
762
+ if param_type == "integer":
763
+ param_schema.type = Type.INTEGER
764
+ elif param_type == "number":
765
+ param_schema.type = Type.NUMBER
766
+ elif param_type == "string":
767
+ param_schema.type = Type.STRING
768
+ elif param_type == "boolean":
769
+ param_schema.type = Type.BOOLEAN
770
+ elif param_type == "array":
771
+ param_schema.type = Type.ARRAY
772
+ if "items" in json_data:
773
+ param_schema.items = GeminiClient._create_gemini_function_declaration_schema(json_data["items"])
774
+ else:
775
+ print("Warning: Array schema missing 'items' definition.")
776
+ elif param_type == "object":
777
+ param_schema.type = Type.OBJECT
778
+ param_schema.properties = {}
779
+ if "properties" in json_data:
780
+ for prop_name, prop_data in json_data["properties"].items():
781
+ param_schema.properties[prop_name] = GeminiClient._create_gemini_function_declaration_schema(
782
+ prop_data
783
+ )
784
+ else:
785
+ print("Warning: Object schema missing 'properties' definition.")
786
+
787
+ elif param_type in ("null", "any"):
788
+ param_schema.type = Type.STRING # Treating these as strings for simplicity
789
+ else:
790
+ print(f"Warning: Unsupported parameter type '{param_type}'.")
791
+
792
+ if "description" in json_data:
793
+ param_schema.description = json_data["description"]
794
+
795
+ return param_schema
796
+
797
+ @staticmethod
798
+ def _create_gemini_function_parameters(function_parameter: dict[str, any]) -> dict[str, any]:
799
+ """Convert function parameters to Gemini format, recursive"""
800
+ function_parameter = GeminiClient._unwrap_references(function_parameter)
801
+
802
+ if "type" in function_parameter:
803
+ function_parameter["type"] = function_parameter["type"].upper()
804
+ # If the schema was created from pydantic BaseModel, it will "title" attribute which needs to be removed
805
+ function_parameter.pop("title", None)
806
+
807
+ # Parameter properties and items
808
+ if "properties" in function_parameter:
809
+ for key in function_parameter["properties"]:
810
+ function_parameter["properties"][key] = GeminiClient._create_gemini_function_parameters(
811
+ function_parameter["properties"][key]
812
+ )
813
+
814
+ if "items" in function_parameter:
815
+ function_parameter["items"] = GeminiClient._create_gemini_function_parameters(function_parameter["items"])
816
+
817
+ # Remove any attributes not needed
818
+ for attr in ["default"]:
819
+ if attr in function_parameter:
820
+ del function_parameter[attr]
821
+
822
+ return function_parameter
823
+
824
+ @staticmethod
825
+ def _to_vertexai_safety_settings(safety_settings):
826
+ """Convert safety settings to VertexAI format if needed,
827
+ like when specifying them in the OAI_CONFIG_LIST
828
+ """
829
+ if isinstance(safety_settings, list) and all([
830
+ isinstance(safety_setting, dict) and not isinstance(safety_setting, VertexAISafetySetting)
831
+ for safety_setting in safety_settings
832
+ ]):
833
+ vertexai_safety_settings = []
834
+ for safety_setting in safety_settings:
835
+ if safety_setting["category"] not in VertexAIHarmCategory.__members__:
836
+ invalid_category = safety_setting["category"]
837
+ logger.error(f"Safety setting category {invalid_category} is invalid")
838
+ elif safety_setting["threshold"] not in VertexAIHarmBlockThreshold.__members__:
839
+ invalid_threshold = safety_setting["threshold"]
840
+ logger.error(f"Safety threshold {invalid_threshold} is invalid")
841
+ else:
842
+ vertexai_safety_setting = VertexAISafetySetting(
843
+ category=safety_setting["category"],
844
+ threshold=safety_setting["threshold"],
845
+ )
846
+ vertexai_safety_settings.append(vertexai_safety_setting)
847
+ return vertexai_safety_settings
848
+ else:
849
+ return safety_settings
850
+
851
+ @staticmethod
852
+ def _to_json_or_str(data: str) -> dict | str:
853
+ try:
854
+ json_data = json.loads(data)
855
+ return json_data
856
+ except (json.JSONDecodeError, ValidationError):
857
+ return data
858
+
859
+
860
+ @require_optional_import(["PIL"], "gemini")
861
+ def get_image_data(image_file: str, use_b64=True) -> bytes:
862
+ if image_file.startswith("http://") or image_file.startswith("https://"):
863
+ response = requests.get(image_file)
864
+ content = response.content
865
+ elif re.match(r"data:image/(?:png|jpeg);base64,", image_file):
866
+ return re.sub(r"data:image/(?:png|jpeg);base64,", "", image_file)
867
+ else:
868
+ image = Image.open(image_file).convert("RGB")
869
+ buffered = BytesIO()
870
+ image.save(buffered, format="PNG")
871
+ content = buffered.getvalue()
872
+
873
+ if use_b64:
874
+ return base64.b64encode(content).decode("utf-8")
875
+ else:
876
+ return content
877
+
878
+
879
+ def _format_json_response(response: Any, original_answer: str) -> str:
880
+ """Formats the JSON response for structured outputs using the format method if it exists."""
881
+ return response.format() if isinstance(response, FormatterProtocol) else original_answer
882
+
883
+
884
+ def calculate_gemini_cost(use_vertexai: bool, input_tokens: int, output_tokens: int, model_name: str) -> float:
885
+ def total_cost_mil(cost_per_mil_input: float, cost_per_mil_output: float):
886
+ # Cost per million
887
+ return cost_per_mil_input * input_tokens / 1e6 + cost_per_mil_output * output_tokens / 1e6
888
+
889
+ def total_cost_k(cost_per_k_input: float, cost_per_k_output: float):
890
+ # Cost per thousand
891
+ return cost_per_k_input * input_tokens / 1e3 + cost_per_k_output * output_tokens / 1e3
892
+
893
+ model_name = model_name.lower()
894
+ up_to_128k = input_tokens <= 128000
895
+ up_to_200k = input_tokens <= 200000
896
+
897
+ if use_vertexai:
898
+ # Vertex AI pricing - based on Text input
899
+ # https://cloud.google.com/vertex-ai/generative-ai/pricing#vertex-ai-pricing
900
+
901
+ if (
902
+ "gemini-2.5-pro-preview-03-25" in model_name
903
+ or "gemini-2.5-pro-exp-03-25" in model_name
904
+ or "gemini-2.5-pro-preview-05-06" in model_name
905
+ ):
906
+ if up_to_200k:
907
+ return total_cost_mil(1.25, 10)
908
+ else:
909
+ return total_cost_mil(2.5, 15)
910
+
911
+ elif "gemini-2.5-flash-preview-04-17" in model_name:
912
+ return total_cost_mil(0.15, 0.6) # NON-THINKING OUTPUT PRICE, $3 FOR THINKING!
913
+
914
+ elif "gemini-2.0-flash-lite" in model_name:
915
+ return total_cost_mil(0.075, 0.3)
916
+
917
+ elif "gemini-2.0-flash" in model_name:
918
+ return total_cost_mil(0.15, 0.6)
919
+
920
+ elif "gemini-1.5-flash" in model_name:
921
+ if up_to_128k:
922
+ return total_cost_k(0.00001875, 0.000075)
923
+ else:
924
+ return total_cost_k(0.0000375, 0.00015)
925
+
926
+ elif "gemini-1.5-pro" in model_name:
927
+ if up_to_128k:
928
+ return total_cost_k(0.0003125, 0.00125)
929
+ else:
930
+ return total_cost_k(0.000625, 0.0025)
931
+
932
+ elif "gemini-1.0-pro" in model_name:
933
+ return total_cost_k(0.000125, 0.00001875)
934
+
935
+ else:
936
+ warnings.warn(
937
+ f"Cost calculation is not implemented for model {model_name}. Cost will be calculated zero.",
938
+ UserWarning,
939
+ )
940
+ return 0
941
+
942
+ else:
943
+ # Non-Vertex AI pricing
944
+
945
+ if (
946
+ "gemini-2.5-pro-preview-03-25" in model_name
947
+ or "gemini-2.5-pro-exp-03-25" in model_name
948
+ or "gemini-2.5-pro-preview-05-06" in model_name
949
+ ):
950
+ # https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-pro-preview
951
+ if up_to_200k:
952
+ return total_cost_mil(1.25, 10)
953
+ else:
954
+ return total_cost_mil(2.5, 15)
955
+
956
+ elif "gemini-2.5-flash-preview-04-17" in model_name:
957
+ # https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-flash
958
+ return total_cost_mil(0.15, 0.6)
959
+
960
+ elif "gemini-2.0-flash-lite" in model_name:
961
+ # https://ai.google.dev/gemini-api/docs/pricing#gemini-2.0-flash-lite
962
+ return total_cost_mil(0.075, 0.3)
963
+
964
+ elif "gemini-2.0-flash" in model_name:
965
+ # https://ai.google.dev/gemini-api/docs/pricing#gemini-2.0-flash
966
+ return total_cost_mil(0.1, 0.4)
967
+
968
+ elif "gemini-1.5-flash-8b" in model_name:
969
+ # https://ai.google.dev/pricing#1_5flash-8B
970
+ if up_to_128k:
971
+ return total_cost_mil(0.0375, 0.15)
972
+ else:
973
+ return total_cost_mil(0.075, 0.3)
974
+
975
+ elif "gemini-1.5-flash" in model_name:
976
+ # https://ai.google.dev/pricing#1_5flash
977
+ if up_to_128k:
978
+ return total_cost_mil(0.075, 0.3)
979
+ else:
980
+ return total_cost_mil(0.15, 0.6)
981
+
982
+ elif "gemini-1.5-pro" in model_name:
983
+ # https://ai.google.dev/pricing#1_5pro
984
+ if up_to_128k:
985
+ return total_cost_mil(1.25, 5.0)
986
+ else:
987
+ return total_cost_mil(2.50, 10.0)
988
+
989
+ elif "gemini-1.0-pro" in model_name:
990
+ # https://ai.google.dev/pricing#1_5pro
991
+ return total_cost_mil(0.50, 1.5)
992
+
993
+ else:
994
+ warnings.warn(
995
+ f"Cost calculation is not implemented for model {model_name}. Cost will be calculated zero.",
996
+ UserWarning,
997
+ )
998
+ return 0