mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
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.
This commit is contained in:
joereyna
parent
c0fc4c4234
commit
91ac86ec17
2 changed files with 101 additions and 133 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue