From 8e2a91c0c15e296f79cd19f6f40dcddd67eb87eb Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 17 Dec 2025 22:16:46 +0530 Subject: [PATCH] fix:mypy errors for litellm_staging_12_17_2025 --- .../chat/guardrail_translation/handler.py | 29 ++-- .../guardrail_translation/handler.py | 1 + .../guardrail_hooks/grayswan/grayswan.py | 143 +++++++++++++++++- .../proxy/response_api_endpoints/endpoints.py | 12 +- 4 files changed, 163 insertions(+), 22 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index b12dee1e3a..6f53bf6571 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -254,10 +254,12 @@ class AnthropicMessagesHandler(BaseTranslation): # Track (content_index, None) for each text # Handle both dict and object responses - if hasattr(response, "get"): - response_content = response.get("content", []) + response_content: List[Any] = [] + if isinstance(response, dict): + response_content = response.get("content", []) or [] elif hasattr(response, "content"): - response_content = response.content or [] + content = getattr(response, "content", None) + response_content = content or [] else: response_content = [] @@ -267,9 +269,10 @@ class AnthropicMessagesHandler(BaseTranslation): # Step 1: Extract all text content and tool calls from response 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 = content_block + 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 @@ -282,7 +285,7 @@ class AnthropicMessagesHandler(BaseTranslation): if block_type in ["text", "tool_use"]: self._extract_output_text_and_images( - content_block=cast(Dict[str, Any], block_dict), + content_block=block_dict, content_idx=content_idx, texts_to_check=texts_to_check, images_to_check=images_to_check, @@ -546,7 +549,11 @@ class AnthropicMessagesHandler(BaseTranslation): Override this method to customize text content detection. """ - response_content = response.get("content", []) + if isinstance(response, dict): + response_content = response.get("content", []) + else: + response_content = getattr(response, "content", None) or [] + if not response_content: return False for content_block in response_content: @@ -607,10 +614,12 @@ class AnthropicMessagesHandler(BaseTranslation): content_idx = cast(int, mapping[0]) # Handle both dict and object responses - if hasattr(response, "get"): - response_content = response.get("content", []) + response_content: List[Any] = [] + if isinstance(response, dict): + response_content = response.get("content", []) or [] elif hasattr(response, "content"): - response_content = response.content or [] + content = getattr(response, "content", None) + response_content = content or [] else: continue @@ -627,7 +636,7 @@ class AnthropicMessagesHandler(BaseTranslation): # Handle both dict and Pydantic object content blocks if isinstance(content_block, dict): if content_block.get("type") == "text": - content_block["text"] = guardrail_response + cast(Dict[str, Any], content_block)["text"] = guardrail_response elif hasattr(content_block, "type") and getattr(content_block, "type", None) == "text": # Update Pydantic object's text attribute if hasattr(content_block, "text"): diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 02899a6703..9b8f15c762 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -30,6 +30,7 @@ Output: response.output is List[GenericResponseOutputItem] where each has: from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast +from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall from pydantic import BaseModel from litellm._logging import verbose_proxy_logger diff --git a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py index 4e9dd63d41..38e75ebabb 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py +++ b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py @@ -6,7 +6,10 @@ from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional from fastapi import HTTPException from litellm._logging import verbose_proxy_logger -from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + ModifyResponseException, +) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, @@ -204,7 +207,7 @@ class GraySwanGuardrail(CustomGuardrail): response_json = await self._call_grayswan_api(payload) # Process response is_output = input_type == "response" - result = self._process_response( + result = self._process_response_internal( response_json=response_json, request_data=request_data, inputs=inputs, @@ -213,6 +216,126 @@ class GraySwanGuardrail(CustomGuardrail): return result + # ------------------------------------------------------------------ + # Legacy Test Interface (for backward compatibility) + # ------------------------------------------------------------------ + + async def run_grayswan_guardrail(self, payload: dict) -> Dict[str, Any]: + """ + Run the GraySwan guardrail on a payload. + + This is a legacy method for testing purposes. + + Args: + payload: The payload to scan + + Returns: + Dict containing the GraySwan API response + """ + response_json = await self._call_grayswan_api(payload) + # Call the legacy response processor (for test compatibility) + self._process_grayswan_response(response_json) + return response_json + + def _process_grayswan_response( + self, + response_json: dict, + data: Optional[dict] = None, + hook_type: Optional[GuardrailEventHooks] = None, + ) -> None: + """ + Legacy method for processing GraySwan API responses. + + This method is maintained for backward compatibility with existing tests. + It handles the test scenarios where responses need to be processed with + knowledge of the request context (pre/during/post call hooks). + + Args: + response_json: Response from GraySwan API + data: Optional request data (for passthrough exceptions) + hook_type: Optional GuardrailEventHooks for determining behavior + """ + violation_score = float(response_json.get("violation", 0.0) or 0.0) + violated_rules = response_json.get("violated_rules", []) + mutation_detected = response_json.get("mutation") + ipi_detected = response_json.get("ipi") + + flagged = violation_score >= self.violation_threshold + if not flagged: + verbose_proxy_logger.debug( + "Gray Swan Guardrail: content passed (score=%s, threshold=%s)", + violation_score, + self.violation_threshold, + ) + return + + verbose_proxy_logger.warning( + "Gray Swan Guardrail: violation score %.3f exceeds threshold %.3f", + violation_score, + self.violation_threshold, + ) + + detection_info = { + "guardrail": "grayswan", + "flagged": True, + "violation_score": violation_score, + "violated_rules": violated_rules, + "mutation": mutation_detected, + "ipi": ipi_detected, + } + + # Determine if this is input (pre-call/during-call) or output (post-call) + if hook_type is not None: + is_input = hook_type in [ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.during_call, + ] + else: + is_input = True + + if self.on_flagged_action == "block": + violation_location = "output" if (not is_input) else "input" + raise HTTPException( + status_code=400, + detail={ + "error": "Blocked by Gray Swan Guardrail", + "violation_location": violation_location, + "violation": violation_score, + "violated_rules": violated_rules, + "mutation": mutation_detected, + "ipi": ipi_detected, + }, + ) + elif self.on_flagged_action == "passthrough": + # For passthrough mode, we need to handle violations + detections = [detection_info] + violation_message = self._format_violation_message( + detections, is_output=not is_input + ) + verbose_proxy_logger.info( + "Gray Swan Guardrail: Passthrough mode - handling violation" + ) + + # If hook_type is provided and in pre/during call, raise exception + if hook_type in [GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call]: + # Raise ModifyResponseException to short-circuit LLM call + if data is None: + data = {} + self.raise_passthrough_exception( + violation_message=violation_message, + request_data=data, + detection_info=detection_info, + ) + elif hook_type == GuardrailEventHooks.post_call: + # For post-call, store detection info in metadata + if data is None: + data = {} + if "metadata" not in data: + data["metadata"] = {} + if "guardrail_detections" not in data["metadata"]: + data["metadata"]["guardrail_detections"] = [] + data["metadata"]["guardrail_detections"].append(detection_info) + # ------------------------------------------------------------------ # Core GraySwan API interaction # ------------------------------------------------------------------ @@ -242,7 +365,7 @@ class GraySwanGuardrail(CustomGuardrail): ) raise GraySwanGuardrailAPIError(str(exc)) from exc - def _process_response( + def _process_response_internal( self, response_json: Dict[str, Any], request_data: dict, @@ -366,18 +489,24 @@ class GraySwanGuardrail(CustomGuardrail): return payload def _format_violation_message( - self, detection_info: dict, is_output: bool = False + self, detection_info: Any, is_output: bool = False ) -> str: """ Format detection info into a user-friendly violation message. Args: - detection_info: Detection info dictionary + detection_info: Can be either: + - A single dict with violation_score, violated_rules, mutation, ipi keys + - A list of such dicts (legacy format) is_output: True if violation is in model output, False if in input Returns: Formatted violation message string """ + # Handle legacy format where detection_info is a list + if isinstance(detection_info, list) and len(detection_info) > 0: + detection_info = detection_info[0] + violation_score = detection_info.get("violation_score", 0.0) violated_rules = detection_info.get("violated_rules", []) mutation = detection_info.get("mutation", False) @@ -397,12 +526,12 @@ class GraySwanGuardrail(CustomGuardrail): if mutation: message_parts.append( - "A potential prompt manipulation/mutation was detected." + "Mutation effort to make the harmful intention disguised was DETECTED." ) if ipi: message_parts.append( - "Indirect prompt injection indicators were detected." + "Indirect Prompt Injection was DETECTED." ) return "\n".join(message_parts) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 50eb729d8e..252b3a7d38 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1,7 +1,7 @@ import asyncio import time -import uuid from typing import Any, AsyncIterator, cast +from uuid import uuid4 from fastapi import APIRouter, Depends, HTTPException, Request, Response @@ -10,7 +10,7 @@ from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing -from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse from litellm.types.responses.main import DeleteResponseResult router = APIRouter() @@ -184,13 +184,15 @@ async def responses_api( violation_text = e.message response_obj = ResponsesAPIResponse( - id=f"resp_{uuid.uuid4()}", + id=f"resp_{uuid4()}", object="response", created_at=int(time.time()), model=e.model or data.get("model"), - output=[{"content": [{"type": "text", "text": violation_text}]}], + output=cast(Any, [{"content": [{"type": "text", "text": violation_text}]}]), status="completed", - usage={"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}, + usage=ResponseAPIUsage( + input_tokens=0, output_tokens=0, total_tokens=0 + ), ) return response_obj except Exception as e: