ant-ai 1.0.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (63) hide show
  1. ant_ai-1.0.0/LICENSE +21 -0
  2. ant_ai-1.0.0/PKG-INFO +148 -0
  3. ant_ai-1.0.0/README.md +116 -0
  4. ant_ai-1.0.0/pyproject.toml +68 -0
  5. ant_ai-1.0.0/src/ant_ai/__init__.py +43 -0
  6. ant_ai-1.0.0/src/ant_ai/a2a/__init__.py +23 -0
  7. ant_ai-1.0.0/src/ant_ai/a2a/agent.py +177 -0
  8. ant_ai-1.0.0/src/ant_ai/a2a/client.py +166 -0
  9. ant_ai-1.0.0/src/ant_ai/a2a/colony.py +230 -0
  10. ant_ai-1.0.0/src/ant_ai/a2a/config.py +54 -0
  11. ant_ai-1.0.0/src/ant_ai/a2a/context_builder.py +139 -0
  12. ant_ai-1.0.0/src/ant_ai/a2a/executor.py +145 -0
  13. ant_ai-1.0.0/src/ant_ai/a2a/server.py +102 -0
  14. ant_ai-1.0.0/src/ant_ai/a2a/session.py +3 -0
  15. ant_ai-1.0.0/src/ant_ai/a2a/translator.py +179 -0
  16. ant_ai-1.0.0/src/ant_ai/a2a/types.py +11 -0
  17. ant_ai-1.0.0/src/ant_ai/agent/__init__.py +7 -0
  18. ant_ai-1.0.0/src/ant_ai/agent/agent.py +32 -0
  19. ant_ai-1.0.0/src/ant_ai/agent/base.py +228 -0
  20. ant_ai-1.0.0/src/ant_ai/agent/loop/__init__.py +9 -0
  21. ant_ai-1.0.0/src/ant_ai/agent/loop/loop.py +191 -0
  22. ant_ai-1.0.0/src/ant_ai/agent/loop/react.py +226 -0
  23. ant_ai-1.0.0/src/ant_ai/core/__init__.py +86 -0
  24. ant_ai-1.0.0/src/ant_ai/core/events.py +153 -0
  25. ant_ai-1.0.0/src/ant_ai/core/exceptions.py +6 -0
  26. ant_ai-1.0.0/src/ant_ai/core/logging.py +94 -0
  27. ant_ai-1.0.0/src/ant_ai/core/message.py +96 -0
  28. ant_ai-1.0.0/src/ant_ai/core/response.py +39 -0
  29. ant_ai-1.0.0/src/ant_ai/core/result.py +118 -0
  30. ant_ai-1.0.0/src/ant_ai/core/types.py +42 -0
  31. ant_ai-1.0.0/src/ant_ai/hooks/__init__.py +21 -0
  32. ant_ai-1.0.0/src/ant_ai/hooks/adapters/__init__.py +3 -0
  33. ant_ai-1.0.0/src/ant_ai/hooks/adapters/guardrails_ai.py +53 -0
  34. ant_ai-1.0.0/src/ant_ai/hooks/layer.py +172 -0
  35. ant_ai-1.0.0/src/ant_ai/hooks/protocol.py +149 -0
  36. ant_ai-1.0.0/src/ant_ai/llm/__init__.py +5 -0
  37. ant_ai-1.0.0/src/ant_ai/llm/integrations/__init__.py +7 -0
  38. ant_ai-1.0.0/src/ant_ai/llm/integrations/lite_llm.py +140 -0
  39. ant_ai-1.0.0/src/ant_ai/llm/integrations/openai_llm.py +110 -0
  40. ant_ai-1.0.0/src/ant_ai/llm/protocol.py +75 -0
  41. ant_ai-1.0.0/src/ant_ai/observer/__init__.py +9 -0
  42. ant_ai-1.0.0/src/ant_ai/observer/composite.py +58 -0
  43. ant_ai-1.0.0/src/ant_ai/observer/integrations/__init__.py +9 -0
  44. ant_ai-1.0.0/src/ant_ai/observer/integrations/langfuse.py +274 -0
  45. ant_ai-1.0.0/src/ant_ai/observer/integrations/log.py +28 -0
  46. ant_ai-1.0.0/src/ant_ai/observer/integrations/otel.py +74 -0
  47. ant_ai-1.0.0/src/ant_ai/observer/obs.py +134 -0
  48. ant_ai-1.0.0/src/ant_ai/observer/protocol.py +76 -0
  49. ant_ai-1.0.0/src/ant_ai/py.typed +0 -0
  50. ant_ai-1.0.0/src/ant_ai/steps/__init__.py +9 -0
  51. ant_ai-1.0.0/src/ant_ai/steps/llm_step.py +101 -0
  52. ant_ai-1.0.0/src/ant_ai/steps/protocol.py +36 -0
  53. ant_ai-1.0.0/src/ant_ai/steps/tool_step.py +208 -0
  54. ant_ai-1.0.0/src/ant_ai/tools/__init__.py +9 -0
  55. ant_ai-1.0.0/src/ant_ai/tools/builtins/__init__.py +9 -0
  56. ant_ai-1.0.0/src/ant_ai/tools/builtins/filesystem_tool.py +111 -0
  57. ant_ai-1.0.0/src/ant_ai/tools/builtins/human_input.py +7 -0
  58. ant_ai-1.0.0/src/ant_ai/tools/builtins/shell_tool.py +108 -0
  59. ant_ai-1.0.0/src/ant_ai/tools/registry.py +81 -0
  60. ant_ai-1.0.0/src/ant_ai/tools/tool.py +462 -0
  61. ant_ai-1.0.0/src/ant_ai/workflow/__init__.py +10 -0
  62. ant_ai-1.0.0/src/ant_ai/workflow/action.py +23 -0
  63. ant_ai-1.0.0/src/ant_ai/workflow/workflow.py +296 -0
ant_ai-1.0.0/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 IDeA Group at IDSIA
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
ant_ai-1.0.0/PKG-INFO ADDED
@@ -0,0 +1,148 @@
1
+ Metadata-Version: 2.4
2
+ Name: ant-ai
3
+ Version: 1.0.0
4
+ Summary: ANT AI
5
+ Author: Cezar Sas, Vincenzo Giuffrida, Sandra Mitrovic, Matteo Salani
6
+ Author-email: Cezar Sas <cezar.sas@supsi.ch>, Vincenzo Giuffrida <vincenzo.giuffrida@supsi.ch>, Sandra Mitrovic <sandra.mitrovic@supsi.ch>, Matteo Salani <matteo.salani@supsi.ch>
7
+ License-Expression: MIT
8
+ License-File: LICENSE
9
+ Classifier: Intended Audience :: Developers
10
+ Classifier: Intended Audience :: Information Technology
11
+ Classifier: Operating System :: OS Independent
12
+ Classifier: Programming Language :: Python
13
+ Classifier: Programming Language :: Python :: 3
14
+ Classifier: Programming Language :: Python :: 3.14
15
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
16
+ Classifier: Topic :: Software Development :: Libraries :: Python Modules
17
+ Classifier: Typing :: Typed
18
+ Requires-Dist: a2a-sdk[sql,http-server,grpc,encryption,telemetry]>=1.0.0
19
+ Requires-Dist: fastapi>=0.136.0
20
+ Requires-Dist: litellm>=1.83.10
21
+ Requires-Dist: loguru>=0.7.3
22
+ Requires-Dist: mcp>=1.27.0
23
+ Requires-Dist: pydantic>=2.12.5
24
+ Requires-Dist: starlette>=1.0.0
25
+ Requires-Dist: uvicorn>=0.44.0
26
+ Requires-Dist: langfuse>=4.3.1 ; extra == 'langfuse'
27
+ Requires-Dist: openai>=2.24.0 ; extra == 'openai'
28
+ Requires-Python: >=3.14
29
+ Provides-Extra: langfuse
30
+ Provides-Extra: openai
31
+ Description-Content-Type: text/markdown
32
+
33
+ <div align="center">
34
+
35
+ <picture>
36
+ <source media="(prefers-color-scheme: dark)" srcset="docs/assets/ant_h_white.png">
37
+ <img alt="ANT AI" src="docs/assets/ant_h_dark.png" height="100">
38
+ </picture>
39
+
40
+ ![Python](https://img.shields.io/badge/python-3.14%2B-4584b6?logo=python&logoColor=white)
41
+ ![License](https://img.shields.io/badge/license-MIT-blue?logo=MIT&logoColor=white-lightgrey)
42
+ [![Coverage main](https://gitlab-core.supsi.ch/dti-idsia/intsys/ant-ai/badges/main/coverage.svg)](https://gitlab-core.supsi.ch/dti-idsia/intsys/ant-ai/-/commits/main)
43
+ [![Docs](https://img.shields.io/badge/docs-mkdocs-526cfe?logo=materialformkdocs&logoColor=white)](https://ant-ai-27f99d.pages-core.supsi.ch)
44
+
45
+ **A lightweight Python framework for building tool-driven AI agents and multi-agent systems.**
46
+
47
+ </div>
48
+
49
+ ---
50
+
51
+ ANT AI provides a composable set of primitives for building production-ready AI agents: a ReAct reasoning loop, a flexible tool system with MCP support, a graph-based workflow engine, and first-class agent-to-agent (A2A) communication via the [A2A protocol](https://github.com/a2aproject/A2A).
52
+
53
+ ## Features
54
+
55
+ - **ReAct agent** — built-in Reason→Act loop with streaming, structured output, and configurable retry logic
56
+ - **Flexible tools** — define tools as decorated functions, class namespaces, or load them directly from any [MCP](https://modelcontextprotocol.io/) server
57
+ - **Workflow engine** — graph-based orchestration with static and conditional edges to sequence agent behaviour predictably
58
+ - **Multi-agent colony** — wire agents together with the A2A protocol; each agent becomes a callable tool to its peers
59
+ - **LLM-agnostic** — ships with [LiteLLM](https://github.com/BerriAI/litellm) and native OpenAI backends; any `ChatLLM`-conforming implementation works
60
+ - **Observability** — structured lifecycle events with [Langfuse](https://langfuse.com/), OpenTelemetry, and log sinks
61
+ - **Lifecycle hooks** — intercept and control every LLM call: pass, block, retry, or substitute results; ships with a [GuardrailsAI](https://www.guardrailsai.com/) adapter
62
+
63
+ ## Installation
64
+
65
+ Requires Python 3.14+. Install with [uv](https://docs.astral.sh/uv/):
66
+
67
+ ```sh
68
+ uv add ant-ai
69
+ ```
70
+
71
+ Or clone and sync for local development:
72
+
73
+ ```sh
74
+ git clone <repo-url>
75
+ cd ant-ai
76
+ uv sync --all-extras
77
+ ```
78
+
79
+ ## Quickstart
80
+
81
+ ### Single agent
82
+
83
+ ```python
84
+ from ant_ai import Agent, Message, State, tool
85
+ from ant_ai.llm.integrations import LiteLLMChat
86
+
87
+ @tool
88
+ def get_weather(city: str) -> str:
89
+ """Return the current weather for a city."""
90
+ return f"Sunny, 22°C in {city}"
91
+
92
+ llm = LiteLLMChat(model="gpt-4o-mini")
93
+
94
+ agent = Agent(
95
+ name="WeatherAgent",
96
+ system_prompt="You are a helpful weather assistant.",
97
+ llm=llm,
98
+ tools=[get_weather],
99
+ )
100
+
101
+ state = State(messages=[Message(role="user", content="What's the weather in Lugano?")])
102
+ answer = agent.invoke(state)
103
+ print(answer)
104
+ ```
105
+
106
+ ### Streaming events
107
+
108
+ ```python
109
+ from ant_ai.core import FinalAnswerEvent
110
+
111
+ async for event in agent.stream(state):
112
+ if isinstance(event, FinalAnswerEvent):
113
+ print(event.content)
114
+ ```
115
+
116
+ ### Structured output
117
+
118
+ ```python
119
+ from pydantic import BaseModel
120
+
121
+ class WeatherReport(BaseModel):
122
+ city: str
123
+ temperature: int
124
+ condition: str
125
+
126
+ answer = agent.invoke(state, response_schema=WeatherReport)
127
+ # answer is a JSON string matching WeatherReport
128
+ ```
129
+
130
+ ## Development
131
+
132
+ ```sh
133
+ # Install dev dependencies and pre-commit hooks
134
+ uv sync --all-extras
135
+ uv run pre-commit install
136
+
137
+ # Run tests
138
+ uv run pytest
139
+
140
+ # Serve docs locally
141
+ uv run mkdocs serve
142
+ ```
143
+
144
+ See [CONTRIBUTING.md](CONTRIBUTING.md) for the full contributing guide, branching model, and review process.
145
+
146
+ ## License
147
+
148
+ This software is licensed under the MIT license. See the [LICENSE](LICENSE) file for details.
ant_ai-1.0.0/README.md ADDED
@@ -0,0 +1,116 @@
1
+ <div align="center">
2
+
3
+ <picture>
4
+ <source media="(prefers-color-scheme: dark)" srcset="docs/assets/ant_h_white.png">
5
+ <img alt="ANT AI" src="docs/assets/ant_h_dark.png" height="100">
6
+ </picture>
7
+
8
+ ![Python](https://img.shields.io/badge/python-3.14%2B-4584b6?logo=python&logoColor=white)
9
+ ![License](https://img.shields.io/badge/license-MIT-blue?logo=MIT&logoColor=white-lightgrey)
10
+ [![Coverage main](https://gitlab-core.supsi.ch/dti-idsia/intsys/ant-ai/badges/main/coverage.svg)](https://gitlab-core.supsi.ch/dti-idsia/intsys/ant-ai/-/commits/main)
11
+ [![Docs](https://img.shields.io/badge/docs-mkdocs-526cfe?logo=materialformkdocs&logoColor=white)](https://ant-ai-27f99d.pages-core.supsi.ch)
12
+
13
+ **A lightweight Python framework for building tool-driven AI agents and multi-agent systems.**
14
+
15
+ </div>
16
+
17
+ ---
18
+
19
+ ANT AI provides a composable set of primitives for building production-ready AI agents: a ReAct reasoning loop, a flexible tool system with MCP support, a graph-based workflow engine, and first-class agent-to-agent (A2A) communication via the [A2A protocol](https://github.com/a2aproject/A2A).
20
+
21
+ ## Features
22
+
23
+ - **ReAct agent** — built-in Reason→Act loop with streaming, structured output, and configurable retry logic
24
+ - **Flexible tools** — define tools as decorated functions, class namespaces, or load them directly from any [MCP](https://modelcontextprotocol.io/) server
25
+ - **Workflow engine** — graph-based orchestration with static and conditional edges to sequence agent behaviour predictably
26
+ - **Multi-agent colony** — wire agents together with the A2A protocol; each agent becomes a callable tool to its peers
27
+ - **LLM-agnostic** — ships with [LiteLLM](https://github.com/BerriAI/litellm) and native OpenAI backends; any `ChatLLM`-conforming implementation works
28
+ - **Observability** — structured lifecycle events with [Langfuse](https://langfuse.com/), OpenTelemetry, and log sinks
29
+ - **Lifecycle hooks** — intercept and control every LLM call: pass, block, retry, or substitute results; ships with a [GuardrailsAI](https://www.guardrailsai.com/) adapter
30
+
31
+ ## Installation
32
+
33
+ Requires Python 3.14+. Install with [uv](https://docs.astral.sh/uv/):
34
+
35
+ ```sh
36
+ uv add ant-ai
37
+ ```
38
+
39
+ Or clone and sync for local development:
40
+
41
+ ```sh
42
+ git clone <repo-url>
43
+ cd ant-ai
44
+ uv sync --all-extras
45
+ ```
46
+
47
+ ## Quickstart
48
+
49
+ ### Single agent
50
+
51
+ ```python
52
+ from ant_ai import Agent, Message, State, tool
53
+ from ant_ai.llm.integrations import LiteLLMChat
54
+
55
+ @tool
56
+ def get_weather(city: str) -> str:
57
+ """Return the current weather for a city."""
58
+ return f"Sunny, 22°C in {city}"
59
+
60
+ llm = LiteLLMChat(model="gpt-4o-mini")
61
+
62
+ agent = Agent(
63
+ name="WeatherAgent",
64
+ system_prompt="You are a helpful weather assistant.",
65
+ llm=llm,
66
+ tools=[get_weather],
67
+ )
68
+
69
+ state = State(messages=[Message(role="user", content="What's the weather in Lugano?")])
70
+ answer = agent.invoke(state)
71
+ print(answer)
72
+ ```
73
+
74
+ ### Streaming events
75
+
76
+ ```python
77
+ from ant_ai.core import FinalAnswerEvent
78
+
79
+ async for event in agent.stream(state):
80
+ if isinstance(event, FinalAnswerEvent):
81
+ print(event.content)
82
+ ```
83
+
84
+ ### Structured output
85
+
86
+ ```python
87
+ from pydantic import BaseModel
88
+
89
+ class WeatherReport(BaseModel):
90
+ city: str
91
+ temperature: int
92
+ condition: str
93
+
94
+ answer = agent.invoke(state, response_schema=WeatherReport)
95
+ # answer is a JSON string matching WeatherReport
96
+ ```
97
+
98
+ ## Development
99
+
100
+ ```sh
101
+ # Install dev dependencies and pre-commit hooks
102
+ uv sync --all-extras
103
+ uv run pre-commit install
104
+
105
+ # Run tests
106
+ uv run pytest
107
+
108
+ # Serve docs locally
109
+ uv run mkdocs serve
110
+ ```
111
+
112
+ See [CONTRIBUTING.md](CONTRIBUTING.md) for the full contributing guide, branching model, and review process.
113
+
114
+ ## License
115
+
116
+ This software is licensed under the MIT license. See the [LICENSE](LICENSE) file for details.
@@ -0,0 +1,68 @@
1
+ [build-system]
2
+ requires = ["uv_build>=0.11.7,<0.12"]
3
+ build-backend = "uv_build"
4
+
5
+ [project]
6
+ name = "ant-ai"
7
+ version = "1.0.0"
8
+ description = "ANT AI"
9
+ readme = "README.md"
10
+ license = "MIT"
11
+ license-files = ["LICENSE"]
12
+ requires-python = ">=3.14"
13
+ authors = [
14
+ { name = "Cezar Sas", email = "cezar.sas@supsi.ch" },
15
+ { name = "Vincenzo Giuffrida", email = "vincenzo.giuffrida@supsi.ch" },
16
+ { name = "Sandra Mitrovic", email = "sandra.mitrovic@supsi.ch" },
17
+ { name = "Matteo Salani", email = "matteo.salani@supsi.ch" },
18
+ ]
19
+
20
+ classifiers = [
21
+ "Intended Audience :: Developers",
22
+ "Intended Audience :: Information Technology",
23
+ "Operating System :: OS Independent",
24
+ "Programming Language :: Python",
25
+ "Programming Language :: Python :: 3",
26
+ "Programming Language :: Python :: 3.14",
27
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
28
+ "Topic :: Software Development :: Libraries :: Python Modules",
29
+ "Typing :: Typed",
30
+ ]
31
+ dependencies = [
32
+ "a2a-sdk[sql,http-server,grpc,encryption,telemetry]>=1.0.0",
33
+ "fastapi>=0.136.0",
34
+ "litellm>=1.83.10",
35
+ "loguru>=0.7.3",
36
+ "mcp>=1.27.0",
37
+ "pydantic>=2.12.5",
38
+ "starlette>=1.0.0",
39
+ "uvicorn>=0.44.0",
40
+ ]
41
+
42
+ [project.optional-dependencies]
43
+ langfuse = ["langfuse>=4.3.1"]
44
+ openai = ["openai>=2.24.0"]
45
+
46
+ [dependency-groups]
47
+ dev = [
48
+ "pre-commit>=4.5.1",
49
+ "pytest>=9.0.3",
50
+ "pytest-asyncio>=1.3.0",
51
+ "pytest-cov>=7.1.0",
52
+ "ruff>=0.15.11",
53
+ "ty>=0.0.32",
54
+ ]
55
+ docs = [
56
+ "griffe-pydantic>=1.3.1",
57
+ "mkdocs>=1.6.1",
58
+ "mkdocs-awesome-nav>=3.3.0",
59
+ "mkdocs-gen-files>=0.6.1",
60
+ "mkdocs-literate-nav>=0.6.3",
61
+ "mkdocs-material>=9.7.6",
62
+ "mkdocs-panzoom-plugin>=0.5.2",
63
+ "mkdocs-section-index>=0.3.12",
64
+ "mkdocstrings[python]>=1.0.4",
65
+ ]
66
+
67
+ [tool.uv]
68
+ exclude-newer = "P2D" # Wait for the package to be up for at least 2 days before updating to it.
@@ -0,0 +1,43 @@
1
+ from ant_ai.agent import Agent, BaseAgent
2
+ from ant_ai.core import (
3
+ AnyEvent,
4
+ AnyMessage,
5
+ ChatLLMResponse,
6
+ Event,
7
+ InvocationContext,
8
+ Message,
9
+ State,
10
+ StepResult,
11
+ configure_logging,
12
+ )
13
+ from ant_ai.observer import CompositeSink, ObservabilitySink, obs
14
+ from ant_ai.tools import Tool, ToolRegistry
15
+ from ant_ai.tools.tool import tool
16
+ from ant_ai.workflow import BaseAction, Workflow
17
+
18
+ __all__ = [
19
+ # agent
20
+ "Agent",
21
+ "BaseAgent",
22
+ # core
23
+ "Message",
24
+ "AnyMessage",
25
+ "Event",
26
+ "AnyEvent",
27
+ "State",
28
+ "InvocationContext",
29
+ "ChatLLMResponse",
30
+ "StepResult",
31
+ "configure_logging",
32
+ # observer
33
+ "obs",
34
+ "ObservabilitySink",
35
+ "CompositeSink",
36
+ # tools
37
+ "Tool",
38
+ "tool",
39
+ "ToolRegistry",
40
+ # workflow
41
+ "Workflow",
42
+ "BaseAction",
43
+ ]
@@ -0,0 +1,23 @@
1
+ from ant_ai.a2a.agent import A2AAgentTool
2
+ from ant_ai.a2a.client import A2AClient, AgentClientError
3
+ from ant_ai.a2a.colony import AgentSpec, Colony
4
+ from ant_ai.a2a.config import A2AConfig
5
+ from ant_ai.a2a.context_builder import HistoryRequestContextBuilder
6
+ from ant_ai.a2a.executor import A2AExecutor
7
+ from ant_ai.a2a.server import A2AServer
8
+ from ant_ai.a2a.session import current_session_id
9
+ from ant_ai.a2a.types import A2AMetadata
10
+
11
+ __all__ = [
12
+ "A2AAgentTool",
13
+ "A2AClient",
14
+ "AgentClientError",
15
+ "Colony",
16
+ "AgentSpec",
17
+ "A2AConfig",
18
+ "HistoryRequestContextBuilder",
19
+ "A2AExecutor",
20
+ "A2AServer",
21
+ "current_session_id",
22
+ "A2AMetadata",
23
+ ]
@@ -0,0 +1,177 @@
1
+ from __future__ import annotations
2
+
3
+ from collections.abc import Awaitable
4
+ from typing import Any, Literal, overload
5
+
6
+ from a2a.types import AgentCard
7
+ from pydantic import Field, PrivateAttr, model_validator
8
+
9
+ from ant_ai.a2a.client import A2AClient
10
+ from ant_ai.a2a.config import A2AConfig
11
+ from ant_ai.a2a.session import current_session_id
12
+ from ant_ai.core.events import ClarificationNeededEvent, FinalAnswerEvent
13
+ from ant_ai.tools.tool import Tool
14
+
15
+
16
+ class A2AAgentTool(Tool):
17
+ """
18
+ A tool that provides an interface to interact with the Agent via the A2A protocol.
19
+ """
20
+
21
+ config: A2AConfig = Field(..., description="Configuration for the A2A client.")
22
+ agent_input_description: str = Field(
23
+ default="Message to send to the agent. Contains all the necessary information to answer the question in a clear way.",
24
+ description="Description of the input to the agent. This description will be used by the agent to generate the prompt for remote agent request.",
25
+ )
26
+
27
+ _a2a: A2AClient | None = PrivateAttr(default=None)
28
+ _initialized: bool = PrivateAttr(default=False)
29
+ _agent_card: AgentCard | None = PrivateAttr(default=None)
30
+ _last_task_id: str | None = PrivateAttr(default=None)
31
+
32
+ @model_validator(mode="after")
33
+ def _set_defaults(self) -> A2AAgentTool:
34
+ return self
35
+
36
+ def _ensure_a2a(self) -> None:
37
+ if self._a2a is None:
38
+ self._a2a = A2AClient(config=self.config)
39
+ self._a2a._agent_card: AgentCard | None = self._agent_card
40
+
41
+ def _init_metadata(self, agent_card: AgentCard) -> None:
42
+ """Set Tool metadata exactly once from an AgentCard."""
43
+ if not self.name:
44
+ self.name: str = agent_card.name
45
+ if not self.description:
46
+ self.description: str = self._create_agent_description(agent_card)
47
+
48
+ self.parameters: dict[str, Any] = {
49
+ "type": "object",
50
+ "properties": {
51
+ "message": {
52
+ "type": "string",
53
+ "description": (self.agent_input_description),
54
+ }
55
+ },
56
+ "required": ["message"],
57
+ }
58
+
59
+ async def _ensure_initialized(self) -> None:
60
+ """Fetch AgentCard (if needed) and set metadata/_func exactly once."""
61
+ if self._initialized:
62
+ return
63
+
64
+ self._ensure_a2a()
65
+ if self._a2a is None:
66
+ raise RuntimeError("A2A client not initialized")
67
+
68
+ agent_card: AgentCard = self._agent_card or await self._a2a.get_agent_card()
69
+ self._agent_card: AgentCard = agent_card
70
+ self._init_metadata(agent_card)
71
+
72
+ self._attach_func()
73
+ self._initialized = True
74
+
75
+ def _attach_func(self) -> None:
76
+ """Attach the call function to the _func (single callable Tool)."""
77
+
78
+ async def _call_remote(message: str) -> str:
79
+ await self._ensure_initialized()
80
+ self._ensure_a2a()
81
+ if self._a2a is None:
82
+ raise RuntimeError("A2A client not initialized")
83
+
84
+ last_text: str = ""
85
+ async for ev in self._a2a.send_message(
86
+ message, context_id=current_session_id.get(None)
87
+ ):
88
+ if ev.content:
89
+ last_text: str = ev.content
90
+ if isinstance(ev, (FinalAnswerEvent, ClarificationNeededEvent)):
91
+ break
92
+ return last_text
93
+
94
+ self._func = _call_remote
95
+
96
+ @overload
97
+ @classmethod
98
+ def from_config(cls, config: A2AConfig, agent_card: AgentCard) -> A2AAgentTool:
99
+ """Creates an A2A agent tool from a configuration and an agent card.
100
+
101
+ Returns:
102
+ An A2A agent tool.
103
+ """
104
+ ...
105
+
106
+ @overload
107
+ @classmethod
108
+ def from_config(
109
+ cls, config: A2AConfig, agent_card: None = None
110
+ ) -> Awaitable[A2AAgentTool]:
111
+ """Creates an A2A agent tool from a configuration.
112
+
113
+ Args:
114
+ config: _description_
115
+ agent_card: _description_. Defaults to None.
116
+
117
+ Returns:
118
+ An awaitable of an A2A agent tool.
119
+ """
120
+ ...
121
+
122
+ @classmethod
123
+ def from_config(
124
+ cls, config: A2AConfig, agent_card: AgentCard | None = None
125
+ ) -> A2AAgentTool | Awaitable[A2AAgentTool]:
126
+ """Creates an A2A agent tool from a configuration and an optional agent card. If no agent card is provided, the tool will be initialized asynchronously and the tool will be returned as an awaitable tool.
127
+
128
+ Args:
129
+ config: _description_
130
+ agent_card: _description_. Defaults to None.
131
+
132
+ Returns:
133
+ An A2A agent tool or an awaitable of an A2A agent tool.
134
+ """
135
+ tool: A2AAgentTool = cls(name=None, description=None, config=config)
136
+
137
+ if agent_card is not None:
138
+ tool._agent_card: AgentCard = agent_card
139
+ tool._ensure_a2a()
140
+ tool._init_metadata(agent_card)
141
+ tool._attach_func()
142
+ tool._initialized = True
143
+ return tool
144
+
145
+ async def _build() -> A2AAgentTool:
146
+ await tool._ensure_initialized()
147
+ return tool
148
+
149
+ return _build()
150
+
151
+ def _create_agent_description(self, agent_card: AgentCard) -> str:
152
+ parts: list[str] = [
153
+ agent_card.description,
154
+ ]
155
+
156
+ if agent_card.skills:
157
+ parts += ["", "### Available Skills", ""]
158
+ for skill in agent_card.skills:
159
+ parts.append(f"**{skill.name}**")
160
+ parts.append(skill.description)
161
+ if skill.tags:
162
+ parts.append(f"Tags: {', '.join(skill.tags)}")
163
+ examples = list(skill.examples) if skill.examples else []
164
+ if examples:
165
+ parts.append("Examples:")
166
+ for ex in examples:
167
+ parts.append(f" - {ex}")
168
+ parts.append("")
169
+
170
+ return "\n".join(parts)
171
+
172
+ @property
173
+ def is_namespace(self) -> Literal[False]:
174
+ return False
175
+
176
+ def _sid(self) -> str:
177
+ return current_session_id.get()