test_a2a_registry_integration

litellm/llms/a2a/chat/transformation.py

@@ -27,6 +27,68 @@ class A2AConfig(BaseConfig):
     Handles transformation between OpenAI and A2A JSON-RPC 2.0 formats.
     """
     
+    @staticmethod
+    def resolve_agent_config_from_registry(
+        model: str,
+        api_base: Optional[str],
+        api_key: Optional[str],
+        headers: Optional[Dict[str, Any]],
+        optional_params: Dict[str, Any],
+    ) -> tuple[Optional[str], Optional[str], Optional[Dict[str, Any]]]:
+        """
+        Resolve agent configuration from registry if model format is "a2a/<agent-name>".
+        
+        Extracts agent name from model string and looks up configuration in the
+        agent registry (if available in proxy context).
+        
+        Args:
+            model: Model string (e.g., "a2a/my-agent")
+            api_base: Explicit api_base (takes precedence over registry)
+            api_key: Explicit api_key (takes precedence over registry)
+            headers: Explicit headers (takes precedence over registry)
+            optional_params: Dict to merge additional litellm_params into
+        
+        Returns:
+            Tuple of (api_base, api_key, headers) with registry values filled in
+        """
+        # Extract agent name from model (e.g., "a2a/my-agent" -> "my-agent")
+        agent_name = model.split("/", 1)[1] if "/" in model else None
+        
+        # Only lookup if agent name exists and some config is missing
+        if not agent_name or (api_base is not None and api_key is not None and headers is not None):
+            return api_base, api_key, headers
+        
+        # Try registry lookup (only available in proxy context)
+        try:
+            from litellm.proxy.agent_endpoints.agent_registry import (
+                global_agent_registry,
+            )
+            
+            agent = global_agent_registry.get_agent_by_name(agent_name)
+            if agent:
+                # Get api_base from agent card URL
+                if api_base is None and agent.agent_card_params:
+                    api_base = agent.agent_card_params.get("url")
+                
+                # Get api_key, headers, and other params from litellm_params
+                if agent.litellm_params:
+                    if api_key is None:
+                        api_key = agent.litellm_params.get("api_key")
+                    
+                    if headers is None:
+                        agent_headers = agent.litellm_params.get("headers")
+                        if agent_headers:
+                            headers = agent_headers
+                    
+                    # Merge other litellm_params (timeout, max_retries, etc.)
+                    for key, value in agent.litellm_params.items():
+                        if key not in ["api_key", "api_base", "headers", "model"] and key not in optional_params:
+                            optional_params[key] = value
+        except ImportError:
+            pass  # Registry not available (not running in proxy context)
+        
+        return api_base, api_key, headers
+    
     def get_supported_openai_params(self, model: str) -> List[str]:
         """Return list of supported OpenAI parameters"""
         return [

litellm/main.py

@@ -2201,14 +2201,24 @@ def completion(  # type: ignore # noqa: PLR0915
             )
         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")
+            # Resolve agent configuration from registry if model format is "a2a/<agent-name>"
+            api_base, api_key, headers = litellm.A2AConfig.resolve_agent_config_from_registry(
+                model=model,
+                api_base=api_base,
+                api_key=api_key,
+                headers=headers,
+                optional_params=optional_params,
             )
-
+            
+            # Fall back to environment variables and defaults
+            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")
+                raise Exception(
+                    "api_base is required for A2A provider. "
+                    "Either provide api_base parameter, set A2A_API_BASE environment variable, "
+                    "or register the agent in the proxy with model='a2a/<agent-name>'."
+                )
 
             headers = headers or litellm.headers
 

tests/test_litellm/test_a2a_registry_lookup.py

@@ -0,0 +1,73 @@
+"""
+Test A2A provider registry lookup functionality.
+
+Maps to: litellm/llms/a2a/chat/transformation.py
+"""
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath("../.."))
+
+import pytest
+
+import litellm
+from litellm.llms.a2a.chat.transformation import A2AConfig
+
+
+def test_resolve_agent_config_from_registry_static_method():
+    """Test the static helper method for registry resolution"""
+    
+    # Test 1: No agent name in model
+    api_base, api_key, headers = A2AConfig.resolve_agent_config_from_registry(
+        model="a2a",
+        api_base="http://test.com",
+        api_key=None,
+        headers=None,
+        optional_params={}
+    )
+    assert api_base == "http://test.com"
+    
+    # Test 2: All params provided - should not lookup registry
+    api_base, api_key, headers = A2AConfig.resolve_agent_config_from_registry(
+        model="a2a/test-agent",
+        api_base="http://explicit.com",
+        api_key="explicit-key",
+        headers={"X-Test": "value"},
+        optional_params={}
+    )
+    assert api_base == "http://explicit.com"
+    assert api_key == "explicit-key"
+
+
+def test_a2a_registry_integration():
+    """Test registry lookup in proxy context"""
+    
+    try:
+        from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
+        from litellm.types.agents import AgentResponse
+
+        # Create test agent
+        test_agent = AgentResponse(
+            agent_id="test-id",
+            agent_name="test-agent",
+            agent_card_params={"url": "http://registry-url.example.com:9999"},
+            litellm_params={"api_key": "registry-key"},
+        )
+        
+        # Register and test
+        original_agents = global_agent_registry.agent_list.copy()
+        global_agent_registry.register_agent(test_agent)
+        
+        try:
+            litellm.completion(
+                model="a2a/test-agent",
+                messages=[{"role": "user", "content": "Hello"}]
+            )
+        except Exception as e:
+            # Should use registry URL (connection error expected)
+            assert "registry-url.example.com" in str(e) or "APIConnectionError" in str(type(e).__name__)
+        finally:
+            global_agent_registry.agent_list = original_agents
+            
+    except ImportError:
+        pytest.skip("Registry not available (not in proxy context)")