diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 632743236c4..2585e9ea621 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,15 +1,15 @@ { "reportAny": { - "limit": 28842 + "limit": 28624 }, "reportArgumentType": { - "limit": 2634 + "limit": 2488 }, "reportAssignmentType": { "limit": 329 }, "reportAttributeAccessIssue": { - "limit": 514 + "limit": 513 }, "reportCallIssue": { "limit": 117 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 9103 + "limit": 9076 }, "reportFunctionMemberAccess": { "limit": 7 @@ -57,7 +57,7 @@ "limit": 5843 }, "reportMissingTypeArgument": { - "limit": 15816 + "limit": 15806 }, "reportMissingTypeStubs": { "limit": 40 @@ -72,7 +72,7 @@ "limit": 0 }, "reportOptionalMemberAccess": { - "limit": 1078 + "limit": 417 }, "reportOptionalOperand": { "limit": 0 @@ -90,7 +90,7 @@ "limit": 8 }, "reportReturnType": { - "limit": 218 + "limit": 205 }, "reportTypedDictNotRequiredAccess": { "limit": 27 @@ -105,16 +105,16 @@ "limit": 113 }, "reportUnknownMemberType": { - "limit": 39838 + "limit": 39831 }, "reportUnknownParameterType": { - "limit": 20237 + "limit": 20235 }, "reportUnknownVariableType": { - "limit": 31383 + "limit": 31358 }, "reportUnnecessaryCast": { - "limit": 122 + "limit": 120 }, "reportUnnecessaryComparison": { "limit": 701 diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py index 339da998d56..78cd7a03b36 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py @@ -6,16 +6,37 @@ This module provides fake streaming by converting non-streaming responses into s """ import asyncio -from collections.abc import AsyncIterator -from typing import Any, Final, cast +from collections.abc import AsyncIterator, Mapping +from typing import Any, Final, TypeAlias, cast from uuid import uuid4 +from pydantic import BaseModel, JsonValue, TypeAdapter + from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, get_async_httpx_client, ) +JsonObject: TypeAlias = dict[str, JsonValue] + +_JSON_OBJECT_ADAPTER: Final = TypeAdapter(JsonObject) + + +def _object_field(data: Mapping[str, JsonValue], key: str) -> JsonObject: + value: Final = data.get(key) + return value if isinstance(value, dict) else {} + + +def _array_field(data: Mapping[str, JsonValue], key: str) -> list[JsonValue]: + value: Final = data.get(key) + return value if isinstance(value, list) else [] + + +def _string_field(data: Mapping[str, JsonValue], key: str, default: str = "") -> str: + value: Final = data.get(key) + return value if isinstance(value, str) else default + class PydanticAITransformation: """ @@ -28,7 +49,7 @@ class PydanticAITransformation: """ @staticmethod - def _remove_none_values(obj: Any) -> Any: + def _remove_none_values(obj: JsonValue) -> JsonValue: """ Recursively remove None values from a dict/list structure. @@ -49,7 +70,7 @@ class PydanticAITransformation: return obj @staticmethod - def _params_to_dict(params: Any) -> dict[str, Any]: + def _params_to_dict(params: object) -> JsonObject: """ Convert params to a dict, handling Pydantic models. @@ -59,17 +80,9 @@ class PydanticAITransformation: Returns: Dict representation of params """ - if hasattr(params, "model_dump"): - # Pydantic v2 model - return params.model_dump(mode="python", exclude_none=True) - elif hasattr(params, "dict"): - # Pydantic v1 model - return params.dict(exclude_none=True) - elif isinstance(params, dict): - return params - else: - # Try to convert to dict - return dict(params) + if isinstance(params, BaseModel): + return _JSON_OBJECT_ADAPTER.validate_python(params.model_dump(mode="python", exclude_none=True)) + return _JSON_OBJECT_ADAPTER.validate_python(params) @staticmethod async def _poll_for_completion( @@ -80,7 +93,7 @@ class PydanticAITransformation: max_attempts: int = 30, poll_interval: float = 0.5, agent_extra_headers: dict[str, str] | None = None, - ) -> dict[str, Any]: + ) -> JsonObject: """ Poll for task completion using tasks/get method. @@ -112,11 +125,11 @@ class PydanticAITransformation: }, ) response.raise_for_status() - poll_data = response.json() + poll_data = _JSON_OBJECT_ADAPTER.validate_json(response.text) - result = poll_data.get("result", {}) - status = result.get("status", {}) - state = status.get("state", "") + result = _object_field(poll_data, "result") + status = _object_field(result, "status") + state = _string_field(status, "state") verbose_logger.debug("Pydantic AI: Poll attempt %s/%s, state=%s", attempt + 1, max_attempts, state) @@ -133,10 +146,10 @@ class PydanticAITransformation: async def _send_and_poll_raw( api_base: str, request_id: str, - params: Any, + params: object, timeout: float = 60.0, agent_extra_headers: dict[str, str] | None = None, - ) -> dict[str, Any]: + ) -> JsonObject: """ Send a request to Pydantic AI agent and return the raw task response. @@ -156,18 +169,21 @@ class PydanticAITransformation: params_dict = PydanticAITransformation._params_to_dict(params) # Remove None values - FastA2A doesn't accept null for optional fields - params_dict = PydanticAITransformation._remove_none_values(params_dict) + cleaned_params: Final = _JSON_OBJECT_ADAPTER.validate_python( + PydanticAITransformation._remove_none_values(params_dict) + ) # Ensure the message has 'kind': 'message' as required by FastA2A/Pydantic AI - if "message" in params_dict: - params_dict["message"]["kind"] = "message" + message_param: Final = cleaned_params.get("message") + if isinstance(message_param, dict): + message_param["kind"] = "message" # Build A2A JSON-RPC request using message/send method for FastA2A compatibility a2a_request: Final = { "jsonrpc": "2.0", "id": request_id, "method": "message/send", - "params": params_dict, + "params": cleaned_params, } # FastA2A uses root endpoint (/) not /messages @@ -189,25 +205,29 @@ class PydanticAITransformation: }, ) response.raise_for_status() - response_data = response.json() + initial_data: Final = _JSON_OBJECT_ADAPTER.validate_json(response.text) # Check if task is already completed - result: Final = response_data.get("result", {}) - status: Final = result.get("status", {}) - state: Final = status.get("state", "") + result: Final = _object_field(initial_data, "result") + status: Final = _object_field(result, "status") + state: Final = _string_field(status, "state") + task_id: Final = _string_field(result, "id") - if state != "completed": - # Need to poll for completion - task_id: Final = result.get("id") - if task_id: - verbose_logger.info("Pydantic AI: Task %s submitted, polling for completion...", task_id) - response_data = await PydanticAITransformation._poll_for_completion( - client=client, - endpoint=endpoint, - task_id=task_id, - request_id=request_id, - agent_extra_headers=agent_extra_headers, - ) + should_poll: Final = state != "completed" and bool(task_id) + if should_poll: + verbose_logger.info("Pydantic AI: Task %s submitted, polling for completion...", task_id) + + response_data: Final = ( + await PydanticAITransformation._poll_for_completion( + client=client, + endpoint=endpoint, + task_id=task_id, + request_id=request_id, + agent_extra_headers=agent_extra_headers, + ) + if should_poll + else initial_data + ) verbose_logger.info("Pydantic AI: Received completed response for request_id=%s", request_id) @@ -217,10 +237,10 @@ class PydanticAITransformation: async def send_non_streaming_request( api_base: str, request_id: str, - params: Any, + params: object, timeout: float = 60.0, agent_extra_headers: dict[str, str] | None = None, - ) -> dict[str, Any]: + ) -> JsonObject: """ Send a non-streaming A2A request to Pydantic AI agent and wait for completion. @@ -253,10 +273,10 @@ class PydanticAITransformation: async def send_and_get_raw_response( api_base: str, request_id: str, - params: Any, + params: object, timeout: float = 60.0, agent_extra_headers: dict[str, str] | None = None, - ) -> dict[str, Any]: + ) -> JsonObject: """ Send a request to Pydantic AI agent and return the raw task response. @@ -282,9 +302,9 @@ class PydanticAITransformation: @staticmethod def _transform_to_a2a_response( - response_data: dict[str, Any], + response_data: JsonObject, request_id: str, - ) -> dict[str, Any]: + ) -> JsonObject: """ Transform Pydantic AI task response to standard A2A non-streaming format. @@ -313,7 +333,7 @@ class PydanticAITransformation: full_text, message_id, parts = PydanticAITransformation._extract_response_text(response_data) # Build standard A2A message - a2a_message: Final = { + a2a_message: Final[JsonObject] = { "kind": "message", "role": "agent", "parts": parts if parts else [{"kind": "text", "text": full_text}], @@ -328,7 +348,7 @@ class PydanticAITransformation: } @staticmethod - def _extract_response_text(response_data: dict[str, Any]) -> tuple[str, str, list]: + def _extract_response_text(response_data: JsonObject) -> tuple[str, str, list[JsonValue]]: """ Extract response text from completed task response. @@ -342,52 +362,55 @@ class PydanticAITransformation: Returns: Tuple of (full_text, message_id, parts) """ - result: Final = response_data.get("result", {}) + result: Final = _object_field(response_data, "result") # Try to extract from artifacts first (preferred for results) - artifacts: Final = result.get("artifacts", []) - if artifacts: - for artifact in artifacts: - parts = artifact.get("parts", []) - for part in parts: - if part.get("kind") == "text": - text = part.get("text", "") - if text: - return text, str(uuid4()), parts + artifacts: Final = _array_field(result, "artifacts") + for artifact in artifacts: + if not isinstance(artifact, dict): + continue + parts = _array_field(artifact, "parts") + for part in parts: + if isinstance(part, dict) and part.get("kind") == "text": + text = _string_field(part, "text") + if text: + return text, str(uuid4()), parts # Fall back to history - get the last agent message - history: Final = result.get("history", []) + history: Final = _array_field(result, "history") for msg in reversed(history): - if msg.get("role") == "agent": - parts = msg.get("parts", []) - message_id = msg.get("messageId", str(uuid4())) - full_text = "" - for part in parts: - if part.get("kind") == "text": - full_text += part.get("text", "") + if isinstance(msg, dict) and msg.get("role") == "agent": + parts = _array_field(msg, "parts") + message_id = _string_field(msg, "messageId", str(uuid4())) + full_text = "".join( + _string_field(part, "text") + for part in parts + if isinstance(part, dict) and part.get("kind") == "text" + ) if full_text: return full_text, message_id, parts # Fall back to message field (original format) - message: Final = result.get("message", {}) + message: Final = _object_field(result, "message") if message: - parts = message.get("parts", []) - message_id = message.get("messageId", str(uuid4())) - full_text = "" - for part in parts: - if part.get("kind") == "text": - full_text += part.get("text", "") - return full_text, message_id, parts + message_parts: Final = _array_field(message, "parts") + fallback_message_id: Final = _string_field(message, "messageId", str(uuid4())) + fallback_text: Final = "".join( + _string_field(part, "text") + for part in message_parts + if isinstance(part, dict) and part.get("kind") == "text" + ) + return fallback_text, fallback_message_id, message_parts return "", str(uuid4()), [] @staticmethod async def fake_streaming_from_response( - response_data: dict[str, Any], + response_data: JsonObject, request_id: str, chunk_size: int = 50, delay_ms: int = 10, - ) -> AsyncIterator[dict[str, Any]]: + ) -> AsyncIterator[JsonObject]: """ Convert a non-streaming A2A response into fake streaming chunks. @@ -410,23 +433,20 @@ class PydanticAITransformation: full_text, message_id, parts = PydanticAITransformation._extract_response_text(response_data) # Extract input message from raw response for history - result: Final = response_data.get("result", {}) - history: Final = result.get("history", []) - input_message = {} - for msg in history: - if msg.get("role") == "user": - input_message = msg - break + result: Final = _object_field(response_data, "result") + history: Final = _array_field(result, "history") + user_messages: Final = [msg for msg in history if isinstance(msg, dict) and msg.get("role") == "user"] + input_message: Final[JsonObject] = user_messages[0] if user_messages else {} # Generate IDs for streaming events task_id: Final = str(uuid4()) context_id: Final = str(uuid4()) artifact_id: Final = str(uuid4()) - input_message_id: Final = input_message.get("messageId", str(uuid4())) + input_message_id: Final = _string_field(input_message, "messageId", str(uuid4())) # 1. Emit initial task event (kind: "task", status: "submitted") # Format matches A2ACompletionBridgeTransformation.create_task_event - task_event: Final = { + task_event: Final[JsonObject] = { "jsonrpc": "2.0", "id": request_id, "result": { @@ -436,7 +456,7 @@ class PydanticAITransformation: "contextId": context_id, "kind": "message", "messageId": input_message_id, - "parts": input_message.get("parts", [{"kind": "text", "text": ""}]), + "parts": _array_field(input_message, "parts") or [{"kind": "text", "text": ""}], "role": "user", "taskId": task_id, } @@ -452,7 +472,7 @@ class PydanticAITransformation: # 2. Emit status update (kind: "status-update", status: "working") # Format matches A2ACompletionBridgeTransformation.create_status_update_event - working_event: Final = { + working_event: Final[JsonObject] = { "jsonrpc": "2.0", "id": request_id, "result": { @@ -478,7 +498,7 @@ class PydanticAITransformation: chunk_text = full_text[i : i + chunk_size] is_last_chunk = (i + chunk_size) >= len(full_text) - artifact_event = { + artifact_event: JsonObject = { "jsonrpc": "2.0", "id": request_id, "result": { @@ -503,7 +523,7 @@ class PydanticAITransformation: await asyncio.sleep(delay_ms / 1000.0) # 4. Emit final status update (kind: "status-update", status: "completed", final: true) - completed_event: Final = { + completed_event: Final[JsonObject] = { "jsonrpc": "2.0", "id": request_id, "result": { diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index d586156b625..87672c14a09 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -8,7 +8,7 @@ import sys import threading import time from collections.abc import Callable, Mapping -from typing import TYPE_CHECKING, Any, Final, Optional +from typing import TYPE_CHECKING, Any, Final, NoReturn, Optional import certifi import httpx @@ -422,7 +422,7 @@ def _safe_read_response(response: httpx.Response, timeout: float | None = None) return b"" -def _raise_masked_sync_error(e: httpx.HTTPStatusError, stream: bool) -> None: +def _raise_masked_sync_error(e: httpx.HTTPStatusError, stream: bool) -> NoReturn: """Raise a MaskedHTTPStatusError for sync HTTP handlers.""" if stream: try: @@ -442,7 +442,7 @@ def _raise_masked_sync_error(e: httpx.HTTPStatusError, stream: bool) -> None: raise MaskedHTTPStatusError(e, message=_text, text=_text) from None -async def _raise_masked_async_error(e: httpx.HTTPStatusError, stream: bool) -> None: +async def _raise_masked_async_error(e: httpx.HTTPStatusError, stream: bool) -> NoReturn: """Raise a MaskedHTTPStatusError for async HTTP handlers.""" if stream: try: diff --git a/litellm/llms/e2b/sandbox/transformation.py b/litellm/llms/e2b/sandbox/transformation.py index 9cd9ade77a4..4928ca0c092 100644 --- a/litellm/llms/e2b/sandbox/transformation.py +++ b/litellm/llms/e2b/sandbox/transformation.py @@ -8,7 +8,7 @@ Talks to e2b's REST API directly over httpx (no e2b SDK dependency): """ import json -from typing import Final, cast +from typing import Final import httpx @@ -68,13 +68,10 @@ class E2BSandboxConfig(BaseSandboxConfig): if metadata: body["metadata"] = metadata - response: Final = cast( - httpx.Response, - await self._http(client).post( - url=f"{base}/sandboxes", - headers={"X-API-Key": key, "Content-Type": "application/json"}, - json=body, - ), + response: Final = await self._http(client).post( + url=f"{base}/sandboxes", + headers={"X-API-Key": key, "Content-Type": "application/json"}, + json=body, ) data: Final = response.json() @@ -117,14 +114,11 @@ class E2BSandboxConfig(BaseSandboxConfig): headers["E2B-Traffic-Access-Token"] = traffic_token url: Final = f"https://{JUPYTER_PORT}-{handle.id}.{handle.domain}/execute" - response: Final = cast( - httpx.Response, - await self._http(client).post( - url=url, - headers=headers, - json={"code": code, "context_id": None, "env_vars": env_vars}, - stream=True, - ), + response: Final = await self._http(client).post( + url=url, + headers=headers, + json={"code": code, "context_id": None, "env_vars": env_vars}, + stream=True, ) lines: Final = await self._read_capped_lines(response) return self._parse_lines(lines) @@ -142,12 +136,9 @@ class E2BSandboxConfig(BaseSandboxConfig): key: Final = api_key or handle._hidden_params.get("api_key") or self.validate_environment() base: Final = api_base or handle._hidden_params.get("api_base") or E2B_API_BASE try: - response: Final = cast( - httpx.Response, - await self._http(client).delete( - url=f"{base}/sandboxes/{handle.id}", - headers={"X-API-Key": key}, - ), + response: Final = await self._http(client).delete( + url=f"{base}/sandboxes/{handle.id}", + headers={"X-API-Key": key}, ) except httpx.HTTPStatusError as e: if e.response.status_code == 404: diff --git a/litellm/llms/opensandbox/sandbox/transformation.py b/litellm/llms/opensandbox/sandbox/transformation.py index 49a7fb08c4a..5126db9bbc9 100644 --- a/litellm/llms/opensandbox/sandbox/transformation.py +++ b/litellm/llms/opensandbox/sandbox/transformation.py @@ -1,7 +1,7 @@ import asyncio import json import time -from typing import Final, cast +from typing import Final import httpx @@ -86,13 +86,10 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): secure_access=secure_access, ) - response: Final = cast( - httpx.Response, - await self._http(client).post( - url=f"{base}/sandboxes", - headers=self._lifecycle_headers(key), - json=body, - ), + response: Final = await self._http(client).post( + url=f"{base}/sandboxes", + headers=self._lifecycle_headers(key), + json=body, ) data: Final = response.json() sandbox_id: Final = str(data["id"]) @@ -182,12 +179,9 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): base: Final = str(handle._hidden_params.get("api_base") or self._api_base(api_base)) key: Final = self._api_key(api_key=api_key, handle=handle) try: - response: Final = cast( - httpx.Response, - await self._http(client).delete( - url=f"{base}/sandboxes/{handle.id}", - headers=self._lifecycle_headers(key), - ), + response: Final = await self._http(client).delete( + url=f"{base}/sandboxes/{handle.id}", + headers=self._lifecycle_headers(key), ) except httpx.HTTPStatusError as e: if e.response.status_code == 404: @@ -245,12 +239,9 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): ) -> None: deadline: Final = time.monotonic() + ready_timeout while True: - response = cast( - httpx.Response, - await self._http(client).get( - url=f"{api_base}/sandboxes/{sandbox_id}", - headers=headers, - ), + response = await self._http(client).get( + url=f"{api_base}/sandboxes/{sandbox_id}", + headers=headers, ) data = response.json() state = self._sandbox_state(data) @@ -306,13 +297,10 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): use_server_proxy: bool, client: AsyncHTTPHandler | None, ) -> tuple[str, dict[str, str]]: - response: Final = cast( - httpx.Response, - await self._http(client).get( - url=f"{api_base}/sandboxes/{sandbox_id}/endpoints/{OPEN_SANDBOX_EXECD_PORT}", - headers=headers, - params={"use_server_proxy": use_server_proxy}, - ), + response: Final = await self._http(client).get( + url=f"{api_base}/sandboxes/{sandbox_id}/endpoints/{OPEN_SANDBOX_EXECD_PORT}", + headers=headers, + params={"use_server_proxy": use_server_proxy}, ) data: Final = response.json() endpoint: Final = data.get("endpoint") @@ -329,15 +317,12 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): client: AsyncHTTPHandler | None, ) -> list[str]: timeout: Final = httpx.Timeout(connect=30.0, read=None, write=30.0, pool=None) - response: Final = cast( - httpx.Response, - await self._http(client).post( - url=url, - headers=headers, - timeout=timeout, - json=body, - stream=True, - ), + response: Final = await self._http(client).post( + url=url, + headers=headers, + timeout=timeout, + json=body, + stream=True, ) return await self._read_capped_lines(response) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 693e3f8e47d..934011bb2af 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -1006,8 +1006,7 @@ async def exchange_token_with_server( headers={"Accept": "application/json", **token_request.headers}, data=token_data, ) - if response is not None: - response.raise_for_status() + response.raise_for_status() except httpx.HTTPStatusError as exc: fault: Final = classify_upstream_token_rejection( exc.response, @@ -1029,11 +1028,6 @@ async def exchange_token_with_server( ) return _bridge_mint_error_response("invalid_refresh") return render_token_fault(fault) - if response is None: - raise HTTPException( - status_code=502, - detail="MCP upstream token endpoint returned no response", - ) token_response = response.json() # Validate token response against server-configured rules before any storage. @@ -1444,16 +1438,10 @@ async def _post_dcr_registration( headers=headers, json=register_data, ) - if response is not None: - response.raise_for_status() + response.raise_for_status() except httpx.HTTPStatusError as exc: status_code, detail = dcr_fault_detail(classify_upstream_dcr_rejection(exc.response, log_context=server_id)) raise HTTPException(status_code=status_code, detail=detail) from exc - if response is None: - raise HTTPException( - status_code=502, - detail="MCP upstream registration endpoint returned no response", - ) return response diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py index d0053fbe0a8..0da03352423 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py @@ -113,8 +113,6 @@ async def post_client_credentials_grant( return TokenEndpointDenied(status_code=status_code, detail=f"token endpoint returned HTTP {status_code}") except Exception as exc: # noqa: BLE001 # any transport failure is the same outcome: unreachable return TokenEndpointUnreachable(detail=str(exc)) - if not isinstance(response, httpx.Response): - return TokenEndpointUnreachable(detail="token endpoint returned no response") try: body: Final = _TOKEN_BODY_ADAPTER.validate_json(response.content) except ValidationError: diff --git a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py index 37e4c72bf96..13f4b619433 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py @@ -778,7 +778,7 @@ class CompresrGuardrail(CustomGuardrail): {"detail": str(e)}, ) return None - if raw_response is None or not 200 <= raw_response.status_code < 300: + if not 200 <= raw_response.status_code < 300: self._handle_compress_failure( "Compresr compression service returned an error", { diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index 61220819d48..d9a98498a3b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -432,16 +432,6 @@ class HeadroomGuardrail(CustomGuardrail): False, {}, ) - if raw_response is None: - return ( - self._handle_compress_failure( - messages, - "Headroom compression service returned no response", - {}, - ), - False, - {}, - ) response: Final[HttpxResponse] = raw_response if response.status_code != 200: @@ -563,7 +553,7 @@ class HeadroomGuardrail(CustomGuardrail): verbose_proxy_logger.warning("Headroom: retrieve failed for hash=%s: %s", hash_value, e) return f"[Headroom: retrieval failed for hash={hash_value}]" - if raw_response is None or raw_response.status_code == 404: + if raw_response.status_code == 404: return f"[Headroom: hash={hash_value} not found or expired]" if raw_response.status_code != 200: diff --git a/litellm/proxy/guardrails/guardrail_hooks/repelloai/repelloai.py b/litellm/proxy/guardrails/guardrail_hooks/repelloai/repelloai.py index 5f73a169215..d582649c137 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/repelloai/repelloai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/repelloai/repelloai.py @@ -202,8 +202,6 @@ class RepelloAIGuardrail(CustomGuardrail): headers={"X-API-Key": self.repelloai_api_key}, json=request, ) - if raw_response is None: - raise ValueError("RepelloAI Argus returned no response") response: Final[HttpxResponse] = raw_response self._raise_for_config_error(response) response.raise_for_status() diff --git a/litellm/proxy/memory/memory_endpoints.py b/litellm/proxy/memory/memory_endpoints.py index 987823d987f..f72bf16b474 100644 --- a/litellm/proxy/memory/memory_endpoints.py +++ b/litellm/proxy/memory/memory_endpoints.py @@ -18,7 +18,7 @@ Scoping: """ import json -from typing import Any, Final +from typing import TYPE_CHECKING, Final from fastapi import APIRouter, Depends, HTTPException, Query @@ -40,10 +40,15 @@ from litellm.types.memory_management import ( MemoryUpdateRequest, ) +if TYPE_CHECKING: + from prisma.models import LiteLLM_MemoryTable + + from litellm.proxy.utils import PrismaClient + router: Final = APIRouter() -def _serialize_metadata_for_prisma(metadata: Any) -> str: +def _serialize_metadata_for_prisma(metadata: object) -> str: """ Encode a `metadata` payload for the `Json?` column. @@ -62,14 +67,14 @@ def _is_admin(user_api_key_dict: UserAPIKeyAuth) -> bool: return user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN -def _visibility_filter(user_api_key_dict: UserAPIKeyAuth) -> dict | None: +def _visibility_filter(user_api_key_dict: UserAPIKeyAuth) -> dict[str, object] | None: """ Prisma `where` fragment restricting rows to those the caller can see. Returns None for admins (no restriction). """ if user_api_key_has_admin_view(user_api_key_dict): return None - ors: Final[list[dict]] = [] + ors: Final[list[dict[str, object]]] = [] if user_api_key_dict.user_id: ors.append({"user_id": user_api_key_dict.user_id}) if user_api_key_dict.team_id: @@ -80,12 +85,12 @@ def _visibility_filter(user_api_key_dict: UserAPIKeyAuth) -> dict | None: return {"OR": ors} -def _row_to_model(row: Any) -> LiteLLM_MemoryRow: +def _row_to_model(row: "LiteLLM_MemoryTable") -> LiteLLM_MemoryRow: return LiteLLM_MemoryRow( memory_id=row.memory_id, key=row.key, value=row.value, - metadata=getattr(row, "metadata", None), + metadata=row.metadata, user_id=row.user_id, team_id=row.team_id, created_at=row.created_at, @@ -95,7 +100,13 @@ def _row_to_model(row: Any) -> LiteLLM_MemoryRow: ) -def _require_prisma(): +def _require_written_row(row: "LiteLLM_MemoryTable | None") -> "LiteLLM_MemoryTable": + if row is None: + raise RuntimeError("Prisma returned no row for a memory write") + return row + + +def _require_prisma() -> "PrismaClient": from litellm.proxy.proxy_server import prisma_client if prisma_client is None: @@ -113,7 +124,9 @@ def _internal_error(log_message: str, exc: Exception, default_detail: str) -> HT return HTTPException(status_code=500, detail=default_detail) -async def _assert_write_access(prisma_client: Any, row: Any, user_api_key_dict: UserAPIKeyAuth) -> None: +async def _assert_write_access( + prisma_client: "PrismaClient", row: "LiteLLM_MemoryTable", user_api_key_dict: UserAPIKeyAuth +) -> None: """ Enforce ownership for mutations (PUT/DELETE). @@ -135,8 +148,8 @@ async def _assert_write_access(prisma_client: Any, row: Any, user_api_key_dict: """ if _is_admin(user_api_key_dict): return - row_user_id: Final = getattr(row, "user_id", None) - row_team_id: Final = getattr(row, "team_id", None) + row_user_id: Final = row.user_id + row_team_id: Final = row.team_id # Personal ownership. if row_user_id and row_user_id == user_api_key_dict.user_id: @@ -153,7 +166,7 @@ async def _assert_write_access(prisma_client: Any, row: Any, user_api_key_dict: ) -async def _is_team_admin_for(prisma_client: Any, user_api_key_dict: UserAPIKeyAuth, team_id: str) -> bool: +async def _is_team_admin_for(prisma_client: "PrismaClient", user_api_key_dict: UserAPIKeyAuth, team_id: str) -> bool: """ True if the caller is a team admin of `team_id`, or an org admin for the team's organization. Mirrors the auth pattern used by team-management @@ -269,7 +282,7 @@ async def create_memory( # `metadata` is a `Json?` column — prisma-client-python rejects raw # Python values, so JSON-encode any non-null payload and omit the field # entirely when None so the column defaults to SQL NULL. - create_data: Final[dict] = { + create_data: Final[dict[str, object]] = { "key": body.key, "value": body.value, "user_id": user_id, @@ -325,14 +338,14 @@ async def list_memory( # top-level "AND" — safer than `dict.update` since future visibility # filters could grow an "OR" key that would clobber this one if merged # by key. - key_filter: Final[dict] = {} + key_filter: Final[dict[str, object]] = {} if key_prefix is not None: key_filter["key"] = {"startsWith": key_prefix} elif key is not None: key_filter["key"] = key vis: Final = _visibility_filter(user_api_key_dict) - where: dict + where: dict[str, object] if vis is None: where = key_filter elif not key_filter: @@ -354,11 +367,13 @@ async def list_memory( return MemoryListResponse(memories=[_row_to_model(r) for r in rows], total=total) -async def _find_memory_for_caller(prisma_client: Any, key: str, user_api_key_dict: UserAPIKeyAuth) -> Any: +async def _find_memory_for_caller( + prisma_client: "PrismaClient", key: str, user_api_key_dict: UserAPIKeyAuth +) -> "LiteLLM_MemoryTable": """Look up a memory row by key, scoped to the caller's visibility.""" - key_filter: Final[dict] = {"key": key} + key_filter: Final[dict[str, object]] = {"key": key} vis: Final = _visibility_filter(user_api_key_dict) - where: Final[dict] = key_filter if vis is None else {"AND": [key_filter, vis]} + where: Final[dict[str, object]] = key_filter if vis is None else {"AND": [key_filter, vis]} rows = await MemoryRepository(prisma_client).table.find_many(where=where, take=1, order={"updated_at": "desc"}) if not rows: raise HTTPException(status_code=404, detail=f"Memory with key '{key}' not found") @@ -415,7 +430,7 @@ async def upsert_memory( fields_sent: Final = body.model_fields_set metadata_in_payload: Final = "metadata" in fields_sent - data: Final[dict] = {} + data: Final[dict[str, object]] = {} if body.value is not None: data["value"] = body.value if metadata_in_payload: @@ -427,7 +442,7 @@ async def upsert_memory( ) data["updated_by"] = user_api_key_dict.user_id - async def _find_existing() -> Any: + async def _find_existing() -> "LiteLLM_MemoryTable | None": """Return the caller-visible row for `key`, or None.""" try: return await _find_memory_for_caller(prisma_client, key, user_api_key_dict) @@ -459,7 +474,7 @@ async def upsert_memory( # Omit `metadata` when None so the column defaults to SQL NULL; # otherwise JSON-encode for Prisma — same pattern as # `create_memory` above. - create_data: Final[dict] = { + create_data: Final[dict[str, object]] = { "key": key, "value": body.value, "user_id": user_id, @@ -496,7 +511,7 @@ async def upsert_memory( except Exception as e: raise _internal_error("Error upserting memory: %s", e, "Internal error updating memory entry.") - return _row_to_model(row) + return _row_to_model(_require_written_row(row)) @router.delete( diff --git a/litellm/rag/ingestion/s3_vectors_ingestion.py b/litellm/rag/ingestion/s3_vectors_ingestion.py index 36f1e4cf480..b21251dd28a 100644 --- a/litellm/rag/ingestion/s3_vectors_ingestion.py +++ b/litellm/rag/ingestion/s3_vectors_ingestion.py @@ -17,7 +17,9 @@ from __future__ import annotations import hashlib import uuid -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final, TypeAlias + +from pydantic import JsonValue, TypeAdapter import litellm from litellm._logging import verbose_logger @@ -29,12 +31,19 @@ from litellm.constants import ( from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, get_async_httpx_client, httpxSpecialProvider, ) from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion +JsonObject: TypeAlias = dict[str, JsonValue] + +_JSON_OBJECT_ADAPTER: Final = TypeAdapter(JsonObject) + if TYPE_CHECKING: + import httpx + from litellm import Router from litellm.types.rag import RAGIngestOptions @@ -66,10 +75,10 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): BaseAWSLLM.__init__(self) # Extract config - self.vector_bucket_name = self.vector_store_config["vector_bucket_name"] - self.index_name = self.vector_store_config.get("index_name") - self.distance_metric = self.vector_store_config.get("distance_metric", S3_VECTORS_DEFAULT_DISTANCE_METRIC) - self.non_filterable_metadata_keys = self.vector_store_config.get( + self.vector_bucket_name: str = self.vector_store_config["vector_bucket_name"] + self.index_name: str | None = self.vector_store_config.get("index_name") + self.distance_metric: str = self.vector_store_config.get("distance_metric", S3_VECTORS_DEFAULT_DISTANCE_METRIC) + self.non_filterable_metadata_keys: list[str] = self.vector_store_config.get( "non_filterable_metadata_keys", S3_VECTORS_DEFAULT_NON_FILTERABLE_METADATA_KEYS, ) @@ -85,7 +94,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): # Create httpx client (similar to s3_v2.py) ssl_verify: Final = self._get_ssl_verify(ssl_verify=self.vector_store_config.get("ssl_verify")) - self.async_httpx_client = get_async_httpx_client( + self.async_httpx_client: AsyncHTTPHandler = get_async_httpx_client( llm_provider=httpxSpecialProvider.RAG, params={"ssl_verify": ssl_verify} if ssl_verify is not None else None, ) @@ -166,7 +175,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): url: str, data: str | None = None, headers: dict[str, str] | None = None, - ) -> Any: + ) -> httpx.Response: """ Helper to sign and execute AWS API requests using httpx + SigV4. @@ -311,7 +320,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): ) # Prepare index configuration per AWS API docs - index_config: Final = { + index_config: Final[dict[str, object]] = { "vectorBucketName": self.vector_bucket_name, "indexName": self.index_name, "dataType": "float32", @@ -336,7 +345,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): verbose_logger.exception("Error creating vector index: %s", e) raise - async def _put_vectors(self, vectors: list[dict[str, Any]]): + async def _put_vectors(self, vectors: list[JsonObject]): """ Call PutVectors API to store vectors in S3 Vectors. @@ -442,10 +451,10 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): raise ValueError(error_msg) # Prepare vectors for PutVectors API - vectors: Final = [] + vectors: Final[list[JsonObject]] = [] for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)): # Build metadata dict - metadata: dict[str, str] = { + metadata: JsonObject = { "source_text": chunk, # Non-filterable (for reference) "chunk_index": str(i), # Filterable } @@ -453,9 +462,9 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): if filename: metadata["filename"] = filename # Filterable - vector_obj = { + vector_obj: JsonObject = { "key": f"{filename}_{i}" if filename else f"chunk_{i}", - "data": {"float32": embedding}, + "data": {"float32": list(embedding)}, "metadata": metadata, } @@ -468,7 +477,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): vector_store_id: Final = f"{self.vector_bucket_name}:{self.index_name}" return vector_store_id, filename - async def query_vector_store(self, vector_store_id: str, query: str, top_k: int = 5) -> dict[str, Any] | None: + async def query_vector_store(self, vector_store_id: str, query: str, top_k: int = 5) -> JsonObject | None: """ Query S3 Vectors using QueryVectors API. @@ -507,18 +516,19 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): response = await self._sign_and_execute_request("POST", url, data=safe_dumps(request_body)) if response.status_code == 200: - results: Final = response.json() - verbose_logger.debug("Query returned %s results", len(results.get("vectors", []))) + results: Final = _JSON_OBJECT_ADAPTER.validate_json(response.text) + vectors: Final = results.get("vectors") + vector_list: Final = vectors if isinstance(vectors, list) else [] + verbose_logger.debug("Query returned %s results", len(vector_list)) - # Check if query terms appear in results - if results.get("vectors"): - for result in results["vectors"]: - metadata = result.get("metadata", {}) - source_text = metadata.get("source_text", "") - if query.lower() in source_text.lower(): - return results + for result in vector_list: + if not isinstance(result, dict): + continue + metadata = result.get("metadata") + source_text = metadata.get("source_text") if isinstance(metadata, dict) else None + if isinstance(source_text, str) and query.lower() in source_text.lower(): + return results - # Return results even if exact match not found return results else: verbose_logger.error("QueryVectors failed with status %s: %s", response.status_code, response.text) diff --git a/litellm/repositories/prisma_protocols.py b/litellm/repositories/prisma_protocols.py index 6aff196ff10..487a98dcbe9 100644 --- a/litellm/repositories/prisma_protocols.py +++ b/litellm/repositories/prisma_protocols.py @@ -10,6 +10,7 @@ from collections.abc import Mapping, Sequence from typing import Protocol, TypeVar RowT_co = TypeVar("RowT_co", covariant=True) +RowT = TypeVar("RowT") class PrismaRecord(Protocol): @@ -26,6 +27,77 @@ class SpendLinkedTable(Protocol[RowT_co]): async def update_many(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> int: ... +class TableActions(Protocol[RowT]): + """Structural view of the generated Prisma actions for a single table.""" + + async def find_many( + self, + *, + where: Mapping[str, object] | None = None, + take: int | None = None, + skip: int | None = None, + order: Mapping[str, str] | Sequence[Mapping[str, str]] | None = None, + include: Mapping[str, object] | None = None, + ) -> list[RowT]: ... + + async def find_first( + self, + *, + where: Mapping[str, object] | None = None, + skip: int | None = None, + order: Mapping[str, str] | Sequence[Mapping[str, str]] | None = None, + include: Mapping[str, object] | None = None, + ) -> RowT | None: ... + + async def find_unique( + self, + *, + where: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> RowT | None: ... + + async def create( + self, + data: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> RowT: ... + + async def create_many( + self, + data: Sequence[Mapping[str, object]], + *, + skip_duplicates: bool | None = None, + ) -> int: ... + + async def update( + self, + data: Mapping[str, object], + where: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> RowT | None: ... + + async def update_many(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> int: ... + + async def upsert( + self, + *, + where: Mapping[str, object], + data: Mapping[str, Mapping[str, object]], + include: Mapping[str, object] | None = None, + ) -> RowT: ... + + async def delete( + self, + *, + where: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> RowT | None: ... + + async def delete_many(self, *, where: Mapping[str, object] | None = None) -> int: ... + + async def count(self, *, where: Mapping[str, object] | None = None) -> int: ... + + class BatchTable(Protocol): def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> None: ... diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index 66e0b6d59e7..cecf448025e 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -7,9 +7,13 @@ These are thin wrappers for tables that do not (yet) need domain-specific query methods; richer repositories live in their own modules. """ -from typing import Any +from typing import TYPE_CHECKING, Any from litellm.proxy.common_utils.config_sync_pubsub import wrap_table_actions_for_config_sync +from litellm.repositories.prisma_protocols import TableActions + +if TYPE_CHECKING: + from prisma import models as prisma_db_models class PrismaTableRepository: @@ -109,6 +113,10 @@ class ManagedFileRepository(PrismaTableRepository): class MemoryRepository(PrismaTableRepository): table_name = "litellm_memorytable" + @property + def table(self) -> "TableActions[prisma_db_models.LiteLLM_MemoryTable]": + return super().table # any-ok: the untyped client wrapper is narrowed to this table's row type + class SearchToolsRepository(PrismaTableRepository): table_name = "litellm_searchtoolstable" diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 60356eda05b..61fff10f37e 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -12,7 +12,7 @@ "limit": 2033 }, "ANN202": { - "limit": 865 + "limit": 864 }, "ANN204": { "limit": 713 @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 1630 + "limit": 1615 }, "ASYNC230": { "limit": 11 @@ -306,7 +306,7 @@ "limit": 0 }, "TID251": { - "limit": 1240 + "limit": 1236 }, "TRY002": { "limit": 528 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index ab8198304bb..fa47f990252 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23256 + "limit": 23248 }, "LIT002": { - "limit": 27213 + "limit": 27200 }, "LIT003": { "limit": 269 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1091 + "limit": 1083 }, "LIT007": { "limit": 0