├── tests
├── __init__.py
├── docs
│ ├── __init__.py
│ └── test_all.py
├── unit
│ ├── __init__.py
│ ├── tools
│ │ ├── __init__.py
│ │ ├── mcp
│ │ │ ├── __init__.py
│ │ │ ├── test_mcp_config.py
│ │ │ └── test_mcp_client.py
│ │ ├── conftest.py
│ │ ├── test_logging.py
│ │ └── test_unit_web_browsing.py
│ ├── frameworks
│ │ ├── __init__.py
│ │ ├── test_agno.py
│ │ ├── test_llama_index.py
│ │ ├── test_google.py
│ │ └── test_smolagents.py
│ ├── conftest.py
│ ├── callbacks
│ │ ├── wrappers
│ │ │ └── test_get_wrapper_and_unwrap.py
│ │ ├── test_console_print_span.py
│ │ └── span_generation
│ │ │ └── test_base_span_generation.py
│ ├── tracing
│ │ └── test_agent_trace.py
│ └── serving
│ │ └── test_envelope_creation.py
├── integration
│ ├── __init__.py
│ ├── a2a
│ │ ├── __init__.py
│ │ ├── test_a2a_serve.py
│ │ └── test_a2a_tool.py
│ ├── mcp
│ │ ├── __init__.py
│ │ └── test_mcp_streamable_http.py
│ ├── tools
│ │ ├── test_composio.py
│ │ └── test_wrap_tools.py
│ ├── frameworks
│ │ ├── test_evaluation.py
│ │ ├── test_thread_safe.py
│ │ └── test_error_handling.py
│ └── conftest.py
├── snapshots
│ └── test_trace.py
└── cookbooks
│ └── test_cookbooks.py
├── src
└── any_agent
│ ├── py.typed
│ ├── testing
│ └── __init__.py
│ ├── frameworks
│ └── __init__.py
│ ├── serving
│ ├── a2a
│ │ ├── __init__.py
│ │ ├── agent_card.py
│ │ ├── envelope.py
│ │ └── server_a2a.py
│ ├── mcp
│ │ ├── __init__.py
│ │ ├── config_mcp.py
│ │ └── server_mcp.py
│ ├── __init__.py
│ └── server_handle.py
│ ├── utils
│ └── __init__.py
│ ├── evaluation
│ ├── __init__.py
│ ├── schemas.py
│ ├── tools.py
│ └── agent_judge.py
│ ├── tools
│ ├── mcp
│ │ └── __init__.py
│ ├── __init__.py
│ ├── user_interaction.py
│ ├── final_output.py
│ ├── composio.py
│ └── web_browsing.py
│ ├── tracing
│ ├── __init__.py
│ └── attributes.py
│ ├── __init__.py
│ ├── callbacks
│ ├── __init__.py
│ ├── span_end.py
│ ├── context.py
│ ├── base.py
│ ├── wrappers
│ │ ├── __init__.py
│ │ ├── agno.py
│ │ ├── smolagents.py
│ │ ├── tinyagent.py
│ │ ├── google.py
│ │ └── llama_index.py
│ ├── span_generation
│ │ ├── __init__.py
│ │ ├── smolagents.py
│ │ ├── openai.py
│ │ ├── tinyagent.py
│ │ ├── agno.py
│ │ └── google.py
│ └── span_print.py
│ └── logging.py
├── demo
├── components
│ ├── __init__.py
│ ├── sidebar.py
│ └── agent_status.py
├── services
│ └── __init__.py
├── requirements.txt
├── README.md
├── restart_space.py
├── tools
│ ├── __init__.py
│ ├── openstreetmap.py
│ └── openmeteo.py
├── Dockerfile
├── config.py
└── constants.py
├── docs
├── api
│ ├── tools.md
│ ├── serving.md
│ ├── agent.md
│ ├── evaluation.md
│ ├── callbacks.md
│ ├── tracing.md
│ ├── config.md
│ └── logging.md
├── images
│ ├── serve_a2a.png
│ ├── any-agent_favicon.png
│ └── any-agent-logo-mark.png
├── assets
│ ├── custom.js
│ └── custom.css
├── agents
│ ├── frameworks
│ │ ├── agno.md
│ │ ├── google_adk.md
│ │ ├── openai.md
│ │ ├── smolagents.md
│ │ ├── llama_index.md
│ │ ├── langchain.md
│ │ └── tinyagent.md
│ └── models.md
└── index.md
├── .gitattributes
├── .codespellrc
├── .github
├── dependabot.yml
└── workflows
│ ├── release.yaml
│ ├── pre-commit-update.yaml
│ ├── lint.yaml
│ ├── tests-unit.yaml
│ ├── tests-docs.yaml
│ ├── publish_demo_to_hf_spaces.yaml
│ ├── docs.yaml
│ ├── tests-cookbook.yaml
│ └── tests-integration.yaml
├── scripts
├── run-mypy.sh
└── wake_up_hf_endpoint.py
├── .pre-commit-config.yaml
├── mkdocs.yml
├── CODE_OF_CONDUCT.md
└── .gitignore
/tests/__init__.py:
--------------------------------------------------------------------------------
1 |
--------------------------------------------------------------------------------
/src/any_agent/py.typed:
--------------------------------------------------------------------------------
1 |
--------------------------------------------------------------------------------
/tests/docs/__init__.py:
--------------------------------------------------------------------------------
1 |
--------------------------------------------------------------------------------
/tests/unit/__init__.py:
--------------------------------------------------------------------------------
1 |
--------------------------------------------------------------------------------
/demo/components/__init__.py:
--------------------------------------------------------------------------------
1 |
--------------------------------------------------------------------------------
/demo/services/__init__.py:
--------------------------------------------------------------------------------
1 |
--------------------------------------------------------------------------------
/tests/unit/tools/__init__.py:
--------------------------------------------------------------------------------
1 |
--------------------------------------------------------------------------------
/src/any_agent/testing/__init__.py:
--------------------------------------------------------------------------------
1 |
--------------------------------------------------------------------------------
/tests/integration/__init__.py:
--------------------------------------------------------------------------------
1 |
--------------------------------------------------------------------------------
/tests/integration/a2a/__init__.py:
--------------------------------------------------------------------------------
1 |
--------------------------------------------------------------------------------
/tests/integration/mcp/__init__.py:
--------------------------------------------------------------------------------
1 |
--------------------------------------------------------------------------------
/tests/unit/frameworks/__init__.py:
--------------------------------------------------------------------------------
1 |
--------------------------------------------------------------------------------
/tests/unit/tools/mcp/__init__.py:
--------------------------------------------------------------------------------
1 |
--------------------------------------------------------------------------------
/src/any_agent/frameworks/__init__.py:
--------------------------------------------------------------------------------
1 |
--------------------------------------------------------------------------------
/src/any_agent/serving/a2a/__init__.py:
--------------------------------------------------------------------------------
1 |
--------------------------------------------------------------------------------
/src/any_agent/serving/mcp/__init__.py:
--------------------------------------------------------------------------------
1 |
--------------------------------------------------------------------------------
/docs/api/tools.md:
--------------------------------------------------------------------------------
1 | # Tools
2 |
3 | ::: any_agent.tools
4 |
--------------------------------------------------------------------------------
/docs/api/serving.md:
--------------------------------------------------------------------------------
1 | # Serving
2 |
3 | ::: any_agent.serving.ServerHandle
4 |
--------------------------------------------------------------------------------
/src/any_agent/utils/__init__.py:
--------------------------------------------------------------------------------
1 | """Utility functions for any-agent."""
2 |
--------------------------------------------------------------------------------
/.gitattributes:
--------------------------------------------------------------------------------
1 | tests/assets/* linguist-vendored
2 | docs/* linguist-vendored
3 |
--------------------------------------------------------------------------------
/demo/requirements.txt:
--------------------------------------------------------------------------------
1 | any-agent[all]>=1.4.0
2 | geocoder
3 | nest_asyncio
4 | streamlit
5 |
--------------------------------------------------------------------------------
/docs/api/agent.md:
--------------------------------------------------------------------------------
1 | ## Agent
2 |
3 | ::: any_agent.AnyAgent
4 |
5 | ::: any_agent.AgentRunError
6 |
--------------------------------------------------------------------------------
/docs/images/serve_a2a.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/mozilla-ai/any-agent/HEAD/docs/images/serve_a2a.png
--------------------------------------------------------------------------------
/docs/images/any-agent_favicon.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/mozilla-ai/any-agent/HEAD/docs/images/any-agent_favicon.png
--------------------------------------------------------------------------------
/docs/api/evaluation.md:
--------------------------------------------------------------------------------
1 | # Evaluation
2 |
3 | ::: any_agent.evaluation.LlmJudge
4 |
5 | ::: any_agent.evaluation.AgentJudge
6 |
--------------------------------------------------------------------------------
/docs/images/any-agent-logo-mark.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/mozilla-ai/any-agent/HEAD/docs/images/any-agent-logo-mark.png
--------------------------------------------------------------------------------
/.codespellrc:
--------------------------------------------------------------------------------
1 | [codespell]
2 | skip = *.git,*.pdf,*.svg
3 | ignore-words-list = assertIn, Pont, checkin, Espace, nieve, Ane, Meu
4 |
--------------------------------------------------------------------------------
/src/any_agent/evaluation/__init__.py:
--------------------------------------------------------------------------------
1 | from .agent_judge import AgentJudge
2 | from .llm_judge import LlmJudge
3 |
4 | __all__ = ["AgentJudge", "LlmJudge"]
5 |
--------------------------------------------------------------------------------
/src/any_agent/tools/mcp/__init__.py:
--------------------------------------------------------------------------------
1 | from .mcp_client import MCPClient
2 | from .smolagents_client import SmolagentsMCPClient
3 |
4 | __all__ = [
5 | "MCPClient",
6 | "SmolagentsMCPClient",
7 | ]
8 |
--------------------------------------------------------------------------------
/docs/api/callbacks.md:
--------------------------------------------------------------------------------
1 | # Callbacks
2 |
3 | ::: any_agent.callbacks.base.Callback
4 |
5 | ::: any_agent.callbacks.context.Context
6 |
7 | ::: any_agent.callbacks.span_print.ConsolePrintSpan
8 |
9 | ::: any_agent.callbacks.get_default_callbacks
10 |
--------------------------------------------------------------------------------
/demo/README.md:
--------------------------------------------------------------------------------
1 | ---
2 | title: Surf Spot Finder
3 | emoji: 🏄🏼♂️
4 | colorFrom: blue
5 | colorTo: indigo
6 | sdk: docker
7 | app_port: 8501
8 | tags:
9 | - streamlit
10 | pinned: false
11 | short_description: Find a surf spot near you
12 | license: apache-2.0
13 | ---
14 |
--------------------------------------------------------------------------------
/src/any_agent/evaluation/schemas.py:
--------------------------------------------------------------------------------
1 | from pydantic import BaseModel
2 |
3 |
4 | class EvaluationOutput(BaseModel):
5 | passed: bool
6 | """Whether the evaluation passed or failed."""
7 |
8 | reasoning: str
9 | """The reasoning for the evaluation."""
10 |
--------------------------------------------------------------------------------
/demo/restart_space.py:
--------------------------------------------------------------------------------
1 | import os
2 |
3 | from huggingface_hub import HfApi
4 |
5 | if __name__ == "__main__":
6 | api = HfApi()
7 | api.restart_space(
8 | repo_id="mozilla-ai/any-agent-demo",
9 | token=os.getenv("HF_TOKEN"),
10 | factory_reboot=True,
11 | )
12 |
--------------------------------------------------------------------------------
/docs/api/tracing.md:
--------------------------------------------------------------------------------
1 | # Tracing
2 |
3 | ::: any_agent.tracing.agent_trace.AgentTrace
4 |
5 | ::: any_agent.tracing.agent_trace.AgentSpan
6 |
7 | ::: any_agent.tracing.agent_trace.CostInfo
8 |
9 | ::: any_agent.tracing.agent_trace.TokenInfo
10 |
11 | ::: any_agent.tracing.attributes.GenAI
12 |
--------------------------------------------------------------------------------
/demo/tools/__init__.py:
--------------------------------------------------------------------------------
1 | from .openmeteo import get_wave_forecast, get_wind_forecast
2 | from .openstreetmap import driving_hours_to_meters, get_area_lat_lon
3 |
4 | __all__ = [
5 | "driving_hours_to_meters",
6 | "get_area_lat_lon",
7 | "get_wave_forecast",
8 | "get_wind_forecast",
9 | ]
10 |
--------------------------------------------------------------------------------
/demo/components/sidebar.py:
--------------------------------------------------------------------------------
1 | import streamlit as st
2 |
3 | from components.inputs import UserInputs, get_user_inputs
4 |
5 |
6 | def ssf_sidebar() -> UserInputs:
7 | st.markdown("### Configuration")
8 | st.markdown("Built using [Any-Agent](https://github.com/mozilla-ai/any-agent)")
9 | return get_user_inputs()
10 |
--------------------------------------------------------------------------------
/docs/assets/custom.js:
--------------------------------------------------------------------------------
1 | // Detect if we're on the tracing page and add a CSS class
2 | document.addEventListener('DOMContentLoaded', function() {
3 | // Check if the current URL contains 'tracing'
4 | if (window.location.pathname.includes('/tracing/')) {
5 | document.body.classList.add('tracing-page');
6 | }
7 | });
8 |
--------------------------------------------------------------------------------
/docs/api/config.md:
--------------------------------------------------------------------------------
1 | # Config
2 |
3 | ::: any_agent.config.AgentConfig
4 |
5 | ::: any_agent.config.MCPStdio
6 |
7 | ::: any_agent.config.MCPStreamableHttp
8 |
9 | ::: any_agent.config.MCPSse
10 |
11 | ::: any_agent.serving.A2AServingConfig
12 |
13 | ::: any_agent.serving.MCPServingConfig
14 |
15 | ::: any_agent.config.AgentFramework
16 |
--------------------------------------------------------------------------------
/.github/dependabot.yml:
--------------------------------------------------------------------------------
1 | version: 2
2 | updates:
3 | - directory: "/"
4 | package-ecosystem: "pip"
5 | schedule:
6 | interval: "weekly"
7 | labels:
8 | - "maintenance"
9 |
10 | - directory: "/"
11 | package-ecosystem: "github-actions"
12 | schedule:
13 | interval: "weekly"
14 | labels:
15 | - "maintenance"
16 |
--------------------------------------------------------------------------------
/src/any_agent/tracing/__init__.py:
--------------------------------------------------------------------------------
1 | # mypy: disable-error-code="attr-defined"
2 | from opentelemetry import trace
3 | from opentelemetry.sdk.trace import TracerProvider
4 |
5 | TRACE_PROVIDER = trace.get_tracer_provider()
6 | if isinstance(TRACE_PROVIDER, trace.ProxyTracerProvider):
7 | TRACE_PROVIDER = TracerProvider()
8 | trace.set_tracer_provider(TRACE_PROVIDER)
9 |
--------------------------------------------------------------------------------
/tests/unit/tools/mcp/test_mcp_config.py:
--------------------------------------------------------------------------------
1 | import warnings
2 |
3 | import pytest
4 |
5 | from any_agent.config import MCPSse
6 |
7 |
8 | def test_sse_deprecation() -> None:
9 | sse1 = MCPSse(url="test.example.com:8888/sse") # noqa: F841
10 | warnings.filterwarnings("error")
11 | with pytest.raises(DeprecationWarning):
12 | sse2 = MCPSse(url="test.example.com:8888/sse") # noqa: F841
13 |
--------------------------------------------------------------------------------
/demo/Dockerfile:
--------------------------------------------------------------------------------
1 | FROM python:3.12-slim
2 |
3 | WORKDIR /app
4 |
5 | RUN apt-get update && apt-get install -y \
6 | build-essential \
7 | curl \
8 | git \
9 | && rm -rf /var/lib/apt/lists/*
10 |
11 | COPY requirements.txt ./
12 | COPY . ./demo/
13 |
14 | RUN pip3 install -r requirements.txt
15 |
16 | EXPOSE 8501
17 |
18 | HEALTHCHECK CMD curl --fail http://localhost:8501/_stcore/health
19 |
20 | ENTRYPOINT ["streamlit", "run", "demo/app.py", "--server.port=8501", "--server.address=0.0.0.0"]
21 |
--------------------------------------------------------------------------------
/scripts/run-mypy.sh:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env bash
2 | # Thank you to https://jaredkhan.com/blog/mypy-pre-commit for this super helpful script!
3 | # This script is called by the pre-commit hook.
4 | set -o errexit
5 |
6 | # Change directory to the project root directory.
7 | cd "$(dirname "$0")/.."
8 |
9 | # Install the dependencies into the mypy env.
10 | # Note that this can take seconds to run.
11 | python -m pip install -U -e '.[all,a2a,composio]' --quiet
12 |
13 | # Run on all files.
14 | python -m mypy src/
15 | python -m mypy tests/
16 |
--------------------------------------------------------------------------------
/docs/agents/frameworks/agno.md:
--------------------------------------------------------------------------------
1 | # Agno
2 |
3 | [https://github.com/agno-agi/agno](https://github.com/agno-agi/agno)
4 |
5 | ## Default Agent Type
6 |
7 | We use [`agno.agent.Agent`](https://docs.agno.com/reference/agents/agent) as default.
8 | Check the reference to find additional supported `agent_args`.
9 |
10 | ## Default Model Type
11 |
12 | We use [`any_llm`](https://mozilla-ai.github.io/any-llm/) as the default model provider.
13 | Check the [AnyLLM documentation](https://mozilla-ai.github.io/any-llm/) for supported providers and `model_args`.
14 |
--------------------------------------------------------------------------------
/tests/snapshots/test_trace.py:
--------------------------------------------------------------------------------
1 | from syrupy.assertion import SnapshotAssertion
2 |
3 | from any_agent.tracing.agent_trace import AgentTrace
4 |
5 |
6 | def test_agent_trace_snapshot(
7 | agent_trace: AgentTrace, snapshot: SnapshotAssertion
8 | ) -> None:
9 | # Snapshot the dict representation (so you see changes in the schema)
10 | # If this assert fails and you decide that you're ok with the new schema,
11 | # you can easily update the snapshot by running:
12 | # pytest tests/snapshots --snapshot-update
13 | assert agent_trace.model_dump() == snapshot(
14 | name=agent_trace.spans[0].context.trace_id
15 | )
16 |
--------------------------------------------------------------------------------
/src/any_agent/__init__.py:
--------------------------------------------------------------------------------
1 | from importlib.metadata import PackageNotFoundError, version
2 |
3 | from .config import AgentConfig, AgentFramework
4 | from .frameworks.any_agent import AgentRunError, AnyAgent
5 | from .tracing.agent_trace import AgentTrace
6 |
7 | try:
8 | __version__ = version("any-agent")
9 | except PackageNotFoundError:
10 | # In the case of local development
11 | # i.e., running directly from the source directory without package being installed
12 | __version__ = "0.0.0-dev"
13 |
14 | __all__ = [
15 | "AgentConfig",
16 | "AgentFramework",
17 | "AgentRunError",
18 | "AgentTrace",
19 | "AnyAgent",
20 | "__version__",
21 | ]
22 |
--------------------------------------------------------------------------------
/src/any_agent/callbacks/__init__.py:
--------------------------------------------------------------------------------
1 | from .base import Callback
2 | from .context import Context
3 | from .span_print import ConsolePrintSpan
4 |
5 | __all__ = ["Callback", "ConsolePrintSpan", "Context"]
6 |
7 |
8 | def get_default_callbacks() -> list[Callback]:
9 | """Return instances of the default callbacks used in any-agent.
10 |
11 | This function is called internally when the user doesn't provide a
12 | value for [`AgentConfig.callbacks`][any_agent.config.AgentConfig.callbacks].
13 |
14 | Returns:
15 | A list of instances containing:
16 |
17 | - [`ConsolePrintSpan`][any_agent.callbacks.span_print.ConsolePrintSpan]
18 |
19 | """
20 | return [ConsolePrintSpan()]
21 |
--------------------------------------------------------------------------------
/src/any_agent/callbacks/span_end.py:
--------------------------------------------------------------------------------
1 | # mypy: disable-error-code="no-untyped-def"
2 | from any_agent.callbacks.base import Callback
3 | from any_agent.callbacks.context import Context
4 |
5 |
6 | def _span_end(context: Context) -> Context:
7 | context.current_span.end()
8 | context.trace.add_span(context.current_span)
9 | return context
10 |
11 |
12 | class SpanEndCallback(Callback):
13 | """End the current span and add it to the corresponding `AgentTrace`."""
14 |
15 | def after_llm_call(self, context: Context, *args, **kwargs) -> Context:
16 | return _span_end(context)
17 |
18 | def after_tool_execution(self, context: Context, *args, **kwargs) -> Context:
19 | return _span_end(context)
20 |
--------------------------------------------------------------------------------
/src/any_agent/tools/__init__.py:
--------------------------------------------------------------------------------
1 | from .a2a import a2a_tool, a2a_tool_async
2 | from .final_output import prepare_final_output
3 | from .mcp.mcp_client import MCPClient
4 | from .user_interaction import (
5 | ask_user_verification,
6 | send_console_message,
7 | show_final_output,
8 | show_plan,
9 | )
10 | from .web_browsing import search_tavily, search_web, visit_webpage
11 | from .wrappers import _wrap_tools
12 |
13 | __all__ = [
14 | "MCPClient",
15 | "_wrap_tools",
16 | "a2a_tool",
17 | "a2a_tool_async",
18 | "ask_user_verification",
19 | "prepare_final_output",
20 | "search_tavily",
21 | "search_web",
22 | "send_console_message",
23 | "show_final_output",
24 | "show_plan",
25 | "visit_webpage",
26 | ]
27 |
--------------------------------------------------------------------------------
/docs/agents/frameworks/google_adk.md:
--------------------------------------------------------------------------------
1 | # Google Agent Development Kit (ADK)
2 |
3 | [https://github.com/google/adk-python](https://github.com/google/adk-python)
4 |
5 | ## Default Agent Type
6 |
7 | We use [`google.adk.agents.llm_agent.LlmAgent`](https://google.github.io/adk-docs/agents/llm-agents/) as default.
8 | Check the reference to find additional supported `agent_args`.
9 |
10 | ## Default Model Type
11 |
12 | We use [`any_llm`](https://mozilla-ai.github.io/any-llm/) as the default model provider.
13 | Check the [AnyLLM documentation](https://mozilla-ai.github.io/any-llm/) for supported providers and `model_args`.
14 |
15 | ## Run args
16 |
17 | Check [`RunConfig`](https://google.github.io/adk-docs/runtime/runconfig/) to find additional supported `AnyAgent.run` args.
18 |
--------------------------------------------------------------------------------
/tests/unit/tools/conftest.py:
--------------------------------------------------------------------------------
1 | import logging
2 | from collections.abc import Generator
3 | from typing import Any
4 | from unittest.mock import patch
5 |
6 | import pytest
7 |
8 |
9 | @pytest.fixture
10 | def local_logger(monkeypatch: Any) -> Generator[logging.Logger, None, None]:
11 | # Create a fresh logger instance for each test
12 | test_logger = logging.getLogger(f"test_logger_{id(object())}")
13 | # Remove all handlers
14 | for handler in test_logger.handlers[:]:
15 | test_logger.removeHandler(handler)
16 | test_logger.handlers.clear()
17 | test_logger.setLevel(logging.NOTSET)
18 | test_logger.propagate = False
19 | # Patch the logger in both import styles
20 | with patch("any_agent.logging.logger", test_logger):
21 | yield test_logger
22 |
--------------------------------------------------------------------------------
/docs/agents/frameworks/openai.md:
--------------------------------------------------------------------------------
1 | # OpenAI Agents SDK
2 |
3 | [https://github.com/openai/openai-agents-python](https://github.com/openai/openai-agents-python)
4 |
5 | ## Default Agent Type
6 |
7 | We use [`agents.Agent`](ttps://openai.github.io/openai-agents-python/ref/agent/#agents.agent.Agent) as default.
8 | Check the reference to find additional supported `agent_args`.
9 |
10 | ## Default Model Type
11 |
12 | We use [`any_llm`](https://mozilla-ai.github.io/any-llm/) as the default model provider.
13 | Check the [AnyLLM documentation](https://mozilla-ai.github.io/any-llm/) for supported providers and `model_args`.
14 |
15 | ## Run args
16 |
17 | Check [`agents.run.Runner.run`](https://openai.github.io/openai-agents-python/ref/run/#agents.run.Runner.run) to find additional supported `AnyAgent.run` args.
18 |
--------------------------------------------------------------------------------
/docs/agents/frameworks/smolagents.md:
--------------------------------------------------------------------------------
1 | # smolagents
2 |
3 | [https://github.com/huggingface/smolagents](https://github.com/huggingface/smolagents)
4 |
5 | ## Default Agent Type
6 |
7 | We use [`smolagents.ToolCallingAgent`](https://huggingface.co/docs/smolagents/reference/agents#smolagents.ToolCallingAgent) as default.
8 | Check the reference to find additional supported `agent_args`.
9 |
10 | ## Default Model Type
11 |
12 | We use [`any_llm`](https://mozilla-ai.github.io/any-llm/) as the default model provider.
13 | Check the [AnyLLM documentation](https://mozilla-ai.github.io/any-llm/) for supported providers and `model_args`.
14 |
15 | ## Run args
16 |
17 | Check [`smolagents.MultiStepAgent.run`](https://huggingface.co/docs/smolagents/main/en/reference/agents#smolagents.MultiStepAgent.run) to find additional supported `AnyAgent.run` args.
18 |
--------------------------------------------------------------------------------
/src/any_agent/serving/mcp/config_mcp.py:
--------------------------------------------------------------------------------
1 | from pydantic import BaseModel, ConfigDict
2 |
3 |
4 | class MCPServingConfig(BaseModel):
5 | """Configuration for serving an agent using the Model Context Protocol (MCP).
6 |
7 | Example:
8 | config = MCPServingConfig(
9 | port=8080,
10 | endpoint="/my-agent",
11 | )
12 |
13 | """
14 |
15 | model_config = ConfigDict(extra="forbid")
16 |
17 | host: str = "localhost"
18 | """Will be passed as argument to `uvicorn.run`."""
19 |
20 | port: int = 5000
21 | """Will be passed as argument to `uvicorn.run`."""
22 |
23 | endpoint: str = "/"
24 | """Will be pass as argument to `Starlette().add_route`"""
25 |
26 | log_level: str = "warning"
27 | """Will be passed as argument to the `uvicorn` server."""
28 |
29 | version: str = "0.1.0"
30 |
--------------------------------------------------------------------------------
/tests/unit/conftest.py:
--------------------------------------------------------------------------------
1 | import os
2 | from collections.abc import Generator
3 |
4 | import pytest
5 |
6 | from any_agent.config import AgentFramework
7 |
8 |
9 | @pytest.fixture(autouse=True)
10 | def mock_api_keys_for_unit_tests(
11 | request: pytest.FixtureRequest,
12 | ) -> Generator[None, None, None]:
13 | """Automatically provide dummy API keys for unit tests to avoid API key requirements."""
14 | if "agent_framework" in request.fixturenames:
15 | agent_framework = request.getfixturevalue("agent_framework")
16 | # Only set dummy API key if we're in a test that uses the agent_framework fixture
17 | # and the framework is OPENAI (which uses any-llm with the AnyLLM.create class-based interface)
18 | if agent_framework == AgentFramework.OPENAI:
19 | os.environ["MISTRAL_API_KEY"] = "dummy-mistral-key-for-unit-tests"
20 | yield # noqa: PT022
21 |
--------------------------------------------------------------------------------
/docs/agents/frameworks/llama_index.md:
--------------------------------------------------------------------------------
1 | # LlamaIndex
2 |
3 | [https://github.com/run-llama/llama_index](https://github.com/run-llama/llama_index)
4 |
5 | ## Default Agent Type
6 |
7 | We use [`llama_index.core.agent.workflow.react_agent.FunctionAgent`](https://docs.llamaindex.ai/en/stable/api_reference/agent/#llama_index.core.agent.workflow.FunctionAgent) as default.
8 | However, this agent requires that tools are used. If no tools are used, any-agent will default to [`llama_index.core.agent.workflow.react_agent.ReActAgent`](https://docs.llamaindex.ai/en/stable/api_reference/agent/#llama_index.core.agent.workflow.ReActAgent).
9 | Check the reference to find additional supported `agent_args`.
10 |
11 | ## Default Model Type
12 |
13 | We use [`any_llm`](https://mozilla-ai.github.io/any-llm/) as the default model provider.
14 | Check the [AnyLLM documentation](https://mozilla-ai.github.io/any-llm/) for supported providers and `model_args`.
15 |
--------------------------------------------------------------------------------
/.github/workflows/release.yaml:
--------------------------------------------------------------------------------
1 | name: Release
2 |
3 | on:
4 | release:
5 | types: [published]
6 | workflow_dispatch:
7 |
8 | jobs:
9 | release:
10 | environment: pypi
11 | permissions:
12 | contents: read
13 | id-token: write
14 | runs-on: ubuntu-latest
15 | steps:
16 | - name: Check out the repository
17 | uses: actions/checkout@v5
18 | with:
19 | fetch-depth: 0
20 |
21 | - name: Set up Python
22 | uses: actions/setup-python@v6
23 | with:
24 | python-version: '3.11'
25 |
26 | - name: Upgrade pip
27 | run: |
28 | pip install --upgrade pip
29 | pip --version
30 |
31 | - name: Install
32 | run: python -m pip install build setuptools
33 |
34 | - name: Build package
35 | run: python -m build
36 |
37 | - name: Upload package
38 | if: github.event_name == 'release'
39 | uses: pypa/gh-action-pypi-publish@release/v1
40 |
--------------------------------------------------------------------------------
/docs/agents/frameworks/langchain.md:
--------------------------------------------------------------------------------
1 | # LangChain
2 |
3 | [https://github.com/langchain-ai/langchain](https://github.com/langchain-ai/langchain)
4 |
5 | [https://github.com/langchain-ai/langgraph](https://github.com/langchain-ai/langgraph)
6 |
7 | ## Default Agent Type
8 |
9 | We use [`langgraph.prebuilt.create_react_agent`](https://langchain-ai.github.io/langgraph/reference/agents/?h=create_rea#langgraph.prebuilt.chat_agent_executor.create_react_agent) as default.
10 | Check the reference to find additional supported `agent_args`.
11 |
12 | ## Default Model Type
13 |
14 | We use [`any_llm`](https://mozilla-ai.github.io/any-llm/) as the default model provider.
15 | Check the [AnyLLM documentation](https://mozilla-ai.github.io/any-llm/) for supported providers and `model_args`.
16 |
17 | ## Run args
18 |
19 | Check [`RunnableConfig`](https://python.langchain.com/api_reference/core/runnables/langchain_core.runnables.config.RunnableConfig.html) to find additional supported `AnyAgent.run` args.
20 |
--------------------------------------------------------------------------------
/docs/api/logging.md:
--------------------------------------------------------------------------------
1 | # Logging with `any-agent`
2 |
3 | `any-agent` comes with a logger powered by [Rich](https://github.com/Textualize/rich)
4 |
5 | ## Quick Start
6 |
7 | By default, logging is set up for you. But if you want to customize it, you can call:
8 |
9 | ```python
10 | from any_agent.logging import setup_logger
11 |
12 | setup_logger()
13 | ```
14 |
15 | ## Customizing the Logger
16 |
17 | View the docstring in [`setup_logger`][any_agent.logging.setup_logger] for a description of the arguments available .
18 |
19 | ### Example: Set Log Level to DEBUG
20 |
21 | ```python
22 | from any_agent.logging import setup_logger
23 | import logging
24 |
25 | setup_logger(level=logging.DEBUG)
26 | ```
27 |
28 | ### Example: Custom Log Format
29 |
30 | ```python
31 | setup_logger(log_format="%(asctime)s - %(levelname)s - %(message)s")
32 | ```
33 |
34 | ### Example: Propagate Logs
35 |
36 | ```python
37 | setup_logger(propagate=True)
38 | ```
39 |
40 | ::: any_agent.logging.setup_logger
41 |
--------------------------------------------------------------------------------
/docs/agents/frameworks/tinyagent.md:
--------------------------------------------------------------------------------
1 | # TinyAgent
2 |
3 | As part of the bare bones library, we provide our own Python implementation based on [HuggingFace Tiny Agents](https://huggingface.co/blog/tiny-agents).
4 |
5 | You can find it in [`any_agent.frameworks.tinyagent`](https://github.com/mozilla-ai/any-agent/blob/main/src/any_agent/frameworks/tinyagent.py).
6 |
7 | ## Examples
8 |
9 | ### Use MCP Tools
10 |
11 | ```python
12 | from any_agent import AnyAgent, AgentConfig
13 | from any_agent.config import MCPStdio
14 |
15 | agent = AnyAgent.create(
16 | "tinyagent",
17 | AgentConfig(
18 | model_id="mistral:mistral-small-latest",
19 | instructions="You must use the available tools to find an answer",
20 | tools=[
21 | MCPStdio(
22 | command="uvx",
23 | args=["duckduckgo-mcp-server"]
24 | )
25 | ]
26 | )
27 | )
28 |
29 | result = agent.run(
30 | "Which Agent Framework is the best??"
31 | )
32 | print(result.final_output)
33 | ```
34 |
--------------------------------------------------------------------------------
/.github/workflows/pre-commit-update.yaml:
--------------------------------------------------------------------------------
1 | name: pre-commit auto-update
2 |
3 | on:
4 | schedule:
5 | - cron: "0 0 1 * *"
6 | workflow_dispatch:
7 |
8 |
9 | permissions:
10 | contents: write
11 | pull-requests: write
12 |
13 | jobs:
14 | pre-commit-update:
15 | timeout-minutes: 30
16 | runs-on: ubuntu-latest
17 |
18 | steps:
19 | - name: Check out the repository
20 | uses: actions/checkout@v5
21 |
22 | - name: Set up Python
23 | uses: actions/setup-python@v6
24 | with:
25 | python-version: '3.11'
26 |
27 | - run: pip install pre-commit
28 |
29 | - run: pre-commit autoupdate
30 |
31 | - uses: peter-evans/create-pull-request@271a8d0340265f705b14b6d32b9829c1cb33d45e # v7.0.8
32 | if: always()
33 | with:
34 | token: ${{ secrets.GITHUB_TOKEN }}
35 | branch: pre-commit-autoupdate
36 | title: pre-commit autoupdate
37 | commit-message: "chore: pre-commit autoupdate"
38 | body: Update pre-commit hooks.
39 |
--------------------------------------------------------------------------------
/src/any_agent/callbacks/context.py:
--------------------------------------------------------------------------------
1 | from __future__ import annotations
2 |
3 | from dataclasses import dataclass
4 | from typing import TYPE_CHECKING, Any
5 |
6 | if TYPE_CHECKING:
7 | from opentelemetry.trace import Span, Tracer
8 |
9 | from any_agent.tracing.agent_trace import AgentTrace
10 |
11 |
12 | @dataclass
13 | class Context:
14 | """Object that will be shared across callbacks.
15 |
16 | Each AnyAgent.run has a separate `Context` available.
17 |
18 | `shared` can be used to store and pass information
19 | across different callbacks.
20 | """
21 |
22 | current_span: Span
23 | """You can use the span in your callbacks to get information consistently across frameworks.
24 |
25 | You can find information about the attributes (available under `current_span.attributes`) in
26 | [Attributes Reference](./tracing.md#any_agent.tracing.attributes).
27 | """
28 |
29 | trace: AgentTrace
30 | tracer: Tracer
31 |
32 | shared: dict[str, Any]
33 | """Can be used to store arbitrary information for sharing across callbacks."""
34 |
--------------------------------------------------------------------------------
/src/any_agent/serving/__init__.py:
--------------------------------------------------------------------------------
1 | from .mcp.config_mcp import MCPServingConfig
2 | from .mcp.server_mcp import (
3 | serve_mcp_async,
4 | )
5 | from .server_handle import ServerHandle
6 |
7 | __all__ = [
8 | "MCPServingConfig",
9 | "ServerHandle",
10 | "serve_mcp_async",
11 | ]
12 |
13 | try:
14 | from .a2a.config_a2a import A2AServingConfig
15 | from .a2a.server_a2a import (
16 | _get_a2a_app_async,
17 | serve_a2a_async,
18 | )
19 |
20 | __all__ += [
21 | "A2AServingConfig",
22 | "_get_a2a_app_async",
23 | "serve_a2a_async",
24 | ]
25 | except ImportError:
26 |
27 | def _raise(*args, **kwargs): # type: ignore[no-untyped-def]
28 | msg = "You need to `pip install 'any-agent[a2a]'` to use this method."
29 | raise ImportError(msg)
30 |
31 | A2AServingConfig = _raise # type: ignore[assignment,misc]
32 | _get_a2a_app_async = _raise
33 | serve_a2a_async = _raise
34 | __all__ += [
35 | "A2AServingConfig",
36 | "_get_a2a_app_async",
37 | "serve_a2a_async",
38 | ]
39 |
--------------------------------------------------------------------------------
/.github/workflows/lint.yaml:
--------------------------------------------------------------------------------
1 | name: Lint
2 |
3 | on:
4 | push:
5 | branches: [main]
6 | pull_request:
7 | workflow_dispatch:
8 |
9 | jobs:
10 | run-linter:
11 | timeout-minutes: 30
12 | runs-on: ubuntu-latest
13 |
14 | steps:
15 | - name: Check out the repository
16 | uses: actions/checkout@v5
17 |
18 | - name: Install the latest version of uv and set the python version to 3.13
19 | uses: astral-sh/setup-uv@v7
20 | with:
21 | python-version: 3.13
22 | activate-environment: true
23 |
24 | - name: Install pre-commit
25 | run: uv sync -U --group lint
26 |
27 | - uses: actions/cache@v4
28 | with:
29 | path: .mypy_cache
30 | key: ${{ runner.os }}-mypy-${{ hashFiles('pyproject.toml') }}
31 | restore-keys: |
32 | ${{ runner.os }}-mypy-
33 |
34 | - uses: actions/cache@v4
35 | with:
36 | path: ~/.cache/pre-commit
37 | key: ${{ runner.os }}-pre-commit-${{ hashFiles('.pre-commit-config.yaml') }}
38 | restore-keys: |
39 | ${{ runner.os }}-pre-commit-
40 |
41 | - name: pre-commit
42 | run: uv run pre-commit run --all-files --verbose
43 |
--------------------------------------------------------------------------------
/.github/workflows/tests-unit.yaml:
--------------------------------------------------------------------------------
1 | name: Unit Tests
2 |
3 | on:
4 | push:
5 | branches: [main]
6 | paths:
7 | - 'src/**'
8 | - 'tests/**'
9 | - '.github/workflows/**'
10 | - 'pyproject.toml'
11 | pull_request:
12 | paths:
13 | - 'src/**'
14 | - 'tests/**'
15 | - '.github/workflows/**'
16 | - 'pyproject.toml'
17 | workflow_dispatch:
18 |
19 | jobs:
20 | run-unit-tests:
21 | timeout-minutes: 30
22 |
23 | strategy:
24 | fail-fast: false
25 | matrix:
26 | os: [ubuntu-latest, macos-latest]
27 | python-version: ['3.11', '3.12', '3.13']
28 |
29 | runs-on: ${{ matrix.os }}
30 |
31 | steps:
32 | - uses: actions/checkout@v5
33 |
34 | - uses: astral-sh/setup-uv@v7
35 | with:
36 | python-version: ${{ matrix.python-version }}
37 | activate-environment: true
38 |
39 | - run: |
40 | uv sync -U --group tests --extra all --extra a2a
41 |
42 | - run: pytest tests/unit -v --cov --cov-report=xml
43 |
44 | - name: Upload coverage reports to Codecov
45 | if: always()
46 | uses: codecov/codecov-action@v5
47 | with:
48 | token: ${{ secrets.CODECOV_TOKEN }}
49 |
--------------------------------------------------------------------------------
/tests/integration/tools/test_composio.py:
--------------------------------------------------------------------------------
1 | import json
2 | import os
3 |
4 | from any_agent import AgentConfig, AgentFramework, AnyAgent
5 | from any_agent.tracing.attributes import GenAI
6 | from any_agent.tools.composio import CallableProvider
7 |
8 | from composio import Composio
9 |
10 |
11 | def test_composio(agent_framework: AgentFramework) -> None:
12 | cpo = Composio(CallableProvider())
13 | tools = cpo.tools.get(
14 | user_id=os.environ["COMPOSIO_USER_ID"],
15 | search="repository issues",
16 | toolkits=["GITHUB"],
17 | limit=10,
18 | )
19 | agent = AnyAgent.create(
20 | "tinyagent",
21 | AgentConfig(
22 | model_id="openai:gpt-4.1-mini",
23 | instructions="You summarize GitHub Issues",
24 | tools=tools,
25 | ),
26 | )
27 | agent_trace = agent.run(
28 | "Summary of open issues with label `callbacks`, no `assignee` in `mozilla-ai/any-agent`"
29 | )
30 | tool_execution = next(s for s in agent_trace.spans if s.is_tool_execution())
31 | assert tool_execution is not None
32 | assert tool_execution.name == "execute_tool GITHUB_LIST_REPOSITORY_ISSUES"
33 | assert tool_execution.attributes[GenAI.OUTPUT] is not None
34 |
--------------------------------------------------------------------------------
/src/any_agent/tools/user_interaction.py:
--------------------------------------------------------------------------------
1 | from rich.prompt import Prompt
2 |
3 | from any_agent.logging import logger
4 |
5 |
6 | def show_plan(plan: str) -> str:
7 | """Show the current plan to the user.
8 |
9 | Args:
10 | plan: The current plan.
11 |
12 | """
13 | logger.info(f"Current plan: {plan}")
14 | return plan
15 |
16 |
17 | def show_final_output(answer: str) -> str:
18 | """Show the final answer to the user.
19 |
20 | Args:
21 | answer: The final answer.
22 |
23 | """
24 | logger.info(f"Final output: {answer}")
25 | return answer
26 |
27 |
28 | def ask_user_verification(query: str) -> str:
29 | """Asks user to verify the given `query`.
30 |
31 | Args:
32 | query: The question that requires verification.
33 |
34 | """
35 | return input(f"{query} => Type your answer here:")
36 |
37 |
38 | def send_console_message(user: str, query: str) -> str:
39 | """Send the specified user a message via console and returns their response.
40 |
41 | Args:
42 | query: The question to ask the user.
43 | user: The user to ask the question to.
44 |
45 | Returns:
46 | str: The user's response.
47 |
48 | """
49 | return Prompt.ask(f"{query}\n{user}")
50 |
--------------------------------------------------------------------------------
/scripts/wake_up_hf_endpoint.py:
--------------------------------------------------------------------------------
1 | import argparse
2 | import time
3 |
4 | from aiohttp.client_exceptions import ClientResponseError
5 | from any_llm.api import completion
6 |
7 | HF_ENDPOINT = "https://y0okp71n85ezo5nr.us-east-1.aws.endpoints.huggingface.cloud/v1/"
8 |
9 |
10 | def wake_up_hf_endpoint(retry: int = 0):
11 | while True:
12 | try:
13 | completion(
14 | model="huggingface:tgi",
15 | messages=[{"role": "user", "content": "Are you awake?"}],
16 | api_base=HF_ENDPOINT,
17 | )
18 | break
19 | except ClientResponseError as e:
20 | if not retry:
21 | print(f"Endpoint not ready, giving up...\n{e}")
22 | return
23 |
24 | print(f"Endpoint not ready, retrying...\n{e}")
25 | time.sleep(retry)
26 |
27 | print("Endpoint ready")
28 |
29 |
30 | if __name__ == "__main__":
31 | parser = argparse.ArgumentParser(description="Wake up Hugging Face endpoint")
32 | parser.add_argument(
33 | "--retry",
34 | type=int,
35 | default=0,
36 | help="Retry interval in seconds (0 means no retry)",
37 | )
38 | args = parser.parse_args()
39 | wake_up_hf_endpoint(retry=args.retry)
40 |
--------------------------------------------------------------------------------
/.github/workflows/tests-docs.yaml:
--------------------------------------------------------------------------------
1 | name: Tests for Docs
2 |
3 | on:
4 | push:
5 | branches: [main]
6 | paths:
7 | - 'src/**'
8 | - 'tests/**'
9 | - 'docs/**'
10 | - '.github/workflows/**'
11 | - 'pyproject.toml'
12 | pull_request:
13 | paths:
14 | - 'src/**'
15 | - 'tests/**'
16 | - 'docs/**'
17 | - '.github/workflows/**'
18 | - 'pyproject.toml'
19 | workflow_dispatch:
20 |
21 | jobs:
22 | run-docs-tests:
23 | timeout-minutes: 30
24 | runs-on: ubuntu-latest
25 |
26 | steps:
27 | - name: Check out the repository
28 | uses: actions/checkout@v5
29 |
30 | - name: Install the latest version of uv and set the python version to 3.13
31 | uses: astral-sh/setup-uv@v7
32 | with:
33 | python-version: 3.13
34 | activate-environment: true
35 |
36 | - name: Install
37 | run: |
38 | uv sync --group tests --extra all --extra a2a --extra composio
39 |
40 | - name: Run Documentation tests
41 | run: pytest tests/docs -v --cov --cov-report=xml
42 |
43 | - name: Upload coverage reports to Codecov
44 | if: always()
45 | uses: codecov/codecov-action@v5
46 | with:
47 | token: ${{ secrets.CODECOV_TOKEN }}
48 |
--------------------------------------------------------------------------------
/.github/workflows/publish_demo_to_hf_spaces.yaml:
--------------------------------------------------------------------------------
1 | name: Publish demo to Hugging Face Spaces
2 |
3 | on:
4 | release:
5 | types: [published]
6 | workflow_dispatch:
7 |
8 | jobs:
9 | sync-space:
10 | runs-on: ubuntu-latest
11 | env:
12 | REPO_NAME: mozilla-ai/any-agent
13 | HTTPS_REPO: https://${{ secrets.HF_USERNAME }}:${{ secrets.HF_TOKEN }}@huggingface.co/spaces/mozilla-ai/any-agent-demo
14 | steps:
15 | - uses: actions/checkout@v5
16 | with:
17 | fetch-depth: 0
18 |
19 | - run: git clone ${{ env.HTTPS_REPO}} hf-space
20 |
21 | - run: |
22 | cp -rT demo hf-space
23 |
24 | - run: |
25 | cd hf-space
26 | git config user.name 'github-actions[bot]'
27 | git config user.email 'github-actions[bot]@users.noreply.github.com'
28 | git add .
29 | git commit -m "Sync with https://github.com/${{ env.REPO_NAME }}"
30 |
31 | - name: Push to Hugging Face
32 | run: |
33 | cd hf-space
34 | git push ${{ env.HTTPS_REPO}} main
35 |
36 | - name: Reboot Space
37 | if: always()
38 | env:
39 | HF_TOKEN: ${{ secrets.HF_TOKEN }}
40 | run: |
41 | pip install huggingface_hub
42 | python demo/restart_space.py
43 |
--------------------------------------------------------------------------------
/src/any_agent/logging.py:
--------------------------------------------------------------------------------
1 | import logging
2 | from typing import Any
3 |
4 | from rich.logging import RichHandler
5 |
6 | logger = logging.getLogger("any_agent")
7 |
8 |
9 | def setup_logger(
10 | level: int = logging.ERROR,
11 | rich_tracebacks: bool = True,
12 | log_format: str | None = None,
13 | propagate: bool = False,
14 | **kwargs: Any,
15 | ) -> None:
16 | """Configure the any_agent logger with the specified settings.
17 |
18 | Args:
19 | level: The logging level to use (default: logging.INFO)
20 | rich_tracebacks: Whether to enable rich tracebacks (default: True)
21 | log_format: Optional custom log format string
22 | propagate: Whether to propagate logs to parent loggers (default: False)
23 | **kwargs: Additional keyword arguments to pass to RichHandler
24 |
25 | """
26 | logger.setLevel(level)
27 | logger.propagate = propagate
28 |
29 | # Remove any existing handlers
30 | for handler in logger.handlers[:]:
31 | logger.removeHandler(handler)
32 |
33 | handler = RichHandler(rich_tracebacks=rich_tracebacks, markup=True, **kwargs)
34 |
35 | if log_format:
36 | formatter = logging.Formatter(log_format)
37 | handler.setFormatter(formatter)
38 |
39 | logger.addHandler(handler)
40 |
41 |
42 | # Set default configuration
43 | setup_logger()
44 |
--------------------------------------------------------------------------------
/demo/config.py:
--------------------------------------------------------------------------------
1 | import os
2 | import tempfile
3 | from datetime import datetime, timedelta
4 | from typing import Annotated
5 |
6 | import geocoder
7 | from pydantic import AfterValidator, BaseModel, ConfigDict, FutureDatetime, PositiveInt
8 | from rich.prompt import Prompt
9 |
10 | from any_agent import AgentFramework
11 | from any_agent.config import AgentConfig
12 | from any_agent.logging import logger
13 |
14 | INPUT_PROMPT_TEMPLATE = """
15 | According to the forecast, what will be the best spot to surf around {LOCATION},
16 | in a {MAX_DRIVING_HOURS} hour driving radius,
17 | at {DATE}?"
18 | """.strip()
19 |
20 |
21 | def validate_prompt(value) -> str:
22 | for placeholder in ("{LOCATION}", "{MAX_DRIVING_HOURS}", "{DATE}"):
23 | if placeholder not in value:
24 | raise ValueError(f"prompt must contain {placeholder}")
25 | return value
26 |
27 |
28 | class Config(BaseModel):
29 | model_config = ConfigDict(extra="forbid")
30 |
31 | location: str
32 | max_driving_hours: PositiveInt
33 | date: FutureDatetime
34 | input_prompt_template: Annotated[str, AfterValidator(validate_prompt)] = (
35 | INPUT_PROMPT_TEMPLATE
36 | )
37 |
38 | framework: AgentFramework
39 |
40 | main_agent: AgentConfig
41 |
42 | evaluation_model: str | None = None
43 | evaluation_criteria: list[dict[str, str]] | None = None
44 |
--------------------------------------------------------------------------------
/src/any_agent/callbacks/base.py:
--------------------------------------------------------------------------------
1 | # mypy: disable-error-code="no-untyped-def"
2 | from __future__ import annotations
3 |
4 | from typing import TYPE_CHECKING
5 |
6 | if TYPE_CHECKING:
7 | from .context import Context
8 |
9 |
10 | class Callback:
11 | """Base class for AnyAgent callbacks."""
12 |
13 | def before_agent_invocation(self, context: Context, *args, **kwargs) -> Context:
14 | """Will be called before the Agent invocation starts."""
15 | return context
16 |
17 | def before_llm_call(self, context: Context, *args, **kwargs) -> Context:
18 | """Will be called before any LLM Call starts."""
19 | return context
20 |
21 | def before_tool_execution(self, context: Context, *args, **kwargs) -> Context:
22 | """Will be called before any Tool Execution starts."""
23 | return context
24 |
25 | def after_agent_invocation(self, context: Context, *args, **kwargs) -> Context:
26 | """Will be called once the Agent invocation ends."""
27 | return context
28 |
29 | def after_llm_call(self, context: Context, *args, **kwargs) -> Context:
30 | """Will be called after any LLM Call is completed."""
31 | return context
32 |
33 | def after_tool_execution(self, context: Context, *args, **kwargs) -> Context:
34 | """Will be called after any Tool Execution is completed."""
35 | return context
36 |
--------------------------------------------------------------------------------
/tests/unit/callbacks/wrappers/test_get_wrapper_and_unwrap.py:
--------------------------------------------------------------------------------
1 | from unittest.mock import MagicMock
2 |
3 | from any_agent import AgentFramework
4 | from any_agent.callbacks.wrappers import _get_wrapper_by_framework
5 |
6 |
7 | async def test_unwrap_before_wrap(agent_framework: AgentFramework) -> None:
8 | wrapper = _get_wrapper_by_framework(agent_framework)
9 | await wrapper.unwrap(MagicMock())
10 |
11 |
12 | async def test_google_instrument_uninstrument() -> None:
13 | """Regression test for https://github.com/mozilla-ai/any-agent/issues/467"""
14 | agent = MagicMock()
15 | agent._agent.before_model_callback = None
16 | agent._agent.after_model_callback = None
17 | agent._agent.before_tool_callback = None
18 | agent._agent.after_tool_callback = None
19 |
20 | wrapper = _get_wrapper_by_framework(AgentFramework.GOOGLE)
21 |
22 | await wrapper.wrap(agent)
23 | assert callable(agent._agent.before_model_callback)
24 | assert callable(agent._agent.after_model_callback)
25 | assert callable(agent._agent.before_tool_callback)
26 | assert callable(agent._agent.after_tool_callback)
27 |
28 | await wrapper.unwrap(agent)
29 | assert agent._agent.before_model_callback is None
30 | assert agent._agent.after_model_callback is None
31 | assert agent._agent.before_tool_callback is None
32 | assert agent._agent.after_tool_callback is None
33 |
--------------------------------------------------------------------------------
/docs/agents/models.md:
--------------------------------------------------------------------------------
1 | # Model Configuration
2 |
3 | ## Overview
4 |
5 | Model configuration in `any-agent` is designed to be consistent across all supported frameworks. We use [`any-llm`](https://mozilla-ai.github.io/any-llm/) as the default model provider, which acts as a unified interface allowing you to use any language model from any provider with the same syntax.
6 |
7 | ## Configuration Parameters
8 |
9 | The model configuration is defined through several parameters in [`AgentConfig`][any_agent.config.AgentConfig]:
10 |
11 | The `model_id` parameter selects which language model your agent will use. The format depends on the provider:
12 |
13 | The `model_args` parameter allows you to pass additional arguments to the model, such as `temperature`, `top_k`, and other provider-specific parameters.
14 |
15 | The `api_base` parameter allows you to specify a custom API endpoint. This is useful when:
16 |
17 | - Using a local model server (e.g., Ollama, llama.cpp, llamafile)
18 | - Routing through a proxy
19 | - Using a self-hosted model endpoint
20 |
21 | The `api_key` parameter allows you to explicitly specify an API key for authentication. By default, `any-llm` will automatically search for common environment variables (like `OPENAI_API_KEY`, `ANTHROPIC_API_KEY`, etc.).å
22 |
23 | See the [AnyLLM Provider Documentation](https://mozilla-ai.github.io/any-llm/providers/) for the complete list of supported providers.å
24 |
--------------------------------------------------------------------------------
/src/any_agent/callbacks/wrappers/__init__.py:
--------------------------------------------------------------------------------
1 | from __future__ import annotations
2 |
3 | from typing import assert_never
4 |
5 | from any_agent import AgentFramework
6 |
7 | from .agno import _AgnoWrapper
8 | from .google import _GoogleADKWrapper
9 | from .langchain import _LangChainWrapper
10 | from .llama_index import _LlamaIndexWrapper
11 | from .openai import _OpenAIAgentsWrapper
12 | from .smolagents import _SmolagentsWrapper
13 | from .tinyagent import _TinyAgentWrapper
14 |
15 |
16 | def _get_wrapper_by_framework(
17 | framework: AgentFramework,
18 | ) -> (
19 | _AgnoWrapper
20 | | _GoogleADKWrapper
21 | | _LangChainWrapper
22 | | _LlamaIndexWrapper
23 | | _OpenAIAgentsWrapper
24 | | _SmolagentsWrapper
25 | | _TinyAgentWrapper
26 | ):
27 | if framework is AgentFramework.AGNO:
28 | return _AgnoWrapper()
29 |
30 | if framework is AgentFramework.GOOGLE:
31 | return _GoogleADKWrapper()
32 |
33 | if framework is AgentFramework.LANGCHAIN:
34 | return _LangChainWrapper()
35 |
36 | if framework is AgentFramework.LLAMA_INDEX:
37 | return _LlamaIndexWrapper()
38 |
39 | if framework is AgentFramework.OPENAI:
40 | return _OpenAIAgentsWrapper()
41 |
42 | if framework is AgentFramework.SMOLAGENTS:
43 | return _SmolagentsWrapper()
44 |
45 | if framework is AgentFramework.TINYAGENT:
46 | return _TinyAgentWrapper()
47 |
48 | assert_never(framework)
49 |
--------------------------------------------------------------------------------
/.github/workflows/docs.yaml:
--------------------------------------------------------------------------------
1 | name: Documentation
2 |
3 | on:
4 | push:
5 | branches: [main]
6 | paths:
7 | - mkdocs.yml
8 | - 'docs/**'
9 | - 'src/**'
10 | - 'scripts/hooks.py'
11 | - 'pyproject.toml'
12 | pull_request:
13 | paths:
14 | - mkdocs.yml
15 | - 'docs/**'
16 | - 'src/**'
17 | - 'scripts/hooks.py'
18 | - 'pyproject.toml'
19 | workflow_dispatch:
20 |
21 | jobs:
22 | docs:
23 | permissions:
24 | contents: write
25 | runs-on: ubuntu-latest
26 | steps:
27 | - name: Check out the repository
28 | uses: actions/checkout@v5
29 | with:
30 | fetch-depth: 0
31 |
32 | - name: Install the latest version of uv and set the python version to 3.13
33 | uses: astral-sh/setup-uv@v7
34 | with:
35 | python-version: 3.13
36 | activate-environment: true
37 |
38 | - name: Configure git
39 | run: |
40 | git config user.name 'github-actions[bot]'
41 | git config user.email 'github-actions[bot]@users.noreply.github.com'
42 |
43 | - name: Install requirements
44 | run: |
45 | uv sync -U --group docs
46 |
47 | - name: Build docs
48 | if: github.event_name == 'pull_request'
49 | run: mkdocs build -s
50 |
51 | - name: Publish docs
52 | if: ${{ github.event_name == 'push' || github.event_name == 'workflow_dispatch' }}
53 | run: mkdocs gh-deploy
54 |
--------------------------------------------------------------------------------
/tests/integration/tools/test_wrap_tools.py:
--------------------------------------------------------------------------------
1 | from collections.abc import Sequence
2 | from typing import Any
3 |
4 | import pytest
5 | from agents.tool import Tool as OpenaiClass
6 | from any_llm.utils.aio import run_async_in_sync
7 | from google.adk.tools import FunctionTool as GoogleClass
8 | from langchain_core.tools import BaseTool as LangchainClass
9 | from llama_index.core.tools import FunctionTool as LlamaindexClass
10 | from smolagents.tools import Tool as SmolagentsClass
11 |
12 | from any_agent import AgentFramework
13 | from any_agent.config import Tool
14 | from any_agent.tools import search_web, visit_webpage
15 | from any_agent.tools.wrappers import _wrap_tools
16 |
17 |
18 | def wrap_sync(
19 | tools: Sequence[Tool],
20 | framework: AgentFramework,
21 | ) -> list[Tool]:
22 | wrapped_tools, _ = run_async_in_sync(_wrap_tools(tools, framework))
23 | return wrapped_tools
24 |
25 |
26 | @pytest.mark.parametrize(
27 | ("framework", "expected_class"),
28 | [
29 | (AgentFramework.GOOGLE, GoogleClass),
30 | (AgentFramework.LANGCHAIN, LangchainClass),
31 | (AgentFramework.LLAMA_INDEX, LlamaindexClass),
32 | (AgentFramework.OPENAI, OpenaiClass),
33 | (AgentFramework.SMOLAGENTS, SmolagentsClass),
34 | ],
35 | )
36 | def test_wrap_tools(framework: AgentFramework, expected_class: Any) -> None:
37 | wrapped_tools = wrap_sync([search_web, visit_webpage], framework)
38 | assert all(isinstance(tool, expected_class) for tool in wrapped_tools)
39 |
--------------------------------------------------------------------------------
/docs/index.md:
--------------------------------------------------------------------------------
1 | # any-agent
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 | `any-agent` is a Python library providing a single interface to different agent frameworks.
10 |
11 | !!! warning
12 |
13 | Compared to traditional code-defined workflows, agent frameworks introduce complexity,
14 | additional security implications to consider, and demand much more computational power.
15 |
16 | Before jumping to use one, carefully consider and evaluate how much value you
17 | would get compared to manually defining a sequence of tools and LLM calls.
18 |
19 | ## Requirements
20 |
21 | - Python 3.11 or newer
22 |
23 | ## Installation
24 |
25 | You can install the bare bones library as follows (only [`TinyAgent`](./agents/frameworks/tinyagent.md) will be available):
26 |
27 | ```bash
28 | pip install any-agent
29 | ```
30 |
31 | Or you can install it with the required dependencies for different frameworks:
32 |
33 | ```bash
34 | pip install any-agent[agno,openai]
35 | ```
36 |
37 | Refer to [pyproject.toml](https://github.com/mozilla-ai/any-agent/blob/main/pyproject.toml) for a list of the options available.
38 |
39 | ## For AI Systems
40 |
41 | This documentation is available in two AI-friendly formats:
42 |
43 | - **[llms.txt](https://mozilla-ai.github.io/any-agent/llms.txt)** - A structured overview with curated links to key documentation sections
44 | - **[llms-full.txt](https://mozilla-ai.github.io/any-agent/llms-full.txt)** - Complete documentation content concatenated into a single file
45 |
--------------------------------------------------------------------------------
/.github/workflows/tests-cookbook.yaml:
--------------------------------------------------------------------------------
1 | name: Cookbook Tests
2 |
3 | on:
4 | push:
5 | branches: [main]
6 | paths:
7 | - 'src/**'
8 | - 'docs/cookbook/**'
9 | - '.github/workflows/**'
10 | - 'pyproject.toml'
11 | workflow_dispatch:
12 |
13 | jobs:
14 | run-cookbook-tests:
15 | timeout-minutes: 30
16 | runs-on: ubuntu-latest
17 |
18 | steps:
19 | - name: Check out the repository
20 | uses: actions/checkout@v5
21 |
22 | - name: Install the latest version of uv and set the python version to 3.13
23 | uses: astral-sh/setup-uv@v7
24 | with:
25 | python-version: 3.13
26 | activate-environment: true
27 |
28 | - name: Install
29 | run: |
30 | uv sync -U --group tests --extra all --extra a2a
31 |
32 | - name: Run Cookbook Tests
33 | env:
34 | HF_TOKEN: ${{ secrets.HF_TOKEN }}
35 | HF_ENDPOINT: ${{ secrets.HF_ENDPOINT }}
36 | MISTRAL_API_KEY: ${{ secrets.MISTRAL_API_KEY }}
37 | TAVILY_API_KEY: ${{ secrets.TAVILY_API_KEY }}
38 | OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
39 | run: |
40 | pytest tests/cookbooks -v -n 4 --cov --cov-report=xml
41 |
42 | - name: Upload finished cookbooks
43 | uses: actions/upload-artifact@v5
44 | if: always()
45 | with:
46 | path: "docs/cookbook/executed_*.ipynb"
47 |
48 | - name: Upload coverage reports to Codecov
49 | if: always()
50 | uses: codecov/codecov-action@v5
51 | with:
52 | token: ${{ secrets.CODECOV_TOKEN }}
53 |
--------------------------------------------------------------------------------
/src/any_agent/callbacks/span_generation/__init__.py:
--------------------------------------------------------------------------------
1 | from __future__ import annotations
2 |
3 | from typing import assert_never
4 |
5 | from any_agent import AgentFramework
6 |
7 | from .agno import _AgnoSpanGeneration
8 | from .google import _GoogleSpanGeneration
9 | from .langchain import _LangchainSpanGeneration
10 | from .llama_index import _LlamaIndexSpanGeneration
11 | from .openai import _OpenAIAgentsSpanGeneration
12 | from .smolagents import _SmolagentsSpanGeneration
13 | from .tinyagent import _TinyAgentSpanGeneration
14 |
15 | SpanGeneration = (
16 | _AgnoSpanGeneration
17 | | _GoogleSpanGeneration
18 | | _LangchainSpanGeneration
19 | | _LlamaIndexSpanGeneration
20 | | _OpenAIAgentsSpanGeneration
21 | | _SmolagentsSpanGeneration
22 | | _TinyAgentSpanGeneration
23 | )
24 |
25 |
26 | def _get_span_generation_callback(
27 | framework: AgentFramework,
28 | ) -> SpanGeneration:
29 | if framework is AgentFramework.AGNO:
30 | return _AgnoSpanGeneration()
31 |
32 | if framework is AgentFramework.GOOGLE:
33 | return _GoogleSpanGeneration()
34 |
35 | if framework is AgentFramework.LANGCHAIN:
36 | return _LangchainSpanGeneration()
37 |
38 | if framework is AgentFramework.LLAMA_INDEX:
39 | return _LlamaIndexSpanGeneration()
40 |
41 | if framework is AgentFramework.OPENAI:
42 | return _OpenAIAgentsSpanGeneration()
43 |
44 | if framework is AgentFramework.SMOLAGENTS:
45 | return _SmolagentsSpanGeneration()
46 |
47 | if framework is AgentFramework.TINYAGENT:
48 | return _TinyAgentSpanGeneration()
49 |
50 | assert_never(framework)
51 |
--------------------------------------------------------------------------------
/demo/tools/openstreetmap.py:
--------------------------------------------------------------------------------
1 | import json
2 |
3 | import requests
4 |
5 |
6 | def get_area_lat_lon(area_name: str) -> tuple[float, float]:
7 | """Get the latitude and longitude of an area from Nominatim.
8 |
9 | Uses the [Nominatim API](https://nominatim.org/release-docs/develop/api/Search/).
10 |
11 | Args:
12 | area_name: The name of the area.
13 |
14 | Returns:
15 | The area found.
16 |
17 | """
18 | response = requests.get(
19 | f"https://nominatim.openstreetmap.org/search?q={area_name}&format=jsonv2",
20 | headers={"User-Agent": "Mozilla/5.0"},
21 | )
22 | response.raise_for_status()
23 | area = json.loads(response.content.decode())
24 | return area[0]["lat"], area[0]["lon"]
25 |
26 |
27 | def driving_hours_to_meters(driving_hours: int) -> int:
28 | """Convert driving hours to meters assuming a 70 km/h average speed.
29 |
30 | Args:
31 | driving_hours: The driving hours.
32 |
33 | Returns:
34 | The distance in meters.
35 |
36 | """
37 | return driving_hours * 70 * 1000
38 |
39 |
40 | def get_lat_lon_center(bounds: dict) -> tuple[float, float]:
41 | """Get the latitude and longitude of the center of a bounding box.
42 |
43 | Args:
44 | bounds: The bounding box.
45 |
46 | ```json
47 | {
48 | "minlat": float,
49 | "minlon": float,
50 | "maxlat": float,
51 | "maxlon": float,
52 | }
53 | ```
54 |
55 | Returns:
56 | The latitude and longitude of the center.
57 |
58 | """
59 | return (
60 | (bounds["minlat"] + bounds["maxlat"]) / 2,
61 | (bounds["minlon"] + bounds["maxlon"]) / 2,
62 | )
63 |
--------------------------------------------------------------------------------
/demo/constants.py:
--------------------------------------------------------------------------------
1 | import os
2 |
3 | from tools import (
4 | get_area_lat_lon,
5 | get_wave_forecast,
6 | get_wind_forecast,
7 | )
8 |
9 | from any_agent.logging import logger
10 | from any_agent.tools.web_browsing import search_tavily, search_web, visit_webpage
11 |
12 | MODEL_OPTIONS = [
13 | "openai:gpt-4.1-nano",
14 | "openai:gpt-4.1-mini",
15 | "openai:gpt-4o",
16 | "gemini:gemini-2.0-flash-lite",
17 | "gemini:gemini-2.0-flash",
18 | ]
19 |
20 | DEFAULT_EVALUATION_MODEL = MODEL_OPTIONS[0]
21 |
22 | DEFAULT_EVALUATION_CRITERIA = [
23 | {
24 | "criteria": "Check if the agent considered at least three surf spot options",
25 | },
26 | {
27 | "criteria": "Check if the agent gathered wind forecasts for each surf spot being evaluated.",
28 | },
29 | {
30 | "criteria": "Check if the agent gathered wave forecasts for each surf spot being evaluated.",
31 | },
32 | {
33 | "criteria": "Check if the agent used any web search tools to explore which surf spots should be considered",
34 | },
35 | {
36 | "criteria": "Check if the final answer contains any description about the weather (air temp, chance of rain, etc) at the chosen location",
37 | },
38 | {
39 | "criteria": "Check if the final answer includes one of the surf spots evaluated by tools",
40 | },
41 | {
42 | "criteria": "Check if the final answer includes information about some alternative surf spots if the user is not satisfied with the chosen one",
43 | },
44 | ]
45 |
46 | DEFAULT_TOOLS = [
47 | get_wind_forecast,
48 | get_wave_forecast,
49 | get_area_lat_lon,
50 | search_web,
51 | visit_webpage,
52 | ]
53 | if os.getenv("TAVILY_API_KEY"):
54 | DEFAULT_TOOLS.append(search_tavily)
55 | else:
56 | logger.warning("TAVILY_API_KEY not set, skipping Tavily search tool")
57 |
--------------------------------------------------------------------------------
/tests/integration/frameworks/test_evaluation.py:
--------------------------------------------------------------------------------
1 | from datetime import datetime
2 |
3 | from any_agent import AgentFramework, AgentTrace
4 | from any_agent.evaluation.agent_judge import AgentJudge
5 | from any_agent.evaluation.llm_judge import LlmJudge
6 | from any_agent.evaluation.schemas import EvaluationOutput
7 | from any_agent.testing.helpers import get_default_agent_model_args
8 |
9 |
10 | def test_llm_judge(agent_trace: AgentTrace) -> None:
11 | llm_judge = LlmJudge(
12 | model_id="openai:gpt-4.1-nano",
13 | model_args={
14 | "temperature": 0.0,
15 | }, # Because it's an llm not agent, the default_model_args are not used
16 | )
17 | result1 = llm_judge.run(
18 | context=str(agent_trace.spans_to_messages()),
19 | question="Do the messages contain the year 2025?",
20 | )
21 | assert isinstance(result1, EvaluationOutput)
22 | assert result1.passed, (
23 | f"Expected agent to call write_file tool, but evaluation failed: {result1.reasoning}"
24 | )
25 |
26 |
27 | def test_agent_judge(agent_trace: AgentTrace) -> None:
28 | agent_judge = AgentJudge(
29 | model_id="openai:gpt-4.1-mini",
30 | model_args=get_default_agent_model_args(AgentFramework.TINYAGENT),
31 | )
32 |
33 | def get_current_year() -> str:
34 | """Get the current year"""
35 | return str(datetime.now().year)
36 |
37 | eval_trace = agent_judge.run(
38 | trace=agent_trace,
39 | question="Did the agent write the year to a file? Grab the messages from the trace and check if the write_file tool was called.",
40 | additional_tools=[get_current_year],
41 | )
42 | result = eval_trace.final_output
43 | assert isinstance(result, EvaluationOutput)
44 | assert result.passed, (
45 | f"Expected agent to write current year to file, but evaluation failed: {result.reasoning}"
46 | )
47 |
--------------------------------------------------------------------------------
/.pre-commit-config.yaml:
--------------------------------------------------------------------------------
1 | repos:
2 | - repo: https://github.com/pre-commit/pre-commit-hooks
3 | rev: "v6.0.0"
4 | hooks:
5 | - id: check-added-large-files
6 | - id: check-case-conflict
7 | - id: check-json
8 | - id: check-merge-conflict
9 | args:
10 | - '--assume-in-merge'
11 | - id: check-toml
12 | - id: check-yaml
13 | - id: end-of-file-fixer
14 | - id: mixed-line-ending
15 | args:
16 | - '--fix=lf'
17 | - id: sort-simple-yaml
18 | - id: trailing-whitespace
19 | exclude: \.ambr$
20 |
21 | - repo: https://github.com/astral-sh/ruff-pre-commit
22 | rev: "v0.14.5"
23 | hooks:
24 | - id: ruff
25 | args:
26 | - "--fix"
27 | - "--exit-non-zero-on-fix"
28 | exclude: src/any_agent/vendor|demo
29 | - id: ruff-format
30 |
31 | - repo: local
32 | # Thanks to https://jaredkhan.com/blog/mypy-pre-commit
33 | # We do not use pre-commit/mirrors-mypy,
34 | # as it is difficult to configure to run
35 | # with the dependencies correctly installed.
36 | hooks:
37 | - id: mypy
38 | name: mypy
39 | entry: "./scripts/run-mypy.sh"
40 | language: python
41 | language_version: python3.13
42 | additional_dependencies:
43 | - mypy==1.15.0
44 | - pip
45 | - types-requests
46 | - types-pyyaml
47 | types: [python]
48 | # use require_serial so that script is only called once per commit
49 | require_serial: true
50 | # Print the number of files as a sanity-check
51 | verbose: true
52 |
53 | - repo: https://github.com/codespell-project/codespell
54 | rev: "v2.4.1"
55 | hooks:
56 | - id: codespell
57 | exclude: CODE_OF_CONDUCT.md|tests/assets
58 |
59 | - repo: https://github.com/kynan/nbstripout
60 | rev: 0.8.2
61 | hooks:
62 | - id: nbstripout
63 |
--------------------------------------------------------------------------------
/tests/unit/tools/test_logging.py:
--------------------------------------------------------------------------------
1 | import logging
2 |
3 | from any_agent.logging import setup_logger
4 |
5 |
6 | def test_setup_logger_sets_level_and_propagate(local_logger: logging.Logger) -> None:
7 | setup_logger(level=logging.INFO, propagate=False)
8 | assert local_logger.level == logging.INFO
9 | assert local_logger.propagate is False
10 | setup_logger(level=logging.DEBUG, propagate=True)
11 | assert local_logger.level == logging.DEBUG
12 | assert local_logger.propagate is True
13 |
14 |
15 | def test_setup_logger_removes_existing_handlers(local_logger: logging.Logger) -> None:
16 | dummy_handler = logging.StreamHandler()
17 | local_logger.addHandler(dummy_handler)
18 | assert dummy_handler in local_logger.handlers
19 | setup_logger()
20 | assert dummy_handler not in local_logger.handlers
21 |
22 |
23 | def test_setup_logger_with_custom_format(local_logger: logging.Logger) -> None:
24 | custom_format = "%(levelname)s: %(message)s"
25 | setup_logger(log_format=custom_format)
26 | handler = local_logger.handlers[0]
27 | assert isinstance(handler.formatter, logging.Formatter)
28 | assert handler.formatter._fmt == custom_format
29 |
30 |
31 | def test_setup_logger_multiple_calls_idempotent(local_logger: logging.Logger) -> None:
32 | setup_logger(level=logging.INFO)
33 | first_handler = local_logger.handlers[0]
34 | setup_logger(level=logging.WARNING)
35 | second_handler = local_logger.handlers[0]
36 | assert first_handler is not second_handler
37 | assert local_logger.level == logging.WARNING
38 |
39 |
40 | def test_setup_logger_edits_global_logger(local_logger: logging.Logger) -> None:
41 | # This test now just checks that setup_logger edits the patched logger
42 | dummy_handler = logging.StreamHandler()
43 | local_logger.addHandler(dummy_handler)
44 | assert dummy_handler in local_logger.handlers
45 | setup_logger()
46 | assert dummy_handler not in local_logger.handlers
47 | assert local_logger.level == logging.ERROR
48 | assert local_logger.propagate is False
49 |
--------------------------------------------------------------------------------
/tests/integration/frameworks/test_thread_safe.py:
--------------------------------------------------------------------------------
1 | import asyncio
2 |
3 | import pytest
4 |
5 | from any_agent import AgentConfig, AgentFramework, AnyAgent
6 | from any_agent.testing.helpers import (
7 | DEFAULT_SMALL_MODEL_ID,
8 | get_default_agent_model_args,
9 | )
10 |
11 |
12 | @pytest.mark.asyncio
13 | async def test_run_agent_concurrently(agent_framework: AgentFramework) -> None:
14 | """When an agent is run concurrently, state from the first run shouldn't bleed into the second run"""
15 |
16 | def mock_capital(query: str) -> str:
17 | """Perform a duckduckgo web search based on your query (think a Google search) then returns the top search results.
18 |
19 | Args:
20 | query (str): The search query to perform.
21 |
22 | Returns:
23 | The top search results.
24 |
25 | """
26 | if "France" in query:
27 | return "The capital of France is Paris."
28 | if "Spain" in query:
29 | return "The capital of Spain is Madrid."
30 | return "No info"
31 |
32 | model_id = DEFAULT_SMALL_MODEL_ID
33 |
34 | agent = await AnyAgent.create_async(
35 | agent_framework,
36 | AgentConfig(
37 | model_id=model_id,
38 | instructions="You must use the tools to find an answer",
39 | model_args=get_default_agent_model_args(agent_framework),
40 | tools=[mock_capital],
41 | ),
42 | )
43 | results = await asyncio.gather(
44 | agent.run_async("What is the capital of France?"),
45 | agent.run_async("What is the capital of Spain?"),
46 | )
47 | outputs = [r.final_output for r in results]
48 | assert all(o is not None for o in outputs)
49 |
50 | assert sum("Paris" in str(o) for o in outputs) == 1
51 | assert sum("Madrid" in str(o) for o in outputs) == 1
52 |
53 | first_spans = results[0].spans
54 | second_spans = results[1].spans
55 | assert second_spans[: len(first_spans)] != first_spans, (
56 | "Spans from the first run should not be in the second"
57 | )
58 |
--------------------------------------------------------------------------------
/src/any_agent/serving/a2a/agent_card.py:
--------------------------------------------------------------------------------
1 | from __future__ import annotations
2 |
3 | import inspect
4 | from typing import TYPE_CHECKING
5 |
6 | from a2a.types import AgentCapabilities, AgentCard, AgentSkill
7 |
8 | from any_agent import AgentFramework
9 |
10 | if TYPE_CHECKING:
11 | from any_agent import AnyAgent
12 | from any_agent.serving.a2a.config_a2a import A2AServingConfig
13 |
14 |
15 | def _get_agent_card(agent: AnyAgent, serving_config: A2AServingConfig) -> AgentCard:
16 | skills = serving_config.skills
17 | if skills is None:
18 | skills = []
19 | for tool in agent._tools:
20 | if hasattr(tool, "name"):
21 | tool_name = tool.name
22 | tool_description = tool.description
23 | elif agent.framework is AgentFramework.LLAMA_INDEX:
24 | tool_name = tool.metadata.name
25 | tool_description = tool.metadata.description
26 | else:
27 | tool_name = tool.__name__
28 | tool_description = inspect.getdoc(tool)
29 | skills.append(
30 | AgentSkill(
31 | id=f"{agent.config.name}-{tool_name}",
32 | name=tool_name,
33 | description=tool_description,
34 | tags=[],
35 | )
36 | )
37 | if agent.config.description is None:
38 | msg = "Agent description is not set. Please set the `description` field in the `AgentConfig`."
39 | raise ValueError(msg)
40 | endpoint = serving_config.endpoint.lstrip("/")
41 | streaming = serving_config.stream_tool_usage
42 | return AgentCard(
43 | name=agent.config.name,
44 | description=agent.config.description,
45 | version=serving_config.version,
46 | default_input_modes=["text"],
47 | default_output_modes=["text"],
48 | url=f"http://{serving_config.host}:{serving_config.port}/{endpoint}",
49 | capabilities=AgentCapabilities(
50 | streaming=streaming, push_notifications=True, state_transition_history=False
51 | ),
52 | skills=skills,
53 | )
54 |
--------------------------------------------------------------------------------
/tests/unit/frameworks/test_agno.py:
--------------------------------------------------------------------------------
1 | from unittest.mock import AsyncMock, MagicMock, patch
2 |
3 | import pytest
4 |
5 | from any_agent import AgentConfig, AgentFramework, AnyAgent
6 |
7 |
8 | def test_load_agno_default() -> None:
9 | mock_agent = MagicMock()
10 | mock_model = MagicMock()
11 |
12 | with (
13 | patch("any_agent.frameworks.agno.Agent", mock_agent),
14 | patch("any_agent.frameworks.agno.DEFAULT_MODEL_TYPE", mock_model),
15 | ):
16 | AnyAgent.create(
17 | AgentFramework.AGNO, AgentConfig(model_id="mistral:mistral-small-latest")
18 | )
19 | mock_agent.assert_called_once_with(
20 | name="any_agent",
21 | instructions=None,
22 | model=mock_model(model="mistral:mistral-small-latest"),
23 | tools=[],
24 | )
25 |
26 |
27 | def test_load_agno_agent_missing() -> None:
28 | with patch("any_agent.frameworks.agno.agno_available", False):
29 | with pytest.raises(ImportError):
30 | AnyAgent.create(
31 | AgentFramework.AGNO,
32 | AgentConfig(model_id="mistral:mistral-small-latest"),
33 | )
34 |
35 |
36 | def test_run_agno_custom_args() -> None:
37 | mock_agent = MagicMock()
38 | # Create a mock response object with the required content attribute
39 | mock_response = MagicMock()
40 | mock_response.content = "mock response"
41 |
42 | # Set up the AsyncMock to return the mock response
43 | mock_agent_instance = AsyncMock()
44 | mock_agent_instance.arun.return_value = mock_response
45 | mock_agent.return_value = mock_agent_instance
46 |
47 | mock_model = MagicMock()
48 |
49 | with (
50 | patch("any_agent.frameworks.agno.Agent", mock_agent),
51 | patch("any_agent.frameworks.agno.DEFAULT_MODEL_TYPE", mock_model),
52 | ):
53 | agent = AnyAgent.create(
54 | AgentFramework.AGNO, AgentConfig(model_id="mistral:mistral-small-latest")
55 | )
56 | result = agent.run("foo", retries=2)
57 |
58 | # Verify the result is as expected
59 | assert isinstance(result.final_output, str)
60 | assert result.final_output == "mock response"
61 |
62 | # Verify the agent was called with the right parameters
63 | mock_agent_instance.arun.assert_called_once_with("foo", retries=2)
64 |
--------------------------------------------------------------------------------
/src/any_agent/callbacks/span_generation/smolagents.py:
--------------------------------------------------------------------------------
1 | # mypy: disable-error-code="no-untyped-def,union-attr"
2 | from __future__ import annotations
3 |
4 | from typing import TYPE_CHECKING, Any
5 |
6 | from any_agent.callbacks.span_generation.base import _SpanGeneration
7 |
8 | if TYPE_CHECKING:
9 | from smolagents.models import ChatMessage
10 | from smolagents.tools import Tool
11 |
12 | from any_agent.callbacks.context import Context
13 |
14 |
15 | class _SmolagentsSpanGeneration(_SpanGeneration):
16 | def before_llm_call(self, context: Context, *args, **kwargs) -> Context:
17 | model_id = context.shared["model_id"]
18 |
19 | messages: list[ChatMessage] = args[0]
20 | input_messages = [
21 | {
22 | "role": message.role.value,
23 | "content": message.content[0]["text"], # type: ignore[index]
24 | }
25 | for message in messages
26 | if message.content
27 | ]
28 |
29 | return self._set_llm_input(context, model_id, input_messages)
30 |
31 | def after_llm_call(self, context: Context, *args, **kwargs) -> Context:
32 | response: ChatMessage = args[0]
33 | output: str | list[dict[str, Any]]
34 | if content := response.content:
35 | output = str(content)
36 | elif tool_calls := response.tool_calls:
37 | output = [
38 | {
39 | "tool.name": tool_call.function.name,
40 | "tool.args": tool_call.function.arguments,
41 | }
42 | for tool_call in tool_calls
43 | ]
44 |
45 | input_tokens = 0
46 | output_tokens = 0
47 | if raw := response.raw:
48 | if token_usage := raw.get("usage", None):
49 | input_tokens = token_usage.get("prompt_tokens", None)
50 | output_tokens = token_usage.get("completion_tokens", None)
51 |
52 | return self._set_llm_output(context, output, input_tokens, output_tokens)
53 |
54 | def before_tool_execution(self, context: Context, *args, **kwargs) -> Context:
55 | tool: Tool = context.shared["original_tool"]
56 |
57 | return self._set_tool_input(
58 | context, name=tool.name, description=tool.description, args=kwargs
59 | )
60 |
61 | def after_tool_execution(self, context: Context, *args, **kwargs) -> Context:
62 | return self._set_tool_output(context, args[0])
63 |
--------------------------------------------------------------------------------
/tests/cookbooks/test_cookbooks.py:
--------------------------------------------------------------------------------
1 | import os
2 | import pathlib
3 | import subprocess
4 |
5 | import pytest
6 |
7 |
8 | @pytest.mark.parametrize(
9 | "notebook_path",
10 | list(pathlib.Path("docs/cookbook").glob("*.ipynb")),
11 | ids=lambda x: x.stem,
12 | )
13 | @pytest.mark.timeout(180)
14 | def test_cookbook_notebook(
15 | notebook_path: pathlib.Path, capsys: pytest.CaptureFixture
16 | ) -> None:
17 | """Test that cookbook notebooks execute without errors using jupyter execute."""
18 | if notebook_path.stem == "mcp_agent":
19 | pytest.skip("See https://github.com/mozilla-ai/any-agent/issues/706")
20 | try:
21 | result = subprocess.run( # noqa: S603
22 | [ # noqa: S607
23 | "jupyter",
24 | "execute",
25 | notebook_path.name,
26 | "--allow-errors",
27 | "--output",
28 | f"executed_{notebook_path.name}",
29 | ],
30 | cwd="docs/cookbook", # Run in cookbook directory like original action
31 | env={
32 | "MISTRAL_API_KEY": os.environ["MISTRAL_API_KEY"],
33 | "TAVILY_API_KEY": os.environ["TAVILY_API_KEY"],
34 | "OPENAI_API_KEY": os.environ["OPENAI_API_KEY"],
35 | "HF_TOKEN": os.environ["HF_TOKEN"],
36 | "HF_ENDPOINT": os.environ["HF_ENDPOINT"],
37 | "PATH": os.environ["PATH"],
38 | "IN_PYTEST": "1", # Signal local_llm and mcp_agent notebooks that we are running an automated test.
39 | },
40 | timeout=170, # Time out slightly earlier so that we can log the output.
41 | capture_output=True,
42 | check=False,
43 | )
44 | except subprocess.TimeoutExpired as e:
45 | # Handle timeout case - log stdout/stderr that were captured before timeout
46 | stdout = e.stdout.decode() if e.stdout else "(no stdout captured)"
47 | stderr = e.stderr.decode() if e.stderr else "(no stderr captured)"
48 | pytest.fail(
49 | f"Notebook {notebook_path.name} timed out after 2 minutes\n stdout: {stdout}\n stderr: {stderr}"
50 | )
51 |
52 | if result.returncode != 0:
53 | stdout = result.stdout.decode() if result.stdout else "(no stdout captured)"
54 | stderr = result.stderr.decode() if result.stderr else "(no stderr captured)"
55 | pytest.fail(
56 | f"Notebook {notebook_path.name} failed with return code {result.returncode}\n stdout: {stdout}\n stderr: {stderr}"
57 | )
58 |
--------------------------------------------------------------------------------
/src/any_agent/evaluation/tools.py:
--------------------------------------------------------------------------------
1 | from collections.abc import Callable
2 | from typing import Any
3 |
4 | from pydantic import BaseModel
5 |
6 | from any_agent.tracing.agent_trace import AgentTrace
7 |
8 |
9 | class TraceTools:
10 | def __init__(self, trace: AgentTrace):
11 | self.trace = trace
12 |
13 | def get_all_tools(self) -> list[Callable[..., Any]]:
14 | """Get all tool functions from this class.
15 |
16 | Returns:
17 | list[callable]: List of all tool functions
18 |
19 | """
20 | # Get all methods that don't start with underscore and aren't get_all_tools
21 | tools = []
22 | for attr_name in dir(self):
23 | if not attr_name.startswith("_") and attr_name != "get_all_tools":
24 | attr = getattr(self, attr_name)
25 | if callable(attr) and attr_name not in ["trace"]:
26 | tools.append(attr)
27 | return tools
28 |
29 | def get_final_output(self) -> str | BaseModel | dict[str, Any] | None:
30 | """Get the final output from the agent trace.
31 |
32 | Returns:
33 | str | BaseModel | None: The final output of the agent
34 |
35 | """
36 | return self.trace.final_output
37 |
38 | def get_tokens_used(self) -> int:
39 | """Get the number of tokens used by the agent as reported by the trace.
40 |
41 | Returns:
42 | int: The number of tokens used by the agent
43 |
44 | """
45 | return self.trace.tokens.total_tokens
46 |
47 | def get_steps_taken(self) -> int:
48 | """Get the number of steps taken by the agent as reported by the trace.
49 |
50 | Returns:
51 | int: The number of steps taken by the agent
52 |
53 | """
54 | return len(self.trace.spans)
55 |
56 | def get_messages_from_trace(self) -> str:
57 | """Get a summary of what happened in each step/span of the agent trace.
58 |
59 | This includes information about the input, output, and tool calls for each step.
60 |
61 | Returns:
62 | str: The evidence of all the spans in the trace
63 |
64 | """
65 | messages = self.trace.spans_to_messages()
66 | evidence = ""
67 | for message in messages:
68 | evidence += f"### {message.role}\n{message.content}\n\n"
69 | return evidence
70 |
71 | def get_duration(self) -> float:
72 | """Get the duration of the agent trace.
73 |
74 | Returns:
75 | float: The duration in seconds of the agent trace
76 |
77 | """
78 | return self.trace.duration.total_seconds()
79 |
--------------------------------------------------------------------------------
/src/any_agent/callbacks/span_generation/openai.py:
--------------------------------------------------------------------------------
1 | # mypy: disable-error-code="no-untyped-def"
2 | from __future__ import annotations
3 |
4 | from typing import TYPE_CHECKING, Any
5 |
6 | from any_agent.callbacks.span_generation.base import _SpanGeneration
7 |
8 | if TYPE_CHECKING:
9 | from agents import FunctionTool, ModelResponse
10 |
11 | from any_agent.callbacks.context import Context
12 |
13 |
14 | class _OpenAIAgentsSpanGeneration(_SpanGeneration):
15 | def before_llm_call(self, context: Context, *args, **kwargs) -> Context:
16 | model_id = context.shared["model_id"]
17 |
18 | user_input = kwargs.get("input", ["No input"])[0]
19 | system_instructions = kwargs.get("system_instructions")
20 | input_messages = [
21 | {"role": "system", "content": system_instructions},
22 | user_input,
23 | ]
24 | return self._set_llm_input(context, model_id, input_messages)
25 |
26 | def after_llm_call(self, context: Context, *args, **kwargs) -> Context:
27 | from openai.types.responses import (
28 | ResponseFunctionToolCall,
29 | ResponseOutputMessage,
30 | ResponseOutputText,
31 | )
32 |
33 | response: ModelResponse = args[0]
34 | if not response.output:
35 | return context
36 |
37 | output: str | list[dict[str, Any]] = ""
38 | if isinstance(response.output[0], ResponseFunctionToolCall):
39 | output = [
40 | {
41 | "tool.name": response.output[0].name,
42 | "tool.args": response.output[0].arguments,
43 | }
44 | ]
45 | elif isinstance(response.output[0], ResponseOutputMessage):
46 | if content := response.output[0].content:
47 | if isinstance(content[0], ResponseOutputText):
48 | output = content[0].text
49 |
50 | input_tokens = 0
51 | output_tokens = 0
52 | if token_usage := response.usage:
53 | input_tokens = token_usage.input_tokens
54 | output_tokens = token_usage.output_tokens
55 |
56 | return self._set_llm_output(context, output, input_tokens, output_tokens)
57 |
58 | def before_tool_execution(self, context: Context, *args, **kwargs) -> Context:
59 | tool: FunctionTool = context.shared["original_tool"]
60 |
61 | return self._set_tool_input(
62 | context, name=tool.name, description=tool.description, args=args[1]
63 | )
64 |
65 | def after_tool_execution(self, context: Context, *args, **kwargs) -> Context:
66 | return self._set_tool_output(context, args[0])
67 |
--------------------------------------------------------------------------------
/tests/integration/mcp/test_mcp_streamable_http.py:
--------------------------------------------------------------------------------
1 | import json
2 | from datetime import datetime
3 | from pathlib import Path
4 | from typing import Any
5 |
6 | import pytest
7 | from pydantic import BaseModel, ConfigDict
8 |
9 | from any_agent import (
10 | AgentConfig,
11 | AgentFramework,
12 | AnyAgent,
13 | )
14 | from any_agent.config import MCPStreamableHttp
15 | from any_agent.testing.helpers import (
16 | DEFAULT_SMALL_MODEL_ID,
17 | get_default_agent_model_args,
18 | group_spans,
19 | )
20 |
21 |
22 | class Step(BaseModel):
23 | number: int
24 | description: str
25 |
26 |
27 | class Steps(BaseModel):
28 | model_config = ConfigDict(extra="forbid")
29 | steps: list[Step]
30 |
31 |
32 | def test_load_and_run_agent_streamable_http(
33 | agent_framework: AgentFramework,
34 | date_streamable_http_server: dict[str, Any],
35 | tmp_path: Path,
36 | ) -> None:
37 | kwargs = {}
38 |
39 | kwargs["model_id"] = DEFAULT_SMALL_MODEL_ID
40 |
41 | tmp_file = "tmp.txt"
42 |
43 | def write_file(text: str) -> None:
44 | """write the text to a file in the tmp_path directory
45 |
46 | Args:
47 | text (str): The text to write to the file.
48 |
49 | Returns:
50 | None
51 | """
52 | (tmp_path / tmp_file).write_text(text)
53 |
54 | tools = [
55 | write_file,
56 | MCPStreamableHttp(
57 | url=date_streamable_http_server["url"],
58 | client_session_timeout_seconds=30,
59 | ),
60 | ]
61 | agent_config = AgentConfig(
62 | tools=tools, # type: ignore[arg-type]
63 | instructions="Use the available tools to answer.",
64 | model_args=get_default_agent_model_args(
65 | agent_framework, model_id=kwargs["model_id"]
66 | ),
67 | output_type=Steps,
68 | **kwargs, # type: ignore[arg-type]
69 | )
70 | agent = AnyAgent.create(agent_framework, agent_config)
71 |
72 | agent_trace = agent.run(
73 | "First, find what year it is in the America/New_York timezone. "
74 | "Then, write the value (single number) to a file. "
75 | "Finally, return a list of the steps you have taken.",
76 | )
77 |
78 | assert isinstance(agent_trace.final_output, Steps)
79 |
80 | assert (tmp_path / tmp_file).read_text() == str(datetime.now().year)
81 | _, _, tool_executions = group_spans(agent_trace.spans)
82 |
83 | assert len(tool_executions) >= 1
84 | tool_args_raw = tool_executions[0].attributes.get("gen_ai.tool.args")
85 | assert tool_args_raw is not None
86 | args = json.loads(tool_args_raw)
87 | assert "timezone" in args
88 |
--------------------------------------------------------------------------------
/tests/docs/test_all.py:
--------------------------------------------------------------------------------
1 | import pathlib
2 | from typing import Any
3 | from unittest.mock import AsyncMock, MagicMock, patch
4 |
5 | import pytest
6 | from mktestdocs import check_md_file
7 |
8 |
9 | # Note the use of `str`, makes for pretty output
10 | # Exclude any files that you have custom mocking for.
11 | @pytest.mark.parametrize(
12 | "fpath",
13 | [f for f in pathlib.Path("docs").glob("**/*.md") if f.name != "evaluation.md"],
14 | ids=str,
15 | )
16 | def test_files_all(fpath: pathlib.Path) -> None:
17 | if fpath.name == "serving.md":
18 | # the serving markdown runs multiple servers in different processes
19 | # which is not supported by this testing.
20 | pytest.skip("Serving.md not supported by docs tester")
21 |
22 | mock_agent = MagicMock()
23 | mock_create = MagicMock(return_value=mock_agent)
24 | mock_a2a_tool = AsyncMock()
25 |
26 | mock_create_async = AsyncMock()
27 | with (
28 | patch("builtins.open", new_callable=MagicMock),
29 | patch("any_agent.AnyAgent.create", mock_create),
30 | patch("any_agent.AnyAgent.create_async", mock_create_async),
31 | patch("any_agent.tools.a2a_tool_async", mock_a2a_tool),
32 | patch("composio.Composio", MagicMock()),
33 | ):
34 | check_md_file(fpath=fpath, memory=True)
35 |
36 |
37 | def test_evaluation_md() -> None:
38 | mock_trace = MagicMock()
39 | mock_trace.tokens.total_tokens = 500
40 | mock_trace.spans = [MagicMock()]
41 | mock_trace.final_output = "Paris"
42 | mock_trace.spans_to_messages.return_value = []
43 |
44 | mock_agent = MagicMock()
45 | mock_agent.run.return_value = mock_trace
46 | mock_create = MagicMock(return_value=mock_agent)
47 |
48 | def mock_run_method(*args: Any, **kwargs: Any) -> Any:
49 | mock_result = MagicMock()
50 | mock_result.passed = True
51 | mock_result.reasoning = "Mock evaluation result"
52 | mock_result.confidence_score = 0.95
53 | mock_result.suggestions = ["Mock suggestion 1", "Mock suggestion 2"]
54 | return mock_result
55 |
56 | mock_judge = MagicMock()
57 | mock_judge.run.side_effect = mock_run_method
58 |
59 | mock_create_async = AsyncMock()
60 | with (
61 | patch("builtins.open", new_callable=MagicMock),
62 | patch("any_agent.AnyAgent.create", mock_create),
63 | patch("any_agent.AnyAgent.create_async", mock_create_async),
64 | patch("any_agent.evaluation.LlmJudge", return_value=mock_judge),
65 | patch("any_agent.evaluation.AgentJudge", return_value=mock_judge),
66 | ):
67 | check_md_file(fpath=pathlib.Path("docs/evaluation.md"), memory=True)
68 |
--------------------------------------------------------------------------------
/src/any_agent/tracing/attributes.py:
--------------------------------------------------------------------------------
1 | from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import (
2 | GEN_AI_AGENT_DESCRIPTION,
3 | GEN_AI_AGENT_NAME,
4 | GEN_AI_OPERATION_NAME,
5 | GEN_AI_OUTPUT_TYPE,
6 | GEN_AI_REQUEST_MODEL,
7 | GEN_AI_TOOL_DESCRIPTION,
8 | GEN_AI_TOOL_NAME,
9 | GEN_AI_USAGE_INPUT_TOKENS,
10 | GEN_AI_USAGE_OUTPUT_TOKENS,
11 | )
12 |
13 |
14 | class GenAI:
15 | """Constants exported for convenience to access span attributes.
16 |
17 | Trying to follow OpenTelemetry's [Semantic Conventions for Generative AI](https://opentelemetry.io/docs/specs/semconv/gen-ai/).
18 |
19 | We import the constants from `opentelemetry.semconv._incubating.attributes.gen_ai_attributes`
20 | whenever is possible.
21 |
22 | We only expose the keys that we currently use in `any-agent`.
23 | """
24 |
25 | AGENT_DESCRIPTION = GEN_AI_AGENT_DESCRIPTION
26 | """Free-form description of the GenAI agent provided by the application."""
27 |
28 | AGENT_NAME = GEN_AI_AGENT_NAME
29 | """Human-readable name of the GenAI agent provided by the application."""
30 |
31 | INPUT_MESSAGES = "gen_ai.input.messages"
32 | """System prompt and user input."""
33 |
34 | OPERATION_NAME = GEN_AI_OPERATION_NAME
35 | """The name of the operation being performed."""
36 |
37 | OUTPUT = "gen_ai.output"
38 | """Used in both LLM Calls and Tool Executions for holding their respective outputs."""
39 |
40 | OUTPUT_TYPE = GEN_AI_OUTPUT_TYPE
41 | """Represents the content type requested by the client."""
42 |
43 | REQUEST_MODEL = GEN_AI_REQUEST_MODEL
44 | """The name of the GenAI model a request is being made to."""
45 |
46 | TOOL_ARGS = "gen_ai.tool.args"
47 | """Arguments passed to the executed tool."""
48 |
49 | TOOL_DESCRIPTION = GEN_AI_TOOL_DESCRIPTION
50 | """The tool description."""
51 |
52 | TOOL_NAME = GEN_AI_TOOL_NAME
53 | """Name of the tool utilized by the agent."""
54 |
55 | USAGE_INPUT_COST = "gen_ai.usage.input_cost"
56 | """Dollars spent for the input of the LLM."""
57 |
58 | USAGE_INPUT_TOKENS = GEN_AI_USAGE_INPUT_TOKENS
59 | """The number of tokens used in the GenAI input (prompt)."""
60 |
61 | USAGE_OUTPUT_COST = "gen_ai.usage.output_cost"
62 | """Dollars spent for the output of the LLM."""
63 |
64 | USAGE_OUTPUT_TOKENS = GEN_AI_USAGE_OUTPUT_TOKENS
65 | """The number of tokens used in the GenAI response (completion)."""
66 |
67 |
68 | class AnyAgentAttributes:
69 | """Span-attribute keys specific to AnyAgent library."""
70 |
71 | VERSION = "any_agent.version"
72 | """The any-agent library version used in the runtime."""
73 |
--------------------------------------------------------------------------------
/tests/unit/frameworks/test_llama_index.py:
--------------------------------------------------------------------------------
1 | from unittest.mock import AsyncMock, MagicMock, patch
2 |
3 | import pytest
4 |
5 | from any_agent import AgentConfig, AgentFramework, AnyAgent
6 |
7 |
8 | def test_load_llama_index_agent_default() -> None:
9 | model_mock = MagicMock()
10 | create_mock = MagicMock()
11 | agent_mock = MagicMock()
12 | create_mock.return_value = agent_mock
13 | tool_mock = MagicMock()
14 | from llama_index.core.tools import FunctionTool
15 |
16 | with (
17 | patch("any_agent.frameworks.llama_index.DEFAULT_AGENT_TYPE", create_mock),
18 | patch("any_agent.frameworks.llama_index.DEFAULT_MODEL_TYPE", model_mock),
19 | patch.object(FunctionTool, "from_defaults", tool_mock),
20 | ):
21 | AnyAgent.create(
22 | AgentFramework.LLAMA_INDEX,
23 | AgentConfig(
24 | model_id="gemini/gemini-2.0-flash",
25 | instructions="You are a helpful assistant",
26 | ),
27 | )
28 |
29 | model_mock.assert_called_once_with(
30 | model="gemini/gemini-2.0-flash",
31 | api_key=None,
32 | api_base=None,
33 | additional_kwargs={},
34 | )
35 | create_mock.assert_called_once_with(
36 | name="any_agent",
37 | llm=model_mock.return_value,
38 | system_prompt="You are a helpful assistant",
39 | description="The main agent",
40 | tools=[],
41 | )
42 |
43 |
44 | def test_load_llama_index_agent_missing() -> None:
45 | with patch("any_agent.frameworks.llama_index.llama_index_available", False):
46 | with pytest.raises(ImportError):
47 | AnyAgent.create(
48 | AgentFramework.LLAMA_INDEX,
49 | AgentConfig(model_id="mistral:mistral-small-latest"),
50 | )
51 |
52 |
53 | def test_run_llama_index_agent_custom_args() -> None:
54 | create_mock = MagicMock()
55 | agent_mock = AsyncMock()
56 | create_mock.return_value = agent_mock
57 | from llama_index.core.tools import FunctionTool
58 |
59 | with (
60 | patch("any_agent.frameworks.llama_index.DEFAULT_AGENT_TYPE", create_mock),
61 | patch("any_agent.frameworks.llama_index.DEFAULT_MODEL_TYPE"),
62 | patch.object(FunctionTool, "from_defaults"),
63 | ):
64 | agent = AnyAgent.create(
65 | AgentFramework.LLAMA_INDEX,
66 | AgentConfig(
67 | model_id="gemini/gemini-2.0-flash",
68 | instructions="You are a helpful assistant",
69 | ),
70 | )
71 | agent.run("foo", timeout=10)
72 | agent_mock.run.assert_called_once_with("foo", timeout=10)
73 |
--------------------------------------------------------------------------------
/src/any_agent/callbacks/span_generation/tinyagent.py:
--------------------------------------------------------------------------------
1 | # mypy: disable-error-code="method-assign,no-untyped-def"
2 | from __future__ import annotations
3 |
4 | from typing import TYPE_CHECKING, Any
5 |
6 | from any_agent.callbacks.span_generation.base import _SpanGeneration
7 |
8 | if TYPE_CHECKING:
9 | from any_llm.types.completion import (
10 | ChatCompletion,
11 | ChatCompletionMessageToolCall,
12 | CompletionUsage,
13 | )
14 |
15 | from any_agent.callbacks.context import Context
16 |
17 |
18 | class _TinyAgentSpanGeneration(_SpanGeneration):
19 | def before_llm_call(self, context: Context, *args, **kwargs) -> Context:
20 | return self._set_llm_input(
21 | context,
22 | model_id=kwargs.get("model", "No model"),
23 | input_messages=kwargs.get("messages", []),
24 | )
25 |
26 | def after_llm_call(self, context: Context, *args, **kwargs) -> Context:
27 | response: ChatCompletion = args[0]
28 |
29 | if not response.choices:
30 | return context
31 |
32 | message = getattr(response.choices[0], "message", None)
33 | if not message:
34 | return context
35 |
36 | output: str | list[dict[str, str]] = ""
37 | if content := getattr(message, "content", None):
38 | output = content
39 |
40 | tool_calls: list[ChatCompletionMessageToolCall] | None
41 | if tool_calls := getattr(message, "tool_calls", None):
42 | output = [
43 | {
44 | "tool.name": getattr(tool_call.function, "name", "No name"),
45 | "tool.args": getattr(tool_call.function, "arguments", "No name"),
46 | }
47 | for tool_call in tool_calls
48 | if tool_call.function
49 | ]
50 |
51 | input_tokens = 0
52 | output_tokens = 0
53 | token_usage: CompletionUsage | None
54 |
55 | if token_usage := getattr(response, "usage", None):
56 | if token_usage:
57 | input_tokens = token_usage.prompt_tokens
58 | output_tokens = token_usage.completion_tokens
59 |
60 | return self._set_llm_output(context, output, input_tokens, output_tokens)
61 |
62 | def before_tool_execution(self, context: Context, *args, **kwargs) -> Context:
63 | request: dict[str, Any] = args[0]
64 |
65 | return self._set_tool_input(
66 | context,
67 | name=request.get("name", "No name"),
68 | args=request.get("arguments", {}),
69 | )
70 |
71 | def after_tool_execution(self, context: Context, *args, **kwargs) -> Context:
72 | return self._set_tool_output(context, args[0])
73 |
--------------------------------------------------------------------------------
/src/any_agent/callbacks/wrappers/agno.py:
--------------------------------------------------------------------------------
1 | # mypy: disable-error-code="method-assign,no-untyped-def,union-attr"
2 | from __future__ import annotations
3 |
4 | from typing import TYPE_CHECKING, Any
5 |
6 | from opentelemetry.trace import get_current_span
7 |
8 | if TYPE_CHECKING:
9 | from any_agent.callbacks.context import Context
10 | from any_agent.frameworks.agno import AgnoAgent
11 |
12 |
13 | class _AgnoWrapper:
14 | def __init__(self) -> None:
15 | self.callback_context: dict[int, Context] = {}
16 | self._original_aprocess_model: Any = None
17 | self._original_arun_function_call: Any = None
18 |
19 | async def wrap(self, agent: AgnoAgent) -> None:
20 | self._original_aprocess_model = agent._agent.model._aprocess_model_response
21 |
22 | async def wrapped_llm_call(*args, **kwargs):
23 | context = self.callback_context[
24 | get_current_span().get_span_context().trace_id
25 | ]
26 | context.shared["model_id"] = agent._agent.model.id
27 |
28 | for callback in agent.config.callbacks:
29 | context = callback.before_llm_call(context, *args, **kwargs)
30 |
31 | result = await self._original_aprocess_model(*args, **kwargs)
32 |
33 | for callback in agent.config.callbacks:
34 | context = callback.after_llm_call(context, result, *args, **kwargs)
35 |
36 | return result
37 |
38 | agent._agent.model._aprocess_model_response = wrapped_llm_call
39 |
40 | self._original_arun_function_call = agent._agent.model.arun_function_call
41 |
42 | async def wrapped_tool_execution(
43 | *args,
44 | **kwargs,
45 | ):
46 | context = self.callback_context[
47 | get_current_span().get_span_context().trace_id
48 | ]
49 |
50 | for callback in agent.config.callbacks:
51 | context = callback.before_tool_execution(context, *args, **kwargs)
52 |
53 | result = await self._original_arun_function_call(*args, **kwargs)
54 |
55 | for callback in agent.config.callbacks:
56 | context = callback.after_tool_execution(
57 | context, result, *args, **kwargs
58 | )
59 |
60 | return result
61 |
62 | agent._agent.model.arun_function_call = wrapped_tool_execution
63 |
64 | async def unwrap(self, agent: AgnoAgent):
65 | if self._original_aprocess_model is not None:
66 | agent._agent.model._aprocess_model_response = self._original_aprocess_model
67 | if self._original_arun_function_call is not None:
68 | agent._agent.model.arun_function_calls = self._original_arun_function_call
69 |
--------------------------------------------------------------------------------
/src/any_agent/serving/server_handle.py:
--------------------------------------------------------------------------------
1 | from __future__ import annotations
2 |
3 | import asyncio
4 | from dataclasses import dataclass
5 | from typing import TYPE_CHECKING
6 |
7 | from any_agent.logging import logger
8 |
9 | if TYPE_CHECKING:
10 | from typing import Any
11 |
12 | from uvicorn import Server as UvicornServer
13 |
14 |
15 | @dataclass
16 | class ServerHandle:
17 | """A handle for managing an async server instance.
18 |
19 | This class provides a clean interface for managing the lifecycle of a server
20 | without requiring manual management of the underlying task and server objects.
21 | """
22 |
23 | task: asyncio.Task[Any]
24 | server: UvicornServer
25 |
26 | async def shutdown(self, timeout_seconds: float = 10.0) -> None:
27 | """Gracefully shutdown the server with a timeout.
28 |
29 | Args:
30 | timeout_seconds: Maximum time to wait for graceful shutdown before forcing cancellation.
31 |
32 | """
33 | if not self.is_running():
34 | return # Already shut down
35 |
36 | self.server.should_exit = True
37 | try:
38 | await asyncio.wait_for(self.task, timeout=timeout_seconds)
39 | except TimeoutError:
40 | logger.warning(
41 | "Server shutdown timed out after %ss, forcing cancellation",
42 | timeout_seconds,
43 | )
44 | self.task.cancel()
45 | try:
46 | await self.task
47 | except asyncio.CancelledError:
48 | pass
49 | except Exception as e:
50 | logger.error("Error during server shutdown: %s", e)
51 | # Still try to cancel the task to clean up
52 | if not self.task.done():
53 | self.task.cancel()
54 | try:
55 | await self.task
56 | except asyncio.CancelledError:
57 | pass
58 |
59 | def is_running(self) -> bool:
60 | """Check if the server is still running.
61 |
62 | Returns:
63 | True if the server task is still running, False otherwise.
64 |
65 | """
66 | return not self.task.done()
67 |
68 | @property
69 | def port(self) -> int:
70 | """Get the port the server is running on.
71 |
72 | If the server port was specified as 0, the port will be the one assigned by the OS.
73 | This helper method is useful to get the actual port that the server is running on.
74 |
75 | Returns:
76 | The port number the server is running on.
77 |
78 | """
79 | port = self.server.servers[0].sockets[0].getsockname()[1]
80 | assert port is not None
81 | assert isinstance(port, int)
82 | return port
83 |
--------------------------------------------------------------------------------
/tests/integration/conftest.py:
--------------------------------------------------------------------------------
1 | import hashlib
2 | import socket
3 |
4 | import pytest
5 |
6 | from any_agent.config import AgentFramework
7 |
8 |
9 | def pytest_addoption(parser: pytest.Parser) -> None:
10 | """
11 | Add custom command-line options to pytest.
12 |
13 | This hook adds the `--update-trace-assets` flag to pytest, which can be used when running integration tests.
14 | When this flag is set, tests that generate trace asset files (aka the integration test that
15 | produces agent traces) will update the asset files in the assets directory.
16 | This is useful when the expected trace output changes and you
17 | want to regenerate the reference files.
18 | """
19 | parser.addoption(
20 | "--update-trace-assets",
21 | action="store_true",
22 | default=False,
23 | help="Update trace asset files instead of asserting equality.",
24 | )
25 |
26 |
27 | def _is_port_available(port: int, host: str = "localhost") -> bool:
28 | """Check if a port is available for binding.
29 |
30 | This isn't a perfect check but it at least tells us if there is absolutely no chance of binding to the port.
31 |
32 | """
33 | with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
34 | try:
35 | sock.bind((host, port))
36 | except OSError:
37 | return False
38 | return True
39 |
40 |
41 | def _get_deterministic_port(test_name: str, framework_name: str) -> int:
42 | """Generate a deterministic port number based on test name and framework.
43 |
44 | This ensures each test gets a unique port that remains consistent across runs.
45 | """
46 | # Create a unique string by combining test name and framework
47 | unique_string = f"{test_name}_{framework_name}"
48 |
49 | # Generate a hash and convert to a port number in the range 6000-9999
50 | hash_value = int(hashlib.md5(unique_string.encode()).hexdigest()[:4], 16) # noqa: S324
51 | port = 6000 + (hash_value % 4000)
52 |
53 | # Ensure the port is available, if not, try nearby ports
54 | original_port = port
55 | attempts = 0
56 | while not _is_port_available(port) and attempts < 50:
57 | port = original_port + attempts + 1
58 | attempts += 1
59 |
60 | if not _is_port_available(port):
61 | msg = f"Could not find an available port starting from {original_port}"
62 | raise RuntimeError(msg)
63 |
64 | return port
65 |
66 |
67 | @pytest.fixture
68 | def test_port(request: pytest.FixtureRequest, agent_framework: AgentFramework) -> int:
69 | """Single fixture that provides a unique, deterministic port for each test."""
70 | test_name = request.node.name
71 | framework_name = agent_framework.value
72 |
73 | return _get_deterministic_port(test_name, framework_name)
74 |
--------------------------------------------------------------------------------
/src/any_agent/callbacks/span_generation/agno.py:
--------------------------------------------------------------------------------
1 | # mypy: disable-error-code="method-assign,no-untyped-def,union-attr"
2 | from __future__ import annotations
3 |
4 | from typing import TYPE_CHECKING, Any
5 |
6 | from any_agent.callbacks.span_generation.base import _SpanGeneration
7 |
8 | if TYPE_CHECKING:
9 | from agno.models.message import Message, MessageMetrics
10 | from agno.tools.function import FunctionCall
11 |
12 | from any_agent.callbacks.context import Context
13 |
14 |
15 | class _AgnoSpanGeneration(_SpanGeneration):
16 | def before_llm_call(self, context: Context, *args, **kwargs):
17 | messages: list[Message] = kwargs.get("messages", [])
18 | input_messages = [
19 | {"role": message.role, "content": str(message.content)}
20 | for message in messages
21 | ]
22 | return self._set_llm_input(context, context.shared["model_id"], input_messages)
23 |
24 | def after_llm_call(self, context: Context, *args, **kwargs) -> Context:
25 | output: str | list[dict[str, Any]] = ""
26 | if assistant_message := kwargs.get("assistant_message"):
27 | if content := getattr(assistant_message, "content", None):
28 | output = str(content)
29 | if tool_calls := getattr(assistant_message, "tool_calls", None):
30 | output = [
31 | {
32 | "tool.name": tool.get("function", {}).get("name", "No name"),
33 | "tool.args": tool.get("function", {}).get(
34 | "arguments", "No args"
35 | ),
36 | }
37 | for tool in tool_calls
38 | ]
39 |
40 | metrics: MessageMetrics | None
41 | input_tokens: int = 0
42 | output_tokens: int = 0
43 | if metrics := getattr(assistant_message, "metrics", None):
44 | input_tokens = metrics.input_tokens
45 | output_tokens = metrics.output_tokens
46 |
47 | context = self._set_llm_output(context, output, input_tokens, output_tokens)
48 |
49 | return context
50 |
51 | def before_tool_execution(self, context: Context, *args, **kwargs) -> Context:
52 | function_call: FunctionCall = args[0]
53 | function = function_call.function
54 |
55 | return self._set_tool_input(
56 | context,
57 | name=function.name,
58 | description=function.description,
59 | args=function_call.arguments,
60 | call_id=function_call.call_id,
61 | )
62 |
63 | def after_tool_execution(self, context: Context, *args, **kwargs) -> Context:
64 | function_call: FunctionCall = args[1]
65 | return self._set_tool_output(context, function_call.result)
66 |
--------------------------------------------------------------------------------
/src/any_agent/tools/final_output.py:
--------------------------------------------------------------------------------
1 | import json
2 | from collections.abc import Callable
3 | from typing import Any
4 |
5 | from pydantic import BaseModel, ValidationError
6 |
7 |
8 | def prepare_final_output(
9 | output_type: type[BaseModel], instructions: str | None = None
10 | ) -> tuple[str, Callable[[str], dict[str, str | bool | dict[str, Any] | list[Any]]]]:
11 | """Prepare instructions and tools for structured output, returning the function directly.
12 |
13 | Args:
14 | output_type: The Pydantic model type for structured output
15 | instructions: Original instructions to modify
16 |
17 | Returns:
18 | Tuple of (modified_instructions, final_output_function)
19 |
20 | """
21 | tool_name = "final_output"
22 | modified_instructions = instructions or ""
23 | modified_instructions += (
24 | f"You must call the {tool_name} tool when finished."
25 | f"The 'answer' argument passed to the {tool_name} tool must be a JSON string that matches the following schema:\n"
26 | f"{output_type.model_json_schema()}"
27 | )
28 |
29 | def final_output_tool(
30 | answer: str,
31 | ) -> dict[str, str | bool | dict[str, Any] | list[Any]]:
32 | # First check if it's valid JSON
33 | try:
34 | parsed_answer = json.loads(answer)
35 | except json.JSONDecodeError as json_err:
36 | return {
37 | "success": False,
38 | "result": f"Invalid JSON format: {json_err}. Please fix the 'answer' parameter so that it is a valid JSON string and call this tool again.",
39 | }
40 | # Then validate against the Pydantic model
41 | try:
42 | output_type.model_validate_json(answer)
43 | except ValidationError as e:
44 | return {
45 | "success": False,
46 | "result": f"Please fix this validation error: {e}. The format must conform to {output_type.model_json_schema()}",
47 | }
48 | else:
49 | return {"success": True, "result": parsed_answer}
50 |
51 | # Set the function name and docstring
52 | final_output_tool.__name__ = tool_name
53 | final_output_tool.__doc__ = f"""This tool is used to validate the final output. It must be called when the final answer is ready in order to ensure that the output is valid.
54 |
55 | Args:
56 | answer: The final output that can be loaded as a Pydantic model. This must be a JSON compatible string that matches the following schema:
57 | {output_type.model_json_schema()}
58 |
59 | Returns:
60 | A dictionary with the following keys:
61 | - success: True if the output is valid, False otherwise.
62 | - result: The final output if success is True, otherwise an error message.
63 |
64 | """
65 |
66 | return modified_instructions, final_output_tool
67 |
--------------------------------------------------------------------------------
/src/any_agent/serving/a2a/envelope.py:
--------------------------------------------------------------------------------
1 | from __future__ import annotations
2 |
3 | from collections.abc import Mapping
4 | from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar
5 |
6 | from a2a.types import TaskState # noqa: TC002
7 | from pydantic import BaseModel, ConfigDict
8 |
9 | if TYPE_CHECKING:
10 | from any_agent import AnyAgent
11 |
12 |
13 | class _DefaultBody(BaseModel):
14 | """Default payload when the user does not supply one."""
15 |
16 | result: str
17 |
18 | model_config = ConfigDict(extra="forbid")
19 |
20 |
21 | # Define a TypeVar for the body type
22 | BodyType = TypeVar("BodyType", bound=BaseModel)
23 |
24 |
25 | class A2AEnvelope(BaseModel, Generic[BodyType]):
26 | """A2A envelope that wraps response data with task status."""
27 |
28 | task_status: Literal[
29 | TaskState.input_required, TaskState.completed, TaskState.failed
30 | ]
31 | """Restricted to the states that are leveraged by our implementation of the A2A protocol.
32 | When we support streaming, the rest of the states can be added and supported."""
33 |
34 | data: BodyType
35 |
36 | model_config = ConfigDict(extra="forbid")
37 |
38 |
39 | def _is_a2a_envelope(typ: type[BaseModel] | None) -> bool:
40 | if typ is None:
41 | return False
42 | fields: Any = getattr(typ, "model_fields", None)
43 |
44 | # We only care about a mapping with the required keys.
45 | if not isinstance(fields, Mapping):
46 | return False
47 |
48 | return "task_status" in fields and "data" in fields
49 |
50 |
51 | def _create_a2a_envelope(body_type: type[BaseModel]) -> type[A2AEnvelope[Any]]:
52 | """Return a *new* Pydantic model that wraps *body_type* with TaskState + data."""
53 | # Ensure body forbids extra keys (OpenAI response_format requirement)
54 | if hasattr(body_type, "model_config"):
55 | body_type.model_config["extra"] = "forbid"
56 | else:
57 | body_type.model_config = ConfigDict(extra="forbid")
58 |
59 | class EnvelopeInstance(A2AEnvelope[body_type]): # type: ignore[valid-type]
60 | pass
61 |
62 | EnvelopeInstance.__name__ = f"{body_type.__name__}Return"
63 | EnvelopeInstance.__qualname__ = f"{body_type.__qualname__}Return"
64 | return EnvelopeInstance
65 |
66 |
67 | async def prepare_agent_for_a2a_async(agent: AnyAgent) -> AnyAgent:
68 | """Async counterpart of :pyfunc:`prepare_agent_for_a2a`.
69 |
70 | This function preserves MCP servers from the original agent to avoid
71 | connection timeouts.
72 | """
73 | if _is_a2a_envelope(agent.config.output_type):
74 | return agent
75 |
76 | body_type = agent.config.output_type or _DefaultBody
77 | new_output_type = _create_a2a_envelope(body_type)
78 |
79 | # Use the new update_output_type method instead of recreating the agent
80 | await agent.update_output_type_async(new_output_type)
81 | return agent
82 |
--------------------------------------------------------------------------------
/tests/unit/tools/test_unit_web_browsing.py:
--------------------------------------------------------------------------------
1 | from typing import Any
2 | from unittest.mock import MagicMock, patch
3 |
4 |
5 | def test_search_tavily_no_api_key(monkeypatch: Any) -> None:
6 | with patch("tavily.tavily.TavilyClient", MagicMock()):
7 | monkeypatch.delenv("TAVILY_API_KEY", raising=False)
8 | from any_agent.tools import search_tavily
9 |
10 | result = search_tavily("test")
11 | assert "environment variable not set" in result
12 |
13 |
14 | def test_search_tavily_success(monkeypatch: Any) -> None:
15 | class FakeClient:
16 | def __init__(self, api_key: str) -> None:
17 | self.api_key = api_key
18 |
19 | def search(self, query: str, include_images: bool = False) -> Any:
20 | return {
21 | "results": [
22 | {
23 | "title": "Test Title",
24 | "url": "http://test.com",
25 | "content": "Test content!",
26 | }
27 | ]
28 | }
29 |
30 | with patch("tavily.tavily.TavilyClient", FakeClient):
31 | monkeypatch.setenv("TAVILY_API_KEY", "fake-key")
32 | from any_agent.tools import search_tavily
33 |
34 | result = search_tavily("test")
35 | assert "Test Title" in result
36 | assert "Test content!" in result
37 |
38 |
39 | def test_search_tavily_with_images(monkeypatch: Any) -> None:
40 | class FakeClient:
41 | def __init__(self, api_key: str) -> None:
42 | self.api_key = api_key
43 |
44 | def search(self, query: str, include_images: bool = False) -> Any:
45 | return {
46 | "results": [
47 | {
48 | "title": "Test Title",
49 | "url": "http://test.com",
50 | "content": "Test content!",
51 | }
52 | ],
53 | "images": ["http://image.com/cat.jpg"],
54 | }
55 |
56 | with patch("tavily.tavily.TavilyClient", FakeClient):
57 | monkeypatch.setenv("TAVILY_API_KEY", "fake-key")
58 | from any_agent.tools import search_tavily
59 |
60 | result = search_tavily("test", include_images=True)
61 | assert "Test Title" in result
62 | assert "Images:" in result
63 | assert "cat.jpg" in result
64 |
65 |
66 | def test_search_tavily_exception(monkeypatch: Any) -> None:
67 | class FakeClient:
68 | def __init__(self, api_key: str) -> None:
69 | self.api_key = api_key
70 |
71 | def search(self, query: str, include_images: bool = False) -> Any:
72 | msg = "Oops!"
73 | raise RuntimeError(msg)
74 |
75 | with patch("tavily.tavily.TavilyClient", FakeClient):
76 | monkeypatch.setenv("TAVILY_API_KEY", "fake-key")
77 | from any_agent.tools import search_tavily
78 |
79 | result = search_tavily("test")
80 | assert "Error performing Tavily search" in result
81 |
--------------------------------------------------------------------------------
/src/any_agent/serving/a2a/server_a2a.py:
--------------------------------------------------------------------------------
1 | from __future__ import annotations
2 |
3 | import asyncio
4 | from typing import TYPE_CHECKING
5 |
6 | import httpx
7 | import uvicorn
8 | from a2a.server.apps import A2AStarletteApplication
9 | from a2a.server.request_handlers import DefaultRequestHandler
10 | from starlette.applications import Starlette
11 | from starlette.routing import Mount
12 |
13 | from any_agent.serving.a2a.context_manager import ContextManager
14 | from any_agent.serving.server_handle import ServerHandle
15 |
16 | from .agent_card import _get_agent_card
17 | from .agent_executor import AnyAgentExecutor
18 | from .envelope import prepare_agent_for_a2a_async
19 |
20 | if TYPE_CHECKING:
21 | from any_agent import AnyAgent
22 | from any_agent.serving import A2AServingConfig
23 |
24 |
25 | async def _get_a2a_app_async(
26 | agent: AnyAgent, serving_config: A2AServingConfig
27 | ) -> A2AStarletteApplication:
28 | agent = await prepare_agent_for_a2a_async(agent)
29 |
30 | agent_card = _get_agent_card(agent, serving_config)
31 | task_manager = ContextManager(serving_config)
32 | push_notification_config_store = serving_config.push_notifier_store_type()
33 | push_notification_sender = serving_config.push_notifier_sender_type(
34 | httpx_client=httpx.AsyncClient(), # type: ignore[call-arg]
35 | config_store=push_notification_config_store,
36 | )
37 |
38 | request_handler = DefaultRequestHandler(
39 | agent_executor=AnyAgentExecutor(
40 | agent, task_manager, serving_config.stream_tool_usage
41 | ),
42 | task_store=serving_config.task_store_type(),
43 | push_config_store=push_notification_config_store,
44 | push_sender=push_notification_sender,
45 | )
46 |
47 | return A2AStarletteApplication(agent_card=agent_card, http_handler=request_handler)
48 |
49 |
50 | def _create_server(
51 | app: A2AStarletteApplication,
52 | host: str,
53 | port: int,
54 | endpoint: str,
55 | log_level: str = "warning",
56 | ) -> uvicorn.Server:
57 | root = endpoint.lstrip("/").rstrip("/")
58 | a2a_app = app.build()
59 | internal_router = Starlette(routes=[Mount(f"/{root}", routes=a2a_app.routes)])
60 |
61 | config = uvicorn.Config(internal_router, host=host, port=port, log_level=log_level)
62 | return uvicorn.Server(config)
63 |
64 |
65 | async def serve_a2a_async(
66 | server: A2AStarletteApplication,
67 | host: str,
68 | port: int,
69 | endpoint: str,
70 | log_level: str = "warning",
71 | ) -> ServerHandle:
72 | """Provide an A2A server to be used in an event loop."""
73 | uv_server = _create_server(server, host, port, endpoint, log_level)
74 | task = asyncio.create_task(uv_server.serve())
75 | while not uv_server.started: # noqa: ASYNC110
76 | await asyncio.sleep(0.1)
77 | if port == 0:
78 | server_port = uv_server.servers[0].sockets[0].getsockname()[1]
79 | server.agent_card.url = f"http://{host}:{server_port}/{endpoint.lstrip('/')}"
80 | return ServerHandle(task=task, server=uv_server)
81 |
--------------------------------------------------------------------------------
/docs/assets/custom.css:
--------------------------------------------------------------------------------
1 | /* Import Roboto font from Google Fonts */
2 | @import url('https://fonts.googleapis.com/css2?family=Roboto:wght@300;400;500;700&display=swap');
3 |
4 | /* Change header title font to Roboto */
5 | .md-header__title,
6 | .md-header__topic {
7 | font-family: 'Roboto', sans-serif !important;
8 | } */
9 |
10 | .jp-OutputArea-output, .output, .cell_output {
11 | max-height: 300px;
12 | }
13 |
14 | /* Hide the default edit icon */
15 | .md-content__button.md-content__button--edit .md-icon {
16 | display: none;
17 | }
18 |
19 | /* Hide the default SVG icon inside the edit button */
20 | .md-content__button.md-icon svg {
21 | display: none;
22 | }
23 |
24 | /* Add the GitHub SVG icon and text before the button content */
25 | .md-content__button.md-icon::before {
26 | content: "";
27 | display: inline-block;
28 | vertical-align: middle;
29 | width: 1.2em; /* Adjust size as needed */
30 | height: 1.2em;
31 | margin-right: 0.4em;
32 | background-image: url('data:image/svg+xml;utf8,');
33 | background-size: contain;
34 | background-repeat: no-repeat;
35 | background-position: left center;
36 | }
37 |
38 | /* Add the text after the icon */
39 | .md-content__button.md-icon::after {
40 | font-size: 1.1em;
41 | color: #24292f;
42 | font-family: inherit;
43 | vertical-align: middle;
44 | margin-left: 0.3em;
45 | }
46 |
47 | /* From https://github.com/squidfunk/mkdocs-material/discussions/6404#discussioncomment-7695243
48 | /* Apply max-width: initial to all pages by default */
49 | .md-main__inner.md-grid {
50 | max-width: initial;
51 | }
52 |
53 | /* The tracing page is designed to show traces in default material theme max-width of 1440px */
54 | body.tracing-page .md-grid {
55 | max-width: 1440px;
56 | }
57 |
--------------------------------------------------------------------------------
/mkdocs.yml:
--------------------------------------------------------------------------------
1 | site_name: any-agent
2 |
3 | repo_url: https://github.com/mozilla-ai/any-agent
4 | repo_name: any-agent
5 |
6 | nav:
7 | - Intro: index.md
8 | - Agents:
9 | - Defining and Running Agents: agents/index.md
10 | - Models: agents/models.md
11 | - Callbacks: agents/callbacks.md
12 | - Frameworks:
13 | - Agno: agents/frameworks/agno.md
14 | - Google ADK: agents/frameworks/google_adk.md
15 | - Langchain: agents/frameworks/langchain.md
16 | - LlamaIndex: agents/frameworks/llama_index.md
17 | - OpenAI Agents SDK: agents/frameworks/openai.md
18 | - smolagents: agents/frameworks/smolagents.md
19 | - TinyAgent: agents/frameworks/tinyagent.md
20 | - Tools: agents/tools.md
21 | - Tracing: tracing.md
22 | - Evaluation: evaluation.md
23 | - Serving: serving.md
24 | - Cookbook:
25 | - Your First Agent: cookbook/your_first_agent.ipynb
26 | - Your First Agent Evaluation: cookbook/your_first_agent_evaluation.ipynb
27 | - Using Callbacks: cookbook/callbacks.ipynb
28 | - MCP Agent: cookbook/mcp_agent.ipynb
29 | - Serve with A2A: cookbook/serve_a2a.ipynb
30 | - Use an Agent as a tool for another agent (A2A): cookbook/a2a_as_tool.ipynb
31 | - Local Agent: cookbook/agent_with_local_llm.ipynb
32 | - API Reference:
33 | - Agent: api/agent.md
34 | - Callbacks: api/callbacks.md
35 | - Config: api/config.md
36 | - Evaluation: api/evaluation.md
37 | - Logging: api/logging.md
38 | - Serving: api/serving.md
39 | - Tools: api/tools.md
40 | - Tracing: api/tracing.md
41 |
42 | theme:
43 | name: material
44 | font:
45 | text: Noto Sans
46 | code: Noto Sans Mono
47 | palette:
48 | - media: "(prefers-color-scheme: light)"
49 | scheme: default
50 | primary: blue gray
51 | toggle:
52 | icon: material/lightbulb
53 | name: Switch to dark mode
54 | - media: "(prefers-color-scheme: dark)"
55 | scheme: slate
56 | primary: grey
57 | toggle:
58 | icon: material/lightbulb-outline
59 | name: Switch to light mode
60 | - accent: blue
61 | logo: images/any-agent-logo-mark.png
62 | favicon: images/any-agent_favicon.png
63 | features:
64 | - content.code.copy
65 | - content.tabs.link
66 | - content.action.edit
67 | - navigation.expand
68 | - navigation.footer
69 |
70 | extra_css:
71 | - assets/custom.css
72 |
73 | extra_javascript:
74 | - assets/custom.js
75 |
76 | markdown_extensions:
77 | - admonition
78 | - pymdownx.highlight:
79 | anchor_linenums: true
80 | line_spans: __span
81 | pygments_lang_class: true
82 | - pymdownx.inlinehilite
83 | - pymdownx.snippets
84 | - pymdownx.superfences
85 | - pymdownx.tabbed:
86 | alternate_style: true
87 |
88 | plugins:
89 | - search
90 | - include-markdown
91 | - mkdocs-jupyter:
92 | no_input: false
93 | show_input: true
94 | - mkdocstrings:
95 | handlers:
96 | python:
97 | options:
98 | show_root_heading: true
99 | heading_level: 3
100 |
101 | hooks:
102 | - scripts/hooks.py
103 |
104 | edit_uri: edit/main/docs/
105 |
--------------------------------------------------------------------------------
/tests/unit/frameworks/test_google.py:
--------------------------------------------------------------------------------
1 | from typing import Any
2 | from unittest.mock import AsyncMock, MagicMock, patch
3 |
4 | import pytest
5 |
6 | from any_agent import AgentConfig, AgentFramework, AnyAgent
7 |
8 |
9 | def test_load_google_default() -> None:
10 | from google.adk.tools import FunctionTool
11 |
12 | mock_agent = MagicMock()
13 | mock_model = MagicMock()
14 | mock_function_tool = MagicMock()
15 |
16 | class MockedFunctionTool(FunctionTool):
17 | def __new__(cls, *args: Any, **kwargs: Any) -> MagicMock:
18 | return mock_function_tool
19 |
20 | with (
21 | patch("any_agent.frameworks.google.LlmAgent", mock_agent),
22 | patch("any_agent.frameworks.google.DEFAULT_MODEL_TYPE", mock_model),
23 | patch("google.adk.tools.FunctionTool", MockedFunctionTool),
24 | ):
25 | AnyAgent.create(
26 | AgentFramework.GOOGLE, AgentConfig(model_id="mistral:mistral-small-latest")
27 | )
28 | mock_agent.assert_called_once_with(
29 | name="any_agent",
30 | instruction="",
31 | model=mock_model(model="mistral:mistral-small-latest"),
32 | tools=[],
33 | output_key="response",
34 | )
35 |
36 |
37 | def test_load_google_agent_missing() -> None:
38 | with patch("any_agent.frameworks.google.adk_available", False):
39 | with pytest.raises(ImportError):
40 | AnyAgent.create(
41 | AgentFramework.GOOGLE,
42 | AgentConfig(model_id="mistral:mistral-small-latest"),
43 | )
44 |
45 |
46 | def test_run_google_custom_args() -> None:
47 | from google.adk.agents.run_config import RunConfig
48 | from google.genai import types
49 |
50 | mock_agent = MagicMock()
51 | mock_runner = MagicMock()
52 | mock_runner.get_tools = AsyncMock()
53 | mock_session = MagicMock()
54 | mock_runner.return_value.session_service.create_session = AsyncMock()
55 | mock_runner.return_value.session_service.get_session = AsyncMock()
56 |
57 | # More explicit mock setup
58 | mock_state = MagicMock()
59 | mock_state.get.return_value = "mock response"
60 | mock_session.state = mock_state
61 | mock_runner.return_value.session_service.get_session.return_value = mock_session
62 |
63 | run_config = RunConfig(max_llm_calls=10)
64 | with (
65 | patch("any_agent.frameworks.google.LlmAgent", mock_agent),
66 | patch("any_agent.frameworks.google.InMemoryRunner", mock_runner),
67 | patch("any_agent.frameworks.google.DEFAULT_MODEL_TYPE"),
68 | patch("google.adk.tools.FunctionTool"),
69 | ):
70 | agent = AnyAgent.create(
71 | AgentFramework.GOOGLE, AgentConfig(model_id="mistral:mistral-small-latest")
72 | )
73 | result = agent.run("foo", user_id="1", session_id="2", run_config=run_config)
74 |
75 | # Verify the result is as expected
76 | assert isinstance(result.final_output, str)
77 | assert result.final_output == "mock response"
78 |
79 | mock_runner.return_value.run_async.assert_called_once_with(
80 | user_id="1",
81 | session_id="2",
82 | new_message=types.Content(role="user", parts=[types.Part(text="foo")]),
83 | run_config=run_config,
84 | )
85 |
--------------------------------------------------------------------------------
/demo/components/agent_status.py:
--------------------------------------------------------------------------------
1 | from collections.abc import Callable
2 | from typing import Any
3 |
4 | from any_agent.callbacks import Callback, Context
5 | from any_agent.tracing.attributes import GenAI
6 |
7 |
8 | class StreamlitStatusCallback(Callback):
9 | """Callback to update Streamlit status with agent progress."""
10 |
11 | def __init__(self, status_callback: Callable[[str], None]):
12 | self.status_callback = status_callback
13 |
14 | def after_llm_call(self, context: Context, *args, **kwargs) -> Context:
15 | """Update status after LLM calls."""
16 | span = context.current_span
17 | input_value = span.attributes.get(GenAI.INPUT_MESSAGES, "")
18 | output_value = span.attributes.get(GenAI.OUTPUT, "")
19 |
20 | self._update_status(span.name, input_value, output_value)
21 | return context
22 |
23 | def after_tool_execution(self, context: Context, *args, **kwargs) -> Context:
24 | """Update status after tool executions."""
25 | span = context.current_span
26 | input_value = span.attributes.get(GenAI.TOOL_ARGS, "")
27 | output_value = span.attributes.get(GenAI.OUTPUT, "")
28 |
29 | self._update_status(span.name, input_value, output_value)
30 | return context
31 |
32 | def _update_status(self, step_name: str, input_value: str, output_value: str):
33 | """Update the Streamlit status with formatted information."""
34 | if input_value:
35 | try:
36 | import json
37 |
38 | parsed_input = json.loads(input_value)
39 | if isinstance(parsed_input, list) and len(parsed_input) > 0:
40 | input_value = str(parsed_input[-1])
41 | except Exception:
42 | pass
43 |
44 | if output_value:
45 | try:
46 | import json
47 |
48 | parsed_output = json.loads(output_value)
49 | if isinstance(parsed_output, list) and len(parsed_output) > 0:
50 | output_value = str(parsed_output[-1])
51 | except Exception:
52 | pass
53 |
54 | max_length = 800
55 | if len(input_value) > max_length:
56 | input_value = f"[Truncated]...{input_value[-max_length:]}"
57 | if len(output_value) > max_length:
58 | output_value = f"[Truncated]...{output_value[-max_length:]}"
59 |
60 | if input_value or output_value:
61 | message = f"Step: {step_name}\n"
62 | if input_value:
63 | message += f"Input: {input_value}\n"
64 | if output_value:
65 | message += f"Output: {output_value}"
66 | else:
67 | message = f"Step: {step_name}"
68 |
69 | self.status_callback(message)
70 |
71 |
72 | def export_logs(agent: Any, callback: Callable[[str], None]) -> None:
73 | """Add a Streamlit status callback to the agent.
74 |
75 | This function adds a custom callback to the agent that will update
76 | the Streamlit status with progress information during agent execution.
77 | """
78 | status_callback = StreamlitStatusCallback(callback)
79 |
80 | if agent.config.callbacks is None:
81 | agent.config.callbacks = []
82 | agent.config.callbacks.append(status_callback)
83 |
--------------------------------------------------------------------------------
/tests/unit/tracing/test_agent_trace.py:
--------------------------------------------------------------------------------
1 | from any_agent.testing.helpers import DEFAULT_SMALL_MODEL_ID
2 | from any_agent.tracing.agent_trace import AgentSpan, AgentTrace
3 | from any_agent.tracing.attributes import GenAI
4 | from any_agent.tracing.otel_types import Resource, SpanContext, SpanKind, Status
5 |
6 |
7 | def create_llm_span(input_tokens: int = 100, output_tokens: int = 50) -> AgentSpan:
8 | """Create a mock LLM span with token usage."""
9 | return AgentSpan(
10 | name=f"call_llm {DEFAULT_SMALL_MODEL_ID}",
11 | kind=SpanKind.INTERNAL,
12 | status=Status(),
13 | context=SpanContext(span_id=123),
14 | attributes={
15 | GenAI.OPERATION_NAME: "call_llm",
16 | GenAI.USAGE_INPUT_TOKENS: input_tokens,
17 | GenAI.USAGE_OUTPUT_TOKENS: output_tokens,
18 | },
19 | links=[],
20 | events=[],
21 | resource=Resource(),
22 | )
23 |
24 |
25 | def test_tokens_and_cost_properties_are_cached() -> None:
26 | """Test that tokens and cost properties are cached after first access."""
27 | trace = AgentTrace()
28 | trace.add_span(create_llm_span(input_tokens=100, output_tokens=50))
29 |
30 | # First access - should compute and cache
31 | tokens1 = trace.tokens
32 | cost1 = trace.cost
33 |
34 | # Second access - should return cached objects
35 | tokens2 = trace.tokens
36 | cost2 = trace.cost
37 |
38 | assert tokens1 is tokens2 # Same object reference indicates caching
39 | assert cost1 is cost2 # Same object reference indicates caching
40 | assert "tokens" in trace.__dict__
41 | assert "cost" in trace.__dict__
42 |
43 |
44 | def test_add_span_invalidates_cache() -> None:
45 | """Test that adding a span invalidates both tokens and cost caches."""
46 | trace = AgentTrace()
47 | trace.add_span(create_llm_span(input_tokens=100, output_tokens=50))
48 |
49 | # Cache the properties
50 | _ = trace.tokens
51 | _ = trace.cost
52 | assert "tokens" in trace.__dict__
53 | assert "cost" in trace.__dict__
54 |
55 | # Add another span - should invalidate cache
56 | trace.add_span(create_llm_span(input_tokens=200, output_tokens=75))
57 |
58 | assert "tokens" not in trace.__dict__
59 | assert "cost" not in trace.__dict__
60 |
61 | # Verify new calculations are correct
62 | tokens = trace.tokens
63 | assert tokens.input_tokens == 300
64 | assert tokens.output_tokens == 125
65 |
66 |
67 | def test_invalidate_cache_method() -> None:
68 | """Test that _invalidate_tokens_and_cost_cache clears both caches."""
69 | trace = AgentTrace()
70 | trace.add_span(create_llm_span())
71 |
72 | # Cache both properties
73 | _ = trace.tokens
74 | _ = trace.cost
75 | assert "tokens" in trace.__dict__
76 | assert "cost" in trace.__dict__
77 |
78 | # Manually invalidate cache
79 | trace._invalidate_tokens_and_cost_cache()
80 |
81 | assert "tokens" not in trace.__dict__
82 | assert "cost" not in trace.__dict__
83 |
84 |
85 | def test_spans_to_messages_handles_empty_spans() -> None:
86 | """Test that spans_to_messages handles traces with no spans."""
87 | empty_trace = AgentTrace()
88 | messages = empty_trace.spans_to_messages()
89 |
90 | assert isinstance(messages, list)
91 | assert len(messages) == 0
92 |
--------------------------------------------------------------------------------
/src/any_agent/callbacks/span_generation/google.py:
--------------------------------------------------------------------------------
1 | # mypy: disable-error-code="no-untyped-def,override,union-attr"
2 | from __future__ import annotations
3 |
4 | from typing import TYPE_CHECKING, Any
5 |
6 | from any_agent.callbacks.span_generation.base import _SpanGeneration
7 |
8 | if TYPE_CHECKING:
9 | from google.adk.models.llm_request import LlmRequest
10 | from google.adk.models.llm_response import LlmResponse
11 | from google.adk.tools.base_tool import BaseTool
12 | from google.adk.tools.tool_context import ToolContext
13 |
14 | from any_agent.callbacks.context import Context
15 |
16 |
17 | class _GoogleSpanGeneration(_SpanGeneration):
18 | def before_llm_call(self, context: Context, *args: Any, **kwargs: Any) -> Context:
19 | llm_request: LlmRequest = kwargs["llm_request"]
20 |
21 | messages = []
22 | if config := llm_request.config:
23 | messages.append(
24 | {
25 | "role": "system",
26 | "content": getattr(config, "system_instruction", "No instructions"),
27 | }
28 | )
29 | if parts := llm_request.contents[0].parts:
30 | messages.append(
31 | {
32 | "role": getattr(llm_request.contents[0], "role", "No role"),
33 | "content": getattr(parts[0], "text", "No content"),
34 | }
35 | )
36 |
37 | return self._set_llm_input(context, str(llm_request.model), messages)
38 |
39 | def after_llm_call(self, context: Context, *args, **kwargs) -> Context:
40 | llm_response: LlmResponse = kwargs["llm_response"]
41 |
42 | content = llm_response.content
43 | output: str | list[dict[str, Any]]
44 | if not content or not content.parts:
45 | output = ""
46 | elif content.parts[0].text:
47 | output = str(content.parts[0].text)
48 | else:
49 | output = [
50 | {
51 | "tool.name": getattr(part.function_call, "name", "No name"),
52 | "tool.args": getattr(part.function_call, "args", "{}"),
53 | }
54 | for part in content.parts
55 | if part.function_call
56 | ]
57 | input_tokens = 0
58 | output_tokens = 0
59 | if resp_meta := llm_response.usage_metadata:
60 | if prompt_tokens := resp_meta.prompt_token_count:
61 | input_tokens = prompt_tokens
62 | if candidates_token := resp_meta.candidates_token_count:
63 | output_tokens = candidates_token
64 | return self._set_llm_output(context, output, input_tokens, output_tokens)
65 |
66 | def before_tool_execution(self, context: Context, *args, **kwargs) -> Context:
67 | tool: BaseTool = kwargs["tool"]
68 | tool_args: dict[str, Any] = kwargs["args"]
69 | tool_context: ToolContext = kwargs["tool_context"]
70 |
71 | return self._set_tool_input(
72 | context,
73 | tool.name,
74 | tool.description,
75 | tool_args,
76 | tool_context.function_call_id,
77 | )
78 |
79 | def after_tool_execution(self, context: Context, *args, **kwargs) -> Context:
80 | return self._set_tool_output(context, kwargs["tool_response"])
81 |
--------------------------------------------------------------------------------
/tests/unit/callbacks/test_console_print_span.py:
--------------------------------------------------------------------------------
1 | from unittest.mock import MagicMock, patch
2 |
3 | import pytest
4 | from opentelemetry.sdk.trace import ReadableSpan
5 |
6 | from any_agent import AgentTrace
7 | from any_agent.callbacks.span_print import ConsolePrintSpan, _get_output_panel
8 |
9 |
10 | @pytest.fixture
11 | def readable_spans(agent_trace: AgentTrace) -> list[ReadableSpan]:
12 | return [span.to_readable_span() for span in agent_trace.spans]
13 |
14 |
15 | def test_console_print_span(
16 | agent_trace: AgentTrace, request: pytest.FixtureRequest
17 | ) -> None:
18 | console_mock = MagicMock()
19 | panel_mock = MagicMock()
20 | markdown_mock = MagicMock()
21 | with (
22 | patch("any_agent.callbacks.span_print.Console", console_mock),
23 | patch("any_agent.callbacks.span_print.Markdown", markdown_mock),
24 | patch("any_agent.callbacks.span_print.Panel", panel_mock),
25 | ):
26 | callback = ConsolePrintSpan()
27 |
28 | context = MagicMock()
29 | for span in agent_trace.spans:
30 | context.current_span = span.to_readable_span()
31 | if span.is_llm_call():
32 | callback.after_llm_call(context)
33 | elif span.is_tool_execution():
34 | callback.after_tool_execution(context)
35 |
36 | console_mock.return_value.print.assert_called()
37 |
38 | # Frameworks that end with a tool call or have JSON final output
39 | if request.node.callspec.id not in (
40 | "AGNO_trace",
41 | "GOOGLE_trace",
42 | "OPENAI_trace",
43 | "SMOLAGENTS_trace",
44 | "TINYAGENT_trace",
45 | ):
46 | panel_mock.assert_any_call(
47 | markdown_mock(agent_trace.final_output),
48 | title="OUTPUT",
49 | style="white",
50 | title_align="left",
51 | )
52 |
53 |
54 | def test_get_output_panel(
55 | readable_spans: list[ReadableSpan], request: pytest.FixtureRequest
56 | ) -> None:
57 | # First LLM call returns JSON
58 | panel_mock = MagicMock()
59 | json_mock = MagicMock()
60 | with (
61 | patch("any_agent.callbacks.span_print.Panel", panel_mock),
62 | patch("any_agent.callbacks.span_print.JSON", json_mock),
63 | ):
64 | _get_output_panel(readable_spans[0])
65 | json_mock.assert_called_once()
66 | panel_mock.assert_called_once()
67 |
68 | if request.node.callspec.id not in ("LLAMA_INDEX_trace",):
69 | # First TOOL execution returns JSON
70 | panel_mock = MagicMock()
71 | json_mock = MagicMock()
72 | with (
73 | patch("any_agent.callbacks.span_print.Panel", panel_mock),
74 | patch("any_agent.callbacks.span_print.JSON", json_mock),
75 | ):
76 | _get_output_panel(readable_spans[1])
77 | json_mock.assert_called_once()
78 | panel_mock.assert_called_once()
79 |
80 | # AGENT invocation has no output
81 | panel_mock = MagicMock()
82 | json_mock = MagicMock()
83 | with (
84 | patch("any_agent.callbacks.span_print.Panel", panel_mock),
85 | patch("any_agent.callbacks.span_print.JSON", json_mock),
86 | ):
87 | _get_output_panel(readable_spans[-1])
88 | json_mock.assert_not_called()
89 | panel_mock.assert_not_called()
90 |
--------------------------------------------------------------------------------
/tests/integration/a2a/test_a2a_serve.py:
--------------------------------------------------------------------------------
1 | from uuid import uuid4
2 |
3 | import pytest
4 | import logging
5 |
6 | # Import your agent and config
7 | from any_agent import AgentConfig, AgentFramework, AnyAgent
8 | from any_agent.logging import setup_logger
9 | from any_agent.serving import A2AServingConfig
10 | from any_agent.testing.helpers import (
11 | DEFAULT_HTTP_KWARGS,
12 | DEFAULT_SMALL_MODEL_ID,
13 | get_default_agent_model_args,
14 | wait_for_server_async,
15 | )
16 | from sse_starlette.sse import AppStatus
17 |
18 | from .conftest import DATE_PROMPT, A2ATestHelpers, a2a_client_from_agent, get_datetime
19 | from ..conftest import _get_deterministic_port # noqa: TID252
20 |
21 |
22 | @pytest.mark.asyncio
23 | async def test_serve_async(
24 | request: pytest.FixtureRequest, a2a_test_helpers: A2ATestHelpers
25 | ) -> None:
26 | agent = await AnyAgent.create_async(
27 | AgentFramework.TINYAGENT,
28 | AgentConfig(
29 | model_id=DEFAULT_SMALL_MODEL_ID,
30 | instructions="Directly answer the question without asking the user for input.",
31 | description="I'm an agent to help.",
32 | model_args=get_default_agent_model_args(AgentFramework.TINYAGENT),
33 | ),
34 | )
35 |
36 | test_port = _get_deterministic_port(
37 | request.node.name, AgentFramework.TINYAGENT.value
38 | )
39 | async with a2a_client_from_agent(agent, A2AServingConfig(port=test_port)) as (
40 | client,
41 | server_url,
42 | ):
43 | await wait_for_server_async(server_url)
44 | request = a2a_test_helpers.create_send_message_request(
45 | text="What is an agent?",
46 | message_id=uuid4().hex,
47 | )
48 | response = await client.send_message(request, http_kwargs=DEFAULT_HTTP_KWARGS)
49 | assert response is not None
50 |
51 |
52 | @pytest.mark.asyncio
53 | async def test_serve_streaming_async(
54 | request: pytest.FixtureRequest, a2a_test_helpers: A2ATestHelpers
55 | ) -> None:
56 | agent = await AnyAgent.create_async(
57 | "tinyagent",
58 | AgentConfig(
59 | model_id=DEFAULT_SMALL_MODEL_ID,
60 | instructions="Use the available tools to obtain additional information to answer the query.",
61 | tools=[get_datetime],
62 | description="I'm an agent to help.",
63 | model_args=get_default_agent_model_args(AgentFramework.TINYAGENT),
64 | ),
65 | )
66 |
67 | test_port = _get_deterministic_port(
68 | request.node.name, AgentFramework.TINYAGENT.value
69 | )
70 |
71 | async with a2a_client_from_agent(
72 | agent, A2AServingConfig(port=test_port, stream_tool_usage=True)
73 | ) as (
74 | client,
75 | server_url,
76 | ):
77 | await wait_for_server_async(server_url)
78 | request = a2a_test_helpers.create_send_streaming_message_request(
79 | text=DATE_PROMPT,
80 | message_id=uuid4().hex,
81 | )
82 | responses = []
83 | async for response in client.send_message_streaming(
84 | request, http_kwargs=DEFAULT_HTTP_KWARGS
85 | ):
86 | responses.append(response)
87 | assert response is not None
88 |
89 | # 4 responses are for tool calls: 2 for get_datetime and 2 for final_answer
90 | assert len(responses) == 6
91 | AppStatus.get_or_create_exit_event().set()
92 |
--------------------------------------------------------------------------------
/src/any_agent/callbacks/wrappers/smolagents.py:
--------------------------------------------------------------------------------
1 | # mypy: disable-error-code="method-assign,misc,no-untyped-call,no-untyped-def,union-attr"
2 | from __future__ import annotations
3 |
4 | from copy import deepcopy
5 | from typing import TYPE_CHECKING, Any
6 |
7 | from opentelemetry.trace import get_current_span
8 |
9 | if TYPE_CHECKING:
10 | from collections.abc import Callable
11 |
12 | from any_agent.callbacks.context import Context
13 | from any_agent.frameworks.smolagents import SmolagentsAgent
14 |
15 |
16 | class _SmolagentsWrapper:
17 | def __init__(self) -> None:
18 | self.callback_context: dict[int, Context] = {}
19 | self._original_llm_call: Callable[..., Any] | None = None
20 | self._original_tools: Any | None = None
21 |
22 | async def wrap(self, agent: SmolagentsAgent) -> None:
23 | self._original_llm_call = agent._agent.model.generate
24 |
25 | def wrap_generate(*args, **kwargs):
26 | context = self.callback_context[
27 | get_current_span().get_span_context().trace_id
28 | ]
29 | context.shared["model_id"] = str(agent._agent.model.model_id)
30 |
31 | for callback in agent.config.callbacks:
32 | context = callback.before_llm_call(context, *args, **kwargs)
33 |
34 | output = self._original_llm_call(*args, **kwargs)
35 |
36 | for callback in agent.config.callbacks:
37 | context = callback.after_llm_call(context, output)
38 |
39 | return output
40 |
41 | agent._agent.model.generate = wrap_generate
42 |
43 | def wrapped_tool_execution(original_tool, original_call, *args, **kwargs):
44 | context = self.callback_context[
45 | get_current_span().get_span_context().trace_id
46 | ]
47 | context.shared["original_tool"] = original_tool
48 |
49 | for callback in agent.config.callbacks:
50 | context = callback.before_tool_execution(context, *args, **kwargs)
51 |
52 | output = original_call(**kwargs)
53 |
54 | for callback in agent.config.callbacks:
55 | context = callback.after_tool_execution(
56 | context, output, *args, **kwargs
57 | )
58 |
59 | return output
60 |
61 | class WrappedToolCall:
62 | def __init__(self, original_tool, original_forward):
63 | self.original_tool = original_tool
64 | self.original_forward = original_forward
65 |
66 | def forward(self, *args, **kwargs):
67 | return wrapped_tool_execution(
68 | self.original_tool, self.original_forward, *args, **kwargs
69 | )
70 |
71 | self._original_tools = deepcopy(agent._agent.tools)
72 | wrapped_tools = {}
73 | for key, tool in agent._agent.tools.items():
74 | original_forward = tool.forward
75 | wrapped = WrappedToolCall(tool, original_forward)
76 | tool.forward = wrapped.forward
77 | wrapped_tools[key] = tool
78 | agent._agent.tools = wrapped_tools
79 |
80 | async def unwrap(self, agent: SmolagentsAgent) -> None:
81 | if self._original_llm_call is not None:
82 | agent._agent.model.generate = self._original_llm_call
83 | if self._original_tools is not None:
84 | agent._agent.tools = self._original_tools
85 |
--------------------------------------------------------------------------------
/src/any_agent/tools/composio.py:
--------------------------------------------------------------------------------
1 | # mypy: disable-error-code="attr-defined,operator,type-arg"
2 | import inspect as i
3 | import typing as t
4 | from types import FunctionType
5 |
6 | try:
7 | from composio.client.types import Tool
8 | from composio.core.provider import AgenticProvider
9 | from composio.core.provider.agentic import AgenticProviderExecuteFn
10 | except ImportError as e:
11 | msg = "Composio is not installed. Please install it with `pip install composio`."
12 | raise ImportError(msg) from e
13 |
14 |
15 | TYPE_MAPPING = {
16 | "string": str,
17 | "integer": int,
18 | "number": float,
19 | "boolean": bool,
20 | "array": list,
21 | "object": dict,
22 | }
23 |
24 |
25 | def _get_parameters(tool: Tool) -> list[i.Parameter]:
26 | parameters = []
27 | if tool.input_parameters and isinstance(tool.input_parameters, dict):
28 | properties = tool.input_parameters.get("properties", {})
29 | required = tool.input_parameters.get("required", [])
30 |
31 | for param_name, param_info in properties.items():
32 | base_param_type = TYPE_MAPPING.get(param_info.get("type", "string"), str)
33 |
34 | if param_name not in required:
35 | param = i.Parameter(
36 | param_name,
37 | i.Parameter.KEYWORD_ONLY,
38 | default=None,
39 | annotation=base_param_type | None,
40 | )
41 | else:
42 | param = i.Parameter(
43 | param_name,
44 | i.Parameter.KEYWORD_ONLY,
45 | annotation=base_param_type,
46 | )
47 | parameters.append(param)
48 |
49 | return parameters
50 |
51 |
52 | class CallableProvider(AgenticProvider[t.Callable, list[t.Callable]], name="callable"):
53 | """Composio toolset for generic callables."""
54 |
55 | __schema_skip_defaults__ = True
56 |
57 | def wrap_tool(
58 | self,
59 | tool: Tool,
60 | execute_tool: AgenticProviderExecuteFn,
61 | ) -> t.Callable:
62 | """Wrap composio tool as a python callable."""
63 | docstring = tool.description
64 | docstring += "\nArgs:"
65 | for _param, _schema in tool.input_parameters["properties"].items():
66 | docstring += "\n "
67 | docstring += _param + ": " + _schema.get("description", _param.title())
68 |
69 | docstring += "\nReturns:"
70 | docstring += "\n A dictionary containing response from the action"
71 |
72 | def _execute(**kwargs: t.Any) -> dict:
73 | return execute_tool(slug=tool.slug, arguments=kwargs)
74 |
75 | function = FunctionType(
76 | code=_execute.__code__,
77 | name=tool.slug,
78 | globals=globals(),
79 | closure=_execute.__closure__,
80 | )
81 |
82 | parameters = _get_parameters(tool)
83 | function.__annotations__ = {p.name: p.annotation for p in parameters} | {
84 | "return": dict
85 | }
86 | function.__signature__ = i.Signature(
87 | parameters=parameters, return_annotation=dict
88 | )
89 | function.__doc__ = docstring
90 | return function
91 |
92 | def wrap_tools(
93 | self,
94 | tools: t.Sequence[Tool],
95 | execute_tool: AgenticProviderExecuteFn,
96 | ) -> list[t.Callable]:
97 | """Wrap composio tools as python functions."""
98 | return [self.wrap_tool(tool, execute_tool) for tool in tools]
99 |
--------------------------------------------------------------------------------
/CODE_OF_CONDUCT.md:
--------------------------------------------------------------------------------
1 | # Contributor Covenant Code of Conduct
2 |
3 | ## Our Pledge
4 |
5 | In the interest of fostering an open and welcoming environment, we as
6 | contributors and maintainers pledge to making participation in our project and
7 | our community a harassment-free experience for everyone, regardless of age, body
8 | size, disability, ethnicity, sex characteristics, gender identity and expression,
9 | level of experience, education, socio-economic status, nationality, personal
10 | appearance, race, religion, or sexual identity and orientation.
11 |
12 | ## Our Standards
13 |
14 | Examples of behavior that contributes to creating a positive environment
15 | include:
16 |
17 | * Using welcoming and inclusive language
18 | * Being respectful of differing viewpoints and experiences
19 | * Gracefully accepting constructive criticism
20 | * Focusing on what is best for the community
21 | * Showing empathy towards other community members
22 |
23 | Examples of unacceptable behavior by participants include:
24 |
25 | * The use of sexualized language or imagery and unwelcome sexual attention or
26 | advances
27 | * Trolling, insulting/derogatory comments, and personal or political attacks
28 | * Public or private harassment
29 | * Publishing others' private information, such as a physical or electronic
30 | address, without explicit permission
31 | * Other conduct which could reasonably be considered inappropriate in a
32 | professional setting
33 |
34 | ## Our Responsibilities
35 |
36 | Project maintainers are responsible for clarifying the standards of acceptable
37 | behavior and are expected to take appropriate and fair corrective action in
38 | response to any instances of unacceptable behavior.
39 |
40 | Project maintainers have the right and responsibility to remove, edit, or
41 | reject comments, commits, code, wiki edits, issues, and other contributions
42 | that are not aligned to this Code of Conduct, or to ban temporarily or
43 | permanently any contributor for other behaviors that they deem inappropriate,
44 | threatening, offensive, or harmful.
45 |
46 | ## Scope
47 |
48 | This Code of Conduct applies both within project spaces and in public spaces
49 | when an individual is representing the project or its community. Examples of
50 | representing a project or community include using an official project e-mail
51 | address, posting via an official social media account, or acting as an appointed
52 | representative at an online or offline event. Representation of a project may be
53 | further defined and clarified by project maintainers.
54 |
55 | ## Enforcement
56 |
57 | Instances of abusive, harassing, or otherwise unacceptable behavior may be
58 | reported by contacting the team at mozilla.ai. All
59 | complaints will be reviewed and investigated and will result in a response that
60 | is deemed necessary and appropriate to the circumstances. The project team is
61 | obligated to maintain confidentiality with regard to the reporter of an incident.
62 | Further details of specific enforcement policies may be posted separately.
63 |
64 | Project maintainers who do not follow or enforce the Code of Conduct in good
65 | faith may face temporary or permanent repercussions as determined by other
66 | members of the project's leadership.
67 |
68 | ## Attribution
69 |
70 | This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4,
71 | available at https://www.contributor-covenant.org/version/1/4/code-of-conduct.html
72 |
73 | [homepage]: https://www.contributor-covenant.org
74 |
75 | For answers to common questions about this code of conduct, see
76 | https://www.contributor-covenant.org/faq
77 |
--------------------------------------------------------------------------------
/src/any_agent/callbacks/wrappers/tinyagent.py:
--------------------------------------------------------------------------------
1 | # mypy: disable-error-code="method-assign,misc,no-untyped-call,no-untyped-def,union-attr"
2 | from __future__ import annotations
3 |
4 | import asyncio
5 | from copy import deepcopy
6 | from typing import TYPE_CHECKING, Any
7 |
8 | from opentelemetry.trace import get_current_span
9 |
10 | if TYPE_CHECKING:
11 | from collections.abc import Callable
12 |
13 | from any_agent.callbacks.context import Context
14 | from any_agent.frameworks.tinyagent import TinyAgent
15 |
16 |
17 | class _TinyAgentWrapper:
18 | def __init__(self) -> None:
19 | self.callback_context: dict[int, Context] = {}
20 | self._original_llm_call: Callable[..., Any] | None = None
21 | self._original_clients: Any | None = None
22 |
23 | async def wrap(self, agent: TinyAgent) -> None:
24 | self._original_llm_call = agent.call_model
25 |
26 | async def wrap_call_model(**kwargs):
27 | context = self.callback_context[
28 | get_current_span().get_span_context().trace_id
29 | ]
30 | for callback in agent.config.callbacks:
31 | result = callback.before_llm_call(context, **kwargs)
32 | if asyncio.iscoroutinefunction(callback.before_llm_call):
33 | context = await result
34 | else:
35 | context = result
36 |
37 | output = await self._original_llm_call(**kwargs)
38 |
39 | for callback in agent.config.callbacks:
40 | result = callback.after_llm_call(context, output)
41 | if asyncio.iscoroutinefunction(callback.after_llm_call):
42 | context = await result
43 | else:
44 | context = result
45 |
46 | return output
47 |
48 | agent.call_model = wrap_call_model
49 |
50 | async def wrapped_tool_execution(original_call, request):
51 | context = self.callback_context[
52 | get_current_span().get_span_context().trace_id
53 | ]
54 | for callback in agent.config.callbacks:
55 | result = callback.before_tool_execution(context, request)
56 | if asyncio.iscoroutinefunction(callback.before_tool_execution):
57 | context = await result
58 | else:
59 | context = result
60 |
61 | output = await original_call(request)
62 |
63 | for callback in agent.config.callbacks:
64 | result = callback.after_tool_execution(context, output)
65 | if asyncio.iscoroutinefunction(callback.after_tool_execution):
66 | context = await result
67 | else:
68 | context = result
69 |
70 | return output
71 |
72 | class WrappedCallTool:
73 | def __init__(self, original_call_tool):
74 | self.original_call_tool = original_call_tool
75 |
76 | async def call_tool(self, request: dict[str, Any]):
77 | return await wrapped_tool_execution(self.original_call_tool, request)
78 |
79 | self._original_clients = deepcopy(agent.clients)
80 | wrapped_tools = {}
81 | for key, tool in agent.clients.items():
82 | wrapped = WrappedCallTool(tool.call_tool)
83 | tool.call_tool = wrapped.call_tool
84 | wrapped_tools[key] = tool
85 | agent.clients = wrapped_tools
86 |
87 | async def unwrap(self, agent: TinyAgent) -> None:
88 | if self._original_llm_call:
89 | agent.call_model = self._original_llm_call
90 | if self._original_clients:
91 | agent.clients = self._original_clients
92 |
--------------------------------------------------------------------------------
/tests/unit/callbacks/span_generation/test_base_span_generation.py:
--------------------------------------------------------------------------------
1 | import json
2 | from unittest.mock import MagicMock, patch
3 |
4 | import pytest
5 | from opentelemetry.trace import StatusCode
6 |
7 | from any_agent.callbacks.span_generation.base import _SpanGeneration
8 | from any_agent.tracing.attributes import GenAI
9 |
10 |
11 | class FooClass:
12 | pass
13 |
14 |
15 | foo_instance = FooClass()
16 |
17 |
18 | @pytest.mark.parametrize(
19 | ("tool_output", "expected_output", "expected_output_type"),
20 | [
21 | ("foo", "foo", "text"),
22 | (json.dumps({"foo": "bar"}), json.dumps({"foo": "bar"}), "json"),
23 | ({"foo": "bar"}, json.dumps({"foo": "bar"}), "json"),
24 | (foo_instance, json.dumps(foo_instance, default=str), "json"),
25 | ],
26 | )
27 | def test_set_tool_output(
28 | tool_output: object, expected_output: str, expected_output_type: str
29 | ) -> None:
30 | context = MagicMock()
31 | _SpanGeneration()._set_tool_output(context, tool_output)
32 |
33 | context.current_span.set_attributes.assert_called_with(
34 | {
35 | GenAI.OUTPUT: expected_output,
36 | GenAI.OUTPUT_TYPE: expected_output_type,
37 | }
38 | )
39 | context.current_span.set_status.assert_called_with(StatusCode.OK)
40 |
41 |
42 | def test_set_tool_output_error() -> None:
43 | error = "Error calling tool: It's a trap!"
44 | context = MagicMock()
45 | status_mock = MagicMock()
46 | with patch("any_agent.callbacks.span_generation.base.Status", status_mock):
47 | _SpanGeneration()._set_tool_output(context, error)
48 |
49 | context.current_span.set_attributes.assert_called_with(
50 | {GenAI.OUTPUT: error, GenAI.OUTPUT_TYPE: "text"}
51 | )
52 | context.current_span.set_status.assert_called_with(
53 | status_mock(status_code=StatusCode.ERROR, description=error)
54 | )
55 |
56 |
57 | def test_set_llm_input() -> None:
58 | context = MagicMock()
59 |
60 | span_generation = _SpanGeneration()
61 | span_generation._set_llm_input(context, model_id="gpt-5", input_messages=[])
62 | context.current_span.set_attribute.assert_called_with(GenAI.INPUT_MESSAGES, "[]")
63 |
64 | # first_llm_call logic should avoid logging input_messages
65 | # on subsequent calls.
66 | span_generation._set_llm_input(context, model_id="gpt-5", input_messages=[])
67 | assert context.current_span.set_attribute.call_count == 1
68 |
69 |
70 | def test_set_llm_output() -> None:
71 | context = MagicMock()
72 |
73 | span_generation = _SpanGeneration()
74 | span_generation._set_llm_output(
75 | context, output="foo", input_tokens=0, output_tokens=0
76 | )
77 | context.current_span.set_attributes.assert_called_with(
78 | {
79 | GenAI.OUTPUT: "foo",
80 | GenAI.OUTPUT_TYPE: "text",
81 | GenAI.USAGE_INPUT_TOKENS: 0,
82 | GenAI.USAGE_OUTPUT_TOKENS: 0,
83 | }
84 | )
85 |
86 | span_generation._set_llm_output(context, output=[], input_tokens=0, output_tokens=0)
87 | context.current_span.set_attributes.assert_called_with(
88 | {
89 | GenAI.OUTPUT: "[]",
90 | GenAI.OUTPUT_TYPE: "json",
91 | GenAI.USAGE_INPUT_TOKENS: 0,
92 | GenAI.USAGE_OUTPUT_TOKENS: 0,
93 | }
94 | )
95 |
96 |
97 | def test_set_tool_input() -> None:
98 | context = MagicMock()
99 |
100 | span_generation = _SpanGeneration()
101 | span_generation._set_tool_input(context, name="foo", args={})
102 | context.current_span.set_attributes.assert_called_with(
103 | {
104 | GenAI.OPERATION_NAME: "execute_tool",
105 | GenAI.TOOL_NAME: "foo",
106 | GenAI.TOOL_ARGS: "{}",
107 | }
108 | )
109 |
--------------------------------------------------------------------------------
/src/any_agent/callbacks/wrappers/google.py:
--------------------------------------------------------------------------------
1 | # mypy: disable-error-code="no-untyped-def,union-attr"
2 | from __future__ import annotations
3 |
4 | from typing import TYPE_CHECKING, Any
5 |
6 | from opentelemetry.trace import get_current_span
7 |
8 | if TYPE_CHECKING:
9 | from any_agent.callbacks.context import Context
10 | from any_agent.frameworks.google import GoogleAgent
11 |
12 |
13 | class _GoogleADKWrapper:
14 | def __init__(self) -> None:
15 | self.callback_context: dict[int, Context] = {}
16 | self._original: dict[str, Any] = {}
17 |
18 | async def wrap(self, agent: GoogleAgent) -> None:
19 | self._original["before_model"] = agent._agent.before_model_callback
20 |
21 | def before_model_callback(*args, **kwargs) -> Any | None:
22 | context = self.callback_context[
23 | get_current_span().get_span_context().trace_id
24 | ]
25 |
26 | for callback in agent.config.callbacks:
27 | context = callback.before_llm_call(context, *args, **kwargs)
28 |
29 | if callable(self._original["before_model"]):
30 | return self._original["before_model"](*args, **kwargs)
31 |
32 | return None
33 |
34 | agent._agent.before_model_callback = before_model_callback
35 |
36 | self._original["after_model"] = agent._agent.after_model_callback
37 |
38 | def after_model_callback(*args, **kwargs) -> Any | None:
39 | context = self.callback_context[
40 | get_current_span().get_span_context().trace_id
41 | ]
42 |
43 | for callback in agent.config.callbacks:
44 | context = callback.after_llm_call(context, *args, **kwargs)
45 |
46 | if callable(self._original["after_model"]):
47 | return self._original["after_model"](*args, **kwargs)
48 |
49 | return None
50 |
51 | agent._agent.after_model_callback = after_model_callback
52 |
53 | self._original["before_tool"] = agent._agent.before_tool_callback
54 |
55 | def before_tool_callback(*args, **kwargs) -> Any | None:
56 | context = self.callback_context[
57 | get_current_span().get_span_context().trace_id
58 | ]
59 |
60 | for callback in agent.config.callbacks:
61 | context = callback.before_tool_execution(context, *args, **kwargs)
62 |
63 | if callable(self._original["before_tool"]):
64 | return self._original["before_tool"](*args, **kwargs)
65 |
66 | return None
67 |
68 | agent._agent.before_tool_callback = before_tool_callback
69 |
70 | self._original["after_tool"] = agent._agent.after_tool_callback
71 |
72 | def after_tool_callback(*args, **kwarg) -> Any | None:
73 | context = self.callback_context[
74 | get_current_span().get_span_context().trace_id
75 | ]
76 |
77 | for callback in agent.config.callbacks:
78 | context = callback.after_tool_execution(context, *args, **kwarg)
79 |
80 | if callable(self._original["after_tool"]):
81 | return self._original["after_tool"](*args, **kwarg)
82 |
83 | return None
84 |
85 | agent._agent.after_tool_callback = after_tool_callback
86 |
87 | async def unwrap(self, agent: GoogleAgent) -> None:
88 | if "before_model" in self._original:
89 | agent._agent.before_model_callback = self._original["before_model"]
90 | if "before_tool" in self._original:
91 | agent._agent.before_tool_callback = self._original["before_tool"]
92 | if "after_model" in self._original:
93 | agent._agent.after_model_callback = self._original["after_model"]
94 | if "after_tool" in self._original:
95 | agent._agent.after_tool_callback = self._original["after_tool"]
96 |
--------------------------------------------------------------------------------
/src/any_agent/callbacks/span_print.py:
--------------------------------------------------------------------------------
1 | # mypy: disable-error-code="arg-type,attr-defined,no-untyped-def,union-attr"
2 | from __future__ import annotations
3 |
4 | import json
5 | from typing import TYPE_CHECKING
6 |
7 | from rich.console import Console, Group
8 | from rich.json import JSON
9 | from rich.markdown import Markdown
10 | from rich.panel import Panel
11 |
12 | from any_agent.callbacks.base import Callback
13 | from any_agent.tracing.attributes import GenAI
14 |
15 | if TYPE_CHECKING:
16 | from opentelemetry.sdk.trace import ReadableSpan
17 |
18 | from any_agent.callbacks.context import Context
19 |
20 |
21 | def _get_output_panel(span: ReadableSpan) -> Panel | None:
22 | if output := span.attributes.get(GenAI.OUTPUT, None):
23 | output_type = span.attributes.get(GenAI.OUTPUT_TYPE, "text")
24 | return Panel(
25 | Markdown(output) if output_type != "json" else JSON(output),
26 | title="OUTPUT",
27 | style="white",
28 | title_align="left",
29 | )
30 | return None
31 |
32 |
33 | class ConsolePrintSpan(Callback):
34 | """Use rich's console to print the `Context.current_span`."""
35 |
36 | def __init__(self, console: Console | None = None) -> None:
37 | """Init the ConsolePrintSpan.
38 |
39 | Args:
40 | console: An optional instance of `rich.console.Console`.
41 | If `None`, a new instance will be used.
42 |
43 | """
44 | self.console = console or Console()
45 |
46 | def after_llm_call(self, context: Context, *args, **kwargs) -> Context:
47 | span = context.current_span
48 |
49 | operation_name = span.attributes.get(GenAI.OPERATION_NAME, "")
50 |
51 | if operation_name != "call_llm":
52 | return context
53 |
54 | panels = []
55 |
56 | if messages := span.attributes.get(GenAI.INPUT_MESSAGES):
57 | panels.append(
58 | Panel(JSON(messages), title="INPUT", style="white", title_align="left")
59 | )
60 |
61 | if output_panel := _get_output_panel(span):
62 | panels.append(output_panel)
63 |
64 | if usage := {
65 | k.replace("gen_ai.usage.", ""): v
66 | for k, v in span.attributes.items()
67 | if "usage" in k
68 | }:
69 | panels.append(
70 | Panel(
71 | JSON(json.dumps(usage)),
72 | title="USAGE",
73 | style="white",
74 | title_align="left",
75 | )
76 | )
77 |
78 | self.console.print(
79 | Panel(
80 | Group(*panels),
81 | title=f"{operation_name.upper()}: {span.attributes.get(GenAI.REQUEST_MODEL)}",
82 | style="yellow",
83 | )
84 | )
85 |
86 | return context
87 |
88 | def after_tool_execution(self, context: Context, *args, **kwargs) -> Context:
89 | span = context.current_span
90 |
91 | operation_name = span.attributes.get(GenAI.OPERATION_NAME, "")
92 |
93 | if operation_name != "execute_tool":
94 | return context
95 |
96 | panels = [
97 | Panel(
98 | JSON(span.attributes.get(GenAI.TOOL_ARGS, "{}")),
99 | title="Input",
100 | style="white",
101 | title_align="left",
102 | )
103 | ]
104 |
105 | if output_panel := _get_output_panel(span):
106 | panels.append(output_panel)
107 |
108 | self.console.print(
109 | Panel(
110 | Group(*panels),
111 | title=f"{operation_name.upper()}: {span.attributes.get(GenAI.TOOL_NAME)}",
112 | style="blue",
113 | )
114 | )
115 | return context
116 |
--------------------------------------------------------------------------------
/src/any_agent/serving/mcp/server_mcp.py:
--------------------------------------------------------------------------------
1 | from __future__ import annotations
2 |
3 | import asyncio
4 | from typing import TYPE_CHECKING, Any
5 |
6 | import mcp.types as mcptypes
7 | import uvicorn
8 | from mcp.server import Server as MCPServer
9 | from mcp.server.sse import SseServerTransport
10 | from pydantic import BaseModel
11 | from starlette.applications import Starlette
12 | from starlette.responses import Response
13 | from starlette.routing import Mount, Route
14 |
15 | from any_agent.serving.server_handle import ServerHandle
16 |
17 | if TYPE_CHECKING:
18 | from starlette.requests import Request
19 |
20 | from any_agent import AnyAgent
21 |
22 |
23 | def _create_mcp_server_instance(agent: AnyAgent) -> MCPServer[Any]:
24 | server = MCPServer[Any]("any-agent-mcp-server")
25 |
26 | @server.list_tools() # type: ignore[no-untyped-call,misc]
27 | async def handle_list_tools() -> list[mcptypes.Tool]:
28 | return [
29 | mcptypes.Tool(
30 | name=f"as-tool-{agent.config.name}",
31 | description=agent.config.description,
32 | inputSchema={
33 | "type": "object",
34 | "required": ["query"],
35 | "properties": {
36 | "query": {
37 | "type": "string",
38 | "description": "The prompt for the agent",
39 | },
40 | },
41 | },
42 | )
43 | ]
44 |
45 | @server.call_tool() # type: ignore[misc]
46 | async def handle_call_tool(
47 | name: str, arguments: dict[str, Any]
48 | ) -> list[mcptypes.TextContent | mcptypes.ImageContent | mcptypes.EmbeddedResource]:
49 | result = await agent.run_async(arguments["query"])
50 | output = result.final_output
51 | if isinstance(output, BaseModel):
52 | serialized_output = output.model_dump_json()
53 | else:
54 | serialized_output = str(output)
55 | return [mcptypes.TextContent(type="text", text=serialized_output)]
56 |
57 | return server
58 |
59 |
60 | def _get_mcp_app(
61 | agent: AnyAgent,
62 | endpoint: str,
63 | ) -> Starlette:
64 | """Provide an MCP server to be used in an event loop."""
65 | root = endpoint.lstrip("/").rstrip("/")
66 | msg_endpoint = f"/{root}/messages/"
67 | sse = SseServerTransport(msg_endpoint)
68 | server = _create_mcp_server_instance(agent)
69 | init_options = server.create_initialization_options()
70 |
71 | async def _handle_sse(request: Request): # type: ignore[no-untyped-def]
72 | async with sse.connect_sse(
73 | request.scope, request.receive, request._send
74 | ) as streams:
75 | await server.run(streams[0], streams[1], init_options)
76 | # Return empty response to avoid NoneType error
77 | # Please check https://github.com/modelcontextprotocol/python-sdk/blob/1eb1bba83c70c3121bce7fc0263e5fac2c3f0520/src/mcp/server/sse.py#L33
78 | return Response()
79 |
80 | routes = [
81 | Route(f"/{root}/sse", endpoint=_handle_sse, methods=["GET"]),
82 | Mount(msg_endpoint, app=sse.handle_post_message),
83 | ]
84 | return Starlette(routes=routes)
85 |
86 |
87 | async def serve_mcp_async(
88 | agent: AnyAgent,
89 | host: str,
90 | port: int,
91 | endpoint: str,
92 | log_level: str = "warning",
93 | ) -> ServerHandle:
94 | """Provide an MCP server to be used in an event loop."""
95 | config = uvicorn.Config(
96 | _get_mcp_app(agent, endpoint), host=host, port=port, log_level=log_level
97 | )
98 | uv_server = uvicorn.Server(config)
99 | task = asyncio.create_task(uv_server.serve())
100 | while not uv_server.started: # noqa: ASYNC110
101 | await asyncio.sleep(0.1)
102 | return ServerHandle(task=task, server=uv_server)
103 |
--------------------------------------------------------------------------------
/tests/unit/frameworks/test_smolagents.py:
--------------------------------------------------------------------------------
1 | from unittest.mock import MagicMock, patch
2 |
3 | import pytest
4 |
5 | from any_agent import AgentConfig, AgentFramework, AnyAgent
6 |
7 |
8 | def test_load_smolagent_default() -> None:
9 | mock_agent = MagicMock()
10 | mock_model = MagicMock()
11 | mock_tool = MagicMock()
12 |
13 | with (
14 | patch("any_agent.frameworks.smolagents.DEFAULT_AGENT_TYPE", mock_agent),
15 | patch("any_agent.frameworks.smolagents.DEFAULT_MODEL_TYPE", mock_model),
16 | patch("smolagents.tool", mock_tool),
17 | ):
18 | AnyAgent.create(
19 | AgentFramework.SMOLAGENTS,
20 | AgentConfig(
21 | model_id="openai:o3-mini",
22 | ),
23 | )
24 |
25 | mock_agent.assert_called_once_with(
26 | name="any_agent",
27 | model=mock_model.return_value,
28 | verbosity_level=-1,
29 | tools=[],
30 | )
31 | mock_model.assert_called_once_with(
32 | model_id="openai:o3-mini", api_base=None, api_key=None
33 | )
34 |
35 |
36 | def test_load_smolagent_with_api_base() -> None:
37 | mock_agent = MagicMock()
38 | mock_model = MagicMock()
39 | mock_tool = MagicMock()
40 |
41 | with (
42 | patch("any_agent.frameworks.smolagents.DEFAULT_AGENT_TYPE", mock_agent),
43 | patch("any_agent.frameworks.smolagents.DEFAULT_MODEL_TYPE", mock_model),
44 | patch("smolagents.tool", mock_tool),
45 | ):
46 | AnyAgent.create(
47 | AgentFramework.SMOLAGENTS,
48 | AgentConfig(
49 | model_id="openai:o3-mini",
50 | model_args={},
51 | api_base="https://custom-api.example.com",
52 | ),
53 | )
54 |
55 | mock_agent.assert_called_once_with(
56 | name="any_agent",
57 | model=mock_model.return_value,
58 | tools=[],
59 | verbosity_level=-1,
60 | )
61 | mock_model.assert_called_once_with(
62 | model_id="openai:o3-mini",
63 | api_base="https://custom-api.example.com",
64 | api_key=None,
65 | )
66 |
67 |
68 | def test_load_smolagents_agent_missing() -> None:
69 | with patch("any_agent.frameworks.smolagents.smolagents_available", False):
70 | with pytest.raises(ImportError):
71 | AnyAgent.create(
72 | AgentFramework.SMOLAGENTS,
73 | AgentConfig(model_id="mistral:mistral-small-latest"),
74 | )
75 |
76 |
77 | def test_load_smolagent_final_answer() -> None:
78 | """Regression test for https://github.com/mozilla-ai/any-agent/issues/662"""
79 | from smolagents import FinalAnswerTool
80 |
81 | mock_model = MagicMock()
82 | mock_tool = MagicMock()
83 |
84 | with (
85 | patch("any_agent.frameworks.smolagents.DEFAULT_MODEL_TYPE", mock_model),
86 | patch("smolagents.tool", mock_tool),
87 | ):
88 | agent = AnyAgent.create(
89 | AgentFramework.SMOLAGENTS,
90 | AgentConfig(
91 | model_id="openai:o3-mini",
92 | ),
93 | )
94 |
95 | assert isinstance(agent._agent.tools["final_answer"], FinalAnswerTool) # type: ignore[attr-defined]
96 |
97 |
98 | def test_run_smolagent_custom_args() -> None:
99 | mock_agent = MagicMock()
100 | mock_agent.return_value = MagicMock()
101 | with (
102 | patch("any_agent.frameworks.smolagents.DEFAULT_AGENT_TYPE", mock_agent),
103 | patch("any_agent.frameworks.smolagents.DEFAULT_MODEL_TYPE"),
104 | patch("smolagents.tool"),
105 | ):
106 | agent = AnyAgent.create(
107 | AgentFramework.SMOLAGENTS,
108 | AgentConfig(
109 | model_id="openai:o3-mini",
110 | ),
111 | )
112 | agent.run("foo", max_steps=30)
113 | mock_agent.return_value.run.assert_called_once_with("foo", max_steps=30)
114 |
--------------------------------------------------------------------------------
/.github/workflows/tests-integration.yaml:
--------------------------------------------------------------------------------
1 | name: Integration Tests
2 |
3 | on:
4 | push:
5 | branches: [main]
6 | paths:
7 | - 'src/**'
8 | - 'tests/**'
9 | - '.github/workflows/**'
10 | - 'pyproject.toml'
11 | workflow_dispatch:
12 |
13 | jobs:
14 | run-integration-tinyagent:
15 | timeout-minutes: 30
16 | runs-on: ubuntu-latest
17 |
18 | steps:
19 | - uses: actions/checkout@v5
20 |
21 | - uses: astral-sh/setup-uv@v7
22 | with:
23 | python-version: 3.13
24 | activate-environment: true
25 |
26 | - run: uv sync -U --group tests
27 |
28 | - name: Run TINYAGENT test
29 | env:
30 | ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
31 | GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }}
32 | MISTRAL_API_KEY: ${{ secrets.MISTRAL_API_KEY }}
33 | OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
34 | XAI_API_KEY: ${{ secrets.XAI_API_KEY }}
35 | run: pytest tests/integration/frameworks -v -k "TINYAGENT" --cov --cov-report=xml --cov-append
36 |
37 | - name: Upload coverage reports to Codecov
38 | if: always()
39 | uses: codecov/codecov-action@v5
40 | with:
41 | token: ${{ secrets.CODECOV_TOKEN }}
42 |
43 | run-integration-all:
44 | timeout-minutes: 30
45 | runs-on: ubuntu-latest
46 |
47 | steps:
48 | - uses: actions/checkout@v5
49 |
50 | - uses: astral-sh/setup-uv@v7
51 | with:
52 | python-version: 3.13
53 | activate-environment: true
54 |
55 | - run: uv sync -U --group tests --extra all --extra composio
56 |
57 | - env:
58 | MISTRAL_API_KEY: ${{ secrets.MISTRAL_API_KEY }}
59 | OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
60 | run: pytest tests/integration/frameworks -n auto -v -k "not TINYAGENT" --cov --cov-report=xml --cov-append
61 |
62 | - name: Run Tool tests
63 | env:
64 | OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
65 | COMPOSIO_API_KEY: ${{ secrets.COMPOSIO_API_KEY }}
66 | COMPOSIO_USER_ID: ${{ secrets.COMPOSIO_USER_ID }}
67 | run: pytest -v tests/integration/tools --cov --cov-report=xml --cov-append
68 |
69 | - name: Run Snapshot tests
70 | run: pytest -v tests/snapshots --cov --cov-report=xml --cov-append
71 |
72 | - name: Upload coverage reports to Codecov
73 | if: always()
74 | uses: codecov/codecov-action@v5
75 | with:
76 | token: ${{ secrets.CODECOV_TOKEN }}
77 |
78 | run-integration-a2a:
79 | timeout-minutes: 30
80 | runs-on: ubuntu-latest
81 |
82 | steps:
83 | - uses: actions/checkout@v5
84 |
85 | - uses: astral-sh/setup-uv@v7
86 | with:
87 | python-version: 3.13
88 | activate-environment: true
89 |
90 | - run: uv sync -U --group tests --extra all --extra a2a
91 |
92 | - env:
93 | MISTRAL_API_KEY: ${{ secrets.MISTRAL_API_KEY }}
94 | run: pytest tests/integration/a2a -n auto -v --cov --cov-report=xml --cov-append
95 |
96 | - name: Upload coverage reports to Codecov
97 | if: always()
98 | uses: codecov/codecov-action@v5
99 | with:
100 | token: ${{ secrets.CODECOV_TOKEN }}
101 |
102 | run-integration-mcp:
103 | timeout-minutes: 30
104 | runs-on: ubuntu-latest
105 |
106 | steps:
107 | - uses: actions/checkout@v5
108 |
109 | - uses: astral-sh/setup-uv@v7
110 | with:
111 | python-version: 3.13
112 | activate-environment: true
113 |
114 | - run: uv sync -U --group tests --extra all
115 |
116 | - env:
117 | MISTRAL_API_KEY: ${{ secrets.MISTRAL_API_KEY }}
118 | run: pytest tests/integration/mcp -n auto -v --cov --cov-report=xml --cov-append
119 |
120 | - name: Upload coverage reports to Codecov
121 | if: always()
122 | uses: codecov/codecov-action@v5
123 | with:
124 | token: ${{ secrets.CODECOV_TOKEN }}
125 |
--------------------------------------------------------------------------------
/demo/tools/openmeteo.py:
--------------------------------------------------------------------------------
1 | import json
2 | from datetime import datetime, timedelta
3 |
4 | import requests
5 |
6 |
7 | def _extract_hourly_data(data: dict) -> list[dict]:
8 | hourly_data = data["hourly"]
9 | result = [
10 | {k: v for k, v in zip(hourly_data.keys(), values, strict=False)}
11 | for values in zip(*hourly_data.values(), strict=False)
12 | ]
13 | return result
14 |
15 |
16 | def _filter_by_date(
17 | date: datetime, hourly_data: list[dict], timedelta: timedelta = timedelta(hours=1)
18 | ):
19 | start_date = date - timedelta
20 | end_date = date + timedelta
21 | return [
22 | item
23 | for item in hourly_data
24 | if start_date <= datetime.fromisoformat(item["time"]) <= end_date
25 | ]
26 |
27 |
28 | def get_wave_forecast(lat: float, lon: float, date: str) -> list[dict]:
29 | """Get wave forecast for given location.
30 |
31 | Forecast will include:
32 |
33 | - wave_direction (degrees)
34 | - wave_height (meters)
35 | - wave_period (seconds)
36 | - sea_level_height_msl (meters)
37 |
38 | Args:
39 | lat: Latitude of the location.
40 | lon: Longitude of the location.
41 | date: Date to filter by in any valid ISO 8601 format.
42 |
43 | Returns:
44 | Hourly data for wave forecast.
45 | Example output:
46 |
47 | ```json
48 | [
49 | {'time': '2025-03-19T09:00', 'winddirection_10m': 140, 'windspeed_10m': 24.5}, {'time': '2025-03-19T10:00', 'winddirection_10m': 140, 'windspeed_10m': 27.1},
50 | {'time': '2025-03-19T10:00', 'winddirection_10m': 140, 'windspeed_10m': 27.1}, {'time': '2025-03-19T11:00', 'winddirection_10m': 141, 'windspeed_10m': 29.2}
51 | ]
52 | ```
53 |
54 | """
55 | url = "https://marine-api.open-meteo.com/v1/marine"
56 | params = {
57 | "latitude": lat,
58 | "longitude": lon,
59 | "hourly": [
60 | "wave_direction",
61 | "wave_height",
62 | "wave_period",
63 | "sea_level_height_msl",
64 | ],
65 | }
66 | response = requests.get(url, params=params)
67 | response.raise_for_status()
68 | data = json.loads(response.content.decode())
69 | hourly_data = _extract_hourly_data(data)
70 | if date is not None:
71 | date = datetime.fromisoformat(date)
72 | hourly_data = _filter_by_date(date, hourly_data)
73 | if len(hourly_data) == 0:
74 | raise ValueError("No data found for the given date")
75 | return hourly_data
76 |
77 |
78 | def get_wind_forecast(lat: float, lon: float, date: str) -> list[dict]:
79 | """Get wind forecast for given location.
80 |
81 | Forecast will include:
82 |
83 | - wind_direction (degrees)
84 | - wind_speed (meters per second)
85 |
86 | Args:
87 | lat: Latitude of the location.
88 | lon: Longitude of the location.
89 | date: Date to filter by in any valid ISO 8601 format.
90 |
91 | Returns:
92 | Hourly data for wind forecast.
93 | Example output:
94 |
95 | ```json
96 | [
97 | {"time": "2025-03-18T22:00", "wind_direction": 196, "wind_speed": 9.6},
98 | {"time": "2025-03-18T23:00", "wind_direction": 183, "wind_speed": 7.9},
99 | ]
100 | ```
101 |
102 | """
103 | url = "https://api.open-meteo.com/v1/forecast"
104 | params = {
105 | "latitude": lat,
106 | "longitude": lon,
107 | "hourly": ["winddirection_10m", "windspeed_10m"],
108 | }
109 | response = requests.get(url, params=params)
110 | response.raise_for_status()
111 | data = json.loads(response.content.decode())
112 | hourly_data = _extract_hourly_data(data)
113 | date = datetime.fromisoformat(date)
114 | hourly_data = _filter_by_date(date, hourly_data)
115 | if len(hourly_data) == 0:
116 | raise ValueError("No data found for the given date")
117 | return hourly_data
118 |
--------------------------------------------------------------------------------
/tests/unit/serving/test_envelope_creation.py:
--------------------------------------------------------------------------------
1 | import pytest
2 | from pydantic import BaseModel
3 |
4 | # Skip entire module if a2a dependencies are not available
5 | pytest.importorskip("a2a")
6 |
7 | from a2a.types import TaskState
8 |
9 | from any_agent.config import AgentConfig, AgentFramework
10 | from any_agent.frameworks.any_agent import AnyAgent
11 | from any_agent.serving.a2a.envelope import (
12 | A2AEnvelope,
13 | _DefaultBody,
14 | _is_a2a_envelope,
15 | prepare_agent_for_a2a_async,
16 | )
17 |
18 |
19 | class CustomOutputType(BaseModel):
20 | custom_field: str
21 | result: str
22 |
23 |
24 | class MockAgent(AnyAgent):
25 | """Mock agent implementation for testing."""
26 |
27 | def __init__(self, config: AgentConfig) -> None:
28 | super().__init__(config)
29 | self._agent = None
30 |
31 | async def _load_agent(self) -> None:
32 | pass
33 |
34 | async def _run_async(self, prompt: str, **kwargs: object) -> str:
35 | return "mock result"
36 |
37 | async def update_output_type_async(
38 | self, output_type: type[BaseModel] | None
39 | ) -> None:
40 | self.config.output_type = output_type
41 |
42 | @property
43 | def framework(self) -> AgentFramework:
44 | from any_agent.config import AgentFramework
45 |
46 | return AgentFramework.TINYAGENT
47 |
48 | @classmethod
49 | def create(cls, framework: object, config: AgentConfig) -> "MockAgent":
50 | return cls(config)
51 |
52 |
53 | @pytest.mark.asyncio
54 | async def test_envelope_created_without_output_type() -> None:
55 | """Test that the envelope is correctly created when the agent is configured without an output_type."""
56 | # Create agent config without output_type
57 | config = AgentConfig(model_id="test-model", description="test agent")
58 | assert config.output_type is None
59 |
60 | # Create mock agent
61 | agent = MockAgent(config)
62 |
63 | # Prepare agent for A2A
64 | prepared_agent = await prepare_agent_for_a2a_async(agent)
65 |
66 | # Verify the envelope was created with default body
67 | assert prepared_agent.config.output_type is not None
68 | assert _is_a2a_envelope(prepared_agent.config.output_type)
69 |
70 | # Verify the envelope wraps _DefaultBody
71 | envelope_instance = prepared_agent.config.output_type(
72 | task_status=TaskState.completed, data=_DefaultBody(result="test result")
73 | )
74 |
75 | assert isinstance(envelope_instance, A2AEnvelope)
76 | assert envelope_instance.task_status == TaskState.completed
77 | assert isinstance(envelope_instance.data, _DefaultBody)
78 | assert envelope_instance.data.result == "test result"
79 |
80 |
81 | @pytest.mark.asyncio
82 | async def test_envelope_created_with_output_type() -> None:
83 | """Test that the envelope is correctly created when an agent is configured with an output_type."""
84 | # Create agent config with custom output_type
85 | config = AgentConfig(
86 | model_id="test-model", description="test agent", output_type=CustomOutputType
87 | )
88 |
89 | # Create mock agent
90 | agent = MockAgent(config)
91 |
92 | # Prepare agent for A2A
93 | prepared_agent = await prepare_agent_for_a2a_async(agent)
94 |
95 | # Verify the envelope was created with custom output type
96 | assert prepared_agent.config.output_type is not None
97 | assert _is_a2a_envelope(prepared_agent.config.output_type)
98 |
99 | # Verify the envelope wraps the custom output type
100 | envelope_instance = prepared_agent.config.output_type(
101 | task_status=TaskState.completed,
102 | data=CustomOutputType(custom_field="test", result="custom result"),
103 | )
104 |
105 | assert isinstance(envelope_instance, A2AEnvelope)
106 | assert envelope_instance.task_status == TaskState.completed
107 | assert isinstance(envelope_instance.data, CustomOutputType)
108 | assert envelope_instance.data.custom_field == "test"
109 | assert envelope_instance.data.result == "custom result"
110 |
--------------------------------------------------------------------------------
/src/any_agent/evaluation/agent_judge.py:
--------------------------------------------------------------------------------
1 | import builtins
2 | from collections.abc import Callable
3 | from typing import Any
4 |
5 | from any_llm.utils.aio import run_async_in_sync
6 | from pydantic import BaseModel
7 |
8 | from any_agent import AgentConfig, AnyAgent
9 | from any_agent.config import AgentFramework
10 | from any_agent.evaluation.schemas import EvaluationOutput
11 | from any_agent.evaluation.tools import TraceTools
12 | from any_agent.tracing.agent_trace import AgentTrace
13 |
14 | INSIDE_NOTEBOOK = hasattr(builtins, "__IPYTHON__")
15 |
16 | AGENT_INSTRUCTIONS = """
17 | You are a helpful assistant that will be used to evaluate the correctness of an agent trace.
18 | Given a specific question regarding the quality of something about the agent,
19 | you may utilize tools as needed in order to check if the trace satisfies the question.
20 |
21 | Whenever you have all the information needed, you must use the `final_answer` tool
22 | to answer with:
23 |
24 | 1. "passed": true or false (true if the trace satisfies the question, false otherwise)
25 | 2. "reasoning": Brief explanation for your decision (2-3 sentences max)"""
26 |
27 |
28 | class AgentJudge:
29 | """An agent that evaluates the correctness of another agent's trace."""
30 |
31 | def __init__(
32 | self,
33 | model_id: str,
34 | framework: AgentFramework = AgentFramework.TINYAGENT,
35 | output_type: type[BaseModel] = EvaluationOutput,
36 | model_args: dict[str, Any] | None = None,
37 | ):
38 | self.model_id = model_id
39 | self.framework = framework
40 | self.model_args = model_args
41 | self.output_type = output_type
42 |
43 | def run(
44 | self,
45 | trace: AgentTrace,
46 | question: str,
47 | additional_tools: list[Callable[[], Any]] | None = None,
48 | ) -> AgentTrace:
49 | """Run the agent judge.
50 |
51 | Args:
52 | trace: The agent trace to evaluate
53 | question: The question to ask the agent
54 | additional_tools: Additional tools to use for the agent
55 |
56 | Returns:
57 | The trace of the evaluation run.
58 | You can access the evaluation result in the `final_output`
59 | property.
60 |
61 | """
62 | if additional_tools is None:
63 | additional_tools = []
64 | return run_async_in_sync(
65 | self.run_async(trace, question, additional_tools),
66 | allow_running_loop=INSIDE_NOTEBOOK,
67 | )
68 |
69 | async def run_async(
70 | self,
71 | trace: AgentTrace,
72 | question: str,
73 | additional_tools: list[Callable[[], Any]] | None = None,
74 | ) -> AgentTrace:
75 | """Run the agent judge asynchronously.
76 |
77 | Args:
78 | trace: The agent trace to evaluate
79 | question: The question to ask the agent
80 | additional_tools: Additional tools to use for the agent
81 | Returns:
82 | The trace of the evaluation run.
83 | You can access the evaluation result in the `final_output`
84 | property.
85 |
86 | """
87 | if additional_tools is None:
88 | additional_tools = []
89 | tooling = TraceTools(trace)
90 |
91 | agent_config = AgentConfig(
92 | model_id=self.model_id,
93 | instructions=AGENT_INSTRUCTIONS.format(
94 | response_schema=self.output_type.model_json_schema()
95 | ),
96 | tools=tooling.get_all_tools() + additional_tools,
97 | output_type=self.output_type,
98 | model_args=self.model_args,
99 | )
100 |
101 | agent = await AnyAgent.create_async(
102 | self.framework,
103 | agent_config=agent_config,
104 | )
105 | agent_trace = await agent.run_async(question)
106 | if not isinstance(agent_trace.final_output, self.output_type):
107 | msg = f"Agent output is not an {self.output_type} instance."
108 | raise ValueError(msg)
109 | return agent_trace
110 |
--------------------------------------------------------------------------------
/src/any_agent/callbacks/wrappers/llama_index.py:
--------------------------------------------------------------------------------
1 | # mypy: disable-error-code="method-assign,misc,no-untyped-call,no-untyped-def,union-attr"
2 | from __future__ import annotations
3 |
4 | from typing import TYPE_CHECKING, Any
5 |
6 | from opentelemetry.trace import get_current_span
7 |
8 | if TYPE_CHECKING:
9 | from any_agent.callbacks.context import Context
10 | from any_agent.frameworks.llama_index import LlamaIndexAgent
11 |
12 |
13 | class _LlamaIndexWrapper:
14 | def __init__(self) -> None:
15 | self.callback_context: dict[int, Context] = {}
16 | self._original_take_step: Any | None = None
17 | self._original_acalls: dict[str, Any] = {}
18 | self._original_llm_call: Any | None = None
19 |
20 | async def wrap(self, agent: LlamaIndexAgent) -> None:
21 | self._original_take_step = agent._agent.take_step
22 |
23 | async def wrap_take_step(*args, **kwargs):
24 | context = self.callback_context[
25 | get_current_span().get_span_context().trace_id
26 | ]
27 | context.shared["model_id"] = getattr(agent._agent.llm, "model", "No model")
28 |
29 | for callback in agent.config.callbacks:
30 | context = callback.before_llm_call(context, *args, **kwargs)
31 |
32 | output = await self._original_take_step( # type: ignore[misc]
33 | *args, **kwargs
34 | )
35 |
36 | for callback in agent.config.callbacks:
37 | context = callback.after_llm_call(context, output)
38 |
39 | return output
40 |
41 | # bypass Pydantic validation because _agent is a BaseModel
42 | agent._agent.model_config["extra"] = "allow"
43 | agent._agent.take_step = wrap_take_step
44 |
45 | async def wrap_tool_execution(original_call, metadata, *args, **kwargs):
46 | context = self.callback_context[
47 | get_current_span().get_span_context().trace_id
48 | ]
49 | context.shared["metadata"] = metadata
50 |
51 | for callback in agent.config.callbacks:
52 | context = callback.before_tool_execution(context, *args, **kwargs)
53 |
54 | output = await original_call(**kwargs)
55 |
56 | for callback in agent.config.callbacks:
57 | context = callback.after_tool_execution(context, output)
58 |
59 | return output
60 |
61 | class WrappedAcall:
62 | def __init__(self, metadata, original_acall):
63 | self.metadata = metadata
64 | self.original_acall = original_acall
65 |
66 | async def acall(self, *args, **kwargs):
67 | return await wrap_tool_execution(
68 | self.original_acall, self.metadata, **kwargs
69 | )
70 |
71 | for tool in agent._agent.tools:
72 | self._original_acalls[str(tool.metadata.name)] = tool.acall
73 | wrapped = WrappedAcall(tool.metadata, tool.acall)
74 | tool.acall = wrapped.acall
75 |
76 | # Wrap call_model to capture any-llm calls during structured output processing
77 | self._original_llm_call = agent.call_model
78 |
79 | async def wrap_call_model(**kwargs):
80 | context = self.callback_context[
81 | get_current_span().get_span_context().trace_id
82 | ]
83 |
84 | for callback in agent.config.callbacks:
85 | context = callback.before_llm_call(context, **kwargs)
86 |
87 | output = await self._original_llm_call(**kwargs)
88 |
89 | for callback in agent.config.callbacks:
90 | context = callback.after_llm_call(context, output)
91 |
92 | return output
93 |
94 | agent.call_model = wrap_call_model
95 |
96 | async def unwrap(self, agent: LlamaIndexAgent) -> None:
97 | if self._original_take_step:
98 | agent._agent.take_step = self._original_take_step
99 | if self._original_acalls:
100 | for tool in agent._agent.tools:
101 | tool.acall = self._original_acalls[str(tool.metadata.name)]
102 | if self._original_llm_call is not None:
103 | agent.call_model = self._original_llm_call
104 |
--------------------------------------------------------------------------------
/src/any_agent/tools/web_browsing.py:
--------------------------------------------------------------------------------
1 | import os
2 | import re
3 |
4 | import requests
5 | from requests.exceptions import RequestException
6 |
7 |
8 | def _truncate_content(content: str, max_length: int) -> str:
9 | if len(content) <= max_length:
10 | return content
11 | return (
12 | content[: max_length // 2]
13 | + f"\n..._This content has been truncated to stay below {max_length} characters_...\n"
14 | + content[-max_length // 2 :]
15 | )
16 |
17 |
18 | def search_web(query: str) -> str:
19 | """Perform a duckduckgo web search based on your query (think a Google search) then returns the top search results.
20 |
21 | Args:
22 | query (str): The search query to perform.
23 |
24 | Returns:
25 | The top search results.
26 |
27 | """
28 | try:
29 | from duckduckgo_search import DDGS # type: ignore[import-not-found]
30 | except ImportError as e:
31 | msg = "You need to `pip install 'duckduckgo_search'` to use this tool"
32 | raise ImportError(msg) from e
33 |
34 | ddgs = DDGS()
35 | results = ddgs.text(query, max_results=10)
36 | return "\n".join(
37 | f"[{result['title']}]({result['href']})\n{result['body']}" for result in results
38 | )
39 |
40 |
41 | def visit_webpage(url: str, timeout: int = 30, max_length: int = 10000) -> str:
42 | """Visits a webpage at the given url and reads its content as a markdown string. Use this to browse webpages.
43 |
44 | Args:
45 | url: The url of the webpage to visit.
46 | timeout: The timeout in seconds for the request.
47 | max_length: The maximum number of characters of text that can be returned (default=10000).
48 | If max_length==-1, text is not truncated and the full webpage is returned.
49 |
50 | """
51 | try:
52 | from markdownify import markdownify # type: ignore[import-not-found]
53 | except ImportError as e:
54 | msg = "You need to `pip install 'markdownify'` to use this tool"
55 | raise ImportError(msg) from e
56 |
57 | try:
58 | response = requests.get(url, timeout=timeout)
59 | response.raise_for_status()
60 |
61 | markdown_content = markdownify(response.text).strip()
62 |
63 | markdown_content = re.sub(r"\n{2,}", "\n", markdown_content)
64 |
65 | if max_length == -1:
66 | return str(markdown_content)
67 | return _truncate_content(markdown_content, max_length)
68 | except RequestException as e:
69 | return f"Error fetching the webpage: {e!s}"
70 | except Exception as e:
71 | return f"An unexpected error occurred: {e!s}"
72 |
73 |
74 | def search_tavily(query: str, include_images: bool = False) -> str:
75 | """Perform a Tavily web search based on your query and return the top search results.
76 |
77 | See https://blog.tavily.com/getting-started-with-the-tavily-search-api for more information.
78 |
79 | Args:
80 | query (str): The search query to perform.
81 | include_images (bool): Whether to include images in the results.
82 |
83 | Returns:
84 | The top search results as a formatted string.
85 |
86 | """
87 | try:
88 | from tavily.tavily import TavilyClient
89 | except ImportError as e:
90 | msg = "You need to `pip install 'tavily-python'` to use this tool"
91 | raise ImportError(msg) from e
92 |
93 | api_key = os.getenv("TAVILY_API_KEY")
94 | if not api_key:
95 | return "TAVILY_API_KEY environment variable not set."
96 | try:
97 | client = TavilyClient(api_key)
98 | response = client.search(query, include_images=include_images)
99 | results = response.get("results", [])
100 | output = []
101 | for result in results:
102 | output.append(
103 | f"[{result.get('title', 'No Title')}]({result.get('url', '#')})\n{result.get('content', '')}"
104 | )
105 | if include_images and "images" in response:
106 | output.append("\nImages:")
107 | for image in response["images"]:
108 | output.append(image)
109 | return "\n\n".join(output) if output else "No results found."
110 | except Exception as e:
111 | return f"Error performing Tavily search: {e!s}"
112 |
--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
1 | # Cookbook generated files
2 | docs/cookbook/sensitive-info
3 |
4 | uv.lock
5 | output/
6 | traces/
7 | # Byte-compiled / optimized / DLL files
8 | __pycache__/
9 | *.py[cod]
10 | *$py.class
11 |
12 | # C extensions
13 | *.so
14 |
15 | # Distribution / packaging
16 | .Python
17 | build/
18 | develop-eggs/
19 | dist/
20 | downloads/
21 | eggs/
22 | .eggs/
23 | lib/
24 | lib64/
25 | parts/
26 | sdist/
27 | var/
28 | wheels/
29 | share/python-wheels/
30 | *.egg-info/
31 | .installed.cfg
32 | *.egg
33 | MANIFEST
34 |
35 | # PyInstaller
36 | # Usually these files are written by a python script from a template
37 | # before PyInstaller builds the exe, so as to inject date/other infos into it.
38 | *.manifest
39 | *.spec
40 |
41 | # Installer logs
42 | pip-log.txt
43 | pip-delete-this-directory.txt
44 |
45 | # Unit test / coverage reports
46 | htmlcov/
47 | .tox/
48 | .nox/
49 | .coverage
50 | .coverage.*
51 | .cache
52 | nosetests.xml
53 | coverage.xml
54 | *.cover
55 | *.py,cover
56 | .hypothesis/
57 | .pytest_cache/
58 | cover/
59 |
60 | # Translations
61 | *.mo
62 | *.pot
63 |
64 | # Django stuff:
65 | *.log
66 | local_settings.py
67 | db.sqlite3
68 | db.sqlite3-journal
69 |
70 | # Flask stuff:
71 | instance/
72 | .webassets-cache
73 |
74 | # Scrapy stuff:
75 | .scrapy
76 |
77 | # Sphinx documentation
78 | docs/_build/
79 |
80 | # PyBuilder
81 | .pybuilder/
82 | target/
83 |
84 | # Jupyter Notebook
85 | .ipynb_checkpoints
86 |
87 | # IPython
88 | profile_default/
89 | ipython_config.py
90 |
91 | # pyenv
92 | # For a library or package, you might want to ignore these files since the code is
93 | # intended to run in multiple environments; otherwise, check them in:
94 | # .python-version
95 |
96 | # pipenv
97 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
98 | # However, in case of collaboration, if having platform-specific dependencies or dependencies
99 | # having no cross-platform support, pipenv may install dependencies that don't work, or not
100 | # install all needed dependencies.
101 | #Pipfile.lock
102 |
103 | # UV
104 | # Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
105 | # This is especially recommended for binary packages to ensure reproducibility, and is more
106 | # commonly ignored for libraries.
107 | #uv.lock
108 |
109 | # poetry
110 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
111 | # This is especially recommended for binary packages to ensure reproducibility, and is more
112 | # commonly ignored for libraries.
113 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
114 | #poetry.lock
115 |
116 | # pdm
117 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
118 | #pdm.lock
119 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
120 | # in version control.
121 | # https://pdm.fming.dev/latest/usage/project/#working-with-version-control
122 | .pdm.toml
123 | .pdm-python
124 | .pdm-build/
125 |
126 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
127 | __pypackages__/
128 |
129 | # Celery stuff
130 | celerybeat-schedule
131 | celerybeat.pid
132 |
133 | # SageMath parsed files
134 | *.sage.py
135 |
136 | # Environments
137 | .env
138 | .venv
139 | env/
140 | venv/
141 | ENV/
142 | env.bak/
143 | venv.bak/
144 |
145 | # Spyder project settings
146 | .spyderproject
147 | .spyproject
148 |
149 | # Rope project settings
150 | .ropeproject
151 |
152 | # mkdocs documentation
153 | /site
154 |
155 | # mypy
156 | .mypy_cache/
157 | .dmypy.json
158 | dmypy.json
159 |
160 | # Pyre type checker
161 | .pyre/
162 |
163 | # pytype static type analyzer
164 | .pytype/
165 |
166 | # Cython debug symbols
167 | cython_debug/
168 |
169 | # PyCharm
170 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can
171 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
172 | # and can be added to the global gitignore or merged into this file. For a more nuclear
173 | # option (not recommended) you can uncomment the following to ignore the entire idea folder.
174 | .idea/
175 |
176 | # Ruff stuff:
177 | .ruff_cache/
178 |
179 | # PyPI configuration file
180 | .pypirc
181 |
182 | .vscode
183 |
--------------------------------------------------------------------------------
/tests/integration/frameworks/test_error_handling.py:
--------------------------------------------------------------------------------
1 | from typing import Any
2 | from unittest.mock import patch
3 |
4 | import pytest
5 |
6 | from any_agent import (
7 | AgentConfig,
8 | AgentFramework,
9 | AgentRunError,
10 | AnyAgent,
11 | )
12 | from any_agent.callbacks import Callback, Context
13 | from any_agent.testing.helpers import (
14 | DEFAULT_SMALL_MODEL_ID,
15 | LLM_IMPORT_PATHS,
16 | get_default_agent_model_args,
17 | )
18 | from any_agent.tracing.otel_types import StatusCode
19 |
20 |
21 | class LimitLLMCalls(Callback):
22 | def __init__(self, max_llm_calls: int) -> None:
23 | self.max_llm_calls = max_llm_calls
24 |
25 | def before_llm_call(self, context: Context, *args: Any, **kwargs: Any) -> Context:
26 | if "n_llm_calls" not in context.shared:
27 | context.shared["n_llm_calls"] = 0
28 |
29 | context.shared["n_llm_calls"] += 1
30 |
31 | if context.shared["n_llm_calls"] > self.max_llm_calls:
32 | msg = "Reached limit of LLM Calls"
33 | raise RuntimeError(msg)
34 |
35 | return context
36 |
37 |
38 | def test_runtime_error(
39 | agent_framework: AgentFramework,
40 | ) -> None:
41 | """An exception not caught by the framework should be caught by us.
42 |
43 | `AnyAgent.run_async` should catch and reraise an `AgentRunError`.
44 |
45 | The `AgentRunError.trace` should be retrieved.
46 | """
47 | kwargs = {}
48 | test_runtime_error_msg = "runtime error trap"
49 |
50 | kwargs["model_id"] = DEFAULT_SMALL_MODEL_ID
51 |
52 | patch_function = LLM_IMPORT_PATHS.get(agent_framework)
53 | if not patch_function:
54 | err_msg = f"No patch function found for agent framework: {agent_framework}"
55 | raise ValueError(err_msg)
56 |
57 | with patch(patch_function) as llm_completion_path:
58 | llm_completion_path.side_effect = RuntimeError(test_runtime_error_msg)
59 | agent_config = AgentConfig(
60 | model_id=kwargs["model_id"],
61 | tools=[],
62 | model_args=get_default_agent_model_args(agent_framework),
63 | )
64 | agent = AnyAgent.create(agent_framework, agent_config)
65 | spans = []
66 | try:
67 | agent.run(
68 | "Write a four-line poem about agent frameworks.",
69 | )
70 | except AgentRunError as are:
71 | spans = are.trace.spans
72 | assert any(
73 | span.status.status_code == StatusCode.ERROR
74 | and span.status.description is not None
75 | and test_runtime_error_msg in span.status.description
76 | for span in spans
77 | )
78 |
79 |
80 | def test_tool_error(
81 | agent_framework: AgentFramework,
82 | ) -> None:
83 | """An exception raised inside a tool will be caught by us.
84 |
85 | We make sure an appropriate Status is set to the tool execution span.
86 | We allow the Agent to try to recover from the tool calling failure.
87 | """
88 | exception_reason = "tool error trap"
89 |
90 | def search_web(query: str) -> str:
91 | """Perform a duckduckgo web search based on your query then returns the top search results.
92 |
93 | Args:
94 | query (str): The search query to perform.
95 |
96 | Returns:
97 | The top search results.
98 |
99 | """
100 | msg = exception_reason
101 | raise ValueError(msg)
102 |
103 | kwargs = {}
104 |
105 | kwargs["model_id"] = DEFAULT_SMALL_MODEL_ID
106 |
107 | agent_config = AgentConfig(
108 | model_id=kwargs["model_id"],
109 | instructions="You must use the available tools to answer questions.",
110 | tools=[search_web],
111 | model_args=get_default_agent_model_args(agent_framework),
112 | callbacks=[LimitLLMCalls(max_llm_calls=5)],
113 | )
114 |
115 | agent = AnyAgent.create(agent_framework, agent_config)
116 | agent_trace = agent.run(
117 | "Check in the web which agent framework is the best. If the tool fails, don't try again, return final answer as failure.",
118 | )
119 | assert any(
120 | span.is_tool_execution()
121 | and span.status.status_code == StatusCode.ERROR
122 | and exception_reason in getattr(span.status, "description", "")
123 | for span in agent_trace.spans
124 | )
125 |
--------------------------------------------------------------------------------
/tests/integration/a2a/test_a2a_tool.py:
--------------------------------------------------------------------------------
1 | from typing import Any
2 | import pytest
3 |
4 | from any_agent import AgentConfig, AgentFramework, AnyAgent
5 | from any_agent.callbacks import Callback, Context
6 | from any_agent.serving import A2AServingConfig
7 | from any_agent.testing.helpers import (
8 | DEFAULT_HTTP_KWARGS,
9 | DEFAULT_SMALL_MODEL_ID,
10 | get_default_agent_model_args,
11 | )
12 | from any_agent.tools import a2a_tool_async
13 | from any_agent.tracing.agent_trace import AgentTrace
14 | from any_agent.tracing.attributes import GenAI
15 |
16 | from .conftest import (
17 | DATE_PROMPT,
18 | a2a_client_from_agent,
19 | assert_contains_current_date_info,
20 | get_datetime,
21 | )
22 |
23 |
24 | class LimitLLMCalls(Callback):
25 | def __init__(self, max_llm_calls: int) -> None:
26 | self.max_llm_calls = max_llm_calls
27 |
28 | def before_llm_call(self, context: Context, *args: Any, **kwargs: Any) -> Context:
29 | if "n_llm_calls" not in context.shared:
30 | context.shared["n_llm_calls"] = 0
31 |
32 | context.shared["n_llm_calls"] += 1
33 |
34 | if context.shared["n_llm_calls"] > self.max_llm_calls:
35 | msg = "Reached limit of LLM Calls"
36 | raise RuntimeError(msg)
37 |
38 | return context
39 |
40 |
41 | def _assert_valid_agent_trace(agent_trace: AgentTrace) -> None:
42 | """Assert that agent_trace is valid and has final output."""
43 | assert isinstance(agent_trace, AgentTrace)
44 | assert agent_trace.final_output
45 |
46 |
47 | def _assert_has_date_agent_tool_call(agent_trace: AgentTrace) -> None:
48 | """Assert that the agent trace contains a tool execution span for the date agent."""
49 | assert any(
50 | span.is_tool_execution()
51 | and span.attributes.get(GenAI.TOOL_NAME, None) == "call_date_agent"
52 | for span in agent_trace.spans
53 | )
54 |
55 |
56 | @pytest.mark.asyncio
57 | async def test_a2a_tool_async(agent_framework: AgentFramework) -> None:
58 | """Tests that an agent contacts another using A2A using the adapter tool.
59 |
60 | Note that there is an issue when using Google ADK: https://github.com/google/adk-python/pull/566
61 | """
62 | skip_reason = {
63 | AgentFramework.SMOLAGENTS: "async a2a is not supported",
64 | }
65 | if agent_framework in skip_reason:
66 | pytest.skip(
67 | f"Framework {agent_framework}, reason: {skip_reason[agent_framework]}"
68 | )
69 |
70 | date_agent_cfg = AgentConfig(
71 | instructions="Use the available tools to obtain additional information to answer the query.",
72 | name="date_agent",
73 | model_id=DEFAULT_SMALL_MODEL_ID,
74 | description="Agent that can return the current date.",
75 | tools=[get_datetime],
76 | model_args=get_default_agent_model_args(agent_framework),
77 | callbacks=[LimitLLMCalls(max_llm_calls=10)],
78 | )
79 | date_agent = await AnyAgent.create_async(
80 | agent_framework=agent_framework,
81 | agent_config=date_agent_cfg,
82 | )
83 |
84 | # Serve the agent and get client
85 | tool_agent_endpoint = "tool_agent"
86 | serving_config = A2AServingConfig(
87 | port=0,
88 | endpoint=f"/{tool_agent_endpoint}",
89 | log_level="info",
90 | )
91 |
92 | async with a2a_client_from_agent(date_agent, serving_config) as (_, server_url):
93 | # Create main agent with A2A tool
94 | main_agent_cfg = AgentConfig(
95 | instructions="Use the available tools to obtain additional information to answer the query.",
96 | description="The orchestrator that can use other agents via tools using the A2A protocol.",
97 | model_id=DEFAULT_SMALL_MODEL_ID,
98 | tools=[await a2a_tool_async(server_url, http_kwargs=DEFAULT_HTTP_KWARGS)],
99 | model_args=get_default_agent_model_args(agent_framework),
100 | callbacks=[LimitLLMCalls(max_llm_calls=10)],
101 | )
102 |
103 | main_agent = await AnyAgent.create_async(
104 | agent_framework=agent_framework,
105 | agent_config=main_agent_cfg,
106 | )
107 |
108 | agent_trace = await main_agent.run_async(DATE_PROMPT)
109 |
110 | _assert_valid_agent_trace(agent_trace)
111 | assert_contains_current_date_info(str(agent_trace.final_output))
112 | _assert_has_date_agent_tool_call(agent_trace)
113 |
--------------------------------------------------------------------------------
/tests/unit/tools/mcp/test_mcp_client.py:
--------------------------------------------------------------------------------
1 | import inspect
2 | from typing import Optional
3 | from unittest.mock import AsyncMock
4 |
5 | import pytest
6 |
7 | from any_agent.config import AgentFramework, MCPStdio
8 | from any_agent.tools.mcp.mcp_client import MCPClient
9 |
10 |
11 | @pytest.mark.asyncio
12 | async def test_create_tool_function_with_complex_schema() -> None:
13 | """Test the core tool function creation with comprehensive parameter handling."""
14 | from mcp.types import Tool as MCPTool
15 |
16 | mock_session = AsyncMock()
17 | mock_result = AsyncMock()
18 | mock_result.content = [AsyncMock()]
19 | mock_result.content[0].text = "Tool executed successfully"
20 | mock_session.call_tool.return_value = mock_result
21 |
22 | config = MCPStdio(command="test", args=[])
23 | client = MCPClient(config=config, framework=AgentFramework.OPENAI)
24 | client._session = mock_session
25 |
26 | complex_tool = MCPTool(
27 | name="complex_search",
28 | description="Search with multiple parameter types",
29 | inputSchema={
30 | "type": "object",
31 | "properties": {
32 | "query": {"type": "string", "description": "Search query string"},
33 | "max_results": {
34 | "type": "integer",
35 | "description": "Maximum number of results",
36 | },
37 | "include_metadata": {
38 | "type": "boolean",
39 | "description": "Include result metadata",
40 | },
41 | "filters": {"type": "object", "description": "Search filters"},
42 | "tags": {"type": "array", "description": "Filter tags"},
43 | "threshold": {"type": "number", "description": "Similarity threshold"},
44 | "optional_param": {
45 | "type": "string",
46 | "description": "Optional parameter",
47 | },
48 | },
49 | "required": ["query", "max_results", "include_metadata"],
50 | },
51 | )
52 |
53 | tool_func = client._create_tool_function(complex_tool)
54 |
55 | assert tool_func.__name__ == "complex_search"
56 | assert tool_func.__doc__ is not None
57 | assert "Search with multiple parameter types" in tool_func.__doc__
58 | assert "query: Search query string" in tool_func.__doc__
59 | assert "max_results: Maximum number of results" in tool_func.__doc__
60 |
61 | sig = inspect.signature(tool_func)
62 | params = sig.parameters
63 |
64 | assert "query" in params
65 | assert params["query"].annotation is str
66 | assert params["query"].default is inspect.Parameter.empty
67 |
68 | assert "max_results" in params
69 | assert params["max_results"].annotation is int
70 | assert params["max_results"].default is inspect.Parameter.empty
71 |
72 | assert "include_metadata" in params
73 | assert params["include_metadata"].annotation is bool
74 | assert params["include_metadata"].default is inspect.Parameter.empty
75 |
76 | assert "optional_param" in params
77 | assert params["optional_param"].annotation == Optional[str] # noqa: UP045
78 | assert params["optional_param"].default is None
79 |
80 | assert params["filters"].annotation == Optional[dict] # noqa: UP045
81 | assert params["tags"].annotation == Optional[list] # noqa: UP045
82 | assert params["threshold"].annotation == Optional[float] # noqa: UP045
83 |
84 | assert sig.return_annotation is str
85 |
86 | result = await tool_func(query="test query", max_results=10, include_metadata=True)
87 | assert result == "Tool executed successfully"
88 | mock_session.call_tool.assert_called_with(
89 | "complex_search",
90 | {"query": "test query", "max_results": 10, "include_metadata": True},
91 | )
92 |
93 | await tool_func(
94 | query="another query",
95 | max_results=5,
96 | include_metadata=False,
97 | optional_param="optional_value",
98 | threshold=0.8,
99 | )
100 | mock_session.call_tool.assert_called_with(
101 | "complex_search",
102 | {
103 | "query": "another query",
104 | "max_results": 5,
105 | "include_metadata": False,
106 | "optional_param": "optional_value",
107 | "threshold": 0.8,
108 | },
109 | )
110 |
111 | client._session = None
112 | error_result = await tool_func(query="test", max_results=1, include_metadata=True)
113 | assert "Error: MCP session not available" in error_result
114 |
--------------------------------------------------------------------------------