mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-22 00:31:44 +00:00
chore(typing): clear basedpyright Any errors in a2a, rag ingestion, memory endpoints
Replace Any seams with JSON value models, Prisma table protocols, and precise httpx return types, and ratchet the lint budgets down Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Devin AI
parent
e1717c5e9c
commit
35261c7cec
16 changed files with 320 additions and 245 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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": {
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
{
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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: ...
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue