tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-05 08:23:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Feat] Use A2A registered agents with /chat/completions (#20362)

Browse Source 
* test_a2a_registry_integration

* fix: render agents on model dropdown on UI

* init append_agents_to_model_group

* route_a2a_agent_request

* is_a2a_agent_model

* route_a2a_agent_request

* fix: error handling

* docs A2A usage

* docs fix

* feat: working A2a streaming

* fix transform
This commit is contained in:
Ishaan Jaff
2026-02-03 15:25:38 -08:00
committed by GitHub
parent 59cab4d2aa
commit d267c69086
14 changed files with 860 additions and 124 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/my-website/docs/a2a.md
+3 -110
View File
@@ -68,116 +68,9 @@ Follow [this guide, to add your pydantic ai agent to LiteLLM Agent Gateway](./pr
## Invoking your Agents
Use the [A2A Python SDK](https://pypi.org/project/a2a-sdk) to invoke agents through LiteLLM.
This example shows how to:
1. **List available agents** - Query `/v1/agents` to see which agents your key can access
2. **Select an agent** - Pick an agent from the list
3. **Invoke via A2A** - Use the A2A protocol to send messages to the agent
```python showLineNumbers title="invoke_a2a_agent.py"
from uuid import uuid4
import httpx
import asyncio
from a2a.client import A2ACardResolver, A2AClient
from a2a.types import MessageSendParams, SendMessageRequest
# === CONFIGURE THESE ===
LITELLM_BASE_URL = "http://localhost:4000" # Your LiteLLM proxy URL
LITELLM_VIRTUAL_KEY = "sk-1234" # Your LiteLLM Virtual Key
# =======================
async def main():
headers = {"Authorization": f"Bearer {LITELLM_VIRTUAL_KEY}"}
async with httpx.AsyncClient(headers=headers) as client:
# Step 1: List available agents
response = await client.get(f"{LITELLM_BASE_URL}/v1/agents")
agents = response.json()
print("Available agents:")
for agent in agents:
print(f" - {agent['agent_name']} (ID: {agent['agent_id']})")
if not agents:
print("No agents available for this key")
return
# Step 2: Select an agent and invoke it
selected_agent = agents[0]
agent_id = selected_agent["agent_id"]
agent_name = selected_agent["agent_name"]
print(f"\nInvoking: {agent_name}")
# Step 3: Use A2A protocol to invoke the agent
base_url = f"{LITELLM_BASE_URL}/a2a/{agent_id}"
resolver = A2ACardResolver(httpx_client=client, base_url=base_url)
agent_card = await resolver.get_agent_card()
a2a_client = A2AClient(httpx_client=client, agent_card=agent_card)
request = SendMessageRequest(
id=str(uuid4()),
params=MessageSendParams(
message={
"role": "user",
"parts": [{"kind": "text", "text": "Hello, what can you do?"}],
"messageId": uuid4().hex,
}
),
)
response = await a2a_client.send_message(request)
print(f"Response: {response.model_dump(mode='json', exclude_none=True, indent=4)}")
if __name__ == "__main__":
asyncio.run(main())
```
### Streaming Responses
For streaming responses, use `send_message_streaming`:
```python showLineNumbers title="invoke_a2a_agent_streaming.py"
from uuid import uuid4
import httpx
import asyncio
from a2a.client import A2ACardResolver, A2AClient
from a2a.types import MessageSendParams, SendStreamingMessageRequest
# === CONFIGURE THESE ===
LITELLM_BASE_URL = "http://localhost:4000" # Your LiteLLM proxy URL
LITELLM_VIRTUAL_KEY = "sk-1234" # Your LiteLLM Virtual Key
LITELLM_AGENT_NAME = "ij-local" # Agent name registered in LiteLLM
# =======================
async def main():
base_url = f"{LITELLM_BASE_URL}/a2a/{LITELLM_AGENT_NAME}"
headers = {"Authorization": f"Bearer {LITELLM_VIRTUAL_KEY}"}
async with httpx.AsyncClient(headers=headers) as httpx_client:
# Resolve agent card and create client
resolver = A2ACardResolver(httpx_client=httpx_client, base_url=base_url)
agent_card = await resolver.get_agent_card()
client = A2AClient(httpx_client=httpx_client, agent_card=agent_card)
# Send a streaming message
request = SendStreamingMessageRequest(
id=str(uuid4()),
params=MessageSendParams(
message={
"role": "user",
"parts": [{"kind": "text", "text": "Hello, what can you do?"}],
"messageId": uuid4().hex,
}
),
)
# Stream the response
async for chunk in client.send_message_streaming(request):
print(chunk.model_dump(mode="json", exclude_none=True))
if __name__ == "__main__":
asyncio.run(main())
```
See the [Invoking A2A Agents](./a2a_invoking_agents) guide to learn how to call your agents using:
- **A2A SDK** - Native A2A protocol with full support for tasks and artifacts
- **OpenAI SDK** - Familiar `/chat/completions` interface with `a2a/` model prefix
## Tracking Agent Logs
docs/my-website/docs/a2a_invoking_agents.md
+280
View File
@@ -0,0 +1,280 @@
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
# Invoking A2A Agents
Learn how to invoke A2A agents through LiteLLM using different methods.
:::tip Deploy Your Own A2A Agent
Want to test with your own agent? Deploy this template A2A agent powered by Google Gemini:
[**shin-bot-litellm/a2a-gemini-agent**](https://github.com/shin-bot-litellm/a2a-gemini-agent) - Simple deployable A2A agent with streaming support
:::
## A2A SDK
Use the [A2A Python SDK](https://pypi.org/project/a2a-sdk) to invoke agents through LiteLLM using the A2A protocol.
### Non-Streaming
This example shows how to:
1. **List available agents** - Query `/v1/agents` to see which agents your key can access
2. **Select an agent** - Pick an agent from the list
3. **Invoke via A2A** - Use the A2A protocol to send messages to the agent
```python showLineNumbers title="invoke_a2a_agent.py"
from uuid import uuid4
import httpx
import asyncio
from a2a.client import A2ACardResolver, A2AClient
from a2a.types import MessageSendParams, SendMessageRequest
# === CONFIGURE THESE ===
LITELLM_BASE_URL = "http://localhost:4000" # Your LiteLLM proxy URL
LITELLM_VIRTUAL_KEY = "sk-1234" # Your LiteLLM Virtual Key
# =======================
async def main():
headers = {"Authorization": f"Bearer {LITELLM_VIRTUAL_KEY}"}
async with httpx.AsyncClient(headers=headers) as client:
# Step 1: List available agents
response = await client.get(f"{LITELLM_BASE_URL}/v1/agents")
agents = response.json()
print("Available agents:")
for agent in agents:
print(f" - {agent['agent_name']} (ID: {agent['agent_id']})")
if not agents:
print("No agents available for this key")
return
# Step 2: Select an agent and invoke it
selected_agent = agents[0]
agent_id = selected_agent["agent_id"]
agent_name = selected_agent["agent_name"]
print(f"\nInvoking: {agent_name}")
# Step 3: Use A2A protocol to invoke the agent
base_url = f"{LITELLM_BASE_URL}/a2a/{agent_id}"
resolver = A2ACardResolver(httpx_client=client, base_url=base_url)
agent_card = await resolver.get_agent_card()
a2a_client = A2AClient(httpx_client=client, agent_card=agent_card)
request = SendMessageRequest(
id=str(uuid4()),
params=MessageSendParams(
message={
"role": "user",
"parts": [{"kind": "text", "text": "Hello, what can you do?"}],
"messageId": uuid4().hex,
}
),
)
response = await a2a_client.send_message(request)
print(f"Response: {response.model_dump(mode='json', exclude_none=True, indent=4)}")
if __name__ == "__main__":
asyncio.run(main())
```
### Streaming
For streaming responses, use `send_message_streaming`:
```python showLineNumbers title="invoke_a2a_agent_streaming.py"
from uuid import uuid4
import httpx
import asyncio
from a2a.client import A2ACardResolver, A2AClient
from a2a.types import MessageSendParams, SendStreamingMessageRequest
# === CONFIGURE THESE ===
LITELLM_BASE_URL = "http://localhost:4000" # Your LiteLLM proxy URL
LITELLM_VIRTUAL_KEY = "sk-1234" # Your LiteLLM Virtual Key
LITELLM_AGENT_NAME = "ij-local" # Agent name registered in LiteLLM
# =======================
async def main():
base_url = f"{LITELLM_BASE_URL}/a2a/{LITELLM_AGENT_NAME}"
headers = {"Authorization": f"Bearer {LITELLM_VIRTUAL_KEY}"}
async with httpx.AsyncClient(headers=headers) as httpx_client:
# Resolve agent card and create client
resolver = A2ACardResolver(httpx_client=httpx_client, base_url=base_url)
agent_card = await resolver.get_agent_card()
client = A2AClient(httpx_client=httpx_client, agent_card=agent_card)
# Send a streaming message
request = SendStreamingMessageRequest(
id=str(uuid4()),
params=MessageSendParams(
message={
"role": "user",
"parts": [{"kind": "text", "text": "Tell me a long story"}],
"messageId": uuid4().hex,
}
),
)
# Stream the response
async for chunk in client.send_message_streaming(request):
print(chunk.model_dump(mode="json", exclude_none=True))
if __name__ == "__main__":
asyncio.run(main())
```
## /chat/completions API (OpenAI SDK)
You can also invoke A2A agents using the familiar OpenAI SDK by using the `a2a/` model prefix.
### Non-Streaming
<Tabs>
<TabItem value="python" label="Python" default>
```python showLineNumbers title="openai_non_streaming.py"
import openai
client = openai.OpenAI(
api_key="sk-1234", # Your LiteLLM Virtual Key
base_url="http://localhost:4000" # Your LiteLLM proxy URL
)
response = client.chat.completions.create(
model="a2a/my-agent", # Use a2a/ prefix with your agent name
messages=[
{"role": "user", "content": "Hello, what can you do?"}
]
)
print(response.choices[0].message.content)
```
</TabItem>
<TabItem value="typescript" label="TypeScript">
```typescript showLineNumbers title="openai_non_streaming.ts"
import OpenAI from 'openai';
const client = new OpenAI({
apiKey: 'sk-1234', // Your LiteLLM Virtual Key
baseURL: 'http://localhost:4000' // Your LiteLLM proxy URL
});
const response = await client.chat.completions.create({
model: 'a2a/my-agent', // Use a2a/ prefix with your agent name
messages: [
{ role: 'user', content: 'Hello, what can you do?' }
]
});
console.log(response.choices[0].message.content);
```
</TabItem>
<TabItem value="curl" label="cURL">
```bash showLineNumbers title="curl_non_streaming.sh"
curl -X POST http://localhost:4000/v1/chat/completions \
-H "Authorization: Bearer sk-1234" \
-H "Content-Type: application/json" \
-d '{
"model": "a2a/my-agent",
"messages": [
{"role": "user", "content": "Hello, what can you do?"}
]
}'
```
</TabItem>
</Tabs>
### Streaming
<Tabs>
<TabItem value="python" label="Python" default>
```python showLineNumbers title="openai_streaming.py"
import openai
client = openai.OpenAI(
api_key="sk-1234", # Your LiteLLM Virtual Key
base_url="http://localhost:4000" # Your LiteLLM proxy URL
)
stream = client.chat.completions.create(
model="a2a/my-agent", # Use a2a/ prefix with your agent name
messages=[
{"role": "user", "content": "Tell me a long story"}
],
stream=True
)
for chunk in stream:
if chunk.choices[0].delta.content:
print(chunk.choices[0].delta.content, end="", flush=True)
```
</TabItem>
<TabItem value="typescript" label="TypeScript">
```typescript showLineNumbers title="openai_streaming.ts"
import OpenAI from 'openai';
const client = new OpenAI({
apiKey: 'sk-1234', // Your LiteLLM Virtual Key
baseURL: 'http://localhost:4000' // Your LiteLLM proxy URL
});
const stream = await client.chat.completions.create({
model: 'a2a/my-agent', // Use a2a/ prefix with your agent name
messages: [
{ role: 'user', content: 'Tell me a long story' }
],
stream: true
});
for await (const chunk of stream) {
const content = chunk.choices[0]?.delta?.content;
if (content) {
process.stdout.write(content);
}
}
```
</TabItem>
<TabItem value="curl" label="cURL">
```bash showLineNumbers title="curl_streaming.sh"
curl -X POST http://localhost:4000/v1/chat/completions \
-H "Authorization: Bearer sk-1234" \
-H "Content-Type: application/json" \
-d '{
"model": "a2a/my-agent",
"messages": [
{"role": "user", "content": "Tell me a long story"}
],
"stream": true
}'
```
</TabItem>
</Tabs>
## Key Differences
| Method | Use Case | Advantages |
|--------|----------|------------|
| **A2A SDK** | Native A2A protocol integration | • Full A2A protocol support<br/>• Access to task states and artifacts<br/>• Context management |
| **OpenAI SDK** | Familiar OpenAI-style interface | • Drop-in replacement for OpenAI calls<br/>• Easier migration from LLM to agent workflows<br/>• Works with existing OpenAI tooling |
:::tip Model Prefix
When using the OpenAI SDK, always prefix your agent name with `a2a/` (e.g., `a2a/my-agent`) to route requests to the A2A agent instead of an LLM provider.
:::
docs/my-website/sidebars.js
+1
View File
@@ -469,6 +469,7 @@ const sidebars = {
label: "/a2a - A2A Agent Gateway",
items: [
"a2a",
"a2a_invoking_agents",
"a2a_cost_tracking",
"a2a_agent_permissions"
],
litellm/llms/a2a/chat/streaming_iterator.py
+1 -1
View File
@@ -94,7 +94,7 @@ class A2AModelResponseIterator(BaseModelResponseIterator):
if state == "completed":
return "stop"
elif state == "failed":
return "error"
return "stop" # Map failed state to 'stop' (valid finish_reason)
# Check for [DONE] marker
if chunk.get("done") is True:
litellm/llms/a2a/chat/transformation.py
+73 -6
View File
@@ -2,10 +2,9 @@
A2A Protocol Transformation for LiteLLM
"""
import uuid
from typing import Any, Dict, Iterator, List, Optional, Union, cast
from typing import Any, Dict, Iterator, List, Optional, Union
import httpx
from pydantic import BaseModel
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
@@ -27,6 +26,68 @@ class A2AConfig(BaseConfig):
Handles transformation between OpenAI and A2A JSON-RPC 2.0 formats.
"""
@staticmethod
def resolve_agent_config_from_registry(
model: str,
api_base: Optional[str],
api_key: Optional[str],
headers: Optional[Dict[str, Any]],
optional_params: Dict[str, Any],
) -> tuple[Optional[str], Optional[str], Optional[Dict[str, Any]]]:
"""
Resolve agent configuration from registry if model format is "a2a/<agent-name>".
Extracts agent name from model string and looks up configuration in the
agent registry (if available in proxy context).
Args:
model: Model string (e.g., "a2a/my-agent")
api_base: Explicit api_base (takes precedence over registry)
api_key: Explicit api_key (takes precedence over registry)
headers: Explicit headers (takes precedence over registry)
optional_params: Dict to merge additional litellm_params into
Returns:
Tuple of (api_base, api_key, headers) with registry values filled in
"""
# Extract agent name from model (e.g., "a2a/my-agent" -> "my-agent")
agent_name = model.split("/", 1)[1] if "/" in model else None
# Only lookup if agent name exists and some config is missing
if not agent_name or (api_base is not None and api_key is not None and headers is not None):
return api_base, api_key, headers
# Try registry lookup (only available in proxy context)
try:
from litellm.proxy.agent_endpoints.agent_registry import (
global_agent_registry,
)
agent = global_agent_registry.get_agent_by_name(agent_name)
if agent:
# Get api_base from agent card URL
if api_base is None and agent.agent_card_params:
api_base = agent.agent_card_params.get("url")
# Get api_key, headers, and other params from litellm_params
if agent.litellm_params:
if api_key is None:
api_key = agent.litellm_params.get("api_key")
if headers is None:
agent_headers = agent.litellm_params.get("headers")
if agent_headers:
headers = agent_headers
# Merge other litellm_params (timeout, max_retries, etc.)
for key, value in agent.litellm_params.items():
if key not in ["api_key", "api_base", "headers", "model"] and key not in optional_params:
optional_params[key] = value
except ImportError:
pass # Registry not available (not running in proxy context)
return api_base, api_key, headers
def get_supported_openai_params(self, model: str) -> List[str]:
"""Return list of supported OpenAI parameters"""
return [
@@ -46,9 +107,14 @@ class A2AConfig(BaseConfig):
"""
Map OpenAI parameters to A2A parameters.
For A2A protocol, we don't need to map most parameters since
they're handled in the transform_request method.
For A2A protocol, we need to map the stream parameter so
transform_request can determine which JSON-RPC method to use.
"""
# Map stream parameter
for param, value in non_default_params.items():
if param == "stream" and value is True:
optional_params["stream"] = value
return optional_params
def validate_environment(
@@ -160,8 +226,9 @@ class A2AConfig(BaseConfig):
# Build JSON-RPC 2.0 request
# For A2A protocol, the method is "message/send" for non-streaming
# and "message/stream" for streaming (handled by optional_params["stream"])
method = "message/stream" if optional_params.get("stream") else "message/send"
# and "message/stream" for streaming
stream = optional_params.get("stream", False)
method = "message/stream" if stream else "message/send"
request_data = {
"jsonrpc": "2.0",
litellm/llms/a2a/common_utils.py
+19 -1
View File
@@ -113,6 +113,8 @@ def extract_text_from_a2a_response(
# 1. Direct message: {"result": {"kind": "message", "parts": [...]}}
# 2. Nested message: {"result": {"message": {"parts": [...]}}}
# 3. Task with artifacts: {"result": {"kind": "task", "artifacts": [{"parts": [...]}]}}
# 4. Task with status message: {"result": {"kind": "task", "status": {"message": {"parts": [...]}}}}
# 5. Streaming artifact-update: {"result": {"kind": "artifact-update", "artifact": {"parts": [...]}}}
# Check if result itself has parts (direct message)
if "parts" in result:
@@ -123,7 +125,23 @@ def extract_text_from_a2a_response(
if message:
return extract_text_from_a2a_message(message, depth=0, max_depth=max_depth)
# Handle task result with artifacts
# Check for streaming artifact-update (singular artifact)
artifact = result.get("artifact")
if artifact and isinstance(artifact, dict):
return extract_text_from_a2a_message(
artifact, depth=0, max_depth=max_depth
)
# Check for task status message (common in Gemini A2A agents)
status = result.get("status", {})
if isinstance(status, dict):
status_message = status.get("message")
if status_message:
return extract_text_from_a2a_message(
status_message, depth=0, max_depth=max_depth
)
# Handle task result with artifacts (plural, array)
artifacts = result.get("artifacts", [])
if artifacts and len(artifacts) > 0:
first_artifact = artifacts[0]
litellm/main.py
+16 -6
View File
@@ -2201,14 +2201,24 @@ def completion( # type: ignore # noqa: PLR0915
)
elif custom_llm_provider == "a2a":
# A2A (Agent-to-Agent) Protocol
api_base = (
api_base
or litellm.api_base
or get_secret_str("A2A_API_BASE")
# Resolve agent configuration from registry if model format is "a2a/<agent-name>"
api_base, api_key, headers = litellm.A2AConfig.resolve_agent_config_from_registry(
model=model,
api_base=api_base,
api_key=api_key,
headers=headers,
optional_params=optional_params,
)
# Fall back to environment variables and defaults
api_base = api_base or litellm.api_base or get_secret_str("A2A_API_BASE")
if api_base is None:
raise Exception("api_base is required for A2A provider")
raise Exception(
"api_base is required for A2A provider. "
"Either provide api_base parameter, set A2A_API_BASE environment variable, "
"or register the agent in the proxy with model='a2a/<agent-name>'."
)
headers = headers or litellm.headers
litellm/proxy/agent_endpoints/a2a_routing.py
+53
View File
@@ -0,0 +1,53 @@
"""
A2A Agent Routing
Handles routing for A2A agents (models with "a2a/<agent-name>" prefix).
Looks up agents in the registry and injects their API base URL.
"""
from typing import Any, Optional
import litellm
from litellm._logging import verbose_proxy_logger
async def route_a2a_agent_request(data: dict, route_type: str) -> Optional[Any]:
"""
Route A2A agent requests directly to litellm with injected API base.
Returns None if not an A2A request (allows normal routing to continue).
"""
# Import here to avoid circular imports
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
from litellm.proxy.route_llm_request import (
ROUTE_ENDPOINT_MAPPING,
ProxyModelNotFoundError,
)
model_name = data.get("model", "")
# Check if this is an A2A agent request
if not isinstance(model_name, str) or not model_name.startswith("a2a/"):
return None
# Extract agent name (e.g., "a2a/my-agent" -> "my-agent")
agent_name = model_name[4:]
# Look up agent in registry
agent = global_agent_registry.get_agent_by_name(agent_name)
if agent is None:
verbose_proxy_logger.error(f"[A2A] Agent '{agent_name}' not found in registry")
route_name = ROUTE_ENDPOINT_MAPPING.get(route_type, route_type)
raise ProxyModelNotFoundError(route=route_name, model_name=model_name)
# Get API base URL from agent config
if not agent.agent_card_params or "url" not in agent.agent_card_params:
verbose_proxy_logger.error(f"[A2A] Agent '{agent_name}' has no URL configured")
route_name = ROUTE_ENDPOINT_MAPPING.get(route_type, route_type)
raise ProxyModelNotFoundError(route=route_name, model_name=model_name)
# Inject API base and route to litellm
data["api_base"] = agent.agent_card_params["url"]
verbose_proxy_logger.debug(f"[A2A] Routing {model_name} to {data['api_base']}")
return getattr(litellm, f"{route_type}")(**data)
litellm/proxy/agent_endpoints/model_list_helpers.py
+96
View File
@@ -0,0 +1,96 @@
"""
Helper functions for appending A2A agents to model lists.
Used by proxy model endpoints to make agents appear in UI alongside models.
"""
from typing import List
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
ModelGroupInfoProxy,
)
async def append_agents_to_model_group(
model_groups: List[ModelGroupInfoProxy],
user_api_key_dict: UserAPIKeyAuth,
) -> List[ModelGroupInfoProxy]:
"""
Append A2A agents to model groups list for UI display.
Converts agents to model format with "a2a/<agent-name>" naming
so they appear in playground and work with LiteLLM routing.
"""
try:
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
AgentRequestHandler,
)
allowed_agent_ids = await AgentRequestHandler.get_allowed_agents(
user_api_key_auth=user_api_key_dict
)
for agent_id in allowed_agent_ids:
agent = global_agent_registry.get_agent_by_id(agent_id)
if agent is not None:
model_groups.append(
ModelGroupInfoProxy(
model_group=f"a2a/{agent.agent_name}",
mode="chat",
providers=["a2a"],
)
)
except Exception as e:
verbose_proxy_logger.debug(
f"Error appending agents to model_group/info: {e}"
)
return model_groups
async def append_agents_to_model_info(
models: List[dict],
user_api_key_dict: UserAPIKeyAuth,
) -> List[dict]:
"""
Append A2A agents to model info list for UI display.
Converts agents to model format with "a2a/<agent-name>" naming
so they appear in models page and work with LiteLLM routing.
"""
try:
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
AgentRequestHandler,
)
allowed_agent_ids = await AgentRequestHandler.get_allowed_agents(
user_api_key_auth=user_api_key_dict
)
for agent_id in allowed_agent_ids:
agent = global_agent_registry.get_agent_by_id(agent_id)
if agent is not None:
models.append({
"model_name": f"a2a/{agent.agent_name}",
"litellm_params": {
"model": f"a2a/{agent.agent_name}",
"custom_llm_provider": "a2a",
},
"model_info": {
"id": agent.agent_id,
"mode": "chat",
"db_model": True,
"created_by": agent.created_by,
"created_at": agent.created_at,
"updated_at": agent.updated_at,
},
})
except Exception as e:
verbose_proxy_logger.debug(
f"Error appending agents to v2/model/info: {e}"
)
return models
litellm/proxy/proxy_server.py
+19
View File
@@ -239,6 +239,10 @@ from litellm.proxy._types import *
from litellm.proxy.agent_endpoints.a2a_endpoints import router as a2a_router
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
from litellm.proxy.agent_endpoints.endpoints import router as agent_endpoints_router
from litellm.proxy.agent_endpoints.model_list_helpers import (
append_agents_to_model_group,
append_agents_to_model_info,
)
from litellm.proxy.analytics_endpoints.analytics_endpoints import (
router as analytics_router,
)
@@ -8616,6 +8620,15 @@ async def model_info_v2(
)
verbose_proxy_logger.debug("all_models: %s", all_models)
# Append A2A agents to models list
all_models = await append_agents_to_model_info(
models=all_models,
user_api_key_dict=user_api_key_dict,
)
# Update total count to include agents
search_total_count = len(all_models)
return _paginate_models_response(
all_models=all_models,
@@ -9456,6 +9469,12 @@ async def model_group_info(
model_groups: List[ModelGroupInfoProxy] = _get_model_group_info(
llm_router=llm_router, all_models_str=all_models_str, model_group=model_group
)
# Append A2A agents to model groups
model_groups = await append_agents_to_model_group(
model_groups=model_groups,
user_api_key_dict=user_api_key_dict,
)
return {"data": model_groups}
litellm/proxy/route_llm_request.py
+11
View File
@@ -12,6 +12,11 @@ else:
LitellmRouter = Any
def _is_a2a_agent_model(model_name: Any) -> bool:
"""Check if the model name is for an A2A agent (a2a/ prefix)."""
return isinstance(model_name, str) and model_name.startswith("a2a/")
ROUTE_ENDPOINT_MAPPING = {
"acompletion": "/chat/completions",
"atext_completion": "/completions",
@@ -322,6 +327,12 @@ async def route_request(
except Exception:
# If router fails (e.g., model not found in router), fall back to direct call
return getattr(litellm, f"{route_type}")(**data)
elif _is_a2a_agent_model(data.get("model", "")):
from litellm.proxy.agent_endpoints.a2a_routing import (
route_a2a_agent_request,
)
return await route_a2a_agent_request(data, route_type)
elif user_model is not None:
return getattr(litellm, f"{route_type}")(**data)
tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py
+110
View File
@@ -0,0 +1,110 @@
"""
Test appending A2A agents to model lists.
Maps to: litellm/proxy/agent_endpoints/model_list_helpers.py
"""
import os
import sys
sys.path.insert(0, os.path.abspath("../../../.."))
from unittest.mock import AsyncMock, Mock, patch
import pytest
from litellm.proxy.agent_endpoints.model_list_helpers import (
append_agents_to_model_group,
append_agents_to_model_info,
)
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.types.agents import AgentResponse
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
ModelGroupInfoProxy,
)
@pytest.mark.asyncio
async def test_append_agents_to_model_group():
"""Test agents are converted to model group format with a2a/ prefix"""
# Mock agent data
mock_agent = AgentResponse(
agent_id="test-agent-id",
agent_name="my-agent",
agent_card_params={"url": "http://example.com"},
litellm_params=None,
)
# Mock AgentRequestHandler at its source location
mock_get_allowed_agents = AsyncMock(return_value=["test-agent-id"])
# Mock global_agent_registry
mock_registry = Mock()
mock_registry.get_agent_by_id = Mock(return_value=mock_agent)
with patch(
"litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents",
mock_get_allowed_agents,
):
with patch(
"litellm.proxy.agent_endpoints.agent_registry.global_agent_registry",
mock_registry,
):
model_groups = []
user_api_key_dict = Mock(spec=UserAPIKeyAuth)
result = await append_agents_to_model_group(
model_groups=model_groups,
user_api_key_dict=user_api_key_dict,
)
# Verify agent was converted with a2a/ prefix
assert len(result) == 1
assert result[0].model_group == "a2a/my-agent"
assert result[0].mode == "chat"
assert result[0].providers == ["a2a"]
@pytest.mark.asyncio
async def test_append_agents_to_model_info():
"""Test agents are converted to model info format with a2a/ prefix"""
# Mock agent data
mock_agent = AgentResponse(
agent_id="agent-123",
agent_name="test-agent",
agent_card_params={"url": "http://example.com"},
litellm_params=None,
created_by="user-123",
)
# Mock AgentRequestHandler at its source location
mock_get_allowed_agents = AsyncMock(return_value=["agent-123"])
# Mock global_agent_registry
mock_registry = Mock()
mock_registry.get_agent_by_id = Mock(return_value=mock_agent)
with patch(
"litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents",
mock_get_allowed_agents,
):
with patch(
"litellm.proxy.agent_endpoints.agent_registry.global_agent_registry",
mock_registry,
):
models = []
user_api_key_dict = Mock(spec=UserAPIKeyAuth)
result = await append_agents_to_model_info(
models=models,
user_api_key_dict=user_api_key_dict,
)
# Verify agent was converted with a2a/ prefix
assert len(result) == 1
assert result[0]["model_name"] == "a2a/test-agent"
assert result[0]["litellm_params"]["model"] == "a2a/test-agent"
assert result[0]["litellm_params"]["custom_llm_provider"] == "a2a"
assert result[0]["model_info"]["id"] == "agent-123"
assert result[0]["model_info"]["mode"] == "chat"
tests/test_litellm/proxy/test_route_a2a_models.py
+105
View File
@@ -0,0 +1,105 @@
"""
Test A2A model routing in proxy.
Maps to: litellm/proxy/agent_endpoints/a2a_routing.py
"""
import os
import sys
sys.path.insert(0, os.path.abspath("../../.."))
from unittest.mock import AsyncMock, Mock, patch
import pytest
from litellm.proxy.agent_endpoints.a2a_routing import route_a2a_agent_request
from litellm.proxy.route_llm_request import route_request
@pytest.mark.asyncio
async def test_route_a2a_model_bypasses_router():
"""Test that a2a/ prefixed models bypass router and go directly to litellm with api_base"""
# Mock data for chat completion with a2a model
data = {
"model": "a2a/test-agent",
"messages": [{"role": "user", "content": "Hello"}],
}
# Mock router that doesn't have the a2a model
mock_router = Mock()
mock_router.model_names = ["gpt-4", "gpt-3.5-turbo"]
mock_router.deployment_names = []
mock_router.has_model_id = Mock(return_value=False)
mock_router.model_group_alias = None
mock_router.router_general_settings = Mock(pass_through_all_models=False)
mock_router.default_deployment = None
mock_router.pattern_router = Mock(patterns=[])
mock_router.map_team_model = Mock(return_value=None)
# Mock agent in registry
from litellm.types.agents import AgentResponse
mock_agent = AgentResponse(
agent_id="test-agent-id",
agent_name="test-agent",
agent_card_params={"url": "http://agent.example.com"},
litellm_params=None,
)
mock_registry = Mock()
mock_registry.get_agent_by_name = Mock(return_value=mock_agent)
# Mock litellm.acompletion to verify it's called
mock_acompletion = AsyncMock(return_value={"id": "test-response"})
with patch("litellm.acompletion", mock_acompletion):
with patch(
"litellm.proxy.agent_endpoints.a2a_routing.global_agent_registry",
mock_registry,
):
result = await route_request(
data=data,
llm_router=mock_router,
user_model=None,
route_type="acompletion",
)
# Verify litellm.acompletion was called with api_base injected
mock_acompletion.assert_called_once()
call_kwargs = mock_acompletion.call_args.kwargs
assert call_kwargs["model"] == "a2a/test-agent"
assert call_kwargs["api_base"] == "http://agent.example.com"
@pytest.mark.asyncio
async def test_route_non_a2a_model_raises_error_if_not_in_router():
"""Test that non-a2a models that aren't in router raise an error"""
# Mock data for chat completion with model not in router
data = {
"model": "unknown-model",
"messages": [{"role": "user", "content": "Hello"}],
}
# Mock router without the model
mock_router = Mock()
mock_router.model_names = ["gpt-4", "gpt-3.5-turbo"]
mock_router.deployment_names = []
mock_router.has_model_id = Mock(return_value=False)
mock_router.model_group_alias = None
mock_router.router_general_settings = Mock(pass_through_all_models=False)
mock_router.default_deployment = None
mock_router.pattern_router = Mock(patterns=[])
mock_router.map_team_model = Mock(return_value=None)
# Should raise ProxyModelNotFoundError
from litellm.proxy.route_llm_request import ProxyModelNotFoundError
with pytest.raises(ProxyModelNotFoundError):
await route_request(
data=data,
llm_router=mock_router,
user_model=None,
route_type="acompletion",
)
tests/test_litellm/test_a2a_registry_lookup.py
+73
View File
@@ -0,0 +1,73 @@
"""
Test A2A provider registry lookup functionality.
Maps to: litellm/llms/a2a/chat/transformation.py
"""
import os
import sys
sys.path.insert(0, os.path.abspath("../.."))
import pytest
import litellm
from litellm.llms.a2a.chat.transformation import A2AConfig
def test_resolve_agent_config_from_registry_static_method():
"""Test the static helper method for registry resolution"""
# Test 1: No agent name in model
api_base, api_key, headers = A2AConfig.resolve_agent_config_from_registry(
model="a2a",
api_base="http://test.com",
api_key=None,
headers=None,
optional_params={}
)
assert api_base == "http://test.com"
# Test 2: All params provided - should not lookup registry
api_base, api_key, headers = A2AConfig.resolve_agent_config_from_registry(
model="a2a/test-agent",
api_base="http://explicit.com",
api_key="explicit-key",
headers={"X-Test": "value"},
optional_params={}
)
assert api_base == "http://explicit.com"
assert api_key == "explicit-key"
def test_a2a_registry_integration():
"""Test registry lookup in proxy context"""
try:
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
from litellm.types.agents import AgentResponse
# Create test agent
test_agent = AgentResponse(
agent_id="test-id",
agent_name="test-agent",
agent_card_params={"url": "http://registry-url.example.com:9999"},
litellm_params={"api_key": "registry-key"},
)
# Register and test
original_agents = global_agent_registry.agent_list.copy()
global_agent_registry.register_agent(test_agent)
try:
litellm.completion(
model="a2a/test-agent",
messages=[{"role": "user", "content": "Hello"}]
)
except Exception as e:
# Should use registry URL (connection error expected)
assert "registry-url.example.com" in str(e) or "APIConnectionError" in str(type(e).__name__)
finally:
global_agent_registry.agent_list = original_agents
except ImportError:
pytest.skip("Registry not available (not in proxy context)")