scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

Merge branch 'main' of https://github.com/BerriAI/litellm into litellm_fix-key-mask

Browse source
This commit is contained in:
Harshit Jain 2026-02-04 03:01:57 +05:30
commit c9267426a5
No known key found for this signature in database
GPG key ID: 36C392CD4415B4CF
23 changed files with 815 additions and 26 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

BIN
enterprise/dist/litellm_enterprise-0.1.29-py3-none-any.whl vendored Normal file
View file

Binary file not shown.

BIN
enterprise/dist/litellm_enterprise-0.1.29.tar.gz vendored Normal file
View file

Binary file not shown.

4
enterprise/pyproject.toml
View file

@ -1,6 +1,6 @@
[tool.poetry]
name = "litellm-enterprise"
version = "0.1.28"
version = "0.1.29"
description = "Package for LiteLLM Enterprise features"
authors = ["BerriAI"]
readme = "README.md"
@ -22,7 +22,7 @@ requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"
[tool.commitizen]
version = "0.1.28"
version = "0.1.29"
version_files = [
"pyproject.toml:version",
"../requirements.txt:litellm-enterprise==",

1
litellm/__init__.py
View file

@ -1378,6 +1378,7 @@ if TYPE_CHECKING:
from .llms.topaz.image_variations.transformation import TopazImageVariationConfig as TopazImageVariationConfig
from litellm.llms.openai.completion.transformation import OpenAITextCompletionConfig as OpenAITextCompletionConfig
from .llms.groq.chat.transformation import GroqChatConfig as GroqChatConfig
from .llms.a2a.chat.transformation import A2AConfig as A2AConfig
from .llms.voyage.embedding.transformation import VoyageEmbeddingConfig as VoyageEmbeddingConfig
from .llms.voyage.embedding.transformation_contextual import VoyageContextualEmbeddingConfig as VoyageContextualEmbeddingConfig
from .llms.infinity.embedding.transformation import InfinityEmbeddingConfig as InfinityEmbeddingConfig

2
litellm/_lazy_imports_registry.py
View file

@ -213,6 +213,7 @@ LLM_CONFIG_NAMES = (
"TopazImageVariationConfig",
"OpenAITextCompletionConfig",
"GroqChatConfig",
"A2AConfig",
"GenAIHubOrchestrationConfig",
"VoyageEmbeddingConfig",
"VoyageContextualEmbeddingConfig",
@ -850,6 +851,7 @@ _LLM_CONFIGS_IMPORT_MAP = {
"OpenAITextCompletionConfig",
),
"GroqChatConfig": (".llms.groq.chat.transformation", "GroqChatConfig"),
"A2AConfig": (".llms.a2a.chat.transformation", "A2AConfig"),
"GenAIHubOrchestrationConfig": (
".llms.sap.chat.transformation",
"GenAIHubOrchestrationConfig",

6
litellm/llms/a2a/__init__.py Normal file
View file

@ -0,0 +1,6 @@
"""
A2A (Agent-to-Agent) Protocol Provider for LiteLLM
"""
from .chat.transformation import A2AConfig
__all__ = ["A2AConfig"]

6
litellm/llms/a2a/chat/__init__.py Normal file
View file

@ -0,0 +1,6 @@
"""
A2A Chat Completion Implementation
"""
from .transformation import A2AConfig
__all__ = ["A2AConfig"]

103
litellm/llms/a2a/chat/streaming_iterator.py Normal file
View file

@ -0,0 +1,103 @@
"""
A2A Streaming Response Iterator
"""
from typing import Optional, Union
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
from litellm.types.utils import GenericStreamingChunk, ModelResponseStream
from ..common_utils import extract_text_from_a2a_response
class A2AModelResponseIterator(BaseModelResponseIterator):
"""
Iterator for parsing A2A streaming responses.
Converts A2A JSON-RPC streaming chunks to OpenAI-compatible format.
"""
def __init__(
self,
streaming_response,
sync_stream: bool,
json_mode: Optional[bool] = False,
model: str = "a2a/agent",
):
super().__init__(
streaming_response=streaming_response,
sync_stream=sync_stream,
json_mode=json_mode,
)
self.model = model
def chunk_parser(self, chunk: dict) -> Union[GenericStreamingChunk, ModelResponseStream]:
"""
Parse A2A streaming chunk to OpenAI format.
A2A chunk format:
{
"jsonrpc": "2.0",
"id": "request-id",
"result": {
"message": {
"parts": [{"kind": "text", "text": "content"}]
}
}
}
Or for tasks:
{
"jsonrpc": "2.0",
"result": {
"kind": "task",
"status": {"state": "running"},
"artifacts": [{"parts": [{"kind": "text", "text": "content"}]}]
}
}
"""
try:
# Extract text from A2A response
text = extract_text_from_a2a_response(chunk)
# Determine finish reason
finish_reason = self._get_finish_reason(chunk)
# Return generic streaming chunk
return GenericStreamingChunk(
text=text,
is_finished=bool(finish_reason),
finish_reason=finish_reason or "",
usage=None,
index=0,
tool_use=None,
)
except Exception:
# Return empty chunk on parse error
return GenericStreamingChunk(
text="",
is_finished=False,
finish_reason="",
usage=None,
index=0,
tool_use=None,
)
def _get_finish_reason(self, chunk: dict) -> Optional[str]:
"""Extract finish reason from A2A chunk"""
result = chunk.get("result", {})
# Check for task completion
if isinstance(result, dict):
status = result.get("status", {})
if isinstance(status, dict):
state = status.get("state")
if state == "completed":
return "stop"
elif state == "failed":
return "error"
# Check for [DONE] marker
if chunk.get("done") is True:
return "stop"
return None

303
litellm/llms/a2a/chat/transformation.py Normal file
View file

@ -0,0 +1,303 @@
"""
A2A Protocol Transformation for LiteLLM
"""
import uuid
from typing import Any, Dict, Iterator, List, Optional, Union, cast
import httpx
from pydantic import BaseModel
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import Choices, Message, ModelResponse
from ..common_utils import (
A2AError,
convert_messages_to_prompt,
extract_text_from_a2a_response,
)
from .streaming_iterator import A2AModelResponseIterator
class A2AConfig(BaseConfig):
"""
Configuration for A2A (Agent-to-Agent) Protocol.
Handles transformation between OpenAI and A2A JSON-RPC 2.0 formats.
"""
def get_supported_openai_params(self, model: str) -> List[str]:
"""Return list of supported OpenAI parameters"""
return [
"stream",
"temperature",
"max_tokens",
"top_p",
]
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
model: str,
drop_params: bool,
) -> dict:
"""
Map OpenAI parameters to A2A parameters.
For A2A protocol, we don't need to map most parameters since
they're handled in the transform_request method.
"""
return optional_params
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 environment and set headers for A2A requests.
Args:
headers: Request headers dict
model: Model name
messages: Messages list
optional_params: Optional parameters
litellm_params: LiteLLM parameters
api_key: API key (optional for A2A)
api_base: API base URL
Returns:
Updated headers dict
"""
# Ensure Content-Type is set to application/json for JSON-RPC 2.0
if "content-type" not in headers and "Content-Type" not in headers:
headers["Content-Type"] = "application/json"
# Add Authorization header if API key is provided
if api_key is not None:
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:
"""
Get the complete A2A agent endpoint URL.
A2A agents use JSON-RPC 2.0 at the base URL, not specific paths.
The method (message/send or message/stream) is specified in the
JSON-RPC request body, not in the URL.
Args:
api_base: Base URL of the A2A agent (e.g., "http://0.0.0.0:9999")
api_key: API key (not used for URL construction)
model: Model name (not used for A2A, agent determined by api_base)
optional_params: Optional parameters
litellm_params: LiteLLM parameters
stream: Whether this is a streaming request (affects JSON-RPC method)
Returns:
Complete URL for the A2A endpoint (base URL)
"""
if api_base is None:
raise ValueError("api_base is required for A2A provider")
# A2A uses JSON-RPC 2.0 at the base URL
# Remove trailing slash for consistency
return api_base.rstrip("/")
def transform_request(
self,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
headers: dict,
) -> dict:
"""
Transform OpenAI request to A2A JSON-RPC 2.0 format.
Args:
model: Model name
messages: List of OpenAI messages
optional_params: Optional parameters
litellm_params: LiteLLM parameters
headers: Request headers
Returns:
A2A JSON-RPC 2.0 request dict
"""
# Generate request ID
request_id = str(uuid.uuid4())
if not messages:
raise ValueError("At least one message is required for A2A completion")
# Convert all messages to maintain conversation history
# Use helper to format conversation with role prefixes
full_context = convert_messages_to_prompt(messages)
# Create single A2A message with full conversation context
a2a_message = {
"role": "user",
"parts": [{"kind": "text", "text": full_context}],
"messageId": str(uuid.uuid4()),
}
# Build JSON-RPC 2.0 request
# For A2A protocol, the method is "message/send" for non-streaming
# and "message/stream" for streaming (handled by optional_params["stream"])
method = "message/stream" if optional_params.get("stream") else "message/send"
request_data = {
"jsonrpc": "2.0",
"id": request_id,
"method": method,
"params": {
"message": a2a_message
}
}
return request_data
def transform_response(
self,
model: str,
raw_response: httpx.Response,
model_response: ModelResponse,
logging_obj: Any,
request_data: dict,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
encoding: Any,
api_key: Optional[str] = None,
json_mode: Optional[bool] = None,
) -> ModelResponse:
"""
Transform A2A JSON-RPC 2.0 response to OpenAI format.
Args:
model: Model name
raw_response: HTTP response from A2A agent
model_response: Model response object to populate
logging_obj: Logging object
request_data: Original request data
messages: Original messages
optional_params: Optional parameters
litellm_params: LiteLLM parameters
encoding: Encoding object
api_key: API key
json_mode: JSON mode flag
Returns:
Populated ModelResponse object
"""
try:
response_json = raw_response.json()
except Exception as e:
raise A2AError(
status_code=raw_response.status_code,
message=f"Failed to parse A2A response: {str(e)}",
headers=dict(raw_response.headers),
)
# Check for JSON-RPC error
if "error" in response_json:
error = response_json["error"]
raise A2AError(
status_code=raw_response.status_code,
message=f"A2A error: {error.get('message', 'Unknown error')}",
headers=dict(raw_response.headers),
)
# Extract text from A2A response
text = extract_text_from_a2a_response(response_json)
# Populate model response
model_response.choices = [
Choices(
finish_reason="stop",
index=0,
message=Message(
content=text,
role="assistant",
),
)
]
# Set model
model_response.model = model
# Set ID from response
model_response.id = response_json.get("id", str(uuid.uuid4()))
return model_response
def get_model_response_iterator(
self,
streaming_response: Union[Iterator, Any],
sync_stream: bool,
json_mode: Optional[bool] = False,
) -> BaseModelResponseIterator:
"""
Get streaming iterator for A2A responses.
Args:
streaming_response: Streaming response iterator
sync_stream: Whether this is a sync stream
json_mode: JSON mode flag
Returns:
A2A streaming iterator
"""
return A2AModelResponseIterator(
streaming_response=streaming_response,
sync_stream=sync_stream,
json_mode=json_mode,
)
def _openai_message_to_a2a_message(self, message: Dict[str, Any]) -> Dict[str, Any]:
"""
Convert OpenAI message to A2A message format.
Args:
message: OpenAI message dict
Returns:
A2A message dict
"""
content = message.get("content", "")
role = message.get("role", "user")
return {
"role": role,
"parts": [{"kind": "text", "text": str(content)}],
"messageId": str(uuid.uuid4()),
}
def get_error_class(
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
) -> BaseLLMException:
"""Return appropriate error class for A2A errors"""
# Convert headers to dict if needed
headers_dict = dict(headers) if isinstance(headers, httpx.Headers) else headers
return A2AError(
status_code=status_code,
message=error_message,
headers=headers_dict,
)

134
litellm/llms/a2a/common_utils.py Normal file
View file

@ -0,0 +1,134 @@
"""
Common utilities for A2A (Agent-to-Agent) Protocol
"""
from typing import Any, Dict, List
from pydantic import BaseModel
from litellm.litellm_core_utils.prompt_templates.common_utils import (
convert_content_list_to_str,
)
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.types.llms.openai import AllMessageValues
class A2AError(BaseLLMException):
"""Base exception for A2A protocol errors"""
def __init__(
self,
status_code: int,
message: str,
headers: Dict[str, Any] = {},
):
super().__init__(
status_code=status_code,
message=message,
headers=headers,
)
def convert_messages_to_prompt(messages: List[AllMessageValues]) -> str:
"""
Convert OpenAI messages to a single prompt string for A2A agent.
Formats each message as "{role}: {content}" and joins with newlines
to preserve conversation history. Handles both string and list content.
Args:
messages: List of OpenAI-format messages
Returns:
Formatted prompt string with full conversation context
"""
conversation_parts = []
for msg in messages:
# Use LiteLLM's helper to extract text from content (handles both str and list)
content_text = convert_content_list_to_str(message=msg)
# Get role
if isinstance(msg, BaseModel):
role = msg.model_dump().get("role", "user")
elif isinstance(msg, dict):
role = msg.get("role", "user")
else:
role = dict(msg).get("role", "user") # type: ignore
if content_text:
conversation_parts.append(f"{role}: {content_text}")
return "\n".join(conversation_parts)
def extract_text_from_a2a_message(
message: Dict[str, Any], depth: int = 0, max_depth: int = 10
) -> str:
"""
Extract text content from A2A message parts.
Args:
message: A2A message dict with 'parts' containing text parts
depth: Current recursion depth (internal use)
max_depth: Maximum recursion depth to prevent infinite loops
Returns:
Concatenated text from all text parts
"""
if message is None or depth >= max_depth:
return ""
parts = message.get("parts", [])
text_parts: List[str] = []
for part in parts:
if part.get("kind") == "text":
text_parts.append(part.get("text", ""))
# Handle nested parts if they exist
elif "parts" in part:
nested_text = extract_text_from_a2a_message(part, depth + 1, max_depth)
if nested_text:
text_parts.append(nested_text)
return " ".join(text_parts)
def extract_text_from_a2a_response(
response_dict: Dict[str, Any], max_depth: int = 10
) -> str:
"""
Extract text content from A2A response result.
Args:
response_dict: A2A response dict with 'result' containing message
max_depth: Maximum recursion depth to prevent infinite loops
Returns:
Text from response message parts
"""
result = response_dict.get("result", {})
if not isinstance(result, dict):
return ""
# A2A response can have different formats:
# 1. Direct message: {"result": {"kind": "message", "parts": [...]}}
# 2. Nested message: {"result": {"message": {"parts": [...]}}}
# 3. Task with artifacts: {"result": {"kind": "task", "artifacts": [{"parts": [...]}]}}
# Check if result itself has parts (direct message)
if "parts" in result:
return extract_text_from_a2a_message(result, depth=0, max_depth=max_depth)
# Check for nested message
message = result.get("message")
if message:
return extract_text_from_a2a_message(message, depth=0, max_depth=max_depth)
# Handle task result with artifacts
artifacts = result.get("artifacts", [])
if artifacts and len(artifacts) > 0:
first_artifact = artifacts[0]
return extract_text_from_a2a_message(
first_artifact, depth=0, max_depth=max_depth
)
return ""

40
litellm/main.py
View file

@ -2199,6 +2199,38 @@ def completion( # type: ignore # noqa: PLR0915
logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements
client=client,
)
elif custom_llm_provider == "a2a":
# A2A (Agent-to-Agent) Protocol
api_base = (
api_base
or litellm.api_base
or get_secret_str("A2A_API_BASE")
)
if api_base is None:
raise Exception("api_base is required for A2A provider")
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,
provider_config=provider_config,
)
elif custom_llm_provider == "gigachat":
# GigaChat - Sber AI's LLM (Russia)
api_key = (
@ -3113,8 +3145,8 @@ def completion( # type: ignore # noqa: PLR0915
api_key
or litellm.api_key
or litellm.openrouter_key
or get_secret("OPENROUTER_API_KEY")
or get_secret("OR_API_KEY")
or get_secret_str("OPENROUTER_API_KEY")
or get_secret_str("OR_API_KEY")
)
openrouter_site_url = get_secret("OR_SITE_URL") or "https://litellm.ai"
@ -4884,8 +4916,8 @@ def embedding( # noqa: PLR0915
api_key
or litellm.api_key
or litellm.openrouter_key
or get_secret("OPENROUTER_API_KEY")
or get_secret("OR_API_KEY")
or get_secret_str("OPENROUTER_API_KEY")
or get_secret_str("OR_API_KEY")
)
openrouter_site_url = get_secret("OR_SITE_URL") or "https://litellm.ai"

28
litellm/model_prices_and_context_window_backup.json
View file

@ -27113,6 +27113,34 @@
"supports_reasoning": true,
"supports_tool_choice": true
},
"together_ai/zai-org/GLM-4.7": {
"input_cost_per_token": 4.5e-07,
"litellm_provider": "together_ai",
"max_input_tokens": 200000,
"max_output_tokens": 200000,
"max_tokens": 200000,
"mode": "chat",
"output_cost_per_token": 2e-06,
"source": "https://www.together.ai/models/glm-4-7",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"together_ai/moonshotai/Kimi-K2.5": {
"input_cost_per_token": 5e-07,
"litellm_provider": "together_ai",
"max_input_tokens": 256000,
"max_output_tokens": 256000,
"max_tokens": 256000,
"mode": "chat",
"output_cost_per_token": 2.8e-06,
"source": "https://www.together.ai/models/kimi-k2-5",
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_reasoning": true
},
"together_ai/moonshotai/Kimi-K2-Instruct-0905": {
"input_cost_per_token": 1e-06,
"litellm_provider": "together_ai",

1
litellm/types/utils.py
View file

@ -3029,6 +3029,7 @@ class LlmProviders(str, Enum):
MISTRAL = "mistral"
MILVUS = "milvus"
GROQ = "groq"
A2A = "a2a"
GIGACHAT = "gigachat"
NVIDIA_NIM = "nvidia_nim"
CEREBRAS = "cerebras"

8
litellm/utils.py
View file

@ -1453,6 +1453,10 @@ def client(original_function): # noqa: PLR0915
logging_obj, kwargs = function_setup(
original_function.__name__, rules_obj, start_time, *args, **kwargs
)
# Type assertion: logging_obj is guaranteed to be non-None after function_setup
assert logging_obj is not None, "logging_obj should not be None after function_setup"
## LOAD CREDENTIALS
load_credentials_from_list(kwargs)
kwargs["litellm_logging_obj"] = logging_obj
@ -1771,6 +1775,9 @@ def client(original_function): # noqa: PLR0915
logging_obj, kwargs = function_setup(
original_function.__name__, rules_obj, start_time, *args, **kwargs
)
# Type assertion: logging_obj is guaranteed to be non-None after function_setup
assert logging_obj is not None, "logging_obj should not be None after function_setup"
modified_kwargs = await async_pre_call_deployment_hook(kwargs, call_type)
if modified_kwargs is not None:
@ -7799,6 +7806,7 @@ class ProviderConfigManager:
# Simple provider mappings (no model parameter needed)
LlmProviders.DEEPSEEK: (lambda: litellm.DeepSeekChatConfig(), False),
LlmProviders.GROQ: (lambda: litellm.GroqChatConfig(), False),
LlmProviders.A2A: (lambda: litellm.A2AConfig(), False),
LlmProviders.BYTEZ: (lambda: litellm.BytezChatConfig(), False),
LlmProviders.DATABRICKS: (lambda: litellm.DatabricksConfig(), False),
LlmProviders.XAI: (lambda: litellm.XAIChatConfig(), False),

20
poetry.lock generated
View file

@ -1,4 +1,4 @@
# This file is automatically @generated by Poetry 2.1.4 and should not be changed by hand.
# This file is automatically @generated by Poetry 2.2.0 and should not be changed by hand.
[[package]]
name = "a2a-sdk"
@ -5704,24 +5704,6 @@ pytest = ">=7.0.0"
[package.extras]
dev = ["black", "flake8", "isort", "mypy"]
[[package]]
name = "pytest-retry"
version = "1.7.0"
description = "Adds the ability to retry flaky tests in CI environments"
optional = false
python-versions = ">=3.9"
groups = ["dev"]
files = [
{file = "pytest_retry-1.7.0-py3-none-any.whl", hash = "sha256:a2dac85b79a4e2375943f1429479c65beb6c69553e7dae6b8332be47a60954f4"},
{file = "pytest_retry-1.7.0.tar.gz", hash = "sha256:f8d52339f01e949df47c11ba9ee8d5b362f5824dff580d3870ec9ae0057df80f"},
]
[package.dependencies]
pytest = ">=7.0.0"
[package.extras]
dev = ["black", "flake8", "isort", "mypy"]
[[package]]
name = "python-dateutil"
version = "2.9.0.post0"

17
provider_endpoints_support.json
View file

@ -32,6 +32,23 @@
}
},
"providers": {
"a2a": {
"display_name": "A2A (Agent-to-Agent) (`a2a`)",
"url": "https://docs.litellm.ai/docs/providers/a2a",
"endpoints": {
"chat_completions": true,
"messages": false,
"responses": false,
"embeddings": false,
"image_generations": false,
"audio_transcriptions": false,
"audio_speech": false,
"moderations": false,
"batches": false,
"rerank": false,
"a2a": false
}
},
"abliteration": {
"display_name": "Abliteration (`abliteration`)",
"url": "https://docs.litellm.ai/docs/providers/abliteration",

2
requirements.txt
View file

@ -73,4 +73,4 @@ pypdf>=6.6.2 # for PDF text extraction in RAG ingestion
########################
# LITELLM ENTERPRISE DEPENDENCIES
########################
litellm-enterprise==0.1.28
litellm-enterprise==0.1.29

1
tests/code_coverage_tests/recursive_detector.py
View file

@ -40,6 +40,7 @@ IGNORE_FUNCTIONS = [
"filter_exceptions_from_params", # max depth set (default 20) to prevent infinite recursion.
"__getattr__", # lazy loading pattern in litellm/__init__.py with proper caching to prevent infinite recursion.
"_validate_inheritance_chain", # max depth set (default 100) to prevent infinite recursion in policy inheritance validation.
"extract_text_from_a2a_message", # max depth set (default 10) to prevent infinite recursion in A2A message parsing.
]

132
tests/llm_translation/test_a2a.py Normal file
View file

@ -0,0 +1,132 @@
"""
Minimal E2E tests for A2A (Agent-to-Agent) Protocol provider.
Tests validate that the endpoint is reachable and can handle both
streaming and non-streaming requests.
"""
import os
import sys
import pytest
sys.path.insert(0, os.path.abspath("../.."))
import litellm
@pytest.mark.asyncio
async def test_a2a_completion_async_non_streaming():
"""
Test A2A provider with async non-streaming request.
Minimal test to validate endpoint reachability.
Note: Requires an A2A agent running at http://0.0.0.0:9999
Set A2A_API_BASE environment variable to use a different endpoint.
"""
api_base = os.environ.get("A2A_API_BASE", "http://0.0.0.0:9999")
try:
response = await litellm.acompletion(
model="a2a/test-agent",
messages=[{"role": "user", "content": "Hello"}],
api_base=api_base,
stream=False,
)
print(f"Response: {response}")
assert response is not None, "Expected non-None response"
print(f"✅ Async non-streaming test passed")
except litellm.exceptions.APIConnectionError as e:
pytest.skip(f"A2A agent not reachable at {api_base}: {e}")
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@pytest.mark.asyncio
async def test_a2a_completion_async_streaming():
"""
Test A2A provider with async streaming request.
Minimal test to validate streaming endpoint reachability.
"""
api_base = os.environ.get("A2A_API_BASE", "http://0.0.0.0:9999")
try:
response = await litellm.acompletion(
model="a2a/test-agent",
messages=[{"role": "user", "content": "Hello"}],
api_base=api_base,
stream=True,
)
chunks = []
async for chunk in response: # type: ignore
chunks.append(chunk)
print(f"Chunk: {chunk}")
assert len(chunks) > 0, "Expected at least one chunk in streaming response"
print(f"✅ Async streaming test passed: received {len(chunks)} chunks")
except litellm.exceptions.APIConnectionError as e:
pytest.skip(f"A2A agent not reachable at {api_base}: {e}")
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_a2a_completion_sync():
"""
Test A2A provider with synchronous non-streaming request.
Minimal test to validate sync endpoint reachability.
"""
api_base = os.environ.get("A2A_API_BASE", "http://0.0.0.0:9999")
try:
response = litellm.completion(
model="a2a/test-agent",
messages=[{"role": "user", "content": "Hello"}],
api_base=api_base,
stream=False,
)
print(f"Response: {response}")
assert response is not None, "Expected non-None response"
print(f"✅ Sync non-streaming test passed")
except litellm.exceptions.APIConnectionError as e:
pytest.skip(f"A2A agent not reachable at {api_base}: {e}")
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_a2a_completion_sync_streaming():
"""
Test A2A provider with synchronous streaming request.
Minimal test to validate sync streaming endpoint reachability.
"""
api_base = os.environ.get("A2A_API_BASE", "http://0.0.0.0:9999")
try:
response = litellm.completion(
model="a2a/test-agent",
messages=[{"role": "user", "content": "Hello"}],
api_base=api_base,
stream=True,
)
chunks = []
for chunk in response: # type: ignore
chunks.append(chunk)
print(f"Chunk: {chunk}")
assert len(chunks) > 0, "Expected at least one chunk in streaming response"
print(f"✅ Sync streaming test passed: received {len(chunks)} chunks")
except litellm.exceptions.APIConnectionError as e:
pytest.skip(f"A2A agent not reachable at {api_base}: {e}")
except Exception as e:
pytest.fail(f"Error occurred: {e}")

1
ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx
View file

@ -400,6 +400,7 @@ const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({ premiumUser, te
all_models_on_proxy={allModelsOnProxy}
getDisplayModelName={getDisplayModelName}
setSelectedModelId={setSelectedModelId}
teams={teams}
/>
</TabPanel>
<ModelRetrySettingsTab

4
ui/litellm-dashboard/src/components/model_dashboard/HealthCheckComponent.tsx
View file

@ -7,6 +7,7 @@ import { healthCheckColumns } from "./health_check_columns";
import { errorPatterns } from "@/utils/errorPatterns";
import { individualModelHealthCheckCall, latestHealthChecksCall } from "../networking";
import { Table as TableInstance } from "@tanstack/react-table";
import { Team } from "../key_team_helpers/key_list";
interface HealthStatus {
status: string;
@ -24,6 +25,7 @@ interface HealthCheckComponentProps {
all_models_on_proxy: string[];
getDisplayModelName: (model: any) => string;
setSelectedModelId?: (modelId: string) => void;
teams?: Team[] | null;
}
const HealthCheckComponent: React.FC<HealthCheckComponentProps> = ({
@ -32,6 +34,7 @@ const HealthCheckComponent: React.FC<HealthCheckComponentProps> = ({
all_models_on_proxy,
getDisplayModelName,
setSelectedModelId,
teams,
}) => {
const [modelHealthStatuses, setModelHealthStatuses] = useState<{ [key: string]: HealthStatus }>({});
const [selectedModelsForHealth, setSelectedModelsForHealth] = useState<string[]>([]);
@ -574,6 +577,7 @@ const HealthCheckComponent: React.FC<HealthCheckComponentProps> = ({
showErrorModal,
showSuccessModal,
setSelectedModelId,
teams,
)}
data={modelData.data.map((model: any) => {
const modelName = model.model_name;

27
ui/litellm-dashboard/src/components/model_dashboard/health_check_columns.tsx
View file

@ -2,6 +2,7 @@ import { ColumnDef } from "@tanstack/react-table";
import { Tooltip, Checkbox } from "antd";
import { Text } from "@tremor/react";
import { InformationCircleIcon, PlayIcon, RefreshIcon } from "@heroicons/react/outline";
import { Team } from "@/components/key_team_helpers/key_list";
interface HealthCheckData {
model_name: string;
@ -42,6 +43,7 @@ export const healthCheckColumns = (
showErrorModal?: (modelName: string, cleanedError: string, fullError: string) => void,
showSuccessModal?: (modelName: string, response: any) => void,
setSelectedModelId?: (modelId: string) => void,
teams?: Team[] | null,
): ColumnDef<HealthCheckData>[] => [
{
header: () => (
@ -100,6 +102,31 @@ export const healthCheckColumns = (
);
},
},
{
header: "Team Alias",
accessorKey: "model_info.team_id",
enableSorting: true,
sortingFn: "alphanumeric",
cell: ({ row }) => {
const model = row.original;
const teamId = model.model_info?.team_id;
if (!teamId) {
return <span className="text-gray-400 text-sm">-</span>;
}
const team = teams?.find((t) => t.team_id === teamId);
const teamAlias = team?.team_alias || teamId;
return (
<div className="text-sm">
<Tooltip title={teamAlias}>
<div className="truncate max-w-[150px]">{teamAlias}</div>
</Tooltip>
</div>
);
},
},
{
header: "Health Status",
accessorKey: "health_status",

1
ui/litellm-dashboard/src/components/templates/model_dashboard.tsx
View file

@ -1368,6 +1368,7 @@ const OldModelDashboard: React.FC<ModelDashboardProps> = ({
all_models_on_proxy={all_models_on_proxy}
getDisplayModelName={getDisplayModelName}
setSelectedModelId={setSelectedModelId}
teams={teams}
/>
</TabPanel>
<TabPanel>