├── 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 | Project logo 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 | --------------------------------------------------------------------------------