diff --git a/litellm/llms/a2a/chat/guardrail_translation/handler.py b/litellm/llms/a2a/chat/guardrail_translation/handler.py index 71b4f1f3b54..5104d07c9e3 100644 --- a/litellm/llms/a2a/chat/guardrail_translation/handler.py +++ b/litellm/llms/a2a/chat/guardrail_translation/handler.py @@ -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, diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index e9f87ba6cae..01f97efed05 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -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.