mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-08 02:24:54 +00:00
[Fix] A2a Gateway - Allow supporting old A2a card formats (#19949)
* fix: LiteLLMA2ACardResolver * fix: LiteLLMA2ACardResolver * feat: .well-known/agent.json * test_card_resolver_fallback_from_new_to_old_path
This commit is contained in:
@@ -0,0 +1,97 @@
|
||||
"""
|
||||
Custom A2A Card Resolver for LiteLLM.
|
||||
|
||||
Extends the A2A SDK's card resolver to support multiple well-known paths.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, Any, Dict, Optional
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from a2a.types import AgentCard
|
||||
|
||||
# Runtime imports with availability check
|
||||
_A2ACardResolver: Any = None
|
||||
AGENT_CARD_WELL_KNOWN_PATH: str = "/.well-known/agent-card.json"
|
||||
PREV_AGENT_CARD_WELL_KNOWN_PATH: str = "/.well-known/agent.json"
|
||||
|
||||
try:
|
||||
from a2a.client import A2ACardResolver as _A2ACardResolver # type: ignore[no-redef]
|
||||
from a2a.utils.constants import ( # type: ignore[no-redef]
|
||||
AGENT_CARD_WELL_KNOWN_PATH,
|
||||
PREV_AGENT_CARD_WELL_KNOWN_PATH,
|
||||
)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
|
||||
class LiteLLMA2ACardResolver(_A2ACardResolver): # type: ignore[misc]
|
||||
"""
|
||||
Custom A2A card resolver that supports multiple well-known paths.
|
||||
|
||||
Extends the base A2ACardResolver to try both:
|
||||
- /.well-known/agent-card.json (standard)
|
||||
- /.well-known/agent.json (previous/alternative)
|
||||
"""
|
||||
|
||||
async def get_agent_card(
|
||||
self,
|
||||
relative_card_path: Optional[str] = None,
|
||||
http_kwargs: Optional[Dict[str, Any]] = None,
|
||||
) -> "AgentCard":
|
||||
"""
|
||||
Fetch the agent card, trying multiple well-known paths.
|
||||
|
||||
First tries the standard path, then falls back to the previous path.
|
||||
|
||||
Args:
|
||||
relative_card_path: Optional path to the agent card endpoint.
|
||||
If None, tries both well-known paths.
|
||||
http_kwargs: Optional dictionary of keyword arguments to pass to httpx.get
|
||||
|
||||
Returns:
|
||||
AgentCard from the A2A agent
|
||||
|
||||
Raises:
|
||||
A2AClientHTTPError or A2AClientJSONError if both paths fail
|
||||
"""
|
||||
# If a specific path is provided, use the parent implementation
|
||||
if relative_card_path is not None:
|
||||
return await super().get_agent_card(
|
||||
relative_card_path=relative_card_path,
|
||||
http_kwargs=http_kwargs,
|
||||
)
|
||||
|
||||
# Try both well-known paths
|
||||
paths = [
|
||||
AGENT_CARD_WELL_KNOWN_PATH,
|
||||
PREV_AGENT_CARD_WELL_KNOWN_PATH,
|
||||
]
|
||||
|
||||
last_error = None
|
||||
for path in paths:
|
||||
try:
|
||||
verbose_logger.debug(
|
||||
f"Attempting to fetch agent card from {self.base_url}{path}"
|
||||
)
|
||||
return await super().get_agent_card(
|
||||
relative_card_path=path,
|
||||
http_kwargs=http_kwargs,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.debug(
|
||||
f"Failed to fetch agent card from {self.base_url}{path}: {e}"
|
||||
)
|
||||
last_error = e
|
||||
continue
|
||||
|
||||
# If we get here, all paths failed - re-raise the last error
|
||||
if last_error is not None:
|
||||
raise last_error
|
||||
|
||||
# This shouldn't happen, but just in case
|
||||
raise Exception(
|
||||
f"Failed to fetch agent card from {self.base_url}. "
|
||||
f"Tried paths: {', '.join(paths)}"
|
||||
)
|
||||
@@ -6,6 +6,7 @@ Provides standalone functions with @client decorator for LiteLLM logging integra
|
||||
|
||||
import asyncio
|
||||
import datetime
|
||||
import uuid
|
||||
from typing import TYPE_CHECKING, Any, AsyncIterator, Coroutine, Dict, Optional, Union
|
||||
|
||||
import litellm
|
||||
@@ -20,7 +21,6 @@ from litellm.llms.custom_httpx.http_handler import (
|
||||
)
|
||||
from litellm.types.agents import LiteLLMSendMessageResponse
|
||||
from litellm.utils import client
|
||||
import uuid
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from a2a.client import A2AClient as A2AClientType
|
||||
@@ -36,13 +36,18 @@ A2ACardResolver: Any = None
|
||||
_A2AClient: Any = None
|
||||
|
||||
try:
|
||||
from a2a.client import A2ACardResolver # type: ignore[no-redef]
|
||||
from a2a.client import A2AClient as _A2AClient # type: ignore[no-redef]
|
||||
|
||||
A2A_SDK_AVAILABLE = True
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# Import our custom card resolver that supports multiple well-known paths
|
||||
from litellm.a2a_protocol.card_resolver import LiteLLMA2ACardResolver
|
||||
|
||||
# Use our custom resolver instead of the default A2A SDK resolver
|
||||
A2ACardResolver = LiteLLMA2ACardResolver
|
||||
|
||||
|
||||
def _set_usage_on_logging_obj(
|
||||
kwargs: Dict[str, Any],
|
||||
|
||||
@@ -119,6 +119,11 @@ async def _handle_stream_message(
|
||||
tags=["[beta] A2A Agents"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
)
|
||||
@router.get(
|
||||
"/a2a/{agent_id}/.well-known/agent.json",
|
||||
tags=["[beta] A2A Agents"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
)
|
||||
async def get_agent_card(
|
||||
agent_id: str,
|
||||
request: Request,
|
||||
@@ -127,6 +132,10 @@ async def get_agent_card(
|
||||
"""
|
||||
Get the agent card for an agent (A2A discovery endpoint).
|
||||
|
||||
Supports both standard paths:
|
||||
- /.well-known/agent-card.json
|
||||
- /.well-known/agent.json
|
||||
|
||||
The URL in the agent card is rewritten to point to the LiteLLM proxy,
|
||||
so all subsequent A2A calls go through LiteLLM for logging and cost tracking.
|
||||
"""
|
||||
|
||||
@@ -0,0 +1,69 @@
|
||||
"""
|
||||
Mock tests for LiteLLMA2ACardResolver.
|
||||
|
||||
Tests that the card resolver tries both old and new well-known paths.
|
||||
"""
|
||||
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_card_resolver_fallback_from_new_to_old_path():
|
||||
"""
|
||||
Test that the card resolver tries the new path (/.well-known/agent-card.json) first,
|
||||
and falls back to the old path (/.well-known/agent.json) if the new path fails.
|
||||
"""
|
||||
# Mock the AgentCard
|
||||
mock_agent_card = MagicMock()
|
||||
mock_agent_card.name = "Test Agent"
|
||||
mock_agent_card.description = "A test agent"
|
||||
|
||||
# Track which paths were called
|
||||
paths_called = []
|
||||
|
||||
# Create a mock base class
|
||||
class MockA2ACardResolver:
|
||||
def __init__(self, base_url):
|
||||
self.base_url = base_url
|
||||
|
||||
async def get_agent_card(self, relative_card_path=None, http_kwargs=None):
|
||||
paths_called.append(relative_card_path)
|
||||
if relative_card_path == "/.well-known/agent-card.json":
|
||||
# First call (new path) fails
|
||||
raise Exception("404 Not Found")
|
||||
else:
|
||||
# Second call (old path) succeeds
|
||||
return mock_agent_card
|
||||
|
||||
# Create mock A2A module
|
||||
mock_a2a_module = MagicMock()
|
||||
mock_a2a_client = MagicMock()
|
||||
mock_a2a_constants = MagicMock()
|
||||
mock_a2a_constants.AGENT_CARD_WELL_KNOWN_PATH = "/.well-known/agent-card.json"
|
||||
mock_a2a_constants.PREV_AGENT_CARD_WELL_KNOWN_PATH = "/.well-known/agent.json"
|
||||
|
||||
with patch.dict(
|
||||
sys.modules,
|
||||
{
|
||||
"a2a": mock_a2a_module,
|
||||
"a2a.client": MagicMock(A2ACardResolver=MockA2ACardResolver),
|
||||
"a2a.utils.constants": mock_a2a_constants,
|
||||
},
|
||||
):
|
||||
# Import after patching
|
||||
from litellm.a2a_protocol.card_resolver import LiteLLMA2ACardResolver
|
||||
|
||||
resolver = LiteLLMA2ACardResolver(base_url="http://test-agent:8000")
|
||||
result = await resolver.get_agent_card()
|
||||
|
||||
# Verify both paths were tried in correct order
|
||||
assert len(paths_called) == 2
|
||||
assert paths_called[0] == "/.well-known/agent-card.json" # New path tried first
|
||||
assert paths_called[1] == "/.well-known/agent.json" # Old path tried second
|
||||
|
||||
# Verify the result
|
||||
assert result == mock_agent_card
|
||||
assert result.name == "Test Agent"
|
||||
Reference in New Issue
Block a user