mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
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
This commit is contained in:
shin-bot-litellm
parent
3f1bda57e2
commit
d09a144d4e
10 changed files with 1560 additions and 0 deletions
|
|
@ -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 (
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
29
litellm/llms/a2a/__init__.py
Normal file
29
litellm/llms/a2a/__init__.py
Normal file
|
|
@ -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"]
|
||||
9
litellm/llms/a2a/chat/__init__.py
Normal file
9
litellm/llms/a2a/chat/__init__.py
Normal file
|
|
@ -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"]
|
||||
409
litellm/llms/a2a/chat/handler.py
Normal file
409
litellm/llms/a2a/chat/handler.py
Normal file
|
|
@ -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()
|
||||
261
litellm/llms/a2a/chat/streaming.py
Normal file
261
litellm/llms/a2a/chat/streaming.py
Normal file
|
|
@ -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
|
||||
560
litellm/llms/a2a/chat/transformation.py
Normal file
560
litellm/llms/a2a/chat/transformation.py
Normal file
|
|
@ -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,
|
||||
)
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -7900,6 +7900,7 @@ class ProviderConfigManager:
|
|||
lambda: ProviderConfigManager._get_langgraph_config(),
|
||||
False,
|
||||
),
|
||||
LlmProviders.A2A_AGENT: (lambda: litellm.A2AAgentConfig(), False),
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
252
tests/local_testing/test_a2a_provider.py
Normal file
252
tests/local_testing/test_a2a_provider.py
Normal file
|
|
@ -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"])
|
||||
Loading…
Add table
Reference in a new issue