[Feat] Use A2A registered agents with /chat/completions - #20362
Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
Greptile OverviewGreptile SummaryThis PR extends A2A agent integration by enabling automatic configuration lookup from the agent registry when using the Changes
The implementation follows existing LiteLLM patterns for provider configuration and maintains backward compatibility with existing A2A usage patterns. Confidence Score: 4/5
|
| Filename | Overview |
|---|---|
| litellm/llms/a2a/chat/transformation.py | Adds registry lookup for A2A agents using model format a2a/<agent-name> |
| litellm/main.py | Integrates registry lookup into completion flow with improved error message |
| tests/test_litellm/test_a2a_registry_lookup.py | Tests registry lookup with unit and integration tests |
Sequence Diagram
sequenceDiagram
participant User
participant completion() as litellm.completion()
participant A2AConfig
participant global_agent_registry
participant Agent Registry
participant A2A Agent
User->>completion(): completion(model="a2a/my-agent", messages=[...])
Note over completion(): custom_llm_provider = "a2a"
completion()->>A2AConfig: resolve_agent_config_from_registry(model, api_base, api_key, headers, optional_params)
A2AConfig->>A2AConfig: Extract agent name from model string
Note over A2AConfig: "a2a/my-agent" → "my-agent"
alt All params provided
A2AConfig-->>completion(): Return provided params (skip registry)
else Some params missing
A2AConfig->>global_agent_registry: get_agent_by_name(agent_name)
global_agent_registry->>Agent Registry: Lookup in agent_list
alt Agent found in registry
Agent Registry-->>global_agent_registry: Return AgentResponse
global_agent_registry-->>A2AConfig: agent object
Note over A2AConfig: Fill missing params:<br/>- api_base from agent_card_params.url<br/>- api_key from litellm_params<br/>- headers from litellm_params<br/>- merge other litellm_params
A2AConfig-->>completion(): Return merged params
else Agent not found
Agent Registry-->>global_agent_registry: None
A2AConfig-->>completion(): Return original params
end
end
completion()->>completion(): Fallback to env vars<br/>(A2A_API_BASE, etc.)
alt api_base is None
completion()->>User: Raise Exception:<br/>"api_base is required"
else api_base exists
completion()->>A2A Agent: Send JSON-RPC request
A2A Agent-->>completion(): Response
completion()-->>User: Return ModelResponse
end
| 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 |
There was a problem hiding this comment.
Short-circuit logic skips registry when all params provided, but partial config (e.g. api_base + api_key but missing headers) won't be enriched from registry
| 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 | |
| if not agent_name or (api_base and api_key and headers): |
Prompt To Fix With AI
This is a comment left during a code review.
Path: litellm/llms/a2a/chat/transformation.py
Line: 58:59
Comment:
Short-circuit logic skips registry when all params provided, but partial config (e.g. `api_base` + `api_key` but missing `headers`) won't be enriched from registry
```suggestion
if not agent_name or (api_base and api_key and headers):
```
How can I resolve this? If you propose a fix, please make it concise.| 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__) |
There was a problem hiding this comment.
Consider mocking the HTTP call to verify registry config is used, rather than relying on error messages containing the URL
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
Prompt To Fix With AI
This is a comment left during a code review.
Path: tests/test_litellm/test_a2a_registry_lookup.py
Line: 62:68
Comment:
Consider mocking the HTTP call to verify registry config is used, rather than relying on error messages containing the URL
<sub>Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!</sub>
How can I resolve this? If you propose a fix, please make it concise.| ) | ||
| 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__) |
Check failure
Code scanning / CodeQL
Incomplete URL substring sanitization
Show autofix suggestion
Hide autofix suggestion
Copilot Autofix
AI 7 months ago
In general, the problem is that the test is asserting the presence of a host substring inside an arbitrary error string, which CodeQL interprets as an unsafe “substring-based URL check.” To avoid this, we should 1) not rely on arbitrary substring checks against the full error string, and 2) instead assert on more structured or clearly safe information. Since this is test code and we must not change external behavior, the best approach is to restructure the assertion so that we either check the exception type (which we already do partially) or, if we still want to confirm the URL, we do it in a clearly non-sanitization context.
The single best way to fix this without changing functionality is to avoid checking "registry-url.example.com" in str(e) and instead assert only on the exception class name indicating a connection error. This keeps the intent of the test (verifying that a network call was attempted and failed) while removing the substring-host check that triggers the CodeQL rule. Concretely, in tests/test_litellm/test_a2a_registry_lookup.py, in test_a2a_registry_integration, modify the except block around line 66–68 so that:
- We no longer check for
"registry-url.example.com" in str(e). - We only assert that the exception type name contains
"APIConnectionError"(or is exactly that type), which is already part of the existing condition.
No new imports or helper methods are needed; we simply change the assertion logic within the shown snippet.
| @@ -64,8 +64,8 @@ | ||
| 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__) | ||
| # Should raise a connection-related error when using the registry URL | ||
| assert "APIConnectionError" in str(type(e).__name__) | ||
| finally: | ||
| global_agent_registry.agent_list = original_agents | ||
|
|
* test_a2a_registry_integration * fix: render agents on model dropdown on UI * init append_agents_to_model_group * route_a2a_agent_request * is_a2a_agent_model * route_a2a_agent_request * fix: error handling * docs A2A usage * docs fix * feat: working A2a streaming * fix transform
[Feat] Use A2A registered agents with /chat/completions
Allows using a2a agents with litellm endpoints /chat/completions, /messages, /responses
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
tests/litellm/directory, Adding at least 1 test is a hard requirement - see detailsmake test-unitCI (LiteLLM team)
Branch creation CI run
Link:
CI run for the last commit
Link:
Merge / cherry-pick CI run
Links:
Type
🆕 New Feature
✅ Test
Changes