scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 79

feat(a2a/): support streaming guardrails

Browse source
This commit is contained in:
Krrish Dholakia 2026-02-06 22:55:48 -08:00
parent f30f883744
commit 11b0958592
7 changed files with 385 additions and 51 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

33
litellm/litellm_core_utils/api_route_to_call_types.py
View file

@ -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:

113
litellm/llms/a2a/chat/guardrail_translation/handler.py
View file

@ -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],

87
litellm/proxy/agent_endpoints/a2a_endpoints.py
View file

@ -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")

70
litellm/proxy/common_request_processing.py
View file

@ -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:

1
litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py
View file

@ -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(

124
litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
View file

@ -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

8
litellm/proxy/utils.py
View file

@ -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: