diff --git a/litellm/a2a_protocol/litellm_completion_bridge/handler.py b/litellm/a2a_protocol/litellm_completion_bridge/handler.py index a62a2b0c724..58e4366ad80 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/handler.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/handler.py @@ -37,6 +37,7 @@ _AGENT_ONLY_PARAMS: Final = frozenset( "agent_id", "agent_card_params", A2A_USER_API_KEY_HASH_PARAM, + "databricks_oauth", } ) @@ -46,6 +47,19 @@ class A2ACompletionBridgeHandler: Static methods for handling A2A requests via LiteLLM completion. """ + @staticmethod + def _merge_stream_values(previous: object, current: object) -> object: + if isinstance(previous, Mapping) and isinstance(current, Mapping): + merged = dict(previous) + for key, value in current.items(): + merged[key] = ( + A2ACompletionBridgeHandler._merge_stream_values(merged[key], value) if key in merged else value + ) + return merged + if isinstance(previous, list) and isinstance(current, list): + return [*previous, *current] + return current + @staticmethod def _build_completion_params( params: dict[str, Any], @@ -59,7 +73,12 @@ class A2ACompletionBridgeHandler: message: Final = params.get("message", {}) # Transform A2A message to OpenAI format - openai_messages: Final = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message) + supplied_messages: Final = params.get("messages") + openai_messages: Final = ( + supplied_messages + if isinstance(supplied_messages, list) + else A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message) + ) # Get completion params custom_llm_provider: Final = litellm_params.get("custom_llm_provider") @@ -84,11 +103,13 @@ class A2ACompletionBridgeHandler: "api_base": api_base, "stream": stream, } + configured_headers: Final[object] = litellm_params.get("extra_headers") or litellm_params.get("headers") # Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.) litellm_params_to_add: Final = { k: v for k, v in litellm_params.items() - if k not in ("model", "custom_llm_provider") and k not in _AGENT_ONLY_PARAMS + if k not in ("model", "custom_llm_provider", "extra_headers", "headers", "api_base", "stream") + and k not in _AGENT_ONLY_PARAMS } completion_params.update(litellm_params_to_add) # Apply forward metadata AFTER the litellm_params merge so the helper @@ -100,10 +121,10 @@ class A2ACompletionBridgeHandler: params=params, ) - if agent_extra_headers: + if agent_extra_headers or configured_headers: completion_params["extra_headers"] = merge_agent_headers( dynamic_headers=agent_extra_headers, - static_headers=completion_params.get("extra_headers"), + static_headers=configured_headers if isinstance(configured_headers, Mapping) else None, ) return completion_params @@ -122,6 +143,7 @@ class A2ACompletionBridgeHandler: litellm_params: dict[str, Any], api_base: str | None = None, agent_extra_headers: dict[str, str] | None = None, + agent_static_headers: Mapping[str, str] | None = None, *, _skip_a2a_provider_routing: bool = False, ) -> dict[str, object]: @@ -135,6 +157,7 @@ class A2ACompletionBridgeHandler: api_base: API base URL from agent_card_params agent_extra_headers: Per-request headers (from x-a2a-{agent}-* rewrite and admin extra_headers) to forward on the upstream HTTP call. + agent_static_headers: Configured headers for provider-specific routing. Returns: A2A SendMessageResponse dict @@ -149,13 +172,24 @@ class A2ACompletionBridgeHandler: if a2a_provider_config is not None: verbose_logger.info("A2A: Using provider config for %s", custom_llm_provider) - return await a2a_provider_config.handle_non_streaming( - request_id=request_id, - params=params, - api_base=api_base, - litellm_params=litellm_params, - agent_extra_headers=agent_extra_headers, - ) + provider_params: Final = dict(params) + if custom_llm_provider == "pydantic_ai_agents" or ( + custom_llm_provider == "bedrock" + and isinstance(litellm_params.get("model"), str) + and "agentcore" in litellm_params["model"] + ): + provider_params.pop("messages", None) + provider_kwargs: Final[dict[str, Any]] = { + "request_id": request_id, + "params": provider_params, + "api_base": api_base, + "litellm_params": litellm_params, + "agent_extra_headers": agent_extra_headers, + "agent_static_headers": agent_static_headers, + } + if litellm_params.get("timeout") is not None: + provider_kwargs["timeout"] = litellm_params["timeout"] + return await a2a_provider_config.handle_non_streaming(**provider_kwargs) completion_params: Final = A2ACompletionBridgeHandler._build_completion_params( params=params, @@ -185,6 +219,7 @@ class A2ACompletionBridgeHandler: litellm_params: dict[str, Any], api_base: str | None = None, agent_extra_headers: dict[str, str] | None = None, + agent_static_headers: Mapping[str, str] | None = None, *, _skip_a2a_provider_routing: bool = False, ) -> AsyncIterator[dict[str, object]]: @@ -204,6 +239,7 @@ class A2ACompletionBridgeHandler: api_base: API base URL from agent_card_params agent_extra_headers: Per-request headers (from x-a2a-{agent}-* rewrite and admin extra_headers) to forward on the upstream HTTP call. + agent_static_headers: Configured headers for provider-specific routing. Yields: A2A streaming response events @@ -218,14 +254,31 @@ class A2ACompletionBridgeHandler: if a2a_provider_config is not None: verbose_logger.info("A2A: Using provider config for %s (streaming)", custom_llm_provider) - async for chunk in a2a_provider_config.handle_streaming( - request_id=request_id, - params=params, - api_base=api_base, - litellm_params=litellm_params, - agent_extra_headers=agent_extra_headers, + provider_params: Final = dict(params) + if custom_llm_provider == "pydantic_ai_agents" or ( + custom_llm_provider == "bedrock" + and isinstance(litellm_params.get("model"), str) + and "agentcore" in litellm_params["model"] ): - yield chunk + provider_params.pop("messages", None) + provider_kwargs: Final[dict[str, Any]] = { + "request_id": request_id, + "params": provider_params, + "api_base": api_base, + "litellm_params": litellm_params, + "agent_extra_headers": agent_extra_headers, + "agent_static_headers": agent_static_headers, + } + if litellm_params.get("timeout") is not None: + provider_kwargs["timeout"] = litellm_params["timeout"] + provider_stream: Final = a2a_provider_config.handle_streaming(**provider_kwargs) + try: + async for chunk in provider_stream: + yield chunk + finally: + close_provider_stream = getattr(provider_stream, "aclose", None) + if close_provider_stream is not None: + await close_provider_stream() return @@ -259,29 +312,104 @@ class A2ACompletionBridgeHandler: # Call litellm.acompletion with streaming response: Final = await A2ACompletionBridgeHandler._acompletion(completion_params) - # 3. Accumulate content and emit artifact update - accumulated_text = "" + # 3. Forward content as artifact updates + accumulated_tool_calls: Final[list[object]] = [] # mutable-ok: collect streaming tool-call deltas + choice_texts: dict[int, str] = {} + choice_tool_calls: dict[int, list[object]] = {} + choice_delta_fields: dict[int, dict[str, object]] = {} + choice_logprobs: dict[int, dict[str, object]] = {} + choice_finish_reasons: dict[int, str] = {} + stream_metadata: dict[str, str] = {} + stream_usage: object | None = None + stream_finish_reason: str | None = None chunk_count = 0 - async for chunk in response: - chunk_count += 1 + try: + async for chunk in response: + chunk_count += 1 - # Extract delta content - content = "" - if chunk is not None and hasattr(chunk, "choices") and chunk.choices: - choice = chunk.choices[0] - if hasattr(choice, "delta") and choice.delta: - content = choice.delta.content or "" + raw_usage = getattr(chunk, "usage", None) + if isinstance(raw_usage, Mapping): + stream_usage = raw_usage + else: + dump_usage = getattr(raw_usage, "model_dump", None) + if callable(dump_usage): + dumped_usage = dump_usage(exclude_none=True) + if isinstance(dumped_usage, Mapping): + stream_usage = dumped_usage + else: + dict_usage = getattr(raw_usage, "dict", None) + if callable(dict_usage): + dumped_usage = dict_usage(exclude_none=True) + if isinstance(dumped_usage, Mapping): + stream_usage = dumped_usage - if content: - accumulated_text += content + for metadata_name in ("system_fingerprint", "service_tier"): + metadata_value = getattr(chunk, metadata_name, None) + if not isinstance(metadata_value, str): + chunk_fields = A2ACompletionBridgeTransformation._model_dump(chunk) + metadata_value = chunk_fields.get(metadata_name) + if isinstance(metadata_value, str) and metadata_value: + stream_metadata[metadata_name] = metadata_value - # Emit artifact update with accumulated content - if accumulated_text: - artifact_event: Final = A2ACompletionBridgeTransformation.create_artifact_update_event( - ctx=ctx, - text=accumulated_text, - ) - yield artifact_event + # Extract delta content + choices = getattr(chunk, "choices", None) if chunk is not None else None + if isinstance(choices, (list, tuple)): + for choice_position, choice in enumerate(choices): + raw_index = getattr(choice, "index", choice_position) + choice_index = raw_index if isinstance(raw_index, int) else choice_position + choice_texts.setdefault(choice_index, "") + raw_finish_reason = getattr(choice, "finish_reason", None) + if isinstance(raw_finish_reason, str) and raw_finish_reason: + choice_finish_reasons[choice_index] = raw_finish_reason + if choice_index == 0 or stream_finish_reason is None: + stream_finish_reason = raw_finish_reason + content = "" + delta = getattr(choice, "delta", None) + if delta: + raw_content = getattr(delta, "content", None) + content = raw_content if isinstance(raw_content, str) else "" + choice_texts[choice_index] += content + tool_calls = getattr(delta, "tool_calls", None) + if isinstance(tool_calls, (list, tuple)): + accumulated_tool_calls.extend(tool_calls) + choice_tool_calls.setdefault(choice_index, []).extend(tool_calls) + delta_fields = A2ACompletionBridgeTransformation._model_dump(delta) + if delta_fields: + choice_fields = choice_delta_fields.setdefault(choice_index, {}) + for field, value in delta_fields.items(): + if field in {"content", "role", "tool_calls"} or value is None: + continue + previous = choice_fields.get(field) + if (isinstance(previous, str) and isinstance(value, str)) or ( + isinstance(previous, list) and isinstance(value, list) + ): + choice_fields[field] = previous + value + elif isinstance(previous, Mapping) and isinstance(value, Mapping): + choice_fields[field] = {**previous, **value} + else: + choice_fields[field] = value + + raw_logprobs = getattr(choice, "logprobs", None) + serialized_logprobs = A2ACompletionBridgeTransformation._model_dump(raw_logprobs) + if serialized_logprobs: + previous_logprobs = choice_logprobs.get(choice_index, {}) + merged_logprobs = A2ACompletionBridgeHandler._merge_stream_values( + previous_logprobs, serialized_logprobs + ) + if isinstance(merged_logprobs, dict): + choice_logprobs[choice_index] = merged_logprobs + + if content: + artifact_event: Final = A2ACompletionBridgeTransformation.create_artifact_update_event( + ctx=ctx, + text=content, + index=choice_index, + ) + yield artifact_event + finally: + close_response = getattr(response, "aclose", None) + if close_response is not None: + await close_response() # 4. Emit final status update (kind: "status-update", status: "completed", final: true) completed_event: Final = A2ACompletionBridgeTransformation.create_status_update_event( @@ -289,6 +417,47 @@ class A2ACompletionBridgeHandler: state="completed", final=True, ) + if accumulated_tool_calls: + completed_event["result"]["tool_calls"] = accumulated_tool_calls + if stream_finish_reason: + completed_event["result"]["finish_reason"] = stream_finish_reason + if stream_usage is not None: + completed_event["usage"] = stream_usage + for metadata_name, metadata_value in stream_metadata.items(): + completed_event[metadata_name] = metadata_value + choice_indices = sorted( + set(choice_texts) + | set(choice_tool_calls) + | set(choice_delta_fields) + | set(choice_logprobs) + | set(choice_finish_reasons) + ) + if choice_indices: + choice_payloads: list[dict[str, object]] = [] + for choice_index in choice_indices: + choice_payload: dict[str, object] = { + "index": choice_index, + "message": { + "kind": "message", + "role": "agent", + "parts": [{"kind": "text", "text": ""}], + **( + {"tool_calls": choice_tool_calls[choice_index]} + if choice_tool_calls.get(choice_index) + else {} + ), + }, + **( + {"finish_reason": choice_finish_reasons[choice_index]} + if choice_index in choice_finish_reasons + else {} + ), + **({"logprobs": choice_logprobs[choice_index]} if choice_index in choice_logprobs else {}), + } + if choice_delta_fields.get(choice_index): + choice_payload["delta"] = choice_delta_fields[choice_index] + choice_payloads.append(choice_payload) + completed_event["result"]["choices"] = choice_payloads yield completed_event verbose_logger.info( diff --git a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py index 15cf77708f9..b74cc0fa11c 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py @@ -34,6 +34,7 @@ class A2AStreamingContext: self.request_id = request_id self.task_id = str(uuid4()) self.context_id = str(uuid4()) + self.artifact_id = str(uuid4()) self.input_message = input_message self.accumulated_text = "" self.has_emitted_task = False @@ -151,6 +152,22 @@ class A2ACompletionBridgeTransformation: return [openai_message] + @staticmethod + def _model_dump(value: Any) -> dict[str, Any]: + if isinstance(value, dict): + return value + dump = getattr(value, "model_dump", None) + if callable(dump): + dumped = dump(exclude_none=True) + if isinstance(dumped, dict): + return dumped + dump = getattr(value, "dict", None) + if callable(dump): + dumped = dump(exclude_none=True) + if isinstance(dumped, dict): + return dumped + return {} + @staticmethod def openai_response_to_a2a_response( response: Any, @@ -166,20 +183,78 @@ class A2ACompletionBridgeTransformation: Returns: A2A SendMessageResponse dict """ - # Extract content from response - content = "" - if hasattr(response, "choices") and response.choices: - choice: Final = response.choices[0] - if hasattr(choice, "message") and choice.message: - content = choice.message.content or "" + serialized_choices: list[dict[str, Any]] = [] + raw_choices: Final = getattr(response, "choices", None) + if raw_choices: + for choice in raw_choices: + raw_message = getattr(choice, "message", None) + message_fields: Final = A2ACompletionBridgeTransformation._model_dump(raw_message) + raw_content = message_fields.get("content") + if raw_content is None: + raw_content = getattr(raw_message, "content", None) + content: Final = raw_content if isinstance(raw_content, str) else "" + message: Final = { + "kind": "message", + "role": "agent", + "parts": [{"kind": "text", "text": content}], + "messageId": uuid4().hex, + } + raw_tool_calls = message_fields.get("tool_calls") + if raw_tool_calls: + message["tool_calls"] = [ + call.model_dump(exclude_none=True) + if hasattr(call, "model_dump") + else call.dict(exclude_none=True) + if hasattr(call, "dict") + else call + for call in raw_tool_calls + ] + for field in ( + "annotations", + "audio", + "function_call", + "images", + "provider_specific_fields", + "reasoning_content", + "reasoning_items", + "refusal", + "thinking_blocks", + ): + value = message_fields.get(field) + if value is not None: + message[field] = value + choice_fields: Final = A2ACompletionBridgeTransformation._model_dump(choice) + finish_reason: Final = choice_fields.get("finish_reason") + if finish_reason is None: + raw_finish_reason = getattr(choice, "finish_reason", None) + finish_reason = raw_finish_reason if isinstance(raw_finish_reason, str) else None + if finish_reason: + message["finish_reason"] = finish_reason + choice_payload: Final[dict[str, Any]] = { + "index": len(serialized_choices), + "message": message, + } + logprobs = choice_fields.get("logprobs") + if logprobs is None: + raw_logprobs = getattr(choice, "logprobs", None) + logprobs = raw_logprobs if isinstance(raw_logprobs, dict) else None + if logprobs is not None: + choice_payload["logprobs"] = logprobs + message["logprobs"] = logprobs + serialized_choices.append(choice_payload) - # Build A2A message - a2a_message: Final = { - "kind": "message", - "role": "agent", - "parts": [{"kind": "text", "text": content}], - "messageId": uuid4().hex, - } + a2a_message: Final = ( + serialized_choices[0]["message"] + if serialized_choices + else { + "kind": "message", + "role": "agent", + "parts": [{"kind": "text", "text": ""}], + "messageId": uuid4().hex, + } + ) + + usage: Final = getattr(response, "usage", None) # Build A2A response a2a_response: Final = { @@ -187,8 +262,16 @@ class A2ACompletionBridgeTransformation: "id": request_id, "result": a2a_message, } + if usage is not None: + a2a_response["usage"] = usage.model_dump(exclude_none=True) if hasattr(usage, "model_dump") else usage + for field in ("system_fingerprint", "service_tier"): + value = getattr(response, field, None) + if value is not None: + a2a_response[field] = value + if len(serialized_choices) > 1: + a2a_response["choices"] = serialized_choices - verbose_logger.debug("OpenAI -> A2A transform: content_length=%s", len(content)) + verbose_logger.debug("OpenAI -> A2A transform: content_length=%s", len(a2a_message["parts"][0]["text"])) return a2a_response @@ -277,6 +360,7 @@ class A2ACompletionBridgeTransformation: def create_artifact_update_event( ctx: A2AStreamingContext, text: str, + index: int | None = None, ) -> dict[str, Any]: """ Create an artifact update event with content. @@ -285,15 +369,18 @@ class A2ACompletionBridgeTransformation: ctx: Streaming context text: The text content for the artifact """ + artifact: Final[dict[str, Any]] = { + "artifactId": ctx.artifact_id, + "name": "response", + "parts": [{"kind": "text", "text": text}], + } + if index is not None: + artifact["index"] = index return { "id": ctx.request_id, "jsonrpc": "2.0", "result": { - "artifact": { - "artifactId": str(uuid4()), - "name": "response", - "parts": [{"kind": "text", "text": text}], - }, + "artifact": artifact, "contextId": ctx.context_id, "kind": "artifact-update", "taskId": ctx.task_id, diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/config.py b/litellm/a2a_protocol/providers/bedrock_agentcore/config.py index 2b37c0c4906..04541d8863e 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/config.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/config.py @@ -38,6 +38,7 @@ class BedrockAgentCoreA2AConfig(BaseA2AProviderConfig): params=params, litellm_params=litellm_params, agent_extra_headers=kwargs.get("agent_extra_headers"), + timeout=kwargs.get("timeout"), ) async def handle_streaming( @@ -58,5 +59,6 @@ class BedrockAgentCoreA2AConfig(BaseA2AProviderConfig): params=params, litellm_params=litellm_params, agent_extra_headers=kwargs.get("agent_extra_headers"), + timeout=kwargs.get("timeout"), ): yield chunk diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py index db57072ca38..64979848110 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py @@ -31,6 +31,7 @@ class BedrockAgentCoreA2AHandler: params: dict[str, Any], litellm_params: dict[str, Any], agent_extra_headers: dict[str, str] | None = None, + timeout: float | None = None, ) -> dict[str, Any]: """ Handle non-streaming A2A request to AgentCore. @@ -62,6 +63,7 @@ class BedrockAgentCoreA2AHandler: url, headers=headers, data=body, + timeout=timeout, ) response.raise_for_status() response_data: Final = response.json() @@ -77,6 +79,7 @@ class BedrockAgentCoreA2AHandler: params: dict[str, Any], litellm_params: dict[str, Any], agent_extra_headers: dict[str, str] | None = None, + timeout: float | None = None, ) -> AsyncIterator[dict[str, Any]]: """ Handle streaming A2A request to AgentCore. @@ -110,6 +113,7 @@ class BedrockAgentCoreA2AHandler: headers=headers, data=body, stream=True, + timeout=timeout, ) response.raise_for_status() diff --git a/litellm/a2a_protocol/providers/langflow/config.py b/litellm/a2a_protocol/providers/langflow/config.py index 54d403f88c0..eec32d4e274 100644 --- a/litellm/a2a_protocol/providers/langflow/config.py +++ b/litellm/a2a_protocol/providers/langflow/config.py @@ -6,6 +6,7 @@ from litellm.a2a_protocol.litellm_completion_bridge.handler import ( A2ACompletionBridgeHandler, ) from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig +from litellm.interactions.agents.utils import merge_agent_headers from litellm.llms.langflow.a2a import merge_a2a_session_into_litellm_params @@ -28,11 +29,16 @@ class LangFlowA2AConfig(BaseA2AProviderConfig): litellm_params = merge_a2a_session_into_litellm_params( litellm_params, params, litellm_params.get(A2A_USER_API_KEY_HASH_PARAM) ) + forwarded_headers = merge_agent_headers( + dynamic_headers=kwargs.get("agent_extra_headers"), + static_headers=kwargs.get("agent_static_headers"), + ) return await A2ACompletionBridgeHandler.handle_non_streaming( request_id=request_id, params=params, litellm_params=litellm_params, api_base=api_base, + agent_extra_headers=forwarded_headers, _skip_a2a_provider_routing=True, ) @@ -51,11 +57,16 @@ class LangFlowA2AConfig(BaseA2AProviderConfig): litellm_params = merge_a2a_session_into_litellm_params( litellm_params, params, litellm_params.get(A2A_USER_API_KEY_HASH_PARAM) ) + forwarded_headers = merge_agent_headers( + dynamic_headers=kwargs.get("agent_extra_headers"), + static_headers=kwargs.get("agent_static_headers"), + ) async for chunk in A2ACompletionBridgeHandler.handle_streaming( request_id=request_id, params=params, litellm_params=litellm_params, api_base=api_base, + agent_extra_headers=forwarded_headers, _skip_a2a_provider_routing=True, ): yield chunk diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py index 024e8c179c2..86d9b4d2eae 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py @@ -6,6 +6,7 @@ This module provides fake streaming by converting non-streaming responses into s """ import asyncio +import time from collections.abc import AsyncIterator, Mapping, Sequence from typing import Any, Final, Protocol, cast, runtime_checkable from uuid import uuid4 @@ -101,6 +102,7 @@ class PydanticAITransformation: max_attempts: int = 30, poll_interval: float = 0.5, agent_extra_headers: dict[str, str] | None = None, + timeout: float | None = None, ) -> dict[str, object]: """ Poll for task completion using tasks/get method. @@ -112,11 +114,16 @@ class PydanticAITransformation: request_id: JSON-RPC request ID max_attempts: Maximum polling attempts poll_interval: Seconds between poll attempts + timeout: Total polling timeout in seconds Returns: Completed task response """ + deadline: Final[float | None] = time.monotonic() + max(timeout, 0.0) if timeout is not None else None for attempt in range(max_attempts): + remaining: float | None = deadline - time.monotonic() if deadline is not None else None + if remaining is not None and remaining <= 0: + break poll_request = { "jsonrpc": "2.0", "id": f"{request_id}-poll-{attempt}", @@ -131,6 +138,7 @@ class PydanticAITransformation: **(agent_extra_headers or {}), "Content-Type": "application/json", }, + timeout=remaining, ) response.raise_for_status() poll_data = _STR_KEY_DICT_ADAPTER.validate_python(response.json()) @@ -146,9 +154,16 @@ class PydanticAITransformation: elif state in ("failed", "canceled"): raise Exception(f"Task {task_id} ended with state: {state}") - await asyncio.sleep(poll_interval) + if deadline is None: + await asyncio.sleep(poll_interval) + continue + remaining = deadline - time.monotonic() + if remaining <= 0: + break + await asyncio.sleep(min(poll_interval, remaining)) - raise TimeoutError(f"Task {task_id} did not complete within {max_attempts * poll_interval} seconds") + timeout_description = timeout if timeout is not None else max_attempts * poll_interval + raise TimeoutError(f"Task {task_id} did not complete within {timeout_description} seconds") @staticmethod async def _send_and_poll_raw( @@ -203,6 +218,7 @@ class PydanticAITransformation: llm_provider=cast(Any, "pydantic_ai_agent"), params={"timeout": timeout}, ) + deadline: Final[float | None] = time.monotonic() + max(timeout, 0.0) if timeout is not None else None response: Final = await client.post( endpoint, json=a2a_request, @@ -210,6 +226,7 @@ class PydanticAITransformation: **(agent_extra_headers or {}), "Content-Type": "application/json", }, + timeout=timeout, ) response.raise_for_status() response_data = _STR_KEY_DICT_ADAPTER.validate_python(response.json()) @@ -230,6 +247,7 @@ class PydanticAITransformation: task_id=task_id, request_id=request_id, agent_extra_headers=agent_extra_headers, + timeout=(max(deadline - time.monotonic(), 0.0) if deadline is not None else None), ) verbose_logger.info("Pydantic AI: Received completed response for request_id=%s", request_id) diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py index ca84d3e07b4..59d6f6003b0 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py @@ -9,6 +9,7 @@ from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig from litellm.a2a_protocol.providers.watsonx_orchestrate.handler import ( WatsonxOrchestrateHandler, ) +from litellm.interactions.agents.utils import merge_agent_headers class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig): @@ -28,10 +29,16 @@ class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig): "litellm_params is required for WatsonxOrchestrateA2AConfig " "(must contain cp4d_host, instance_id, wxo_agent_id, api_key)" ) + forwarded_headers: Final = merge_agent_headers( + dynamic_headers=kwargs.get("agent_extra_headers"), + static_headers=kwargs.get("agent_static_headers"), + ) return await WatsonxOrchestrateHandler.handle_non_streaming( request_id=request_id, params=params, litellm_params=litellm_params, + static_headers=forwarded_headers, + timeout=kwargs.get("timeout"), ) async def handle_streaming( @@ -48,9 +55,15 @@ class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig): "litellm_params is required for WatsonxOrchestrateA2AConfig " "(must contain cp4d_host, instance_id, wxo_agent_id, api_key)" ) + forwarded_headers: Final = merge_agent_headers( + dynamic_headers=kwargs.get("agent_extra_headers"), + static_headers=kwargs.get("agent_static_headers"), + ) async for chunk in WatsonxOrchestrateHandler.handle_streaming( request_id=request_id, params=params, litellm_params=litellm_params, + static_headers=forwarded_headers, + timeout=kwargs.get("timeout"), ): yield chunk diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py index c66b07c321c..f4eb65715fa 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py @@ -6,7 +6,7 @@ import asyncio import hashlib import json import time -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Mapping from typing import Any, Final, NamedTuple, Protocol import httpx @@ -26,9 +26,34 @@ _IBM_CLOUD_IAM_URL: Final = "https://iam.cloud.ibm.com/identity/token" _POLL_INTERVAL_S: Final = 2.0 _MAX_POLL_ATTEMPTS: Final = 90 _TOKEN_CACHE_TTL_BUFFER_S: Final = 60 +_WXO_RESERVED_HEADERS: Final = frozenset({"accept", "authorization", "content-type"}) _token_cache: Final[dict[str, tuple[str, float]]] = {} +def _build_wxo_headers( + token: str, + accept: str, + static_headers: Mapping[str, str] | None = None, +) -> dict[str, str]: + headers: dict[str, str] = ( + { + key: value + for key, value in static_headers.items() + if isinstance(key, str) and isinstance(value, str) and key.lower() not in _WXO_RESERVED_HEADERS + } + if isinstance(static_headers, Mapping) + else {} + ) + headers.update( + { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + "Accept": accept, + } + ) + return headers + + class WXORequestParams(NamedTuple): cp4d_host: str instance_id: str @@ -185,12 +210,25 @@ class WatsonxOrchestrateHandler: client: AsyncHTTPHandler, max_attempts: int = _MAX_POLL_ATTEMPTS, interval_s: float = _POLL_INTERVAL_S, + timeout: float | None = None, ) -> _WXORun: url: Final = f"{base_url}/v1/orchestrate/runs/{run_id}" + deadline: float | None = time.monotonic() + max(timeout, 0) if timeout is not None else None for attempt in range(max_attempts): - await asyncio.sleep(interval_s) - response = await client.get(url, headers=auth_headers) + if deadline is not None: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise asyncio.TimeoutError(f"WXO run '{run_id}' exceeded timeout of {timeout}s") + if interval_s > 0: + await asyncio.sleep(min(interval_s, remaining)) + remaining = deadline - time.monotonic() + if remaining <= 0: + raise asyncio.TimeoutError(f"WXO run '{run_id}' exceeded timeout of {timeout}s") + response = await asyncio.wait_for(client.get(url, headers=auth_headers), timeout=remaining) + else: + await asyncio.sleep(interval_s) + response = await client.get(url, headers=auth_headers) response.raise_for_status() result = WatsonxOrchestrateHandler._run_body(response) status = result.get("status", "") @@ -208,6 +246,7 @@ class WatsonxOrchestrateHandler: base_url: str, auth_headers: dict[str, str], client: AsyncHTTPHandler, + timeout: float | None = None, ) -> _WXORun: status = run_data.get("status", "") if status not in WatsonxOrchestrateTransformation.TERMINAL_STATES: @@ -219,6 +258,7 @@ class WatsonxOrchestrateHandler: run_id=run_id, auth_headers=auth_headers, client=client, + timeout=timeout, ) status = run_data.get("status", "") @@ -228,23 +268,28 @@ class WatsonxOrchestrateHandler: return run_data @staticmethod - async def _accumulate_wxo_sse_text(response: Any) -> str: - source: Final[_WXOView] = {"sse_source": response} - accumulated_text = "" - async for line in source["sse_source"].aiter_lines(): - if not line.startswith("data:"): - continue - data_str = line[5:].strip() - if not data_str or data_str == "[DONE]": - continue - try: - event = WatsonxOrchestrateHandler._decode_run_event(data_str) - except json.JSONDecodeError: - continue - chunk_text = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(event) - if chunk_text: - accumulated_text += chunk_text - return accumulated_text + async def _accumulate_wxo_sse_text(response: Any, timeout: float | None = None) -> str: + async def _collect() -> str: + source: Final[_WXOView] = {"sse_source": response} + accumulated_text = "" + async for line in source["sse_source"].aiter_lines(): + if not line.startswith("data:"): + continue + data_str = line[5:].strip() + if not data_str or data_str == "[DONE]": + continue + try: + event = WatsonxOrchestrateHandler._decode_run_event(data_str) + except json.JSONDecodeError: + continue + chunk_text = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(event) + if chunk_text: + accumulated_text += chunk_text + return accumulated_text + + if timeout is None: + return await _collect() + return await asyncio.wait_for(_collect(), timeout=max(timeout, 0)) @staticmethod def _extract_litellm_params(litellm_params: WXOLitellmParams) -> WXORequestParams: @@ -277,10 +322,12 @@ class WatsonxOrchestrateHandler: request_id: str, params: dict[str, object], litellm_params: WXOLitellmParams, + static_headers: Mapping[str, str] | None = None, + timeout: float | None = None, ) -> dict[str, object]: wxo: Final = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params) - client: Final = WatsonxOrchestrateHandler._http_client(timeout=90.0) + client: Final = WatsonxOrchestrateHandler._http_client(timeout=timeout if timeout is not None else 90.0) token: Final = await WatsonxOrchestrateHandler._get_bearer_token( cp4d_host=wxo.cp4d_host, auth_mode=wxo.auth_mode, @@ -289,11 +336,11 @@ class WatsonxOrchestrateHandler: client=client, ) base_url: Final = WatsonxOrchestrateTransformation.get_api_base_url(wxo.cp4d_host, wxo.instance_id) - auth_headers: Final = { - "Authorization": f"Bearer {token}", - "Content-Type": "application/json", - "Accept": "application/json", - } + auth_headers: Final = _build_wxo_headers( + token=token, + accept="application/json", + static_headers=static_headers, + ) text: Final = WatsonxOrchestrateTransformation.extract_text_from_a2a_params(params) body: Final = WatsonxOrchestrateTransformation.build_wxo_run_body( @@ -314,6 +361,7 @@ class WatsonxOrchestrateHandler: base_url=base_url, auth_headers=auth_headers, client=client, + timeout=timeout, ) response_text: Final = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(run_data) @@ -326,10 +374,12 @@ class WatsonxOrchestrateHandler: litellm_params: WXOLitellmParams, chunk_size: int = 50, delay_ms: int = 10, + static_headers: Mapping[str, str] | None = None, + timeout: float | None = None, ) -> AsyncIterator[dict[str, object]]: wxo: Final = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params) - client: Final = WatsonxOrchestrateHandler._http_client(timeout=120.0) + client: Final = WatsonxOrchestrateHandler._http_client(timeout=timeout if timeout is not None else 120.0) token: Final = await WatsonxOrchestrateHandler._get_bearer_token( cp4d_host=wxo.cp4d_host, auth_mode=wxo.auth_mode, @@ -338,11 +388,11 @@ class WatsonxOrchestrateHandler: client=client, ) base_url: Final = WatsonxOrchestrateTransformation.get_api_base_url(wxo.cp4d_host, wxo.instance_id) - auth_headers: Final = { - "Authorization": f"Bearer {token}", - "Content-Type": "application/json", - "Accept": "text/event-stream, application/json", - } + auth_headers: Final = _build_wxo_headers( + token=token, + accept="text/event-stream, application/json", + static_headers=static_headers, + ) text: Final = WatsonxOrchestrateTransformation.extract_text_from_a2a_params(params) body: Final = WatsonxOrchestrateTransformation.build_wxo_run_body( wxo_agent_id=wxo.wxo_agent_id, text=text, thread_id=wxo.thread_id @@ -366,6 +416,8 @@ class WatsonxOrchestrateHandler: request_id=request_id, params=params, litellm_params=litellm_params, + static_headers=static_headers, + timeout=timeout, ) response_text: Final = WatsonxOrchestrateTransformation.extract_text_from_a2a_message_response(result) async for chunk in WatsonxOrchestrateTransformation.fake_streaming_from_text( @@ -387,10 +439,11 @@ class WatsonxOrchestrateHandler: base_url=base_url, auth_headers=auth_headers, client=client, + timeout=timeout, ) accumulated_text = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(result) else: - accumulated_text = await WatsonxOrchestrateHandler._accumulate_wxo_sse_text(response) + accumulated_text = await WatsonxOrchestrateHandler._accumulate_wxo_sse_text(response, timeout=timeout) async for chunk in WatsonxOrchestrateTransformation.fake_streaming_from_text( text=accumulated_text, diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 1e0b778d244..9b1eb3f8c29 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1226,6 +1226,15 @@ class CustomStreamWrapper: completion_obj: dict[str, Any], ) -> _ProviderChunkResult: response_obj: dict[str, Any] = {} + if isinstance(chunk, ModelResponseStream) and self.custom_llm_provider == "a2a": + model_response = chunk + model_response.model = self.model + finish_reasons = [getattr(choice, "finish_reason", None) for choice in chunk.choices] + if finish_reasons and all(isinstance(reason, str) and reason for reason in finish_reasons): + self.received_finish_reason = finish_reasons[0] + self.sent_last_chunk = True + return _ProviderChunkEarlyReturn(model_response) + if ( isinstance(chunk, ModelResponseStream) and self.custom_llm_provider is not None @@ -1264,7 +1273,13 @@ class CustomStreamWrapper: if not _chunk_has_content and (not isinstance(chunk, dict) or "provider_specific_fields" not in chunk): raise StopIteration anthropic_response_obj: Final[GChunk] = cast(GChunk, chunk) + chunk_id = anthropic_response_obj.get("id") + if isinstance(chunk_id, str) and chunk_id.strip(): + model_response = self.set_model_id(chunk_id, model_response) completion_obj["content"] = anthropic_response_obj["text"] + chunk_index = anthropic_response_obj.get("index") + if isinstance(chunk_index, int): + model_response.choices[0].index = chunk_index if anthropic_response_obj["is_finished"]: self.received_finish_reason = anthropic_response_obj["finish_reason"] @@ -1279,7 +1294,8 @@ class CustomStreamWrapper: ) if "tool_use" in anthropic_response_obj and anthropic_response_obj["tool_use"] is not None: - completion_obj["tool_calls"] = [anthropic_response_obj["tool_use"]] + tool_use = anthropic_response_obj["tool_use"] + completion_obj["tool_calls"] = tool_use if isinstance(tool_use, list) else [tool_use] if ( "provider_specific_fields" in anthropic_response_obj @@ -2611,6 +2627,8 @@ def convert_generic_chunk_to_model_response_stream( ) -> ModelResponseStream: from litellm.types.utils import Delta + tool_use = chunk.get("tool_use", None) + tool_calls = tool_use if isinstance(tool_use, list) else [tool_use] if tool_use is not None else None model_response_stream: Final = ModelResponseStream( id=str(uuid.uuid4()), model="", @@ -2619,7 +2637,7 @@ def convert_generic_chunk_to_model_response_stream( index=chunk.get("index", 0), delta=Delta( content=chunk["text"], - tool_calls=chunk.get("tool_use", None), + tool_calls=tool_calls, ), ) ], diff --git a/litellm/llms/a2a/chat/streaming_iterator.py b/litellm/llms/a2a/chat/streaming_iterator.py index 1983c18a6b3..c4f8d1b2edd 100644 --- a/litellm/llms/a2a/chat/streaming_iterator.py +++ b/litellm/llms/a2a/chat/streaming_iterator.py @@ -2,12 +2,16 @@ A2A Streaming Response Iterator """ -from typing import Final +from collections.abc import Mapping +from typing import Any, Final +from uuid import uuid4 +import litellm from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator -from litellm.types.utils import GenericStreamingChunk, ModelResponseStream +from litellm.types.llms.openai import ChatCompletionToolCallChunk +from litellm.types.utils import Delta, GenericStreamingChunk, ModelResponseStream, StreamingChoices -from ..common_utils import extract_text_from_a2a_response +from ..common_utils import A2AError, extract_text_from_a2a_response class A2AModelResponseIterator(BaseModelResponseIterator): @@ -30,6 +34,7 @@ class A2AModelResponseIterator(BaseModelResponseIterator): json_mode=json_mode, ) self.model = model + self.response_id: str | None = None def chunk_parser(self, chunk: dict) -> GenericStreamingChunk | ModelResponseStream: """ @@ -56,26 +61,161 @@ class A2AModelResponseIterator(BaseModelResponseIterator): } } """ + if "error" in chunk: + error_value: Final = chunk["error"] + error_message: Final = ( + error_value.get("message") + if isinstance(error_value, dict) and isinstance(error_value.get("message"), str) + else str(error_value) + ) + raise A2AError(status_code=500, message=f"A2A error: {error_message}") + try: + if self.response_id is None: + raw_response_id = chunk.get("id") + raw_result = chunk.get("result") + if not isinstance(raw_response_id, str) and isinstance(raw_result, Mapping): + raw_response_id = raw_result.get("id") + self.response_id = ( + raw_response_id + if isinstance(raw_response_id, str) and raw_response_id.strip() + else f"chatcmpl-{uuid4().hex}" + ) # Extract text from A2A response - text: Final = extract_text_from_a2a_response(chunk) + result: Final = chunk.get("result", {}) + chunk_index = 0 + if isinstance(result, Mapping): + artifact = result.get("artifact") + if isinstance(artifact, Mapping) and isinstance(artifact.get("index"), int): + chunk_index = artifact["index"] + choices = result.get("choices") + if isinstance(choices, list) and choices and isinstance(choices[0], Mapping): + raw_index = choices[0].get("index") + if isinstance(raw_index, int): + chunk_index = raw_index + status: Final = result.get("status", {}) if isinstance(result, Mapping) else {} + is_working_status: Final = ( + isinstance(result, Mapping) + and result.get("kind") == "status-update" + and isinstance(status, Mapping) + and status.get("state") == "working" + ) + text: Final = "" if is_working_status else extract_text_from_a2a_response(chunk) + provider_fields: dict[str, object] = {} + provider_fields.update( + { + key: value + for key, value in chunk.items() + if key in {"system_fingerprint", "service_tier"} and value is not None + } + ) + if isinstance(result, Mapping) and not is_working_status: + control_fields = { + "artifacts", + "choices", + "contextId", + "final", + "finish_reason", + "history", + "id", + "kind", + "message", + "parts", + "status", + "taskId", + "tool_calls", + "usage", + } + provider_fields.update( + {key: value for key, value in result.items() if key not in control_fields and value is not None} + ) + choices = result.get("choices") + if isinstance(choices, list) and choices: + first_choice = choices[0] + if isinstance(first_choice, Mapping): + provider_fields.update( + { + key: value + for key, value in first_choice.items() + if key not in {"index", "message", "finish_reason"} and value is not None + } + ) + first_message = first_choice.get("message") + if isinstance(first_message, Mapping): + provider_fields.update( + { + key: value + for key, value in first_message.items() + if key not in {"kind", "role", "parts", "tool_calls"} and value is not None + } + ) # Determine finish reason finish_reason: Final = self._get_finish_reason(chunk) + tool_calls: Final = self._get_tool_calls(chunk) + usage: Final = self._get_usage(chunk) + + if isinstance(result, Mapping): + choices = result.get("choices") + if isinstance(choices, list) and choices: + streaming_choices: list[StreamingChoices] = [] + for choice_position, raw_choice in enumerate(choices): + if not isinstance(raw_choice, Mapping): + continue + raw_index = raw_choice.get("index", choice_position) + choice_index = raw_index if isinstance(raw_index, int) else choice_position + delta_fields: dict[str, Any] = {} + raw_delta = raw_choice.get("delta") + if isinstance(raw_delta, Mapping): + delta_fields.update(raw_delta) + raw_message = raw_choice.get("message") + if isinstance(raw_message, Mapping): + message_text = extract_text_from_a2a_response({"result": {"message": raw_message}}) + if message_text and "content" not in delta_fields: + delta_fields["content"] = message_text + message_tool_calls = raw_message.get("tool_calls") + if message_tool_calls and "tool_calls" not in delta_fields: + delta_fields["tool_calls"] = message_tool_calls + raw_finish_reason = raw_choice.get("finish_reason") + choice_finish_reason = ( + raw_finish_reason + if isinstance(raw_finish_reason, str) and raw_finish_reason + else finish_reason + ) + streaming_choices.append( + StreamingChoices( + index=choice_index, + delta=Delta(**delta_fields), + finish_reason=choice_finish_reason, + logprobs=raw_choice.get("logprobs"), + ) + ) + if streaming_choices: + return ModelResponseStream( + choices=streaming_choices, + id=self.response_id, + usage=usage, + provider_specific_fields=provider_fields or None, + ) # Return generic streaming chunk return GenericStreamingChunk( text=text, - is_finished=bool(finish_reason), - finish_reason=finish_reason or "", - usage=None, - index=0, - tool_use=None, + id=self.response_id, + is_finished=bool(finish_reason or tool_calls), + finish_reason=finish_reason or ("tool_calls" if tool_calls else ""), + usage=usage, + index=chunk_index, + tool_use=tool_calls, + provider_specific_fields=provider_fields or None, ) except Exception: # Return empty chunk on parse error + if self.response_id is None: + self.response_id = f"chatcmpl-{uuid4().hex}" return GenericStreamingChunk( text="", + id=self.response_id, is_finished=False, finish_reason="", usage=None, @@ -83,12 +223,25 @@ class A2AModelResponseIterator(BaseModelResponseIterator): tool_use=None, ) + def _handle_string_chunk(self, str_line: str | dict) -> GenericStreamingChunk | ModelResponseStream: + if isinstance(str_line, dict): + return self.chunk_parser(chunk=str_line) + return super()._handle_string_chunk(str_line=str_line) + def _get_finish_reason(self, chunk: dict) -> str | None: """Extract finish reason from A2A chunk""" result: Final = chunk.get("result", {}) # Check for task completion if isinstance(result, dict): + explicit_finish_reason: Final = result.get("finish_reason") + if isinstance(explicit_finish_reason, str) and explicit_finish_reason: + return explicit_finish_reason + message: Final = result.get("message") + if isinstance(message, dict): + message_finish_reason: Final = message.get("finish_reason") + if isinstance(message_finish_reason, str) and message_finish_reason: + return message_finish_reason status: Final = result.get("status", {}) if isinstance(status, dict): state: Final = status.get("state") @@ -102,3 +255,67 @@ class A2AModelResponseIterator(BaseModelResponseIterator): return "stop" return None + + def _get_usage(self, chunk: dict) -> object | None: + raw_usage: object | None = chunk.get("usage") + result: Final = chunk.get("result", {}) + if raw_usage is None and isinstance(result, dict): + raw_usage = result.get("usage") + if raw_usage is None: + return None + if isinstance(raw_usage, Mapping): + try: + return litellm.Usage(**raw_usage) + except Exception: + return raw_usage + if hasattr(raw_usage, "model_dump"): + try: + return litellm.Usage(**raw_usage.model_dump(exclude_none=True)) + except Exception: + return raw_usage + return raw_usage + + def _get_tool_calls(self, chunk: dict) -> ChatCompletionToolCallChunk | list[ChatCompletionToolCallChunk] | None: + result: Final = chunk.get("result", {}) + if not isinstance(result, dict): + return None + tool_calls = result.get("tool_calls") + if isinstance(tool_calls, list) and tool_calls: + return self._serialize_tool_calls(tool_calls) + message = result.get("message") + if isinstance(message, dict) and isinstance(message.get("tool_calls"), list) and message["tool_calls"]: + return self._serialize_tool_calls(message["tool_calls"]) + return None + + @staticmethod + def _serialize_tool_call(tool_call: object) -> ChatCompletionToolCallChunk | None: + if isinstance(tool_call, dict): + return tool_call + if hasattr(tool_call, "model_dump"): + return tool_call.model_dump(exclude_none=True) + if hasattr(tool_call, "dict"): + return tool_call.dict(exclude_none=True) + return None + + @classmethod + def _serialize_tool_calls( + cls, tool_calls: list[object] + ) -> ChatCompletionToolCallChunk | list[ChatCompletionToolCallChunk] | None: + serialized: Final = [ + tool_call_value + for tool_call in tool_calls + if (tool_call_value := cls._serialize_tool_call(tool_call)) is not None + ] + if len(serialized) == 1: + return serialized[0] + return serialized or None + + async def aclose(self) -> None: + streaming_response = self.streaming_response + self.streaming_response = None + try: + await super().aclose() + finally: + close_stream = getattr(streaming_response, "aclose", None) + if close_stream is not None: + await close_stream() diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index bd02cfdf907..e113be3ce7f 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -850,17 +850,30 @@ async def invoke_agent_a2a( agent_extra_headers=agent_extra_headers, ) + post_call_succeeded = False try: response = await proxy_logging_obj.post_call_success_hook( user_api_key_dict=user_api_key_dict, data=data, response=response, ) + post_call_succeeded = True + except HTTPException as e: + try: + await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=e, + request_data=data, + ) + except Exception: + pass + raise finally: _enqueue_fn: Final = getattr(logging_obj, "_enqueue_deferred_logging", None) if _enqueue_fn is not None: logging_obj._enqueue_deferred_logging = None - _enqueue_fn() + if post_call_succeeded: + _enqueue_fn(response) response_dict: Final[dict[str, Any]] = ( response.model_dump(mode="json", exclude_none=True) diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py index 8a795214750..7a5e06330da 100644 --- a/litellm/proxy/agent_endpoints/a2a_routing.py +++ b/litellm/proxy/agent_endpoints/a2a_routing.py @@ -5,20 +5,507 @@ Handles routing for A2A agents (models with "a2a/" prefix). Looks up agents in the registry and injects their API base URL. """ -from typing import Any, Final +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Final +from uuid import uuid4 from fastapi import HTTPException +from pydantic import TypeAdapter +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_proxy_logger +from litellm.interactions.agents.utils import merge_agent_headers +from litellm.llms.a2a.common_utils import A2AError, convert_messages_to_prompt, extract_text_from_a2a_response from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import Choices, CustomPricingLiteLLMParams, Message, ModelResponse + +if TYPE_CHECKING: + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + +_OBJECT_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) +_HEADERS_ADAPTER: Final = TypeAdapter(dict[str, str]) +_MESSAGES_ADAPTER: Final = TypeAdapter(list[AllMessageValues]) +_FORWARDED_REQUEST_PARAMS: Final = frozenset( + { + "audio", + "frequency_penalty", + "functions", + "function_call", + "guided_json", + "include_server_side_tool_invocations", + "logit_bias", + "logprobs", + "guardrails", + "max_completion_tokens", + "max_tokens", + "modalities", + "n", + "parallel_tool_calls", + "prediction", + "presence_penalty", + "reasoning_effort", + "response_format", + "seed", + "service_tier", + "safety_identifier", + "stop", + "store", + "stream_options", + "temperature", + "thinking", + "timeout", + "tool_choice", + "tools", + "top_logprobs", + "top_p", + "user", + "verbosity", + "web_search_options", + "output_config", + "prompt_cache_key", + } +) +_A2A_PRICING_PARAMS: Final = frozenset({"cost_per_query", "response_cost"}) | frozenset( + CustomPricingLiteLLMParams.model_fields +) + + +def _get_agent_request_headers(data: Mapping[str, object]) -> dict[str, str]: + proxy_request: Final = data.get("proxy_server_request") + raw_headers: object = proxy_request.get("headers") if isinstance(proxy_request, Mapping) else None + if not isinstance(raw_headers, Mapping): + metadata: Final = data.get("metadata") + if not isinstance(metadata, Mapping): + metadata = data.get("litellm_metadata") + raw_headers = metadata.get("headers") if isinstance(metadata, Mapping) else None + return ( + {str(key).lower(): str(value) for key, value in raw_headers.items()} if isinstance(raw_headers, Mapping) else {} + ) + + +class _A2ATextPart(TypedDict): + kind: ReadOnly[str] + text: ReadOnly[str] + + +class _A2AMessage(TypedDict): + role: ReadOnly[str] + parts: ReadOnly[tuple[_A2ATextPart, ...]] + messageId: ReadOnly[str] + contextId: ReadOnly[str | None] + + +class _A2AParams(TypedDict): + message: ReadOnly[_A2AMessage] + messages: ReadOnly[list[AllMessageValues]] + + +async def _route_registered_provider( + data: Mapping[str, object], + model_name: str, + api_base: str | None, + litellm_params: Mapping[str, object], + static_headers: Mapping[str, str] | None, + dynamic_headers: Mapping[str, str] | None = None, +) -> ModelResponse | CustomStreamWrapper: + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2ACompletionBridgeHandler, + ) + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + from litellm.llms.a2a.chat.streaming_iterator import A2AModelResponseIterator + + raw_messages: Final = data.get("messages") + messages: Final = _MESSAGES_ADAPTER.validate_python(raw_messages) + stream: Final = data.get("stream") is True + request_id: Final = str(uuid4()) + raw_session_id: Final = data.get("litellm_session_id") + metadata: Final = data.get("metadata") + session_id: Final = ( + raw_session_id + if isinstance(raw_session_id, str) + else metadata.get("session_id") + if isinstance(metadata, Mapping) and isinstance(metadata.get("session_id"), str) + else None + ) + params: Final[_A2AParams] = { + "message": { + "role": "user", + "parts": ({"kind": "text", "text": convert_messages_to_prompt(messages)},), + "messageId": str(uuid4()), + "contextId": session_id, + }, + "messages": messages, + } + provider_params: Final = { + **_OBJECT_DICT_ADAPTER.validate_python(litellm_params), + **{key: data[key] for key in _FORWARDED_REQUEST_PARAMS if key in data and data[key] is not None}, + } + registered_provider: Final = litellm_params.get("custom_llm_provider") + registered_model: Final = litellm_params.get("model") + native_provider: Final = registered_provider == "pydantic_ai_agents" or ( + registered_provider == "bedrock" and isinstance(registered_model, str) and "agentcore" in registered_model + ) + bridge_params: Final = _OBJECT_DICT_ADAPTER.validate_python( + {"message": params["message"]} if native_provider else params + ) + configured_headers: Final = litellm_params.get("extra_headers") or litellm_params.get("headers") + configured_headers_dict: Final = ( + _HEADERS_ADAPTER.validate_python(configured_headers) if isinstance(configured_headers, dict) else None + ) + agent_extra_headers: Final = merge_agent_headers( + dynamic_headers=merge_agent_headers( + dynamic_headers=dynamic_headers, + static_headers=configured_headers_dict, + ), + static_headers=static_headers, + ) + if agent_extra_headers: + provider_params["extra_headers"] = agent_extra_headers + + logging_obj: Final = data.get("litellm_logging_obj") + if isinstance(logging_obj, Logging): + provider_params["no-log"] = True + provider_model: Final = litellm_params.get("model") + if isinstance(provider_model, str): + logging_obj.model_call_details["model"] = provider_model + logging_obj.model_call_details.setdefault("litellm_params", {})["model"] = provider_model + provider_name: Final = litellm_params.get("custom_llm_provider") + if isinstance(provider_name, str): + logging_obj.model_call_details["custom_llm_provider"] = provider_name + pricing_params = { + key: litellm_params[key] + for key in _A2A_PRICING_PARAMS + if key in litellm_params and litellm_params[key] is not None + } + if pricing_params: + logging_obj.litellm_params.update(pricing_params) + logging_obj.model_call_details["litellm_params"].update(pricing_params) + logging_obj.custom_pricing = True + + if stream: + streaming_response: Final = A2ACompletionBridgeHandler.handle_streaming( + request_id=request_id, + params=bridge_params, + litellm_params=provider_params, + api_base=api_base, + agent_extra_headers=agent_extra_headers, + agent_static_headers=static_headers, + ) + completion_stream: Final = A2AModelResponseIterator( + streaming_response=streaming_response, + sync_stream=False, + model=model_name, + ) + if not isinstance(logging_obj, Logging): + raise TypeError("litellm_logging_obj is required for streaming A2A requests") + return CustomStreamWrapper( + completion_stream=completion_stream, + model=model_name, + custom_llm_provider="a2a", + logging_obj=logging_obj, + stream_options=data.get("stream_options"), + ) + + response: Final = await A2ACompletionBridgeHandler.handle_non_streaming( + request_id=request_id, + params=bridge_params, + litellm_params=provider_params, + api_base=api_base, + agent_extra_headers=agent_extra_headers, + agent_static_headers=static_headers, + ) + error_value: Final = response.get("error") + if isinstance(error_value, dict): + error: Final = _OBJECT_DICT_ADAPTER.validate_python(error_value) + error_message: Final = error.get("message") + raise A2AError( + status_code=500, + message=f"A2A error: {error_message if isinstance(error_message, str) else 'Unknown error'}", + ) + + result: Final = response.get("result") + result_dict: Final = result if isinstance(result, Mapping) else {} + nested_message: Final = result_dict.get("message") + response_message: Final = nested_message if isinstance(nested_message, Mapping) else result_dict + response_choices: Final = response.get("choices") + choice_payloads: Final = response_choices if isinstance(response_choices, list) else result_dict.get("choices") + + def _serialize_value(value: object) -> object: + if hasattr(value, "model_dump"): + return value.model_dump(exclude_none=True) + if hasattr(value, "dict"): + return value.dict(exclude_none=True) + return value + + def _build_message(message_payload: Mapping[str, object], content: str) -> Message: + message_kwargs: dict[str, object] = { + "content": content, + "role": "assistant", + } + raw_tool_calls = message_payload.get("tool_calls") + if isinstance(raw_tool_calls, list): + message_kwargs["tool_calls"] = raw_tool_calls + for field in ( + "audio", + "annotations", + "function_call", + "images", + "provider_specific_fields", + "reasoning_content", + "reasoning_items", + "refusal", + "thinking_blocks", + ): + value = message_payload.get(field) + if value is not None: + message_kwargs[field] = _serialize_value(value) + return Message(**message_kwargs) + + if isinstance(choice_payloads, list) and choice_payloads: + model_choices = [] + for choice_index, choice in enumerate(choice_payloads): + choice_mapping: Mapping[str, object] = choice if isinstance(choice, Mapping) else {} + raw_message = choice_mapping.get("message") + message_payload: Mapping[str, object] = raw_message if isinstance(raw_message, Mapping) else choice_mapping + choice_kwargs: dict[str, object] = { + "finish_reason": ( + choice_mapping.get("finish_reason") + if isinstance(choice_mapping.get("finish_reason"), str) + else message_payload.get("finish_reason") + if isinstance(message_payload.get("finish_reason"), str) + else "stop" + ), + "index": choice_mapping.get("index", choice_index) + if isinstance(choice_mapping.get("index", choice_index), int) + else choice_index, + "message": _build_message( + message_payload, + extract_text_from_a2a_response({"result": message_payload}), + ), + } + raw_logprobs = choice_mapping.get("logprobs", message_payload.get("logprobs")) + if raw_logprobs is not None: + choice_kwargs["logprobs"] = _serialize_value(raw_logprobs) + model_choices.append(Choices(**choice_kwargs)) + else: + tool_calls: Final = response_message.get("tool_calls") + normalized_tool_calls: Final = tool_calls if isinstance(tool_calls, list) else None + finish_reason: Final = response_message.get("finish_reason") + text: Final = extract_text_from_a2a_response(response) + choice_kwargs = { + "finish_reason": ( + finish_reason if isinstance(finish_reason, str) else "tool_calls" if normalized_tool_calls else "stop" + ), + "index": 0, + "message": _build_message(response_message, text), + } + raw_logprobs = response_message.get("logprobs") + if raw_logprobs is not None: + choice_kwargs["logprobs"] = _serialize_value(raw_logprobs) + model_choices = [Choices(**choice_kwargs)] + model_response: Final = ModelResponse( + id=str(response.get("id") or request_id), + model=model_name, + choices=model_choices, + system_fingerprint=response.get("system_fingerprint") + if isinstance(response.get("system_fingerprint"), str) + else None, + service_tier=response.get("service_tier") if isinstance(response.get("service_tier"), str) else None, + ) + raw_usage: Final = response.get("usage") + usage = litellm.Usage(**raw_usage) if isinstance(raw_usage, Mapping) else raw_usage + if usage is None and native_provider: + try: + from litellm.utils import token_counter + + prompt_tokens: Final = token_counter(model="gpt-3.5-turbo", messages=messages) + completion_tokens: Final = token_counter( + model="gpt-3.5-turbo", + text=extract_text_from_a2a_response(response), + count_response_tokens=True, + ) + usage = litellm.Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ) + except Exception: # noqa: BLE001 - token estimation must not fail the response + pass + if usage is not None: + model_response.usage = usage + if isinstance(logging_obj, Logging): + logging_obj.model_call_details["usage"] = usage + + if isinstance(logging_obj, Logging): + + def _enqueue_logging(final_response: ModelResponse | None = None) -> None: + asyncio.create_task( + logging_obj.dispatch_success_handlers( + final_response if final_response is not None else model_response, + cache_hit=False, + prefer_async_handlers=True, + ) + ) + + logging_obj._enqueue_deferred_logging = _enqueue_logging + + return model_response + + +def _merge_agent_guardrails( + data: Mapping[str, object], + agent_guardrails: object, +) -> Mapping[str, object]: + if not agent_guardrails: + return data + + configured_guardrails: list[object] = agent_guardrails if isinstance(agent_guardrails, list) else [agent_guardrails] + metadata_key: Final = "litellm_metadata" if "litellm_metadata" in data else "metadata" + metadata = data.get(metadata_key) + metadata_guardrails = metadata.get("guardrails") if isinstance(metadata, dict) else None + root_guardrails = data.get("guardrails") + existing_guardrails: list[object] = [] + for value in (metadata_guardrails, root_guardrails): + if isinstance(value, list): + existing_guardrails.extend(value) + elif value: + existing_guardrails.append(value) + + merged_guardrails = existing_guardrails + [ + guardrail for guardrail in configured_guardrails if guardrail not in existing_guardrails + ] + if isinstance(data, dict): + data["guardrails"] = merged_guardrails + if isinstance(metadata, dict): + data[metadata_key] = {**metadata, "guardrails": merged_guardrails} + return data + + merged_data = dict(data) + merged_data["guardrails"] = merged_guardrails + if isinstance(metadata, dict): + merged_data[metadata_key] = {**metadata, "guardrails": merged_guardrails} + return merged_data + + +async def merge_a2a_agent_guardrails_before_hooks(data: Mapping[str, object]) -> Mapping[str, object]: + model_name: Final = data.get("model") + if not isinstance(model_name, str) or not model_name.startswith("a2a/"): + return data + + from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through + + agent = await get_agent_with_read_through(model_name[4:]) + if agent is None or not agent.litellm_params: + return data + return _merge_agent_guardrails(data, agent.litellm_params.get("guardrails")) + + +async def authorize_a2a_agent_before_hooks( + data: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth | None, +) -> Mapping[str, object]: + model_name: Final = data.get("model") + if not isinstance(model_name, str) or not model_name.startswith("a2a/"): + return data + + from litellm.proxy.agent_endpoints.auth.agent_permission_handler import AgentRequestHandler + from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through + + agent = await get_agent_with_read_through(model_name[4:]) + if agent is None: + return data + + is_admin: Final = user_api_key_dict is not None and ( + user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN + or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value + ) + if not is_admin: + is_allowed: Final = await AgentRequestHandler.is_agent_allowed( + agent_id=agent.agent_id, + user_api_key_auth=user_api_key_dict, + ) + if not is_allowed: + raise HTTPException( + status_code=403, + detail=f"Agent '{agent.agent_name}' is not allowed for your key/team. Contact proxy admin for access.", + ) + + if (agent.litellm_params or {}).get("require_trace_id_on_calls_to_agent"): + _enforce_inbound_trace_id(data, agent.agent_id) + + if isinstance(data, dict): + data["agent_id"] = agent.agent_id + metadata = data.get("metadata") + if not isinstance(metadata, dict): + metadata = {} + data["metadata"] = metadata + metadata["agent_id"] = agent.agent_id + return data + + +def _get_agent_dynamic_headers( + data: Mapping[str, object], + agent_id: str, + agent_name: str, + extra_headers: list[str] | None, +) -> dict[str, str]: + normalized_headers: Final = _get_agent_request_headers(data) + + dynamic_headers: dict[str, str] = {} + for header_name in extra_headers or []: + header_name_str: Final = str(header_name) + if header_name_str.lower().startswith("x-litellm-"): + continue + value: Final = normalized_headers.get(header_name_str.lower()) + if value is not None: + dynamic_headers[header_name_str] = value + + for alias in (agent_id.lower(), agent_name.lower()): + prefix: Final = f"x-a2a-{alias}-" + for key, value in normalized_headers.items(): + if key.startswith(prefix): + header_name: Final = key[len(prefix) :] + if header_name and not header_name.lower().startswith("x-litellm-"): + dynamic_headers[header_name] = value + return dynamic_headers + + +def _get_agent_identity_headers( + user_api_key_dict: UserAPIKeyAuth | None, + trace_id: object | None = None, +) -> dict[str, str]: + headers: dict[str, str] = {} + if user_api_key_dict is not None and user_api_key_dict.user_id: + headers["X-LiteLLM-User-Id"] = user_api_key_dict.user_id + if user_api_key_dict is not None and user_api_key_dict.team_id: + headers["X-LiteLLM-Team-Id"] = user_api_key_dict.team_id + if trace_id: + headers["X-LiteLLM-Trace-Id"] = str(trace_id) + return headers + + +def _enforce_inbound_trace_id(data: Mapping[str, object], agent_id: str) -> None: + from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers + + if not get_chain_id_from_headers(_get_agent_request_headers(data)): + raise HTTPException( + status_code=400, + detail=f"Agent '{agent_id}' requires x-litellm-trace-id header on all inbound requests.", + ) async def route_a2a_agent_request( - data: dict, + data: Mapping[str, object], route_type: str, user_api_key_dict: UserAPIKeyAuth | None = None, -) -> Any | None: +) -> Awaitable[object] | None: """ Route A2A agent requests directly to litellm with injected API base. @@ -68,14 +555,75 @@ async def route_a2a_agent_request( detail=f"Agent '{agent_name}' is not allowed for your key/team. Contact proxy admin for access.", ) + if (agent.litellm_params or {}).get("require_trace_id_on_calls_to_agent"): + _enforce_inbound_trace_id(data, agent.agent_id) + # Get API base URL from agent config - if not agent.agent_card_params or "url" not in agent.agent_card_params: + agent_card_params: Final = agent.agent_card_params + agent_url: Final = agent_card_params.get("url") if agent_card_params else None + registered_params_value: Final = agent.litellm_params + registered_provider_value: Final = ( + registered_params_value.get("custom_llm_provider") if registered_params_value else None + ) + registered_provider: Final = registered_provider_value if isinstance(registered_provider_value, str) else None + from litellm.a2a_protocol.litellm_completion_bridge.handler import A2A_USER_API_KEY_HASH_PARAM + + registered_params_for_route: Final[Mapping[str, object]] = ( + { + **registered_params_value, + A2A_USER_API_KEY_HASH_PARAM: user_api_key_dict.api_key, + } + if registered_params_value and user_api_key_dict is not None and user_api_key_dict.api_key + else registered_params_value or {} + ) + configured_api_base: Final = registered_params_value.get("api_base") if registered_params_value else None + api_base: Final = configured_api_base if isinstance(configured_api_base, str) and configured_api_base else agent_url + cardless_provider: Final = registered_provider is not None and registered_provider != "a2a" + has_configured_api_base: Final = isinstance(configured_api_base, str) and bool(configured_api_base) + if (not isinstance(agent_url, str) or not agent_url) and not has_configured_api_base and not cardless_provider: verbose_proxy_logger.error("[A2A] Agent '%s' has no URL configured", agent_name) route_name = ROUTE_ENDPOINT_MAPPING.get(route_type, route_type) raise ProxyModelNotFoundError(route=route_name, model_name=model_name, retryable_with_model_read_through=False) - # Inject API base and route to litellm - data["api_base"] = agent.agent_card_params["url"] - verbose_proxy_logger.debug("[A2A] Routing %s to %s", model_name, data["api_base"]) + routed_data: Final = _merge_agent_guardrails( + data=data, + agent_guardrails=registered_params_value.get("guardrails") if registered_params_value else None, + ) + registered_dynamic_headers: Final = _get_agent_dynamic_headers( + data=routed_data, + agent_id=agent.agent_id, + agent_name=agent.agent_name, + extra_headers=agent.extra_headers, + ) + registered_static_headers: Mapping[str, str] | None = agent.static_headers + if registered_params_value and registered_params_value.get("databricks_oauth"): + from litellm.proxy.agent_endpoints.databricks_oauth import resolve_databricks_app_auth_header - return getattr(litellm, f"{route_type}")(**data) + databricks_headers = await resolve_databricks_app_auth_header(dict(registered_params_value)) + registered_static_headers = merge_agent_headers( + dynamic_headers=registered_static_headers, + static_headers=databricks_headers, + ) + registered_static_headers = merge_agent_headers( + dynamic_headers=registered_static_headers, + static_headers=_get_agent_identity_headers(user_api_key_dict, data.get("litellm_trace_id")), + ) + if ( + registered_provider + and registered_provider != "a2a" + and route_type == "acompletion" + and registered_params_value is not None + ): + verbose_proxy_logger.debug("[A2A] Routing %s through %s", model_name, registered_provider) + return _route_registered_provider( + data=routed_data, + model_name=model_name, + api_base=api_base, + litellm_params=registered_params_for_route, + static_headers=registered_static_headers, + dynamic_headers=registered_dynamic_headers, + ) + + completion_data: Final = MappingProxyType({**routed_data, "api_base": api_base}) + verbose_proxy_logger.debug("[A2A] Routing %s to %s", model_name, api_base) + return getattr(litellm, f"{route_type}")(**completion_data) # pyright: ignore[reportAny] # dynamic SDK route diff --git a/litellm/proxy/agent_endpoints/model_list_helpers.py b/litellm/proxy/agent_endpoints/model_list_helpers.py index 4a88644bae5..3fa62d0275c 100644 --- a/litellm/proxy/agent_endpoints/model_list_helpers.py +++ b/litellm/proxy/agent_endpoints/model_list_helpers.py @@ -4,6 +4,8 @@ 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 Final + from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth from litellm.types.proxy.management_endpoints.model_management_endpoints import ( @@ -33,11 +35,14 @@ async def append_agents_to_model_group( for agent_id in allowed_agent_ids: agent = global_agent_registry.get_agent_by_id(agent_id) if agent is not None: + agent_params: Final = agent.litellm_params + provider_value: Final = agent_params.get("custom_llm_provider") if agent_params else None + custom_llm_provider: Final = provider_value if isinstance(provider_value, str) else "a2a" model_groups.append( ModelGroupInfoProxy( model_group=f"a2a/{agent.agent_name}", mode="chat", - providers=["a2a"], + providers=[custom_llm_provider], ) ) case _: @@ -70,12 +75,15 @@ async def append_agents_to_model_info( for agent_id in allowed_agent_ids: agent = global_agent_registry.get_agent_by_id(agent_id) if agent is not None: + agent_params: Final = agent.litellm_params + provider_value: Final = agent_params.get("custom_llm_provider") if agent_params else None + custom_llm_provider: Final = provider_value if isinstance(provider_value, str) else "a2a" models.append( { "model_name": f"a2a/{agent.agent_name}", "litellm_params": { "model": f"a2a/{agent.agent_name}", - "custom_llm_provider": "a2a", + "custom_llm_provider": custom_llm_provider, }, "model_info": { "id": agent.agent_id, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index ff6c8d1b1f8..9740d8cc71f 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1876,6 +1876,11 @@ class ProxyBaseLLMRequestProcessing: ## LOGGING OBJECT ## - initialize logging object for logging success/failure events for call ## IMPORTANT Note: - initialize this before running pre-call checks. Ensures we log rejected requests to langfuse. + from litellm.proxy.agent_endpoints.a2a_routing import ( + authorize_a2a_agent_before_hooks, + merge_a2a_agent_guardrails_before_hooks, + ) + logging_obj, self.data = litellm.utils.function_setup( original_function=route_type, rules_obj=litellm.utils.Rules(), @@ -1885,6 +1890,13 @@ class ProxyBaseLLMRequestProcessing: self.data["litellm_logging_obj"] = logging_obj + self.data = await authorize_a2a_agent_before_hooks( + data=self.data, + user_api_key_dict=user_api_key_dict, + ) + + self.data = await merge_a2a_agent_guardrails_before_hooks(self.data) + # Merge model-level guardrails before pre_call_hook so DB/UI-configured # guardrails actually execute on pre_call. Without this, guardrails set # via litellm_params.guardrails are only honored on post_call paths @@ -1900,12 +1912,32 @@ class ProxyBaseLLMRequestProcessing: trust_client_model_info=False, ) + authorized_model = self.data.get("model") self.data = await proxy_logging_obj.pre_call_hook( user_api_key_dict=user_api_key_dict, data=self.data, call_type=route_type, ) + if self.data.get("model") != authorized_model: + self.data = await authorize_a2a_agent_before_hooks( + data=self.data, + user_api_key_dict=user_api_key_dict, + ) + self.data = await merge_a2a_agent_guardrails_before_hooks(self.data) + self.data = _check_and_merge_model_level_guardrails( + data=self.data, + llm_router=llm_router, + trust_client_model_info=False, + ) + if isinstance(self.data.get("model"), str) and self.data["model"].startswith("a2a/"): + self.data = await proxy_logging_obj.pre_call_hook( + user_api_key_dict=user_api_key_dict, + data=self.data, + call_type=route_type, + guardrails_only=True, + ) + # Refresh AFTER pre_call_hook: guardrails (e.g. Presidio PII masking) may # have mutated `self.data` in place, and the audit-trail snapshot taken in # add_litellm_data_to_request predates that mutation. @@ -2539,6 +2571,7 @@ class ProxyBaseLLMRequestProcessing: ProxyBaseLLMRequestProcessing._flush_deferred_async_logging( logging_obj=logging_obj, exception_raised=_exception_raised, + response=response, ) # Streaming cleanup: if an exception occurred AND the deferred @@ -3023,6 +3056,7 @@ class ProxyBaseLLMRequestProcessing: def _flush_deferred_async_logging( logging_obj: Any, exception_raised: bool, + response: Any | None = None, ) -> None: """ Fire the deferred async-success closure stored by wrapper_async, then @@ -3053,7 +3087,7 @@ class ProxyBaseLLMRequestProcessing: if exception_raised: return try: - _enqueue_fn() + _enqueue_fn(response) if response is not None else _enqueue_fn() except Exception as e: verbose_proxy_logger.exception("Error firing deferred logging: %s", e) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index f0319a7c664..e9bc8b916e3 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -323,7 +323,8 @@ class ModelInfo(ModelInfoBase, total=False): class GenericStreamingChunk(TypedDict, total=False): text: Required[str] - tool_use: ChatCompletionToolCallChunk | None + id: str + tool_use: ChatCompletionToolCallChunk | list[ChatCompletionToolCallChunk] | None is_finished: Required[bool] finish_reason: Required[str] usage: Required[ChatCompletionUsageBlock | None] diff --git a/litellm/utils.py b/litellm/utils.py index 5cd9bfc5f32..7782f535f69 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1928,11 +1928,11 @@ def client(original_function): if not _is_litellm_internal_call: if getattr(logging_obj, "_defer_async_logging", False): - def _enqueue_deferred_logging() -> None: + def _enqueue_deferred_logging(final_response=None) -> None: asyncio.create_task( _client_async_logging_helper( logging_obj=logging_obj, - result=result, + result=final_response if final_response is not None else result, start_time=start_time, end_time=end_time, is_completion_with_fallbacks=is_completion_with_fallbacks, diff --git a/tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py b/tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py index 5503a5668bf..a959386aae9 100644 --- a/tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py +++ b/tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py @@ -10,10 +10,9 @@ Verifies that: """ import json - -import pytest from unittest.mock import AsyncMock, MagicMock, patch +import pytest SAMPLE_ARN = "arn:aws:bedrock-agentcore:us-west-2:123456789:runtime/my_agent" SAMPLE_MODEL = f"bedrock/agentcore/{SAMPLE_ARN}" @@ -482,6 +481,7 @@ class TestHandlerIntegration: api_base=None, litellm_params=SAMPLE_LITELLM_PARAMS, agent_extra_headers=None, + agent_static_headers=None, ) @pytest.mark.asyncio diff --git a/tests/test_litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py b/tests/test_litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py index c5626afa954..8b40b64987f 100644 --- a/tests/test_litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py +++ b/tests/test_litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py @@ -4,10 +4,10 @@ Tests for Pydantic AI agents transformation. Tests the helper functions and response transformation without making real API calls. """ +from unittest.mock import AsyncMock, MagicMock, patch import pytest - from litellm.a2a_protocol.providers.pydantic_ai_agents.transformation import ( PydanticAITransformation, ) @@ -94,3 +94,31 @@ class TestPydanticAITransformation: assert result["result"]["kind"] == "message" assert result["result"]["role"] == "agent" assert result["result"]["parts"][0]["text"] == "The answer is 4." + + +@pytest.mark.asyncio +async def test_poll_for_completion_honors_request_timeout(): + client = MagicMock() + response = MagicMock() + response.json.return_value = { + "result": {"status": {"state": "working"}, "id": "task-1"} + } + response.raise_for_status.return_value = None + client.post = AsyncMock(return_value=response) + + with patch( + "litellm.a2a_protocol.providers.pydantic_ai_agents.transformation.time.monotonic", + side_effect=[100.0, 100.0, 100.02], + ): + with pytest.raises(TimeoutError, match=r"0\.01 seconds"): + await PydanticAITransformation._poll_for_completion( + client=client, + endpoint="http://example.test", + task_id="task-1", + request_id="req-1", + poll_interval=1.0, + timeout=0.01, + ) + + client.post.assert_awaited_once() + assert client.post.await_args.kwargs["timeout"] == pytest.approx(0.01) diff --git a/tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py b/tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py index 43dfdaba02d..7dadfb9c063 100644 --- a/tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py +++ b/tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py @@ -6,9 +6,11 @@ from pathlib import Path import httpx import pytest - from litellm.a2a_protocol.providers.config_manager import A2AProviderConfigManager from litellm.a2a_protocol.providers.watsonx_orchestrate import handler as wxo_handler +from litellm.a2a_protocol.providers.watsonx_orchestrate.config import ( + WatsonxOrchestrateA2AConfig, +) from litellm.a2a_protocol.providers.watsonx_orchestrate.handler import ( WatsonxOrchestrateHandler, ) @@ -46,6 +48,12 @@ class _SSELines: yield line +class _HangingSSELines: + async def aiter_lines(self): + await asyncio.Event().wait() + yield "" + + class _InvalidJsonStreamResponse: headers = {"content-type": "application/json"} @@ -231,6 +239,12 @@ async def test_accumulate_wxo_sse_text_ignores_non_dict_json_events(): ) +@pytest.mark.asyncio +async def test_accumulate_wxo_sse_text_respects_timeout(): + with pytest.raises(asyncio.TimeoutError): + await WatsonxOrchestrateHandler._accumulate_wxo_sse_text(_HangingSSELines(), timeout=0.001) + + @pytest.mark.asyncio async def test_short_lived_tokens_are_not_served_from_cache(): client = _ShortTtlTokenClient() @@ -343,6 +357,50 @@ async def test_poll_run_raises_asyncio_timeout_when_never_terminal(): assert client.get_calls == 2 +def test_build_wxo_headers_preserves_auth_headers(): + headers = wxo_handler._build_wxo_headers( + token="token", + accept="application/json", + static_headers={ + "x-tenant-id": "tenant-1", + "Authorization": "caller-token", + "content-type": "text/plain", + }, + ) + + assert headers["x-tenant-id"] == "tenant-1" + assert headers["Authorization"] == "Bearer token" + assert headers["Content-Type"] == "application/json" + assert headers["Accept"] == "application/json" + assert "content-type" not in headers + + +@pytest.mark.asyncio +async def test_wxo_config_forwards_headers_and_timeout(monkeypatch): + captured = {} + + async def fake_handle_non_streaming(**kwargs): + captured.update(kwargs) + return {"result": {}} + + monkeypatch.setattr(WatsonxOrchestrateHandler, "handle_non_streaming", fake_handle_non_streaming) + + await WatsonxOrchestrateA2AConfig().handle_non_streaming( + request_id="req-1", + params={}, + litellm_params={"model": "agent"}, + agent_extra_headers={"x-request-id": "request-1"}, + agent_static_headers={"x-tenant-id": "tenant-1"}, + timeout=12, + ) + + assert captured["static_headers"] == { + "x-request-id": "request-1", + "x-tenant-id": "tenant-1", + } + assert captured["timeout"] == 12 + + @pytest.mark.asyncio async def test_handle_streaming_polls_non_sse_json_until_complete(monkeypatch): client = _JsonStreamClient({"status": "running", "run_id": "run-1"}) diff --git a/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py b/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py index 1b3e5f86020..7314ec43b46 100644 --- a/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py +++ b/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py @@ -12,6 +12,8 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from litellm.types.utils import Choices, Message, ModelResponse + class TestA2AStreamingTransformation: """Test the A2A streaming transformation creates proper events.""" @@ -26,9 +28,7 @@ class TestA2AStreamingTransformation: "parts": [{"text": "Reply to ticket #4823"}], "metadata": {"skillId": "draft_reply"}, } - openai_messages = ( - A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message) - ) + openai_messages = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message) # Metadata is forwarded on the run payload only, not duplicated on messages. assert "metadata" not in openai_messages[0] @@ -174,10 +174,7 @@ class TestA2AStreamingTransformation: assert "artifactId" in event["result"]["artifact"] assert event["result"]["artifact"]["name"] == "response" assert event["result"]["artifact"]["parts"][0]["kind"] == "text" - assert ( - event["result"]["artifact"]["parts"][0]["text"] - == "Hello, I am an AI assistant." - ) + assert event["result"]["artifact"]["parts"][0]["text"] == "Hello, I am an AI assistant." @pytest.mark.asyncio @@ -197,6 +194,8 @@ async def test_handle_streaming_emits_proper_events(): mock_chunk2.choices = [MagicMock()] mock_chunk2.choices[0].delta = MagicMock() mock_chunk2.choices[0].delta.content = " world" + mock_chunk2.choices[0].finish_reason = "length" + mock_chunk2.usage = {"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5} async def mock_streaming_response(): yield mock_chunk1 @@ -222,8 +221,8 @@ async def test_handle_streaming_emits_proper_events(): ): events.append(event) - # Should have 4 events: task, working, artifact, completed - assert len(events) == 4 + # Should have 5 events: task, working, two artifacts, completed + assert len(events) == 5 # Event 1: task submitted assert events[0]["result"]["kind"] == "task" @@ -234,14 +233,299 @@ async def test_handle_streaming_emits_proper_events(): assert events[1]["result"]["status"]["state"] == "working" assert events[1]["result"]["final"] is False - # Event 3: artifact update with accumulated content + # Event 3: first artifact update assert events[2]["result"]["kind"] == "artifact-update" - assert events[2]["result"]["artifact"]["parts"][0]["text"] == "Hello world" + assert events[2]["result"]["artifact"]["parts"][0]["text"] == "Hello" - # Event 4: status completed - assert events[3]["result"]["kind"] == "status-update" - assert events[3]["result"]["status"]["state"] == "completed" - assert events[3]["result"]["final"] is True + # Event 4: second artifact update + assert events[3]["result"]["kind"] == "artifact-update" + assert events[3]["result"]["artifact"]["parts"][0]["text"] == " world" + assert ( + events[2]["result"]["artifact"]["artifactId"] + == events[3]["result"]["artifact"]["artifactId"] + ) + + # Event 5: status completed + assert events[4]["result"]["kind"] == "status-update" + assert events[4]["result"]["status"]["state"] == "completed" + assert events[4]["result"]["final"] is True + assert events[4]["result"]["finish_reason"] == "length" + assert events[4]["usage"]["total_tokens"] == 5 + + +def test_build_completion_params_keeps_bridge_routing_fields(): + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2ACompletionBridgeHandler, + ) + + params = A2ACompletionBridgeHandler._build_completion_params( + params={"message": {"role": "user", "parts": []}}, + litellm_params={ + "custom_llm_provider": "openai", + "model": "agent", + "api_base": "https://untrusted.example", + "stream": False, + }, + api_base="https://configured.example", + agent_extra_headers=None, + stream=True, + ) + + assert params["api_base"] == "https://configured.example" + assert params["stream"] is True + + +def test_build_completion_params_drops_proxy_only_databricks_oauth(): + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2ACompletionBridgeHandler, + ) + + params = A2ACompletionBridgeHandler._build_completion_params( + params={"message": {"role": "user", "parts": []}}, + litellm_params={ + "custom_llm_provider": "databricks", + "model": "agent", + "databricks_oauth": {"client_id": "id"}, + }, + api_base="https://configured.example", + agent_extra_headers=None, + stream=False, + ) + + assert "databricks_oauth" not in params + + +@pytest.mark.asyncio +async def test_handle_streaming_accumulates_logprobs_and_provider_metadata(): + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2ACompletionBridgeHandler, + ) + + chunks = [] + for token in ("a", "b"): + choice = MagicMock() + choice.index = 0 + choice.finish_reason = None + choice.delta.content = token + choice.logprobs = {"content": [{"token": token}]} + chunk = MagicMock() + chunk.choices = [choice] + chunk.system_fingerprint = "fp-1" + chunk.service_tier = "scale" + chunks.append(chunk) + chunks[-1].choices[0].finish_reason = "stop" + + async def mock_streaming_response(): + for chunk in chunks: + yield chunk + + with patch("litellm.acompletion", new_callable=AsyncMock) as mock_acompletion: + mock_acompletion.return_value = mock_streaming_response() + events = [ + event + async for event in A2ACompletionBridgeHandler.handle_streaming( + request_id="req-metadata", + params={"message": {"role": "user", "parts": []}}, + litellm_params={"custom_llm_provider": "openai", "model": "agent"}, + ) + ] + + result = events[-1] + assert result["system_fingerprint"] == "fp-1" + assert result["service_tier"] == "scale" + assert result["result"]["choices"][0]["logprobs"]["content"] == [ + {"token": "a"}, + {"token": "b"}, + ] + + +@pytest.mark.asyncio +async def test_handle_streaming_preserves_multiple_choices(): + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2ACompletionBridgeHandler, + ) + + mock_chunk = MagicMock() + first_choice = MagicMock() + first_choice.index = 0 + first_choice.finish_reason = None + first_choice.delta.content = "first" + second_choice = MagicMock() + second_choice.index = 1 + second_choice.finish_reason = "length" + second_choice.delta.content = "second" + mock_chunk.choices = [first_choice, second_choice] + + 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() + events = [ + event + async for event in A2ACompletionBridgeHandler.handle_streaming( + request_id="req-choices", + params={"message": {"role": "user", "parts": []}}, + litellm_params={"custom_llm_provider": "langgraph", "model": "agent", "n": 2}, + ) + ] + + choices = events[-1]["result"]["choices"] + assert [choice["index"] for choice in choices] == [0, 1] + assert [choice["message"]["parts"][0]["text"] for choice in choices] == ["", ""] + assert choices[1]["finish_reason"] == "length" + + +@pytest.mark.asyncio +async def test_handle_streaming_preserves_non_text_delta_fields(): + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2ACompletionBridgeHandler, + ) + + delta = MagicMock() + delta.content = "" + delta.tool_calls = None + delta.model_dump.return_value = { + "audio": {"data": "abc"}, + "reasoning_content": "thinking", + "provider_specific_fields": {"trace_id": "trace-1"}, + } + choice = MagicMock() + choice.index = 0 + choice.finish_reason = "stop" + choice.delta = delta + choice.logprobs = {"content": []} + chunk = MagicMock() + chunk.choices = [choice] + + async def mock_streaming_response(): + yield chunk + + with patch("litellm.acompletion", new_callable=AsyncMock) as mock_acompletion: + mock_acompletion.return_value = mock_streaming_response() + events = [ + event + async for event in A2ACompletionBridgeHandler.handle_streaming( + request_id="req-fields", + params={"message": {"role": "user", "parts": []}}, + litellm_params={"custom_llm_provider": "langgraph", "model": "agent"}, + ) + ] + + result = events[-1]["result"] + choice_result = result["choices"][0] + assert choice_result["delta"] == { + "audio": {"data": "abc"}, + "reasoning_content": "thinking", + "provider_specific_fields": {"trace_id": "trace-1"}, + } + assert choice_result["logprobs"] == {"content": []} + + +@pytest.mark.asyncio +async def test_provider_config_receives_full_message_history(): + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2ACompletionBridgeHandler, + ) + + provider_config = MagicMock() + provider_config.handle_non_streaming = AsyncMock(return_value={"result": {}}) + messages = [ + {"role": "system", "content": "Be concise"}, + {"role": "user", "content": "Hello"}, + ] + params = { + "message": {"role": "user", "parts": []}, + "messages": messages, + } + + with patch( + "litellm.a2a_protocol.litellm_completion_bridge.handler.A2AProviderConfigManager.get_provider_config", + return_value=provider_config, + ): + await A2ACompletionBridgeHandler.handle_non_streaming( + request_id="req-1", + params=params, + litellm_params={"custom_llm_provider": "langflow", "model": "flow"}, + ) + + assert provider_config.handle_non_streaming.await_args.kwargs["params"]["messages"] == messages + + +@pytest.mark.asyncio +async def test_native_provider_config_drops_internal_message_history(): + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2ACompletionBridgeHandler, + ) + + provider_config = MagicMock() + provider_config.handle_non_streaming = AsyncMock(return_value={"result": {}}) + params = { + "message": {"role": "user", "parts": []}, + "messages": [{"role": "user", "content": "Hello"}], + } + + with patch( + "litellm.a2a_protocol.litellm_completion_bridge.handler.A2AProviderConfigManager.get_provider_config", + return_value=provider_config, + ): + await A2ACompletionBridgeHandler.handle_non_streaming( + request_id="req-native", + params=params, + litellm_params={"custom_llm_provider": "pydantic_ai_agents", "model": "agent"}, + ) + + assert provider_config.handle_non_streaming.await_args.kwargs["params"] == { + "message": params["message"] + } + + +def test_response_transform_preserves_audio_and_logprobs(): + from litellm.a2a_protocol.litellm_completion_bridge.transformation import ( + A2ACompletionBridgeTransformation, + ) + + response = ModelResponse( + id="resp-1", + model="test-model", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message( + content="hello", + role="assistant", + audio={"data": "abc", "expires_at": 1, "transcript": "hello"}, + ), + logprobs={"content": []}, + ) + ], + ) + + transformed = A2ACompletionBridgeTransformation.openai_response_to_a2a_response(response) + + assert transformed["result"]["audio"]["data"] == "abc" + assert transformed["result"]["logprobs"] == {"content": []} + + +def test_response_transform_preserves_refusal(): + from litellm.a2a_protocol.litellm_completion_bridge.transformation import ( + A2ACompletionBridgeTransformation, + ) + + message = MagicMock() + message.model_dump.return_value = { + "content": None, + "refusal": "I cannot help with that request.", + } + choice = MagicMock(message=message) + choice.model_dump.return_value = {"finish_reason": "stop"} + response = MagicMock(choices=[choice], usage=None) + + transformed = A2ACompletionBridgeTransformation.openai_response_to_a2a_response(response) + + assert transformed["result"]["refusal"] == "I cannot help with that request." + assert transformed["result"]["parts"] == [{"kind": "text", "text": ""}] @pytest.mark.asyncio diff --git a/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py b/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py new file mode 100644 index 00000000000..1643604ad84 --- /dev/null +++ b/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py @@ -0,0 +1,223 @@ +"""Tests for the A2A chat streaming iterator.""" + +import pytest + +from litellm.llms.a2a.chat.streaming_iterator import A2AModelResponseIterator +from litellm.llms.a2a.common_utils import A2AError +from litellm.types.utils import Delta + + +@pytest.mark.asyncio +async def test_async_iterator_accepts_decoded_a2a_events(): + async def _events(): + yield { + "jsonrpc": "2.0", + "result": { + "kind": "artifact-update", + "artifact": {"parts": [{"kind": "text", "text": "Hello"}]}, + }, + } + + iterator = A2AModelResponseIterator( + streaming_response=_events(), + sync_stream=False, + ) + + chunk = await iterator.__aiter__().__anext__() + + assert chunk["text"] == "Hello" + + +def test_chunk_parser_reuses_response_id_for_idless_artifacts(): + iterator = A2AModelResponseIterator(streaming_response=[], sync_stream=False) + first = iterator.chunk_parser( + {"result": {"kind": "artifact-update", "artifact": {"parts": [{"kind": "text", "text": "one"}]}}} + ) + second = iterator.chunk_parser( + {"result": {"kind": "artifact-update", "artifact": {"parts": [{"kind": "text", "text": "two"}]}}} + ) + + assert first["id"] == second["id"] + + +@pytest.mark.asyncio +async def test_async_iterator_ignores_status_message_text(): + async def _events(): + yield { + "jsonrpc": "2.0", + "result": { + "kind": "status-update", + "status": { + "state": "working", + "message": {"parts": [{"kind": "text", "text": "Processing request..."}]}, + }, + }, + } + + iterator = A2AModelResponseIterator(streaming_response=_events(), sync_stream=False) + + chunk = await iterator.__aiter__().__anext__() + + assert chunk["text"] == "" + + +@pytest.mark.asyncio +async def test_async_iterator_preserves_non_text_fields(): + async def _events(): + yield { + "jsonrpc": "2.0", + "result": { + "kind": "status-update", + "status": {"state": "completed"}, + "audio": {"data": "abc"}, + "reasoning_content": "thinking", + "logprobs": {"content": []}, + }, + } + + iterator = A2AModelResponseIterator(streaming_response=_events(), sync_stream=False) + + chunk = await iterator.__aiter__().__anext__() + + assert chunk["provider_specific_fields"] == { + "audio": {"data": "abc"}, + "reasoning_content": "thinking", + "logprobs": {"content": []}, + } + + +@pytest.mark.asyncio +async def test_async_iterator_preserves_tool_calls(): + tool_calls = [ + { + "id": "call-1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ] + + async def _events(): + yield {"jsonrpc": "2.0", "result": {"tool_calls": tool_calls}} + + iterator = A2AModelResponseIterator(streaming_response=_events(), sync_stream=False) + + chunk = await iterator.__aiter__().__anext__() + + assert chunk["tool_use"] == tool_calls[0] + assert chunk["finish_reason"] == "tool_calls" + + +@pytest.mark.asyncio +async def test_async_iterator_preserves_parallel_tool_calls(): + tool_calls = [ + { + "id": "call-1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + }, + { + "id": "call-2", + "type": "function", + "function": {"name": "write", "arguments": "{}"}, + }, + ] + + async def _events(): + yield {"jsonrpc": "2.0", "result": {"tool_calls": tool_calls}} + + iterator = A2AModelResponseIterator(streaming_response=_events(), sync_stream=False) + + chunk = await iterator.__aiter__().__anext__() + + assert chunk["tool_use"] == tool_calls + + +@pytest.mark.asyncio +async def test_async_iterator_preserves_every_terminal_choice(): + async def _events(): + yield { + "jsonrpc": "2.0", + "result": { + "kind": "status-update", + "status": {"state": "completed"}, + "choices": [ + { + "index": 0, + "message": {"parts": [{"kind": "text", "text": "first"}]}, + "finish_reason": "stop", + }, + { + "index": 1, + "message": {"parts": [{"kind": "text", "text": "second"}]}, + "finish_reason": "length", + }, + ], + }, + } + + iterator = A2AModelResponseIterator(streaming_response=_events(), sync_stream=False) + chunk = await iterator.__aiter__().__anext__() + + assert [choice.index for choice in chunk.choices] == [0, 1] + assert [choice.delta.content for choice in chunk.choices] == ["first", "second"] + assert [choice.finish_reason for choice in chunk.choices] == ["stop", "length"] + + +@pytest.mark.asyncio +async def test_async_iterator_serializes_delta_tool_calls_and_usage(): + delta = Delta( + tool_calls=[ + { + "id": "call-1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ] + ) + + async def _events(): + yield { + "jsonrpc": "2.0", + "usage": {"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5}, + "result": { + "tool_calls": [delta.tool_calls[0]], + "finish_reason": "length", + }, + } + + iterator = A2AModelResponseIterator(streaming_response=_events(), sync_stream=False) + chunk = await iterator.__aiter__().__anext__() + + assert chunk["tool_use"]["id"] == "call-1" + assert chunk["finish_reason"] == "length" + assert chunk["usage"].total_tokens == 5 + + +@pytest.mark.asyncio +async def test_async_iterator_closes_nested_stream(): + closed = False + + async def _events(): + nonlocal closed + try: + yield {"jsonrpc": "2.0", "result": {"kind": "artifact-update"}} + raise AssertionError("stream should be closed before a second event") + finally: + closed = True + + iterator = A2AModelResponseIterator(streaming_response=_events(), sync_stream=False) + await iterator.__aiter__().__anext__() + await iterator.aclose() + + assert closed is True + + +@pytest.mark.asyncio +async def test_async_iterator_propagates_jsonrpc_errors(): + async def _events(): + yield {"jsonrpc": "2.0", "error": {"code": -32000, "message": "agent failed"}} + + iterator = A2AModelResponseIterator(streaming_response=_events(), sync_stream=False) + + with pytest.raises(A2AError, match="agent failed"): + await iterator.__aiter__().__anext__() diff --git a/tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py b/tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py index 939ab1cab40..6ce22ed1a3f 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py @@ -4,8 +4,6 @@ Test appending A2A agents to model lists. Maps to: litellm/proxy/agent_endpoints/model_list_helpers.py """ - - from unittest.mock import AsyncMock, Mock, patch import pytest @@ -66,6 +64,32 @@ async def test_append_agents_to_model_group(): assert result[0].providers == ["a2a"] +@pytest.mark.asyncio +async def test_append_agents_to_model_group_preserves_registered_provider(): + agent = AgentResponse( + agent_id="agent-123", + agent_name="test-agent", + agent_card_params={"url": "http://example.com"}, + litellm_params={"custom_llm_provider": "pydantic_ai_agents"}, + ) + registry = Mock() + registry.get_agent_by_id = Mock(return_value=agent) + + with ( + patch( + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access", + AsyncMock(return_value=RestrictedAgentAccess(frozenset({"agent-123"}))), + ), + patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", registry), + ): + result = await append_agents_to_model_group( + model_groups=[], + user_api_key_dict=Mock(spec=UserAPIKeyAuth), + ) + + assert result[0].providers == ["pydantic_ai_agents"] + + @pytest.mark.asyncio async def test_append_agents_to_model_info(): """Test agents are converted to model info format with a2a/ prefix""" @@ -109,3 +133,32 @@ async def test_append_agents_to_model_info(): assert result[0]["litellm_params"]["custom_llm_provider"] == "a2a" assert result[0]["model_info"]["id"] == "agent-123" assert result[0]["model_info"]["mode"] == "chat" + + +@pytest.mark.asyncio +async def test_append_agents_to_model_info_preserves_registered_provider(): + agent = AgentResponse( + agent_id="agent-123", + agent_name="test-agent", + agent_card_params={"url": "http://example.com"}, + litellm_params={"custom_llm_provider": "pydantic_ai_agents"}, + ) + registry = Mock() + registry.get_agent_by_id = Mock(return_value=agent) + + with ( + patch( # test-quality-ok: access resolution is outside model-list assembly + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access", + AsyncMock(return_value=RestrictedAgentAccess(frozenset({"agent-123"}))), + ), + patch( # test-quality-ok: registry output drives model-list assembly + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", + registry, + ), + ): + result = await append_agents_to_model_info( + models=[], + user_api_key_dict=Mock(spec=UserAPIKeyAuth), + ) + + assert result[0]["litellm_params"]["custom_llm_provider"] == "pydantic_ai_agents" diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 71d4666416d..9614a2bfa3c 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -6505,6 +6505,71 @@ class TestPerRequestModelGroupAlias: assert merged_for == ["group-b"] + @pytest.mark.asyncio + async def test_a2a_reroute_runs_target_guardrails(self, monkeypatch): + processing_obj = ProxyBaseLLMRequestProcessing( + data={"model": "source", "messages": [{"role": "user", "content": "hello"}]} + ) + mock_request = MagicMock(spec=Request) + mock_request.headers = {} + + async def mock_add_litellm_data_to_request(*args, **kwargs): + return kwargs.get("data", {}) + + hook_modes = [] + + async def reroute_once(user_api_key_dict, data, call_type, guardrails_only=False): + hook_modes.append(guardrails_only) + if len(hook_modes) == 1: + data["model"] = "a2a/agent" + return data + + async def passthrough(data, user_api_key_dict): + return data + + async def passthrough_data_only(data): + return data + + mock_proxy_logging_obj = MagicMock(spec=ProxyLogging) + mock_proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=reroute_once) + monkeypatch.setattr( + litellm.proxy.common_request_processing, + "add_litellm_data_to_request", + mock_add_litellm_data_to_request, + ) + monkeypatch.setattr( + litellm.utils, + "function_setup", + lambda original_function, rules_obj, start_time, **kwargs: (MagicMock(), kwargs), + ) + monkeypatch.setattr( + litellm.proxy.common_request_processing, + "_check_and_merge_model_level_guardrails", + lambda data, llm_router, trust_client_model_info=True: data, + ) + monkeypatch.setattr( + "litellm.proxy.agent_endpoints.a2a_routing.authorize_a2a_agent_before_hooks", + passthrough, + ) + monkeypatch.setattr( + "litellm.proxy.agent_endpoints.a2a_routing.merge_a2a_agent_guardrails_before_hooks", + passthrough_data_only, + ) + + returned_data, _ = await processing_obj.common_processing_pre_call_logic( + request=mock_request, + general_settings={}, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="hash"), + proxy_logging_obj=mock_proxy_logging_obj, + proxy_config=None, + route_type="acompletion", + llm_router=None, + ) + + assert returned_data["model"] == "a2a/agent" + assert hook_modes == [False, True] + + class TestInjectCostIntoUsageDict: @staticmethod def _expected_cost(model, prompt_tokens, completion_tokens): diff --git a/tests/test_litellm/proxy/test_route_a2a_models.py b/tests/test_litellm/proxy/test_route_a2a_models.py index 35308474949..f0d0ff084d6 100644 --- a/tests/test_litellm/proxy/test_route_a2a_models.py +++ b/tests/test_litellm/proxy/test_route_a2a_models.py @@ -4,13 +4,17 @@ Test A2A model routing in proxy. Maps to: litellm/proxy/agent_endpoints/a2a_routing.py """ - - from unittest.mock import AsyncMock, Mock, patch import pytest +from fastapi import HTTPException -from litellm.proxy.agent_endpoints.a2a_routing import route_a2a_agent_request +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.agent_endpoints.a2a_routing import ( + _route_registered_provider, + merge_a2a_agent_guardrails_before_hooks, + route_a2a_agent_request, +) from litellm.proxy.route_llm_request import route_request @@ -58,7 +62,7 @@ async def test_route_a2a_model_bypasses_router(): "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", mock_registry, ): - result = await route_request( + await route_request( data=data, llm_router=mock_router, user_model=None, @@ -72,6 +76,632 @@ async def test_route_a2a_model_bypasses_router(): assert call_kwargs["api_base"] == "http://agent.example.com" +@pytest.mark.asyncio +async def test_route_a2a_model_uses_registered_provider(): + from litellm.types.agents import AgentResponse + + agent = AgentResponse( + agent_id="test-agent-id", + agent_name="test-agent", + agent_card_params={"url": "http://agent.example.com"}, + litellm_params={ + "custom_llm_provider": "pydantic_ai_agents", + "guardrails": ["agent-guardrail"], + }, + static_headers={"Authorization": "Bearer static"}, + extra_headers=["X-Tenant"], + ) + data = { + "model": "a2a/test-agent", + "messages": [{"role": "user", "content": "Hello"}], + "guardrails": ["request-guardrail"], + "max_tokens": 32, + "temperature": 0.2, + "timeout": 12.0, + "tools": [{"type": "function", "function": {"name": "lookup"}}], + "output_config": {"format": "json"}, + "prompt_cache_key": "cache-key", + "safety_identifier": "safety-id", + "proxy_server_request": { + "headers": { + "x-tenant": "tenant-1", + "x-a2a-test-agent-x-run": "run-1", + } + }, + } + bridge_response = { + "jsonrpc": "2.0", + "id": "request-id", + "result": { + "kind": "message", + "role": "agent", + "parts": [{"kind": "text", "text": "Hello back"}], + "refusal": "I cannot complete that request.", + "messageId": "message-id", + }, + } + + with ( + patch( # test-quality-ok: registry lookup is the routing seam + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ), + patch( # test-quality-ok: access control is outside this routing test + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed", + AsyncMock(return_value=True), + ), + patch( # test-quality-ok: provider dispatch is the tested seam + "litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming", + AsyncMock(return_value=bridge_response), + ) as bridge, + patch( # test-quality-ok: generic dispatch must stay unused + "litellm.acompletion", AsyncMock() + ) as generic_completion, + ): + call = await route_a2a_agent_request(data, "acompletion") + response = await call + + bridge.assert_awaited_once() + generic_completion.assert_not_called() + assert response.choices[0].message.content == "Hello back" + assert response.choices[0].message.refusal == "I cannot complete that request." + bridge_kwargs = bridge.await_args.kwargs + assert bridge_kwargs["litellm_params"]["max_tokens"] == 32 + assert bridge_kwargs["litellm_params"]["temperature"] == 0.2 + assert bridge_kwargs["litellm_params"]["timeout"] == 12.0 + assert bridge_kwargs["litellm_params"]["tools"] == data["tools"] + assert bridge_kwargs["litellm_params"]["output_config"] == data["output_config"] + assert bridge_kwargs["litellm_params"]["prompt_cache_key"] == data["prompt_cache_key"] + assert bridge_kwargs["litellm_params"]["safety_identifier"] == data["safety_identifier"] + assert bridge_kwargs["litellm_params"]["guardrails"] == ["request-guardrail", "agent-guardrail"] + assert bridge_kwargs["litellm_params"]["extra_headers"] == { + "X-Tenant": "tenant-1", + "x-run": "run-1", + "Authorization": "Bearer static", + } + + +@pytest.mark.asyncio +async def test_route_a2a_cardless_bedrock_agentcore_uses_registered_model(): + from litellm.types.agents import AgentResponse + + agent = AgentResponse( + agent_id="test-agent-id", + agent_name="test-agent", + agent_card_params={}, + litellm_params={ + "custom_llm_provider": "bedrock", + "model": "bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:123:runtime/test", + }, + ) + bridge_response = { + "jsonrpc": "2.0", + "id": "request-id", + "result": {"kind": "message", "parts": [{"kind": "text", "text": "Hello back"}]}, + } + + with ( + patch( + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ), + patch( + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed", + AsyncMock(return_value=True), + ), + patch( + "litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming", + AsyncMock(return_value=bridge_response), + ) as bridge, + ): + call = await route_a2a_agent_request( + {"model": "a2a/test-agent", "messages": [{"role": "user", "content": "Hello"}]}, + "acompletion", + ) + await call + + assert bridge.await_args.kwargs["api_base"] is None + + +@pytest.mark.asyncio +async def test_route_a2a_registered_provider_uses_configured_api_base_without_card_url(): + from litellm.types.agents import AgentResponse + + agent = AgentResponse( + agent_id="test-agent-id", + agent_name="test-agent", + agent_card_params={}, + litellm_params={ + "custom_llm_provider": "langflow", + "model": "flow", + "api_base": "https://flow.example.com", + }, + ) + bridge_response = { + "jsonrpc": "2.0", + "id": "request-id", + "result": {"kind": "message", "parts": [{"kind": "text", "text": "Hello back"}]}, + } + + with ( + patch( + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ), + patch( + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed", + AsyncMock(return_value=True), + ), + patch( + "litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming", + AsyncMock(return_value=bridge_response), + ) as bridge, + ): + call = await route_a2a_agent_request( + {"model": "a2a/test-agent", "messages": [{"role": "user", "content": "Hello"}]}, + "acompletion", + ) + await call + + assert bridge.await_args.kwargs["api_base"] == "https://flow.example.com" + + +@pytest.mark.asyncio +async def test_registered_provider_response_preserves_multiple_choices(): + from litellm.types.agents import AgentResponse + + agent = AgentResponse( + agent_id="test-agent-id", + agent_name="test-agent", + agent_card_params={"url": "http://agent.example.com"}, + litellm_params={"custom_llm_provider": "pydantic_ai_agents"}, + ) + bridge_response = { + "jsonrpc": "2.0", + "id": "request-id", + "choices": [ + {"index": 0, "message": {"parts": [{"kind": "text", "text": "first"}]}}, + {"index": 1, "message": {"parts": [{"kind": "text", "text": "second"}]}}, + ], + "result": {}, + } + + with ( + patch( + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ), + patch( + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed", + AsyncMock(return_value=True), + ), + patch( + "litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming", + AsyncMock(return_value=bridge_response), + ), + ): + call = await route_a2a_agent_request( + {"model": "a2a/test-agent", "messages": [{"role": "user", "content": "Hello"}]}, + "acompletion", + ) + response = await call + + assert [choice.message.content for choice in response.choices] == ["first", "second"] + + +@pytest.mark.asyncio +async def test_route_a2a_cardless_watsonx_orchestrate_uses_registered_model(): + from litellm.types.agents import AgentResponse + + agent = AgentResponse( + agent_id="test-agent-id", + agent_name="test-agent", + agent_card_params={}, + litellm_params={ + "custom_llm_provider": "watsonx_orchestrate", + "model": "agent", + "cp4d_host": "https://wxo.example.com", + "instance_id": "instance", + }, + ) + bridge_response = { + "jsonrpc": "2.0", + "id": "request-id", + "result": {"kind": "message", "parts": [{"kind": "text", "text": "Hello back"}]}, + } + + with ( + patch( + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ), + patch( + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed", + AsyncMock(return_value=True), + ), + patch( + "litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming", + AsyncMock(return_value=bridge_response), + ) as bridge, + ): + call = await route_a2a_agent_request( + {"model": "a2a/test-agent", "messages": [{"role": "user", "content": "Hello"}]}, + "acompletion", + ) + await call + + assert bridge.await_args.kwargs["api_base"] is None + + +@pytest.mark.asyncio +async def test_route_a2a_registered_provider_preserves_identity_headers(): + from litellm.types.agents import AgentResponse + + agent = AgentResponse( + agent_id="test-agent-id", + agent_name="test-agent", + agent_card_params={"url": "http://agent.example.com"}, + litellm_params={"custom_llm_provider": "pydantic_ai_agents"}, + ) + bridge_response = { + "jsonrpc": "2.0", + "id": "request-id", + "result": {"kind": "message", "parts": [{"kind": "text", "text": "Hello back"}]}, + } + + with ( + patch( + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ), + patch( + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed", + AsyncMock(return_value=True), + ), + patch( + "litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming", + AsyncMock(return_value=bridge_response), + ) as bridge, + ): + call = await route_a2a_agent_request( + { + "model": "a2a/test-agent", + "messages": [{"role": "user", "content": "Hello"}], + "proxy_server_request": { + "headers": { + "x-a2a-test-agent-x-litellm-user-id": "attacker", + "x-a2a-test-agent-x-litellm-team-id": "attacker-team", + } + }, + }, + "acompletion", + user_api_key_dict=UserAPIKeyAuth(user_id="trusted-user", team_id="trusted-team"), + ) + await call + + headers = bridge.await_args.kwargs["agent_extra_headers"] + assert headers["X-LiteLLM-User-Id"] == "trusted-user" + assert headers["X-LiteLLM-Team-Id"] == "trusted-team" + assert "x-litellm-user-id" not in {key.lower() for key in headers if key != "X-LiteLLM-User-Id"} + + +@pytest.mark.asyncio +async def test_route_a2a_registered_provider_preserves_messages_and_session(): + from litellm.a2a_protocol.litellm_completion_bridge.handler import A2A_USER_API_KEY_HASH_PARAM + from litellm.types.agents import AgentResponse + + agent = AgentResponse( + agent_id="test-agent-id", + agent_name="test-agent", + agent_card_params={"url": "http://agent.example.com"}, + litellm_params={"custom_llm_provider": "langflow", "model": "flow"}, + ) + data = { + "model": "a2a/test-agent", + "messages": [ + {"role": "system", "content": "Be concise"}, + {"role": "user", "content": "Hello"}, + ], + "litellm_session_id": "session-1", + } + bridge_response = { + "jsonrpc": "2.0", + "id": "request-id", + "result": {"kind": "message", "parts": [{"kind": "text", "text": "Hello back"}]}, + } + + with ( + patch( + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ), + patch( + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed", + AsyncMock(return_value=True), + ), + patch( + "litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming", + AsyncMock(return_value=bridge_response), + ) as bridge, + ): + call = await route_a2a_agent_request( + data, + "acompletion", + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + ) + await call + + bridge_kwargs = bridge.await_args.kwargs + assert bridge_kwargs["params"]["messages"] == data["messages"] + assert bridge_kwargs["params"]["message"]["contextId"] == "session-1" + assert bridge_kwargs["litellm_params"][A2A_USER_API_KEY_HASH_PARAM] == "hashed-key" + + +@pytest.mark.asyncio +async def test_route_a2a_requires_inbound_trace_id(): + from litellm.types.agents import AgentResponse + + agent = AgentResponse( + agent_id="test-agent-id", + agent_name="test-agent", + agent_card_params={"url": "http://agent.example.com"}, + litellm_params={ + "custom_llm_provider": "pydantic_ai_agents", + "require_trace_id_on_calls_to_agent": True, + }, + ) + + with ( + patch( + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ), + patch( + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed", + AsyncMock(return_value=True), + ), + ): + with pytest.raises(HTTPException, match="requires x-litellm-trace-id"): + await route_a2a_agent_request( + {"model": "a2a/test-agent", "messages": [{"role": "user", "content": "Hello"}]}, + "acompletion", + ) + + +@pytest.mark.asyncio +async def test_route_a2a_resolves_databricks_oauth_headers(): + from litellm.types.agents import AgentResponse + + agent = AgentResponse( + agent_id="test-agent-id", + agent_name="test-agent", + agent_card_params={"url": "http://agent.example.com"}, + litellm_params={"custom_llm_provider": "databricks", "databricks_oauth": {"client_id": "id"}}, + ) + bridge_response = { + "jsonrpc": "2.0", + "id": "request-id", + "result": {"kind": "message", "parts": [{"kind": "text", "text": "Hello back"}]}, + } + + with ( + patch( + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ), + patch( + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed", + AsyncMock(return_value=True), + ), + patch( + "litellm.proxy.agent_endpoints.databricks_oauth.resolve_databricks_app_auth_header", + AsyncMock(return_value={"Authorization": "Bearer minted"}), + ), + patch( + "litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming", + AsyncMock(return_value=bridge_response), + ) as bridge, + ): + call = await route_a2a_agent_request( + {"model": "a2a/test-agent", "messages": [{"role": "user", "content": "Hello"}]}, + "acompletion", + ) + await call + + assert bridge.await_args.kwargs["agent_extra_headers"]["Authorization"] == "Bearer minted" + + +@pytest.mark.asyncio +async def test_registered_provider_response_preserves_tool_calls(): + from litellm.types.agents import AgentResponse + + agent = AgentResponse( + agent_id="test-agent-id", + agent_name="test-agent", + agent_card_params={"url": "http://agent.example.com"}, + litellm_params={"custom_llm_provider": "pydantic_ai_agents"}, + ) + bridge_response = { + "jsonrpc": "2.0", + "id": "request-id", + "result": { + "kind": "message", + "parts": [], + "tool_calls": [ + { + "id": "call-1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ], + "finish_reason": "tool_calls", + }, + } + + with ( + patch( + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ), + patch( + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed", + AsyncMock(return_value=True), + ), + patch( + "litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming", + AsyncMock(return_value=bridge_response), + ), + ): + call = await route_a2a_agent_request( + {"model": "a2a/test-agent", "messages": [{"role": "user", "content": "Hello"}]}, + "acompletion", + ) + response = await call + + assert response.choices[0].finish_reason == "tool_calls" + assert response.choices[0].message.tool_calls[0].id == "call-1" + + +@pytest.mark.asyncio +async def test_a2a_agent_guardrails_merge_before_hooks(): + from litellm.types.agents import AgentResponse + + agent = AgentResponse( + agent_id="test-agent-id", + agent_name="test-agent", + agent_card_params={"url": "http://agent.example.com"}, + litellm_params={"guardrails": ["agent-guardrail"]}, + ) + with patch( + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ): + merged = await merge_a2a_agent_guardrails_before_hooks( + {"model": "a2a/test-agent", "guardrails": ["request-guardrail"]} + ) + + assert merged["guardrails"] == ["request-guardrail", "agent-guardrail"] + + +@pytest.mark.asyncio +async def test_route_a2a_stream_uses_registered_provider(): + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.agents import AgentResponse + + agent = AgentResponse( + agent_id="test-agent-id", + agent_name="test-agent", + agent_card_params={"url": "http://agent.example.com"}, + litellm_params={"custom_llm_provider": "pydantic_ai_agents"}, + ) + logging_obj = Mock(spec=Logging) + logging_obj.model_call_details = {} + data = { + "model": "a2a/test-agent", + "messages": [{"role": "user", "content": "Hello"}], + "stream": True, + "litellm_logging_obj": logging_obj, + } + provider_stream = object() + completion_stream = object() + wrapper = object() + + with ( + patch( # test-quality-ok: registry lookup is the routing seam + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ), + patch( # test-quality-ok: access control is outside this routing test + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed", + AsyncMock(return_value=True), + ), + patch( # test-quality-ok: provider dispatch is the tested seam + "litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_streaming", + Mock(return_value=provider_stream), + ) as bridge, + patch( # test-quality-ok: iterator wiring is the tested seam + "litellm.llms.a2a.chat.streaming_iterator.A2AModelResponseIterator", + Mock(return_value=completion_stream), + ), + patch( # test-quality-ok: wrapper wiring is the tested seam + "litellm.litellm_core_utils.streaming_handler.CustomStreamWrapper", + Mock(return_value=wrapper), + ) as stream_wrapper, + patch( # test-quality-ok: generic dispatch must stay unused + "litellm.acompletion", AsyncMock() + ) as generic_completion, + ): + call = await route_a2a_agent_request(data, "acompletion") + response = await call + + bridge.assert_called_once() + stream_wrapper.assert_called_once_with( + completion_stream=completion_stream, + model="a2a/test-agent", + custom_llm_provider="a2a", + logging_obj=logging_obj, + stream_options=None, + ) + generic_completion.assert_not_called() + assert response is wrapper + + +@pytest.mark.asyncio +async def test_registered_provider_logging_uses_provider_model_for_builtin_pricing(): + class FakeLogging: + def __init__(self) -> None: + self.model_call_details = {"litellm_params": {}} + self.litellm_params = self.model_call_details["litellm_params"] + self.custom_pricing = False + + logging_obj = FakeLogging() + response = {"result": {"message": {"parts": [{"kind": "text", "text": "hello"}]}}} + with ( + patch("litellm.litellm_core_utils.litellm_logging.Logging", FakeLogging), + patch( + "litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming", + AsyncMock(return_value=response), + ), + ): + await _route_registered_provider( + data={ + "messages": [{"role": "user", "content": "hello"}], + "litellm_logging_obj": logging_obj, + }, + model_name="a2a/agent", + api_base="https://provider.example", + litellm_params={"model": "gpt-4o", "custom_llm_provider": "openai"}, + static_headers=None, + ) + + assert logging_obj.model_call_details["model"] == "gpt-4o" + assert logging_obj.model_call_details["custom_llm_provider"] == "openai" + assert logging_obj.model_call_details["litellm_params"]["model"] == "gpt-4o" + + +@pytest.mark.asyncio +async def test_native_registered_provider_estimates_usage_when_missing(monkeypatch): + response = {"result": {"message": {"parts": [{"kind": "text", "text": "hello"}]}}} + counter = Mock(side_effect=[3, 2]) + monkeypatch.setattr("litellm.utils.token_counter", counter) + + with patch( + "litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming", + AsyncMock(return_value=response), + ): + result = await _route_registered_provider( + data={"messages": [{"role": "user", "content": "hello"}]}, + model_name="a2a/agent", + api_base="https://provider.example", + litellm_params={"model": "agent", "custom_llm_provider": "pydantic_ai_agents"}, + static_headers=None, + ) + + assert result.usage.prompt_tokens == 3 + assert result.usage.completion_tokens == 2 + assert result.usage.total_tokens == 5 + + @pytest.mark.asyncio async def test_route_non_a2a_model_raises_error_if_not_in_router(): """Test that non-a2a models that aren't in router raise an error""" @@ -141,7 +771,7 @@ def _router_without_models(): @pytest.mark.asyncio async def test_route_a2a_model_read_through_recovers_agent_created_on_sibling_replica(monkeypatch): - import litellm.proxy.proxy_server as proxy_server + from litellm.proxy import proxy_server from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry agent_name = "a2a-sibling-replica-agent"