From d09a144d4eb22d4277348105f4a9c4f96f4edf75 Mon Sep 17 00:00:00 2001 From: shin-bot-litellm Date: Sat, 31 Jan 2026 07:50:34 +0000 Subject: [PATCH] litellm_feat: add A2A (Agent-to-Agent) protocol as a provider This PR adds support for calling A2A-compliant agents through the standard LiteLLM completion API by transforming OpenAI chat completion requests to A2A protocol and vice versa. Features: - Full transformation between OpenAI chat completion format and A2A protocol - Support for both streaming and non-streaming completions - Proper handling of A2A message/task/artifact responses - Integration with LiteLLM's provider infrastructure Usage: response = litellm.completion( model='a2a_agent/my-agent', messages=[{'role': 'user', 'content': 'Hello!'}], api_base='http://localhost:9999', # A2A agent endpoint ) For streaming: response = litellm.completion( model='a2a_agent/my-agent', messages=[{'role': 'user', 'content': 'Hello!'}], api_base='http://localhost:9999', stream=True, ) A2A Protocol Reference: https://github.com/a2aproject/A2A --- litellm/__init__.py | 1 + litellm/_lazy_imports_registry.py | 5 + litellm/llms/a2a/__init__.py | 29 ++ litellm/llms/a2a/chat/__init__.py | 9 + litellm/llms/a2a/chat/handler.py | 409 +++++++++++++++++ litellm/llms/a2a/chat/streaming.py | 261 +++++++++++ litellm/llms/a2a/chat/transformation.py | 560 +++++++++++++++++++++++ litellm/main.py | 33 ++ litellm/utils.py | 1 + tests/local_testing/test_a2a_provider.py | 252 ++++++++++ 10 files changed, 1560 insertions(+) create mode 100644 litellm/llms/a2a/__init__.py create mode 100644 litellm/llms/a2a/chat/__init__.py create mode 100644 litellm/llms/a2a/chat/handler.py create mode 100644 litellm/llms/a2a/chat/streaming.py create mode 100644 litellm/llms/a2a/chat/transformation.py create mode 100644 tests/local_testing/test_a2a_provider.py diff --git a/litellm/__init__.py b/litellm/__init__.py index a74a79635f0..1e898d55e8f 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1496,6 +1496,7 @@ if TYPE_CHECKING: from .llms.lemonade.chat.transformation import LemonadeChatConfig as LemonadeChatConfig from .llms.snowflake.embedding.transformation import SnowflakeEmbeddingConfig as SnowflakeEmbeddingConfig from .llms.amazon_nova.chat.transformation import AmazonNovaChatConfig as AmazonNovaChatConfig + from .llms.a2a.chat.transformation import A2AAgentConfig as A2AAgentConfig from litellm.caching.llm_caching_handler import LLMClientCache from litellm.types.llms.bedrock import COHERE_EMBEDDING_INPUT_TYPES from litellm.types.utils import ( diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index 0e52e9a59eb..05d0a4118f1 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -305,6 +305,7 @@ LLM_CONFIG_NAMES = ( "LemonadeChatConfig", "SnowflakeEmbeddingConfig", "AmazonNovaChatConfig", + "A2AAgentConfig", ) # Types that support lazy loading via _lazy_import_types @@ -1119,6 +1120,10 @@ _LLM_CONFIGS_IMPORT_MAP = { ".llms.amazon_nova.chat.transformation", "AmazonNovaChatConfig", ), + "A2AAgentConfig": ( + ".llms.a2a.chat.transformation", + "A2AAgentConfig", + ), } # Import map for utils module lazy imports diff --git a/litellm/llms/a2a/__init__.py b/litellm/llms/a2a/__init__.py new file mode 100644 index 00000000000..ca5949ec232 --- /dev/null +++ b/litellm/llms/a2a/__init__.py @@ -0,0 +1,29 @@ +""" +A2A (Agent-to-Agent) Provider for LiteLLM + +This provider enables calling A2A-compliant agents through the standard LiteLLM completion API. +It transforms OpenAI chat completion requests to A2A protocol and vice versa. + +Usage: + import litellm + + response = litellm.completion( + model="a2a_agent/my-agent", + messages=[{"role": "user", "content": "Hello!"}], + api_base="http://localhost:9999", # A2A agent endpoint + ) + +For streaming: + response = litellm.completion( + model="a2a_agent/my-agent", + messages=[{"role": "user", "content": "Hello!"}], + api_base="http://localhost:9999", + stream=True, + ) + for chunk in response: + print(chunk.choices[0].delta.content) +""" + +from litellm.llms.a2a.chat.transformation import A2AAgentConfig + +__all__ = ["A2AAgentConfig"] diff --git a/litellm/llms/a2a/chat/__init__.py b/litellm/llms/a2a/chat/__init__.py new file mode 100644 index 00000000000..0b5ce12717f --- /dev/null +++ b/litellm/llms/a2a/chat/__init__.py @@ -0,0 +1,9 @@ +""" +A2A Chat Completion Transformation + +Handles transformation between OpenAI chat completion format and A2A protocol. +""" + +from litellm.llms.a2a.chat.transformation import A2AAgentConfig + +__all__ = ["A2AAgentConfig"] diff --git a/litellm/llms/a2a/chat/handler.py b/litellm/llms/a2a/chat/handler.py new file mode 100644 index 00000000000..9a5dc9ff6ac --- /dev/null +++ b/litellm/llms/a2a/chat/handler.py @@ -0,0 +1,409 @@ +""" +A2A Agent HTTP Handler + +Handles making HTTP requests to A2A agents and processing responses. +""" + +import json +from typing import TYPE_CHECKING, Any, AsyncIterator, Callable, Dict, Iterator, List, Optional, Tuple, Union + +import httpx + +import litellm +from litellm._logging import verbose_logger +from litellm.llms.base import BaseLLM +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + HTTPHandler, + get_async_httpx_client, +) +from litellm.types.utils import ModelResponse, ModelResponseStream + +from .streaming import A2AStreamingIterator, create_streaming_response +from .transformation import A2AAgentConfig, A2AAgentError + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + +async def make_a2a_call_async( + client: Optional[AsyncHTTPHandler], + api_base: str, + headers: dict, + data: str, + timeout: Optional[Union[float, httpx.Timeout]], + stream: bool = False, +) -> Tuple[Any, httpx.Headers]: + """ + Make an async HTTP call to an A2A agent. + + Args: + client: AsyncHTTPHandler instance + api_base: The A2A agent endpoint URL + headers: Request headers + data: JSON-encoded request body + timeout: Request timeout + stream: Whether to stream the response + + Returns: + Tuple of (response content or iterator, response headers) + """ + if client is None: + client = litellm.module_level_aclient + + try: + response = await client.post( + api_base, + headers=headers, + data=data, + stream=stream, + timeout=timeout, + ) + except httpx.HTTPStatusError as e: + error_headers = getattr(e, "headers", None) + error_response = getattr(e, "response", None) + if error_headers is None and error_response: + error_headers = getattr(error_response, "headers", None) + + error_body = "" + if error_response: + try: + error_body = await error_response.aread() + error_body = error_body.decode("utf-8") if isinstance(error_body, bytes) else error_body + except Exception: + pass + + raise A2AAgentError( + status_code=e.response.status_code, + message=f"A2A agent request failed: {error_body or str(e)}", + headers=error_headers, + ) + except Exception as e: + for exception in litellm.LITELLM_EXCEPTION_TYPES: + if isinstance(e, exception): + raise e + raise A2AAgentError(status_code=500, message=str(e)) + + if stream: + return response.aiter_lines(), response.headers + else: + return response, response.headers + + +def make_a2a_call_sync( + client: Optional[HTTPHandler], + api_base: str, + headers: dict, + data: str, + timeout: Optional[Union[float, httpx.Timeout]], + stream: bool = False, +) -> Tuple[Any, httpx.Headers]: + """ + Make a sync HTTP call to an A2A agent. + + Args: + client: HTTPHandler instance + api_base: The A2A agent endpoint URL + headers: Request headers + data: JSON-encoded request body + timeout: Request timeout + stream: Whether to stream the response + + Returns: + Tuple of (response content or iterator, response headers) + """ + if client is None: + client = litellm.module_level_client + + try: + response = client.post( + api_base, + headers=headers, + data=data, + stream=stream, + timeout=timeout, + ) + except httpx.HTTPStatusError as e: + error_headers = getattr(e, "headers", None) + error_response = getattr(e, "response", None) + if error_headers is None and error_response: + error_headers = getattr(error_response, "headers", None) + + error_body = "" + if error_response: + try: + error_body = error_response.read() + error_body = error_body.decode("utf-8") if isinstance(error_body, bytes) else error_body + except Exception: + pass + + raise A2AAgentError( + status_code=e.response.status_code, + message=f"A2A agent request failed: {error_body or str(e)}", + headers=error_headers, + ) + except Exception as e: + for exception in litellm.LITELLM_EXCEPTION_TYPES: + if isinstance(e, exception): + raise e + raise A2AAgentError(status_code=500, message=str(e)) + + if stream: + return response.iter_lines(), response.headers + else: + return response, response.headers + + +class A2AAgentChatCompletion(BaseLLM): + """ + Handler for A2A Agent chat completions. + + This class handles making requests to A2A agents and transforming + responses to OpenAI format. + """ + + def __init__(self) -> None: + super().__init__() + self.config = A2AAgentConfig() + + async def async_completion( + self, + model: str, + messages: List[Dict[str, Any]], + api_base: str, + model_response: ModelResponse, + print_verbose: Callable, + timeout: Union[float, httpx.Timeout], + client: Optional[AsyncHTTPHandler], + encoding: Any, + api_key: Optional[str], + logging_obj: "LiteLLMLoggingObj", + optional_params: dict, + litellm_params: dict, + headers: dict, + custom_llm_provider: str, + stream: bool = False, + ): + """ + Async completion call to A2A agent. + """ + # Validate and get headers + headers = self.config.validate_environment( + headers=headers, + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + api_key=api_key, + api_base=api_base, + ) + + # Get complete URL + complete_url = self.config.get_complete_url( + api_base=api_base, + api_key=api_key, + model=model, + optional_params=optional_params, + litellm_params=litellm_params, + stream=stream, + ) + + # Transform request to A2A format + request_data = self.config.transform_request( + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + headers=headers, + ) + + verbose_logger.info(f"A2A async completion: url={complete_url}, stream={stream}") + + if stream: + return await self._async_streaming_completion( + client=client, + api_base=complete_url, + headers=headers, + data=json.dumps(request_data), + timeout=timeout, + model=model, + logging_obj=logging_obj, + ) + else: + response, response_headers = await make_a2a_call_async( + client=client, + api_base=complete_url, + headers=headers, + data=json.dumps(request_data), + timeout=timeout, + stream=False, + ) + + # Transform response + return self.config.transform_response( + model=model, + raw_response=response, + model_response=model_response, + logging_obj=logging_obj, + api_key=api_key, + request_data=request_data, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + encoding=encoding, + ) + + async def _async_streaming_completion( + self, + client: Optional[AsyncHTTPHandler], + api_base: str, + headers: dict, + data: str, + timeout: Optional[Union[float, httpx.Timeout]], + model: str, + logging_obj: "LiteLLMLoggingObj", + ) -> AsyncIterator[ModelResponseStream]: + """ + Handle async streaming completion. + """ + response_iterator, response_headers = await make_a2a_call_async( + client=client, + api_base=api_base, + headers=headers, + data=data, + timeout=timeout, + stream=True, + ) + + # Create streaming iterator + streaming_iterator = A2AStreamingIterator( + streaming_response=response_iterator, + sync_stream=False, + ) + + # Yield transformed chunks + async for chunk in streaming_iterator: + yield create_streaming_response(chunk=chunk, model=model) + + def completion( + self, + model: str, + messages: List[Dict[str, Any]], + api_base: str, + model_response: ModelResponse, + print_verbose: Callable, + timeout: Union[float, httpx.Timeout], + client: Optional[HTTPHandler], + encoding: Any, + api_key: Optional[str], + logging_obj: "LiteLLMLoggingObj", + optional_params: dict, + litellm_params: dict, + headers: dict, + custom_llm_provider: str, + stream: bool = False, + ): + """ + Sync completion call to A2A agent. + """ + # Validate and get headers + headers = self.config.validate_environment( + headers=headers, + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + api_key=api_key, + api_base=api_base, + ) + + # Get complete URL + complete_url = self.config.get_complete_url( + api_base=api_base, + api_key=api_key, + model=model, + optional_params=optional_params, + litellm_params=litellm_params, + stream=stream, + ) + + # Transform request to A2A format + request_data = self.config.transform_request( + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + headers=headers, + ) + + verbose_logger.info(f"A2A sync completion: url={complete_url}, stream={stream}") + + if stream: + return self._sync_streaming_completion( + client=client, + api_base=complete_url, + headers=headers, + data=json.dumps(request_data), + timeout=timeout, + model=model, + logging_obj=logging_obj, + ) + else: + response, response_headers = make_a2a_call_sync( + client=client, + api_base=complete_url, + headers=headers, + data=json.dumps(request_data), + timeout=timeout, + stream=False, + ) + + # Transform response + return self.config.transform_response( + model=model, + raw_response=response, + model_response=model_response, + logging_obj=logging_obj, + api_key=api_key, + request_data=request_data, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + encoding=encoding, + ) + + def _sync_streaming_completion( + self, + client: Optional[HTTPHandler], + api_base: str, + headers: dict, + data: str, + timeout: Optional[Union[float, httpx.Timeout]], + model: str, + logging_obj: "LiteLLMLoggingObj", + ) -> Iterator[ModelResponseStream]: + """ + Handle sync streaming completion. + """ + response_iterator, response_headers = make_a2a_call_sync( + client=client, + api_base=api_base, + headers=headers, + data=data, + timeout=timeout, + stream=True, + ) + + # Create streaming iterator + streaming_iterator = A2AStreamingIterator( + streaming_response=response_iterator, + sync_stream=True, + ) + + # Yield transformed chunks + for chunk in streaming_iterator: + yield create_streaming_response(chunk=chunk, model=model) + + +# Create singleton instance +a2a_agent_chat_completion = A2AAgentChatCompletion() diff --git a/litellm/llms/a2a/chat/streaming.py b/litellm/llms/a2a/chat/streaming.py new file mode 100644 index 00000000000..009e075ff70 --- /dev/null +++ b/litellm/llms/a2a/chat/streaming.py @@ -0,0 +1,261 @@ +""" +A2A Streaming Iterator + +Handles transformation of A2A Server-Sent Events (SSE) to OpenAI streaming format. + +A2A Streaming Events: +1. Task event (kind: "task") - Initial task creation with status "submitted" +2. Status update (kind: "status-update") - Status changes (working, completed) +3. Artifact update (kind: "artifact-update") - Content/artifact delivery + +OpenAI Streaming Format: + data: {"id":"chatcmpl-xxx","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":null}]} +""" + +import json +import time +from typing import Any, AsyncIterator, Dict, Iterator, Optional, Union +from uuid import uuid4 + +from litellm._logging import verbose_logger +from litellm.types.utils import ( + Delta, + GenericStreamingChunk, + ModelResponseStream, + StreamingChoices, + Usage, +) + + +class A2AStreamingIterator: + """ + Iterator that transforms A2A streaming events to OpenAI streaming format. + + Handles both sync and async iteration. + """ + + def __init__( + self, + streaming_response: Union[Iterator[str], AsyncIterator[str]], + sync_stream: bool, + json_mode: bool = False, + ): + self.streaming_response = streaming_response + self.sync_stream = sync_stream + self.json_mode = json_mode + + self._response_id = f"chatcmpl-{uuid4().hex[:8]}" + self._created = int(time.time()) + self._accumulated_text = "" + self._is_finished = False + self._task_id: Optional[str] = None + self._context_id: Optional[str] = None + + def __iter__(self) -> "A2AStreamingIterator": + return self + + def __aiter__(self) -> "A2AStreamingIterator": + return self + + def __next__(self) -> GenericStreamingChunk: + if self._is_finished: + raise StopIteration + + try: + # Get next line from sync iterator + line = next(self.streaming_response) # type: ignore + return self._process_line(line) + except StopIteration: + self._is_finished = True + raise + + async def __anext__(self) -> GenericStreamingChunk: + if self._is_finished: + raise StopAsyncIteration + + try: + # Get next line from async iterator + line = await self.streaming_response.__anext__() # type: ignore + return self._process_line(line) + except StopAsyncIteration: + self._is_finished = True + raise + + def _process_line(self, line: str) -> GenericStreamingChunk: + """Process a single SSE line and transform to OpenAI format.""" + line = line.strip() + + # Skip empty lines and SSE comments + if not line or line.startswith(":"): + return self._create_empty_chunk() + + # Handle SSE data format + if line.startswith("data:"): + data = line[5:].strip() + + # Handle "[DONE]" marker + if data == "[DONE]": + self._is_finished = True + return self._create_final_chunk() + + try: + event_data = json.loads(data) + return self._transform_a2a_event(event_data) + except json.JSONDecodeError as e: + verbose_logger.debug(f"Failed to parse A2A SSE data: {e}") + return self._create_empty_chunk() + + # Try parsing as raw JSON (some implementations don't use SSE format) + try: + event_data = json.loads(line) + return self._transform_a2a_event(event_data) + except json.JSONDecodeError: + return self._create_empty_chunk() + + def _transform_a2a_event(self, event: Dict[str, Any]) -> GenericStreamingChunk: + """Transform an A2A streaming event to OpenAI chunk format.""" + result = event.get("result", {}) + + # Determine event kind + kind = result.get("kind", "") + + if kind == "task": + # Initial task event + self._task_id = result.get("id") + self._context_id = result.get("contextId") + return self._create_empty_chunk() + + elif kind == "status-update": + # Status update event + status = result.get("status", {}) + state = status.get("state", "").lower().replace("task_state_", "") + final = result.get("final", False) + + # Check for message in status + status_message = status.get("message", {}) + text = self._extract_text_from_parts(status_message.get("parts", [])) + + if state == "completed" or final: + self._is_finished = True + return self._create_chunk(text="", is_finished=True, finish_reason="stop") + + if text: + return self._create_chunk(text=text) + + return self._create_empty_chunk() + + elif kind == "artifact-update": + # Artifact content event + artifact = result.get("artifact", {}) + text = self._extract_text_from_parts(artifact.get("parts", [])) + + if text: + self._accumulated_text += text + return self._create_chunk(text=text) + + return self._create_empty_chunk() + + elif kind == "message": + # Direct message response + message = result.get("message", result) + text = self._extract_text_from_parts(message.get("parts", [])) + + if text: + self._accumulated_text += text + return self._create_chunk(text=text) + + return self._create_empty_chunk() + + else: + # Try to extract any text content + text = "" + + # Check for message at top level + if "message" in result: + text = self._extract_text_from_parts(result["message"].get("parts", [])) + + # Check for artifact at top level + elif "artifact" in result: + text = self._extract_text_from_parts(result["artifact"].get("parts", [])) + + if text: + self._accumulated_text += text + return self._create_chunk(text=text) + + return self._create_empty_chunk() + + def _extract_text_from_parts(self, parts: list) -> str: + """Extract text content from A2A parts.""" + text_parts = [] + for part in parts: + if part.get("kind") == "text": + text_parts.append(part.get("text", "")) + return "".join(text_parts) + + def _create_chunk( + self, + text: str = "", + is_finished: bool = False, + finish_reason: Optional[str] = None, + ) -> GenericStreamingChunk: + """Create an OpenAI-format streaming chunk.""" + return GenericStreamingChunk( + text=text, + is_finished=is_finished, + finish_reason=finish_reason or "", + usage=None, + index=0, + ) + + def _create_empty_chunk(self) -> GenericStreamingChunk: + """Create an empty chunk (for non-content events).""" + return GenericStreamingChunk( + text="", + is_finished=False, + finish_reason="", + usage=None, + index=0, + ) + + def _create_final_chunk(self) -> GenericStreamingChunk: + """Create the final chunk marking end of stream.""" + return GenericStreamingChunk( + text="", + is_finished=True, + finish_reason="stop", + usage=None, + index=0, + ) + + +def create_streaming_response( + chunk: GenericStreamingChunk, + model: str, + response_id: Optional[str] = None, +) -> ModelResponseStream: + """ + Create a ModelResponseStream from a GenericStreamingChunk. + + This formats the chunk in the standard OpenAI streaming response format. + """ + if response_id is None: + response_id = f"chatcmpl-{uuid4().hex[:8]}" + + response = ModelResponseStream( + id=response_id, + object="chat.completion.chunk", + created=int(time.time()), + model=model, + choices=[ + StreamingChoices( + index=chunk.get("index", 0), + delta=Delta( + role="assistant" if not chunk.get("is_finished") else None, + content=chunk.get("text") or None, + ), + finish_reason=chunk.get("finish_reason") if chunk.get("is_finished") else None, + ) + ], + ) + + return response diff --git a/litellm/llms/a2a/chat/transformation.py b/litellm/llms/a2a/chat/transformation.py new file mode 100644 index 00000000000..73c7327d72a --- /dev/null +++ b/litellm/llms/a2a/chat/transformation.py @@ -0,0 +1,560 @@ +""" +A2A Agent Chat Transformation + +Transforms OpenAI chat completion API requests/responses to/from A2A protocol. + +A2A Protocol Reference: https://github.com/a2aproject/A2A + +OpenAI Message Format: + {"role": "user", "content": "Hello!"} + +A2A Message Format: + { + "role": "user", + "parts": [{"kind": "text", "text": "Hello!"}], + "messageId": "abc123" + } +""" + +import json +import time +from typing import ( + TYPE_CHECKING, + Any, + AsyncIterator, + Dict, + Iterator, + List, + Literal, + Optional, + Tuple, + Union, +) +from uuid import uuid4 + +import httpx + +import litellm +from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import map_finish_reason +from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import ( + Choices, + Delta, + Message, + ModelResponse, + StreamingChoices, + Usage, +) + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + +class A2AAgentError(BaseLLMException): + """Exception for A2A Agent errors.""" + + pass + + +class A2AAgentConfig(BaseConfig): + """ + Configuration class for A2A Agent provider. + + Handles transformation between OpenAI chat completion format and A2A protocol. + + Reference: https://github.com/a2aproject/A2A/blob/main/docs/specification.md + """ + + # Default values + frequency_penalty: Optional[float] = None + max_tokens: Optional[int] = None + presence_penalty: Optional[float] = None + stop: Optional[Union[str, List[str]]] = None + temperature: Optional[float] = None + top_p: Optional[float] = None + + def __init__( + self, + frequency_penalty: Optional[float] = None, + max_tokens: Optional[int] = None, + presence_penalty: Optional[float] = None, + stop: Optional[Union[str, List[str]]] = None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + ) -> None: + locals_ = locals().copy() + for key, value in locals_.items(): + if key != "self" and value is not None: + setattr(self.__class__, key, value) + + @property + def custom_llm_provider(self) -> str: + return "a2a_agent" + + def get_supported_openai_params(self, model: str) -> List[str]: + """ + Return the list of supported OpenAI parameters. + + A2A agents may support a subset of these depending on their implementation. + """ + return [ + "stream", + "max_tokens", + "temperature", + "top_p", + "stop", + "user", + ] + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + """ + Map OpenAI parameters to A2A compatible format. + + A2A doesn't have direct parameter equivalents for most LLM params, + but we preserve them in case the underlying agent uses them. + """ + supported_params = self.get_supported_openai_params(model) + + for param, value in non_default_params.items(): + if param in supported_params: + optional_params[param] = value + elif not drop_params: + # Pass through unsupported params if not dropping + optional_params[param] = value + + return optional_params + + def _get_openai_compatible_provider_info( + self, api_base: Optional[str], api_key: Optional[str] + ) -> Tuple[Optional[str], Optional[str]]: + """Get the API base and key for A2A agent.""" + api_base = api_base or get_secret_str("A2A_AGENT_API_BASE") + api_key = api_key or get_secret_str("A2A_AGENT_API_KEY") + return api_base, api_key + + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + """ + Validate and set up the request headers for A2A. + """ + api_base, api_key = self._get_openai_compatible_provider_info(api_base, api_key) + + if not api_base: + raise A2AAgentError( + status_code=400, + message="api_base is required for A2A agent calls. Set via api_base parameter or A2A_AGENT_API_BASE env var.", + ) + + # Set up headers + headers = headers or {} + headers["Content-Type"] = "application/json" + + # Add authorization if API key is provided + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + + return headers + + def get_complete_url( + self, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: dict, + litellm_params: dict, + stream: Optional[bool] = None, + ) -> str: + """ + Construct the complete URL for the A2A agent endpoint. + + A2A uses JSON-RPC at the agent's base URL. + """ + api_base, _ = self._get_openai_compatible_provider_info(api_base, api_key) + + if not api_base: + raise A2AAgentError( + status_code=400, + message="api_base is required for A2A agent calls.", + ) + + # Strip trailing slash + api_base = api_base.rstrip("/") + + return api_base + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: Union[dict, httpx.Headers], + ) -> A2AAgentError: + """Return the appropriate error class for A2A errors.""" + return A2AAgentError( + status_code=status_code, + message=error_message, + headers=dict(headers) if headers else None, + ) + + # ======================================================================== + # Request Transformation: OpenAI -> A2A + # ======================================================================== + + def transform_request( + self, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + headers: dict, + ) -> Dict[str, Any]: + """ + Transform OpenAI chat completion request to A2A SendMessageRequest format. + + OpenAI format: + { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}] + } + + A2A format (JSON-RPC): + { + "jsonrpc": "2.0", + "method": "message/send", + "id": "request-id", + "params": { + "message": { + "role": "user", + "parts": [{"kind": "text", "text": "Hello"}], + "messageId": "msg-id" + } + } + } + """ + stream = optional_params.get("stream", False) + + # Transform messages to A2A format + a2a_message = self._transform_messages_to_a2a(messages) + + # Build A2A JSON-RPC request + request_id = str(uuid4()) + method = "message/stream" if stream else "message/send" + + a2a_request = { + "jsonrpc": "2.0", + "method": method, + "id": request_id, + "params": { + "message": a2a_message, + }, + } + + # Add configuration if applicable + config = {} + if optional_params.get("blocking") is not None: + config["blocking"] = optional_params["blocking"] + if config: + a2a_request["params"]["configuration"] = config + + verbose_logger.debug(f"A2A request: {json.dumps(a2a_request, indent=2)}") + + return a2a_request + + def _transform_messages_to_a2a( + self, messages: List[AllMessageValues] + ) -> Dict[str, Any]: + """ + Transform OpenAI messages to a single A2A message. + + A2A typically works with single messages in a conversation context. + We'll combine all messages into context and use the last user message. + """ + # Find the last user message + last_user_message = None + for msg in reversed(messages): + if msg.get("role") == "user": + last_user_message = msg + break + + if last_user_message is None: + # If no user message, use the last message + last_user_message = messages[-1] if messages else {"role": "user", "content": ""} + + # Transform to A2A parts format + content = last_user_message.get("content", "") + + # Handle content that might be a list (multimodal) + parts = [] + if isinstance(content, list): + for item in content: + if isinstance(item, dict): + if item.get("type") == "text": + parts.append({ + "kind": "text", + "text": item.get("text", ""), + }) + elif item.get("type") == "image_url": + # A2A supports file parts for images + image_url = item.get("image_url", {}) + url = image_url.get("url", "") if isinstance(image_url, dict) else str(image_url) + parts.append({ + "kind": "file", + "file": { + "uri": url, + "mimeType": "image/*", + }, + }) + elif isinstance(item, str): + parts.append({"kind": "text", "text": item}) + else: + parts.append({"kind": "text", "text": str(content)}) + + # Build A2A message + a2a_message = { + "role": self._map_role_to_a2a(last_user_message.get("role", "user")), + "parts": parts, + "messageId": uuid4().hex, + } + + return a2a_message + + def _map_role_to_a2a(self, openai_role: str) -> str: + """Map OpenAI role to A2A role.""" + role_mapping = { + "user": "user", + "assistant": "agent", + "system": "user", # A2A doesn't have system role, treat as user context + "tool": "user", + "function": "user", + } + return role_mapping.get(openai_role, "user") + + # ======================================================================== + # Response Transformation: A2A -> OpenAI + # ======================================================================== + + def transform_response( + self, + model: str, + raw_response: httpx.Response, + model_response: ModelResponse, + logging_obj: "LiteLLMLoggingObj", + api_key: Optional[str], + request_data: dict, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + encoding: Any, + json_mode: bool = False, + ) -> ModelResponse: + """ + Transform A2A response to OpenAI chat completion format. + + A2A response format: + { + "jsonrpc": "2.0", + "id": "request-id", + "result": { + "task": {...} or + "message": { + "role": "agent", + "parts": [{"kind": "text", "text": "Hello!"}], + "messageId": "msg-id" + } + } + } + + OpenAI response format: + { + "id": "chatcmpl-xxx", + "object": "chat.completion", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "Hello!"}, + "finish_reason": "stop" + }], + "usage": {...} + } + """ + try: + response_json = raw_response.json() + except json.JSONDecodeError as e: + raise A2AAgentError( + status_code=500, + message=f"Failed to parse A2A response: {str(e)}", + headers=dict(raw_response.headers), + ) + + verbose_logger.debug(f"A2A response: {json.dumps(response_json, indent=2)}") + + # Check for JSON-RPC error + if "error" in response_json: + error = response_json["error"] + raise A2AAgentError( + status_code=error.get("code", 500), + message=error.get("message", "Unknown A2A error"), + headers=dict(raw_response.headers), + ) + + result = response_json.get("result", {}) + + # Extract content from A2A response + content = self._extract_content_from_a2a_result(result) + finish_reason = self._determine_finish_reason(result) + + # Build OpenAI response + model_response.id = response_json.get("id", f"chatcmpl-{uuid4().hex[:8]}") + model_response.object = "chat.completion" + model_response.created = int(time.time()) + model_response.model = model + + model_response.choices = [ + Choices( + index=0, + message=Message( + role="assistant", + content=content, + ), + finish_reason=finish_reason, + ) + ] + + # Estimate token usage (A2A doesn't provide this) + prompt_tokens = self._estimate_tokens(messages) + completion_tokens = self._estimate_tokens_from_text(content) + + model_response.usage = Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ) + + return model_response + + def _extract_content_from_a2a_result(self, result: Dict[str, Any]) -> str: + """Extract text content from A2A result.""" + content_parts = [] + + # Check the kind of result + result_kind = result.get("kind", "") + + # Handle direct message response (kind: "message") + # The result itself contains role and parts when kind is "message" + if result_kind == "message" or (result.get("role") and "parts" in result): + parts = result.get("parts", []) + for part in parts: + if part.get("kind") == "text": + content_parts.append(part.get("text", "")) + + # Check for nested message response (some implementations) + elif "message" in result: + message = result["message"] + parts = message.get("parts", []) + for part in parts: + if part.get("kind") == "text": + content_parts.append(part.get("text", "")) + + # Check for task with artifacts + elif "task" in result or result_kind == "task": + task = result.get("task", result) if "task" in result else result + artifacts = task.get("artifacts", []) + for artifact in artifacts: + for part in artifact.get("parts", []): + if part.get("kind") == "text": + content_parts.append(part.get("text", "")) + + # Also check task status message + status = task.get("status", {}) + status_message = status.get("message", {}) + for part in status_message.get("parts", []): + if part.get("kind") == "text": + content_parts.append(part.get("text", "")) + + # Check for artifacts at root level + elif "artifact" in result: + artifact = result["artifact"] + for part in artifact.get("parts", []): + if part.get("kind") == "text": + content_parts.append(part.get("text", "")) + + # Fallback: check for parts directly on result + elif "parts" in result: + for part in result.get("parts", []): + if part.get("kind") == "text": + content_parts.append(part.get("text", "")) + + return "\n".join(content_parts) if content_parts else "" + + def _determine_finish_reason(self, result: Dict[str, Any]) -> str: + """Determine the finish reason from A2A result.""" + # Check task status + if "task" in result: + status = result["task"].get("status", {}) + state = status.get("state", "") + + state_mapping = { + "completed": "stop", + "failed": "stop", + "canceled": "stop", + "rejected": "stop", + "input_required": "stop", # Needs more input + "working": "stop", + "submitted": "stop", + } + return state_mapping.get(state.lower().replace("task_state_", ""), "stop") + + return "stop" + + def _estimate_tokens(self, messages: List[AllMessageValues]) -> int: + """Estimate token count for messages.""" + total = 0 + for msg in messages: + content = msg.get("content", "") + if isinstance(content, str): + total += len(content) // 4 # Rough estimate + elif isinstance(content, list): + for item in content: + if isinstance(item, dict) and "text" in item: + total += len(item["text"]) // 4 + return max(total, 1) + + def _estimate_tokens_from_text(self, text: str) -> int: + """Estimate token count from text.""" + return max(len(text) // 4, 1) + + # ======================================================================== + # Streaming Support + # ======================================================================== + + def get_model_response_iterator( + self, + streaming_response: Union[Iterator[str], AsyncIterator[str]], + sync_stream: bool, + json_mode: bool = False, + ): + """ + Get an iterator that transforms A2A streaming events to OpenAI format. + """ + from litellm.llms.a2a.chat.streaming import A2AStreamingIterator + + return A2AStreamingIterator( + streaming_response=streaming_response, + sync_stream=sync_stream, + json_mode=json_mode, + ) diff --git a/litellm/main.py b/litellm/main.py index 13361c644cb..6a52bb94162 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -4232,6 +4232,39 @@ def completion( # type: ignore # noqa: PLR0915 client=client, ) + elif custom_llm_provider == "a2a_agent": + # A2A (Agent-to-Agent) Protocol Provider + from litellm.llms.a2a.chat.transformation import A2AAgentConfig + + ( + api_base, + api_key, + ) = A2AAgentConfig()._get_openai_compatible_provider_info( + api_base=api_base or litellm.api_base, + api_key=api_key or litellm.api_key, + ) + + headers = headers or litellm.headers + + response = base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider=custom_llm_provider, + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, + client=client, + ) + else: raise LiteLLMUnknownProvider( model=model, custom_llm_provider=custom_llm_provider diff --git a/litellm/utils.py b/litellm/utils.py index 7c4eec7ba32..f53e347859f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7900,6 +7900,7 @@ class ProviderConfigManager: lambda: ProviderConfigManager._get_langgraph_config(), False, ), + LlmProviders.A2A_AGENT: (lambda: litellm.A2AAgentConfig(), False), } @staticmethod diff --git a/tests/local_testing/test_a2a_provider.py b/tests/local_testing/test_a2a_provider.py new file mode 100644 index 00000000000..584c34ba4dd --- /dev/null +++ b/tests/local_testing/test_a2a_provider.py @@ -0,0 +1,252 @@ +""" +Tests for the A2A (Agent-to-Agent) Provider + +This test file verifies the A2A provider integration with LiteLLM. + +To run these tests: +1. Start an A2A agent (e.g., the helloworld sample agent on port 9999) +2. Run: pytest tests/local_testing/test_a2a_provider.py -v + +For the helloworld agent: + cd a2a-samples/samples/python/agents/helloworld + python __main__.py + +Environment variables: + A2A_AGENT_API_BASE: Base URL for the A2A agent (default: http://localhost:9999) +""" + +import os +import json +import pytest +from unittest.mock import MagicMock, patch, AsyncMock +import httpx + +import litellm +from litellm.llms.a2a.chat.transformation import A2AAgentConfig, A2AAgentError + + +# Test fixtures +@pytest.fixture +def a2a_config(): + return A2AAgentConfig() + + +@pytest.fixture +def sample_openai_messages(): + return [ + {"role": "user", "content": "Hello!"} + ] + + +@pytest.fixture +def sample_a2a_response(): + return { + "jsonrpc": "2.0", + "id": "test-123", + "result": { + "kind": "message", + "messageId": "msg-123", + "parts": [ + {"kind": "text", "text": "Hello World"} + ], + "role": "agent" + } + } + + +class TestA2AAgentConfig: + """Test the A2AAgentConfig class.""" + + def test_custom_llm_provider(self, a2a_config): + """Test that the custom_llm_provider is correctly set.""" + assert a2a_config.custom_llm_provider == "a2a_agent" + + def test_get_supported_openai_params(self, a2a_config): + """Test that supported params are returned.""" + params = a2a_config.get_supported_openai_params("test-model") + assert "stream" in params + assert "max_tokens" in params + assert "temperature" in params + + def test_validate_environment_missing_api_base(self, a2a_config): + """Test that validation fails when api_base is missing.""" + with pytest.raises(A2AAgentError): + a2a_config.validate_environment( + headers={}, + model="test-model", + messages=[], + optional_params={}, + litellm_params={}, + api_key=None, + api_base=None, + ) + + def test_validate_environment_with_api_base(self, a2a_config): + """Test that validation succeeds with api_base.""" + headers = a2a_config.validate_environment( + headers={}, + model="test-model", + messages=[], + optional_params={}, + litellm_params={}, + api_key="test-key", + api_base="http://localhost:9999", + ) + assert "Content-Type" in headers + assert headers["Content-Type"] == "application/json" + assert "Authorization" in headers + + +class TestA2AMessageTransformation: + """Test OpenAI to A2A message transformation.""" + + def test_transform_simple_message(self, a2a_config, sample_openai_messages): + """Test transforming a simple user message.""" + request = a2a_config.transform_request( + model="test-model", + messages=sample_openai_messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert request["jsonrpc"] == "2.0" + assert request["method"] == "message/send" + assert "params" in request + assert "message" in request["params"] + + message = request["params"]["message"] + assert message["role"] == "user" + assert len(message["parts"]) == 1 + assert message["parts"][0]["kind"] == "text" + assert message["parts"][0]["text"] == "Hello!" + + def test_transform_streaming_message(self, a2a_config, sample_openai_messages): + """Test transforming a message with streaming enabled.""" + request = a2a_config.transform_request( + model="test-model", + messages=sample_openai_messages, + optional_params={"stream": True}, + litellm_params={}, + headers={}, + ) + + assert request["method"] == "message/stream" + + def test_transform_multipart_content(self, a2a_config): + """Test transforming messages with multipart content.""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What's in this image?"}, + {"type": "image_url", "image_url": {"url": "http://example.com/image.jpg"}} + ] + } + ] + + request = a2a_config.transform_request( + model="test-model", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + + message = request["params"]["message"] + assert len(message["parts"]) == 2 + assert message["parts"][0]["kind"] == "text" + assert message["parts"][1]["kind"] == "file" + + +class TestA2AResponseTransformation: + """Test A2A to OpenAI response transformation.""" + + def test_extract_message_content(self, a2a_config, sample_a2a_response): + """Test extracting content from A2A message response.""" + result = sample_a2a_response["result"] + content = a2a_config._extract_content_from_a2a_result(result) + assert content == "Hello World" + + def test_extract_task_artifact_content(self, a2a_config): + """Test extracting content from A2A task with artifacts.""" + result = { + "kind": "task", + "id": "task-123", + "status": {"state": "completed"}, + "artifacts": [ + { + "artifactId": "artifact-1", + "parts": [{"kind": "text", "text": "Task result"}] + } + ] + } + content = a2a_config._extract_content_from_a2a_result(result) + assert content == "Task result" + + def test_determine_finish_reason_completed(self, a2a_config): + """Test determining finish reason for completed task.""" + result = { + "task": { + "status": {"state": "completed"} + } + } + reason = a2a_config._determine_finish_reason(result) + assert reason == "stop" + + def test_determine_finish_reason_input_required(self, a2a_config): + """Test determining finish reason for input_required task.""" + result = { + "task": { + "status": {"state": "input_required"} + } + } + reason = a2a_config._determine_finish_reason(result) + assert reason == "stop" + + +class TestA2AIntegration: + """Integration tests requiring a running A2A agent.""" + + @pytest.mark.skip(reason="Requires running A2A agent") + def test_non_streaming_completion(self): + """Test non-streaming completion with real A2A agent.""" + api_base = os.environ.get("A2A_AGENT_API_BASE", "http://localhost:9999") + + response = litellm.completion( + model="a2a_agent/test-agent", + messages=[{"role": "user", "content": "Hello!"}], + api_base=api_base, + ) + + assert response.id is not None + assert len(response.choices) > 0 + assert response.choices[0].message.content is not None + assert response.choices[0].finish_reason == "stop" + + @pytest.mark.skip(reason="Requires running A2A agent") + def test_streaming_completion(self): + """Test streaming completion with real A2A agent.""" + api_base = os.environ.get("A2A_AGENT_API_BASE", "http://localhost:9999") + + response = litellm.completion( + model="a2a_agent/test-agent", + messages=[{"role": "user", "content": "Hello!"}], + api_base=api_base, + stream=True, + ) + + content = "" + finish_reason = None + for chunk in response: + if chunk.choices and chunk.choices[0].delta.content: + content += chunk.choices[0].delta.content + if chunk.choices and chunk.choices[0].finish_reason: + finish_reason = chunk.choices[0].finish_reason + + assert content != "" + assert finish_reason == "stop" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])