[Feat] A2A Gateway - allow adding Azure Foundry Agents on UI (#17909)
* add CostConfigFields * add CostConfigFields * add output_cost_per_token * refactor table * add agent cost view * add azure foundry fields * add foundry logo * fix: clean error * fix utils * fix agent edi * add easter egg * fix order * test_handle_streaming_forwards_api_key * fix forward api key down * fix a2a send msg * add A2a comparison on compare playground * fix chat ui * fix bedrock agentcore stream
This commit is contained in:
Ishaan Jaff
GitHub
parent
bede40a90d
commit
3054b6ea60
@ -55,6 +55,7 @@ class A2ACompletionBridgeHandler:
|
||||
# Get completion params
|
||||
custom_llm_provider = litellm_params.get("custom_llm_provider")
|
||||
model = litellm_params.get("model", "agent")
|
||||
api_key = litellm_params.get("api_key")
|
||||
|
||||
# Build full model string if provider specified
|
||||
# Skip prepending if model already starts with the provider prefix
|
||||
@ -72,6 +73,7 @@ class A2ACompletionBridgeHandler:
|
||||
model=full_model,
|
||||
messages=openai_messages,
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
stream=False,
|
||||
)
|
||||
|
||||
@ -127,6 +129,7 @@ class A2ACompletionBridgeHandler:
|
||||
# Get completion params
|
||||
custom_llm_provider = litellm_params.get("custom_llm_provider")
|
||||
model = litellm_params.get("model", "agent")
|
||||
api_key = litellm_params.get("api_key")
|
||||
|
||||
# Build full model string if provider specified
|
||||
# Skip prepending if model already starts with the provider prefix
|
||||
@ -157,6 +160,7 @@ class A2ACompletionBridgeHandler:
|
||||
model=full_model,
|
||||
messages=openai_messages,
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
|
||||
@ -5,7 +5,7 @@ Handles Server-Sent Events (SSE) streaming responses from AgentCore.
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import TYPE_CHECKING
|
||||
from typing import TYPE_CHECKING, Any, Optional
|
||||
|
||||
import httpx
|
||||
|
||||
@ -19,262 +19,234 @@ if TYPE_CHECKING:
|
||||
|
||||
|
||||
class AgentCoreSSEStreamIterator:
|
||||
"""Iterator for AgentCore SSE streaming responses. Supports both sync and async iteration."""
|
||||
"""
|
||||
Iterator for AgentCore SSE streaming responses.
|
||||
Supports both sync and async iteration.
|
||||
|
||||
CRITICAL: The line iterators are created lazily on first access and reused.
|
||||
We must NOT create new iterators in __aiter__/__iter__ because
|
||||
CustomStreamWrapper calls __aiter__ on every call to its __anext__,
|
||||
which would create new iterators and cause StreamConsumed errors.
|
||||
"""
|
||||
|
||||
def __init__(self, response: httpx.Response, model: str):
|
||||
self.response = response
|
||||
self.model = model
|
||||
self.finished = False
|
||||
self.line_iterator = None
|
||||
self.async_line_iterator = None
|
||||
self._sync_iter: Any = None
|
||||
self._async_iter: Any = None
|
||||
self._sync_iter_initialized = False
|
||||
self._async_iter_initialized = False
|
||||
|
||||
def __iter__(self):
|
||||
"""Initialize sync iteration."""
|
||||
self.line_iterator = self.response.iter_lines()
|
||||
"""Initialize sync iteration - create iterator lazily on first call only."""
|
||||
if not self._sync_iter_initialized:
|
||||
self._sync_iter = iter(self.response.iter_lines())
|
||||
self._sync_iter_initialized = True
|
||||
return self
|
||||
|
||||
def __aiter__(self):
|
||||
"""Initialize async iteration."""
|
||||
self.async_line_iterator = self.response.aiter_lines()
|
||||
"""Initialize async iteration - create iterator lazily on first call only."""
|
||||
if not self._async_iter_initialized:
|
||||
self._async_iter = self.response.aiter_lines().__aiter__()
|
||||
self._async_iter_initialized = True
|
||||
return self
|
||||
|
||||
def __next__(self) -> ModelResponse:
|
||||
"""Sync iteration - parse SSE events and yield ModelResponse chunks."""
|
||||
def _parse_sse_line(self, line: str) -> Optional[ModelResponse]:
|
||||
"""
|
||||
Parse a single SSE line and return a ModelResponse chunk if applicable.
|
||||
|
||||
AgentCore SSE format:
|
||||
- data: {"event": {"contentBlockDelta": {"delta": {"text": "..."}}}}
|
||||
- data: {"event": {"metadata": {"usage": {...}}}}
|
||||
- data: {"message": {...}}
|
||||
"""
|
||||
line = line.strip()
|
||||
if not line or not line.startswith("data:"):
|
||||
return None
|
||||
|
||||
json_str = line[5:].strip()
|
||||
if not json_str:
|
||||
return None
|
||||
|
||||
try:
|
||||
if self.line_iterator is None:
|
||||
data = json.loads(json_str)
|
||||
|
||||
# Skip non-dict data (some lines contain Python repr strings)
|
||||
if not isinstance(data, dict):
|
||||
return None
|
||||
|
||||
# Process content delta events
|
||||
if "event" in data and isinstance(data["event"], dict):
|
||||
event_payload = data["event"]
|
||||
content_block_delta = event_payload.get("contentBlockDelta")
|
||||
|
||||
if content_block_delta:
|
||||
delta = content_block_delta.get("delta", {})
|
||||
text = delta.get("text", "")
|
||||
|
||||
if text:
|
||||
# Return chunk with text
|
||||
chunk = ModelResponse(
|
||||
id=f"chatcmpl-{uuid.uuid4()}",
|
||||
created=0,
|
||||
model=self.model,
|
||||
object="chat.completion.chunk",
|
||||
)
|
||||
|
||||
chunk.choices = [
|
||||
StreamingChoices(
|
||||
finish_reason=None,
|
||||
index=0,
|
||||
delta=Delta(content=text, role="assistant"),
|
||||
)
|
||||
]
|
||||
|
||||
return chunk
|
||||
|
||||
# Check for metadata/usage - this signals the end
|
||||
metadata = event_payload.get("metadata")
|
||||
if metadata and "usage" in metadata:
|
||||
chunk = ModelResponse(
|
||||
id=f"chatcmpl-{uuid.uuid4()}",
|
||||
created=0,
|
||||
model=self.model,
|
||||
object="chat.completion.chunk",
|
||||
)
|
||||
|
||||
chunk.choices = [
|
||||
StreamingChoices(
|
||||
finish_reason="stop",
|
||||
index=0,
|
||||
delta=Delta(),
|
||||
)
|
||||
]
|
||||
|
||||
usage_data: AgentCoreUsage = metadata["usage"] # type: ignore
|
||||
setattr(
|
||||
chunk,
|
||||
"usage",
|
||||
Usage(
|
||||
prompt_tokens=usage_data.get("inputTokens", 0),
|
||||
completion_tokens=usage_data.get("outputTokens", 0),
|
||||
total_tokens=usage_data.get("totalTokens", 0),
|
||||
),
|
||||
)
|
||||
|
||||
self.finished = True
|
||||
return chunk
|
||||
|
||||
# Check for final message (alternative finish signal)
|
||||
if "message" in data and isinstance(data["message"], dict):
|
||||
if not self.finished:
|
||||
chunk = ModelResponse(
|
||||
id=f"chatcmpl-{uuid.uuid4()}",
|
||||
created=0,
|
||||
model=self.model,
|
||||
object="chat.completion.chunk",
|
||||
)
|
||||
|
||||
chunk.choices = [
|
||||
StreamingChoices(
|
||||
finish_reason="stop",
|
||||
index=0,
|
||||
delta=Delta(),
|
||||
)
|
||||
]
|
||||
|
||||
self.finished = True
|
||||
return chunk
|
||||
|
||||
except json.JSONDecodeError:
|
||||
verbose_logger.debug(f"Skipping non-JSON SSE line: {line[:100]}")
|
||||
|
||||
return None
|
||||
|
||||
def _create_final_chunk(self) -> ModelResponse:
|
||||
"""Create a final chunk to signal stream completion."""
|
||||
chunk = ModelResponse(
|
||||
id=f"chatcmpl-{uuid.uuid4()}",
|
||||
created=0,
|
||||
model=self.model,
|
||||
object="chat.completion.chunk",
|
||||
)
|
||||
|
||||
chunk.choices = [
|
||||
StreamingChoices(
|
||||
finish_reason="stop",
|
||||
index=0,
|
||||
delta=Delta(),
|
||||
)
|
||||
]
|
||||
|
||||
return chunk
|
||||
|
||||
def __next__(self) -> ModelResponse:
|
||||
"""
|
||||
Sync iteration - parse SSE events and yield ModelResponse chunks.
|
||||
|
||||
Uses next() on the stored iterator to properly resume between calls.
|
||||
"""
|
||||
try:
|
||||
if self._sync_iter is None:
|
||||
raise StopIteration
|
||||
for line in self.line_iterator:
|
||||
line = line.strip()
|
||||
|
||||
if not line or not line.startswith('data:'):
|
||||
continue
|
||||
|
||||
# Extract JSON from SSE line
|
||||
json_str = line[5:].strip()
|
||||
if not json_str:
|
||||
continue
|
||||
|
||||
|
||||
# Keep getting lines until we have a result to return
|
||||
while True:
|
||||
try:
|
||||
data = json.loads(json_str)
|
||||
|
||||
# Skip non-dict data
|
||||
if not isinstance(data, dict):
|
||||
continue
|
||||
|
||||
# Process content delta events
|
||||
if "event" in data and isinstance(data["event"], dict):
|
||||
event_payload = data["event"]
|
||||
content_block_delta = event_payload.get("contentBlockDelta")
|
||||
|
||||
if content_block_delta:
|
||||
delta = content_block_delta.get("delta", {})
|
||||
text = delta.get("text", "")
|
||||
|
||||
if text:
|
||||
# Yield chunk with text
|
||||
chunk = ModelResponse(
|
||||
id=f"chatcmpl-{uuid.uuid4()}",
|
||||
created=0,
|
||||
model=self.model,
|
||||
object="chat.completion.chunk",
|
||||
)
|
||||
|
||||
chunk.choices = [
|
||||
StreamingChoices(
|
||||
finish_reason=None,
|
||||
index=0,
|
||||
delta=Delta(content=text, role="assistant"),
|
||||
)
|
||||
]
|
||||
|
||||
return chunk
|
||||
|
||||
# Check for metadata/usage
|
||||
metadata = event_payload.get("metadata")
|
||||
if metadata and "usage" in metadata:
|
||||
# This is the final chunk with usage
|
||||
chunk = ModelResponse(
|
||||
id=f"chatcmpl-{uuid.uuid4()}",
|
||||
created=0,
|
||||
model=self.model,
|
||||
object="chat.completion.chunk",
|
||||
)
|
||||
|
||||
chunk.choices = [
|
||||
StreamingChoices(
|
||||
finish_reason="stop",
|
||||
index=0,
|
||||
delta=Delta(),
|
||||
)
|
||||
]
|
||||
|
||||
usage_data: AgentCoreUsage = metadata["usage"] # type: ignore
|
||||
setattr(chunk, "usage", Usage(
|
||||
prompt_tokens=usage_data.get("inputTokens", 0),
|
||||
completion_tokens=usage_data.get("outputTokens", 0),
|
||||
total_tokens=usage_data.get("totalTokens", 0),
|
||||
))
|
||||
|
||||
self.finished = True
|
||||
return chunk
|
||||
|
||||
# Check for final message (alternative finish signal)
|
||||
if "message" in data and isinstance(data["message"], dict):
|
||||
if not self.finished:
|
||||
chunk = ModelResponse(
|
||||
id=f"chatcmpl-{uuid.uuid4()}",
|
||||
created=0,
|
||||
model=self.model,
|
||||
object="chat.completion.chunk",
|
||||
)
|
||||
|
||||
chunk.choices = [
|
||||
StreamingChoices(
|
||||
finish_reason="stop",
|
||||
index=0,
|
||||
delta=Delta(),
|
||||
)
|
||||
]
|
||||
|
||||
self.finished = True
|
||||
return chunk
|
||||
|
||||
except json.JSONDecodeError:
|
||||
verbose_logger.debug(f"Skipping non-JSON SSE line: {line[:100]}")
|
||||
continue
|
||||
|
||||
# Stream ended naturally
|
||||
raise StopIteration
|
||||
line = next(self._sync_iter)
|
||||
except StopIteration:
|
||||
# Stream ended - send final chunk if not already finished
|
||||
if not self.finished:
|
||||
self.finished = True
|
||||
return self._create_final_chunk()
|
||||
raise
|
||||
|
||||
result = self._parse_sse_line(line)
|
||||
if result is not None:
|
||||
return result
|
||||
|
||||
except StopIteration:
|
||||
raise
|
||||
except httpx.StreamConsumed:
|
||||
# This is expected when the stream has been fully consumed
|
||||
raise StopIteration
|
||||
except httpx.StreamClosed:
|
||||
# This is expected when the stream is closed
|
||||
raise StopIteration
|
||||
except Exception as e:
|
||||
verbose_logger.error(f"Error in AgentCore SSE stream: {str(e)}")
|
||||
raise StopIteration
|
||||
|
||||
async def __anext__(self) -> ModelResponse:
|
||||
"""Async iteration - parse SSE events and yield ModelResponse chunks."""
|
||||
"""
|
||||
Async iteration - parse SSE events and yield ModelResponse chunks.
|
||||
|
||||
Uses __anext__() on the stored iterator to properly resume between calls.
|
||||
"""
|
||||
try:
|
||||
if self.async_line_iterator is None:
|
||||
if self._async_iter is None:
|
||||
raise StopAsyncIteration
|
||||
async for line in self.async_line_iterator:
|
||||
line = line.strip()
|
||||
|
||||
if not line or not line.startswith('data:'):
|
||||
continue
|
||||
|
||||
# Extract JSON from SSE line
|
||||
json_str = line[5:].strip()
|
||||
if not json_str:
|
||||
continue
|
||||
|
||||
|
||||
# Keep getting lines until we have a result to return
|
||||
while True:
|
||||
try:
|
||||
data = json.loads(json_str)
|
||||
|
||||
# Skip non-dict data
|
||||
if not isinstance(data, dict):
|
||||
continue
|
||||
|
||||
# Process content delta events
|
||||
if "event" in data and isinstance(data["event"], dict):
|
||||
event_payload = data["event"]
|
||||
content_block_delta = event_payload.get("contentBlockDelta")
|
||||
|
||||
if content_block_delta:
|
||||
delta = content_block_delta.get("delta", {})
|
||||
text = delta.get("text", "")
|
||||
|
||||
if text:
|
||||
# Yield chunk with text
|
||||
chunk = ModelResponse(
|
||||
id=f"chatcmpl-{uuid.uuid4()}",
|
||||
created=0,
|
||||
model=self.model,
|
||||
object="chat.completion.chunk",
|
||||
)
|
||||
|
||||
chunk.choices = [
|
||||
StreamingChoices(
|
||||
finish_reason=None,
|
||||
index=0,
|
||||
delta=Delta(content=text, role="assistant"),
|
||||
)
|
||||
]
|
||||
|
||||
return chunk
|
||||
|
||||
# Check for metadata/usage
|
||||
metadata = event_payload.get("metadata")
|
||||
if metadata and "usage" in metadata:
|
||||
# This is the final chunk with usage
|
||||
chunk = ModelResponse(
|
||||
id=f"chatcmpl-{uuid.uuid4()}",
|
||||
created=0,
|
||||
model=self.model,
|
||||
object="chat.completion.chunk",
|
||||
)
|
||||
|
||||
chunk.choices = [
|
||||
StreamingChoices(
|
||||
finish_reason="stop",
|
||||
index=0,
|
||||
delta=Delta(),
|
||||
)
|
||||
]
|
||||
|
||||
usage_data: AgentCoreUsage = metadata["usage"] # type: ignore
|
||||
setattr(chunk, "usage", Usage(
|
||||
prompt_tokens=usage_data.get("inputTokens", 0),
|
||||
completion_tokens=usage_data.get("outputTokens", 0),
|
||||
total_tokens=usage_data.get("totalTokens", 0),
|
||||
))
|
||||
|
||||
self.finished = True
|
||||
return chunk
|
||||
|
||||
# Check for final message (alternative finish signal)
|
||||
if "message" in data and isinstance(data["message"], dict):
|
||||
if not self.finished:
|
||||
chunk = ModelResponse(
|
||||
id=f"chatcmpl-{uuid.uuid4()}",
|
||||
created=0,
|
||||
model=self.model,
|
||||
object="chat.completion.chunk",
|
||||
)
|
||||
|
||||
chunk.choices = [
|
||||
StreamingChoices(
|
||||
finish_reason="stop",
|
||||
index=0,
|
||||
delta=Delta(),
|
||||
)
|
||||
]
|
||||
|
||||
self.finished = True
|
||||
return chunk
|
||||
|
||||
except json.JSONDecodeError:
|
||||
verbose_logger.debug(f"Skipping non-JSON SSE line: {line[:100]}")
|
||||
continue
|
||||
|
||||
# Stream ended naturally
|
||||
raise StopAsyncIteration
|
||||
line = await self._async_iter.__anext__()
|
||||
except StopAsyncIteration:
|
||||
# Stream ended - send final chunk if not already finished
|
||||
if not self.finished:
|
||||
self.finished = True
|
||||
return self._create_final_chunk()
|
||||
raise
|
||||
|
||||
result = self._parse_sse_line(line)
|
||||
if result is not None:
|
||||
return result
|
||||
|
||||
except StopAsyncIteration:
|
||||
raise
|
||||
except httpx.StreamConsumed:
|
||||
# This is expected when the stream has been fully consumed
|
||||
raise StopAsyncIteration
|
||||
except httpx.StreamClosed:
|
||||
# This is expected when the stream is closed
|
||||
raise StopAsyncIteration
|
||||
except Exception as e:
|
||||
verbose_logger.error(f"Error in AgentCore SSE stream: {str(e)}")
|
||||
raise StopAsyncIteration
|
||||
|
||||
|
||||
@ -71,6 +71,49 @@
|
||||
"litellm_params_template": {
|
||||
"custom_llm_provider": "bedrock"
|
||||
}
|
||||
},
|
||||
{
|
||||
"agent_type": "azure_ai_foundry",
|
||||
"agent_type_display_name": "Azure AI Foundry",
|
||||
"description": "Connect to Microsoft Azure AI Foundry agents",
|
||||
"logo_url": "/assets/logos/azure_ai_foundry.png",
|
||||
"inherit_credentials_from_provider": "Azure AI",
|
||||
"model_template": "azure_ai/agents/{agent_id}",
|
||||
"credential_fields": [
|
||||
{
|
||||
"key": "agent_id",
|
||||
"label": "Agent ID",
|
||||
"placeholder": "asst_abc123",
|
||||
"tooltip": "The agent/assistant ID from your Azure AI Foundry project (e.g., asst_abc123)",
|
||||
"required": true,
|
||||
"field_type": "text",
|
||||
"default_value": null,
|
||||
"include_in_litellm_params": false
|
||||
},
|
||||
{
|
||||
"key": "api_base",
|
||||
"label": "Azure AI API Base",
|
||||
"placeholder": "https://your-project.services.ai.azure.com",
|
||||
"tooltip": "The base URL for your Azure AI Foundry project endpoint",
|
||||
"required": true,
|
||||
"field_type": "text",
|
||||
"default_value": null,
|
||||
"include_in_litellm_params": true
|
||||
},
|
||||
{
|
||||
"key": "api_key",
|
||||
"label": "Azure AI API Key",
|
||||
"placeholder": null,
|
||||
"tooltip": "API key for authenticating with your Azure AI Foundry project",
|
||||
"required": true,
|
||||
"field_type": "password",
|
||||
"default_value": null,
|
||||
"include_in_litellm_params": true
|
||||
}
|
||||
],
|
||||
"litellm_params_template": {
|
||||
"custom_llm_provider": "azure_ai"
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@ -41,16 +41,16 @@ def test_bedrock_agentcore_basic(model):
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model", [
|
||||
"bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/non_stream_agent-mdfwS2DlAu", # non-streaming invocation
|
||||
"bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC", # streaming invocation
|
||||
"bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_13sf6-cALnp38iZD", # streaming invocation
|
||||
]
|
||||
)
|
||||
async def test_bedrock_agentcore_with_streaming(model):
|
||||
"""
|
||||
Test AgentCore with streaming
|
||||
"""
|
||||
print("running streming test for model=", model)
|
||||
#litellm._turn_on_debug()
|
||||
response = litellm.completion(
|
||||
response = await litellm.acompletion(
|
||||
model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC",
|
||||
messages=[
|
||||
{
|
||||
@ -61,7 +61,7 @@ async def test_bedrock_agentcore_with_streaming(model):
|
||||
stream=True,
|
||||
)
|
||||
|
||||
for chunk in response:
|
||||
async for chunk in response:
|
||||
print("chunk=", chunk)
|
||||
|
||||
|
||||
|
||||
@ -157,3 +157,93 @@ async def test_handle_streaming_emits_proper_events():
|
||||
assert events[3]["result"]["status"]["state"] == "completed"
|
||||
assert events[3]["result"]["final"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_streaming_forwards_api_key():
|
||||
"""Test that handle_streaming forwards api_key from litellm_params to acompletion."""
|
||||
from litellm.a2a_protocol.litellm_completion_bridge.handler import (
|
||||
A2ACompletionBridgeHandler,
|
||||
)
|
||||
|
||||
mock_chunk = MagicMock()
|
||||
mock_chunk.choices = [MagicMock()]
|
||||
mock_chunk.choices[0].delta = MagicMock()
|
||||
mock_chunk.choices[0].delta.content = "Response"
|
||||
|
||||
async def mock_streaming_response():
|
||||
yield mock_chunk
|
||||
|
||||
with patch("litellm.acompletion", new_callable=AsyncMock) as mock_acompletion:
|
||||
mock_acompletion.return_value = mock_streaming_response()
|
||||
|
||||
params = {
|
||||
"message": {
|
||||
"role": "user",
|
||||
"parts": [{"kind": "text", "text": "Hi"}],
|
||||
"messageId": "msg-123",
|
||||
}
|
||||
}
|
||||
|
||||
events = []
|
||||
async for event in A2ACompletionBridgeHandler.handle_streaming(
|
||||
request_id="req-456",
|
||||
params=params,
|
||||
litellm_params={
|
||||
"custom_llm_provider": "azure_ai",
|
||||
"model": "agents/asst_123",
|
||||
"api_key": "test-api-key-12345",
|
||||
},
|
||||
api_base="https://example.azure.com/",
|
||||
):
|
||||
events.append(event)
|
||||
|
||||
# Verify acompletion was called with api_key
|
||||
mock_acompletion.assert_called_once()
|
||||
call_kwargs = mock_acompletion.call_args.kwargs
|
||||
assert call_kwargs["api_key"] == "test-api-key-12345"
|
||||
assert call_kwargs["api_base"] == "https://example.azure.com/"
|
||||
assert call_kwargs["model"] == "azure_ai/agents/asst_123"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_non_streaming_forwards_api_key():
|
||||
"""Test that handle_non_streaming forwards api_key from litellm_params to acompletion."""
|
||||
from litellm.a2a_protocol.litellm_completion_bridge.handler import (
|
||||
A2ACompletionBridgeHandler,
|
||||
)
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.choices = [MagicMock()]
|
||||
mock_response.choices[0].message = MagicMock()
|
||||
mock_response.choices[0].message.content = "Hello!"
|
||||
mock_response.id = "resp-123"
|
||||
|
||||
with patch("litellm.acompletion", new_callable=AsyncMock) as mock_acompletion:
|
||||
mock_acompletion.return_value = mock_response
|
||||
|
||||
params = {
|
||||
"message": {
|
||||
"role": "user",
|
||||
"parts": [{"kind": "text", "text": "Hi"}],
|
||||
"messageId": "msg-123",
|
||||
}
|
||||
}
|
||||
|
||||
await A2ACompletionBridgeHandler.handle_non_streaming(
|
||||
request_id="req-456",
|
||||
params=params,
|
||||
litellm_params={
|
||||
"custom_llm_provider": "azure_ai",
|
||||
"model": "agents/asst_456",
|
||||
"api_key": "my-secret-api-key",
|
||||
},
|
||||
api_base="https://my-azure.com/",
|
||||
)
|
||||
|
||||
# Verify acompletion was called with api_key
|
||||
mock_acompletion.assert_called_once()
|
||||
call_kwargs = mock_acompletion.call_args.kwargs
|
||||
assert call_kwargs["api_key"] == "my-secret-api-key"
|
||||
assert call_kwargs["api_base"] == "https://my-azure.com/"
|
||||
assert call_kwargs["model"] == "azure_ai/agents/asst_456"
|
||||
|
||||
|
||||
BIN
ui/litellm-dashboard/public/assets/logos/azure_ai_foundry.png
Normal file
BIN
ui/litellm-dashboard/public/assets/logos/azure_ai_foundry.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 26 KiB |
@ -91,14 +91,14 @@ const AgentsPanel: React.FC<AgentsPanelProps> = ({ accessToken, userRole }) => {
|
||||
|
||||
return (
|
||||
<div className="w-full mx-auto flex-auto overflow-y-auto m-8 p-2">
|
||||
<div className="flex justify-between items-center mb-4">
|
||||
<div className="flex-col gap-2">
|
||||
<h1 className="text-2xl font-bold">Agents</h1>
|
||||
<p className="text-sm text-gray-600">List of A2A-spec agents that are available to be used in your organization. Go to AI Hub, to make agents public.</p>
|
||||
<div className="flex flex-col gap-2 mb-4">
|
||||
<h1 className="text-2xl font-bold">Agents</h1>
|
||||
<p className="text-sm text-gray-600">List of A2A-spec agents that are available to be used in your organization. Go to AI Hub, to make agents public.</p>
|
||||
<div className="mt-2">
|
||||
<Button onClick={handleAddAgent} disabled={!accessToken}>
|
||||
+ Add New Agent
|
||||
</Button>
|
||||
</div>
|
||||
<Button onClick={handleAddAgent} disabled={!accessToken}>
|
||||
+ Add New Agent
|
||||
</Button>
|
||||
</div>
|
||||
|
||||
{selectedAgentId ? (
|
||||
|
||||
@ -2,11 +2,13 @@ import React, { useState, useEffect } from "react";
|
||||
import { Card, Title, Text, Button as TremorButton, Tab, TabGroup, TabList, TabPanel, TabPanels} from "@tremor/react";
|
||||
import { Form, Input, Button as AntButton, message, Spin, Descriptions } from "antd";
|
||||
import { ArrowLeftIcon } from "@heroicons/react/outline";
|
||||
import { getAgentInfo, patchAgentCall } from "../networking";
|
||||
import { getAgentInfo, patchAgentCall, getAgentCreateMetadata, AgentCreateInfo } from "../networking";
|
||||
import { Agent } from "./types";
|
||||
import AgentFormFields from "./agent_form_fields";
|
||||
import DynamicAgentFormFields, { buildDynamicAgentData } from "./dynamic_agent_form_fields";
|
||||
import { buildAgentDataFromForm, parseAgentForForm } from "./agent_config";
|
||||
import AgentCostView from "./agent_cost_view";
|
||||
import { detectAgentType, parseDynamicAgentForForm } from "./agent_type_utils";
|
||||
|
||||
interface AgentInfoViewProps {
|
||||
agentId: string;
|
||||
@ -26,6 +28,20 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({
|
||||
const [isEditing, setIsEditing] = useState(false);
|
||||
const [isSaving, setIsSaving] = useState(false);
|
||||
const [form] = Form.useForm();
|
||||
const [agentTypeMetadata, setAgentTypeMetadata] = useState<AgentCreateInfo[]>([]);
|
||||
const [detectedAgentType, setDetectedAgentType] = useState<string>("a2a");
|
||||
|
||||
useEffect(() => {
|
||||
const fetchMetadata = async () => {
|
||||
try {
|
||||
const metadata = await getAgentCreateMetadata();
|
||||
setAgentTypeMetadata(metadata);
|
||||
} catch (error) {
|
||||
console.error("Error fetching agent metadata:", error);
|
||||
}
|
||||
};
|
||||
fetchMetadata();
|
||||
}, []);
|
||||
|
||||
useEffect(() => {
|
||||
fetchAgentInfo();
|
||||
@ -38,7 +54,22 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({
|
||||
try {
|
||||
const data = await getAgentInfo(accessToken, agentId);
|
||||
setAgent(data);
|
||||
|
||||
// Detect agent type
|
||||
const agentType = detectAgentType(data);
|
||||
setDetectedAgentType(agentType);
|
||||
|
||||
// Parse form values based on agent type
|
||||
if (agentType === "a2a") {
|
||||
form.setFieldsValue(parseAgentForForm(data));
|
||||
} else {
|
||||
const typeInfo = agentTypeMetadata.find(t => t.agent_type === agentType);
|
||||
if (typeInfo) {
|
||||
form.setFieldsValue(parseDynamicAgentForForm(data, typeInfo));
|
||||
} else {
|
||||
form.setFieldsValue(parseAgentForForm(data));
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
console.error("Error fetching agent info:", error);
|
||||
message.error("Failed to load agent information");
|
||||
@ -47,12 +78,38 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({
|
||||
}
|
||||
};
|
||||
|
||||
// Re-parse form when metadata is loaded
|
||||
useEffect(() => {
|
||||
if (agent && agentTypeMetadata.length > 0) {
|
||||
const agentType = detectAgentType(agent);
|
||||
if (agentType !== "a2a") {
|
||||
const typeInfo = agentTypeMetadata.find(t => t.agent_type === agentType);
|
||||
if (typeInfo) {
|
||||
form.setFieldsValue(parseDynamicAgentForForm(agent, typeInfo));
|
||||
}
|
||||
}
|
||||
}
|
||||
}, [agentTypeMetadata, agent]);
|
||||
|
||||
const selectedAgentTypeInfo = agentTypeMetadata.find(t => t.agent_type === detectedAgentType);
|
||||
|
||||
const handleUpdate = async (values: any) => {
|
||||
if (!accessToken || !agent) return;
|
||||
|
||||
setIsSaving(true);
|
||||
try {
|
||||
const updateData = buildAgentDataFromForm(values, agent);
|
||||
let updateData: any;
|
||||
|
||||
if (detectedAgentType === "a2a") {
|
||||
updateData = buildAgentDataFromForm(values, agent);
|
||||
} else if (selectedAgentTypeInfo) {
|
||||
updateData = buildDynamicAgentData(values, selectedAgentTypeInfo);
|
||||
// Preserve the agent_name from form
|
||||
updateData.agent_name = values.agent_name;
|
||||
} else {
|
||||
updateData = buildAgentDataFromForm(values, agent);
|
||||
}
|
||||
|
||||
await patchAgentCall(accessToken, agentId, updateData);
|
||||
message.success("Agent updated successfully");
|
||||
setIsEditing(false);
|
||||
@ -192,7 +249,13 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({
|
||||
<Input value={agent.agent_id} disabled />
|
||||
</Form.Item>
|
||||
|
||||
{detectedAgentType === "a2a" ? (
|
||||
<AgentFormFields showAgentName={true} />
|
||||
) : selectedAgentTypeInfo ? (
|
||||
<DynamicAgentFormFields agentTypeInfo={selectedAgentTypeInfo} />
|
||||
) : (
|
||||
<AgentFormFields showAgentName={true} />
|
||||
)}
|
||||
|
||||
<div className="flex justify-end gap-2 mt-6">
|
||||
<AntButton onClick={() => {
|
||||
|
||||
@ -51,18 +51,18 @@ const AgentTable: React.FC<AgentTableProps> = ({
|
||||
cell: ({ row }) => {
|
||||
const agent = row.original;
|
||||
const name = agent.agent_name || "";
|
||||
return (
|
||||
return (
|
||||
<div className="flex items-center gap-2">
|
||||
<Tooltip title={name}>
|
||||
<Button
|
||||
size="xs"
|
||||
variant="light"
|
||||
<Button
|
||||
size="xs"
|
||||
variant="light"
|
||||
className="font-mono text-blue-500 bg-blue-50 hover:bg-blue-100 text-xs font-normal px-2 py-0.5 text-left overflow-hidden truncate min-w-[200px] justify-start"
|
||||
onClick={() => onAgentClick(agent.agent_id)}
|
||||
>
|
||||
onClick={() => onAgentClick(agent.agent_id)}
|
||||
>
|
||||
{name}
|
||||
</Button>
|
||||
</Tooltip>
|
||||
</Button>
|
||||
</Tooltip>
|
||||
<Tooltip title="Copy Agent ID">
|
||||
<CopyOutlined
|
||||
onClick={(e) => {
|
||||
@ -201,12 +201,12 @@ const AgentTable: React.FC<AgentTableProps> = ({
|
||||
<TableCell colSpan={columns.length} className="h-8 text-center">
|
||||
<div className="text-center text-gray-500">
|
||||
<p>No agents found. Create one to get started.</p>
|
||||
</div>
|
||||
</TableCell>
|
||||
</div>
|
||||
</TableCell>
|
||||
</TableRow>
|
||||
)}
|
||||
</TableBody>
|
||||
</Table>
|
||||
</TableBody>
|
||||
</Table>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
|
||||
@ -0,0 +1,67 @@
|
||||
import { Agent } from "./types";
|
||||
import { AgentCreateInfo } from "../networking";
|
||||
|
||||
/**
|
||||
* Detects the agent type from an agent's litellm_params.
|
||||
* Returns the agent_type string (e.g., "langgraph", "azure_ai_foundry", "bedrock_agentcore", or "a2a")
|
||||
*/
|
||||
export const detectAgentType = (agent: Agent): string => {
|
||||
const model = agent.litellm_params?.model || "";
|
||||
const customProvider = agent.litellm_params?.custom_llm_provider;
|
||||
|
||||
// Check by custom_llm_provider first
|
||||
if (customProvider === "langgraph") return "langgraph";
|
||||
if (customProvider === "azure_ai") return "azure_ai_foundry";
|
||||
if (customProvider === "bedrock") return "bedrock_agentcore";
|
||||
|
||||
// Check by model prefix
|
||||
if (model.startsWith("langgraph/")) return "langgraph";
|
||||
if (model.startsWith("azure_ai/agents/")) return "azure_ai_foundry";
|
||||
if (model.startsWith("bedrock/agentcore/")) return "bedrock_agentcore";
|
||||
|
||||
// Default to a2a
|
||||
return "a2a";
|
||||
};
|
||||
|
||||
/**
|
||||
* Parses agent data for dynamic form fields (non-A2A agents).
|
||||
* Extracts values from litellm_params based on the agent type metadata.
|
||||
*/
|
||||
export const parseDynamicAgentForForm = (
|
||||
agent: Agent,
|
||||
agentTypeInfo: AgentCreateInfo
|
||||
): Record<string, any> => {
|
||||
const values: Record<string, any> = {
|
||||
agent_name: agent.agent_name,
|
||||
description: agent.agent_card_params?.description || "",
|
||||
};
|
||||
|
||||
// Extract credential field values from litellm_params
|
||||
for (const field of agentTypeInfo.credential_fields) {
|
||||
if (field.include_in_litellm_params !== false) {
|
||||
values[field.key] = agent.litellm_params?.[field.key] || field.default_value || "";
|
||||
} else {
|
||||
// For fields not in litellm_params (like agent_id), try to extract from model string
|
||||
if (agentTypeInfo.model_template && agent.litellm_params?.model) {
|
||||
const model = agent.litellm_params.model;
|
||||
const templateParts = agentTypeInfo.model_template.split("/");
|
||||
const modelParts = model.split("/");
|
||||
|
||||
// Find the placeholder position and extract the value
|
||||
templateParts.forEach((part, index) => {
|
||||
if (part === `{${field.key}}` && modelParts[index]) {
|
||||
values[field.key] = modelParts[index];
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Extract cost configuration
|
||||
values.cost_per_query = agent.litellm_params?.cost_per_query;
|
||||
values.input_cost_per_token = agent.litellm_params?.input_cost_per_token;
|
||||
values.output_cost_per_token = agent.litellm_params?.output_cost_per_token;
|
||||
|
||||
return values;
|
||||
};
|
||||
|
||||
@ -165,7 +165,16 @@ const Navbar: React.FC<NavbarProps> = ({
|
||||
|
||||
<div className="flex items-center">
|
||||
<Link href="/" className="flex items-center">
|
||||
<img src={imageUrl} alt="LiteLLM Brand" className="h-10 w-auto" />
|
||||
<div className="relative">
|
||||
<img src={imageUrl} alt="LiteLLM Brand" className="h-10 w-auto" />
|
||||
<span
|
||||
className="absolute -top-1 -right-2 text-lg animate-bounce"
|
||||
style={{ animationDuration: '2s' }}
|
||||
title="Happy Holidays!"
|
||||
>
|
||||
🎄
|
||||
</span>
|
||||
</div>
|
||||
</Link>
|
||||
{version && (
|
||||
<a
|
||||
|
||||
@ -1860,6 +1860,25 @@ const ChatUI: React.FC<ChatUIProps> = ({
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Suggested prompts - show when chat is empty and not loading */}
|
||||
{chatHistory.length === 0 && !isLoading && (
|
||||
<div className="flex items-center gap-2 mb-3 overflow-x-auto">
|
||||
{(endpointType === EndpointType.A2A_AGENTS
|
||||
? ["What can you help me with?", "Tell me about yourself", "What tasks can you perform?"]
|
||||
: ["Write me a poem", "Explain quantum computing", "Draft a polite email requesting a meeting"]
|
||||
).map((prompt) => (
|
||||
<button
|
||||
key={prompt}
|
||||
type="button"
|
||||
className="shrink-0 rounded-full border border-gray-200 px-3 py-1 text-xs font-medium text-gray-600 transition-colors hover:bg-blue-50 hover:border-blue-300 hover:text-blue-600 cursor-pointer"
|
||||
onClick={() => setInputMessage(prompt)}
|
||||
>
|
||||
{prompt}
|
||||
</button>
|
||||
))}
|
||||
</div>
|
||||
)}
|
||||
|
||||
<div className="flex items-center gap-2">
|
||||
<div className="flex items-center flex-1 bg-white border border-gray-300 rounded-xl px-3 py-1 min-h-[44px]">
|
||||
{/* Left: attachment and code interpreter icons */}
|
||||
|
||||
@ -11,11 +11,25 @@ import type { TokenUsage } from "../chat_ui/ResponseMetrics";
|
||||
import type { MessageType, VectorStoreSearchResponse } from "../chat_ui/types";
|
||||
import { makeOpenAIChatCompletionRequest } from "../llm_calls/chat_completion";
|
||||
import { fetchAvailableModels } from "../llm_calls/fetch_models";
|
||||
import { Agent, fetchAvailableAgents } from "../llm_calls/fetch_agents";
|
||||
import { makeA2AStreamMessageRequest } from "../llm_calls/a2a_send_message";
|
||||
import { ComparisonPanel } from "./components/ComparisonPanel";
|
||||
import { MessageInput } from "./components/MessageInput";
|
||||
import {
|
||||
EndpointId,
|
||||
EndpointIdType,
|
||||
getAvailableEndpoints,
|
||||
getEndpointConfig,
|
||||
isAgentEndpoint,
|
||||
hasValidSelection,
|
||||
getComparisonSelection,
|
||||
modelOptionsToSelectorOptions,
|
||||
agentOptionsToSelectorOptions,
|
||||
} from "./endpoint_config";
|
||||
export interface ComparisonInstance {
|
||||
id: string;
|
||||
model: string;
|
||||
agent: string;
|
||||
messages: MessageType[];
|
||||
isLoading: boolean;
|
||||
tags: string[];
|
||||
@ -38,12 +52,13 @@ const GENERIC_FOLLOW_UPS = [
|
||||
"What are the next steps?",
|
||||
];
|
||||
const SUGGESTED_PROMPTS = ["Write me a poem", "Explain quantum computing", "Draft a polite email requesting a meeting"];
|
||||
const DEFAULT_ENDPOINT = "/v1/chat/completions";
|
||||
const DEFAULT_ENDPOINT = EndpointId.CHAT_COMPLETIONS;
|
||||
export default function CompareUI({ accessToken, disabledPersonalKeyCreation }: CompareUIProps) {
|
||||
const [comparisons, setComparisons] = useState<ComparisonInstance[]>([
|
||||
{
|
||||
id: "1",
|
||||
model: "",
|
||||
agent: "",
|
||||
messages: [],
|
||||
isLoading: false,
|
||||
tags: [],
|
||||
@ -58,6 +73,7 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
|
||||
{
|
||||
id: "2",
|
||||
model: "",
|
||||
agent: "",
|
||||
messages: [],
|
||||
isLoading: false,
|
||||
tags: [],
|
||||
@ -71,7 +87,18 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
|
||||
},
|
||||
]);
|
||||
const [modelOptions, setModelOptions] = useState<string[]>([]);
|
||||
const [agentOptions, setAgentOptions] = useState<Agent[]>([]);
|
||||
const [isLoadingModels, setIsLoadingModels] = useState(false);
|
||||
const [isLoadingAgents, setIsLoadingAgents] = useState(false);
|
||||
const [selectedEndpoint, setSelectedEndpoint] = useState<EndpointIdType>(DEFAULT_ENDPOINT);
|
||||
|
||||
// Derived state from endpoint config
|
||||
const endpointConfig = getEndpointConfig(selectedEndpoint);
|
||||
const isA2AMode = isAgentEndpoint(selectedEndpoint);
|
||||
const selectorOptions = isA2AMode
|
||||
? agentOptionsToSelectorOptions(agentOptions)
|
||||
: modelOptionsToSelectorOptions(modelOptions);
|
||||
const isLoadingOptions = isA2AMode ? isLoadingAgents : isLoadingModels;
|
||||
const [inputValue, setInputValue] = useState("");
|
||||
const [uploadedFile, setUploadedFile] = useState<File | null>(null);
|
||||
const [uploadedFilePreviewUrl, setUploadedFilePreviewUrl] = useState<string | null>(null);
|
||||
@ -134,6 +161,37 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
|
||||
active = false;
|
||||
};
|
||||
}, [effectiveApiKey]);
|
||||
|
||||
// Fetch agents when A2A mode is selected
|
||||
useEffect(() => {
|
||||
let active = true;
|
||||
const loadAgents = async () => {
|
||||
if (!effectiveApiKey || !isA2AMode) {
|
||||
setAgentOptions([]);
|
||||
return;
|
||||
}
|
||||
setIsLoadingAgents(true);
|
||||
try {
|
||||
const agents = await fetchAvailableAgents(effectiveApiKey);
|
||||
if (!active) return;
|
||||
setAgentOptions(agents);
|
||||
} catch (error) {
|
||||
console.error("CompareUI: failed to fetch agents", error);
|
||||
if (active) {
|
||||
setAgentOptions([]);
|
||||
}
|
||||
} finally {
|
||||
if (active) {
|
||||
setIsLoadingAgents(false);
|
||||
}
|
||||
}
|
||||
};
|
||||
loadAgents();
|
||||
return () => {
|
||||
active = false;
|
||||
};
|
||||
}, [effectiveApiKey, isA2AMode]);
|
||||
|
||||
useEffect(() => {
|
||||
if (modelOptions.length === 0) {
|
||||
return;
|
||||
@ -160,10 +218,12 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
|
||||
if (comparisons.length >= maxComparisons) {
|
||||
return;
|
||||
}
|
||||
const fallback = modelOptions[comparisons.length % (modelOptions.length || 1)] ?? "";
|
||||
const fallbackModel = modelOptions[comparisons.length % (modelOptions.length || 1)] ?? "";
|
||||
const fallbackAgent = agentOptions[comparisons.length % (agentOptions.length || 1)]?.agent_name ?? "";
|
||||
const newComparison: ComparisonInstance = {
|
||||
id: Date.now().toString(),
|
||||
model: fallback,
|
||||
model: fallbackModel,
|
||||
agent: fallbackAgent,
|
||||
messages: [],
|
||||
isLoading: false,
|
||||
tags: [],
|
||||
@ -430,8 +490,9 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
|
||||
if (targetComparisons.length === 0) {
|
||||
return;
|
||||
}
|
||||
if (targetComparisons.some((comparison) => !comparison.model)) {
|
||||
NotificationsManager.fromBackend("Select a model before sending a message.");
|
||||
// Validate selection based on endpoint type
|
||||
if (targetComparisons.some((comparison) => !hasValidSelection(comparison, selectedEndpoint))) {
|
||||
NotificationsManager.fromBackend(endpointConfig.validationMessage);
|
||||
return;
|
||||
}
|
||||
|
||||
@ -450,6 +511,8 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
|
||||
{
|
||||
id: string;
|
||||
model: string;
|
||||
agent: string;
|
||||
inputMessage: string;
|
||||
traceId: string;
|
||||
tags: string[];
|
||||
vectorStores: string[];
|
||||
@ -472,6 +535,8 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
|
||||
preparedTargets.set(comparison.id, {
|
||||
id: comparison.id,
|
||||
model: comparison.model,
|
||||
agent: comparison.agent,
|
||||
inputMessage: trimmed,
|
||||
traceId,
|
||||
tags: comparison.tags,
|
||||
vectorStores: comparison.vectorStores,
|
||||
@ -508,26 +573,55 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
|
||||
const guardrails = prepared.guardrails.length > 0 ? prepared.guardrails : undefined;
|
||||
const comparison = comparisons.find((c) => c.id === prepared.id);
|
||||
const useAdvancedParams = comparison?.useAdvancedParams ?? false;
|
||||
makeOpenAIChatCompletionRequest(
|
||||
prepared.apiChatHistory,
|
||||
(chunk, model) => appendAssistantChunk(prepared.id, chunk, model),
|
||||
prepared.model,
|
||||
effectiveApiKey,
|
||||
tags,
|
||||
undefined,
|
||||
(content) => appendReasoningContent(prepared.id, content),
|
||||
(time) => updateTimingDataForComparison(prepared.id, time),
|
||||
(usage) => updateUsageDataForComparison(prepared.id, usage),
|
||||
prepared.traceId,
|
||||
vectorStoreIds,
|
||||
guardrails,
|
||||
undefined,
|
||||
undefined,
|
||||
(searchResults) => updateSearchResultsForComparison(prepared.id, searchResults),
|
||||
useAdvancedParams ? prepared.temperature : undefined,
|
||||
useAdvancedParams ? prepared.maxTokens : undefined,
|
||||
(latency) => updateTotalLatencyForComparison(prepared.id, latency),
|
||||
)
|
||||
|
||||
// Use A2A or chat completion based on endpoint
|
||||
const requestPromise = isA2AMode
|
||||
? makeA2AStreamMessageRequest(
|
||||
prepared.agent,
|
||||
prepared.inputMessage,
|
||||
(text, model) => {
|
||||
// A2A sends full accumulated text, so replace instead of append
|
||||
setComparisons((prev) =>
|
||||
prev.map((c) => {
|
||||
if (c.id !== prepared.id) return c;
|
||||
const messages = [...c.messages];
|
||||
const last = messages[messages.length - 1];
|
||||
if (last && last.role === "assistant") {
|
||||
messages[messages.length - 1] = { ...last, content: text, model: last.model ?? model };
|
||||
} else {
|
||||
messages.push({ role: "assistant", content: text, model });
|
||||
}
|
||||
return { ...c, messages };
|
||||
}),
|
||||
);
|
||||
},
|
||||
effectiveApiKey,
|
||||
undefined,
|
||||
(time) => updateTimingDataForComparison(prepared.id, time),
|
||||
(latency) => updateTotalLatencyForComparison(prepared.id, latency),
|
||||
)
|
||||
: makeOpenAIChatCompletionRequest(
|
||||
prepared.apiChatHistory,
|
||||
(chunk, model) => appendAssistantChunk(prepared.id, chunk, model),
|
||||
prepared.model,
|
||||
effectiveApiKey,
|
||||
tags,
|
||||
undefined,
|
||||
(content) => appendReasoningContent(prepared.id, content),
|
||||
(time) => updateTimingDataForComparison(prepared.id, time),
|
||||
(usage) => updateUsageDataForComparison(prepared.id, usage),
|
||||
prepared.traceId,
|
||||
vectorStoreIds,
|
||||
guardrails,
|
||||
undefined,
|
||||
undefined,
|
||||
(searchResults) => updateSearchResultsForComparison(prepared.id, searchResults),
|
||||
useAdvancedParams ? prepared.temperature : undefined,
|
||||
useAdvancedParams ? prepared.maxTokens : undefined,
|
||||
(latency) => updateTotalLatencyForComparison(prepared.id, latency),
|
||||
);
|
||||
|
||||
requestPromise
|
||||
.catch((error) => {
|
||||
const errorMessage = error instanceof Error ? error.message : String(error);
|
||||
console.error("CompareUI: failed to fetch response", error);
|
||||
@ -618,11 +712,20 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
|
||||
</div>
|
||||
<div className="flex items-center gap-2">
|
||||
<span className="text-sm font-medium text-gray-600">Endpoint</span>
|
||||
<Tooltip title="Other endpoints will be available soon">
|
||||
<Select value={DEFAULT_ENDPOINT} disabled className="w-56">
|
||||
<Select.Option value={DEFAULT_ENDPOINT}>{DEFAULT_ENDPOINT}</Select.Option>
|
||||
</Select>
|
||||
</Tooltip>
|
||||
<Select
|
||||
value={selectedEndpoint}
|
||||
onChange={(value) => setSelectedEndpoint(value as EndpointIdType)}
|
||||
className="w-56"
|
||||
>
|
||||
{getAvailableEndpoints().map((endpoint) => (
|
||||
<Select.Option
|
||||
key={endpoint.value}
|
||||
value={endpoint.value}
|
||||
>
|
||||
{endpoint.label}
|
||||
</Select.Option>
|
||||
))}
|
||||
</Select>
|
||||
</div>
|
||||
<div className="flex items-center gap-3">
|
||||
<Button onClick={clearAllChats} disabled={!hasMessages} icon={<ClearOutlined />}>
|
||||
@ -654,8 +757,9 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
|
||||
onUpdate={(updates, options) => updateComparison(comparison.id, updates, options)}
|
||||
onRemove={() => removeComparison(comparison.id)}
|
||||
canRemove={comparisons.length > 1}
|
||||
modelOptions={modelOptions}
|
||||
isLoadingModels={isLoadingModels}
|
||||
selectorOptions={selectorOptions}
|
||||
isLoadingOptions={isLoadingOptions}
|
||||
endpointConfig={endpointConfig}
|
||||
apiKey={effectiveApiKey}
|
||||
/>
|
||||
))}
|
||||
@ -695,10 +799,10 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
|
||||
) : isAnyComparisonLoading ? (
|
||||
<span className="flex items-center gap-2 text-sm text-gray-500">
|
||||
<span className="h-2 w-2 rounded-full bg-blue-500 animate-pulse" aria-hidden />
|
||||
Gathering responses from all models...
|
||||
{endpointConfig.loadingMessage}
|
||||
</span>
|
||||
) : (
|
||||
<span className="text-sm text-gray-500">Send a prompt to compare models</span>
|
||||
<span className="text-sm text-gray-500">{endpointConfig.inputPlaceholder}</span>
|
||||
)}
|
||||
</div>
|
||||
{uploadedFile && (
|
||||
|
||||
@ -2,11 +2,13 @@ import { Settings, X } from "lucide-react";
|
||||
import { useState } from "react";
|
||||
import { ComparisonInstance } from "../CompareUI";
|
||||
import { MessageDisplay } from "./MessageDisplay";
|
||||
import { ModelSelector } from "./ModelSelector";
|
||||
import { UnifiedSelector } from "./UnifiedSelector";
|
||||
import TagSelector from "../../../tag_management/TagSelector";
|
||||
import VectorStoreSelector from "../../../vector_store_management/VectorStoreSelector";
|
||||
import GuardrailSelector from "../../../guardrails/GuardrailSelector";
|
||||
import { Checkbox, Divider, Popover, Slider } from "antd";
|
||||
import { SelectorOption, EndpointConfig, isAgentEndpoint, getComparisonSelection } from "../endpoint_config";
|
||||
|
||||
interface ComparisonPanelProps {
|
||||
comparison: ComparisonInstance;
|
||||
onUpdate: (
|
||||
@ -15,8 +17,9 @@ interface ComparisonPanelProps {
|
||||
) => void;
|
||||
onRemove: () => void;
|
||||
canRemove: boolean;
|
||||
modelOptions: string[];
|
||||
isLoadingModels: boolean;
|
||||
selectorOptions: SelectorOption[];
|
||||
isLoadingOptions: boolean;
|
||||
endpointConfig: EndpointConfig;
|
||||
apiKey: string;
|
||||
}
|
||||
export function ComparisonPanel({
|
||||
@ -24,10 +27,13 @@ export function ComparisonPanel({
|
||||
onUpdate,
|
||||
onRemove,
|
||||
canRemove,
|
||||
modelOptions,
|
||||
isLoadingModels,
|
||||
selectorOptions,
|
||||
isLoadingOptions,
|
||||
endpointConfig,
|
||||
apiKey,
|
||||
}: ComparisonPanelProps) {
|
||||
const isA2AMode = isAgentEndpoint(endpointConfig.id);
|
||||
const currentSelection = getComparisonSelection(comparison, endpointConfig.id);
|
||||
const [popoverVisible, setPopoverVisible] = useState(false);
|
||||
|
||||
const handleSyncChange = (checked: boolean) => {
|
||||
@ -194,14 +200,13 @@ export function ComparisonPanel({
|
||||
<div className="bg-white first:border-l-0 border-l border-gray-200 flex flex-col min-h-0">
|
||||
<div className="border-b flex items-center justify-between gap-3 px-4 py-3">
|
||||
<div className="flex items-center gap-3 flex-1">
|
||||
<ModelSelector
|
||||
value={comparison.model}
|
||||
models={modelOptions}
|
||||
loading={isLoadingModels}
|
||||
onChange={(model) =>
|
||||
onUpdate({
|
||||
model,
|
||||
})
|
||||
<UnifiedSelector
|
||||
value={currentSelection}
|
||||
options={selectorOptions}
|
||||
loading={isLoadingOptions}
|
||||
config={endpointConfig}
|
||||
onChange={(value) =>
|
||||
onUpdate(isA2AMode ? { agent: value } : { model: value })
|
||||
}
|
||||
/>
|
||||
<div className="flex items-center gap-2">
|
||||
|
||||
@ -0,0 +1,48 @@
|
||||
/**
|
||||
* Unified selector component that handles both model and agent selection
|
||||
* based on the current endpoint configuration.
|
||||
*/
|
||||
|
||||
import { Select, Spin } from "antd";
|
||||
import { SelectorOption, EndpointConfig } from "../endpoint_config";
|
||||
|
||||
interface UnifiedSelectorProps {
|
||||
value: string;
|
||||
options: SelectorOption[];
|
||||
loading: boolean;
|
||||
config: EndpointConfig;
|
||||
onChange: (value: string) => void;
|
||||
}
|
||||
|
||||
export function UnifiedSelector({
|
||||
value,
|
||||
options,
|
||||
loading,
|
||||
config,
|
||||
onChange,
|
||||
}: UnifiedSelectorProps) {
|
||||
return (
|
||||
<Select
|
||||
value={value || undefined}
|
||||
placeholder={loading ? `Loading ${config.selectorLabel.toLowerCase()}s...` : config.selectorPlaceholder}
|
||||
onChange={onChange}
|
||||
loading={loading}
|
||||
showSearch
|
||||
filterOption={(input, option) =>
|
||||
(option?.label ?? "").toLowerCase().includes(input.toLowerCase())
|
||||
}
|
||||
options={options}
|
||||
className="w-48"
|
||||
notFoundContent={
|
||||
loading ? (
|
||||
<div className="flex items-center justify-center py-2">
|
||||
<Spin size="small" />
|
||||
</div>
|
||||
) : (
|
||||
`No ${config.selectorLabel.toLowerCase()}s available`
|
||||
)
|
||||
}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
@ -0,0 +1,138 @@
|
||||
/**
|
||||
* Endpoint configuration for CompareUI.
|
||||
* Add new endpoints here to extend the comparison functionality.
|
||||
*/
|
||||
|
||||
import { Agent } from "../llm_calls/fetch_agents";
|
||||
|
||||
// Endpoint identifiers
|
||||
export const EndpointId = {
|
||||
CHAT_COMPLETIONS: "/v1/chat/completions",
|
||||
A2A_AGENTS: "/a2a",
|
||||
// Future endpoints:
|
||||
// RESPONSES: "/v1/responses",
|
||||
// ANTHROPIC: "/v1/messages",
|
||||
} as const;
|
||||
|
||||
export type EndpointIdType = (typeof EndpointId)[keyof typeof EndpointId];
|
||||
|
||||
// Selector type determines what the user picks (model, agent, etc.)
|
||||
export type SelectorType = "model" | "agent";
|
||||
|
||||
export interface EndpointConfig {
|
||||
id: EndpointIdType;
|
||||
label: string;
|
||||
selectorType: SelectorType;
|
||||
selectorLabel: string;
|
||||
selectorPlaceholder: string;
|
||||
inputPlaceholder: string;
|
||||
loadingMessage: string;
|
||||
validationMessage: string;
|
||||
}
|
||||
|
||||
// Endpoint configurations
|
||||
export const ENDPOINT_CONFIGS: Record<EndpointIdType, EndpointConfig> = {
|
||||
[EndpointId.CHAT_COMPLETIONS]: {
|
||||
id: EndpointId.CHAT_COMPLETIONS,
|
||||
label: "/v1/chat/completions",
|
||||
selectorType: "model",
|
||||
selectorLabel: "Model",
|
||||
selectorPlaceholder: "Select a model",
|
||||
inputPlaceholder: "Send a prompt to compare models",
|
||||
loadingMessage: "Gathering responses from all models...",
|
||||
validationMessage: "Select a model before sending a message.",
|
||||
},
|
||||
[EndpointId.A2A_AGENTS]: {
|
||||
id: EndpointId.A2A_AGENTS,
|
||||
label: "/a2a (Agents)",
|
||||
selectorType: "agent",
|
||||
selectorLabel: "Agent",
|
||||
selectorPlaceholder: "Select an agent",
|
||||
inputPlaceholder: "Send a message to compare agents",
|
||||
loadingMessage: "Gathering responses from all agents...",
|
||||
validationMessage: "Select an agent before sending a message.",
|
||||
},
|
||||
};
|
||||
|
||||
// Get list of available endpoints for the dropdown
|
||||
export const getAvailableEndpoints = () =>
|
||||
Object.values(ENDPOINT_CONFIGS).map((config) => ({
|
||||
value: config.id,
|
||||
label: config.label,
|
||||
}));
|
||||
|
||||
// Helper to get config for an endpoint
|
||||
export const getEndpointConfig = (endpointId: EndpointIdType): EndpointConfig =>
|
||||
ENDPOINT_CONFIGS[endpointId];
|
||||
|
||||
// Helper to check if endpoint uses agents
|
||||
export const isAgentEndpoint = (endpointId: EndpointIdType): boolean =>
|
||||
ENDPOINT_CONFIGS[endpointId].selectorType === "agent";
|
||||
|
||||
// Helper to check if endpoint uses models
|
||||
export const isModelEndpoint = (endpointId: EndpointIdType): boolean =>
|
||||
ENDPOINT_CONFIGS[endpointId].selectorType === "model";
|
||||
|
||||
// Selector options type - unified interface for both models and agents
|
||||
export interface SelectorOption {
|
||||
value: string;
|
||||
label: string;
|
||||
}
|
||||
|
||||
// Convert model options to unified format
|
||||
export const modelOptionsToSelectorOptions = (models: string[]): SelectorOption[] =>
|
||||
models.map((model) => ({ value: model, label: model }));
|
||||
|
||||
// Convert agent options to unified format
|
||||
export const agentOptionsToSelectorOptions = (agents: Agent[]): SelectorOption[] =>
|
||||
agents.map((agent) => ({
|
||||
value: agent.agent_name,
|
||||
label: agent.agent_name || agent.agent_id,
|
||||
}));
|
||||
|
||||
// Get the selected value field name based on endpoint
|
||||
export const getSelectionFieldName = (endpointId: EndpointIdType): "model" | "agent" =>
|
||||
isAgentEndpoint(endpointId) ? "agent" : "model";
|
||||
|
||||
// Get the current selection from a comparison based on endpoint
|
||||
export const getComparisonSelection = (
|
||||
comparison: { model: string; agent: string },
|
||||
endpointId: EndpointIdType
|
||||
): string => (isAgentEndpoint(endpointId) ? comparison.agent : comparison.model);
|
||||
|
||||
// Check if comparison has a valid selection for the endpoint
|
||||
export const hasValidSelection = (
|
||||
comparison: { model: string; agent: string },
|
||||
endpointId: EndpointIdType
|
||||
): boolean => {
|
||||
const selection = getComparisonSelection(comparison, endpointId);
|
||||
return Boolean(selection && selection.trim());
|
||||
};
|
||||
|
||||
/**
|
||||
* To add a new endpoint:
|
||||
*
|
||||
* 1. Add the endpoint ID to EndpointId const
|
||||
* 2. Add configuration to ENDPOINT_CONFIGS
|
||||
* 3. If the endpoint uses a new selector type (not model or agent):
|
||||
* - Add the type to SelectorType
|
||||
* - Add fetch logic in CompareUI.tsx
|
||||
* - Add conversion function (e.g., xxxOptionsToSelectorOptions)
|
||||
* 4. Add request handling in CompareUI.tsx handleSendMessage
|
||||
*
|
||||
* Example for adding /v1/responses endpoint:
|
||||
*
|
||||
* EndpointId.RESPONSES = "/v1/responses"
|
||||
*
|
||||
* ENDPOINT_CONFIGS[EndpointId.RESPONSES] = {
|
||||
* id: EndpointId.RESPONSES,
|
||||
* label: "/v1/responses",
|
||||
* selectorType: "model",
|
||||
* selectorLabel: "Model",
|
||||
* selectorPlaceholder: "Select a model",
|
||||
* inputPlaceholder: "Send a prompt to compare responses",
|
||||
* loadingMessage: "Gathering responses...",
|
||||
* validationMessage: "Select a model before sending.",
|
||||
* }
|
||||
*/
|
||||
|
||||
@ -337,7 +337,6 @@ export const makeA2AStreamMessageRequest = async (
|
||||
if (artifact.parts && Array.isArray(artifact.parts)) {
|
||||
for (const part of artifact.parts) {
|
||||
if (part.kind === "text" && part.text) {
|
||||
// Accumulate actual response content
|
||||
accumulatedText += part.text;
|
||||
onTextUpdate(accumulatedText, `a2a_agent/${agentId}`);
|
||||
}
|
||||
@ -358,17 +357,11 @@ export const makeA2AStreamMessageRequest = async (
|
||||
}
|
||||
}
|
||||
// Handle status-update chunks (progress messages like "Processing request...")
|
||||
// Only show these temporarily if we haven't received actual content yet
|
||||
else if (chunkKind === "status-update" && result.status?.message?.parts) {
|
||||
// Skip status messages once we have real content
|
||||
if (!accumulatedText) {
|
||||
for (const part of result.status.message.parts) {
|
||||
if (part.kind === "text" && part.text) {
|
||||
// Show as temporary status - will be replaced when real content arrives
|
||||
onTextUpdate(part.text, `a2a_agent/${agentId}`);
|
||||
}
|
||||
}
|
||||
}
|
||||
// These are metadata/status updates, not actual response content
|
||||
// We skip showing them in the chat UI - they're captured in metadata instead
|
||||
else if (chunkKind === "status-update") {
|
||||
// Status updates are handled via metadata extraction, not shown as text
|
||||
// This prevents "Processing request..." from appearing in the response
|
||||
}
|
||||
// Direct parts array (fallback)
|
||||
else if (result.parts && Array.isArray(result.parts)) {
|
||||
@ -381,10 +374,16 @@ export const makeA2AStreamMessageRequest = async (
|
||||
}
|
||||
}
|
||||
|
||||
// Handle JSON-RPC error response
|
||||
if (chunk.error) {
|
||||
throw new Error(chunk.error.message);
|
||||
const errorMessage = chunk.error.message || "Unknown A2A error";
|
||||
throw new Error(errorMessage);
|
||||
}
|
||||
} catch (parseError) {
|
||||
// Re-throw if it's an actual error we threw (not a parse error)
|
||||
if (parseError instanceof Error && parseError.message && !parseError.message.includes("JSON")) {
|
||||
throw parseError;
|
||||
}
|
||||
// Only warn if it's not a JSON parse error on an empty/partial line
|
||||
if (line.trim().length > 0) {
|
||||
console.warn("Failed to parse A2A streaming chunk:", line, parseError);
|
||||
|
||||
Loading…
Reference in New Issue
Block a user