scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

Merge e4cd222fd9 into febbe38326

Browse source
This commit is contained in:
Hill Patel 2026-11-06 13:58:09 +00:00 • committed by GitHub
commit b7b53c563d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 397 additions and 94 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.github/workflows/test-unit.yml vendored
View file

@ -753,6 +753,7 @@ jobs:
tests/unit/responses/test_metadata_codex_callback.py
tests/unit/responses/test_no_duplicate_spend_logs.py
tests/unit/responses/test_null_test_fix.py
tests/unit/responses/test_regression_issue_28553.py
tests/unit/responses/test_responses_api_bridge_flag.py
tests/unit/responses/test_responses_api_lifecycle.py
tests/unit/responses/test_responses_api_request_body.py

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

@ -12,12 +12,14 @@ A2A Protocol Format:
import json
from collections.abc import Sequence
from itertools import chain
from typing import TYPE_CHECKING, Any, Final, Optional
from pydantic import ConfigDict, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.llms.a2a.common_utils import serialize_a2a_data_part
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
StreamingScanKey,
@ -37,6 +39,8 @@ class _A2ATextPart(TypedDict, total=False):
kind: ReadOnly[str]
text: ReadOnly[str]
data: ReadOnly[object]
parts: ReadOnly[Sequence["_A2ATextPart"]]
class A2AGuardrailHandler(BaseTranslation):
@ -48,8 +52,10 @@ class A2AGuardrailHandler(BaseTranslation):
2. Process output responses (post-call hook) - extracts text from A2A response parts
A2A Message Format:
- Input: params.message.parts[].text (where kind == "text")
- Output: result.message.parts[].text or result.artifacts[].parts[].text
- Input: params.message.parts[].text (where kind == "text") or
params.message.parts[].data (where kind == "data")
- Output: result.message.parts[].text or result.artifacts[].parts[].text,
and the "data" equivalents of both
"""
async def process_input_messages(
@ -80,20 +86,28 @@ class A2AGuardrailHandler(BaseTranslation):
verbose_proxy_logger.debug("A2A: No parts in message, skipping guardrail")
return data
texts_to_check: Final[list[str]] = []
text_part_indices: Final[list[int]] = [] # Track which parts contain text
def _scan_input_part(part_idx: int, part: dict[str, Any]) -> tuple[str, int, str] | None:
kind: Final = part.get("kind")
if kind == "text":
text: Final = part.get("text", "")
return (text, part_idx, "text") if text else None
if kind == "data":
part_data: Final = part.get("data")
return (serialize_a2a_data_part(part_data), part_idx, "data") if part_data is not None else None
return None
# Step 1: Extract text from all text parts
for part_idx, part in enumerate(parts):
if part.get("kind") == "text":
text = part.get("text", "")
if text:
texts_to_check.append(text)
text_part_indices.append(part_idx)
# Extract text from all text parts, and serialized data from all data parts
scanned: Final = tuple(
entry for part_idx, part in enumerate(parts) if (entry := _scan_input_part(part_idx, part)) is not None
)
texts_to_check: Final = tuple(text for text, _, _ in scanned)
# Track which parts contain scannable content, and which field to write
# the guardrailed value back to ("text" or "data")
part_mappings: Final = tuple((part_idx, field) for _, part_idx, field in scanned)
# Step 2: Apply guardrail to all texts in batch
if texts_to_check:
inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check)
inputs: Final = GenericGuardrailAPIInputs(texts=list(texts_to_check))
# Pass the structured A2A message to guardrails
inputs["structured_messages"] = [message]
@ -113,9 +127,9 @@ class A2AGuardrailHandler(BaseTranslation):
guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
# Step 3: Apply guardrailed text back to original parts
if guardrailed_texts and len(guardrailed_texts) == len(text_part_indices):
for task_idx, part_idx in enumerate(text_part_indices):
parts[part_idx]["text"] = guardrailed_texts[task_idx]
if guardrailed_texts and len(guardrailed_texts) == len(part_mappings):
for task_idx, (part_idx, field) in enumerate(part_mappings):
parts[part_idx][field] = guardrailed_texts[task_idx]
verbose_proxy_logger.debug("A2A: Processed input message: %s", message)
@ -163,18 +177,12 @@ class A2AGuardrailHandler(BaseTranslation):
verbose_proxy_logger.debug("A2A: No result in response, skipping guardrail")
return response
# Find all text-containing parts in the response
texts_to_check: Final[list[str]] = []
# Each mapping is (path_to_parts_list, part_index)
# path_to_parts_list is a tuple of keys to navigate to the parts list
task_mappings: Final[list[tuple[tuple[str, ...], int]]] = []
# Extract texts from all possible locations
self._extract_texts_from_result(
result=result,
texts_to_check=texts_to_check,
task_mappings=task_mappings,
)
# Find all text-containing parts in the response. Each scanned entry is
# (text, path_to_parts_list, part_index, field); path_to_parts_list is a
# tuple of keys to navigate to the parts list.
scanned: Final = self._extract_texts_from_result(result=result)
texts_to_check: Final = tuple(text for text, _, _, _ in scanned)
task_mappings: Final = tuple((path, part_idx, field) for _, path, part_idx, field in scanned)
if not texts_to_check:
verbose_proxy_logger.debug("A2A: No text content in response")
@ -195,7 +203,7 @@ class A2AGuardrailHandler(BaseTranslation):
if user_metadata:
request_data["litellm_metadata"] = user_metadata
inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check)
inputs: Final = GenericGuardrailAPIInputs(texts=list(texts_to_check))
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
@ -208,11 +216,12 @@ class A2AGuardrailHandler(BaseTranslation):
# Step 3: Apply guardrailed text back to original response
if guardrailed_texts and len(guardrailed_texts) == len(task_mappings):
for task_idx, (path, part_idx) in enumerate(task_mappings):
for task_idx, (path, part_idx, field) in enumerate(task_mappings):
self._apply_text_to_path(
result=result,
path=path,
part_idx=part_idx,
field=field,
text=guardrailed_texts[task_idx],
)
@ -284,31 +293,27 @@ class A2AGuardrailHandler(BaseTranslation):
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:
scanned = self._extract_texts_from_result(result=result)
if not scanned:
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):
for task_idx, (_, path, part_idx, field) in enumerate(scanned):
text = guardrailed_text if task_idx == 0 else ""
self._apply_text_to_path(
result=result,
path=path,
part_idx=part_idx,
field=field,
text=text,
)
else:
for path, part_idx in mappings:
for _, path, part_idx, field in scanned:
self._apply_text_to_path(
result=result,
path=path,
part_idx=part_idx,
field=field,
text="",
)
@ -365,9 +370,7 @@ class A2AGuardrailHandler(BaseTranslation):
def _extract_texts_from_result(
self,
result: dict[str, Any],
texts_to_check: list[str],
task_mappings: list[tuple[tuple[str, ...], int]],
) -> None:
) -> tuple[tuple[str, tuple[str, ...], int, str], ...]:
"""
Extract text from all possible locations in an A2A result.
@ -377,80 +380,100 @@ class A2AGuardrailHandler(BaseTranslation):
3. Task with artifacts: {"artifacts": [{"parts": [...]}]}
4. Task with status message: {"status": {"message": {"parts": [...]}}}
5. Streaming artifact-update: {"artifact": {"parts": [...]}}
Returns a tuple of (text, path_to_parts_list, part_index, field) entries.
"""
# Case 1: Direct parts in result (direct message)
if "parts" in result:
self._extract_texts_from_parts(
parts=result["parts"],
path=("parts",),
texts_to_check=texts_to_check,
task_mappings=task_mappings,
)
case1: Final = (
self._extract_texts_from_parts(parts=result["parts"], path=("parts",)) if "parts" in result else ()
)
# Case 2: Nested message
message: Final = result.get("message")
if message and isinstance(message, dict) and "parts" in message:
self._extract_texts_from_parts(
parts=message["parts"],
path=("message", "parts"),
texts_to_check=texts_to_check,
task_mappings=task_mappings,
)
case2: Final = (
self._extract_texts_from_parts(parts=message["parts"], path=("message", "parts"))
if message and isinstance(message, dict) and "parts" in message
else ()
)
# Case 3: Streaming artifact-update (singular artifact)
artifact: Final = result.get("artifact")
if artifact and isinstance(artifact, dict) and "parts" in artifact:
self._extract_texts_from_parts(
parts=artifact["parts"],
path=("artifact", "parts"),
texts_to_check=texts_to_check,
task_mappings=task_mappings,
)
case3: Final = (
self._extract_texts_from_parts(parts=artifact["parts"], path=("artifact", "parts"))
if artifact and isinstance(artifact, dict) and "parts" in artifact
else ()
)
# Case 4: Task with status message
status: Final = result.get("status", {})
if isinstance(status, dict):
status_message: Final = status.get("message")
if status_message and isinstance(status_message, dict) and "parts" in status_message:
self._extract_texts_from_parts(
parts=status_message["parts"],
path=("status", "message", "parts"),
texts_to_check=texts_to_check,
task_mappings=task_mappings,
)
status_message: Final = status.get("message") if isinstance(status, dict) else None
case4: Final = (
self._extract_texts_from_parts(parts=status_message["parts"], path=("status", "message", "parts"))
if isinstance(status_message, dict) and "parts" in status_message
else ()
)
# Case 5: Task with artifacts (plural, array)
artifacts: Final = result.get("artifacts", [])
if artifacts and isinstance(artifacts, list):
for artifact_idx, art in enumerate(artifacts):
if isinstance(art, dict) and "parts" in art:
self._extract_texts_from_parts(
parts=art["parts"],
path=("artifacts", str(artifact_idx), "parts"),
texts_to_check=texts_to_check,
task_mappings=task_mappings,
)
case5: Final = (
tuple(
chain.from_iterable(
self._extract_texts_from_parts(parts=art["parts"], path=("artifacts", str(artifact_idx), "parts"))
for artifact_idx, art in enumerate(artifacts)
if isinstance(art, dict) and "parts" in art
)
)
if artifacts and isinstance(artifacts, list)
else ()
)
return case1 + case2 + case3 + case4 + case5
def _extract_texts_from_parts(
self,
parts: Sequence[_A2ATextPart],
path: tuple[str, ...],
texts_to_check: list[str],
task_mappings: list[tuple[tuple[str, ...], int]],
) -> None:
"""Extract text from message parts."""
for part_idx, part in enumerate(parts):
if part.get("kind") == "text":
text = part.get("text", "")
if text:
texts_to_check.append(text)
task_mappings.append((path, part_idx))
depth: int = 0,
max_depth: int = 10,
) -> tuple[tuple[str, tuple[str, ...], int, str], ...]:
"""
Extract text from message parts, serialized data from data parts, and
recurse into any part that itself carries a nested "parts" list.
Mirrors `extract_text_from_a2a_message`'s handling (including the
recursion depth guard) so the two stay in sync. Returns a tuple of
(text, path_to_parts_list, part_index, field) entries.
"""
if depth >= max_depth:
return ()
def _scan(part_idx: int, part: _A2ATextPart) -> tuple[tuple[str, tuple[str, ...], int, str], ...]:
kind: Final = part.get("kind")
if kind == "text":
text: Final = part.get("text", "")
return ((text, path, part_idx, "text"),) if text else ()
if kind == "data":
part_data: Final = part.get("data")
if part_data is None:
return ()
return ((serialize_a2a_data_part(part_data), path, part_idx, "data"),)
if "parts" in part:
return self._extract_texts_from_parts(
parts=part["parts"],
path=path + (str(part_idx), "parts"),
depth=depth + 1,
max_depth=max_depth,
)
return ()
return tuple(chain.from_iterable(_scan(part_idx, part) for part_idx, part in enumerate(parts)))
def _apply_text_to_path(
self,
result: dict[str | int, Any],
path: tuple[str, ...],
part_idx: int,
field: str,
text: str,
) -> None:
"""Apply guardrailed text back to the specified path in the result."""
@ -463,5 +486,5 @@ class A2AGuardrailHandler(BaseTranslation):
else:
current = current[key]
# Update the text in the part
current[part_idx]["text"] = text
# Update the guardrailed value in the part
current[part_idx][field] = text

1
litellm/llms/a2a/chat/transformation.py
View file

@ -413,6 +413,7 @@ class A2AConfig(BaseConfig):
role: Final = message.get("role", "user")
return {
"kind": "message",
"role": role,
"parts": [{"kind": "text", "text": str(content)}],
"messageId": str(uuid.uuid4()),

21
litellm/llms/a2a/common_utils.py
View file

@ -2,6 +2,7 @@
Common utilities for A2A (Agent-to-Agent) Protocol
"""
import json
from collections.abc import Awaitable, Callable, Mapping
from typing import Any, Final
@ -67,6 +68,19 @@ def convert_messages_to_prompt(messages: list[AllMessageValues]) -> str:
return "\n".join(conversation_parts)
def serialize_a2a_data_part(data: Any) -> str:
"""
Serialize an A2A ``data``-kind part's payload to text.
Used both to build the flattened completion text shown to callers and to
extract guardrail-scannable text, so the two stay in sync.
"""
try:
return json.dumps(data, ensure_ascii=False)
except (TypeError, ValueError):
return str(data)
def extract_text_from_a2a_message(message: dict[str, Any], depth: int = 0, max_depth: int = 10) -> str:
"""
Extract text content from A2A message parts.
@ -86,8 +100,13 @@ def extract_text_from_a2a_message(message: dict[str, Any], depth: int = 0, max_d
text_parts: Final[list[str]] = []
for part in parts:
if part.get("kind") == "text":
kind = part.get("kind")
if kind == "text":
text_parts.append(part.get("text", ""))
elif kind == "data":
data = part.get("data")
if data is not None:
text_parts.append(serialize_a2a_data_part(data))
# Handle nested parts if they exist
elif "parts" in part:
nested_text = extract_text_from_a2a_message(part, depth + 1, max_depth)

5
ruff-strict.toml
View file

@ -35,6 +35,11 @@ external = [
# hook signatures inherit `Any` for `response` from CustomLogger, so narrowing them here
# would break the override rather than describe it.
"litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py" = ["ANN401"]
# `serialize_a2a_data_part`'s `data` param is the payload of an A2A `data`-kind part,
# which the A2A protocol defines as arbitrary JSON-serializable content with no fixed
# shape. `object` was not viable: json.dumps requires the narrower JSON-value types,
# so typing it `object` just moves the Any into a cast at every call site instead.
"litellm/llms/a2a/common_utils.py" = ["ANN401"]
[lint.mccabe]
max-complexity = 15

1
tests/code_coverage_tests/recursive_detector.py
View file

@ -48,6 +48,7 @@ IGNORE_FUNCTIONS = [
"_validate_inheritance_chain", # max depth set (default 100) to prevent infinite recursion in policy inheritance validation.
"_basic_json_schema_validate", # max depth set.
"extract_text_from_a2a_message", # max depth set (default 10) to prevent infinite recursion in A2A message parsing.
"_extract_texts_from_parts", # max depth set (default 10), mirrors extract_text_from_a2a_message's recursion guard.
"_convert_to_json_serializable_dict", # max depth set (default 20) and circular reference protection to prevent infinite recursion.
"dict", # max depth set. _LiteLLMParamsDictView.dict() calls builtin dict(), not itself.
"_read_image_bytes", # max depth set.

63
tests/unit/a2a_protocol/test_regression_issue_28577.py Normal file
View file

@ -0,0 +1,63 @@
from litellm.llms.a2a.chat.transformation import A2AConfig
from litellm.llms.a2a.common_utils import extract_text_from_a2a_message
def test_regression_issue_28577_a2a_discriminator():
"""
Test that A2A transformation adds the mandatory 'kind': 'message' discriminator.
Fixes Bug 1 in #28577.
"""
config = A2AConfig()
messages = [{"role": "user", "content": "ping"}]
# transform_request creates the A2A JSON-RPC payload
request_data = config.transform_request(
model="a2a/demo",
messages=messages,
optional_params={},
litellm_params={},
headers={},
)
# Check Bug 1: message.kind missing
a2a_message = request_data["params"]["message"]
assert a2a_message["kind"] == "message"
assert a2a_message["role"] == "user"
assert "parts" in a2a_message
def test_regression_issue_28577_a2a_data_serialization():
"""
Test that A2A common_utils handle kind: 'data' parts by serializing them.
Fixes Bug 2 in #28577.
"""
message_with_data = {
"kind": "message",
"role": "assistant",
"parts": [{"kind": "data", "data": {"result": {"msg": "pong"}}}],
"messageId": "msg-123",
}
text = extract_text_from_a2a_message(message_with_data)
assert '"result": {"msg": "pong"}' in text
def test_regression_issue_28577_a2a_blocking_param():
"""
Test that A2A requests include configuration.blocking: True.
Fixes Bug 3 in #28577 (async task unblocking).
"""
config = A2AConfig()
messages = [{"role": "user", "content": "ping"}]
request_data = config.transform_request(
model="a2a/demo",
messages=messages,
optional_params={},
litellm_params={},
headers={},
)
# Check Bug 3 fix: configuration.blocking = True
assert "configuration" in request_data["params"]
assert request_data["params"]["configuration"]["blocking"] is True

167
tests/unit/llms/a2a/chat/guardrail_translation/test_handler.py Normal file
View file

@ -0,0 +1,167 @@
"""
Unit tests for A2A Protocol Guardrail Translation Handler
Regression coverage for the "data"-kind part guardrail bypass: A2A responses
can carry structured content in `kind: "data"` parts, which
`extract_text_from_a2a_message` (used to build the completion text callers
see) folds into the final text, but the guardrail handler previously only
inspected `kind: "text"` parts, so guarded output checks were skipped for
that content path.
"""
from typing import Any, Literal, Optional
import pytest
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.llms.a2a.chat.guardrail_translation.handler import A2AGuardrailHandler
from litellm.types.utils import GenericGuardrailAPIInputs
class MockGuardrail(CustomGuardrail):
"""Mock guardrail that uppercases text so we can assert exactly what was scanned and where the result landed."""
def __init__(self, guardrail_name: str = "test"):
super().__init__(guardrail_name=guardrail_name)
self.last_inputs: Optional[GenericGuardrailAPIInputs] = None
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional[Any] = None,
) -> GenericGuardrailAPIInputs:
self.last_inputs = inputs
texts = inputs.get("texts", [])
return {"texts": [text.upper() for text in texts]}
@pytest.mark.asyncio
async def test_process_output_response_scans_data_parts():
"""A `kind: data` part in the output must be sent to the guardrail and the
guardrailed value written back into `data`, not silently skipped."""
handler = A2AGuardrailHandler()
guardrail = MockGuardrail()
response = {
"result": {
"kind": "message",
"parts": [
{"kind": "text", "text": "hello"},
{"kind": "data", "data": {"secret": "leak-me"}},
],
}
}
result = await handler.process_output_response(
response=response,
guardrail_to_apply=guardrail,
)
# The data part's serialized content must have reached the guardrail.
assert guardrail.last_inputs is not None
scanned_texts = guardrail.last_inputs["texts"]
assert any("leak-me" in t for t in scanned_texts)
# The guardrailed (uppercased) value must be written back into "data",
# and the part must remain a "data" part, not be silently dropped or
# converted into an unguarded pass-through.
data_part = result["result"]["parts"][1]
assert data_part["kind"] == "data"
assert "LEAK-ME" in data_part["data"]
# The text part must still be guardrailed as before (no regression).
text_part = result["result"]["parts"][0]
assert text_part["text"] == "HELLO"
@pytest.mark.asyncio
async def test_process_output_response_data_only_still_scanned():
"""A response with ONLY a data part (no text parts at all) must not be
skipped as "no text content in response"."""
handler = A2AGuardrailHandler()
guardrail = MockGuardrail()
response = {
"result": {
"kind": "message",
"parts": [{"kind": "data", "data": {"result": {"msg": "pong"}}}],
}
}
result = await handler.process_output_response(
response=response,
guardrail_to_apply=guardrail,
)
assert guardrail.last_inputs is not None
assert guardrail.last_inputs["texts"]
assert "PONG" in result["result"]["parts"][0]["data"]
@pytest.mark.asyncio
async def test_process_output_response_scans_nested_parts():
"""A part that itself carries a nested "parts" list (grouping sub-parts)
must be recursed into, matching extract_text_from_a2a_message's own
recursion, instead of being silently skipped as neither text nor data."""
handler = A2AGuardrailHandler()
guardrail = MockGuardrail()
response = {
"result": {
"kind": "message",
"parts": [
{
"kind": "group",
"parts": [
{"kind": "text", "text": "hello"},
{"kind": "data", "data": {"secret": "leak-me"}},
],
},
],
}
}
result = await handler.process_output_response(
response=response,
guardrail_to_apply=guardrail,
)
assert guardrail.last_inputs is not None
scanned_texts = guardrail.last_inputs["texts"]
assert "hello" in scanned_texts
assert any("leak-me" in t for t in scanned_texts)
nested_parts = result["result"]["parts"][0]["parts"]
assert nested_parts[0]["text"] == "HELLO"
assert "LEAK-ME" in nested_parts[1]["data"]
@pytest.mark.asyncio
async def test_process_input_messages_scans_data_parts():
"""The same bypass existed on the request/input side of the handler."""
handler = A2AGuardrailHandler()
guardrail = MockGuardrail()
data = {
"params": {
"message": {
"kind": "message",
"role": "user",
"parts": [{"kind": "data", "data": {"secret": "leak-me"}}],
}
}
}
result = await handler.process_input_messages(
data=data,
guardrail_to_apply=guardrail,
)
assert guardrail.last_inputs is not None
assert any("leak-me" in t for t in guardrail.last_inputs["texts"])
data_part = result["params"]["message"]["parts"][0]
assert data_part["kind"] == "data"
assert "LEAK-ME" in data_part["data"]

23
tests/unit/responses/test_regression_issue_28553.py Normal file
View file

@ -0,0 +1,23 @@
import pytest
from litellm.proxy.common_request_processing import _stream_usage_tracking_updates
@pytest.mark.parametrize(
"route_type,expect_injected",
[
("acompletion", True),
("atext_completion", False),
("aresponses", False),
("aget_responses", False),
("arealtime", False),
],
)
def test_stream_usage_only_injected_for_whitelisted_route_types(route_type, expect_injected):
updates = _stream_usage_tracking_updates(
data={"stream": True, "model": "gpt-4o"},
general_settings={},
route_type=route_type,
supports_stream_options=lambda: True,
)
assert ("stream_options" in updates) is expect_injected