mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
test agnts
This commit is contained in:
parent
5c7f4636b1
commit
3621e512f6
3 changed files with 55 additions and 77 deletions
|
|
@ -1,8 +1,7 @@
|
|||
"""
|
||||
Custom A2A Card Resolver for LiteLLM.
|
||||
|
||||
Extends the A2A SDK's card resolver to support multiple well-known paths
|
||||
and fixes agent card URL issues where the card contains internal/localhost URLs.
|
||||
Extends the A2A SDK's card resolver to support multiple well-known paths.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, Any, Dict, Optional
|
||||
|
|
@ -30,7 +29,7 @@ except ImportError:
|
|||
|
||||
def is_localhost_or_internal_url(url: Optional[str]) -> bool:
|
||||
"""
|
||||
Check if a URL is a localhost or internal URL that should be replaced.
|
||||
Check if a URL is a localhost or internal URL.
|
||||
|
||||
This detects common development URLs that are accidentally left in
|
||||
agent cards when deploying to production.
|
||||
|
|
@ -39,17 +38,7 @@ def is_localhost_or_internal_url(url: Optional[str]) -> bool:
|
|||
url: The URL to check
|
||||
|
||||
Returns:
|
||||
True if the URL is localhost/internal and should be replaced
|
||||
|
||||
Examples:
|
||||
>>> is_localhost_or_internal_url("http://localhost:8000/")
|
||||
True
|
||||
>>> is_localhost_or_internal_url("http://0.0.0.0:8001/")
|
||||
True
|
||||
>>> is_localhost_or_internal_url("https://my-agent.example.com/")
|
||||
False
|
||||
>>> is_localhost_or_internal_url(None)
|
||||
False
|
||||
True if the URL is localhost/internal
|
||||
"""
|
||||
if not url:
|
||||
return False
|
||||
|
|
@ -59,50 +48,40 @@ def is_localhost_or_internal_url(url: Optional[str]) -> bool:
|
|||
return any(pattern in url_lower for pattern in LOCALHOST_URL_PATTERNS)
|
||||
|
||||
|
||||
def fix_agent_card_url(agent_card: "AgentCard", base_url: str) -> "AgentCard":
|
||||
"""
|
||||
Fix the agent card URL if it contains a localhost/internal address.
|
||||
|
||||
Many A2A agents are deployed with agent cards that contain internal URLs
|
||||
like "http://0.0.0.0:8001/" or "http://localhost:8000/". This function
|
||||
replaces such URLs with the provided base_url.
|
||||
|
||||
Args:
|
||||
agent_card: The agent card to fix
|
||||
base_url: The base URL to use as replacement
|
||||
|
||||
Returns:
|
||||
The agent card with the URL fixed if necessary
|
||||
"""
|
||||
card_url = getattr(agent_card, "url", None)
|
||||
|
||||
if card_url and is_localhost_or_internal_url(card_url):
|
||||
# Normalize base_url to ensure it ends with /
|
||||
fixed_url = base_url.rstrip("/") + "/"
|
||||
agent_card.url = fixed_url
|
||||
|
||||
return agent_card
|
||||
|
||||
|
||||
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) and /.well-known/agent.json (previous/alternative)
|
||||
- Fix agent card URLs that contain localhost/internal addresses by replacing them with the original base_url
|
||||
|
||||
This fixes a common issue where agent cards are deployed with internal URLs
|
||||
(e.g., "http://0.0.0.0:8001/") that don't work when accessed from external clients.
|
||||
|
||||
Extends the base A2ACardResolver to try both:
|
||||
- /.well-known/agent-card.json (standard)
|
||||
- /.well-known/agent.json (previous/alternative)
|
||||
"""
|
||||
|
||||
def _fix_agent_card_url(self, agent_card: "AgentCard") -> "AgentCard":
|
||||
"""
|
||||
Fix the agent card URL if it contains a localhost/internal address.
|
||||
|
||||
Many A2A agents are deployed with agent cards that contain internal URLs
|
||||
like "http://0.0.0.0:8001/" or "http://localhost:8000/". This method
|
||||
replaces such URLs with the original base_url used to fetch the card.
|
||||
|
||||
Args:
|
||||
agent_card: The agent card to fix
|
||||
|
||||
Returns:
|
||||
The agent card with the URL fixed if necessary
|
||||
"""
|
||||
card_url = getattr(agent_card, "url", None)
|
||||
|
||||
if card_url and is_localhost_or_internal_url(card_url):
|
||||
# Normalize base_url to ensure it ends with /
|
||||
fixed_url = self.base_url.rstrip("/") + "/"
|
||||
|
||||
verbose_logger.warning(
|
||||
f"Agent card contains localhost/internal URL '{card_url}'. "
|
||||
f"Replacing with base_url '{fixed_url}'. "
|
||||
f"Consider updating the agent's configuration to use the correct public URL."
|
||||
)
|
||||
|
||||
# Create a new agent card with the fixed URL
|
||||
# We need to handle this carefully to preserve all other fields
|
||||
agent_card.url = fixed_url
|
||||
|
||||
return agent_card
|
||||
|
||||
|
||||
async def get_agent_card(
|
||||
self,
|
||||
relative_card_path: Optional[str] = None,
|
||||
|
|
@ -110,57 +89,54 @@ class LiteLLMA2ACardResolver(_A2ACardResolver): # type: ignore[misc]
|
|||
) -> "AgentCard":
|
||||
"""
|
||||
Fetch the agent card, trying multiple well-known paths.
|
||||
|
||||
|
||||
First tries the standard path, then falls back to the previous path.
|
||||
Also fixes agent card URLs that contain localhost/internal addresses.
|
||||
|
||||
|
||||
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 (with URL fixed if necessary)
|
||||
|
||||
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:
|
||||
agent_card = await super().get_agent_card(
|
||||
return await super().get_agent_card(
|
||||
relative_card_path=relative_card_path,
|
||||
http_kwargs=http_kwargs,
|
||||
)
|
||||
return self._fix_agent_card_url(agent_card)
|
||||
|
||||
|
||||
# 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}"
|
||||
)
|
||||
agent_card = await super().get_agent_card(
|
||||
return await super().get_agent_card(
|
||||
relative_card_path=path,
|
||||
http_kwargs=http_kwargs,
|
||||
)
|
||||
return self._fix_agent_card_url(agent_card)
|
||||
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}. "
|
||||
|
|
|
|||
|
|
@ -1,6 +1,11 @@
|
|||
"""
|
||||
Simple A2A agent tests - non-streaming and streaming.
|
||||
|
||||
These tests validate the localhost URL retry logic: if an A2A agent's card
|
||||
contains a localhost/internal URL (e.g., http://0.0.0.0:8001/), the request
|
||||
will fail with a connection error. LiteLLM detects this and automatically
|
||||
retries using the original api_base URL instead.
|
||||
|
||||
Requires A2A_AGENT_URL environment variable to be set.
|
||||
|
||||
Run with:
|
||||
|
|
|
|||
|
|
@ -9,7 +9,10 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
|||
|
||||
import pytest
|
||||
|
||||
from litellm.a2a_protocol.card_resolver import is_localhost_or_internal_url
|
||||
from litellm.a2a_protocol.card_resolver import (
|
||||
fix_agent_card_url,
|
||||
is_localhost_or_internal_url,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -83,19 +86,13 @@ def test_is_localhost_or_internal_url():
|
|||
|
||||
|
||||
def test_fix_agent_card_url_replaces_localhost():
|
||||
"""Test that _fix_agent_card_url replaces localhost URLs with base_url."""
|
||||
from litellm.a2a_protocol.card_resolver import LiteLLMA2ACardResolver
|
||||
|
||||
"""Test that fix_agent_card_url replaces localhost URLs with base_url."""
|
||||
# Create mock agent card with localhost URL
|
||||
mock_card = MagicMock()
|
||||
mock_card.url = "http://0.0.0.0:8001/"
|
||||
|
||||
# Create resolver instance without calling __init__
|
||||
resolver = LiteLLMA2ACardResolver.__new__(LiteLLMA2ACardResolver)
|
||||
resolver.base_url = "https://my-public-agent.example.com"
|
||||
|
||||
# Fix the URL
|
||||
result = resolver._fix_agent_card_url(mock_card)
|
||||
result = fix_agent_card_url(mock_card, "https://my-public-agent.example.com")
|
||||
|
||||
# Verify localhost URL was replaced with base_url
|
||||
assert result.url == "https://my-public-agent.example.com/"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue