agno 2.1.2__py3-none-any.whl → 2.3.13__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- agno/agent/agent.py +5540 -2273
- agno/api/api.py +2 -0
- agno/api/os.py +1 -1
- agno/compression/__init__.py +3 -0
- agno/compression/manager.py +247 -0
- agno/culture/__init__.py +3 -0
- agno/culture/manager.py +956 -0
- agno/db/async_postgres/__init__.py +3 -0
- agno/db/base.py +689 -6
- agno/db/dynamo/dynamo.py +933 -37
- agno/db/dynamo/schemas.py +174 -10
- agno/db/dynamo/utils.py +63 -4
- agno/db/firestore/firestore.py +831 -9
- agno/db/firestore/schemas.py +51 -0
- agno/db/firestore/utils.py +102 -4
- agno/db/gcs_json/gcs_json_db.py +660 -12
- agno/db/gcs_json/utils.py +60 -26
- agno/db/in_memory/in_memory_db.py +287 -14
- agno/db/in_memory/utils.py +60 -2
- agno/db/json/json_db.py +590 -14
- agno/db/json/utils.py +60 -26
- agno/db/migrations/manager.py +199 -0
- agno/db/migrations/v1_to_v2.py +43 -13
- agno/db/migrations/versions/__init__.py +0 -0
- agno/db/migrations/versions/v2_3_0.py +938 -0
- agno/db/mongo/__init__.py +15 -1
- agno/db/mongo/async_mongo.py +2760 -0
- agno/db/mongo/mongo.py +879 -11
- agno/db/mongo/schemas.py +42 -0
- agno/db/mongo/utils.py +80 -8
- agno/db/mysql/__init__.py +2 -1
- agno/db/mysql/async_mysql.py +2912 -0
- agno/db/mysql/mysql.py +946 -68
- agno/db/mysql/schemas.py +72 -10
- agno/db/mysql/utils.py +198 -7
- agno/db/postgres/__init__.py +2 -1
- agno/db/postgres/async_postgres.py +2579 -0
- agno/db/postgres/postgres.py +942 -57
- agno/db/postgres/schemas.py +81 -18
- agno/db/postgres/utils.py +164 -2
- agno/db/redis/redis.py +671 -7
- agno/db/redis/schemas.py +50 -0
- agno/db/redis/utils.py +65 -7
- agno/db/schemas/__init__.py +2 -1
- agno/db/schemas/culture.py +120 -0
- agno/db/schemas/evals.py +1 -0
- agno/db/schemas/memory.py +17 -2
- agno/db/singlestore/schemas.py +63 -0
- agno/db/singlestore/singlestore.py +949 -83
- agno/db/singlestore/utils.py +60 -2
- agno/db/sqlite/__init__.py +2 -1
- agno/db/sqlite/async_sqlite.py +2911 -0
- agno/db/sqlite/schemas.py +62 -0
- agno/db/sqlite/sqlite.py +965 -46
- agno/db/sqlite/utils.py +169 -8
- agno/db/surrealdb/__init__.py +3 -0
- agno/db/surrealdb/metrics.py +292 -0
- agno/db/surrealdb/models.py +334 -0
- agno/db/surrealdb/queries.py +71 -0
- agno/db/surrealdb/surrealdb.py +1908 -0
- agno/db/surrealdb/utils.py +147 -0
- agno/db/utils.py +2 -0
- agno/eval/__init__.py +10 -0
- agno/eval/accuracy.py +75 -55
- agno/eval/agent_as_judge.py +861 -0
- agno/eval/base.py +29 -0
- agno/eval/performance.py +16 -7
- agno/eval/reliability.py +28 -16
- agno/eval/utils.py +35 -17
- agno/exceptions.py +27 -2
- agno/filters.py +354 -0
- agno/guardrails/prompt_injection.py +1 -0
- agno/hooks/__init__.py +3 -0
- agno/hooks/decorator.py +164 -0
- agno/integrations/discord/client.py +1 -1
- agno/knowledge/chunking/agentic.py +13 -10
- agno/knowledge/chunking/fixed.py +4 -1
- agno/knowledge/chunking/semantic.py +9 -4
- agno/knowledge/chunking/strategy.py +59 -15
- agno/knowledge/embedder/fastembed.py +1 -1
- agno/knowledge/embedder/nebius.py +1 -1
- agno/knowledge/embedder/ollama.py +8 -0
- agno/knowledge/embedder/openai.py +8 -8
- agno/knowledge/embedder/sentence_transformer.py +6 -2
- agno/knowledge/embedder/vllm.py +262 -0
- agno/knowledge/knowledge.py +1618 -318
- agno/knowledge/reader/base.py +6 -2
- agno/knowledge/reader/csv_reader.py +8 -10
- agno/knowledge/reader/docx_reader.py +5 -6
- agno/knowledge/reader/field_labeled_csv_reader.py +16 -20
- agno/knowledge/reader/json_reader.py +5 -4
- agno/knowledge/reader/markdown_reader.py +8 -8
- agno/knowledge/reader/pdf_reader.py +17 -19
- agno/knowledge/reader/pptx_reader.py +101 -0
- agno/knowledge/reader/reader_factory.py +32 -3
- agno/knowledge/reader/s3_reader.py +3 -3
- agno/knowledge/reader/tavily_reader.py +193 -0
- agno/knowledge/reader/text_reader.py +22 -10
- agno/knowledge/reader/web_search_reader.py +1 -48
- agno/knowledge/reader/website_reader.py +10 -10
- agno/knowledge/reader/wikipedia_reader.py +33 -1
- agno/knowledge/types.py +1 -0
- agno/knowledge/utils.py +72 -7
- agno/media.py +22 -6
- agno/memory/__init__.py +14 -1
- agno/memory/manager.py +544 -83
- agno/memory/strategies/__init__.py +15 -0
- agno/memory/strategies/base.py +66 -0
- agno/memory/strategies/summarize.py +196 -0
- agno/memory/strategies/types.py +37 -0
- agno/models/aimlapi/aimlapi.py +17 -0
- agno/models/anthropic/claude.py +515 -40
- agno/models/aws/bedrock.py +102 -21
- agno/models/aws/claude.py +131 -274
- agno/models/azure/ai_foundry.py +41 -19
- agno/models/azure/openai_chat.py +39 -8
- agno/models/base.py +1249 -525
- agno/models/cerebras/cerebras.py +91 -21
- agno/models/cerebras/cerebras_openai.py +21 -2
- agno/models/cohere/chat.py +40 -6
- agno/models/cometapi/cometapi.py +18 -1
- agno/models/dashscope/dashscope.py +2 -3
- agno/models/deepinfra/deepinfra.py +18 -1
- agno/models/deepseek/deepseek.py +69 -3
- agno/models/fireworks/fireworks.py +18 -1
- agno/models/google/gemini.py +877 -80
- agno/models/google/utils.py +22 -0
- agno/models/groq/groq.py +51 -18
- agno/models/huggingface/huggingface.py +17 -6
- agno/models/ibm/watsonx.py +16 -6
- agno/models/internlm/internlm.py +18 -1
- agno/models/langdb/langdb.py +13 -1
- agno/models/litellm/chat.py +44 -9
- agno/models/litellm/litellm_openai.py +18 -1
- agno/models/message.py +28 -5
- agno/models/meta/llama.py +47 -14
- agno/models/meta/llama_openai.py +22 -17
- agno/models/mistral/mistral.py +8 -4
- agno/models/nebius/nebius.py +6 -7
- agno/models/nvidia/nvidia.py +20 -3
- agno/models/ollama/chat.py +24 -8
- agno/models/openai/chat.py +104 -29
- agno/models/openai/responses.py +101 -81
- agno/models/openrouter/openrouter.py +60 -3
- agno/models/perplexity/perplexity.py +17 -1
- agno/models/portkey/portkey.py +7 -6
- agno/models/requesty/requesty.py +24 -4
- agno/models/response.py +73 -2
- agno/models/sambanova/sambanova.py +20 -3
- agno/models/siliconflow/siliconflow.py +19 -2
- agno/models/together/together.py +20 -3
- agno/models/utils.py +254 -8
- agno/models/vercel/v0.py +20 -3
- agno/models/vertexai/__init__.py +0 -0
- agno/models/vertexai/claude.py +190 -0
- agno/models/vllm/vllm.py +19 -14
- agno/models/xai/xai.py +19 -2
- agno/os/app.py +549 -152
- agno/os/auth.py +190 -3
- agno/os/config.py +23 -0
- agno/os/interfaces/a2a/router.py +8 -11
- agno/os/interfaces/a2a/utils.py +1 -1
- agno/os/interfaces/agui/router.py +18 -3
- agno/os/interfaces/agui/utils.py +152 -39
- agno/os/interfaces/slack/router.py +55 -37
- agno/os/interfaces/slack/slack.py +9 -1
- agno/os/interfaces/whatsapp/router.py +0 -1
- agno/os/interfaces/whatsapp/security.py +3 -1
- agno/os/mcp.py +110 -52
- agno/os/middleware/__init__.py +2 -0
- agno/os/middleware/jwt.py +676 -112
- agno/os/router.py +40 -1478
- agno/os/routers/agents/__init__.py +3 -0
- agno/os/routers/agents/router.py +599 -0
- agno/os/routers/agents/schema.py +261 -0
- agno/os/routers/evals/evals.py +96 -39
- agno/os/routers/evals/schemas.py +65 -33
- agno/os/routers/evals/utils.py +80 -10
- agno/os/routers/health.py +10 -4
- agno/os/routers/knowledge/knowledge.py +196 -38
- agno/os/routers/knowledge/schemas.py +82 -22
- agno/os/routers/memory/memory.py +279 -52
- agno/os/routers/memory/schemas.py +46 -17
- agno/os/routers/metrics/metrics.py +20 -8
- agno/os/routers/metrics/schemas.py +16 -16
- agno/os/routers/session/session.py +462 -34
- agno/os/routers/teams/__init__.py +3 -0
- agno/os/routers/teams/router.py +512 -0
- agno/os/routers/teams/schema.py +257 -0
- agno/os/routers/traces/__init__.py +3 -0
- agno/os/routers/traces/schemas.py +414 -0
- agno/os/routers/traces/traces.py +499 -0
- agno/os/routers/workflows/__init__.py +3 -0
- agno/os/routers/workflows/router.py +624 -0
- agno/os/routers/workflows/schema.py +75 -0
- agno/os/schema.py +256 -693
- agno/os/scopes.py +469 -0
- agno/os/utils.py +514 -36
- agno/reasoning/anthropic.py +80 -0
- agno/reasoning/gemini.py +73 -0
- agno/reasoning/openai.py +5 -0
- agno/reasoning/vertexai.py +76 -0
- agno/run/__init__.py +6 -0
- agno/run/agent.py +155 -32
- agno/run/base.py +55 -3
- agno/run/requirement.py +181 -0
- agno/run/team.py +125 -38
- agno/run/workflow.py +72 -18
- agno/session/agent.py +102 -89
- agno/session/summary.py +56 -15
- agno/session/team.py +164 -90
- agno/session/workflow.py +405 -40
- agno/table.py +10 -0
- agno/team/team.py +3974 -1903
- agno/tools/dalle.py +2 -4
- agno/tools/eleven_labs.py +23 -25
- agno/tools/exa.py +21 -16
- agno/tools/file.py +153 -23
- agno/tools/file_generation.py +16 -10
- agno/tools/firecrawl.py +15 -7
- agno/tools/function.py +193 -38
- agno/tools/gmail.py +238 -14
- agno/tools/google_drive.py +271 -0
- agno/tools/googlecalendar.py +36 -8
- agno/tools/googlesheets.py +20 -5
- agno/tools/jira.py +20 -0
- agno/tools/mcp/__init__.py +10 -0
- agno/tools/mcp/mcp.py +331 -0
- agno/tools/mcp/multi_mcp.py +347 -0
- agno/tools/mcp/params.py +24 -0
- agno/tools/mcp_toolbox.py +3 -3
- agno/tools/models/nebius.py +5 -5
- agno/tools/models_labs.py +20 -10
- agno/tools/nano_banana.py +151 -0
- agno/tools/notion.py +204 -0
- agno/tools/parallel.py +314 -0
- agno/tools/postgres.py +76 -36
- agno/tools/redshift.py +406 -0
- agno/tools/scrapegraph.py +1 -1
- agno/tools/shopify.py +1519 -0
- agno/tools/slack.py +18 -3
- agno/tools/spotify.py +919 -0
- agno/tools/tavily.py +146 -0
- agno/tools/toolkit.py +25 -0
- agno/tools/workflow.py +8 -1
- agno/tools/yfinance.py +12 -11
- agno/tracing/__init__.py +12 -0
- agno/tracing/exporter.py +157 -0
- agno/tracing/schemas.py +276 -0
- agno/tracing/setup.py +111 -0
- agno/utils/agent.py +938 -0
- agno/utils/cryptography.py +22 -0
- agno/utils/dttm.py +33 -0
- agno/utils/events.py +151 -3
- agno/utils/gemini.py +15 -5
- agno/utils/hooks.py +118 -4
- agno/utils/http.py +113 -2
- agno/utils/knowledge.py +12 -5
- agno/utils/log.py +1 -0
- agno/utils/mcp.py +92 -2
- agno/utils/media.py +187 -1
- agno/utils/merge_dict.py +3 -3
- agno/utils/message.py +60 -0
- agno/utils/models/ai_foundry.py +9 -2
- agno/utils/models/claude.py +49 -14
- agno/utils/models/cohere.py +9 -2
- agno/utils/models/llama.py +9 -2
- agno/utils/models/mistral.py +4 -2
- agno/utils/print_response/agent.py +109 -16
- agno/utils/print_response/team.py +223 -30
- agno/utils/print_response/workflow.py +251 -34
- agno/utils/streamlit.py +1 -1
- agno/utils/team.py +98 -9
- agno/utils/tokens.py +657 -0
- agno/vectordb/base.py +39 -7
- agno/vectordb/cassandra/cassandra.py +21 -5
- agno/vectordb/chroma/chromadb.py +43 -12
- agno/vectordb/clickhouse/clickhousedb.py +21 -5
- agno/vectordb/couchbase/couchbase.py +29 -5
- agno/vectordb/lancedb/lance_db.py +92 -181
- agno/vectordb/langchaindb/langchaindb.py +24 -4
- agno/vectordb/lightrag/lightrag.py +17 -3
- agno/vectordb/llamaindex/llamaindexdb.py +25 -5
- agno/vectordb/milvus/milvus.py +50 -37
- agno/vectordb/mongodb/__init__.py +7 -1
- agno/vectordb/mongodb/mongodb.py +36 -30
- agno/vectordb/pgvector/pgvector.py +201 -77
- agno/vectordb/pineconedb/pineconedb.py +41 -23
- agno/vectordb/qdrant/qdrant.py +67 -54
- agno/vectordb/redis/__init__.py +9 -0
- agno/vectordb/redis/redisdb.py +682 -0
- agno/vectordb/singlestore/singlestore.py +50 -29
- agno/vectordb/surrealdb/surrealdb.py +31 -41
- agno/vectordb/upstashdb/upstashdb.py +34 -6
- agno/vectordb/weaviate/weaviate.py +53 -14
- agno/workflow/__init__.py +2 -0
- agno/workflow/agent.py +299 -0
- agno/workflow/condition.py +120 -18
- agno/workflow/loop.py +77 -10
- agno/workflow/parallel.py +231 -143
- agno/workflow/router.py +118 -17
- agno/workflow/step.py +609 -170
- agno/workflow/steps.py +73 -6
- agno/workflow/types.py +96 -21
- agno/workflow/workflow.py +2039 -262
- {agno-2.1.2.dist-info → agno-2.3.13.dist-info}/METADATA +201 -66
- agno-2.3.13.dist-info/RECORD +613 -0
- agno/tools/googlesearch.py +0 -98
- agno/tools/mcp.py +0 -679
- agno/tools/memori.py +0 -339
- agno-2.1.2.dist-info/RECORD +0 -543
- {agno-2.1.2.dist-info → agno-2.3.13.dist-info}/WHEEL +0 -0
- {agno-2.1.2.dist-info → agno-2.3.13.dist-info}/licenses/LICENSE +0 -0
- {agno-2.1.2.dist-info → agno-2.3.13.dist-info}/top_level.txt +0 -0
|
@@ -0,0 +1,347 @@
|
|
|
1
|
+
import weakref
|
|
2
|
+
from contextlib import AsyncExitStack
|
|
3
|
+
from dataclasses import asdict
|
|
4
|
+
from datetime import timedelta
|
|
5
|
+
from types import TracebackType
|
|
6
|
+
from typing import List, Literal, Optional, Union
|
|
7
|
+
|
|
8
|
+
from agno.tools import Toolkit
|
|
9
|
+
from agno.tools.function import Function
|
|
10
|
+
from agno.tools.mcp.params import SSEClientParams, StreamableHTTPClientParams
|
|
11
|
+
from agno.utils.log import log_debug, log_error, log_info, log_warning
|
|
12
|
+
from agno.utils.mcp import get_entrypoint_for_tool, prepare_command
|
|
13
|
+
|
|
14
|
+
try:
|
|
15
|
+
from mcp import ClientSession, StdioServerParameters
|
|
16
|
+
from mcp.client.sse import sse_client
|
|
17
|
+
from mcp.client.stdio import get_default_environment, stdio_client
|
|
18
|
+
from mcp.client.streamable_http import streamablehttp_client
|
|
19
|
+
except (ImportError, ModuleNotFoundError):
|
|
20
|
+
raise ImportError("`mcp` not installed. Please install using `pip install mcp`")
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
class MultiMCPTools(Toolkit):
|
|
24
|
+
"""
|
|
25
|
+
A toolkit for integrating multiple Model Context Protocol (MCP) servers with Agno agents.
|
|
26
|
+
This allows agents to access tools, resources, and prompts exposed by MCP servers.
|
|
27
|
+
|
|
28
|
+
Can be used in three ways:
|
|
29
|
+
1. Direct initialization with a ClientSession
|
|
30
|
+
2. As an async context manager with StdioServerParameters
|
|
31
|
+
3. As an async context manager with SSE or Streamable HTTP endpoints
|
|
32
|
+
"""
|
|
33
|
+
|
|
34
|
+
def __init__(
|
|
35
|
+
self,
|
|
36
|
+
commands: Optional[List[str]] = None,
|
|
37
|
+
urls: Optional[List[str]] = None,
|
|
38
|
+
urls_transports: Optional[List[Literal["sse", "streamable-http"]]] = None,
|
|
39
|
+
*,
|
|
40
|
+
env: Optional[dict[str, str]] = None,
|
|
41
|
+
server_params_list: Optional[
|
|
42
|
+
list[Union[SSEClientParams, StdioServerParameters, StreamableHTTPClientParams]]
|
|
43
|
+
] = None,
|
|
44
|
+
timeout_seconds: int = 10,
|
|
45
|
+
client=None,
|
|
46
|
+
include_tools: Optional[list[str]] = None,
|
|
47
|
+
exclude_tools: Optional[list[str]] = None,
|
|
48
|
+
refresh_connection: bool = False,
|
|
49
|
+
allow_partial_failure: bool = False,
|
|
50
|
+
**kwargs,
|
|
51
|
+
):
|
|
52
|
+
"""
|
|
53
|
+
Initialize the MCP toolkit.
|
|
54
|
+
|
|
55
|
+
Args:
|
|
56
|
+
commands: List of commands to run to start the servers. Should be used in conjunction with env.
|
|
57
|
+
urls: List of URLs for SSE and/or Streamable HTTP endpoints.
|
|
58
|
+
urls_transports: List of transports to use for the given URLs.
|
|
59
|
+
server_params_list: List of StdioServerParameters or SSEClientParams or StreamableHTTPClientParams for creating new sessions.
|
|
60
|
+
env: The environment variables to pass to the servers. Should be used in conjunction with commands.
|
|
61
|
+
client: The underlying MCP client (optional, used to prevent garbage collection).
|
|
62
|
+
timeout_seconds: Timeout in seconds for managing timeouts for Client Session if Agent or Tool doesn't respond.
|
|
63
|
+
include_tools: Optional list of tool names to include (if None, includes all).
|
|
64
|
+
exclude_tools: Optional list of tool names to exclude (if None, excludes none).
|
|
65
|
+
allow_partial_failure: If True, allows toolkit to initialize even if some MCP servers fail to connect. If False, any failure will raise an exception.
|
|
66
|
+
refresh_connection: If True, the connection and tools will be refreshed on each run
|
|
67
|
+
"""
|
|
68
|
+
super().__init__(name="MultiMCPTools", **kwargs)
|
|
69
|
+
|
|
70
|
+
if urls_transports is not None:
|
|
71
|
+
if "sse" in urls_transports:
|
|
72
|
+
log_info("SSE as a standalone transport is deprecated. Please use Streamable HTTP instead.")
|
|
73
|
+
|
|
74
|
+
if urls is not None:
|
|
75
|
+
if urls_transports is None:
|
|
76
|
+
log_warning(
|
|
77
|
+
"The default transport 'streamable-http' will be used. You can explicitly set the transports by providing the urls_transports parameter."
|
|
78
|
+
)
|
|
79
|
+
else:
|
|
80
|
+
if len(urls) != len(urls_transports):
|
|
81
|
+
raise ValueError("urls and urls_transports must be of the same length")
|
|
82
|
+
|
|
83
|
+
# Set these after `__init__` to bypass the `_check_tools_filters`
|
|
84
|
+
# beacuse tools are not available until `initialize()` is called.
|
|
85
|
+
self.include_tools = include_tools
|
|
86
|
+
self.exclude_tools = exclude_tools
|
|
87
|
+
self.refresh_connection = refresh_connection
|
|
88
|
+
|
|
89
|
+
if server_params_list is None and commands is None and urls is None:
|
|
90
|
+
raise ValueError("Either server_params_list or commands or urls must be provided")
|
|
91
|
+
|
|
92
|
+
self.server_params_list: List[Union[SSEClientParams, StdioServerParameters, StreamableHTTPClientParams]] = (
|
|
93
|
+
server_params_list or []
|
|
94
|
+
)
|
|
95
|
+
self.timeout_seconds = timeout_seconds
|
|
96
|
+
self.commands: Optional[List[str]] = commands
|
|
97
|
+
self.urls: Optional[List[str]] = urls
|
|
98
|
+
# Merge provided env with system env
|
|
99
|
+
if env is not None:
|
|
100
|
+
env = {
|
|
101
|
+
**get_default_environment(),
|
|
102
|
+
**env,
|
|
103
|
+
}
|
|
104
|
+
else:
|
|
105
|
+
env = get_default_environment()
|
|
106
|
+
|
|
107
|
+
if commands is not None:
|
|
108
|
+
for command in commands:
|
|
109
|
+
parts = prepare_command(command)
|
|
110
|
+
cmd = parts[0]
|
|
111
|
+
arguments = parts[1:] if len(parts) > 1 else []
|
|
112
|
+
self.server_params_list.append(StdioServerParameters(command=cmd, args=arguments, env=env))
|
|
113
|
+
|
|
114
|
+
if urls is not None:
|
|
115
|
+
if urls_transports is not None:
|
|
116
|
+
for url, transport in zip(urls, urls_transports):
|
|
117
|
+
if transport == "streamable-http":
|
|
118
|
+
self.server_params_list.append(StreamableHTTPClientParams(url=url))
|
|
119
|
+
else:
|
|
120
|
+
self.server_params_list.append(SSEClientParams(url=url))
|
|
121
|
+
else:
|
|
122
|
+
for url in urls:
|
|
123
|
+
self.server_params_list.append(StreamableHTTPClientParams(url=url))
|
|
124
|
+
|
|
125
|
+
self._async_exit_stack = AsyncExitStack()
|
|
126
|
+
|
|
127
|
+
self._client = client
|
|
128
|
+
|
|
129
|
+
self._initialized = False
|
|
130
|
+
self._connection_task = None
|
|
131
|
+
self._successful_connections = 0
|
|
132
|
+
self._sessions: list[ClientSession] = []
|
|
133
|
+
|
|
134
|
+
self.allow_partial_failure = allow_partial_failure
|
|
135
|
+
|
|
136
|
+
def cleanup():
|
|
137
|
+
"""Cancel active connections"""
|
|
138
|
+
if self._connection_task and not self._connection_task.done():
|
|
139
|
+
self._connection_task.cancel()
|
|
140
|
+
|
|
141
|
+
# Setup cleanup logic before the instance is garbage collected
|
|
142
|
+
self._cleanup_finalizer = weakref.finalize(self, cleanup)
|
|
143
|
+
|
|
144
|
+
@property
|
|
145
|
+
def initialized(self) -> bool:
|
|
146
|
+
return self._initialized
|
|
147
|
+
|
|
148
|
+
async def is_alive(self) -> bool:
|
|
149
|
+
try:
|
|
150
|
+
for session in self._sessions:
|
|
151
|
+
await session.send_ping()
|
|
152
|
+
return True
|
|
153
|
+
except (RuntimeError, BaseException):
|
|
154
|
+
return False
|
|
155
|
+
|
|
156
|
+
async def connect(self, force: bool = False):
|
|
157
|
+
"""Initialize a MultiMCPTools instance and connect to the MCP servers"""
|
|
158
|
+
|
|
159
|
+
if force:
|
|
160
|
+
# Clean up the session and context so we force a new connection
|
|
161
|
+
self._sessions = []
|
|
162
|
+
self._successful_connections = 0
|
|
163
|
+
self._initialized = False
|
|
164
|
+
self._connection_task = None
|
|
165
|
+
|
|
166
|
+
if self._initialized:
|
|
167
|
+
return
|
|
168
|
+
|
|
169
|
+
try:
|
|
170
|
+
await self._connect()
|
|
171
|
+
except (RuntimeError, BaseException) as e:
|
|
172
|
+
log_error(f"Failed to connect to {str(self)}: {e}")
|
|
173
|
+
|
|
174
|
+
@classmethod
|
|
175
|
+
async def create_and_connect(
|
|
176
|
+
cls,
|
|
177
|
+
commands: Optional[List[str]] = None,
|
|
178
|
+
urls: Optional[List[str]] = None,
|
|
179
|
+
urls_transports: Optional[List[Literal["sse", "streamable-http"]]] = None,
|
|
180
|
+
*,
|
|
181
|
+
env: Optional[dict[str, str]] = None,
|
|
182
|
+
server_params_list: Optional[
|
|
183
|
+
List[Union[SSEClientParams, StdioServerParameters, StreamableHTTPClientParams]]
|
|
184
|
+
] = None,
|
|
185
|
+
timeout_seconds: int = 5,
|
|
186
|
+
client=None,
|
|
187
|
+
include_tools: Optional[list[str]] = None,
|
|
188
|
+
exclude_tools: Optional[list[str]] = None,
|
|
189
|
+
refresh_connection: bool = False,
|
|
190
|
+
**kwargs,
|
|
191
|
+
) -> "MultiMCPTools":
|
|
192
|
+
"""Initialize a MultiMCPTools instance and connect to the MCP servers"""
|
|
193
|
+
instance = cls(
|
|
194
|
+
commands=commands,
|
|
195
|
+
urls=urls,
|
|
196
|
+
urls_transports=urls_transports,
|
|
197
|
+
env=env,
|
|
198
|
+
server_params_list=server_params_list,
|
|
199
|
+
timeout_seconds=timeout_seconds,
|
|
200
|
+
client=client,
|
|
201
|
+
include_tools=include_tools,
|
|
202
|
+
exclude_tools=exclude_tools,
|
|
203
|
+
refresh_connection=refresh_connection,
|
|
204
|
+
**kwargs,
|
|
205
|
+
)
|
|
206
|
+
|
|
207
|
+
await instance._connect()
|
|
208
|
+
return instance
|
|
209
|
+
|
|
210
|
+
async def _connect(self) -> None:
|
|
211
|
+
"""Connects to the MCP servers and initializes the tools"""
|
|
212
|
+
if self._initialized:
|
|
213
|
+
return
|
|
214
|
+
|
|
215
|
+
server_connection_errors = []
|
|
216
|
+
|
|
217
|
+
for server_params in self.server_params_list:
|
|
218
|
+
try:
|
|
219
|
+
# Handle stdio connections
|
|
220
|
+
if isinstance(server_params, StdioServerParameters):
|
|
221
|
+
stdio_transport = await self._async_exit_stack.enter_async_context(stdio_client(server_params))
|
|
222
|
+
read, write = stdio_transport
|
|
223
|
+
session = await self._async_exit_stack.enter_async_context(
|
|
224
|
+
ClientSession(read, write, read_timeout_seconds=timedelta(seconds=self.timeout_seconds))
|
|
225
|
+
)
|
|
226
|
+
await self.initialize(session)
|
|
227
|
+
self._successful_connections += 1
|
|
228
|
+
|
|
229
|
+
# Handle SSE connections
|
|
230
|
+
elif isinstance(server_params, SSEClientParams):
|
|
231
|
+
client_connection = await self._async_exit_stack.enter_async_context(
|
|
232
|
+
sse_client(**asdict(server_params))
|
|
233
|
+
)
|
|
234
|
+
read, write = client_connection
|
|
235
|
+
session = await self._async_exit_stack.enter_async_context(ClientSession(read, write))
|
|
236
|
+
await self.initialize(session)
|
|
237
|
+
self._successful_connections += 1
|
|
238
|
+
|
|
239
|
+
# Handle Streamable HTTP connections
|
|
240
|
+
elif isinstance(server_params, StreamableHTTPClientParams):
|
|
241
|
+
client_connection = await self._async_exit_stack.enter_async_context(
|
|
242
|
+
streamablehttp_client(**asdict(server_params))
|
|
243
|
+
)
|
|
244
|
+
read, write = client_connection[0:2]
|
|
245
|
+
session = await self._async_exit_stack.enter_async_context(ClientSession(read, write))
|
|
246
|
+
await self.initialize(session)
|
|
247
|
+
self._successful_connections += 1
|
|
248
|
+
|
|
249
|
+
except Exception as e:
|
|
250
|
+
if not self.allow_partial_failure:
|
|
251
|
+
raise ValueError(f"MCP connection failed: {e}")
|
|
252
|
+
|
|
253
|
+
log_error(f"Failed to initialize MCP server with params {server_params}: {e}")
|
|
254
|
+
server_connection_errors.append(str(e))
|
|
255
|
+
continue
|
|
256
|
+
|
|
257
|
+
if self._successful_connections > 0:
|
|
258
|
+
await self.build_tools()
|
|
259
|
+
|
|
260
|
+
if self._successful_connections == 0 and server_connection_errors:
|
|
261
|
+
raise ValueError(f"All MCP connections failed: {server_connection_errors}")
|
|
262
|
+
|
|
263
|
+
if not self._initialized and self._successful_connections > 0:
|
|
264
|
+
self._initialized = True
|
|
265
|
+
|
|
266
|
+
async def close(self) -> None:
|
|
267
|
+
"""Close the MCP connections and clean up resources"""
|
|
268
|
+
if not self._initialized:
|
|
269
|
+
return
|
|
270
|
+
|
|
271
|
+
try:
|
|
272
|
+
await self._async_exit_stack.aclose()
|
|
273
|
+
self._sessions = []
|
|
274
|
+
self._successful_connections = 0
|
|
275
|
+
|
|
276
|
+
except (RuntimeError, BaseException) as e:
|
|
277
|
+
log_error(f"Failed to close MCP connections: {e}")
|
|
278
|
+
|
|
279
|
+
self._initialized = False
|
|
280
|
+
|
|
281
|
+
async def __aenter__(self) -> "MultiMCPTools":
|
|
282
|
+
"""Enter the async context manager."""
|
|
283
|
+
try:
|
|
284
|
+
await self._connect()
|
|
285
|
+
except (RuntimeError, BaseException) as e:
|
|
286
|
+
log_error(f"Failed to connect to {str(self)}: {e}")
|
|
287
|
+
return self
|
|
288
|
+
|
|
289
|
+
async def __aexit__(
|
|
290
|
+
self,
|
|
291
|
+
exc_type: Union[type[BaseException], None],
|
|
292
|
+
exc_val: Union[BaseException, None],
|
|
293
|
+
exc_tb: Union[TracebackType, None],
|
|
294
|
+
):
|
|
295
|
+
"""Exit the async context manager."""
|
|
296
|
+
await self._async_exit_stack.aclose()
|
|
297
|
+
self._initialized = False
|
|
298
|
+
self._successful_connections = 0
|
|
299
|
+
|
|
300
|
+
async def build_tools(self) -> None:
|
|
301
|
+
for session in self._sessions:
|
|
302
|
+
# Get the list of tools from the MCP server
|
|
303
|
+
available_tools = await session.list_tools()
|
|
304
|
+
|
|
305
|
+
# Filter tools based on include/exclude lists
|
|
306
|
+
filtered_tools = []
|
|
307
|
+
for tool in available_tools.tools:
|
|
308
|
+
if self.exclude_tools and tool.name in self.exclude_tools:
|
|
309
|
+
continue
|
|
310
|
+
if self.include_tools is None or tool.name in self.include_tools:
|
|
311
|
+
filtered_tools.append(tool)
|
|
312
|
+
|
|
313
|
+
# Register the tools with the toolkit
|
|
314
|
+
for tool in filtered_tools:
|
|
315
|
+
try:
|
|
316
|
+
# Get an entrypoint for the tool
|
|
317
|
+
entrypoint = get_entrypoint_for_tool(tool, session)
|
|
318
|
+
|
|
319
|
+
# Create a Function for the tool
|
|
320
|
+
f = Function(
|
|
321
|
+
name=tool.name,
|
|
322
|
+
description=tool.description,
|
|
323
|
+
parameters=tool.inputSchema,
|
|
324
|
+
entrypoint=entrypoint,
|
|
325
|
+
# Set skip_entrypoint_processing to True to avoid processing the entrypoint
|
|
326
|
+
skip_entrypoint_processing=True,
|
|
327
|
+
)
|
|
328
|
+
|
|
329
|
+
# Register the Function with the toolkit
|
|
330
|
+
self.functions[f.name] = f
|
|
331
|
+
log_debug(f"Function: {f.name} registered with {self.name}")
|
|
332
|
+
except Exception as e:
|
|
333
|
+
log_error(f"Failed to register tool {tool.name}: {e}")
|
|
334
|
+
raise
|
|
335
|
+
|
|
336
|
+
async def initialize(self, session: ClientSession) -> None:
|
|
337
|
+
"""Initialize the MCP toolkit by getting available tools from the MCP server"""
|
|
338
|
+
|
|
339
|
+
try:
|
|
340
|
+
# Initialize the session if not already initialized
|
|
341
|
+
await session.initialize()
|
|
342
|
+
|
|
343
|
+
self._sessions.append(session)
|
|
344
|
+
self._initialized = True
|
|
345
|
+
except Exception as e:
|
|
346
|
+
log_error(f"Failed to get MCP tools: {e}")
|
|
347
|
+
raise
|
agno/tools/mcp/params.py
ADDED
|
@@ -0,0 +1,24 @@
|
|
|
1
|
+
from dataclasses import dataclass
|
|
2
|
+
from datetime import timedelta
|
|
3
|
+
from typing import Any, Dict, Optional
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
@dataclass
|
|
7
|
+
class SSEClientParams:
|
|
8
|
+
"""Parameters for SSE client connection."""
|
|
9
|
+
|
|
10
|
+
url: str
|
|
11
|
+
headers: Optional[Dict[str, Any]] = None
|
|
12
|
+
timeout: Optional[float] = 5
|
|
13
|
+
sse_read_timeout: Optional[float] = 60 * 5
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
@dataclass
|
|
17
|
+
class StreamableHTTPClientParams:
|
|
18
|
+
"""Parameters for Streamable HTTP client connection."""
|
|
19
|
+
|
|
20
|
+
url: str
|
|
21
|
+
headers: Optional[Dict[str, Any]] = None
|
|
22
|
+
timeout: Optional[timedelta] = timedelta(seconds=30)
|
|
23
|
+
sse_read_timeout: Optional[timedelta] = timedelta(seconds=60 * 5)
|
|
24
|
+
terminate_on_close: Optional[bool] = None
|
agno/tools/mcp_toolbox.py
CHANGED
|
@@ -35,6 +35,7 @@ class MCPToolbox(MCPTools, metaclass=MCPToolsMeta):
|
|
|
35
35
|
tool_name: Optional[str] = None,
|
|
36
36
|
headers: Optional[Dict[str, Any]] = None,
|
|
37
37
|
transport: Literal["stdio", "sse", "streamable-http"] = "streamable-http",
|
|
38
|
+
append_mcp_to_url: bool = True,
|
|
38
39
|
**kwargs,
|
|
39
40
|
):
|
|
40
41
|
"""Initialize MCPToolbox with filtering capabilities.
|
|
@@ -45,11 +46,10 @@ class MCPToolbox(MCPTools, metaclass=MCPToolsMeta):
|
|
|
45
46
|
tool_name (Optional[str], optional): Single tool name to load. Defaults to None.
|
|
46
47
|
headers (Optional[Dict[str, Any]], optional): Headers for toolbox-core client requests. Defaults to None.
|
|
47
48
|
transport (Literal["stdio", "sse", "streamable-http"], optional): MCP transport protocol. Defaults to "streamable-http".
|
|
49
|
+
append_mcp_to_url (bool, optional): Whether to append "/mcp" to the URL if it doesn't end with it. Defaults to True.
|
|
48
50
|
|
|
49
51
|
"""
|
|
50
|
-
|
|
51
|
-
# Ensure the URL ends in "/mcp" as expected
|
|
52
|
-
if not url.endswith("/mcp"):
|
|
52
|
+
if append_mcp_to_url and not url.endswith("/mcp"):
|
|
53
53
|
url = url + "/mcp"
|
|
54
54
|
|
|
55
55
|
super().__init__(url=url, transport=transport, **kwargs)
|
agno/tools/models/nebius.py
CHANGED
|
@@ -12,12 +12,12 @@ from agno.utils.log import log_error, log_warning
|
|
|
12
12
|
|
|
13
13
|
|
|
14
14
|
class NebiusTools(Toolkit):
|
|
15
|
-
"""Tools for interacting with Nebius
|
|
15
|
+
"""Tools for interacting with Nebius Token Factory's text-to-image API"""
|
|
16
16
|
|
|
17
17
|
def __init__(
|
|
18
18
|
self,
|
|
19
19
|
api_key: Optional[str] = None,
|
|
20
|
-
base_url: str = "https://api.
|
|
20
|
+
base_url: str = "https://api.tokenfactory.nebius.com/v1",
|
|
21
21
|
image_model: str = "black-forest-labs/flux-schnell",
|
|
22
22
|
image_quality: Optional[str] = "standard",
|
|
23
23
|
image_size: Optional[str] = "1024x1024",
|
|
@@ -26,11 +26,11 @@ class NebiusTools(Toolkit):
|
|
|
26
26
|
all: bool = False,
|
|
27
27
|
**kwargs,
|
|
28
28
|
):
|
|
29
|
-
"""Initialize Nebius
|
|
29
|
+
"""Initialize Nebius Token Factory text-to-image tools.
|
|
30
30
|
|
|
31
31
|
Args:
|
|
32
32
|
api_key: Nebius API key. If not provided, will look for NEBIUS_API_KEY environment variable.
|
|
33
|
-
base_url: The base URL for the Nebius
|
|
33
|
+
base_url: The base URL for the Nebius Token Factory API. This should be configured according to Nebius's documentation.
|
|
34
34
|
image_model: The model to use for generation. Options include:
|
|
35
35
|
- "black-forest-labs/flux-schnell" (fastest)
|
|
36
36
|
- "black-forest-labs/flux-dev" (balanced)
|
|
@@ -69,7 +69,7 @@ class NebiusTools(Toolkit):
|
|
|
69
69
|
agent: Agent,
|
|
70
70
|
prompt: str,
|
|
71
71
|
) -> ToolResult:
|
|
72
|
-
"""Generate images based on a text prompt using Nebius
|
|
72
|
+
"""Generate images based on a text prompt using Nebius Token Factory.
|
|
73
73
|
|
|
74
74
|
Args:
|
|
75
75
|
agent: The agent instance for adding images
|
agno/tools/models_labs.py
CHANGED
|
@@ -4,10 +4,8 @@ from os import getenv
|
|
|
4
4
|
from typing import Any, Dict, List, Optional, Union
|
|
5
5
|
from uuid import uuid4
|
|
6
6
|
|
|
7
|
-
from agno.agent import Agent
|
|
8
7
|
from agno.media import Audio, Image, Video
|
|
9
8
|
from agno.models.response import FileType
|
|
10
|
-
from agno.team import Team
|
|
11
9
|
from agno.tools import Toolkit
|
|
12
10
|
from agno.tools.function import ToolResult
|
|
13
11
|
from agno.utils.log import log_debug, log_info, logger
|
|
@@ -22,12 +20,14 @@ MODELS_LAB_URLS = {
|
|
|
22
20
|
"MP4": "https://modelslab.com/api/v6/video/text2video",
|
|
23
21
|
"MP3": "https://modelslab.com/api/v6/voice/music_gen",
|
|
24
22
|
"GIF": "https://modelslab.com/api/v6/video/text2video",
|
|
23
|
+
"WAV": "https://modelslab.com/api/v6/voice/sfx",
|
|
25
24
|
}
|
|
26
25
|
|
|
27
26
|
MODELS_LAB_FETCH_URLS = {
|
|
28
27
|
"MP4": "https://modelslab.com/api/v6/video/fetch",
|
|
29
28
|
"MP3": "https://modelslab.com/api/v6/voice/fetch",
|
|
30
29
|
"GIF": "https://modelslab.com/api/v6/video/fetch",
|
|
30
|
+
"WAV": "https://modelslab.com/api/v6/voice/fetch",
|
|
31
31
|
}
|
|
32
32
|
|
|
33
33
|
|
|
@@ -78,6 +78,13 @@ class ModelsLabTools(Toolkit):
|
|
|
78
78
|
"output_type": self.file_type.value,
|
|
79
79
|
}
|
|
80
80
|
base_payload |= video_template # Use |= instead of update()
|
|
81
|
+
elif self.file_type == FileType.WAV:
|
|
82
|
+
sfx_template = {
|
|
83
|
+
"duration": 10,
|
|
84
|
+
"output_format": "wav",
|
|
85
|
+
"temp": False,
|
|
86
|
+
}
|
|
87
|
+
base_payload |= sfx_template # Use |= instead of update()
|
|
81
88
|
else:
|
|
82
89
|
audio_template = {
|
|
83
90
|
"base64": False,
|
|
@@ -101,7 +108,7 @@ class ModelsLabTools(Toolkit):
|
|
|
101
108
|
elif self.file_type == FileType.GIF:
|
|
102
109
|
image_artifact = Image(id=str(media_id), url=media_url)
|
|
103
110
|
artifacts["images"].append(image_artifact)
|
|
104
|
-
elif self.file_type
|
|
111
|
+
elif self.file_type in [FileType.MP3, FileType.WAV]:
|
|
105
112
|
audio_artifact = Audio(id=str(media_id), url=media_url)
|
|
106
113
|
artifacts["audios"].append(audio_artifact)
|
|
107
114
|
|
|
@@ -131,7 +138,7 @@ class ModelsLabTools(Toolkit):
|
|
|
131
138
|
|
|
132
139
|
return False
|
|
133
140
|
|
|
134
|
-
def generate_media(self,
|
|
141
|
+
def generate_media(self, prompt: str) -> ToolResult:
|
|
135
142
|
"""Generate media (video, image, or audio) given a prompt."""
|
|
136
143
|
if not self.api_key:
|
|
137
144
|
return ToolResult(content="Please set the MODELS_LAB_API_KEY")
|
|
@@ -157,7 +164,6 @@ class ModelsLabTools(Toolkit):
|
|
|
157
164
|
return ToolResult(content=f"Error: {result['error']}")
|
|
158
165
|
|
|
159
166
|
eta = result.get("eta")
|
|
160
|
-
url_links = result.get("future_links")
|
|
161
167
|
media_id = str(uuid4())
|
|
162
168
|
|
|
163
169
|
# Collect all media artifacts
|
|
@@ -165,17 +171,21 @@ class ModelsLabTools(Toolkit):
|
|
|
165
171
|
all_videos = []
|
|
166
172
|
all_audios = []
|
|
167
173
|
|
|
174
|
+
if self.file_type == FileType.WAV:
|
|
175
|
+
url_links = result.get("output", [])
|
|
176
|
+
else:
|
|
177
|
+
url_links = result.get("future_links")
|
|
168
178
|
for media_url in url_links:
|
|
169
179
|
artifacts = self._create_media_artifacts(media_id, media_url, str(eta))
|
|
170
180
|
all_images.extend(artifacts["images"])
|
|
171
181
|
all_videos.extend(artifacts["videos"])
|
|
172
182
|
all_audios.extend(artifacts["audios"])
|
|
173
183
|
|
|
174
|
-
|
|
175
|
-
|
|
176
|
-
|
|
177
|
-
|
|
178
|
-
|
|
184
|
+
if self.wait_for_completion and isinstance(eta, int):
|
|
185
|
+
if self._wait_for_media(media_id, eta):
|
|
186
|
+
log_info("Media generation completed successfully")
|
|
187
|
+
else:
|
|
188
|
+
logger.warning("Media generation timed out")
|
|
179
189
|
|
|
180
190
|
# Return ToolResult with appropriate media artifacts
|
|
181
191
|
return ToolResult(
|
|
@@ -0,0 +1,151 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import os
|
|
4
|
+
from io import BytesIO
|
|
5
|
+
from typing import Any, List, Optional
|
|
6
|
+
from uuid import uuid4
|
|
7
|
+
|
|
8
|
+
from agno.media import Image
|
|
9
|
+
from agno.tools import Toolkit
|
|
10
|
+
from agno.tools.function import ToolResult
|
|
11
|
+
from agno.utils.log import log_debug, logger
|
|
12
|
+
|
|
13
|
+
try:
|
|
14
|
+
from google import genai
|
|
15
|
+
from google.genai import types
|
|
16
|
+
from PIL import Image as PILImage
|
|
17
|
+
|
|
18
|
+
except ImportError as exc:
|
|
19
|
+
missing = []
|
|
20
|
+
try:
|
|
21
|
+
from google.genai import types
|
|
22
|
+
except ImportError:
|
|
23
|
+
missing.append("google-genai")
|
|
24
|
+
|
|
25
|
+
try:
|
|
26
|
+
from PIL import Image as PILImage
|
|
27
|
+
except ImportError:
|
|
28
|
+
missing.append("Pillow")
|
|
29
|
+
|
|
30
|
+
raise ImportError(
|
|
31
|
+
f"Missing required package(s): {', '.join(missing)}. Install using: pip install {' '.join(missing)}"
|
|
32
|
+
) from exc
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
# Note: Expand this list as new models become supported by the Google Content Generation API.
|
|
36
|
+
ALLOWED_MODELS = ["gemini-2.5-flash-image"]
|
|
37
|
+
ALLOWED_RATIOS = ["1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9", "21:9"]
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
class NanoBananaTools(Toolkit):
|
|
41
|
+
def __init__(
|
|
42
|
+
self,
|
|
43
|
+
model: str = "gemini-2.5-flash-image",
|
|
44
|
+
aspect_ratio: str = "1:1",
|
|
45
|
+
api_key: Optional[str] = None,
|
|
46
|
+
enable_create_image: bool = True,
|
|
47
|
+
**kwargs,
|
|
48
|
+
):
|
|
49
|
+
self.model = model
|
|
50
|
+
self.aspect_ratio = aspect_ratio
|
|
51
|
+
self.api_key = api_key or os.getenv("GOOGLE_API_KEY")
|
|
52
|
+
|
|
53
|
+
# Validate model
|
|
54
|
+
if model not in ALLOWED_MODELS:
|
|
55
|
+
raise ValueError(f"Invalid model '{model}'. Supported: {', '.join(ALLOWED_MODELS)}")
|
|
56
|
+
|
|
57
|
+
if self.aspect_ratio not in ALLOWED_RATIOS:
|
|
58
|
+
raise ValueError(f"Invalid aspect_ratio '{self.aspect_ratio}'. Supported: {', '.join(ALLOWED_RATIOS)}")
|
|
59
|
+
|
|
60
|
+
if not self.api_key:
|
|
61
|
+
raise ValueError("GOOGLE_API_KEY not set. Export it: `export GOOGLE_API_KEY=<your-key>`")
|
|
62
|
+
|
|
63
|
+
tools: List[Any] = []
|
|
64
|
+
if enable_create_image:
|
|
65
|
+
tools.append(self.create_image)
|
|
66
|
+
|
|
67
|
+
super().__init__(name="nano_banana", tools=tools, **kwargs)
|
|
68
|
+
|
|
69
|
+
def create_image(self, prompt: str) -> ToolResult:
|
|
70
|
+
"""Generate an image from a text prompt."""
|
|
71
|
+
try:
|
|
72
|
+
client = genai.Client(api_key=self.api_key)
|
|
73
|
+
log_debug(f"NanoBanana generating image with prompt: {prompt}")
|
|
74
|
+
|
|
75
|
+
cfg = types.GenerateContentConfig(
|
|
76
|
+
response_modalities=["IMAGE"],
|
|
77
|
+
image_config=types.ImageConfig(aspect_ratio=self.aspect_ratio),
|
|
78
|
+
)
|
|
79
|
+
|
|
80
|
+
response = client.models.generate_content(
|
|
81
|
+
model=self.model,
|
|
82
|
+
contents=[prompt], # type: ignore
|
|
83
|
+
config=cfg,
|
|
84
|
+
)
|
|
85
|
+
|
|
86
|
+
generated_images: List[Image] = []
|
|
87
|
+
response_str = ""
|
|
88
|
+
|
|
89
|
+
if not hasattr(response, "candidates") or not response.candidates:
|
|
90
|
+
logger.warning("No candidates in response")
|
|
91
|
+
return ToolResult(content="No images were generated in the response")
|
|
92
|
+
|
|
93
|
+
# Process each candidate
|
|
94
|
+
for candidate in response.candidates:
|
|
95
|
+
if not hasattr(candidate, "content") or not candidate.content or not candidate.content.parts:
|
|
96
|
+
continue
|
|
97
|
+
|
|
98
|
+
for part in candidate.content.parts:
|
|
99
|
+
if hasattr(part, "text") and part.text:
|
|
100
|
+
response_str += part.text + "\n"
|
|
101
|
+
|
|
102
|
+
if hasattr(part, "inline_data") and part.inline_data:
|
|
103
|
+
try:
|
|
104
|
+
# Extract image data from the blob
|
|
105
|
+
image_data = part.inline_data.data
|
|
106
|
+
mime_type = getattr(part.inline_data, "mime_type", "image/png")
|
|
107
|
+
|
|
108
|
+
if image_data:
|
|
109
|
+
pil_img = PILImage.open(BytesIO(image_data))
|
|
110
|
+
|
|
111
|
+
# Save to buffer with proper format
|
|
112
|
+
buffer = BytesIO()
|
|
113
|
+
image_format = "PNG" if "png" in mime_type.lower() else "JPEG"
|
|
114
|
+
pil_img.save(buffer, format=image_format)
|
|
115
|
+
buffer.seek(0)
|
|
116
|
+
|
|
117
|
+
agno_img = Image(
|
|
118
|
+
id=str(uuid4()),
|
|
119
|
+
content=buffer.getvalue(),
|
|
120
|
+
original_prompt=prompt,
|
|
121
|
+
)
|
|
122
|
+
generated_images.append(agno_img)
|
|
123
|
+
|
|
124
|
+
log_debug(f"Successfully processed image with ID: {agno_img.id}")
|
|
125
|
+
response_str += f"Image generated successfully (ID: {agno_img.id}).\n"
|
|
126
|
+
|
|
127
|
+
except Exception as img_exc:
|
|
128
|
+
logger.error(f"Failed to process image data: {img_exc}")
|
|
129
|
+
response_str += f"Failed to process image: {img_exc}\n"
|
|
130
|
+
|
|
131
|
+
if hasattr(response, "usage_metadata") and response.usage_metadata:
|
|
132
|
+
log_debug(
|
|
133
|
+
f"Token usage - Prompt: {response.usage_metadata.prompt_token_count}, "
|
|
134
|
+
f"Response: {response.usage_metadata.candidates_token_count}, "
|
|
135
|
+
f"Total: {response.usage_metadata.total_token_count}"
|
|
136
|
+
)
|
|
137
|
+
|
|
138
|
+
if generated_images:
|
|
139
|
+
return ToolResult(
|
|
140
|
+
content=response_str.strip() or "Image(s) generated successfully",
|
|
141
|
+
images=generated_images,
|
|
142
|
+
)
|
|
143
|
+
else:
|
|
144
|
+
return ToolResult(
|
|
145
|
+
content=response_str.strip() or "No images were generated",
|
|
146
|
+
images=None,
|
|
147
|
+
)
|
|
148
|
+
|
|
149
|
+
except Exception as exc:
|
|
150
|
+
logger.error(f"NanoBanana image generation failed: {exc}")
|
|
151
|
+
return ToolResult(content=f"Error generating image: {str(exc)}")
|