mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-04 02:31:27 +00:00
fix(a2a): forward signature_verifier and type the resolver fallback
Annotating the absent-SDK fallback as Any meant type checkers could not validate the subclass at all, which is what the previous commit's fallback relied on. Declare the base under TYPE_CHECKING and fall back to object only at runtime, matching the idiom already used in types/utils.py and llms/bedrock.py With the real base class visible, basedpyright immediately caught that get_agent_card drops the SDK's signature_verifier parameter, so any caller passing it hit a TypeError. Forward it to both super() calls basedpyright on this file goes from 25 errors to 14 Signed-off-by: Aryan Gorde <aryangorde8@gmail.com>
This commit is contained in:
parent
327e9ffd72
commit
246c27bac1
2 changed files with 46 additions and 3 deletions
|
|
@ -4,21 +4,26 @@ Custom A2A Card Resolver for LiteLLM.
|
|||
Extends the A2A SDK's card resolver to support multiple well-known paths.
|
||||
"""
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import LOCALHOST_URL_PATTERNS
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from a2a.client import A2ACardResolver as _A2ACardResolver
|
||||
from a2a.types import AgentCard
|
||||
else:
|
||||
try:
|
||||
from a2a.client import A2ACardResolver as _A2ACardResolver
|
||||
except ImportError:
|
||||
_A2ACardResolver = object
|
||||
|
||||
# Runtime imports with availability check
|
||||
_A2ACardResolver: Any = object
|
||||
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
|
||||
from a2a.utils.constants import (
|
||||
AGENT_CARD_WELL_KNOWN_PATH,
|
||||
PREV_AGENT_CARD_WELL_KNOWN_PATH,
|
||||
|
|
@ -115,6 +120,7 @@ class LiteLLMA2ACardResolver(_A2ACardResolver):
|
|||
self,
|
||||
relative_card_path: str | None = None,
|
||||
http_kwargs: dict[str, Any] | None = None,
|
||||
signature_verifier: Callable[["AgentCard"], None] | None = None,
|
||||
) -> "AgentCard":
|
||||
"""
|
||||
Fetch the agent card, trying multiple well-known paths.
|
||||
|
|
@ -125,6 +131,8 @@ class LiteLLMA2ACardResolver(_A2ACardResolver):
|
|||
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
|
||||
signature_verifier: Optional callable forwarded to the SDK to verify
|
||||
the resolved card's signature.
|
||||
|
||||
Returns:
|
||||
AgentCard from the A2A agent
|
||||
|
|
@ -137,6 +145,7 @@ class LiteLLMA2ACardResolver(_A2ACardResolver):
|
|||
return await super().get_agent_card(
|
||||
relative_card_path=relative_card_path,
|
||||
http_kwargs=http_kwargs,
|
||||
signature_verifier=signature_verifier,
|
||||
)
|
||||
|
||||
# Try both well-known paths
|
||||
|
|
@ -152,6 +161,7 @@ class LiteLLMA2ACardResolver(_A2ACardResolver):
|
|||
return await super().get_agent_card(
|
||||
relative_card_path=path,
|
||||
http_kwargs=http_kwargs,
|
||||
signature_verifier=signature_verifier,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.debug("Failed to fetch agent card from %s%s: %s", self.base_url, path, e)
|
||||
|
|
|
|||
|
|
@ -76,7 +76,7 @@ async def test_card_resolver_fallback_from_new_to_old_path():
|
|||
|
||||
# Create a mock for the parent's get_agent_card method
|
||||
async def mock_parent_get_agent_card(
|
||||
self, relative_card_path=None, http_kwargs=None
|
||||
self, relative_card_path=None, http_kwargs=None, signature_verifier=None
|
||||
):
|
||||
paths_called.append(relative_card_path)
|
||||
if relative_card_path == "/.well-known/agent-card.json":
|
||||
|
|
@ -111,6 +111,39 @@ async def test_card_resolver_fallback_from_new_to_old_path():
|
|||
assert result.name == "Test Agent"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_agent_card_forwards_signature_verifier():
|
||||
"""
|
||||
The SDK's get_agent_card accepts a signature_verifier. Dropping it from this override
|
||||
made the call raise TypeError for any caller passing it, so it must be forwarded.
|
||||
"""
|
||||
received = {}
|
||||
|
||||
async def mock_parent_get_agent_card(
|
||||
self, relative_card_path=None, http_kwargs=None, signature_verifier=None
|
||||
):
|
||||
received["signature_verifier"] = signature_verifier
|
||||
return MagicMock()
|
||||
|
||||
def verifier(card):
|
||||
return None
|
||||
|
||||
with patch.object(
|
||||
LiteLLMA2ACardResolver.__bases__[0],
|
||||
"get_agent_card",
|
||||
mock_parent_get_agent_card,
|
||||
):
|
||||
resolver = LiteLLMA2ACardResolver(
|
||||
httpx_client=MagicMock(), base_url="http://test-agent:8000"
|
||||
)
|
||||
await resolver.get_agent_card(
|
||||
relative_card_path="/.well-known/agent-card.json",
|
||||
signature_verifier=verifier,
|
||||
)
|
||||
|
||||
assert received["signature_verifier"] is verifier
|
||||
|
||||
|
||||
def test_is_localhost_or_internal_url():
|
||||
"""Test that localhost/internal URLs are correctly detected."""
|
||||
# Should return True for localhost variants
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue