diff --git a/docs/my-website/docs/a2a.md b/docs/my-website/docs/a2a.md index a7e8b52d99..b1166a7809 100644 --- a/docs/my-website/docs/a2a.md +++ b/docs/my-website/docs/a2a.md @@ -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 diff --git a/docs/my-website/docs/a2a_invoking_agents.md b/docs/my-website/docs/a2a_invoking_agents.md new file mode 100644 index 0000000000..3bb248e456 --- /dev/null +++ b/docs/my-website/docs/a2a_invoking_agents.md @@ -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 + + + + +```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) +``` + + + + +```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); +``` + + + + +```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?"} + ] + }' +``` + + + + +### Streaming + + + + +```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) +``` + + + + +```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); + } +} +``` + + + + +```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 + }' +``` + + + + +## Key Differences + +| Method | Use Case | Advantages | +|--------|----------|------------| +| **A2A SDK** | Native A2A protocol integration | • Full A2A protocol support
• Access to task states and artifacts
• Context management | +| **OpenAI SDK** | Familiar OpenAI-style interface | • Drop-in replacement for OpenAI calls
• Easier migration from LLM to agent workflows
• 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. + +::: diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index d932b6af25..98c8ee6eaa 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -469,6 +469,7 @@ const sidebars = { label: "/a2a - A2A Agent Gateway", items: [ "a2a", + "a2a_invoking_agents", "a2a_cost_tracking", "a2a_agent_permissions" ], diff --git a/litellm/llms/a2a/chat/streaming_iterator.py b/litellm/llms/a2a/chat/streaming_iterator.py index 84b6fffaa3..4b689414dd 100644 --- a/litellm/llms/a2a/chat/streaming_iterator.py +++ b/litellm/llms/a2a/chat/streaming_iterator.py @@ -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: diff --git a/litellm/llms/a2a/chat/transformation.py b/litellm/llms/a2a/chat/transformation.py index 243fba6371..163cd5ab22 100644 --- a/litellm/llms/a2a/chat/transformation.py +++ b/litellm/llms/a2a/chat/transformation.py @@ -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/". + + 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", diff --git a/litellm/llms/a2a/common_utils.py b/litellm/llms/a2a/common_utils.py index 4c7da78b42..116e120540 100644 --- a/litellm/llms/a2a/common_utils.py +++ b/litellm/llms/a2a/common_utils.py @@ -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] diff --git a/litellm/main.py b/litellm/main.py index 60c889ab1d..6d6bef81c2 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -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/" + 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/'." + ) headers = headers or litellm.headers diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py new file mode 100644 index 0000000000..8e4c705df2 --- /dev/null +++ b/litellm/proxy/agent_endpoints/a2a_routing.py @@ -0,0 +1,53 @@ +""" +A2A Agent Routing + +Handles routing for A2A agents (models with "a2a/" 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) diff --git a/litellm/proxy/agent_endpoints/model_list_helpers.py b/litellm/proxy/agent_endpoints/model_list_helpers.py new file mode 100644 index 0000000000..c640300bb8 --- /dev/null +++ b/litellm/proxy/agent_endpoints/model_list_helpers.py @@ -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/" 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/" 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 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8f433bfa48..b06071481e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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} diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index c6a93164d4..39ef5fdd1d 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -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) diff --git a/tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py b/tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py new file mode 100644 index 0000000000..92cd3d9ad6 --- /dev/null +++ b/tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py @@ -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" diff --git a/tests/test_litellm/proxy/test_route_a2a_models.py b/tests/test_litellm/proxy/test_route_a2a_models.py new file mode 100644 index 0000000000..e8d89c0c7e --- /dev/null +++ b/tests/test_litellm/proxy/test_route_a2a_models.py @@ -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", + ) diff --git a/tests/test_litellm/test_a2a_registry_lookup.py b/tests/test_litellm/test_a2a_registry_lookup.py new file mode 100644 index 0000000000..9938f10a43 --- /dev/null +++ b/tests/test_litellm/test_a2a_registry_lookup.py @@ -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)")