From 11b0958592519a2501ac2b86ff637e5946cb48fd Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 6 Feb 2026 22:55:48 -0800 Subject: [PATCH] feat(a2a/): support streaming guardrails --- .../api_route_to_call_types.py | 33 ++++- .../a2a/chat/guardrail_translation/handler.py | 113 ++++++++++++++++ .../proxy/agent_endpoints/a2a_endpoints.py | 87 ++++++++++-- litellm/proxy/common_request_processing.py | 70 +++++++--- .../custom_code/custom_code_guardrail.py | 1 - .../unified_guardrail/unified_guardrail.py | 124 +++++++++++++++--- litellm/proxy/utils.py | 8 ++ 7 files changed, 385 insertions(+), 51 deletions(-) diff --git a/litellm/litellm_core_utils/api_route_to_call_types.py b/litellm/litellm_core_utils/api_route_to_call_types.py index 4146ff6d6a6..2ae9986ce94 100644 --- a/litellm/litellm_core_utils/api_route_to_call_types.py +++ b/litellm/litellm_core_utils/api_route_to_call_types.py @@ -3,6 +3,9 @@ Dictionary mapping API routes to their corresponding CallTypes in LiteLLM. This dictionary maps each API endpoint to the CallTypes that can be used for that route. Each route can have both async (prefixed with 'a') and sync call types. + +Route patterns may contain placeholders like {agent_id}, {model}, {batch_id}; these +match a single path segment when resolving call types for a concrete path. """ from typing import List, Optional @@ -10,17 +13,43 @@ from typing import List, Optional from litellm.types.utils import API_ROUTE_TO_CALL_TYPES, CallTypes +def _route_matches_pattern(route: str, pattern: str) -> bool: + """ + Return True if the concrete route matches the pattern. + Pattern segments like {param} match any single path segment. + """ + route_parts = route.strip("/").split("/") + pattern_parts = pattern.strip("/").split("/") + if len(route_parts) != len(pattern_parts): + return False + for r, p in zip(route_parts, pattern_parts): + if p.startswith("{") and p.endswith("}"): + continue + if r != p: + return False + return True + + def get_call_types_for_route(route: str) -> Optional[List[CallTypes]]: """ Get the list of CallTypes for a given API route. + Supports both exact keys and dynamic patterns (e.g. /a2a/my-agent/message/send + matches /a2a/{agent_id}/message/send). + Args: - route: API route path (e.g., "/chat/completions") + route: API route path (e.g., "/chat/completions" or "/a2a/my-pydantic-agent/message/send") Returns: List of CallTypes for that route, or None if route not found """ - return API_ROUTE_TO_CALL_TYPES.get(route, None) + exact = API_ROUTE_TO_CALL_TYPES.get(route, None) + if exact is not None: + return exact + for pattern, call_types in API_ROUTE_TO_CALL_TYPES.items(): + if _route_matches_pattern(route, pattern): + return call_types + return None def get_routes_for_call_type(call_type: CallTypes) -> list: diff --git a/litellm/llms/a2a/chat/guardrail_translation/handler.py b/litellm/llms/a2a/chat/guardrail_translation/handler.py index 770453f2def..fbd1da749c2 100644 --- a/litellm/llms/a2a/chat/guardrail_translation/handler.py +++ b/litellm/llms/a2a/chat/guardrail_translation/handler.py @@ -10,6 +10,7 @@ A2A Protocol Format: - Output: JSON-RPC 2.0 with result containing message/artifact parts """ +import json from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union from litellm._logging import verbose_proxy_logger @@ -206,6 +207,118 @@ class A2AGuardrailHandler(BaseTranslation): response["result"] = result return response + async def process_output_streaming_response( + self, + responses_so_far: List[Any], + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, + user_api_key_dict: Optional["UserAPIKeyAuth"] = None, + ) -> List[Any]: + """ + Process A2A streaming output by applying guardrails to accumulated text. + + responses_so_far can be a list of JSON-RPC 2.0 objects (dict or NDJSON str), e.g.: + - task with history, status-update, artifact-update (with result.artifact.parts), + - then status-update (final). Text is extracted from result.artifact.parts, + result.message.parts, result.parts, etc., concatenated in order, guardrailed once, + then the combined guardrailed text is written into the first chunk that had text + and all other text parts in other chunks are cleared (in-place). + """ + from litellm.llms.a2a.common_utils import extract_text_from_a2a_response + + # Parse each item; keep alignment with responses_so_far (None where unparseable) + parsed: List[Optional[Dict[str, Any]]] = [None] * len(responses_so_far) + for i, item in enumerate(responses_so_far): + if isinstance(item, dict): + obj = item + elif isinstance(item, str): + try: + obj = json.loads(item.strip()) + except (json.JSONDecodeError, TypeError): + continue + else: + continue + if isinstance(obj.get("result"), dict): + parsed[i] = obj + + valid_parsed = [(i, obj) for i, obj in enumerate(parsed) if obj is not None] + if not valid_parsed: + return responses_so_far + + # Collect text from each chunk in order (by original index in responses_so_far) + text_parts: List[str] = [] + chunk_indices_with_text: List[int] = [] # indices into valid_parsed + for idx, (orig_i, obj) in enumerate(valid_parsed): + t = extract_text_from_a2a_response(obj) + if t: + text_parts.append(t) + chunk_indices_with_text.append(orig_i) + + combined_text = "".join(text_parts) + if not combined_text: + return responses_so_far + + request_data: dict = {"responses_so_far": responses_so_far} + user_metadata = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) + if user_metadata: + request_data["litellm_metadata"] = user_metadata + + inputs = GenericGuardrailAPIInputs(texts=[combined_text]) + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="response", + logging_obj=litellm_logging_obj, + ) + guardrailed_texts = guardrailed_inputs.get("texts", []) + if not guardrailed_texts: + return responses_so_far + guardrailed_text = guardrailed_texts[0] + + # Find first chunk (by original index) that has text; put full guardrailed text there and clear rest + first_chunk_with_text: Optional[int] = ( + chunk_indices_with_text[0] if chunk_indices_with_text else None + ) + + for orig_i, obj in valid_parsed: + result = obj.get("result", {}) + if not isinstance(result, dict): + continue + texts_in_chunk: List[str] = [] + mappings: List[Tuple[Tuple[str, ...], int]] = [] + self._extract_texts_from_result( + result=result, + texts_to_check=texts_in_chunk, + task_mappings=mappings, + ) + if not mappings: + continue + if orig_i == first_chunk_with_text: + # Put full guardrailed text in first text part; clear others + for task_idx, (path, part_idx) in enumerate(mappings): + text = guardrailed_text if task_idx == 0 else "" + self._apply_text_to_path( + result=result, + path=path, + part_idx=part_idx, + text=text, + ) + else: + for path, part_idx in mappings: + self._apply_text_to_path( + result=result, + path=path, + part_idx=part_idx, + text="", + ) + + # Write back to responses_so_far where we had NDJSON strings + for i, item in enumerate(responses_so_far): + if isinstance(item, str) and parsed[i] is not None: + responses_so_far[i] = json.dumps(parsed[i]) + "\n" + + return responses_so_far + def _extract_texts_from_result( self, result: Dict[str, Any], diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index f1ae1f64f00..7f30277ebca 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -6,7 +6,7 @@ The A2A SDK can point to LiteLLM's URL and invoke agents registered with LiteLLM """ import json -from typing import Optional +from typing import Any, Optional from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import JSONResponse, StreamingResponse @@ -54,14 +54,21 @@ async def _handle_stream_message( agent_id: Optional[str] = None, metadata: Optional[dict] = None, proxy_server_request: Optional[dict] = None, + *, + user_api_key_dict: Optional[UserAPIKeyAuth] = None, + request_data: Optional[dict] = None, + proxy_logging_obj: Optional[Any] = None, ) -> StreamingResponse: - """Handle message/stream method via SDK functions.""" + """Handle message/stream method via SDK functions. + + When user_api_key_dict, request_data, and proxy_logging_obj are provided, + uses common_request_processing.async_streaming_data_generator with NDJSON + serializers so proxy hooks and cost injection apply. + """ from litellm.a2a_protocol import asend_message_streaming from litellm.a2a_protocol.main import A2A_SDK_AVAILABLE - # Check is handled in invoke_agent_a2a, but if called directly: if not A2A_SDK_AVAILABLE: - # Return a streaming response that yields an error async def _error_stream(): yield json.dumps( { @@ -78,29 +85,82 @@ async def _handle_stream_message( from a2a.types import MessageSendParams, SendStreamingMessageRequest + use_proxy_hooks = ( + user_api_key_dict is not None + and request_data is not None + and proxy_logging_obj is not None + ) + async def stream_response(): try: a2a_request = SendStreamingMessageRequest( id=request_id, params=MessageSendParams(**params), ) - async for chunk in asend_message_streaming( + a2a_stream = asend_message_streaming( request=a2a_request, api_base=api_base, litellm_params=litellm_params, agent_id=agent_id, metadata=metadata, proxy_server_request=proxy_server_request, - ): - # Chunk may be dict or object depending on bridge vs standard path - if hasattr(chunk, "model_dump"): - yield json.dumps( - chunk.model_dump(mode="json", exclude_none=True) + ) + + if use_proxy_hooks and user_api_key_dict is not None and request_data is not None and proxy_logging_obj is not None: + from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + ) + + def _ndjson_chunk(chunk: Any) -> str: + if hasattr(chunk, "model_dump"): + obj = chunk.model_dump(mode="json", exclude_none=True) + else: + obj = chunk + return json.dumps(obj) + "\n" + + def _ndjson_error(proxy_exc: Any) -> str: + return json.dumps( + { + "jsonrpc": "2.0", + "id": request_id, + "error": { + "code": -32603, + "message": getattr( + proxy_exc, "message", f"Streaming error: {proxy_exc!s}" + ), + }, + } ) + "\n" - else: - yield json.dumps(chunk) + "\n" + + async for line in ProxyBaseLLMRequestProcessing.async_streaming_data_generator( + response=a2a_stream, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + proxy_logging_obj=proxy_logging_obj, + serialize_chunk=_ndjson_chunk, + serialize_error=_ndjson_error, + ): + yield line + else: + async for chunk in a2a_stream: + if hasattr(chunk, "model_dump"): + yield json.dumps( + chunk.model_dump(mode="json", exclude_none=True) + ) + "\n" + else: + yield json.dumps(chunk) + "\n" except Exception as e: verbose_proxy_logger.exception(f"Error streaming A2A response: {e}") + if use_proxy_hooks and proxy_logging_obj is not None and user_api_key_dict is not None and request_data is not None: + transformed_exception = await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=e, + request_data=request_data, + ) + if transformed_exception is not None: + e = transformed_exception + if isinstance(e, HTTPException): + raise yield json.dumps( { "jsonrpc": "2.0", @@ -346,6 +406,9 @@ async def invoke_agent_a2a( agent_id=agent.agent_id, metadata=data.get("metadata", {}), proxy_server_request=data.get("proxy_server_request"), + user_api_key_dict=user_api_key_dict, + request_data=data, + proxy_logging_obj=proxy_logging_obj, ) else: return _jsonrpc_error(request_id, -32601, f"Method '{method}' not found") diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index a02bc7f9e51..f0fa5e44b05 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -43,6 +43,11 @@ from litellm.proxy.utils import ProxyLogging from litellm.router import Router from litellm.types.utils import ServerToolUse +# Type alias for streaming chunk serializer (chunk after hooks + cost injection -> wire format) +StreamChunkSerializer = Callable[[Any], str] +# Type alias for streaming error serializer (ProxyException -> wire format) +StreamErrorSerializer = Callable[[ProxyException], str] + if TYPE_CHECKING: from litellm.proxy.proxy_server import ProxyConfig as _ProxyConfig @@ -1193,22 +1198,24 @@ class ProxyBaseLLMRequestProcessing: return chunk @staticmethod - async def async_sse_data_generator( - response, + async def async_streaming_data_generator( + response: Any, user_api_key_dict: UserAPIKeyAuth, request_data: dict, proxy_logging_obj: ProxyLogging, - ): + *, + serialize_chunk: StreamChunkSerializer, + serialize_error: StreamErrorSerializer, + ) -> AsyncGenerator[str, None]: """ - Anthropic /messages and Google /generateContent streaming data generator require SSE events + Shared streaming data generator: runs proxy iterator hook, per-chunk hook, + cost injection, then yields chunks via serialize_chunk; on exception runs + failure hook and yields via serialize_error. Use for SSE or NDJSON. """ - verbose_proxy_logger.debug("inside generator") try: str_so_far = "" - async for ( - chunk - ) in proxy_logging_obj.async_post_call_streaming_iterator_hook( + async for chunk in proxy_logging_obj.async_post_call_streaming_iterator_hook( user_api_key_dict=user_api_key_dict, response=response, request_data=request_data, @@ -1216,7 +1223,6 @@ class ProxyBaseLLMRequestProcessing: verbose_proxy_logger.debug( "async_data_generator: received streaming chunk - {}".format(chunk) ) - ### CALL HOOKS ### - modify outgoing data chunk = await proxy_logging_obj.async_post_call_streaming_hook( user_api_key_dict=user_api_key_dict, response=chunk, @@ -1227,30 +1233,34 @@ class ProxyBaseLLMRequestProcessing: if isinstance(chunk, (ModelResponse, ModelResponseStream)): response_str = litellm.get_response_string(response_obj=chunk) str_so_far += response_str + elif hasattr(chunk, "model_dump"): + try: + d = chunk.model_dump(mode="json", exclude_none=True) + if isinstance(d, dict): + str_so_far += str(d.get("content", "")) + except Exception: + pass + elif isinstance(chunk, dict): + str_so_far += str(chunk.get("content", "")) - # Inject cost into Anthropic-style SSE usage for /v1/messages for any provider model_name = request_data.get("model", "") chunk = ( ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( chunk, model_name ) ) - - # Format chunk using helper function - yield ProxyBaseLLMRequestProcessing.return_sse_chunk(chunk) + yield serialize_chunk(chunk) except Exception as e: verbose_proxy_logger.exception( "litellm.proxy.proxy_server.async_data_generator(): Exception occured - {}".format( str(e) ) ) - # Allow callbacks to transform the error response transformed_exception = await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=request_data, ) - # Use transformed exception if callback returned one, otherwise use original if transformed_exception is not None: e = transformed_exception verbose_proxy_logger.debug( @@ -1259,18 +1269,36 @@ class ProxyBaseLLMRequestProcessing: if isinstance(e, HTTPException): raise e - else: - error_traceback = traceback.format_exc() - error_msg = f"{str(e)}\n\n{error_traceback}" - + error_traceback = traceback.format_exc() + error_msg = f"{str(e)}\n\n{error_traceback}" proxy_exception = ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), param=getattr(e, "param", "None"), code=getattr(e, "status_code", 500), ) - error_returned = json.dumps({"error": proxy_exception.to_dict()}) - yield f"{STREAM_SSE_DATA_PREFIX}{error_returned}\n\n" + yield serialize_error(proxy_exception) + + @staticmethod + async def async_sse_data_generator( + response: Any, + user_api_key_dict: UserAPIKeyAuth, + request_data: dict, + proxy_logging_obj: ProxyLogging, + ) -> AsyncGenerator[str, None]: + """ + Anthropic /messages and Google /generateContent streaming data generator require SSE events. + Delegates to async_streaming_data_generator with SSE serializers. + """ + async for chunk in ProxyBaseLLMRequestProcessing.async_streaming_data_generator( + response=response, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + proxy_logging_obj=proxy_logging_obj, + serialize_chunk=ProxyBaseLLMRequestProcessing.return_sse_chunk, + serialize_error=lambda proxy_exc: f"{STREAM_SSE_DATA_PREFIX}{json.dumps({'error': proxy_exc.to_dict()})}\n\n", + ): + yield chunk @staticmethod def _process_chunk_with_cost_injection(chunk: Any, model_name: str) -> Any: diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py index 66b80c10f18..b2ef495be8a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py @@ -216,7 +216,6 @@ class CustomCodeGuardrail(CustomGuardrail): HTTPException: If content is blocked CustomCodeExecutionError: If execution fails """ - if self._compiled_function is None: if self._compile_error: raise CustomCodeExecutionError( diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index fe69e1caf99..84bbf6d20e1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -7,8 +7,11 @@ Unified Guardrail, leveraging LiteLLM's /applyGuardrail endpoint """ import copy +import json from typing import Any, AsyncGenerator, List, Optional, Union +from fastapi import HTTPException + from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.cost_calculator import _infer_call_type @@ -18,9 +21,34 @@ from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_fo from litellm.llms import load_guardrail_translation_mappings from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.utils import CallTypes, CallTypesLiteral, ModelResponseStream +from litellm.types.utils import CallTypes, CallTypesLiteral + +# Call types that use NDJSON streaming (A2A); guardrail HTTPException is emitted as in-stream error +A2A_CALL_TYPES = (CallTypes.asend_message, CallTypes.send_message) GUARDRAIL_NAME = "unified_llm_guardrails" + + +def _get_a2a_request_id( + responses_so_far: List[Any], request_data: dict +) -> Optional[str]: + """Get JSON-RPC request id from first A2A chunk or request body for in-stream error reporting.""" + for item in responses_so_far: + if isinstance(item, dict) and "id" in item: + return item.get("id") + if isinstance(item, str): + try: + obj = json.loads(item.strip()) + if isinstance(obj, dict) and "id" in obj: + return obj.get("id") + except (json.JSONDecodeError, TypeError): + continue + body = request_data.get("body") or request_data.get("data") or {} + if isinstance(body, dict): + return body.get("id") + return None + + endpoint_guardrail_translation_mappings = None @@ -232,7 +260,7 @@ class UnifiedLLMGuardrails(CustomLogger): user_api_key_dict: UserAPIKeyAuth, response: Any, request_data: dict, - ) -> AsyncGenerator[ModelResponseStream, None]: + ) -> AsyncGenerator[Any, None]: """ Passes the entire stream to the guardrail @@ -363,13 +391,47 @@ class UnifiedLLMGuardrails(CustomLogger): CallTypes(call_type) ]() - await endpoint_translation.process_output_streaming_response( - responses_so_far=responses_so_far, - guardrail_to_apply=guardrail_to_apply, - litellm_logging_obj=request_data.get("litellm_logging_obj"), - user_api_key_dict=user_api_key_dict, - ) - + try: + await endpoint_translation.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=request_data.get("litellm_logging_obj"), + user_api_key_dict=user_api_key_dict, + ) + except HTTPException as e: + # Response already started (we already yielded chunks); cannot send 400. + # For A2A (NDJSON), yield an in-stream JSON-RPC error so the client sees it. + if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES: + request_id = _get_a2a_request_id(responses_so_far, request_data) + detail = ( + e.detail + if isinstance(e.detail, dict) + else {"message": str(e.detail)} + ) + error_chunk = ( + json.dumps( + { + "jsonrpc": "2.0", + "id": request_id, + "error": { + "code": -32603, + "message": detail.get( + "error", + detail.get("message", str(e.detail)), + ), + "data": { + k: v + for k, v in detail.items() + if k not in ("error", "message") + }, + }, + } + ) + + "\n" + ) + yield error_chunk + return + raise yield original_item else: yield item @@ -389,9 +451,41 @@ class UnifiedLLMGuardrails(CustomLogger): CallTypes(call_type) ]() - await endpoint_translation.process_output_streaming_response( - responses_so_far=responses_so_far, - guardrail_to_apply=guardrail_to_apply, - litellm_logging_obj=request_data.get("litellm_logging_obj"), - user_api_key_dict=user_api_key_dict, - ) + try: + await endpoint_translation.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=request_data.get("litellm_logging_obj"), + user_api_key_dict=user_api_key_dict, + ) + except HTTPException as e: + if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES: + request_id = _get_a2a_request_id(responses_so_far, request_data) + detail = ( + e.detail + if isinstance(e.detail, dict) + else {"message": str(e.detail)} + ) + error_chunk = ( + json.dumps( + { + "jsonrpc": "2.0", + "id": request_id, + "error": { + "code": -32603, + "message": detail.get( + "error", detail.get("message", str(e.detail)) + ), + "data": { + k: v + for k, v in detail.items() + if k not in ("error", "message") + }, + }, + } + ) + + "\n" + ) + yield error_chunk + else: + raise diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 53cd2db19cd..6c8db8d7d99 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1981,6 +1981,10 @@ class ProxyLogging: ) return merged_headers + def is_a2a_streaming_response(self, response: dict) -> bool: + expected_keys = ["jsonrpc", "id", "result"] + return all(key in response for key in expected_keys) + async def async_post_call_streaming_hook( self, data: dict, @@ -2001,6 +2005,10 @@ class ProxyLogging: response_str: Optional[str] = None if isinstance(response, (ModelResponse, ModelResponseStream)): response_str = litellm.get_response_string(response_obj=response) + elif isinstance(response, dict) and self.is_a2a_streaming_response(response): + from litellm.llms.a2a.common_utils import extract_text_from_a2a_response + + response_str = extract_text_from_a2a_response(response) if response_str is not None: for callback in litellm.callbacks: try: