diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index cc0b260998f..06a1b92e1b0 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -13,7 +13,7 @@ Pattern Overview: """ import asyncio -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, cast +from typing import TYPE_CHECKING, Any, Coroutine, Dict, List, Optional, Tuple, cast from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation @@ -49,7 +49,7 @@ class AnthropicMessagesHandler(BaseTranslation): if messages is None: return data - tasks = [] + tasks: List[Coroutine[Any, Any, str]] = [] task_mappings: List[Tuple[int, Optional[int]]] = [] # Track (message_index, content_index) for each task # content_index is None for string content, int for list content @@ -166,7 +166,7 @@ class AnthropicMessagesHandler(BaseTranslation): ) return response - tasks = [] + tasks: List[Coroutine[Any, Any, str]] = [] task_mappings: List[Tuple[int, Optional[int]]] = [] # Track (choice_index, content_index) for each task diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 12d8cfad23a..ec25f491d8e 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -15,7 +15,7 @@ This pattern can be replicated for other message formats (e.g., Anthropic). """ import asyncio -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, cast +from typing import TYPE_CHECKING, Any, Coroutine, Dict, List, Optional, Tuple, cast import litellm from litellm._logging import verbose_proxy_logger @@ -50,7 +50,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if messages is None: return data - tasks = [] + tasks: List[Coroutine[Any, Any, str]] = [] task_mappings: List[Tuple[int, Optional[int]]] = [] # Track (message_index, content_index) for each task # content_index is None for string content, int for list content @@ -168,7 +168,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): ) return response - tasks = [] + tasks: List[Coroutine[Any, Any, str]] = [] task_mappings: List[Tuple[int, Optional[int]]] = [] # Track (choice_index, content_index) for each task diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 40b1d73ddcb..fdac13176b1 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -29,7 +29,7 @@ Output: response.output is List[GenericResponseOutputItem] where each has: """ import asyncio -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast +from typing import TYPE_CHECKING, Any, Coroutine, List, Optional, Tuple, Union, cast from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation @@ -79,7 +79,7 @@ class OpenAIResponsesHandler(BaseTranslation): if not isinstance(input_data, list): return data - tasks = [] + tasks: List[Coroutine[Any, Any, str]] = [] task_mappings: List[Tuple[int, Optional[int]]] = [] # Track (message_index, content_index) for each task # content_index is None for string content, int for list content @@ -113,9 +113,9 @@ class OpenAIResponsesHandler(BaseTranslation): async def _extract_input_text_and_create_tasks( self, - message: Dict[str, Any], + message: Any, # Can be Dict[str, Any] or ResponseInputParam msg_idx: int, - tasks: List, + tasks: List[Coroutine[Any, Any, str]], task_mappings: List[Tuple[int, Optional[int]]], guardrail_to_apply: "CustomGuardrail", ) -> None: @@ -144,7 +144,7 @@ class OpenAIResponsesHandler(BaseTranslation): async def _apply_guardrail_responses_to_input( self, - messages: List[Dict[str, Any]], + messages: Any, # Can be List[Dict[str, Any]] or ResponseInputParam responses: List[str], task_mappings: List[Tuple[int, Optional[int]]], ) -> None: @@ -200,7 +200,7 @@ class OpenAIResponsesHandler(BaseTranslation): ) return response - tasks = [] + tasks: List[Coroutine[Any, Any, str]] = [] task_mappings: List[Tuple[int, int]] = [] # Track (output_item_index, content_index) for each task diff --git a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py index 00e30d3e5ba..fb93dec47b7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py +++ b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py @@ -216,7 +216,7 @@ class GraySwanGuardrail(CustomGuardrail): verbose_proxy_logger.debug("GraySwan Guardrail: post-call hook triggered") - response_dict = response.model_dump() if hasattr(response, "model_dump") else {} + response_dict = response.model_dump() if hasattr(response, "model_dump") else {} # type: ignore[union-attr] response_messages = [ msg if isinstance(msg, dict) else msg.model_dump() for choice in response_dict.get("choices", []) diff --git a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py index 537b6038b53..1646e39b24d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py +++ b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py @@ -284,7 +284,7 @@ class PillarGuardrail(CustomGuardrail): verbose_proxy_logger.debug("Pillar Guardrail: Post-call hook") # Extract response messages in the format Pillar expects - response_dict = response.model_dump() if hasattr(response, "model_dump") else {} + response_dict = response.model_dump() if hasattr(response, "model_dump") else {} # type: ignore[union-attr] response_messages = [ choice.get("message") for choice in response_dict.get("choices", []) diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 216de57769e..c59e91aa1ef 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -926,6 +926,7 @@ async def test_model_connection( "batch", "rerank", "realtime", + "responses", "ocr", ] ] = fastapi.Body("chat", description="The mode to test the model with"),