aru-code 0.11.2__tar.gz → 0.12.0__tar.gz
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.
- {aru_code-0.11.2/aru_code.egg-info → aru_code-0.12.0}/PKG-INFO +1 -1
- aru_code-0.12.0/aru/__init__.py +1 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/aru/agent_factory.py +2 -2
- {aru_code-0.11.2 → aru_code-0.12.0}/aru/agents/executor.py +1 -1
- {aru_code-0.11.2 → aru_code-0.12.0}/aru/agents/planner.py +1 -1
- {aru_code-0.11.2 → aru_code-0.12.0}/aru/config.py +12 -2
- {aru_code-0.11.2 → aru_code-0.12.0}/aru/context.py +29 -2
- {aru_code-0.11.2 → aru_code-0.12.0}/aru/providers.py +80 -4
- {aru_code-0.11.2 → aru_code-0.12.0}/aru/runtime.py +3 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/aru/tools/codebase.py +66 -5
- aru_code-0.12.0/aru/tools/mcp_client.py +283 -0
- {aru_code-0.11.2 → aru_code-0.12.0/aru_code.egg-info}/PKG-INFO +1 -1
- {aru_code-0.11.2 → aru_code-0.12.0}/pyproject.toml +1 -1
- {aru_code-0.11.2 → aru_code-0.12.0}/tests/test_config.py +3 -3
- {aru_code-0.11.2 → aru_code-0.12.0}/tests/test_mcp_client.py +21 -22
- aru_code-0.11.2/aru/__init__.py +0 -1
- aru_code-0.11.2/aru/tools/mcp_client.py +0 -158
- {aru_code-0.11.2 → aru_code-0.12.0}/LICENSE +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/README.md +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/aru/agents/__init__.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/aru/agents/base.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/aru/cli.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/aru/commands.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/aru/completers.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/aru/display.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/aru/permissions.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/aru/runner.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/aru/session.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/aru/tools/__init__.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/aru/tools/ast_tools.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/aru/tools/gitignore.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/aru/tools/ranker.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/aru/tools/tasklist.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/aru_code.egg-info/SOURCES.txt +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/aru_code.egg-info/dependency_links.txt +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/aru_code.egg-info/entry_points.txt +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/aru_code.egg-info/requires.txt +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/aru_code.egg-info/top_level.txt +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/setup.cfg +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/tests/test_agents_base.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/tests/test_ast_tools.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/tests/test_cli.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/tests/test_cli_advanced.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/tests/test_cli_base.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/tests/test_cli_completers.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/tests/test_cli_new.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/tests/test_cli_run_cli.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/tests/test_cli_session.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/tests/test_cli_shell.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/tests/test_codebase.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/tests/test_context.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/tests/test_executor.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/tests/test_gitignore.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/tests/test_main.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/tests/test_permissions.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/tests/test_planner.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/tests/test_providers.py +0 -0
- {aru_code-0.11.2 → aru_code-0.12.0}/tests/test_ranker.py +0 -0
|
@@ -0,0 +1 @@
|
|
|
1
|
+
__version__ = "0.12.0"
|
|
@@ -33,7 +33,7 @@ def create_general_agent(
|
|
|
33
33
|
compression_manager=CompressionManager(
|
|
34
34
|
model=create_model(get_ctx().small_model_ref, max_tokens=1024),
|
|
35
35
|
compress_tool_results=True,
|
|
36
|
-
compress_tool_results_limit=
|
|
36
|
+
compress_tool_results_limit=25,
|
|
37
37
|
),
|
|
38
38
|
tool_call_limit=20,
|
|
39
39
|
)
|
|
@@ -67,7 +67,7 @@ def create_custom_agent_instance(agent_def: CustomAgent, session: Session,
|
|
|
67
67
|
compression_manager=CompressionManager(
|
|
68
68
|
model=create_model(get_ctx().small_model_ref, max_tokens=1024),
|
|
69
69
|
compress_tool_results=True,
|
|
70
|
-
compress_tool_results_limit=
|
|
70
|
+
compress_tool_results_limit=25,
|
|
71
71
|
),
|
|
72
72
|
tool_call_limit=agent_def.max_turns or 20,
|
|
73
73
|
)
|
|
@@ -53,7 +53,7 @@ def create_executor(model_ref: str = "anthropic/claude-sonnet-4-5", extra_instru
|
|
|
53
53
|
compression_manager=_SafeCompressionManager(
|
|
54
54
|
model=create_model(get_ctx().small_model_ref, max_tokens=2048),
|
|
55
55
|
compress_tool_results=True,
|
|
56
|
-
compress_tool_results_limit=
|
|
56
|
+
compress_tool_results_limit=25,
|
|
57
57
|
),
|
|
58
58
|
tool_call_limit=None,
|
|
59
59
|
)
|
|
@@ -79,7 +79,7 @@ def create_planner(model_ref: str = "anthropic/claude-sonnet-4-5", extra_instruc
|
|
|
79
79
|
compression_manager=CompressionManager(
|
|
80
80
|
model=create_model(get_ctx().small_model_ref, max_tokens=1024),
|
|
81
81
|
compress_tool_results=True,
|
|
82
|
-
compress_tool_results_limit=
|
|
82
|
+
compress_tool_results_limit=25,
|
|
83
83
|
),
|
|
84
84
|
tool_call_limit=20,
|
|
85
85
|
)
|
|
@@ -173,8 +173,9 @@ class AgentConfig:
|
|
|
173
173
|
lightweight: If True, skip README.md and skill catalog to save tokens.
|
|
174
174
|
"""
|
|
175
175
|
parts = []
|
|
176
|
-
|
|
177
|
-
|
|
176
|
+
# README.md is no longer included by default — it's written for humans
|
|
177
|
+
# (badges, install instructions, contributing guides) and wastes ~2K tokens
|
|
178
|
+
# per turn. AGENTS.md is the proper place for model-facing context.
|
|
178
179
|
if self.agents_md:
|
|
179
180
|
parts.append(f"## Project Instructions (AGENTS.md)\n\n{self.agents_md}")
|
|
180
181
|
if self.rules_instructions:
|
|
@@ -195,6 +196,15 @@ class AgentConfig:
|
|
|
195
196
|
lines.append(f"- `/{name}{hint}`: {skill.description}")
|
|
196
197
|
parts.append("\n".join(lines))
|
|
197
198
|
|
|
199
|
+
# Include MCP tool catalog (lazy mode — lightweight text instead of full schemas)
|
|
200
|
+
try:
|
|
201
|
+
from aru.runtime import get_ctx
|
|
202
|
+
catalog_text = get_ctx().mcp_catalog_text
|
|
203
|
+
if catalog_text and not lightweight:
|
|
204
|
+
parts.append(catalog_text)
|
|
205
|
+
except LookupError:
|
|
206
|
+
pass
|
|
207
|
+
|
|
198
208
|
return "\n\n".join(parts)
|
|
199
209
|
|
|
200
210
|
|
|
@@ -24,18 +24,45 @@ TRUNCATE_KEEP_START = 350 # lines to keep from the start
|
|
|
24
24
|
TRUNCATE_KEEP_END = 100 # lines to keep from the end
|
|
25
25
|
|
|
26
26
|
# Compaction: trigger when cumulative input tokens exceed this fraction of model limit
|
|
27
|
-
COMPACTION_THRESHOLD_RATIO = 0.
|
|
27
|
+
COMPACTION_THRESHOLD_RATIO = 0.85
|
|
28
28
|
# Default model context limits (input tokens)
|
|
29
29
|
MODEL_CONTEXT_LIMITS: dict[str, int] = {
|
|
30
|
+
# Anthropic
|
|
30
31
|
"claude-sonnet-4-5-20250929": 200_000,
|
|
31
32
|
"claude-sonnet-4-20250514": 200_000,
|
|
32
33
|
"claude-haiku-4-5-20251001": 200_000,
|
|
33
34
|
"claude-opus-4-20250514": 200_000,
|
|
34
35
|
"claude-opus-4-6": 1_000_000,
|
|
35
36
|
"claude-sonnet-4-6": 1_000_000,
|
|
37
|
+
# OpenAI
|
|
36
38
|
"gpt-4o": 128_000,
|
|
37
39
|
"gpt-4o-mini": 128_000,
|
|
38
|
-
"
|
|
40
|
+
"gpt-4.1": 1_000_000,
|
|
41
|
+
"gpt-4.1-mini": 1_000_000,
|
|
42
|
+
"gpt-4.1-nano": 1_000_000,
|
|
43
|
+
"o3": 200_000,
|
|
44
|
+
"o3-mini": 200_000,
|
|
45
|
+
"o4-mini": 200_000,
|
|
46
|
+
# Qwen (AlibabaCloud)
|
|
47
|
+
"qwen3-plus": 128_000,
|
|
48
|
+
"qwen3.6-plus": 128_000,
|
|
49
|
+
"qwen-plus": 128_000,
|
|
50
|
+
"qwen-max": 128_000,
|
|
51
|
+
"qwen-turbo": 128_000,
|
|
52
|
+
"qwen3-coder-plus": 128_000,
|
|
53
|
+
# DeepSeek
|
|
54
|
+
"deepseek-chat": 128_000,
|
|
55
|
+
"deepseek-reasoner": 128_000,
|
|
56
|
+
# Meta Llama (common Ollama/Groq)
|
|
57
|
+
"llama3.1": 128_000,
|
|
58
|
+
"llama-3.1-70b-versatile": 128_000,
|
|
59
|
+
"llama-3.3-70b-versatile": 128_000,
|
|
60
|
+
"llama4-scout": 512_000,
|
|
61
|
+
# Google Gemini (OpenRouter)
|
|
62
|
+
"gemini-2.5-pro": 1_000_000,
|
|
63
|
+
"gemini-2.5-flash": 1_000_000,
|
|
64
|
+
# Fallback
|
|
65
|
+
"default": 128_000,
|
|
39
66
|
}
|
|
40
67
|
|
|
41
68
|
COMPACTION_TEMPLATE = """\
|
|
@@ -6,10 +6,13 @@ Maps provider names to Agno model classes and handles provider-specific configur
|
|
|
6
6
|
|
|
7
7
|
from __future__ import annotations
|
|
8
8
|
|
|
9
|
+
import logging
|
|
9
10
|
import os
|
|
10
11
|
from dataclasses import dataclass, field
|
|
11
12
|
from typing import Any
|
|
12
13
|
|
|
14
|
+
logger = logging.getLogger("aru.providers")
|
|
15
|
+
|
|
13
16
|
|
|
14
17
|
# ---------------------------------------------------------------------------
|
|
15
18
|
# Built-in provider definitions
|
|
@@ -150,7 +153,7 @@ def load_providers_from_config(config_data: dict[str, Any]):
|
|
|
150
153
|
"ollama": {
|
|
151
154
|
"base_url": "http://localhost:11434",
|
|
152
155
|
"models": {
|
|
153
|
-
"deepseek-coder-v2": {"id": "deepseek-coder-v2:latest"}
|
|
156
|
+
"deepseek-coder-v2": {"id": "deepseek-coder-v2:latest", "context_limit": 128000}
|
|
154
157
|
}
|
|
155
158
|
},
|
|
156
159
|
"my-custom": {
|
|
@@ -158,18 +161,28 @@ def load_providers_from_config(config_data: dict[str, Any]):
|
|
|
158
161
|
"name": "My Custom Provider",
|
|
159
162
|
"api_key_env": "MY_API_KEY",
|
|
160
163
|
"base_url": "https://my-api.example.com/v1",
|
|
164
|
+
"context_limit": 128000,
|
|
161
165
|
"models": {
|
|
162
|
-
"my-model": {"id": "my-model-v1"}
|
|
166
|
+
"my-model": {"id": "my-model-v1", "context_limit": 64000}
|
|
163
167
|
}
|
|
164
168
|
}
|
|
165
169
|
}
|
|
166
170
|
}
|
|
171
|
+
|
|
172
|
+
context_limit can be set per-model or per-provider (provider-level serves as
|
|
173
|
+
default for all its models). Values are merged into MODEL_CONTEXT_LIMITS so
|
|
174
|
+
compaction triggers at the correct threshold.
|
|
167
175
|
"""
|
|
176
|
+
from aru.context import MODEL_CONTEXT_LIMITS
|
|
177
|
+
|
|
168
178
|
providers_data = config_data.get("providers", {})
|
|
169
179
|
for key, pdata in providers_data.items():
|
|
170
180
|
if not isinstance(pdata, dict):
|
|
171
181
|
continue
|
|
172
182
|
|
|
183
|
+
# Provider-level context_limit (applies to all models as default)
|
|
184
|
+
provider_context_limit = pdata.get("context_limit")
|
|
185
|
+
|
|
173
186
|
# If this extends a built-in, start from that base
|
|
174
187
|
existing = _providers.get(key)
|
|
175
188
|
if existing:
|
|
@@ -200,6 +213,23 @@ def load_providers_from_config(config_data: dict[str, Any]):
|
|
|
200
213
|
if "type" in pdata:
|
|
201
214
|
_providers[key].options["_provider_type"] = pdata["type"]
|
|
202
215
|
|
|
216
|
+
# Register context_limit values into MODEL_CONTEXT_LIMITS
|
|
217
|
+
models_data = pdata.get("models", {})
|
|
218
|
+
for model_name, model_cfg in models_data.items():
|
|
219
|
+
if not isinstance(model_cfg, dict):
|
|
220
|
+
continue
|
|
221
|
+
limit = model_cfg.get("context_limit") or provider_context_limit
|
|
222
|
+
if isinstance(limit, int) and limit > 0:
|
|
223
|
+
model_id = model_cfg.get("id", model_name)
|
|
224
|
+
MODEL_CONTEXT_LIMITS[model_id] = limit
|
|
225
|
+
|
|
226
|
+
# If provider has context_limit but no per-model overrides, register default model
|
|
227
|
+
if isinstance(provider_context_limit, int) and provider_context_limit > 0:
|
|
228
|
+
provider_obj = _providers.get(key)
|
|
229
|
+
if provider_obj and provider_obj.default_model:
|
|
230
|
+
default_id = _get_actual_model_id(provider_obj, provider_obj.default_model)
|
|
231
|
+
MODEL_CONTEXT_LIMITS.setdefault(default_id, provider_context_limit)
|
|
232
|
+
|
|
203
233
|
|
|
204
234
|
# ---------------------------------------------------------------------------
|
|
205
235
|
# Model resolution
|
|
@@ -300,6 +330,38 @@ def create_model(
|
|
|
300
330
|
)
|
|
301
331
|
|
|
302
332
|
|
|
333
|
+
def _make_cached_openai_chat_class():
|
|
334
|
+
"""Create a CachedOpenAIChat subclass (lazy import to avoid top-level dependency)."""
|
|
335
|
+
from agno.models.openai import OpenAIChat
|
|
336
|
+
from agno.models.message import Message
|
|
337
|
+
|
|
338
|
+
class CachedOpenAIChat(OpenAIChat):
|
|
339
|
+
"""OpenAIChat subclass that injects cache_control into system messages.
|
|
340
|
+
|
|
341
|
+
DashScope (Qwen) and other OpenAI-compatible APIs support explicit prompt caching
|
|
342
|
+
via cache_control: {"type": "ephemeral"} on content blocks. This subclass
|
|
343
|
+
automatically adds that marker to system messages so the provider can cache
|
|
344
|
+
the system prompt between turns (up to 90% cost reduction on cached tokens).
|
|
345
|
+
"""
|
|
346
|
+
|
|
347
|
+
def _format_message(self, message: Message, compress_tool_results: bool = False):
|
|
348
|
+
formatted = super()._format_message(message, compress_tool_results)
|
|
349
|
+
|
|
350
|
+
if message.role == "system" and isinstance(formatted.get("content"), str):
|
|
351
|
+
text = formatted["content"]
|
|
352
|
+
formatted["content"] = [
|
|
353
|
+
{
|
|
354
|
+
"type": "text",
|
|
355
|
+
"text": text,
|
|
356
|
+
"cache_control": {"type": "ephemeral"},
|
|
357
|
+
}
|
|
358
|
+
]
|
|
359
|
+
|
|
360
|
+
return formatted
|
|
361
|
+
|
|
362
|
+
return CachedOpenAIChat
|
|
363
|
+
|
|
364
|
+
|
|
303
365
|
def _create_provider_model(
|
|
304
366
|
provider_type: str,
|
|
305
367
|
provider: ProviderConfig,
|
|
@@ -322,7 +384,6 @@ def _create_provider_model(
|
|
|
322
384
|
return Claude(**params)
|
|
323
385
|
|
|
324
386
|
elif provider_type == "openai":
|
|
325
|
-
from agno.models.openai import OpenAIChat
|
|
326
387
|
api_key = _resolve_api_key(provider)
|
|
327
388
|
params = {"id": model_id, "max_tokens": max_tokens}
|
|
328
389
|
if api_key:
|
|
@@ -338,6 +399,10 @@ def _create_provider_model(
|
|
|
338
399
|
"model": "assistant",
|
|
339
400
|
}
|
|
340
401
|
params.update(kwargs)
|
|
402
|
+
if cache_system_prompt:
|
|
403
|
+
CachedOpenAIChat = _make_cached_openai_chat_class()
|
|
404
|
+
return CachedOpenAIChat(**params)
|
|
405
|
+
from agno.models.openai import OpenAIChat
|
|
341
406
|
return OpenAIChat(**params)
|
|
342
407
|
|
|
343
408
|
elif provider_type == "ollama":
|
|
@@ -382,14 +447,25 @@ def _create_provider_model(
|
|
|
382
447
|
|
|
383
448
|
else:
|
|
384
449
|
# Fallback: try OpenAI-compatible (works for many providers)
|
|
385
|
-
from agno.models.openai import OpenAIChat
|
|
386
450
|
api_key = _resolve_api_key(provider)
|
|
387
451
|
params = {"id": model_id, "max_tokens": max_tokens}
|
|
388
452
|
if api_key:
|
|
389
453
|
params["api_key"] = api_key
|
|
390
454
|
if provider.base_url:
|
|
391
455
|
params["base_url"] = provider.base_url
|
|
456
|
+
if provider.options.get("use_system_role"):
|
|
457
|
+
params["role_map"] = {
|
|
458
|
+
"system": "system",
|
|
459
|
+
"user": "user",
|
|
460
|
+
"assistant": "assistant",
|
|
461
|
+
"tool": "tool",
|
|
462
|
+
"model": "assistant",
|
|
463
|
+
}
|
|
392
464
|
params.update(kwargs)
|
|
465
|
+
if cache_system_prompt:
|
|
466
|
+
CachedOpenAIChat = _make_cached_openai_chat_class()
|
|
467
|
+
return CachedOpenAIChat(**params)
|
|
468
|
+
from agno.models.openai import OpenAIChat
|
|
393
469
|
return OpenAIChat(**params)
|
|
394
470
|
|
|
395
471
|
|
|
@@ -1211,16 +1211,77 @@ def _update_delegate_task_docstring():
|
|
|
1211
1211
|
delegate_task.__doc__ = base_doc
|
|
1212
1212
|
|
|
1213
1213
|
|
|
1214
|
-
async def load_mcp_tools():
|
|
1215
|
-
"""Initialize MCP servers and
|
|
1214
|
+
async def load_mcp_tools(eager: bool = False):
|
|
1215
|
+
"""Initialize MCP servers and expose their tools to agents.
|
|
1216
|
+
|
|
1217
|
+
Args:
|
|
1218
|
+
eager: If True, inject each MCP tool as its own Agno Function (legacy mode).
|
|
1219
|
+
If False (default), inject a single gateway tool + lightweight catalog
|
|
1220
|
+
in the system prompt — saves thousands of tokens per turn.
|
|
1221
|
+
"""
|
|
1216
1222
|
from aru.tools.mcp_client import init_mcp
|
|
1217
1223
|
try:
|
|
1218
|
-
|
|
1219
|
-
if
|
|
1220
|
-
|
|
1224
|
+
manager = await init_mcp()
|
|
1225
|
+
if manager is None or not manager.catalog:
|
|
1226
|
+
return
|
|
1227
|
+
|
|
1228
|
+
tool_count = len(manager.catalog)
|
|
1229
|
+
|
|
1230
|
+
if eager:
|
|
1231
|
+
# Legacy: each MCP tool = one Agno Function (expensive)
|
|
1232
|
+
mcp_tools = manager.get_eager_tools()
|
|
1233
|
+
get_ctx().console.print(f"[dim]Loaded {tool_count} tools from MCP servers (eager mode).[/dim]")
|
|
1221
1234
|
for t in mcp_tools:
|
|
1222
1235
|
ALL_TOOLS.append(t)
|
|
1223
1236
|
EXECUTOR_TOOLS.append(t)
|
|
1224
1237
|
GENERAL_TOOLS.append(t)
|
|
1238
|
+
else:
|
|
1239
|
+
# Lazy: single gateway tool + text catalog
|
|
1240
|
+
gateway = _build_mcp_gateway(manager)
|
|
1241
|
+
ALL_TOOLS.append(gateway)
|
|
1242
|
+
EXECUTOR_TOOLS.append(gateway)
|
|
1243
|
+
GENERAL_TOOLS.append(gateway)
|
|
1244
|
+
# Store catalog text for injection into system prompt
|
|
1245
|
+
get_ctx().mcp_catalog_text = manager.get_catalog_text()
|
|
1246
|
+
get_ctx().console.print(f"[dim]Loaded {tool_count} tools from MCP servers.[/dim]")
|
|
1247
|
+
|
|
1225
1248
|
except Exception as e:
|
|
1226
1249
|
get_ctx().console.print(f"[dim]Failed to load MCP tools: {e}[/dim]")
|
|
1250
|
+
|
|
1251
|
+
|
|
1252
|
+
def _build_mcp_gateway(manager):
|
|
1253
|
+
"""Build the single gateway Function that routes to any MCP tool."""
|
|
1254
|
+
from agno.tools import Function
|
|
1255
|
+
|
|
1256
|
+
async def use_mcp_tool(tool_name: str, arguments: dict | None = None) -> str:
|
|
1257
|
+
"""Call an external MCP tool by name.
|
|
1258
|
+
|
|
1259
|
+
Use this to invoke any tool from the MCP Tools catalog listed in your instructions.
|
|
1260
|
+
Pass the exact tool name and its arguments as a JSON object.
|
|
1261
|
+
|
|
1262
|
+
Args:
|
|
1263
|
+
tool_name: The MCP tool name (e.g. "github__search_repositories").
|
|
1264
|
+
arguments: The arguments to pass to the tool as key-value pairs.
|
|
1265
|
+
"""
|
|
1266
|
+
return await manager.call_tool(tool_name, arguments)
|
|
1267
|
+
|
|
1268
|
+
return Function(
|
|
1269
|
+
name="use_mcp_tool",
|
|
1270
|
+
description="Call an external MCP tool by name. See the MCP Tools catalog in your instructions for available tools and their parameters.",
|
|
1271
|
+
parameters={
|
|
1272
|
+
"type": "object",
|
|
1273
|
+
"properties": {
|
|
1274
|
+
"tool_name": {
|
|
1275
|
+
"type": "string",
|
|
1276
|
+
"description": "The MCP tool name from the catalog (e.g. 'github__search_repositories')"
|
|
1277
|
+
},
|
|
1278
|
+
"arguments": {
|
|
1279
|
+
"type": "object",
|
|
1280
|
+
"description": "Arguments to pass to the tool as key-value pairs",
|
|
1281
|
+
"additionalProperties": True
|
|
1282
|
+
}
|
|
1283
|
+
},
|
|
1284
|
+
"required": ["tool_name"]
|
|
1285
|
+
},
|
|
1286
|
+
entrypoint=use_mcp_tool,
|
|
1287
|
+
)
|
|
@@ -0,0 +1,283 @@
|
|
|
1
|
+
"""Model Context Protocol (MCP) client manager and tool generation.
|
|
2
|
+
|
|
3
|
+
Supports two modes for exposing MCP tools to agents:
|
|
4
|
+
- **Eager** (legacy): Each MCP tool becomes its own Agno Function with full JSON Schema.
|
|
5
|
+
Sends all tool schemas in every request — expensive with many tools.
|
|
6
|
+
- **Lazy** (default): A single gateway tool `use_mcp_tool` replaces all individual tools.
|
|
7
|
+
The tool catalog (name + description) is injected as lightweight text in the system prompt.
|
|
8
|
+
Full schema resolution happens only when the model invokes a specific tool.
|
|
9
|
+
"""
|
|
10
|
+
|
|
11
|
+
from __future__ import annotations
|
|
12
|
+
|
|
13
|
+
import asyncio
|
|
14
|
+
import json
|
|
15
|
+
import os
|
|
16
|
+
from contextlib import AsyncExitStack
|
|
17
|
+
from dataclasses import dataclass, field
|
|
18
|
+
|
|
19
|
+
from agno.tools import Function
|
|
20
|
+
from mcp.client.stdio import stdio_client, StdioServerParameters
|
|
21
|
+
from mcp.client.session import ClientSession
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
@dataclass
|
|
25
|
+
class McpToolEntry:
|
|
26
|
+
"""Lightweight catalog entry for a discovered MCP tool."""
|
|
27
|
+
name: str # safe_name: "server__tool_name"
|
|
28
|
+
description: str # "[server] original description"
|
|
29
|
+
parameters: dict # full JSON Schema (only used on invocation)
|
|
30
|
+
server_name: str # originating MCP server
|
|
31
|
+
original_name: str # tool name as the MCP server knows it
|
|
32
|
+
session: ClientSession = field(repr=False)
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
class McpSessionManager:
|
|
36
|
+
"""Manages MCP server subprocesses and active client sessions."""
|
|
37
|
+
|
|
38
|
+
def __init__(self, config_path: str = "arc.mcp.json"):
|
|
39
|
+
self.config_path = config_path
|
|
40
|
+
self._exit_stack = AsyncExitStack()
|
|
41
|
+
self.sessions: dict[str, ClientSession] = {}
|
|
42
|
+
self.catalog: dict[str, McpToolEntry] = {}
|
|
43
|
+
|
|
44
|
+
async def initialize(self):
|
|
45
|
+
"""Read config and spawn all MCP servers concurrently."""
|
|
46
|
+
if not os.path.exists(self.config_path):
|
|
47
|
+
return
|
|
48
|
+
|
|
49
|
+
with open(self.config_path, "r", encoding="utf-8") as f:
|
|
50
|
+
try:
|
|
51
|
+
config = json.load(f)
|
|
52
|
+
except json.JSONDecodeError:
|
|
53
|
+
print(f"[Warning] Failed to parse {self.config_path}")
|
|
54
|
+
return
|
|
55
|
+
|
|
56
|
+
servers = config.get("mcpServers", {})
|
|
57
|
+
tasks = []
|
|
58
|
+
for name, svr_config in servers.items():
|
|
59
|
+
cmd = svr_config.get("command")
|
|
60
|
+
if not cmd:
|
|
61
|
+
continue
|
|
62
|
+
tasks.append(self._start_server(name, svr_config))
|
|
63
|
+
|
|
64
|
+
if tasks:
|
|
65
|
+
await asyncio.gather(*tasks)
|
|
66
|
+
|
|
67
|
+
async def _start_server(self, name: str, svr_config: dict):
|
|
68
|
+
"""Start a single MCP server and register its session."""
|
|
69
|
+
cmd = svr_config.get("command")
|
|
70
|
+
args = svr_config.get("args", [])
|
|
71
|
+
env = svr_config.get("env", None)
|
|
72
|
+
|
|
73
|
+
server_params = StdioServerParameters(
|
|
74
|
+
command=cmd,
|
|
75
|
+
args=args,
|
|
76
|
+
env={**os.environ.copy(), **env} if env else None
|
|
77
|
+
)
|
|
78
|
+
|
|
79
|
+
try:
|
|
80
|
+
read_stream, write_stream = await self._exit_stack.enter_async_context(
|
|
81
|
+
stdio_client(server_params)
|
|
82
|
+
)
|
|
83
|
+
session = await self._exit_stack.enter_async_context(
|
|
84
|
+
ClientSession(read_stream, write_stream)
|
|
85
|
+
)
|
|
86
|
+
|
|
87
|
+
await session.initialize()
|
|
88
|
+
self.sessions[name] = session
|
|
89
|
+
except Exception as e:
|
|
90
|
+
print(f"[Warning] Failed to start MCP server '{name}': {e}")
|
|
91
|
+
|
|
92
|
+
async def discover_tools(self) -> int:
|
|
93
|
+
"""Fetch all tools from connected servers and populate the catalog.
|
|
94
|
+
|
|
95
|
+
Returns the number of tools discovered.
|
|
96
|
+
"""
|
|
97
|
+
async def _fetch(server_name: str, session: ClientSession) -> list[McpToolEntry]:
|
|
98
|
+
try:
|
|
99
|
+
result = await session.list_tools()
|
|
100
|
+
entries = []
|
|
101
|
+
for tool in result.tools:
|
|
102
|
+
safe_name = f"{server_name}__{tool.name}".replace("-", "_")
|
|
103
|
+
entries.append(McpToolEntry(
|
|
104
|
+
name=safe_name,
|
|
105
|
+
description=f"[{server_name}] {tool.description or ''}",
|
|
106
|
+
parameters=tool.inputSchema,
|
|
107
|
+
server_name=server_name,
|
|
108
|
+
original_name=tool.name,
|
|
109
|
+
session=session,
|
|
110
|
+
))
|
|
111
|
+
return entries
|
|
112
|
+
except Exception as e:
|
|
113
|
+
print(f"[Warning] Failed to fetch tools from MCP server '{server_name}': {e}")
|
|
114
|
+
return []
|
|
115
|
+
|
|
116
|
+
results = await asyncio.gather(
|
|
117
|
+
*[_fetch(name, sess) for name, sess in self.sessions.items()]
|
|
118
|
+
)
|
|
119
|
+
for entries in results:
|
|
120
|
+
for entry in entries:
|
|
121
|
+
self.catalog[entry.name] = entry
|
|
122
|
+
|
|
123
|
+
return len(self.catalog)
|
|
124
|
+
|
|
125
|
+
async def call_tool(self, tool_name: str, arguments: dict | None = None) -> str:
|
|
126
|
+
"""Execute an MCP tool by its safe name."""
|
|
127
|
+
entry = self.catalog.get(tool_name)
|
|
128
|
+
if entry is None:
|
|
129
|
+
available = ", ".join(sorted(self.catalog.keys()))
|
|
130
|
+
return f"Error: Unknown MCP tool '{tool_name}'. Available: {available}"
|
|
131
|
+
|
|
132
|
+
try:
|
|
133
|
+
result = await entry.session.call_tool(entry.original_name, arguments=arguments or {})
|
|
134
|
+
output = []
|
|
135
|
+
for content in result.content:
|
|
136
|
+
if hasattr(content, "text"):
|
|
137
|
+
output.append(content.text)
|
|
138
|
+
if result.isError:
|
|
139
|
+
return f"Error from {entry.original_name}: " + "\n".join(output)
|
|
140
|
+
return "\n".join(output)
|
|
141
|
+
except Exception as e:
|
|
142
|
+
return f"Error executing {entry.original_name} on {entry.server_name}: {e}"
|
|
143
|
+
|
|
144
|
+
def get_catalog_text(self) -> str:
|
|
145
|
+
"""Build a lightweight text catalog of available MCP tools.
|
|
146
|
+
|
|
147
|
+
This text is injected into the system prompt so the model knows
|
|
148
|
+
which tools exist — without the cost of full JSON Schema per tool.
|
|
149
|
+
"""
|
|
150
|
+
if not self.catalog:
|
|
151
|
+
return ""
|
|
152
|
+
|
|
153
|
+
lines = ["## MCP Tools (external)\n"]
|
|
154
|
+
lines.append("Call these via `use_mcp_tool(tool_name=\"<name>\", arguments={...})`.\n")
|
|
155
|
+
|
|
156
|
+
# Group by server
|
|
157
|
+
by_server: dict[str, list[McpToolEntry]] = {}
|
|
158
|
+
for entry in self.catalog.values():
|
|
159
|
+
by_server.setdefault(entry.server_name, []).append(entry)
|
|
160
|
+
|
|
161
|
+
for server, entries in sorted(by_server.items()):
|
|
162
|
+
lines.append(f"### {server}")
|
|
163
|
+
for entry in sorted(entries, key=lambda e: e.name):
|
|
164
|
+
desc = entry.description.split("] ", 1)[-1] if "] " in entry.description else entry.description
|
|
165
|
+
# Include parameter names as hints
|
|
166
|
+
props = entry.parameters.get("properties", {})
|
|
167
|
+
if props:
|
|
168
|
+
param_hints = ", ".join(props.keys())
|
|
169
|
+
lines.append(f"- `{entry.name}({param_hints})`: {desc}")
|
|
170
|
+
else:
|
|
171
|
+
lines.append(f"- `{entry.name}()`: {desc}")
|
|
172
|
+
lines.append("")
|
|
173
|
+
|
|
174
|
+
return "\n".join(lines)
|
|
175
|
+
|
|
176
|
+
def get_eager_tools(self) -> list[Function]:
|
|
177
|
+
"""Create individual Agno Functions for each MCP tool (legacy eager mode)."""
|
|
178
|
+
functions = []
|
|
179
|
+
for entry in self.catalog.values():
|
|
180
|
+
async def mcp_caller(*, _entry=entry, **kwargs) -> str:
|
|
181
|
+
return await self.call_tool(_entry.name, kwargs)
|
|
182
|
+
|
|
183
|
+
mcp_caller.__name__ = entry.name
|
|
184
|
+
|
|
185
|
+
functions.append(Function(
|
|
186
|
+
name=entry.name,
|
|
187
|
+
description=entry.description,
|
|
188
|
+
parameters=entry.parameters,
|
|
189
|
+
entrypoint=mcp_caller,
|
|
190
|
+
))
|
|
191
|
+
return functions
|
|
192
|
+
|
|
193
|
+
# -- Backward-compatible API (used by tests and eager mode) --
|
|
194
|
+
|
|
195
|
+
async def get_tools(self) -> list[Function]:
|
|
196
|
+
"""Fetch tools and return as Agno Functions (legacy API).
|
|
197
|
+
|
|
198
|
+
Calls discover_tools() if catalog is empty, then returns eager functions.
|
|
199
|
+
"""
|
|
200
|
+
if not self.catalog:
|
|
201
|
+
await self.discover_tools()
|
|
202
|
+
return self.get_eager_tools()
|
|
203
|
+
|
|
204
|
+
def _create_agno_function(self, server_name: str, session: ClientSession, tool) -> Function:
|
|
205
|
+
"""Create a single Agno Function from an MCP tool (legacy API)."""
|
|
206
|
+
safe_name = f"{server_name}__{tool.name}".replace("-", "_")
|
|
207
|
+
description = f"[{server_name}] {tool.description or ''}"
|
|
208
|
+
original_name = tool.name
|
|
209
|
+
|
|
210
|
+
async def mcp_caller(**kwargs) -> str:
|
|
211
|
+
try:
|
|
212
|
+
result = await session.call_tool(original_name, arguments=kwargs)
|
|
213
|
+
output = []
|
|
214
|
+
for content in result.content:
|
|
215
|
+
if hasattr(content, "text"):
|
|
216
|
+
output.append(content.text)
|
|
217
|
+
if result.isError:
|
|
218
|
+
return f"Error from {original_name}: " + "\n".join(output)
|
|
219
|
+
return "\n".join(output)
|
|
220
|
+
except Exception as e:
|
|
221
|
+
return f"Error executing {original_name} on {server_name}: {e}"
|
|
222
|
+
|
|
223
|
+
mcp_caller.__name__ = safe_name
|
|
224
|
+
|
|
225
|
+
return Function(
|
|
226
|
+
name=safe_name,
|
|
227
|
+
description=description,
|
|
228
|
+
parameters=tool.inputSchema,
|
|
229
|
+
entrypoint=mcp_caller,
|
|
230
|
+
)
|
|
231
|
+
|
|
232
|
+
async def cleanup(self):
|
|
233
|
+
"""Close all active MCP client sessions and terminate server subprocesses."""
|
|
234
|
+
try:
|
|
235
|
+
await self._exit_stack.aclose()
|
|
236
|
+
except (RuntimeError, Exception):
|
|
237
|
+
pass
|
|
238
|
+
|
|
239
|
+
|
|
240
|
+
# Global Singleton manager to be used entirely inside aru's async loops
|
|
241
|
+
_manager: McpSessionManager | None = None
|
|
242
|
+
|
|
243
|
+
|
|
244
|
+
async def init_mcp() -> McpSessionManager | None:
|
|
245
|
+
"""Initialize MCP servers, discover tools, and return the manager.
|
|
246
|
+
|
|
247
|
+
Returns None if no MCP config is found.
|
|
248
|
+
"""
|
|
249
|
+
global _manager
|
|
250
|
+
if _manager is None:
|
|
251
|
+
config_path = None
|
|
252
|
+
for path in [
|
|
253
|
+
".aru/mcp_servers.json",
|
|
254
|
+
"aru.mcp.json",
|
|
255
|
+
".mcp.json",
|
|
256
|
+
"mcp.json"
|
|
257
|
+
]:
|
|
258
|
+
if os.path.exists(path):
|
|
259
|
+
config_path = path
|
|
260
|
+
break
|
|
261
|
+
|
|
262
|
+
if config_path:
|
|
263
|
+
_manager = McpSessionManager(config_path=config_path)
|
|
264
|
+
await _manager.initialize()
|
|
265
|
+
await _manager.discover_tools()
|
|
266
|
+
else:
|
|
267
|
+
_manager = McpSessionManager(config_path="")
|
|
268
|
+
return None
|
|
269
|
+
|
|
270
|
+
return _manager
|
|
271
|
+
|
|
272
|
+
|
|
273
|
+
def get_mcp_manager() -> McpSessionManager | None:
|
|
274
|
+
"""Return the global MCP manager (None if not initialized)."""
|
|
275
|
+
return _manager
|
|
276
|
+
|
|
277
|
+
|
|
278
|
+
async def cleanup_mcp():
|
|
279
|
+
"""Cleanup global manager."""
|
|
280
|
+
global _manager
|
|
281
|
+
if _manager:
|
|
282
|
+
await _manager.cleanup()
|
|
283
|
+
_manager = None
|
|
@@ -109,11 +109,11 @@ class TestDataClasses:
|
|
|
109
109
|
result = config.get_extra_instructions()
|
|
110
110
|
assert result == ""
|
|
111
111
|
|
|
112
|
-
def
|
|
112
|
+
def test_agent_config_get_extra_instructions_readme_not_included(self):
|
|
113
|
+
"""README.md is no longer included in the system prompt by default."""
|
|
113
114
|
config = AgentConfig(readme_md="# My Project\n\nDescription")
|
|
114
115
|
result = config.get_extra_instructions()
|
|
115
|
-
assert "
|
|
116
|
-
assert "# My Project" in result
|
|
116
|
+
assert "README" not in result
|
|
117
117
|
|
|
118
118
|
def test_agent_config_get_extra_instructions_agents_md_only(self):
|
|
119
119
|
config = AgentConfig(agents_md="Follow these rules")
|
|
@@ -996,15 +996,14 @@ class TestGlobalFunctions:
|
|
|
996
996
|
"""Test init_mcp when no config file exists."""
|
|
997
997
|
# Change to temp directory with no config files
|
|
998
998
|
monkeypatch.chdir(tmp_path)
|
|
999
|
-
|
|
999
|
+
|
|
1000
1000
|
# Reset global manager
|
|
1001
1001
|
import aru.tools.mcp_client as mcp_module
|
|
1002
1002
|
mcp_module._manager = None
|
|
1003
|
-
|
|
1004
|
-
|
|
1005
|
-
|
|
1006
|
-
assert
|
|
1007
|
-
assert mcp_module._manager is not None
|
|
1003
|
+
|
|
1004
|
+
result = await init_mcp()
|
|
1005
|
+
|
|
1006
|
+
assert result is None
|
|
1008
1007
|
|
|
1009
1008
|
@pytest.mark.asyncio
|
|
1010
1009
|
async def test_init_mcp_with_config(self, tmp_path, monkeypatch):
|
|
@@ -1021,12 +1020,12 @@ class TestGlobalFunctions:
|
|
|
1021
1020
|
mcp_module._manager = None
|
|
1022
1021
|
|
|
1023
1022
|
with patch.object(McpSessionManager, 'initialize', new_callable=AsyncMock) as mock_init, \
|
|
1024
|
-
patch.object(McpSessionManager, '
|
|
1025
|
-
|
|
1026
|
-
|
|
1027
|
-
|
|
1023
|
+
patch.object(McpSessionManager, 'discover_tools', new_callable=AsyncMock, return_value=0):
|
|
1024
|
+
|
|
1025
|
+
result = await init_mcp()
|
|
1026
|
+
|
|
1028
1027
|
mock_init.assert_called_once()
|
|
1029
|
-
assert isinstance(
|
|
1028
|
+
assert isinstance(result, McpSessionManager)
|
|
1030
1029
|
|
|
1031
1030
|
@pytest.mark.asyncio
|
|
1032
1031
|
async def test_init_mcp_config_priority(self, tmp_path, monkeypatch):
|
|
@@ -1043,10 +1042,10 @@ class TestGlobalFunctions:
|
|
|
1043
1042
|
mcp_module._manager = None
|
|
1044
1043
|
|
|
1045
1044
|
with patch.object(McpSessionManager, 'initialize', new_callable=AsyncMock), \
|
|
1046
|
-
patch.object(McpSessionManager, '
|
|
1047
|
-
|
|
1045
|
+
patch.object(McpSessionManager, 'discover_tools', new_callable=AsyncMock, return_value=0):
|
|
1046
|
+
|
|
1048
1047
|
await init_mcp()
|
|
1049
|
-
|
|
1048
|
+
|
|
1050
1049
|
# Should use .aru/mcp_servers.json (first in priority)
|
|
1051
1050
|
assert mcp_module._manager.config_path == ".aru/mcp_servers.json"
|
|
1052
1051
|
|
|
@@ -1060,16 +1059,16 @@ class TestGlobalFunctions:
|
|
|
1060
1059
|
mcp_module._manager = None
|
|
1061
1060
|
|
|
1062
1061
|
with patch.object(McpSessionManager, 'initialize', new_callable=AsyncMock), \
|
|
1063
|
-
patch.object(McpSessionManager, '
|
|
1064
|
-
|
|
1062
|
+
patch.object(McpSessionManager, 'discover_tools', new_callable=AsyncMock, return_value=0):
|
|
1063
|
+
|
|
1065
1064
|
# First call
|
|
1066
1065
|
await init_mcp()
|
|
1067
1066
|
first_manager = mcp_module._manager
|
|
1068
|
-
|
|
1067
|
+
|
|
1069
1068
|
# Second call
|
|
1070
1069
|
await init_mcp()
|
|
1071
1070
|
second_manager = mcp_module._manager
|
|
1072
|
-
|
|
1071
|
+
|
|
1073
1072
|
# Should be the same instance
|
|
1074
1073
|
assert first_manager is second_manager
|
|
1075
1074
|
|
|
@@ -1122,12 +1121,12 @@ class TestGlobalFunctions:
|
|
|
1122
1121
|
mcp_module._manager = None
|
|
1123
1122
|
|
|
1124
1123
|
with patch.object(McpSessionManager, '_start_server', new_callable=AsyncMock), \
|
|
1125
|
-
patch.object(McpSessionManager, '
|
|
1126
|
-
|
|
1124
|
+
patch.object(McpSessionManager, 'discover_tools', new_callable=AsyncMock, return_value=0):
|
|
1125
|
+
|
|
1127
1126
|
# Initialize
|
|
1128
|
-
|
|
1127
|
+
result = await init_mcp()
|
|
1129
1128
|
assert mcp_module._manager is not None
|
|
1130
|
-
|
|
1129
|
+
|
|
1131
1130
|
# Cleanup
|
|
1132
1131
|
await cleanup_mcp()
|
|
1133
1132
|
assert mcp_module._manager is None
|
aru_code-0.11.2/aru/__init__.py
DELETED
|
@@ -1 +0,0 @@
|
|
|
1
|
-
__version__ = "0.11.2"
|
|
@@ -1,158 +0,0 @@
|
|
|
1
|
-
"""Model Context Protocol (MCP) client manager and tool generation."""
|
|
2
|
-
|
|
3
|
-
from __future__ import annotations
|
|
4
|
-
|
|
5
|
-
import asyncio
|
|
6
|
-
import json
|
|
7
|
-
import os
|
|
8
|
-
from contextlib import AsyncExitStack
|
|
9
|
-
|
|
10
|
-
from agno.tools import Function
|
|
11
|
-
from mcp.client.stdio import stdio_client, StdioServerParameters
|
|
12
|
-
from mcp.client.session import ClientSession
|
|
13
|
-
|
|
14
|
-
|
|
15
|
-
class McpSessionManager:
|
|
16
|
-
"""Manages MCP server subprocesses and active client sessions."""
|
|
17
|
-
|
|
18
|
-
def __init__(self, config_path: str = "arc.mcp.json"):
|
|
19
|
-
self.config_path = config_path
|
|
20
|
-
self._exit_stack = AsyncExitStack()
|
|
21
|
-
self.sessions: dict[str, ClientSession] = {}
|
|
22
|
-
|
|
23
|
-
async def initialize(self):
|
|
24
|
-
"""Read config and spawn all MCP servers concurrently."""
|
|
25
|
-
if not os.path.exists(self.config_path):
|
|
26
|
-
return
|
|
27
|
-
|
|
28
|
-
with open(self.config_path, "r", encoding="utf-8") as f:
|
|
29
|
-
try:
|
|
30
|
-
config = json.load(f)
|
|
31
|
-
except json.JSONDecodeError:
|
|
32
|
-
print(f"[Warning] Failed to parse {self.config_path}")
|
|
33
|
-
return
|
|
34
|
-
|
|
35
|
-
servers = config.get("mcpServers", {})
|
|
36
|
-
tasks = []
|
|
37
|
-
for name, svr_config in servers.items():
|
|
38
|
-
cmd = svr_config.get("command")
|
|
39
|
-
if not cmd:
|
|
40
|
-
continue
|
|
41
|
-
tasks.append(self._start_server(name, svr_config))
|
|
42
|
-
|
|
43
|
-
if tasks:
|
|
44
|
-
await asyncio.gather(*tasks)
|
|
45
|
-
|
|
46
|
-
async def _start_server(self, name: str, svr_config: dict):
|
|
47
|
-
"""Start a single MCP server and register its session."""
|
|
48
|
-
cmd = svr_config.get("command")
|
|
49
|
-
args = svr_config.get("args", [])
|
|
50
|
-
env = svr_config.get("env", None)
|
|
51
|
-
|
|
52
|
-
server_params = StdioServerParameters(
|
|
53
|
-
command=cmd,
|
|
54
|
-
args=args,
|
|
55
|
-
env={**os.environ.copy(), **env} if env else None
|
|
56
|
-
)
|
|
57
|
-
|
|
58
|
-
try:
|
|
59
|
-
read_stream, write_stream = await self._exit_stack.enter_async_context(
|
|
60
|
-
stdio_client(server_params)
|
|
61
|
-
)
|
|
62
|
-
session = await self._exit_stack.enter_async_context(
|
|
63
|
-
ClientSession(read_stream, write_stream)
|
|
64
|
-
)
|
|
65
|
-
|
|
66
|
-
await session.initialize()
|
|
67
|
-
self.sessions[name] = session
|
|
68
|
-
except Exception as e:
|
|
69
|
-
print(f"[Warning] Failed to start MCP server '{name}': {e}")
|
|
70
|
-
|
|
71
|
-
async def get_tools(self) -> list[Function]:
|
|
72
|
-
"""Fetch all tools from connected servers concurrently and convert to Agno Functions."""
|
|
73
|
-
|
|
74
|
-
async def _fetch(server_name: str, session: ClientSession) -> list[Function]:
|
|
75
|
-
try:
|
|
76
|
-
result = await session.list_tools()
|
|
77
|
-
return [self._create_agno_function(server_name, session, tool) for tool in result.tools]
|
|
78
|
-
except Exception as e:
|
|
79
|
-
print(f"[Warning] Failed to fetch tools from MCP server '{server_name}': {e}")
|
|
80
|
-
return []
|
|
81
|
-
|
|
82
|
-
results = await asyncio.gather(
|
|
83
|
-
*[_fetch(name, sess) for name, sess in self.sessions.items()]
|
|
84
|
-
)
|
|
85
|
-
return [tool for tools in results for tool in tools]
|
|
86
|
-
|
|
87
|
-
def _create_agno_function(self, server_name: str, session: ClientSession, tool) -> Function:
|
|
88
|
-
"""Dynamically create an Agno Function that routes to the remote MCP tool."""
|
|
89
|
-
|
|
90
|
-
# We need to capture 'session' and 'tool.name' cleanly.
|
|
91
|
-
# Python's default arguments trick captures loop variables.
|
|
92
|
-
async def mcp_caller(**kwargs) -> str:
|
|
93
|
-
try:
|
|
94
|
-
result = await session.call_tool(tool.name, arguments=kwargs)
|
|
95
|
-
# Parse MCP ToolResultContent
|
|
96
|
-
output = []
|
|
97
|
-
for content in result.content:
|
|
98
|
-
if hasattr(content, "text"):
|
|
99
|
-
output.append(content.text)
|
|
100
|
-
if result.isError:
|
|
101
|
-
return f"Error from {tool.name}: " + "\n".join(output)
|
|
102
|
-
return "\n".join(output)
|
|
103
|
-
except Exception as e:
|
|
104
|
-
return f"Error executing {tool.name} on {server_name}: {e}"
|
|
105
|
-
|
|
106
|
-
# Assign __name__ to the callable for Agno's internal representation
|
|
107
|
-
safe_name = f"{server_name}__{tool.name}".replace("-", "_")
|
|
108
|
-
mcp_caller.__name__ = safe_name
|
|
109
|
-
|
|
110
|
-
return Function(
|
|
111
|
-
name=safe_name,
|
|
112
|
-
description=f"[{server_name}] {tool.description or ''}",
|
|
113
|
-
parameters=tool.inputSchema,
|
|
114
|
-
entrypoint=mcp_caller
|
|
115
|
-
)
|
|
116
|
-
|
|
117
|
-
async def cleanup(self):
|
|
118
|
-
"""Close all active MCP client sessions and terminate server subprocesses."""
|
|
119
|
-
try:
|
|
120
|
-
await self._exit_stack.aclose()
|
|
121
|
-
except (RuntimeError, Exception):
|
|
122
|
-
pass
|
|
123
|
-
|
|
124
|
-
|
|
125
|
-
# Global Singleton manager to be used entirely inside aru's async loops
|
|
126
|
-
_manager: McpSessionManager | None = None
|
|
127
|
-
|
|
128
|
-
async def init_mcp() -> list[Function]:
|
|
129
|
-
"""Initialize MCP servers and return the loaded Agno functions."""
|
|
130
|
-
global _manager
|
|
131
|
-
if _manager is None:
|
|
132
|
-
config_path = None
|
|
133
|
-
for path in [
|
|
134
|
-
".aru/mcp_servers.json",
|
|
135
|
-
"aru.mcp.json",
|
|
136
|
-
".mcp.json",
|
|
137
|
-
"mcp.json"
|
|
138
|
-
]:
|
|
139
|
-
if os.path.exists(path):
|
|
140
|
-
config_path = path
|
|
141
|
-
break
|
|
142
|
-
|
|
143
|
-
if config_path:
|
|
144
|
-
_manager = McpSessionManager(config_path=config_path)
|
|
145
|
-
await _manager.initialize()
|
|
146
|
-
else:
|
|
147
|
-
# Create an empty manager so cleanup doesn't fail, but return no tools
|
|
148
|
-
_manager = McpSessionManager(config_path="")
|
|
149
|
-
return []
|
|
150
|
-
|
|
151
|
-
return await _manager.get_tools()
|
|
152
|
-
|
|
153
|
-
async def cleanup_mcp():
|
|
154
|
-
"""Cleanup global manager."""
|
|
155
|
-
global _manager
|
|
156
|
-
if _manager:
|
|
157
|
-
await _manager.cleanup()
|
|
158
|
-
_manager = None
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|