scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 86

fix(a2a): rewrite guardrail text extraction functionally to satisfy LIT002 budget

Browse source 
_extract_texts_from_result/_extract_texts_from_parts now return tuples of
scanned entries instead of mutating shared out-parameter lists, matching
the repo's no-mutation convention and staying under the type-discipline
gate's mutable-collection-construction budget. Two GenericGuardrailAPIInputs
call sites convert the resulting tuple to a list at that external boundary
(marked mutable-ok, since the TypedDict itself is typed List[str]).

Also allowlists the now-legitimately-recursive _extract_texts_from_parts
in the CI's recursive-function detector, matching the pattern already used
for its sibling extract_text_from_a2a_message (bounded max_depth=10 guard).
This commit is contained in:
STiFLeR7 2026-08-04 10:21:36 +05:30
parent a93b5ee3f2
commit 050dcb166d
No known key found for this signature in database
2 changed files with 61 additions and 86 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -82,28 +82,30 @@ class A2AGuardrailHandler(BaseTranslation):
verbose_proxy_logger.debug("A2A: No parts in message, skipping guardrail")
return data
texts_to_check: Final[list[str]] = []
# Track which parts contain scannable content, and which field to write
# the guardrailed value back to ("text" or "data")
part_mappings: Final[list[tuple[int, str]]] = []
# Step 1: Extract text from all text parts, and serialized data from all data parts
for part_idx, part in enumerate(parts):
def _scan_input_part(part_idx: int, part: dict[str, Any]) -> tuple[str, int, str] | None:
kind = part.get("kind")
if kind == "text":
text = part.get("text", "")
if text:
texts_to_check.append(text)
part_mappings.append((part_idx, "text"))
elif kind == "data":
return (text, part_idx, "text") if text else None
if kind == "data":
part_data = part.get("data")
if part_data is not None:
texts_to_check.append(serialize_a2a_data_part(part_data))
part_mappings.append((part_idx, "data"))
return (serialize_a2a_data_part(part_data), part_idx, "data") if part_data is not None else None
return None
# 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) # mutable-ok: GenericGuardrailAPIInputs.texts is typed List[str]
)
# Pass the structured A2A message to guardrails
inputs["structured_messages"] = [message]
@ -173,18 +175,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, str]]] = []
# 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")
@ -205,7 +201,9 @@ 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) # mutable-ok: GenericGuardrailAPIInputs.texts is typed List[str]
)
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
@ -295,18 +293,12 @@ class A2AGuardrailHandler(BaseTranslation):
result = obj.get("result", {})
if not isinstance(result, dict):
continue
texts_in_chunk: Final[list[str]] = []
mappings: Final[list[tuple[tuple[str, ...], int, str]]] = []
self._extract_texts_from_result(
result=result,
texts_to_check=texts_in_chunk,
task_mappings=mappings,
)
if not mappings:
scanned: Final = 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, field) 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,
@ -316,7 +308,7 @@ class A2AGuardrailHandler(BaseTranslation):
text=text,
)
else:
for path, part_idx, field in mappings:
for _, path, part_idx, field in scanned:
self._apply_text_to_path(
result=result,
path=path,
@ -378,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, str]],
) -> None:
) -> tuple[tuple[str, tuple[str, ...], int, str], ...]:
"""
Extract text from all possible locations in an A2A result.
@ -390,46 +380,32 @@ 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.
"""
entries: tuple[tuple[str, tuple[str, ...], int, str], ...] = ()
# 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,
)
entries += self._extract_texts_from_parts(parts=result["parts"], path=("parts",))
# 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,
)
entries += self._extract_texts_from_parts(parts=message["parts"], path=("message", "parts"))
# 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,
)
entries += self._extract_texts_from_parts(parts=artifact["parts"], path=("artifact", "parts"))
# 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,
entries += self._extract_texts_from_parts(
parts=status_message["parts"], path=("status", "message", "parts")
)
# Case 5: Task with artifacts (plural, array)
@ -437,52 +413,50 @@ class A2AGuardrailHandler(BaseTranslation):
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,
entries += self._extract_texts_from_parts(
parts=art["parts"], path=("artifacts", str(artifact_idx), "parts")
)
return entries
def _extract_texts_from_parts(
self,
parts: Sequence[_A2ATextPart],
path: tuple[str, ...],
texts_to_check: list[str],
task_mappings: list[tuple[tuple[str, ...], int, str]],
depth: int = 0,
max_depth: int = 10,
) -> None:
) -> 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.
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
for part_idx, part in enumerate(parts):
return ()
def _scan(part_idx: int, part: dict[str, Any]) -> tuple[tuple[str, tuple[str, ...], int, str], ...]:
kind = part.get("kind")
if kind == "text":
text = part.get("text", "")
if text:
texts_to_check.append(text)
task_mappings.append((path, part_idx, "text"))
elif kind == "data":
return ((text, path, part_idx, "text"),) if text else ()
if kind == "data":
part_data = part.get("data")
if part_data is not None:
texts_to_check.append(serialize_a2a_data_part(part_data))
task_mappings.append((path, part_idx, "data"))
elif "parts" in part:
self._extract_texts_from_parts(
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"),
texts_to_check=texts_to_check,
task_mappings=task_mappings,
depth=depth + 1,
max_depth=max_depth,
)
return ()
return tuple(entry for part_idx, part in enumerate(parts) for entry in _scan(part_idx, part))
def _apply_text_to_path(
self,

1
tests/code_coverage_tests/recursive_detector.py
View file

@ -43,6 +43,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.