mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
init append_agents_to_model_group
This commit is contained in:
Ishaan Jaffer
parent
f0f4922bc3
commit
6e151cbc9b
2 changed files with 206 additions and 0 deletions
96
litellm/proxy/agent_endpoints/model_list_helpers.py
Normal file
96
litellm/proxy/agent_endpoints/model_list_helpers.py
Normal file
|
|
@ -0,0 +1,96 @@
|
|||
"""
|
||||
Helper functions for appending A2A agents to model lists.
|
||||
|
||||
Used by proxy model endpoints to make agents appear in UI alongside models.
|
||||
"""
|
||||
from typing import List
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
|
||||
ModelGroupInfoProxy,
|
||||
)
|
||||
|
||||
|
||||
async def append_agents_to_model_group(
|
||||
model_groups: List[ModelGroupInfoProxy],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> List[ModelGroupInfoProxy]:
|
||||
"""
|
||||
Append A2A agents to model groups list for UI display.
|
||||
|
||||
Converts agents to model format with "a2a/<agent-name>" naming
|
||||
so they appear in playground and work with LiteLLM routing.
|
||||
"""
|
||||
try:
|
||||
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
|
||||
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
|
||||
AgentRequestHandler,
|
||||
)
|
||||
|
||||
allowed_agent_ids = await AgentRequestHandler.get_allowed_agents(
|
||||
user_api_key_auth=user_api_key_dict
|
||||
)
|
||||
|
||||
for agent_id in allowed_agent_ids:
|
||||
agent = global_agent_registry.get_agent_by_id(agent_id)
|
||||
if agent is not None:
|
||||
model_groups.append(
|
||||
ModelGroupInfoProxy(
|
||||
model_group=f"a2a/{agent.agent_name}",
|
||||
mode="chat",
|
||||
providers=["a2a"],
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug(
|
||||
f"Error appending agents to model_group/info: {e}"
|
||||
)
|
||||
|
||||
return model_groups
|
||||
|
||||
|
||||
async def append_agents_to_model_info(
|
||||
models: List[dict],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> List[dict]:
|
||||
"""
|
||||
Append A2A agents to model info list for UI display.
|
||||
|
||||
Converts agents to model format with "a2a/<agent-name>" naming
|
||||
so they appear in models page and work with LiteLLM routing.
|
||||
"""
|
||||
try:
|
||||
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
|
||||
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
|
||||
AgentRequestHandler,
|
||||
)
|
||||
|
||||
allowed_agent_ids = await AgentRequestHandler.get_allowed_agents(
|
||||
user_api_key_auth=user_api_key_dict
|
||||
)
|
||||
|
||||
for agent_id in allowed_agent_ids:
|
||||
agent = global_agent_registry.get_agent_by_id(agent_id)
|
||||
if agent is not None:
|
||||
models.append({
|
||||
"model_name": f"a2a/{agent.agent_name}",
|
||||
"litellm_params": {
|
||||
"model": f"a2a/{agent.agent_name}",
|
||||
"custom_llm_provider": "a2a",
|
||||
},
|
||||
"model_info": {
|
||||
"id": agent.agent_id,
|
||||
"mode": "chat",
|
||||
"db_model": True,
|
||||
"created_by": agent.created_by,
|
||||
"created_at": agent.created_at,
|
||||
"updated_at": agent.updated_at,
|
||||
},
|
||||
})
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug(
|
||||
f"Error appending agents to v2/model/info: {e}"
|
||||
)
|
||||
|
||||
return models
|
||||
|
|
@ -0,0 +1,110 @@
|
|||
"""
|
||||
Test appending A2A agents to model lists.
|
||||
|
||||
Maps to: litellm/proxy/agent_endpoints/model_list_helpers.py
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../../.."))
|
||||
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.agent_endpoints.model_list_helpers import (
|
||||
append_agents_to_model_group,
|
||||
append_agents_to_model_info,
|
||||
)
|
||||
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
|
||||
from litellm.types.agents import AgentResponse
|
||||
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
|
||||
ModelGroupInfoProxy,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_append_agents_to_model_group():
|
||||
"""Test agents are converted to model group format with a2a/ prefix"""
|
||||
|
||||
# Mock agent data
|
||||
mock_agent = AgentResponse(
|
||||
agent_id="test-agent-id",
|
||||
agent_name="my-agent",
|
||||
agent_card_params={"url": "http://example.com"},
|
||||
litellm_params=None,
|
||||
)
|
||||
|
||||
# Mock AgentRequestHandler at its source location
|
||||
mock_get_allowed_agents = AsyncMock(return_value=["test-agent-id"])
|
||||
|
||||
# Mock global_agent_registry
|
||||
mock_registry = Mock()
|
||||
mock_registry.get_agent_by_id = Mock(return_value=mock_agent)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents",
|
||||
mock_get_allowed_agents,
|
||||
):
|
||||
with patch(
|
||||
"litellm.proxy.agent_endpoints.agent_registry.global_agent_registry",
|
||||
mock_registry,
|
||||
):
|
||||
model_groups = []
|
||||
user_api_key_dict = Mock(spec=UserAPIKeyAuth)
|
||||
|
||||
result = await append_agents_to_model_group(
|
||||
model_groups=model_groups,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
# Verify agent was converted with a2a/ prefix
|
||||
assert len(result) == 1
|
||||
assert result[0].model_group == "a2a/my-agent"
|
||||
assert result[0].mode == "chat"
|
||||
assert result[0].providers == ["a2a"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_append_agents_to_model_info():
|
||||
"""Test agents are converted to model info format with a2a/ prefix"""
|
||||
|
||||
# Mock agent data
|
||||
mock_agent = AgentResponse(
|
||||
agent_id="agent-123",
|
||||
agent_name="test-agent",
|
||||
agent_card_params={"url": "http://example.com"},
|
||||
litellm_params=None,
|
||||
created_by="user-123",
|
||||
)
|
||||
|
||||
# Mock AgentRequestHandler at its source location
|
||||
mock_get_allowed_agents = AsyncMock(return_value=["agent-123"])
|
||||
|
||||
# Mock global_agent_registry
|
||||
mock_registry = Mock()
|
||||
mock_registry.get_agent_by_id = Mock(return_value=mock_agent)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents",
|
||||
mock_get_allowed_agents,
|
||||
):
|
||||
with patch(
|
||||
"litellm.proxy.agent_endpoints.agent_registry.global_agent_registry",
|
||||
mock_registry,
|
||||
):
|
||||
models = []
|
||||
user_api_key_dict = Mock(spec=UserAPIKeyAuth)
|
||||
|
||||
result = await append_agents_to_model_info(
|
||||
models=models,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
# Verify agent was converted with a2a/ prefix
|
||||
assert len(result) == 1
|
||||
assert result[0]["model_name"] == "a2a/test-agent"
|
||||
assert result[0]["litellm_params"]["model"] == "a2a/test-agent"
|
||||
assert result[0]["litellm_params"]["custom_llm_provider"] == "a2a"
|
||||
assert result[0]["model_info"]["id"] == "agent-123"
|
||||
assert result[0]["model_info"]["mode"] == "chat"
|
||||
Loading…
Add table
Reference in a new issue