From 91ac86ec17ae67f6c9f975dcbd5ab7d8cbf10fff Mon Sep 17 00:00:00 2001 From: joereyna Date: Tue, 24 Mar 2026 10:40:44 -0700 Subject: [PATCH] refactor: extract helper methods to fix PLR0915 in guardrail handlers Both process_output_streaming_response (a2a) and process_output_response (anthropic) exceeded ruff's PLR0915 limit of 50 statements (~54 each). Root cause: lint-dev skips ruff entirely, so these slipped past local and CI checks on the dev path. - a2a handler: extract JSON-parse loop into _parse_a2a_response_items() - anthropic handler: extract inputs-builder block into _build_output_guardrail_inputs() Both methods now sit at ~44 statements, well under the 50-statement limit. --- .../a2a/chat/guardrail_translation/handler.py | 45 ++--- .../chat/guardrail_translation/handler.py | 189 ++++++++---------- 2 files changed, 101 insertions(+), 133 deletions(-) diff --git a/litellm/llms/a2a/chat/guardrail_translation/handler.py b/litellm/llms/a2a/chat/guardrail_translation/handler.py index 3d6037b1f8f..3fb432efa63 100644 --- a/litellm/llms/a2a/chat/guardrail_translation/handler.py +++ b/litellm/llms/a2a/chat/guardrail_translation/handler.py @@ -234,13 +234,25 @@ class A2AGuardrailHandler(BaseTranslation): 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). """ - parsed, valid_parsed = self._parse_streaming_responses(responses_so_far) + 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 = self._parse_a2a_response_items(responses_so_far) + + valid_parsed = [(i, obj) for i, obj in enumerate(parsed) if obj is not None] if not valid_parsed: return responses_so_far - combined_text, chunk_indices_with_text = self._collect_text_from_parsed_chunks( - valid_parsed - ) + # 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 @@ -313,11 +325,11 @@ class A2AGuardrailHandler(BaseTranslation): return responses_so_far - def _parse_streaming_responses( + def _parse_a2a_response_items( self, responses_so_far: List[Any], - ) -> Tuple[List[Optional[Dict[str, Any]]], List[Tuple[int, Dict[str, Any]]]]: - """Parse JSON-RPC items, returning aligned parsed list and valid entries.""" + ) -> List[Optional[Dict[str, Any]]]: + """Parse each item in responses_so_far into dicts, keeping index alignment.""" parsed: List[Optional[Dict[str, Any]]] = [None] * len(responses_so_far) for i, item in enumerate(responses_so_far): if isinstance(item, dict): @@ -331,24 +343,7 @@ class A2AGuardrailHandler(BaseTranslation): 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] - return parsed, valid_parsed - - def _collect_text_from_parsed_chunks( - self, - valid_parsed: List[Tuple[int, Dict[str, Any]]], - ) -> Tuple[str, List[int]]: - """Collect text from parsed chunks, returning combined text and indices.""" - from litellm.llms.a2a.common_utils import extract_text_from_a2a_response - - text_parts: List[str] = [] - chunk_indices_with_text: List[int] = [] - 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) - return "".join(text_parts), chunk_indices_with_text + return parsed def _extract_texts_from_result( self, diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 2d1ca4b6e30..1c36e21de4e 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -294,34 +294,74 @@ class AnthropicMessagesHandler(BaseTranslation): images_to_check: List[str] = [] tool_calls_to_check: List[ChatCompletionToolCallChunk] = [] task_mappings: List[Tuple[int, Optional[int]]] = [] + # Track (content_index, None) for each text + + # Handle both dict and object responses + response_content: List[Any] = [] + if isinstance(response, dict): + response_content = response.get("content", []) or [] + elif hasattr(response, "content"): + content = getattr(response, "content", None) + response_content = content or [] + else: + response_content = [] - response_content = self._get_response_content(response) if not response_content: return response # Step 1: Extract all text content and tool calls from response - self._extract_from_content_blocks( - response_content, - texts_to_check, - images_to_check, - task_mappings, - tool_calls_to_check, - ) + for content_idx, content_block in enumerate(response_content): + # Handle both dict and Pydantic object content blocks + block_dict: Dict[str, Any] = {} + if isinstance(content_block, dict): + block_type = content_block.get("type") + block_dict = cast(Dict[str, Any], content_block) + elif hasattr(content_block, "type"): + block_type = getattr(content_block, "type", None) + # Convert Pydantic object to dict for processing + if hasattr(content_block, "model_dump"): + block_dict = content_block.model_dump() + else: + block_dict = { + "type": block_type, + "text": getattr(content_block, "text", None), + } + else: + continue + + if block_type in ["text", "tool_use"]: + self._extract_output_text_and_images( + content_block=block_dict, + content_idx=content_idx, + texts_to_check=texts_to_check, + images_to_check=images_to_check, + task_mappings=task_mappings, + tool_calls_to_check=tool_calls_to_check, + ) # Step 2: Apply guardrail to all texts in batch if texts_to_check or tool_calls_to_check: - request_data = self._prepare_request_data( - request_data, - response, - user_api_key_dict, - key="response", - ) + # Use the real request_data if provided (proxy path), otherwise + # create a standalone dict (SDK / direct-call path). + if request_data is None: + request_data = {"response": response} + else: + if "response" not in request_data: + request_data["response"] = response - inputs = self._build_guardrail_inputs( - texts_to_check, - images_to_check, - tool_calls_to_check, - response, + # Add user API key metadata with prefixed keys + if "litellm_metadata" not in request_data: + user_metadata = self.transform_user_api_key_dict_to_metadata( + user_api_key_dict + ) + if user_metadata: + request_data["litellm_metadata"] = user_metadata + + inputs = self._build_output_guardrail_inputs( + response=response, + texts_to_check=texts_to_check, + images_to_check=images_to_check, + tool_calls_to_check=tool_calls_to_check, ) guardrailed_inputs = await guardrail_to_apply.apply_guardrail( @@ -346,6 +386,28 @@ class AnthropicMessagesHandler(BaseTranslation): return response + def _build_output_guardrail_inputs( + self, + response: Any, + texts_to_check: List[str], + images_to_check: List[str], + tool_calls_to_check: List[ChatCompletionToolCallChunk], + ) -> GenericGuardrailAPIInputs: + """Build guardrail inputs dict, including model info extracted from the response.""" + inputs = GenericGuardrailAPIInputs(texts=texts_to_check) + if images_to_check: + inputs["images"] = images_to_check + if tool_calls_to_check: + inputs["tool_calls"] = tool_calls_to_check + response_model = None + if isinstance(response, dict): + response_model = response.get("model") + elif hasattr(response, "model"): + response_model = getattr(response, "model", None) + if response_model: + inputs["model"] = response_model + return inputs + async def process_output_streaming_response( self, responses_so_far: List[Any], @@ -410,95 +472,6 @@ class AnthropicMessagesHandler(BaseTranslation): ) return responses_so_far - def _prepare_request_data( - self, - request_data: Optional[dict], - response: Any, - user_api_key_dict: Optional[Any], - key: str, - ) -> dict: - """Ensure request_data has the response/responses_so_far key and metadata.""" - if request_data is None: - request_data = {key: response} - else: - if key not in request_data: - request_data[key] = response - - if "litellm_metadata" not in request_data: - user_metadata = self.transform_user_api_key_dict_to_metadata( - user_api_key_dict - ) - if user_metadata: - request_data["litellm_metadata"] = user_metadata - return request_data - - @staticmethod - def _get_response_content(response: Any) -> List[Any]: - """Extract content list from a dict or object response.""" - if isinstance(response, dict): - return response.get("content", []) or [] - elif hasattr(response, "content"): - return getattr(response, "content", None) or [] - return [] - - def _extract_from_content_blocks( - self, - response_content: List[Any], - texts_to_check: List[str], - images_to_check: List[str], - task_mappings: List[Tuple[int, Optional[int]]], - tool_calls_to_check: List["ChatCompletionToolCallChunk"], - ) -> None: - """Extract text, images, and tool calls from content blocks.""" - for content_idx, content_block in enumerate(response_content): - block_dict: Dict[str, Any] = {} - if isinstance(content_block, dict): - block_type = content_block.get("type") - block_dict = cast(Dict[str, Any], content_block) - elif hasattr(content_block, "type"): - block_type = getattr(content_block, "type", None) - if hasattr(content_block, "model_dump"): - block_dict = content_block.model_dump() - else: - block_dict = { - "type": block_type, - "text": getattr(content_block, "text", None), - } - else: - continue - - if block_type in ["text", "tool_use"]: - self._extract_output_text_and_images( - content_block=block_dict, - content_idx=content_idx, - texts_to_check=texts_to_check, - images_to_check=images_to_check, - task_mappings=task_mappings, - tool_calls_to_check=tool_calls_to_check, - ) - - @staticmethod - def _build_guardrail_inputs( - texts_to_check: List[str], - images_to_check: List[str], - tool_calls_to_check: List["ChatCompletionToolCallChunk"], - response: Any, - ) -> "GenericGuardrailAPIInputs": - """Build GenericGuardrailAPIInputs with optional images, tool calls, model.""" - inputs = GenericGuardrailAPIInputs(texts=texts_to_check) - if images_to_check: - inputs["images"] = images_to_check - if tool_calls_to_check: - inputs["tool_calls"] = tool_calls_to_check - response_model = None - if isinstance(response, dict): - response_model = response.get("model") - elif hasattr(response, "model"): - response_model = getattr(response, "model", None) - if response_model: - inputs["model"] = response_model - return inputs - def get_streaming_string_so_far(self, responses_so_far: List[Any]) -> str: """ Parse streaming responses and extract accumulated text content.