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.
- {ag2-0.9.1a1.dist-info → ag2-0.9.2.dist-info}/METADATA +272 -75
- ag2-0.9.2.dist-info/RECORD +406 -0
- {ag2-0.9.1a1.dist-info → ag2-0.9.2.dist-info}/WHEEL +1 -2
- autogen/__init__.py +89 -0
- autogen/_website/__init__.py +3 -0
- autogen/_website/generate_api_references.py +427 -0
- autogen/_website/generate_mkdocs.py +1174 -0
- autogen/_website/notebook_processor.py +476 -0
- autogen/_website/process_notebooks.py +656 -0
- autogen/_website/utils.py +412 -0
- autogen/agentchat/__init__.py +44 -0
- autogen/agentchat/agent.py +182 -0
- autogen/agentchat/assistant_agent.py +85 -0
- autogen/agentchat/chat.py +309 -0
- autogen/agentchat/contrib/__init__.py +5 -0
- autogen/agentchat/contrib/agent_eval/README.md +7 -0
- autogen/agentchat/contrib/agent_eval/agent_eval.py +108 -0
- autogen/agentchat/contrib/agent_eval/criterion.py +43 -0
- autogen/agentchat/contrib/agent_eval/critic_agent.py +44 -0
- autogen/agentchat/contrib/agent_eval/quantifier_agent.py +39 -0
- autogen/agentchat/contrib/agent_eval/subcritic_agent.py +45 -0
- autogen/agentchat/contrib/agent_eval/task.py +42 -0
- autogen/agentchat/contrib/agent_optimizer.py +429 -0
- autogen/agentchat/contrib/capabilities/__init__.py +5 -0
- autogen/agentchat/contrib/capabilities/agent_capability.py +20 -0
- autogen/agentchat/contrib/capabilities/generate_images.py +301 -0
- autogen/agentchat/contrib/capabilities/teachability.py +393 -0
- autogen/agentchat/contrib/capabilities/text_compressors.py +66 -0
- autogen/agentchat/contrib/capabilities/tools_capability.py +22 -0
- autogen/agentchat/contrib/capabilities/transform_messages.py +93 -0
- autogen/agentchat/contrib/capabilities/transforms.py +566 -0
- autogen/agentchat/contrib/capabilities/transforms_util.py +122 -0
- autogen/agentchat/contrib/capabilities/vision_capability.py +214 -0
- autogen/agentchat/contrib/captainagent/__init__.py +9 -0
- autogen/agentchat/contrib/captainagent/agent_builder.py +790 -0
- autogen/agentchat/contrib/captainagent/captainagent.py +512 -0
- autogen/agentchat/contrib/captainagent/tool_retriever.py +335 -0
- autogen/agentchat/contrib/captainagent/tools/README.md +44 -0
- autogen/agentchat/contrib/captainagent/tools/__init__.py +5 -0
- autogen/agentchat/contrib/captainagent/tools/data_analysis/calculate_correlation.py +40 -0
- autogen/agentchat/contrib/captainagent/tools/data_analysis/calculate_skewness_and_kurtosis.py +28 -0
- autogen/agentchat/contrib/captainagent/tools/data_analysis/detect_outlier_iqr.py +28 -0
- autogen/agentchat/contrib/captainagent/tools/data_analysis/detect_outlier_zscore.py +28 -0
- autogen/agentchat/contrib/captainagent/tools/data_analysis/explore_csv.py +21 -0
- autogen/agentchat/contrib/captainagent/tools/data_analysis/shapiro_wilk_test.py +30 -0
- autogen/agentchat/contrib/captainagent/tools/information_retrieval/arxiv_download.py +27 -0
- autogen/agentchat/contrib/captainagent/tools/information_retrieval/arxiv_search.py +53 -0
- autogen/agentchat/contrib/captainagent/tools/information_retrieval/extract_pdf_image.py +53 -0
- autogen/agentchat/contrib/captainagent/tools/information_retrieval/extract_pdf_text.py +38 -0
- autogen/agentchat/contrib/captainagent/tools/information_retrieval/get_wikipedia_text.py +21 -0
- autogen/agentchat/contrib/captainagent/tools/information_retrieval/get_youtube_caption.py +34 -0
- autogen/agentchat/contrib/captainagent/tools/information_retrieval/image_qa.py +60 -0
- autogen/agentchat/contrib/captainagent/tools/information_retrieval/optical_character_recognition.py +61 -0
- autogen/agentchat/contrib/captainagent/tools/information_retrieval/perform_web_search.py +47 -0
- autogen/agentchat/contrib/captainagent/tools/information_retrieval/scrape_wikipedia_tables.py +33 -0
- autogen/agentchat/contrib/captainagent/tools/information_retrieval/transcribe_audio_file.py +21 -0
- autogen/agentchat/contrib/captainagent/tools/information_retrieval/youtube_download.py +35 -0
- autogen/agentchat/contrib/captainagent/tools/math/calculate_circle_area_from_diameter.py +21 -0
- autogen/agentchat/contrib/captainagent/tools/math/calculate_day_of_the_week.py +18 -0
- autogen/agentchat/contrib/captainagent/tools/math/calculate_fraction_sum.py +28 -0
- autogen/agentchat/contrib/captainagent/tools/math/calculate_matrix_power.py +31 -0
- autogen/agentchat/contrib/captainagent/tools/math/calculate_reflected_point.py +16 -0
- autogen/agentchat/contrib/captainagent/tools/math/complex_numbers_product.py +25 -0
- autogen/agentchat/contrib/captainagent/tools/math/compute_currency_conversion.py +23 -0
- autogen/agentchat/contrib/captainagent/tools/math/count_distinct_permutations.py +27 -0
- autogen/agentchat/contrib/captainagent/tools/math/evaluate_expression.py +28 -0
- autogen/agentchat/contrib/captainagent/tools/math/find_continuity_point.py +34 -0
- autogen/agentchat/contrib/captainagent/tools/math/fraction_to_mixed_numbers.py +39 -0
- autogen/agentchat/contrib/captainagent/tools/math/modular_inverse_sum.py +23 -0
- autogen/agentchat/contrib/captainagent/tools/math/simplify_mixed_numbers.py +36 -0
- autogen/agentchat/contrib/captainagent/tools/math/sum_of_digit_factorials.py +15 -0
- autogen/agentchat/contrib/captainagent/tools/math/sum_of_primes_below.py +15 -0
- autogen/agentchat/contrib/captainagent/tools/requirements.txt +10 -0
- autogen/agentchat/contrib/captainagent/tools/tool_description.tsv +34 -0
- autogen/agentchat/contrib/gpt_assistant_agent.py +526 -0
- autogen/agentchat/contrib/graph_rag/__init__.py +9 -0
- autogen/agentchat/contrib/graph_rag/document.py +29 -0
- autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py +170 -0
- autogen/agentchat/contrib/graph_rag/falkor_graph_rag_capability.py +103 -0
- autogen/agentchat/contrib/graph_rag/graph_query_engine.py +53 -0
- autogen/agentchat/contrib/graph_rag/graph_rag_capability.py +63 -0
- autogen/agentchat/contrib/graph_rag/neo4j_graph_query_engine.py +268 -0
- autogen/agentchat/contrib/graph_rag/neo4j_graph_rag_capability.py +83 -0
- autogen/agentchat/contrib/graph_rag/neo4j_native_graph_query_engine.py +210 -0
- autogen/agentchat/contrib/graph_rag/neo4j_native_graph_rag_capability.py +93 -0
- autogen/agentchat/contrib/img_utils.py +397 -0
- autogen/agentchat/contrib/llamaindex_conversable_agent.py +117 -0
- autogen/agentchat/contrib/llava_agent.py +187 -0
- autogen/agentchat/contrib/math_user_proxy_agent.py +464 -0
- autogen/agentchat/contrib/multimodal_conversable_agent.py +125 -0
- autogen/agentchat/contrib/qdrant_retrieve_user_proxy_agent.py +324 -0
- autogen/agentchat/contrib/rag/__init__.py +10 -0
- autogen/agentchat/contrib/rag/chromadb_query_engine.py +272 -0
- autogen/agentchat/contrib/rag/llamaindex_query_engine.py +198 -0
- autogen/agentchat/contrib/rag/mongodb_query_engine.py +329 -0
- autogen/agentchat/contrib/rag/query_engine.py +74 -0
- autogen/agentchat/contrib/retrieve_assistant_agent.py +56 -0
- autogen/agentchat/contrib/retrieve_user_proxy_agent.py +703 -0
- autogen/agentchat/contrib/society_of_mind_agent.py +199 -0
- autogen/agentchat/contrib/swarm_agent.py +1425 -0
- autogen/agentchat/contrib/text_analyzer_agent.py +79 -0
- autogen/agentchat/contrib/vectordb/__init__.py +5 -0
- autogen/agentchat/contrib/vectordb/base.py +232 -0
- autogen/agentchat/contrib/vectordb/chromadb.py +315 -0
- autogen/agentchat/contrib/vectordb/couchbase.py +407 -0
- autogen/agentchat/contrib/vectordb/mongodb.py +550 -0
- autogen/agentchat/contrib/vectordb/pgvectordb.py +928 -0
- autogen/agentchat/contrib/vectordb/qdrant.py +320 -0
- autogen/agentchat/contrib/vectordb/utils.py +126 -0
- autogen/agentchat/contrib/web_surfer.py +303 -0
- autogen/agentchat/conversable_agent.py +4023 -0
- autogen/agentchat/group/__init__.py +64 -0
- autogen/agentchat/group/available_condition.py +91 -0
- autogen/agentchat/group/context_condition.py +77 -0
- autogen/agentchat/group/context_expression.py +238 -0
- autogen/agentchat/group/context_str.py +41 -0
- autogen/agentchat/group/context_variables.py +192 -0
- autogen/agentchat/group/group_tool_executor.py +202 -0
- autogen/agentchat/group/group_utils.py +591 -0
- autogen/agentchat/group/handoffs.py +244 -0
- autogen/agentchat/group/llm_condition.py +93 -0
- autogen/agentchat/group/multi_agent_chat.py +237 -0
- autogen/agentchat/group/on_condition.py +58 -0
- autogen/agentchat/group/on_context_condition.py +54 -0
- autogen/agentchat/group/patterns/__init__.py +18 -0
- autogen/agentchat/group/patterns/auto.py +159 -0
- autogen/agentchat/group/patterns/manual.py +176 -0
- autogen/agentchat/group/patterns/pattern.py +288 -0
- autogen/agentchat/group/patterns/random.py +106 -0
- autogen/agentchat/group/patterns/round_robin.py +117 -0
- autogen/agentchat/group/reply_result.py +26 -0
- autogen/agentchat/group/speaker_selection_result.py +41 -0
- autogen/agentchat/group/targets/__init__.py +4 -0
- autogen/agentchat/group/targets/group_chat_target.py +132 -0
- autogen/agentchat/group/targets/group_manager_target.py +151 -0
- autogen/agentchat/group/targets/transition_target.py +413 -0
- autogen/agentchat/group/targets/transition_utils.py +6 -0
- autogen/agentchat/groupchat.py +1694 -0
- autogen/agentchat/realtime/__init__.py +3 -0
- autogen/agentchat/realtime/experimental/__init__.py +20 -0
- autogen/agentchat/realtime/experimental/audio_adapters/__init__.py +8 -0
- autogen/agentchat/realtime/experimental/audio_adapters/twilio_audio_adapter.py +148 -0
- autogen/agentchat/realtime/experimental/audio_adapters/websocket_audio_adapter.py +139 -0
- autogen/agentchat/realtime/experimental/audio_observer.py +42 -0
- autogen/agentchat/realtime/experimental/clients/__init__.py +15 -0
- autogen/agentchat/realtime/experimental/clients/gemini/__init__.py +7 -0
- autogen/agentchat/realtime/experimental/clients/gemini/client.py +274 -0
- autogen/agentchat/realtime/experimental/clients/oai/__init__.py +8 -0
- autogen/agentchat/realtime/experimental/clients/oai/base_client.py +220 -0
- autogen/agentchat/realtime/experimental/clients/oai/rtc_client.py +243 -0
- autogen/agentchat/realtime/experimental/clients/oai/utils.py +48 -0
- autogen/agentchat/realtime/experimental/clients/realtime_client.py +190 -0
- autogen/agentchat/realtime/experimental/function_observer.py +85 -0
- autogen/agentchat/realtime/experimental/realtime_agent.py +158 -0
- autogen/agentchat/realtime/experimental/realtime_events.py +42 -0
- autogen/agentchat/realtime/experimental/realtime_observer.py +100 -0
- autogen/agentchat/realtime/experimental/realtime_swarm.py +475 -0
- autogen/agentchat/realtime/experimental/websockets.py +21 -0
- autogen/agentchat/realtime_agent/__init__.py +21 -0
- autogen/agentchat/user_proxy_agent.py +111 -0
- autogen/agentchat/utils.py +206 -0
- autogen/agents/__init__.py +3 -0
- autogen/agents/contrib/__init__.py +10 -0
- autogen/agents/contrib/time/__init__.py +8 -0
- autogen/agents/contrib/time/time_reply_agent.py +73 -0
- autogen/agents/contrib/time/time_tool_agent.py +51 -0
- autogen/agents/experimental/__init__.py +27 -0
- autogen/agents/experimental/deep_research/__init__.py +7 -0
- autogen/agents/experimental/deep_research/deep_research.py +52 -0
- autogen/agents/experimental/discord/__init__.py +7 -0
- autogen/agents/experimental/discord/discord.py +66 -0
- autogen/agents/experimental/document_agent/__init__.py +19 -0
- autogen/agents/experimental/document_agent/chroma_query_engine.py +316 -0
- autogen/agents/experimental/document_agent/docling_doc_ingest_agent.py +118 -0
- autogen/agents/experimental/document_agent/document_agent.py +461 -0
- autogen/agents/experimental/document_agent/document_conditions.py +50 -0
- autogen/agents/experimental/document_agent/document_utils.py +380 -0
- autogen/agents/experimental/document_agent/inmemory_query_engine.py +220 -0
- autogen/agents/experimental/document_agent/parser_utils.py +130 -0
- autogen/agents/experimental/document_agent/url_utils.py +426 -0
- autogen/agents/experimental/reasoning/__init__.py +7 -0
- autogen/agents/experimental/reasoning/reasoning_agent.py +1178 -0
- autogen/agents/experimental/slack/__init__.py +7 -0
- autogen/agents/experimental/slack/slack.py +73 -0
- autogen/agents/experimental/telegram/__init__.py +7 -0
- autogen/agents/experimental/telegram/telegram.py +77 -0
- autogen/agents/experimental/websurfer/__init__.py +7 -0
- autogen/agents/experimental/websurfer/websurfer.py +62 -0
- autogen/agents/experimental/wikipedia/__init__.py +7 -0
- autogen/agents/experimental/wikipedia/wikipedia.py +90 -0
- autogen/browser_utils.py +309 -0
- autogen/cache/__init__.py +10 -0
- autogen/cache/abstract_cache_base.py +75 -0
- autogen/cache/cache.py +203 -0
- autogen/cache/cache_factory.py +88 -0
- autogen/cache/cosmos_db_cache.py +144 -0
- autogen/cache/disk_cache.py +102 -0
- autogen/cache/in_memory_cache.py +58 -0
- autogen/cache/redis_cache.py +123 -0
- autogen/code_utils.py +596 -0
- autogen/coding/__init__.py +22 -0
- autogen/coding/base.py +119 -0
- autogen/coding/docker_commandline_code_executor.py +268 -0
- autogen/coding/factory.py +47 -0
- autogen/coding/func_with_reqs.py +202 -0
- autogen/coding/jupyter/__init__.py +23 -0
- autogen/coding/jupyter/base.py +36 -0
- autogen/coding/jupyter/docker_jupyter_server.py +167 -0
- autogen/coding/jupyter/embedded_ipython_code_executor.py +182 -0
- autogen/coding/jupyter/import_utils.py +82 -0
- autogen/coding/jupyter/jupyter_client.py +231 -0
- autogen/coding/jupyter/jupyter_code_executor.py +160 -0
- autogen/coding/jupyter/local_jupyter_server.py +172 -0
- autogen/coding/local_commandline_code_executor.py +405 -0
- autogen/coding/markdown_code_extractor.py +45 -0
- autogen/coding/utils.py +56 -0
- autogen/doc_utils.py +34 -0
- autogen/events/__init__.py +7 -0
- autogen/events/agent_events.py +1013 -0
- autogen/events/base_event.py +99 -0
- autogen/events/client_events.py +167 -0
- autogen/events/helpers.py +36 -0
- autogen/events/print_event.py +46 -0
- autogen/exception_utils.py +73 -0
- autogen/extensions/__init__.py +5 -0
- autogen/fast_depends/__init__.py +16 -0
- autogen/fast_depends/_compat.py +80 -0
- autogen/fast_depends/core/__init__.py +14 -0
- autogen/fast_depends/core/build.py +225 -0
- autogen/fast_depends/core/model.py +576 -0
- autogen/fast_depends/dependencies/__init__.py +15 -0
- autogen/fast_depends/dependencies/model.py +29 -0
- autogen/fast_depends/dependencies/provider.py +39 -0
- autogen/fast_depends/library/__init__.py +10 -0
- autogen/fast_depends/library/model.py +46 -0
- autogen/fast_depends/py.typed +6 -0
- autogen/fast_depends/schema.py +66 -0
- autogen/fast_depends/use.py +280 -0
- autogen/fast_depends/utils.py +187 -0
- autogen/formatting_utils.py +83 -0
- autogen/function_utils.py +13 -0
- autogen/graph_utils.py +178 -0
- autogen/import_utils.py +526 -0
- autogen/interop/__init__.py +22 -0
- autogen/interop/crewai/__init__.py +7 -0
- autogen/interop/crewai/crewai.py +88 -0
- autogen/interop/interoperability.py +71 -0
- autogen/interop/interoperable.py +46 -0
- autogen/interop/langchain/__init__.py +8 -0
- autogen/interop/langchain/langchain_chat_model_factory.py +155 -0
- autogen/interop/langchain/langchain_tool.py +82 -0
- autogen/interop/litellm/__init__.py +7 -0
- autogen/interop/litellm/litellm_config_factory.py +179 -0
- autogen/interop/pydantic_ai/__init__.py +7 -0
- autogen/interop/pydantic_ai/pydantic_ai.py +168 -0
- autogen/interop/registry.py +69 -0
- autogen/io/__init__.py +15 -0
- autogen/io/base.py +151 -0
- autogen/io/console.py +56 -0
- autogen/io/processors/__init__.py +12 -0
- autogen/io/processors/base.py +21 -0
- autogen/io/processors/console_event_processor.py +56 -0
- autogen/io/run_response.py +293 -0
- autogen/io/thread_io_stream.py +63 -0
- autogen/io/websockets.py +213 -0
- autogen/json_utils.py +43 -0
- autogen/llm_config.py +382 -0
- autogen/logger/__init__.py +11 -0
- autogen/logger/base_logger.py +128 -0
- autogen/logger/file_logger.py +261 -0
- autogen/logger/logger_factory.py +42 -0
- autogen/logger/logger_utils.py +57 -0
- autogen/logger/sqlite_logger.py +523 -0
- autogen/math_utils.py +339 -0
- autogen/mcp/__init__.py +7 -0
- autogen/mcp/__main__.py +78 -0
- autogen/mcp/mcp_client.py +208 -0
- autogen/mcp/mcp_proxy/__init__.py +19 -0
- autogen/mcp/mcp_proxy/fastapi_code_generator_helpers.py +63 -0
- autogen/mcp/mcp_proxy/mcp_proxy.py +581 -0
- autogen/mcp/mcp_proxy/operation_grouping.py +158 -0
- autogen/mcp/mcp_proxy/operation_renaming.py +114 -0
- autogen/mcp/mcp_proxy/patch_fastapi_code_generator.py +98 -0
- autogen/mcp/mcp_proxy/security.py +400 -0
- autogen/mcp/mcp_proxy/security_schema_visitor.py +37 -0
- autogen/messages/__init__.py +7 -0
- autogen/messages/agent_messages.py +948 -0
- autogen/messages/base_message.py +107 -0
- autogen/messages/client_messages.py +171 -0
- autogen/messages/print_message.py +49 -0
- autogen/oai/__init__.py +53 -0
- autogen/oai/anthropic.py +714 -0
- autogen/oai/bedrock.py +628 -0
- autogen/oai/cerebras.py +299 -0
- autogen/oai/client.py +1444 -0
- autogen/oai/client_utils.py +169 -0
- autogen/oai/cohere.py +479 -0
- autogen/oai/gemini.py +998 -0
- autogen/oai/gemini_types.py +155 -0
- autogen/oai/groq.py +305 -0
- autogen/oai/mistral.py +303 -0
- autogen/oai/oai_models/__init__.py +11 -0
- autogen/oai/oai_models/_models.py +16 -0
- autogen/oai/oai_models/chat_completion.py +87 -0
- autogen/oai/oai_models/chat_completion_audio.py +32 -0
- autogen/oai/oai_models/chat_completion_message.py +86 -0
- autogen/oai/oai_models/chat_completion_message_tool_call.py +37 -0
- autogen/oai/oai_models/chat_completion_token_logprob.py +63 -0
- autogen/oai/oai_models/completion_usage.py +60 -0
- autogen/oai/ollama.py +643 -0
- autogen/oai/openai_utils.py +881 -0
- autogen/oai/together.py +370 -0
- autogen/retrieve_utils.py +491 -0
- autogen/runtime_logging.py +160 -0
- autogen/token_count_utils.py +267 -0
- autogen/tools/__init__.py +20 -0
- autogen/tools/contrib/__init__.py +9 -0
- autogen/tools/contrib/time/__init__.py +7 -0
- autogen/tools/contrib/time/time.py +41 -0
- autogen/tools/dependency_injection.py +254 -0
- autogen/tools/experimental/__init__.py +48 -0
- autogen/tools/experimental/browser_use/__init__.py +7 -0
- autogen/tools/experimental/browser_use/browser_use.py +161 -0
- autogen/tools/experimental/crawl4ai/__init__.py +7 -0
- autogen/tools/experimental/crawl4ai/crawl4ai.py +153 -0
- autogen/tools/experimental/deep_research/__init__.py +7 -0
- autogen/tools/experimental/deep_research/deep_research.py +328 -0
- autogen/tools/experimental/duckduckgo/__init__.py +7 -0
- autogen/tools/experimental/duckduckgo/duckduckgo_search.py +109 -0
- autogen/tools/experimental/google/__init__.py +14 -0
- autogen/tools/experimental/google/authentication/__init__.py +11 -0
- autogen/tools/experimental/google/authentication/credentials_hosted_provider.py +43 -0
- autogen/tools/experimental/google/authentication/credentials_local_provider.py +91 -0
- autogen/tools/experimental/google/authentication/credentials_provider.py +35 -0
- autogen/tools/experimental/google/drive/__init__.py +9 -0
- autogen/tools/experimental/google/drive/drive_functions.py +124 -0
- autogen/tools/experimental/google/drive/toolkit.py +88 -0
- autogen/tools/experimental/google/model.py +17 -0
- autogen/tools/experimental/google/toolkit_protocol.py +19 -0
- autogen/tools/experimental/google_search/__init__.py +8 -0
- autogen/tools/experimental/google_search/google_search.py +93 -0
- autogen/tools/experimental/google_search/youtube_search.py +181 -0
- autogen/tools/experimental/messageplatform/__init__.py +17 -0
- autogen/tools/experimental/messageplatform/discord/__init__.py +7 -0
- autogen/tools/experimental/messageplatform/discord/discord.py +288 -0
- autogen/tools/experimental/messageplatform/slack/__init__.py +7 -0
- autogen/tools/experimental/messageplatform/slack/slack.py +391 -0
- autogen/tools/experimental/messageplatform/telegram/__init__.py +7 -0
- autogen/tools/experimental/messageplatform/telegram/telegram.py +275 -0
- autogen/tools/experimental/perplexity/__init__.py +7 -0
- autogen/tools/experimental/perplexity/perplexity_search.py +260 -0
- autogen/tools/experimental/reliable/__init__.py +10 -0
- autogen/tools/experimental/reliable/reliable.py +1316 -0
- autogen/tools/experimental/tavily/__init__.py +7 -0
- autogen/tools/experimental/tavily/tavily_search.py +183 -0
- autogen/tools/experimental/web_search_preview/__init__.py +7 -0
- autogen/tools/experimental/web_search_preview/web_search_preview.py +114 -0
- autogen/tools/experimental/wikipedia/__init__.py +7 -0
- autogen/tools/experimental/wikipedia/wikipedia.py +287 -0
- autogen/tools/function_utils.py +411 -0
- autogen/tools/tool.py +187 -0
- autogen/tools/toolkit.py +86 -0
- autogen/types.py +29 -0
- autogen/version.py +7 -0
- templates/client_template/main.jinja2 +69 -0
- templates/config_template/config.jinja2 +7 -0
- templates/main.jinja2 +61 -0
- ag2-0.9.1a1.dist-info/RECORD +0 -6
- ag2-0.9.1a1.dist-info/top_level.txt +0 -1
- {ag2-0.9.1a1.dist-info → ag2-0.9.2.dist-info/licenses}/LICENSE +0 -0
- {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
|