mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-25 05:54:53 +00:00
Merge 0c6d41fa4d into 714d84acb6
This commit is contained in:
Mateo Wang
GitHub
commit
034cee8c02
9 changed files with 1254 additions and 43 deletions
|
|
@ -43,7 +43,7 @@ _ADDRESSED_RESPONSE_ID_FIELD_BY_CALL_TYPE: Final[Mapping[str, str]] = MappingPro
|
|||
}
|
||||
)
|
||||
|
||||
_UNMANAGED_RESPONSE_ID_DETAIL: Final = (
|
||||
UNMANAGED_RESPONSE_ID_DETAIL: Final = (
|
||||
"Forbidden. This response id was not issued by this proxy, so the proxy cannot tell who owns it. "
|
||||
"To let keys address responses this proxy did not issue, set "
|
||||
"general_settings::allow_unmanaged_response_ids to True in the config.yaml file."
|
||||
|
|
@ -57,7 +57,7 @@ def _proxy_general_settings() -> Mapping[str, object]:
|
|||
return general_settings
|
||||
|
||||
|
||||
def _proxy_signing_key() -> str | None:
|
||||
def proxy_signing_key() -> str | None:
|
||||
import os
|
||||
|
||||
from litellm.proxy.proxy_server import master_key
|
||||
|
|
@ -66,6 +66,20 @@ def _proxy_signing_key() -> str | None:
|
|||
return master_key if salt_key is None else salt_key
|
||||
|
||||
|
||||
def unmanaged_response_ids_allowed(
|
||||
general_settings: Mapping[str, object],
|
||||
signing_key: str | None,
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
) -> bool:
|
||||
if general_settings.get("disable_responses_id_security", False):
|
||||
return True
|
||||
if general_settings.get("allow_unmanaged_response_ids", False):
|
||||
return True
|
||||
if signing_key is None:
|
||||
return True
|
||||
return user_api_key_dict.user_role in _PROXY_ADMIN_ROLES
|
||||
|
||||
|
||||
_RESPONSE_PAYLOAD_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
|
|
@ -119,7 +133,7 @@ class ResponsesIDSecurity(CustomLogger):
|
|||
def __init__(
|
||||
self,
|
||||
general_settings_reader: Callable[[], Mapping[str, object]] = _proxy_general_settings,
|
||||
signing_key_reader: Callable[[], str | None] = _proxy_signing_key,
|
||||
signing_key_reader: Callable[[], str | None] = proxy_signing_key,
|
||||
) -> None:
|
||||
self._general_settings_reader: Final = general_settings_reader
|
||||
self._signing_key_reader: Final = signing_key_reader
|
||||
|
|
@ -155,21 +169,10 @@ class ResponsesIDSecurity(CustomLogger):
|
|||
self.check_user_access_to_response_id(user_id, team_id, user_api_key_dict)
|
||||
return original_response_id
|
||||
|
||||
if self._unmanaged_response_ids_allowed(user_api_key_dict):
|
||||
if unmanaged_response_ids_allowed(self._general_settings_reader(), self._get_signing_key(), user_api_key_dict):
|
||||
return response_id
|
||||
|
||||
raise HTTPException(status_code=403, detail=_UNMANAGED_RESPONSE_ID_DETAIL)
|
||||
|
||||
def _unmanaged_response_ids_allowed(self, user_api_key_dict: "UserAPIKeyAuth") -> bool:
|
||||
general_settings: Final = self._general_settings_reader()
|
||||
|
||||
if general_settings.get("disable_responses_id_security", False):
|
||||
return True
|
||||
if general_settings.get("allow_unmanaged_response_ids", False):
|
||||
return True
|
||||
if self._get_signing_key() is None:
|
||||
return True
|
||||
return user_api_key_dict.user_role in _PROXY_ADMIN_ROLES
|
||||
raise HTTPException(status_code=403, detail=UNMANAGED_RESPONSE_ID_DETAIL)
|
||||
|
||||
def check_user_access_to_response_id(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -25,7 +25,7 @@ import httpx
|
|||
import openai
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, Response, WebSocket
|
||||
from fastapi.responses import StreamingResponse
|
||||
from pydantic import ConfigDict, TypeAdapter, ValidationError
|
||||
from pydantic import ConfigDict, JsonValue, TypeAdapter, ValidationError
|
||||
from starlette.websockets import WebSocketState
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
|
|
@ -55,6 +55,7 @@ from litellm.llms.azure.passthrough.transformation import (
|
|||
foreign_azure_deployment,
|
||||
is_azure_body_model_inference_endpoint,
|
||||
)
|
||||
from litellm.llms.base_llm.passthrough.transformation import strip_leading_model_segment
|
||||
from litellm.llms.bedrock.passthrough.transformation import is_bedrock_streaming_endpoint
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.llms.deepgram.common_utils import (
|
||||
|
|
@ -112,6 +113,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
|||
HttpPassThroughEndpointHelpers,
|
||||
create_pass_through_route,
|
||||
create_websocket_passthrough_route,
|
||||
refuse_unissued_response_ids,
|
||||
websocket_passthrough_request,
|
||||
)
|
||||
from litellm.proxy.utils import ProxyLogging as ProxyLoggingType
|
||||
|
|
@ -143,7 +145,9 @@ from litellm.utils import ProviderConfigManager
|
|||
from .passthrough_endpoint_router import PassthroughEndpointRouter
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.proxy.proxy_server import ProxyConfig as _ProxyConfig
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.router import Router
|
||||
|
||||
ProxyConfig = _ProxyConfig # rebind-ok: conditional type alias
|
||||
|
|
@ -2168,9 +2172,103 @@ async def _relay_response_bytes(upstream: httpx.Response) -> AsyncGenerator[byte
|
|||
await upstream.aclose()
|
||||
|
||||
|
||||
async def _relay_upstream_response(upstream: httpx.Response) -> Response:
|
||||
_RELAYED_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _RouterRelayManagedIds:
|
||||
provider: str
|
||||
prisma_client: PrismaClient
|
||||
managed_files_hook: CustomLogger
|
||||
user_api_key_dict: UserAPIKeyAuth
|
||||
method: str
|
||||
route: str
|
||||
|
||||
async def resolve_request(
|
||||
self, endpoint: str, request_body: Mapping[str, object]
|
||||
) -> tuple[str, Mapping[str, object]]:
|
||||
from litellm.proxy.pass_through_endpoints.managed_id_rewriter import rewrite_body_ids, rewrite_path_ids
|
||||
|
||||
await refuse_unissued_response_ids(
|
||||
endpoint, request_body, self.provider, self.prisma_client, self.user_api_key_dict
|
||||
)
|
||||
resolved_endpoint: Final = await rewrite_path_ids(
|
||||
endpoint, self.provider, self.user_api_key_dict, self.prisma_client, self.managed_files_hook
|
||||
)
|
||||
resolved_body: Final = await rewrite_body_ids(
|
||||
dict(request_body), self.provider, self.user_api_key_dict, self.prisma_client, self.managed_files_hook
|
||||
)
|
||||
return resolved_endpoint, request_body if resolved_body is None else resolved_body
|
||||
|
||||
async def own_response_body(self, content: bytes, status_code: int) -> bytes:
|
||||
from litellm.proxy.pass_through_endpoints.managed_id_rewriter import rewrite_response_ids
|
||||
|
||||
if status_code >= 300:
|
||||
return content
|
||||
try:
|
||||
body: Final = _RELAYED_JSON_OBJECT.validate_json(content)
|
||||
except ValidationError:
|
||||
return content
|
||||
owned: Final = await rewrite_response_ids(
|
||||
provider=self.provider,
|
||||
method=self.method,
|
||||
route=self.route,
|
||||
body=body,
|
||||
user_api_key_dict=self.user_api_key_dict,
|
||||
prisma_client=self.prisma_client,
|
||||
managed_files_hook=self.managed_files_hook,
|
||||
)
|
||||
return content if owned is body else json.dumps(owned).encode("utf-8")
|
||||
|
||||
def own_stream(self, stream: AsyncGenerator[bytes, None]) -> AsyncGenerator[bytes, None]:
|
||||
from litellm.proxy.pass_through_endpoints.managed_id_rewriter import rewrite_streamed_response_ids
|
||||
|
||||
return rewrite_streamed_response_ids(
|
||||
stream=stream,
|
||||
provider=self.provider,
|
||||
method=self.method,
|
||||
route=self.route,
|
||||
user_api_key_dict=self.user_api_key_dict,
|
||||
prisma_client=self.prisma_client,
|
||||
)
|
||||
|
||||
|
||||
def _router_relay_managed_ids(
|
||||
provider: str, method: str, route: str, user_api_key_dict: UserAPIKeyAuth
|
||||
) -> _RouterRelayManagedIds | None:
|
||||
from litellm.proxy.proxy_server import general_settings_view, prisma_client, proxy_logging_obj
|
||||
|
||||
if "responses" not in route.split("/"):
|
||||
return None
|
||||
managed_files_hook: Final = proxy_logging_obj.get_proxy_hook("managed_files")
|
||||
if (
|
||||
not general_settings_view().get("passthrough_managed_object_ids", False)
|
||||
or prisma_client is None
|
||||
or managed_files_hook is None
|
||||
):
|
||||
return None
|
||||
return _RouterRelayManagedIds(
|
||||
provider=provider,
|
||||
prisma_client=prisma_client,
|
||||
managed_files_hook=managed_files_hook,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
method=method,
|
||||
route=route,
|
||||
)
|
||||
|
||||
|
||||
def _owned_relay_stream(
|
||||
stream: AsyncGenerator[bytes, None], managed_ids: _RouterRelayManagedIds | None
|
||||
) -> AsyncGenerator[bytes, None]:
|
||||
return stream if managed_ids is None else managed_ids.own_stream(stream)
|
||||
|
||||
|
||||
async def _relay_upstream_response(
|
||||
upstream: httpx.Response, managed_ids: _RouterRelayManagedIds | None = None
|
||||
) -> Response:
|
||||
content: Final = await upstream.aread()
|
||||
return Response(
|
||||
content=await upstream.aread(),
|
||||
content=content if managed_ids is None else await managed_ids.own_response_body(content, upstream.status_code),
|
||||
status_code=upstream.status_code,
|
||||
headers=HttpPassThroughEndpointHelpers.get_response_headers(headers=upstream.headers, custom_headers=None),
|
||||
)
|
||||
|
|
@ -2184,6 +2282,7 @@ async def _relay_azure_router_model(
|
|||
request_body: Mapping[str, object],
|
||||
is_streaming_request: bool,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
managed_ids: _RouterRelayManagedIds | None,
|
||||
) -> Response:
|
||||
foreign_deployment: Final = foreign_azure_deployment(
|
||||
endpoint, model, lambda: _models_served_by_group(llm_router, model)
|
||||
|
|
@ -2202,6 +2301,7 @@ async def _relay_azure_router_model(
|
|||
request_body=request_body,
|
||||
is_streaming_request=is_streaming_request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
managed_ids=managed_ids,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -2213,6 +2313,7 @@ async def _relay_router_model(
|
|||
request_body: Mapping[str, object],
|
||||
is_streaming_request: bool,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
managed_ids: _RouterRelayManagedIds | None,
|
||||
) -> Response:
|
||||
try:
|
||||
result: Final = await llm_router.allm_passthrough_route(
|
||||
|
|
@ -2235,13 +2336,13 @@ async def _relay_router_model(
|
|||
return await _relay_upstream_response(upstream_error.response)
|
||||
|
||||
if not is_streaming_request:
|
||||
return await _relay_upstream_response(cast(httpx.Response, result))
|
||||
return await _relay_upstream_response(cast(httpx.Response, result), managed_ids)
|
||||
|
||||
if inspect.isasyncgen(result):
|
||||
sse_headers: Final = {"content-type": "text/event-stream"}
|
||||
return StreamingResponse(
|
||||
content=wrap_passthrough_sse_bytes_with_keepalive_pings(
|
||||
stream=_relay_upstream_bytes(result),
|
||||
stream=_owned_relay_stream(_relay_upstream_bytes(result), managed_ids),
|
||||
ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds,
|
||||
upstream_headers=sse_headers,
|
||||
),
|
||||
|
|
@ -2252,7 +2353,7 @@ async def _relay_router_model(
|
|||
upstream_stream: Final = cast(AsyncPassthroughStreamingResponse, result)
|
||||
return StreamingResponse(
|
||||
content=wrap_passthrough_sse_bytes_with_keepalive_pings(
|
||||
stream=_relay_upstream_bytes(upstream_stream),
|
||||
stream=_owned_relay_stream(_relay_upstream_bytes(upstream_stream), managed_ids),
|
||||
ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds,
|
||||
upstream_headers=upstream_stream.headers,
|
||||
),
|
||||
|
|
@ -2317,6 +2418,7 @@ async def relay_nvidia_nim_request(
|
|||
request_body=request_body,
|
||||
is_streaming_request=is_streaming_request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
managed_ids=None,
|
||||
),
|
||||
ping_interval_seconds=(litellm.sse_keepalive_ping_interval_seconds if is_streaming_request else None),
|
||||
)
|
||||
|
|
@ -2336,16 +2438,37 @@ async def _relay_azure_body_model_group(
|
|||
model: Final = _optional_str(request_body.get("model"))
|
||||
if model is None or not is_passthrough_request_using_router_model(request_body, llm_router):
|
||||
return None
|
||||
return await _open_azure_router_relay(llm_router, model, endpoint, request, request_body, user_api_key_dict)
|
||||
|
||||
|
||||
async def _open_azure_router_relay(
|
||||
llm_router: litellm.Router,
|
||||
model: str,
|
||||
endpoint: str,
|
||||
request: Request,
|
||||
request_body: Mapping[str, object],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> Response:
|
||||
is_streaming_request: Final = is_passthrough_request_streaming(request_body)
|
||||
managed_ids: Final = _router_relay_managed_ids(
|
||||
LlmProviders.AZURE.value,
|
||||
request.method,
|
||||
f"/azure/{strip_leading_model_segment(endpoint, (model,))}",
|
||||
user_api_key_dict,
|
||||
)
|
||||
relayed_endpoint, relayed_body = (
|
||||
(endpoint, request_body) if managed_ids is None else await managed_ids.resolve_request(endpoint, request_body)
|
||||
)
|
||||
return await open_sse_before_first_byte(
|
||||
_relay_azure_router_model(
|
||||
llm_router=llm_router,
|
||||
model=model,
|
||||
endpoint=endpoint,
|
||||
endpoint=relayed_endpoint,
|
||||
request=request,
|
||||
request_body=request_body,
|
||||
request_body=relayed_body,
|
||||
is_streaming_request=is_streaming_request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
managed_ids=managed_ids,
|
||||
),
|
||||
ping_interval_seconds=(litellm.sse_keepalive_ping_interval_seconds if is_streaming_request else None),
|
||||
)
|
||||
|
|
@ -2399,21 +2522,8 @@ async def azure_proxy_route(
|
|||
)
|
||||
|
||||
if is_router_model:
|
||||
request_body = await get_request_body(request)
|
||||
is_streaming_request = is_passthrough_request_streaming(request_body)
|
||||
return await open_sse_before_first_byte(
|
||||
_relay_azure_router_model(
|
||||
llm_router=llm_router,
|
||||
model=part,
|
||||
endpoint=endpoint,
|
||||
request=request,
|
||||
request_body=request_body,
|
||||
is_streaming_request=is_streaming_request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
),
|
||||
ping_interval_seconds=(
|
||||
litellm.sse_keepalive_ping_interval_seconds if is_streaming_request else None
|
||||
),
|
||||
return await _open_azure_router_relay(
|
||||
llm_router, part, endpoint, request, await get_request_body(request), user_api_key_dict
|
||||
)
|
||||
elif is_vector_store_index:
|
||||
# get the api key from the provider config
|
||||
|
|
|
|||
|
|
@ -33,6 +33,7 @@ from __future__ import annotations
|
|||
import json
|
||||
import re
|
||||
from collections.abc import AsyncGenerator, Callable, Mapping, Sequence
|
||||
from itertools import pairwise
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Final,
|
||||
|
|
@ -51,9 +52,10 @@ from litellm.llms.base_llm.managed_resources.isolation import (
|
|||
can_access_resource,
|
||||
resolve_resource_owner_id,
|
||||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy._types import DB_CONNECTION_ERROR_TYPES, UserAPIKeyAuth
|
||||
from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit
|
||||
from litellm.proxy.common_utils.sse_keepalive import split_complete_sse_frames
|
||||
from litellm.repositories.chunked_in import count_in
|
||||
from litellm.repositories.managed_file_repository import ManagedFileRepository
|
||||
from litellm.repositories.table_repositories import ManagedObjectRepository
|
||||
from litellm.types.llms.openai import BATCH_GUARDRAIL_RESPONSE_FIELD, OpenAIFileObject
|
||||
|
|
@ -70,6 +72,7 @@ from litellm.types.passthrough_endpoints.managed_id_rewriter import (
|
|||
PrismaWhere,
|
||||
PrismaWhereValue,
|
||||
ResourceKind,
|
||||
ResponseIdIssuance,
|
||||
SortOrder,
|
||||
)
|
||||
|
||||
|
|
@ -178,8 +181,10 @@ BUILTIN_OUTPUT_ID_FIELD_MAP: Final[dict[_MapKey, list[_FieldSpec]]] = {
|
|||
# Prefixes that live in the *file* table rather than the object table.
|
||||
_FILE_PREFIXES: Final[frozenset[str]] = frozenset({"file-"})
|
||||
|
||||
_RESPONSE_ID_PREFIX: Final = "resp_"
|
||||
|
||||
# Raw provider-ID prefixes that live in the object table (batches, responses).
|
||||
_OBJECT_PREFIXES: Final[frozenset[str]] = frozenset({"batch_", "resp_"})
|
||||
_OBJECT_PREFIXES: Final[frozenset[str]] = frozenset({"batch_", _RESPONSE_ID_PREFIX})
|
||||
|
||||
# Guards request-body rewriting against stack exhaustion from adversarially
|
||||
# deep payloads. Real OpenAI files/batches bodies nest only a few levels.
|
||||
|
|
@ -466,6 +471,38 @@ async def _guard_raw_provider_id(
|
|||
raise HTTPException(status_code=404, detail="Managed resource not found.")
|
||||
|
||||
|
||||
def _addressed_raw_response_ids(path: str, body: Mapping[str, object]) -> frozenset[str]:
|
||||
segments: Final = tuple(unquote(seg) for seg in path.split("/"))
|
||||
path_ids: Final = frozenset(
|
||||
seg for parent, seg in pairwise(segments) if parent == "responses" and seg.startswith(_RESPONSE_ID_PREFIX)
|
||||
)
|
||||
previous_response_id: Final = body.get("previous_response_id")
|
||||
if isinstance(previous_response_id, str) and previous_response_id.startswith(_RESPONSE_ID_PREFIX):
|
||||
return path_ids | {previous_response_id}
|
||||
return path_ids
|
||||
|
||||
|
||||
async def addressed_response_id_issuance(
|
||||
path: str,
|
||||
body: Mapping[str, object],
|
||||
provider: str,
|
||||
prisma_client: PrismaClient | None,
|
||||
) -> ResponseIdIssuance:
|
||||
from prisma.errors import PrismaError
|
||||
|
||||
if prisma_client is None:
|
||||
return "issued"
|
||||
object_ids: Final = tuple(f"passthrough:{provider}:{raw_id}" for raw_id in _addressed_raw_response_ids(path, body))
|
||||
if not object_ids:
|
||||
return "issued"
|
||||
try:
|
||||
issued_count: Final = await count_in(_object_table(prisma_client), "model_object_id", object_ids)
|
||||
except (PrismaError, *DB_CONNECTION_ERROR_TYPES):
|
||||
verbose_proxy_logger.warning("managed_id_rewriter: response id ownership lookup failed", exc_info=True)
|
||||
return "unverified"
|
||||
return "issued" if issued_count == len(object_ids) else "unissued"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# OUTPUT path — helpers for minting and storing managed IDs
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -817,7 +854,6 @@ async def rewrite_response_ids(
|
|||
return mutated if changed else body
|
||||
|
||||
|
||||
_RESPONSE_ID_PREFIX: Final = "resp_"
|
||||
_STREAMED_RESPONSE_ID_SPEC: Final[_FieldSpec] = ("id", _RESPONSE_ID_PREFIX)
|
||||
_SSE_DATA_PREFIX: Final = "data:"
|
||||
_SSE_EVENT_ADAPTER: Final = TypeAdapter(Mapping[str, JsonValue])
|
||||
|
|
|
|||
|
|
@ -30,6 +30,7 @@ from pydantic import TypeAdapter
|
|||
from starlette.datastructures import UploadFile as StarletteUploadFile
|
||||
from starlette.routing import BaseRoute, Route
|
||||
from starlette.websockets import WebSocketState
|
||||
from typing_extensions import assert_never
|
||||
from websockets.asyncio.client import connect
|
||||
from websockets.exceptions import (
|
||||
ConnectionClosedError,
|
||||
|
|
@ -153,6 +154,7 @@ from .upstream_usage_headers import (
|
|||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
router: Final = APIRouter()
|
||||
|
||||
|
|
@ -161,6 +163,7 @@ pass_through_endpoint_logging: Final = PassThroughEndpointLogging()
|
|||
_METADATA_KEYS: Final = frozenset(("litellm_metadata", "metadata"))
|
||||
_DECODED_JSON: Final = TypeAdapter(object)
|
||||
_KEPT_OUT_OF_LITELLM_PARAMS: Final = _METADATA_KEYS | frozenset(wire_names(ProxyRequestState))
|
||||
_REQUEST_BODY_VIEW: Final = TypeAdapter(Mapping[str, object])
|
||||
_PRICING_MODEL_OVERRIDE_PARAMS: Final = frozenset(("base_model",))
|
||||
|
||||
|
||||
|
|
@ -1374,6 +1377,14 @@ async def pass_through_request(
|
|||
prisma_client as _passthrough_prisma,
|
||||
)
|
||||
|
||||
await refuse_unissued_response_ids(
|
||||
url.path,
|
||||
_REQUEST_BODY_VIEW.validate_python(_parsed_body or {}),
|
||||
_managed_id_provider,
|
||||
_passthrough_prisma,
|
||||
user_api_key_dict,
|
||||
)
|
||||
|
||||
_original_path: Final = url.path
|
||||
_original_query_params: Final = requested_query_params
|
||||
_original_body: Final = _parsed_body
|
||||
|
|
@ -2918,6 +2929,41 @@ def _own_streamed_managed_ids(
|
|||
)
|
||||
|
||||
|
||||
async def refuse_unissued_response_ids(
|
||||
path: str,
|
||||
body: Mapping[str, object],
|
||||
managed_id_provider: str,
|
||||
prisma_client: "PrismaClient | None",
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> None:
|
||||
from litellm.proxy.hooks.responses_id_security import (
|
||||
UNMANAGED_RESPONSE_ID_DETAIL,
|
||||
proxy_signing_key,
|
||||
unmanaged_response_ids_allowed,
|
||||
)
|
||||
from litellm.proxy.pass_through_endpoints.managed_id_rewriter import (
|
||||
addressed_response_id_issuance,
|
||||
)
|
||||
from litellm.proxy.proxy_server import general_settings_view
|
||||
|
||||
if unmanaged_response_ids_allowed(general_settings_view(), proxy_signing_key(), user_api_key_dict):
|
||||
return
|
||||
issuance: Final = await addressed_response_id_issuance(path, body, managed_id_provider, prisma_client)
|
||||
match issuance:
|
||||
case "issued":
|
||||
return
|
||||
case "unissued":
|
||||
raise HTTPException(status_code=403, detail=UNMANAGED_RESPONSE_ID_DETAIL)
|
||||
case "unverified":
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail="The proxy could not check who owns this response id because its database lookup failed. "
|
||||
"Retry the request.",
|
||||
)
|
||||
case _:
|
||||
assert_never(issuance)
|
||||
|
||||
|
||||
def _should_buffer_passthrough_response(response: httpx.Response) -> bool:
|
||||
"""
|
||||
Decide from the response headers whether the body must be read into memory.
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ if TYPE_CHECKING:
|
|||
|
||||
SortOrder: TypeAlias = Literal["asc", "desc"]
|
||||
ResourceKind: TypeAlias = Literal["files", "batches"]
|
||||
ResponseIdIssuance: TypeAlias = Literal["issued", "unissued", "unverified"]
|
||||
|
||||
PrismaWhereValue: TypeAlias = (
|
||||
"str | int | bool | datetime | None | Mapping[str, PrismaWhereValue] | Sequence[PrismaWhereValue]"
|
||||
|
|
@ -80,6 +81,8 @@ class ManagedObjectTable(ManagedTable[ManagedObjectRow], Protocol):
|
|||
|
||||
async def upsert(self, *, where: PrismaWhere, data: Mapping[str, ManagedRowData]) -> ManagedObjectRow: ...
|
||||
|
||||
async def count(self, *, where: Mapping[str, object]) -> int: ...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class ManagedFileIdReader(Protocol):
|
||||
|
|
|
|||
|
|
@ -0,0 +1,533 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import uuid
|
||||
from collections.abc import Iterator, Mapping
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal, assert_never
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
import httpx
|
||||
import psycopg
|
||||
import pytest
|
||||
from integration._support.client import Gateway, Scenario
|
||||
from integration._support.database import scratch_database
|
||||
from integration._support.process import graceful_stop_seconds, owned_proxy, owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from integration._support.worker_readiness import wait_for_worker_readiness
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state]
|
||||
|
||||
_UPSTREAM_KEY: Final = "synthetic-openai-key"
|
||||
_MODEL: Final = "gpt-5"
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_MANAGED: Final = "general_settings:\n passthrough_managed_object_ids: true\n"
|
||||
_ALLOW_UNMANAGED: Final = (
|
||||
"general_settings:\n passthrough_managed_object_ids: true\n allow_unmanaged_response_ids: true\n"
|
||||
)
|
||||
_UNMANAGED_DETAIL: Final = (
|
||||
"Forbidden. This response id was not issued by this proxy, so the proxy cannot tell who owns it. "
|
||||
"To let keys address responses this proxy did not issue, set "
|
||||
"general_settings::allow_unmanaged_response_ids to True in the config.yaml file."
|
||||
)
|
||||
_LOOKUP_FAILED_DETAIL: Final = (
|
||||
"The proxy could not check who owns this response id because its database lookup failed. Retry the request."
|
||||
)
|
||||
_ACCESS_DENIED: Final = "Access denied to managed resource."
|
||||
_NOT_FOUND: Final = "Managed resource not found."
|
||||
_TAKE_OWNER_TABLE_OFFLINE: Final = (
|
||||
'ALTER TABLE "LiteLLM_ManagedObjectTable" RENAME TO "LiteLLM_ManagedObjectTable_offline"'
|
||||
)
|
||||
_BRING_OWNER_TABLE_BACK: Final = (
|
||||
'ALTER TABLE "LiteLLM_ManagedObjectTable_offline" RENAME TO "LiteLLM_ManagedObjectTable"'
|
||||
)
|
||||
|
||||
_SEARCH_STORE: Final = "integration-search-store"
|
||||
_SEARCH_KEY: Final = "synthetic-search-key"
|
||||
_SEARCH_RESULTS: Final = Reply(body=json.dumps({"value": []}).encode())
|
||||
|
||||
Route = Literal[
|
||||
"retrieve",
|
||||
"input_items",
|
||||
"delete",
|
||||
"cancel",
|
||||
"create",
|
||||
"compact",
|
||||
"input_tokens",
|
||||
"openai_passthrough",
|
||||
"percent_encoded",
|
||||
]
|
||||
|
||||
|
||||
def _vendor_id() -> str:
|
||||
return f"resp_{uuid.uuid4().hex}"
|
||||
|
||||
|
||||
def _response(identity: str) -> bytes:
|
||||
return json.dumps(
|
||||
{
|
||||
"id": identity,
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": "completed",
|
||||
"model": _MODEL,
|
||||
"output": [],
|
||||
}
|
||||
).encode()
|
||||
|
||||
|
||||
def _requested_vendor_id(request: Request) -> str | None:
|
||||
if not request.body.startswith(b"{"):
|
||||
return None
|
||||
requested: Final = _JSON_OBJECT.validate_json(request.body).get("input")
|
||||
return requested if isinstance(requested, str) and requested.startswith("resp_") else None
|
||||
|
||||
|
||||
def _respond(request: Request) -> Reply:
|
||||
addressed: Final = next(
|
||||
(segment for segment in urlsplit(request.target).path.split("/") if segment.startswith("resp_")), None
|
||||
)
|
||||
return Reply(body=_response(addressed or _requested_vendor_id(request) or _vendor_id()))
|
||||
|
||||
|
||||
def _forwarded(wire: Wire) -> tuple[tuple[str, str], ...]:
|
||||
return tuple((request.method, urlsplit(request.target).path) for request in wire.drain())
|
||||
|
||||
|
||||
def _send(candidate: Gateway, key: str, method: str, path: str, body: bytes | None = None) -> httpx.Response:
|
||||
json_header: Final = {"Content-Type": "application/json"} if body is not None else {}
|
||||
return candidate.client.request(
|
||||
method, path, content=body, headers={"Authorization": f"Bearer {key}", **json_header}
|
||||
)
|
||||
|
||||
|
||||
def _continue(previous_response_id: str) -> bytes:
|
||||
return json.dumps({"model": _MODEL, "input": "continue", "previous_response_id": previous_response_id}).encode()
|
||||
|
||||
|
||||
def _user_key(scenario: Scenario) -> str:
|
||||
return scenario.key(user_id=scenario.user())
|
||||
|
||||
|
||||
def _create(candidate: Gateway, key: str, vendor_id: str, prefix: str = "/openai") -> str:
|
||||
created: Final = _send(
|
||||
candidate, key, "POST", f"{prefix}/responses", json.dumps({"model": _MODEL, "input": vendor_id}).encode()
|
||||
)
|
||||
assert created.status_code == 200, created.text
|
||||
issued: Final = created.json()["id"]
|
||||
assert isinstance(issued, str) and issued != vendor_id, created.text
|
||||
return issued
|
||||
|
||||
|
||||
def _assert_unmanaged_refusal(response: httpx.Response) -> None:
|
||||
assert response.status_code == 403, response.text
|
||||
assert response.json() == {
|
||||
"error": {"message": _UNMANAGED_DETAIL, "type": "permission_error", "param": None, "code": "403"}
|
||||
}, response.text
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _proxy(
|
||||
gateway: Gateway,
|
||||
directory: Path,
|
||||
upstream: Wire,
|
||||
settings: str = _MANAGED,
|
||||
overrides: Mapping[str, str] = MappingProxyType({}),
|
||||
) -> Iterator[Gateway]:
|
||||
config: Final = directory / f"passthrough-{uuid.uuid4().hex}.yaml"
|
||||
config.write_text(settings)
|
||||
with owned_proxy(
|
||||
gateway,
|
||||
directory,
|
||||
{"OPENAI_API_BASE": upstream.url, "OPENAI_API_KEY": _UPSTREAM_KEY, **overrides},
|
||||
config=config,
|
||||
) as candidate:
|
||||
yield candidate
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _owner_table_offline(database_url: str) -> Iterator[None]:
|
||||
with psycopg.connect(database_url, autocommit=True) as admin:
|
||||
admin.execute(_TAKE_OWNER_TABLE_OFFLINE)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
admin.execute(_BRING_OWNER_TABLE_BACK)
|
||||
|
||||
|
||||
def _search_settings(search: Wire) -> str:
|
||||
return (
|
||||
_MANAGED
|
||||
+ "vector_store_registry:\n"
|
||||
+ f" - vector_store_name: {_SEARCH_STORE}\n"
|
||||
+ " litellm_params:\n"
|
||||
+ f" vector_store_id: {_SEARCH_STORE}\n"
|
||||
+ " custom_llm_provider: azure_ai\n"
|
||||
+ f" api_key: {_SEARCH_KEY}\n"
|
||||
+ f" api_base: {search.url}\n"
|
||||
)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _search_index(gateway: Gateway) -> Iterator[str]:
|
||||
index: Final = f"integration-index-{uuid.uuid4().hex}"
|
||||
gateway.post(
|
||||
"/v1/indexes",
|
||||
{"index_name": index, "litellm_params": {"vector_store_index": index, "vector_store_name": _SEARCH_STORE}},
|
||||
)
|
||||
try:
|
||||
yield index
|
||||
finally:
|
||||
with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as admin:
|
||||
admin.execute('DELETE FROM "LiteLLM_ManagedVectorStoreIndexTable" WHERE index_name = %s', (index,))
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _search_rig(gateway: Gateway, directory: Path) -> Iterator[tuple[Gateway, Wire, str, str]]:
|
||||
with (
|
||||
wire_server(lambda _request: _SEARCH_RESULTS) as search,
|
||||
wire_server(_respond) as openai_upstream,
|
||||
_search_index(gateway) as index,
|
||||
_proxy(gateway, directory, openai_upstream, settings=_search_settings(search)) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
):
|
||||
scenario.model(model=f"azure/{_MODEL}", api_base=openai_upstream.url, api_key=_UPSTREAM_KEY)
|
||||
reader: Final = scenario.key(
|
||||
user_id=scenario.user(),
|
||||
metadata={"allowed_vector_store_indexes": [{"index_name": index, "index_permissions": ["read"]}]},
|
||||
)
|
||||
yield candidate, search, index, reader
|
||||
assert _forwarded(openai_upstream) == ()
|
||||
|
||||
|
||||
def _call(route: Route, vendor_id: str) -> tuple[str, str, bytes | None]:
|
||||
match route:
|
||||
case "retrieve":
|
||||
return ("GET", f"/openai/responses/{vendor_id}", None)
|
||||
case "input_items":
|
||||
return ("GET", f"/openai/responses/{vendor_id}/input_items", None)
|
||||
case "delete":
|
||||
return ("DELETE", f"/openai/responses/{vendor_id}", None)
|
||||
case "cancel":
|
||||
return ("POST", f"/openai/responses/{vendor_id}/cancel", b"{}")
|
||||
case "create":
|
||||
return ("POST", "/openai/responses", _continue(vendor_id))
|
||||
case "compact":
|
||||
return ("POST", "/openai/responses/compact", _continue(vendor_id))
|
||||
case "input_tokens":
|
||||
return ("POST", "/openai/responses/input_tokens", _continue(vendor_id))
|
||||
case "openai_passthrough":
|
||||
return ("GET", f"/openai_passthrough/v1/responses/{vendor_id}", None)
|
||||
case "percent_encoded":
|
||||
return ("GET", f"/openai/responses/%72{vendor_id[1:]}", None)
|
||||
case _:
|
||||
assert_never(route)
|
||||
|
||||
|
||||
def test_the_owner_keeps_using_the_gateway_id_and_another_key_cannot(gateway: Gateway, tmp_path: Path) -> None:
|
||||
vendor_id: Final = _vendor_id()
|
||||
with wire_server(_respond) as upstream, gateway.scenario() as scenario:
|
||||
owner: Final = _user_key(scenario)
|
||||
other: Final = _user_key(scenario)
|
||||
with _proxy(gateway, tmp_path, upstream) as candidate:
|
||||
issued: Final = _create(candidate, owner, vendor_id)
|
||||
assert _forwarded(upstream) == (("POST", "/responses"),)
|
||||
|
||||
retrieved: Final = _send(candidate, owner, "GET", f"/openai/responses/{issued}")
|
||||
assert retrieved.status_code == 200, retrieved.text
|
||||
assert retrieved.json()["id"] == issued, retrieved.text
|
||||
assert _forwarded(upstream) == (("GET", f"/responses/{vendor_id}"),)
|
||||
|
||||
continued: Final = _send(candidate, owner, "POST", "/openai/responses", _continue(issued))
|
||||
assert continued.status_code == 200, continued.text
|
||||
sent: Final = upstream.drain()
|
||||
assert [(request.method, request.target) for request in sent] == [("POST", "/responses")], sent
|
||||
assert _JSON_OBJECT.validate_json(sent[0].body)["previous_response_id"] == vendor_id, sent[0].body
|
||||
|
||||
refused: Final = _send(candidate, other, "GET", f"/openai/responses/{issued}")
|
||||
assert refused.status_code == 403, refused.text
|
||||
assert refused.json()["error"]["message"] == _ACCESS_DENIED, refused.text
|
||||
assert _forwarded(upstream) == ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"route",
|
||||
(
|
||||
"retrieve",
|
||||
"input_items",
|
||||
"delete",
|
||||
"cancel",
|
||||
"create",
|
||||
"compact",
|
||||
"input_tokens",
|
||||
"openai_passthrough",
|
||||
"percent_encoded",
|
||||
),
|
||||
)
|
||||
def test_a_vendor_id_the_gateway_never_issued_is_refused_before_it_reaches_openai(
|
||||
gateway: Gateway, tmp_path: Path, route: Route
|
||||
) -> None:
|
||||
method, path, body = _call(route, _vendor_id())
|
||||
with wire_server(_respond) as upstream, gateway.scenario() as scenario:
|
||||
key: Final = _user_key(scenario)
|
||||
with _proxy(gateway, tmp_path, upstream) as candidate:
|
||||
_assert_unmanaged_refusal(_send(candidate, key, method, path, body))
|
||||
assert _forwarded(upstream) == ()
|
||||
|
||||
|
||||
def test_another_keys_vendor_id_is_not_found_and_never_forwarded(gateway: Gateway, tmp_path: Path) -> None:
|
||||
vendor_id: Final = _vendor_id()
|
||||
with wire_server(_respond) as upstream, gateway.scenario() as scenario:
|
||||
owner: Final = _user_key(scenario)
|
||||
other: Final = _user_key(scenario)
|
||||
with _proxy(gateway, tmp_path, upstream) as candidate:
|
||||
_create(candidate, owner, vendor_id)
|
||||
assert _forwarded(upstream) == (("POST", "/responses"),)
|
||||
refusals: Final = tuple(
|
||||
_send(candidate, other, *_call(route, vendor_id)) for route in ("retrieve", "create")
|
||||
)
|
||||
assert [response.status_code for response in refusals] == [404, 404], [r.text for r in refusals]
|
||||
assert _forwarded(upstream) == ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exemption", ("admin", "allow_unmanaged", "flag_off"))
|
||||
def test_exempt_callers_reach_openai_with_the_vendor_id(
|
||||
gateway: Gateway, tmp_path: Path, exemption: Literal["admin", "allow_unmanaged", "flag_off"]
|
||||
) -> None:
|
||||
vendor_id: Final = _vendor_id()
|
||||
settings: Final = {"admin": _MANAGED, "allow_unmanaged": _ALLOW_UNMANAGED, "flag_off": "general_settings: {}\n"}[
|
||||
exemption
|
||||
]
|
||||
with wire_server(_respond) as upstream, gateway.scenario() as scenario:
|
||||
key: Final = gateway.key if exemption == "admin" else _user_key(scenario)
|
||||
with _proxy(gateway, tmp_path, upstream, settings) as candidate:
|
||||
retrieved: Final = _send(candidate, key, "GET", f"/openai/responses/{vendor_id}")
|
||||
assert retrieved.status_code == 200, retrieved.text
|
||||
continued: Final = _send(candidate, key, "POST", "/openai/responses", _continue(vendor_id))
|
||||
assert continued.status_code == 200, continued.text
|
||||
sent: Final = upstream.drain()
|
||||
assert [(request.method, request.target) for request in sent] == [
|
||||
("GET", f"/responses/{vendor_id}"),
|
||||
("POST", "/responses"),
|
||||
], sent
|
||||
assert _JSON_OBJECT.validate_json(sent[1].body)["previous_response_id"] == vendor_id, sent[1].body
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"value",
|
||||
(123, ["resp_listed"], "", None),
|
||||
ids=("int", "list", "empty", "null"),
|
||||
)
|
||||
def test_a_previous_response_id_that_is_not_an_id_string_is_forwarded_unchanged(
|
||||
gateway: Gateway, tmp_path: Path, value: JsonValue
|
||||
) -> None:
|
||||
body: Final = {"model": _MODEL, "input": "hostile", "previous_response_id": value}
|
||||
with wire_server(_respond) as upstream, gateway.scenario() as scenario:
|
||||
key: Final = _user_key(scenario)
|
||||
with _proxy(gateway, tmp_path, upstream) as candidate:
|
||||
response: Final = _send(candidate, key, "POST", "/openai/responses", json.dumps(body).encode())
|
||||
assert response.status_code == 200, response.text
|
||||
sent: Final = upstream.drain()
|
||||
assert [(request.method, request.target) for request in sent] == [("POST", "/responses")], sent
|
||||
assert _JSON_OBJECT.validate_json(sent[0].body) == body, sent[0].body
|
||||
|
||||
|
||||
def test_an_oversized_vendor_id_is_refused_before_it_reaches_openai(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with wire_server(_respond) as upstream, gateway.scenario() as scenario:
|
||||
key: Final = _user_key(scenario)
|
||||
with _proxy(gateway, tmp_path, upstream) as candidate:
|
||||
_assert_unmanaged_refusal(
|
||||
_send(candidate, key, "POST", "/openai/responses", _continue("resp_" + "a" * 5000))
|
||||
)
|
||||
assert _forwarded(upstream) == ()
|
||||
|
||||
|
||||
def test_a_duplicated_previous_response_id_is_judged_and_forwarded_as_one_value(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
owned_vendor_id: Final = _vendor_id()
|
||||
foreign_vendor_id: Final = _vendor_id()
|
||||
with wire_server(_respond) as upstream, gateway.scenario() as scenario:
|
||||
key: Final = _user_key(scenario)
|
||||
with _proxy(gateway, tmp_path, upstream) as candidate:
|
||||
issued: Final = _create(candidate, key, owned_vendor_id)
|
||||
upstream.drain()
|
||||
vendor_last: Final = (
|
||||
'{"model":"%s","input":"twice","previous_response_id":"%s","previous_response_id":"%s"}'
|
||||
% (_MODEL, issued, foreign_vendor_id)
|
||||
).encode()
|
||||
_assert_unmanaged_refusal(_send(candidate, key, "POST", "/openai/responses", vendor_last))
|
||||
assert _forwarded(upstream) == ()
|
||||
|
||||
issued_last: Final = (
|
||||
'{"model":"%s","input":"twice","previous_response_id":"%s","previous_response_id":"%s"}'
|
||||
% (_MODEL, foreign_vendor_id, issued)
|
||||
).encode()
|
||||
continued: Final = _send(candidate, key, "POST", "/openai/responses", issued_last)
|
||||
assert continued.status_code == 200, continued.text
|
||||
sent: Final = upstream.drain()
|
||||
assert [(request.method, request.target) for request in sent] == [("POST", "/responses")], sent
|
||||
assert sent[0].body.count(b"previous_response_id") == 1, sent[0].body
|
||||
assert _JSON_OBJECT.validate_json(sent[0].body)["previous_response_id"] == owned_vendor_id, sent[0].body
|
||||
|
||||
|
||||
def test_a_failed_owner_lookup_answers_a_retryable_503_and_forwards_nothing(gateway: Gateway, tmp_path: Path) -> None:
|
||||
vendor_id: Final = _vendor_id()
|
||||
with (
|
||||
scratch_database("passthrough_lookup") as database_url,
|
||||
wire_server(_respond) as upstream,
|
||||
_proxy(gateway, tmp_path, upstream, overrides={"DATABASE_URL": database_url}) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
):
|
||||
key: Final = _user_key(scenario)
|
||||
with _owner_table_offline(database_url):
|
||||
failed: Final = _send(candidate, key, "GET", f"/openai/responses/{vendor_id}")
|
||||
assert failed.status_code == 503, failed.text
|
||||
assert failed.json()["error"]["message"] == _LOOKUP_FAILED_DETAIL, failed.text
|
||||
assert _forwarded(upstream) == ()
|
||||
_assert_unmanaged_refusal(_send(candidate, key, "GET", f"/openai/responses/{vendor_id}"))
|
||||
assert _forwarded(upstream) == ()
|
||||
|
||||
|
||||
def test_the_model_list_relay_hands_out_and_checks_gateway_ids(gateway: Gateway, tmp_path: Path) -> None:
|
||||
vendor_id: Final = _vendor_id()
|
||||
with (
|
||||
wire_server(_respond) as azure,
|
||||
wire_server(_respond) as openai_upstream,
|
||||
_proxy(gateway, tmp_path, openai_upstream) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
):
|
||||
group: Final = scenario.model(model=f"azure/{_MODEL}", api_base=azure.url, api_key=_UPSTREAM_KEY)
|
||||
wait_for_worker_readiness(candidate, group, 1, 1)
|
||||
owner: Final = _user_key(scenario)
|
||||
other: Final = _user_key(scenario)
|
||||
relay: Final = f"/azure/{group}/openai/v1"
|
||||
issued: Final = _create(candidate, owner, vendor_id, relay)
|
||||
assert _forwarded(azure) == (("POST", "/openai/v1/responses"),)
|
||||
|
||||
retrieved: Final = _send(candidate, owner, "GET", f"{relay}/responses/{issued}")
|
||||
assert retrieved.status_code == 200, retrieved.text
|
||||
assert retrieved.json()["id"] == issued, retrieved.text
|
||||
assert _forwarded(azure) == (("GET", f"/openai/v1/responses/{vendor_id}"),)
|
||||
|
||||
refusals: Final = tuple(
|
||||
_send(candidate, other, "GET", path)
|
||||
for path in (
|
||||
f"/azure/{group}/openai/v1/responses/{_vendor_id()}",
|
||||
f"/azure_ai/{group}/openai/v1/responses/{_vendor_id()}",
|
||||
f"{relay}/responses/{issued}",
|
||||
f"{relay}/responses/{vendor_id}",
|
||||
)
|
||||
)
|
||||
assert [(response.status_code, response.json()) for response in refusals] == [
|
||||
(403, {"detail": _UNMANAGED_DETAIL}),
|
||||
(403, {"detail": _UNMANAGED_DETAIL}),
|
||||
(403, {"detail": _ACCESS_DENIED}),
|
||||
(404, {"detail": _NOT_FOUND}),
|
||||
], [response.text for response in refusals]
|
||||
assert _forwarded(azure) == ()
|
||||
assert _forwarded(openai_upstream) == ()
|
||||
|
||||
|
||||
def test_an_azure_ai_search_query_reaches_the_index_unchanged(gateway: Gateway, tmp_path: Path) -> None:
|
||||
query: Final = {"search": "mango", "top": 3}
|
||||
with _search_rig(gateway, tmp_path) as (candidate, search, index, reader):
|
||||
answered: Final = _send(
|
||||
candidate, reader, "POST", f"/azure_ai/indexes/{index}/docs/search", json.dumps(query).encode()
|
||||
)
|
||||
assert answered.status_code == 200, answered.text
|
||||
assert answered.json() == {"value": []}, answered.text
|
||||
received: Final = search.drain()
|
||||
assert [(request.method, urlsplit(request.target).path) for request in received] == [
|
||||
("POST", f"/indexes/{index}/docs/search")
|
||||
]
|
||||
assert [_JSON_OBJECT.validate_json(request.body) for request in received] == [query]
|
||||
assert [request.headers.get("api-key") for request in received] == [_SEARCH_KEY]
|
||||
|
||||
|
||||
def test_an_azure_ai_search_query_carrying_a_vendor_response_id_is_refused(gateway: Gateway, tmp_path: Path) -> None:
|
||||
query: Final = {"search": "mango", "previous_response_id": _vendor_id()}
|
||||
with _search_rig(gateway, tmp_path) as (candidate, search, index, reader):
|
||||
_assert_unmanaged_refusal(
|
||||
_send(candidate, reader, "POST", f"/azure_ai/indexes/{index}/docs/search", json.dumps(query).encode())
|
||||
)
|
||||
assert _forwarded(search) == ()
|
||||
|
||||
|
||||
def test_the_model_list_relay_passes_error_and_non_object_bodies_through_unchanged(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
vendor_id: Final = _vendor_id()
|
||||
refused: Final = json.dumps(
|
||||
{"error": {"message": "invalid input", "type": "invalid_request_error"}, "id": vendor_id}
|
||||
)
|
||||
listed: Final = json.dumps([{"id": vendor_id, "object": "response"}])
|
||||
replies: Final = MappingProxyType(
|
||||
{"refused": Reply(status=400, body=refused.encode()), "listed": Reply(body=listed.encode())}
|
||||
)
|
||||
with (
|
||||
wire_server(lambda request: replies[_JSON_OBJECT.validate_json(request.body)["input"]]) as azure,
|
||||
wire_server(_respond) as openai_upstream,
|
||||
_proxy(gateway, tmp_path, openai_upstream) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
):
|
||||
group: Final = scenario.model(model=f"azure/{_MODEL}", api_base=azure.url, api_key=_UPSTREAM_KEY)
|
||||
wait_for_worker_readiness(candidate, group, 1, 1)
|
||||
key: Final = _user_key(scenario)
|
||||
answers: Final = tuple(
|
||||
_send(
|
||||
candidate,
|
||||
key,
|
||||
"POST",
|
||||
f"/azure/{group}/openai/v1/responses",
|
||||
json.dumps({"model": _MODEL, "input": kind}).encode(),
|
||||
)
|
||||
for kind in ("refused", "listed")
|
||||
)
|
||||
assert [(answer.status_code, answer.text) for answer in answers] == [(400, refused), (200, listed)]
|
||||
assert _forwarded(azure) == (("POST", "/openai/v1/responses"), ("POST", "/openai/v1/responses"))
|
||||
assert _forwarded(openai_upstream) == ()
|
||||
|
||||
|
||||
def test_every_verdict_holds_through_a_burst_after_a_worker_is_replaced(gateway: Gateway, tmp_path: Path) -> None:
|
||||
vendor_ids: Final = tuple(_vendor_id() for _ in range(10))
|
||||
config: Final = tmp_path / "passthrough-burst.yaml"
|
||||
config.write_text(_MANAGED)
|
||||
with (
|
||||
wire_server(_respond) as upstream,
|
||||
gateway.scenario() as scenario,
|
||||
owned_proxy_process(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{"OPENAI_API_BASE": upstream.url, "OPENAI_API_KEY": _UPSTREAM_KEY},
|
||||
config=config,
|
||||
workers=2,
|
||||
) as owned,
|
||||
):
|
||||
owner: Final = _user_key(scenario)
|
||||
other: Final = _user_key(scenario)
|
||||
issued: Final = tuple(_create(owned.gateway, owner, vendor_id) for vendor_id in vendor_ids)
|
||||
upstream.drain()
|
||||
workers: Final = owned.worker_pids()
|
||||
assert len(workers) >= 2, workers
|
||||
owned.kill_worker(min(workers))
|
||||
|
||||
calls: Final = (
|
||||
*((owner, f"/openai/responses/{identity}", 200) for identity in issued),
|
||||
*((other, f"/openai/responses/{_vendor_id()}", 403) for _ in range(10)),
|
||||
*((other, f"/openai/responses/{vendor_id}", 404) for vendor_id in vendor_ids),
|
||||
)
|
||||
|
||||
def call(entry: tuple[str, str, int]) -> tuple[int, int]:
|
||||
key, path, expected = entry
|
||||
return expected, _send(owned.gateway, key, "GET", path).status_code
|
||||
|
||||
with ThreadPoolExecutor(max_workers=10) as pool:
|
||||
outcomes: Final = tuple(pool.map(call, calls))
|
||||
assert all(expected == observed for expected, observed in outcomes), outcomes
|
||||
assert sorted(_forwarded(upstream)) == sorted(("GET", f"/responses/{vendor_id}") for vendor_id in vendor_ids)
|
||||
|
|
@ -7503,6 +7503,189 @@ class TestAzureBodyModelGroupRelay:
|
|||
assert route["target"] == "https://operator.openai.azure.com/openai/v1/fine_tuning/jobs"
|
||||
|
||||
|
||||
_VENDOR_RESPONSE_ID: Final = "resp_0000000000000000vendor"
|
||||
|
||||
|
||||
class _ManagedObjectStore:
|
||||
def __init__(self) -> None:
|
||||
self.rows: list[SimpleNamespace] = []
|
||||
self.litellm_managedobjecttable = SimpleNamespace(
|
||||
find_first=AsyncMock(side_effect=self._find_first),
|
||||
count=AsyncMock(side_effect=self._count),
|
||||
upsert=AsyncMock(side_effect=self._upsert),
|
||||
update=AsyncMock(),
|
||||
)
|
||||
|
||||
@property
|
||||
def db(self) -> "_ManagedObjectStore":
|
||||
return self
|
||||
|
||||
async def _find_first(self, where: Mapping[str, str]) -> SimpleNamespace | None:
|
||||
((column, value),) = where.items()
|
||||
return next((row for row in self.rows if getattr(row, column) == value), None)
|
||||
|
||||
async def _count(self, where: Mapping[str, Mapping[str, list[str]]]) -> int:
|
||||
((column, condition),) = where.items()
|
||||
return sum(1 for row in self.rows if getattr(row, column) in condition["in"])
|
||||
|
||||
async def _upsert(self, where: Mapping[str, str], data: Mapping[str, Mapping[str, str | None]]) -> None:
|
||||
self.rows.append(SimpleNamespace(**data["create"]))
|
||||
|
||||
|
||||
class TestAzureRouterRelayManagedIds:
|
||||
def _install(
|
||||
self,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
bodies: Iterator[dict],
|
||||
upstream: _AzureRelayUpstream = _azure_relay_json_upstream,
|
||||
general_settings: Mapping[str, object] = MappingProxyType({}),
|
||||
) -> tuple[list[dict], _ManagedObjectStore]:
|
||||
import litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints as ep
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
captured: list[dict] = []
|
||||
store: Final = _ManagedObjectStore()
|
||||
|
||||
async def next_request_body(_request: Request) -> dict:
|
||||
return next(bodies)
|
||||
|
||||
monkeypatch.setattr(proxy_server, "llm_router", _AzureBodyModelGroupRouter(captured, upstream))
|
||||
monkeypatch.setattr(ep, "get_request_body", next_request_body)
|
||||
monkeypatch.setattr(
|
||||
proxy_server, "general_settings", {"passthrough_managed_object_ids": True, **general_settings}
|
||||
)
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", store)
|
||||
monkeypatch.setattr(proxy_server, "master_key", "sk-test-master-key")
|
||||
monkeypatch.setattr(proxy_server, "proxy_logging_obj", SimpleNamespace(get_proxy_hook=lambda name: object()))
|
||||
monkeypatch.delenv("AZURE_API_BASE", raising=False)
|
||||
return captured, store
|
||||
|
||||
def _request(self, method: str) -> Request:
|
||||
request = MagicMock(spec=Request)
|
||||
request.method = method
|
||||
request.headers = {"content-type": "application/json"}
|
||||
request.query_params = {"api-version": "2025-03-01-preview"}
|
||||
return request
|
||||
|
||||
async def _call(self, method: str, endpoint: str, user_id: str) -> Response:
|
||||
return await azure_proxy_route(
|
||||
endpoint=endpoint,
|
||||
request=self._request(method),
|
||||
fastapi_response=MagicMock(spec=Response),
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key=f"hashed-{user_id}", user_id=user_id),
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("method", "endpoint", "body"),
|
||||
[
|
||||
("POST", "openai/v1/responses", {"model": "gpt-5.4", "previous_response_id": _VENDOR_RESPONSE_ID}),
|
||||
(
|
||||
"POST",
|
||||
"gpt-5.4/openai/v1/responses",
|
||||
{"model": "gpt-5.4", "previous_response_id": _VENDOR_RESPONSE_ID},
|
||||
),
|
||||
("GET", f"gpt-5.4/openai/v1/responses/{_VENDOR_RESPONSE_ID}", {}),
|
||||
("POST", "gpt-5.4/openai/v1/responses/compact", {"previous_response_id": _VENDOR_RESPONSE_ID}),
|
||||
],
|
||||
ids=["body_model_create", "path_model_create", "path_model_retrieve", "path_model_compact"],
|
||||
)
|
||||
async def test_router_relay_refuses_a_response_id_the_pass_through_never_issued(
|
||||
self, monkeypatch, method, endpoint, body
|
||||
):
|
||||
captured, _ = self._install(monkeypatch, iter([body]))
|
||||
|
||||
with pytest.raises(HTTPException) as refusal:
|
||||
await self._call(method, endpoint, "user-2")
|
||||
|
||||
assert refusal.value.status_code == 403
|
||||
assert "allow_unmanaged_response_ids" in refusal.value.detail
|
||||
assert captured == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_router_relay_lets_a_vendor_id_through_with_allow_unmanaged_response_ids(self, monkeypatch):
|
||||
body = {"model": "gpt-5.4", "previous_response_id": _VENDOR_RESPONSE_ID}
|
||||
captured, _ = self._install(monkeypatch, iter([body]), general_settings={"allow_unmanaged_response_ids": True})
|
||||
|
||||
result = await self._call("POST", "openai/v1/responses", "user-2")
|
||||
|
||||
assert result.status_code == 200
|
||||
assert captured[0]["json"]["previous_response_id"] == _VENDOR_RESPONSE_ID
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_id_the_router_relay_issued_keeps_working_for_its_owner_only(self, monkeypatch):
|
||||
from litellm.proxy.pass_through_endpoints.managed_id_codec import decode
|
||||
|
||||
captured, store = self._install(monkeypatch, iter([{"model": "gpt-5.4", "input": "ping"}, {}, {}, {}]))
|
||||
|
||||
created = await self._call("POST", "openai/v1/responses", "user-1")
|
||||
issued_id = json.loads(created.body)["id"]
|
||||
payload = decode(issued_id)
|
||||
assert payload is not None
|
||||
assert (payload.provider, payload.raw_provider_id) == ("azure", "resp_1")
|
||||
assert [(row.model_object_id, row.created_by) for row in store.rows] == [("passthrough:azure:resp_1", "user-1")]
|
||||
|
||||
await self._call("GET", f"gpt-5.4/openai/v1/responses/{issued_id}", "user-1")
|
||||
assert captured[-1]["endpoint"] == "gpt-5.4/openai/v1/responses/resp_1"
|
||||
|
||||
await self._call("GET", "gpt-5.4/openai/v1/responses/resp_1", "user-1")
|
||||
assert captured[-1]["endpoint"] == "gpt-5.4/openai/v1/responses/resp_1"
|
||||
|
||||
relayed_before_the_stranger: Final = len(captured)
|
||||
with pytest.raises(HTTPException) as refusal:
|
||||
await self._call("GET", f"gpt-5.4/openai/v1/responses/{issued_id}", "user-2")
|
||||
assert refusal.value.status_code == 403
|
||||
assert len(captured) == relayed_before_the_stranger
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("method", "endpoint", "upstream_body"),
|
||||
[
|
||||
("GET", "gpt-5.4/openai/v1/batches/batch_1", {"id": "batch_1", "object": "batch"}),
|
||||
("POST", "gpt-5.4/openai/v1/batches/batch_1/cancel", {"id": "batch_1", "object": "batch"}),
|
||||
],
|
||||
)
|
||||
async def test_router_relay_leaves_ids_off_the_responses_routes_as_they_were(
|
||||
self, monkeypatch, method, endpoint, upstream_body
|
||||
):
|
||||
async def upstream() -> httpx.Response:
|
||||
return httpx.Response(200, json=upstream_body)
|
||||
|
||||
captured, store = self._install(monkeypatch, iter([{}]), upstream)
|
||||
|
||||
result = await self._call(method, endpoint, "user-1")
|
||||
|
||||
assert json.loads(result.body) == upstream_body
|
||||
assert captured[-1]["endpoint"] == endpoint
|
||||
assert store.rows == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_streamed_create_through_the_router_relay_hands_out_an_owned_id(self, monkeypatch):
|
||||
from litellm.proxy.pass_through_endpoints.managed_id_codec import decode
|
||||
|
||||
async def upstream_events() -> AsyncIterator[bytes]:
|
||||
yield b'event: response.created\ndata: {"type":"response.created","response":{"id":"resp_1"}}\n\n'
|
||||
yield b'event: response.completed\ndata: {"type":"response.completed","response":{"id":"resp_1"}}\n\n'
|
||||
|
||||
async def streaming_upstream() -> AsyncIterator[bytes]:
|
||||
return upstream_events()
|
||||
|
||||
_, store = self._install(
|
||||
monkeypatch, iter([{"model": "gpt-5.4", "input": "ping", "stream": True}]), streaming_upstream
|
||||
)
|
||||
|
||||
result = await self._call("POST", "openai/v1/responses", "user-1")
|
||||
|
||||
assert isinstance(result, StreamingResponse)
|
||||
streamed = b"".join([chunk async for chunk in result.body_iterator]).decode()
|
||||
assert "resp_1" not in streamed
|
||||
((row),) = store.rows
|
||||
assert row.created_by == "user-1"
|
||||
assert streamed.count(row.unified_object_id) == 2
|
||||
payload = decode(row.unified_object_id)
|
||||
assert payload is not None and payload.raw_provider_id == "resp_1"
|
||||
|
||||
|
||||
AZURE_SPEECH_SHORT_AUDIO_ENDPOINT: Final = "/speech/recognition/conversation/cognitiveservices/v1"
|
||||
AZURE_SPEECH_BATCH_ENDPOINT: Final = "/speechtotext/v3.2/transcriptions"
|
||||
AZURE_SPEECH_FAST_ENDPOINT: Final = "/speechtotext/transcriptions:transcribe"
|
||||
|
|
|
|||
|
|
@ -3,7 +3,9 @@ import json
|
|||
from collections.abc import AsyncIterator, Iterable
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from prisma.errors import PrismaError
|
||||
|
||||
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
|
||||
from litellm.proxy.pass_through_endpoints.managed_id_codec import(
|
||||
|
|
@ -17,6 +19,7 @@ from litellm.proxy.pass_through_endpoints.managed_id_rewriter import(
|
|||
_MAX_RAW_ID_GUARD_LOOKUPS,
|
||||
_passthrough_provider_marker,
|
||||
_resolve_one,
|
||||
addressed_response_id_issuance,
|
||||
is_passthrough_list_route,
|
||||
list_passthrough_ids_from_db,
|
||||
rewrite_body_ids,
|
||||
|
|
@ -1709,6 +1712,127 @@ class TestRawProviderIdGuardBudget:
|
|||
assert exc_info.value.status_code == 400
|
||||
assert pc.db.litellm_managedfiletable.find_many.call_count == _MAX_RAW_ID_GUARD_LOOKUPS
|
||||
|
||||
|
||||
UNISSUED_RESPONSE_ID = "resp_0000000000000000unissued"
|
||||
|
||||
|
||||
def _issued_count(issued: int) -> MagicMock:
|
||||
pc = _prisma_client_passthrough_managed()
|
||||
pc.db.litellm_managedobjecttable.count = AsyncMock(return_value=issued)
|
||||
return pc
|
||||
|
||||
|
||||
class TestAddressedResponseIdIssuance:
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("path", "body"),
|
||||
[
|
||||
(f"/v1/responses/{UNISSUED_RESPONSE_ID}", {}),
|
||||
(f"/v1/responses/{UNISSUED_RESPONSE_ID}/input_items", {}),
|
||||
(f"/v1/responses/{UNISSUED_RESPONSE_ID}/cancel", {}),
|
||||
(f"/openai/responses/{UNISSUED_RESPONSE_ID}", {}),
|
||||
("/v1/responses/compact", {"model": "gpt-5.6", "previous_response_id": UNISSUED_RESPONSE_ID}),
|
||||
("/v1/responses", {"model": "gpt-5.6", "previous_response_id": UNISSUED_RESPONSE_ID}),
|
||||
],
|
||||
)
|
||||
async def test_response_id_without_a_passthrough_row_is_unissued(self, path, body):
|
||||
pc = _issued_count(0)
|
||||
|
||||
assert await addressed_response_id_issuance(path, body, "openai", pc) == "unissued"
|
||||
pc.db.litellm_managedobjecttable.count.assert_awaited_once_with(
|
||||
where={"model_object_id": {"in": [f"passthrough:openai:{UNISSUED_RESPONSE_ID}"]}}
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_response_id_with_a_passthrough_row_is_issued(self):
|
||||
pc = _issued_count(1)
|
||||
|
||||
assert (
|
||||
await addressed_response_id_issuance(f"/v1/responses/{UNISSUED_RESPONSE_ID}", {}, "azure", pc)
|
||||
== "issued"
|
||||
)
|
||||
pc.db.litellm_managedobjecttable.count.assert_awaited_once_with(
|
||||
where={"model_object_id": {"in": [f"passthrough:azure:{UNISSUED_RESPONSE_ID}"]}}
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_any_unissued_id_among_several_is_unissued(self):
|
||||
issued_id = "resp_1111111111111111issued"
|
||||
pc = _prisma_client_passthrough_managed()
|
||||
pc.db.litellm_managedobjecttable.count = AsyncMock(
|
||||
side_effect=lambda where: sum(key.endswith(issued_id) for key in where["model_object_id"]["in"])
|
||||
)
|
||||
|
||||
assert (
|
||||
await addressed_response_id_issuance(
|
||||
f"/v1/responses/{issued_id}", {"previous_response_id": UNISSUED_RESPONSE_ID}, "openai", pc
|
||||
)
|
||||
== "unissued"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("path", "body"),
|
||||
[
|
||||
("/v1/responses", {"model": "gpt-5.6", "input": f"what does {UNISSUED_RESPONSE_ID} mean?"}),
|
||||
("/v1/responses", {"input": [{"role": "user", "content": UNISSUED_RESPONSE_ID}]}),
|
||||
(f"/v1/files/{UNISSUED_RESPONSE_ID}", {}),
|
||||
("/v1/responses", {"previous_response_id": "msg_not_a_response"}),
|
||||
("/v1/responses", {}),
|
||||
],
|
||||
)
|
||||
async def test_requests_that_address_no_response_id_skip_the_lookup(self, path, body):
|
||||
pc = _issued_count(0)
|
||||
|
||||
assert await addressed_response_id_issuance(path, body, "openai", pc) == "issued"
|
||||
pc.db.litellm_managedobjecttable.count.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_managed_ids_are_left_to_the_managed_id_resolver(self):
|
||||
managed_id = new_managed_id("openai", UNISSUED_RESPONSE_ID)
|
||||
pc = _issued_count(0)
|
||||
|
||||
assert (
|
||||
await addressed_response_id_issuance(
|
||||
f"/v1/responses/{managed_id}", {"previous_response_id": managed_id}, "openai", pc
|
||||
)
|
||||
== "issued"
|
||||
)
|
||||
pc.db.litellm_managedobjecttable.count.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("lookup_error", [httpx.ConnectError("db down"), PrismaError("query failed")])
|
||||
async def test_failed_lookup_is_unverified_rather_than_issued_or_unissued(self, lookup_error):
|
||||
pc = _prisma_client_passthrough_managed()
|
||||
pc.db.litellm_managedobjecttable.count = AsyncMock(side_effect=lookup_error)
|
||||
|
||||
assert (
|
||||
await addressed_response_id_issuance(f"/v1/responses/{UNISSUED_RESPONSE_ID}", {}, "openai", pc)
|
||||
== "unverified"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_path_naming_many_response_ids_costs_one_lookup(self):
|
||||
path = "/v1" + "".join(f"/responses/resp_{n:040d}" for n in range(500))
|
||||
pc = _issued_count(0)
|
||||
|
||||
assert await addressed_response_id_issuance(path, {}, "openai", pc) == "unissued"
|
||||
pc.db.litellm_managedobjecttable.count.assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_without_a_database_the_pass_through_issues_only_raw_ids_so_none_is_refused(self):
|
||||
assert (
|
||||
await addressed_response_id_issuance(f"/v1/responses/{UNISSUED_RESPONSE_ID}", {}, "openai", None)
|
||||
== "issued"
|
||||
)
|
||||
assert (
|
||||
await addressed_response_id_issuance(
|
||||
"/v1/responses", {"previous_response_id": UNISSUED_RESPONSE_ID}, "openai", None
|
||||
)
|
||||
== "issued"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_drain_logging_worker", "_vcr_outcome_gate")
|
||||
class TestFlagOff:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -35,7 +35,7 @@ from litellm.integrations.custom_guardrail import CustomGuardrail
|
|||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy._lazy_features import LazyFeature, attach_lazy_features
|
||||
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
|
||||
from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
|
||||
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import router as llm_passthrough_router
|
||||
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
||||
|
|
@ -3678,6 +3678,179 @@ async def test_pass_through_request_merge_query_params_rewrites_managed_ids_on_t
|
|||
}
|
||||
|
||||
|
||||
_UNISSUED_RESPONSE_ID: Final = "resp_0000000000000000unissued"
|
||||
|
||||
|
||||
def _passthrough_object_rows(owner_user_id: str | None) -> SimpleNamespace:
|
||||
row: Final = None if owner_user_id is None else SimpleNamespace(created_by=owner_user_id, team_id=None)
|
||||
return SimpleNamespace(
|
||||
db=SimpleNamespace(
|
||||
litellm_managedobjecttable=SimpleNamespace(
|
||||
find_first=AsyncMock(return_value=row),
|
||||
count=AsyncMock(return_value=0 if row is None else 1),
|
||||
upsert=AsyncMock(),
|
||||
update=AsyncMock(),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def _run_managed_responses_pass_through(
|
||||
method: str,
|
||||
target: str,
|
||||
body: dict[str, object] | None,
|
||||
general_settings: Mapping[str, object],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
prisma_client: SimpleNamespace,
|
||||
) -> tuple[str, int]:
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
|
||||
recorded_requests: Final[list[httpx.Request]] = []
|
||||
|
||||
def transport_handler(upstream_request: httpx.Request) -> httpx.Response:
|
||||
recorded_requests.append(upstream_request)
|
||||
return httpx.Response(200, json={"object": "list", "data": []})
|
||||
|
||||
real_handler: Final = get_async_httpx_client(
|
||||
llm_provider=httpxSpecialProvider.PassThroughEndpoint,
|
||||
params={"timeout": resolve_pass_through_request_timeout(None)},
|
||||
)
|
||||
cache_dict: Final = litellm.in_memory_llm_clients_cache.cache_dict
|
||||
cache_key: Final = next(key for key, cached in cache_dict.items() if cached is real_handler)
|
||||
cache_dict[cache_key] = SimpleNamespace(client=httpx.AsyncClient(transport=httpx.MockTransport(transport_handler)))
|
||||
|
||||
mock_request: Final = MagicMock(spec=Request)
|
||||
mock_request.method = method
|
||||
mock_request.headers = Headers({"content-type": "application/json"})
|
||||
mock_request.query_params = QueryParams("")
|
||||
mock_request.body = AsyncMock(return_value=b"")
|
||||
|
||||
mock_proxy_logging: Final = MagicMock()
|
||||
mock_proxy_logging.pre_call_hook = AsyncMock(
|
||||
side_effect=lambda user_api_key_dict, data, call_type, endpoint_type=None: data
|
||||
)
|
||||
mock_proxy_logging.post_call_failure_hook = AsyncMock()
|
||||
mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={})
|
||||
mock_proxy_logging.get_proxy_hook = MagicMock(
|
||||
return_value=_FakeManagedFilesHook(SimpleNamespace(created_by="user-1", team_id=None))
|
||||
)
|
||||
|
||||
try:
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
{"passthrough_managed_object_ids": True, **general_settings},
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.prisma_client", prisma_client),
|
||||
patch("litellm.proxy.proxy_server.master_key", "sk-test-master-key"),
|
||||
):
|
||||
response: Final = await pass_through_request(
|
||||
request=mock_request,
|
||||
target=target,
|
||||
custom_headers={},
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
custom_body=body,
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
return str(response.status_code), len(recorded_requests)
|
||||
except ProxyException as refusal:
|
||||
return f"{refusal.code} {refusal.message}", len(recorded_requests)
|
||||
finally:
|
||||
cache_dict[cache_key] = real_handler
|
||||
|
||||
|
||||
_RESPONSE_ID_ROUTES: Final = (
|
||||
("GET", f"https://api.openai.com/v1/responses/{_UNISSUED_RESPONSE_ID}", None),
|
||||
("GET", f"https://api.openai.com/v1/responses/{_UNISSUED_RESPONSE_ID}/input_items", None),
|
||||
("DELETE", f"https://api.openai.com/v1/responses/{_UNISSUED_RESPONSE_ID}", None),
|
||||
("POST", f"https://api.openai.com/v1/responses/{_UNISSUED_RESPONSE_ID}/cancel", None),
|
||||
(
|
||||
"POST",
|
||||
"https://api.openai.com/v1/responses/input_tokens",
|
||||
{"model": "gpt-5.6", "previous_response_id": _UNISSUED_RESPONSE_ID},
|
||||
),
|
||||
(
|
||||
"POST",
|
||||
"https://api.openai.com/v1/responses/compact",
|
||||
{"model": "gpt-5.6", "previous_response_id": _UNISSUED_RESPONSE_ID},
|
||||
),
|
||||
(
|
||||
"POST",
|
||||
"https://api.openai.com/v1/responses",
|
||||
{"model": "gpt-5.6", "previous_response_id": _UNISSUED_RESPONSE_ID},
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("method", "target", "body"), _RESPONSE_ID_ROUTES)
|
||||
async def test_managed_pass_through_refuses_a_response_id_it_never_issued(method, target, body):
|
||||
outcome, upstream_calls = await _run_managed_responses_pass_through(
|
||||
method, target, body, {}, UserAPIKeyAuth(user_id="user-2"), _passthrough_object_rows(None)
|
||||
)
|
||||
|
||||
assert outcome.startswith("403 ")
|
||||
assert "allow_unmanaged_response_ids" in outcome
|
||||
assert upstream_calls == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("method", "target", "body"), _RESPONSE_ID_ROUTES)
|
||||
@pytest.mark.parametrize(
|
||||
("general_settings", "user_api_key_dict", "owner_user_id"),
|
||||
[
|
||||
({"allow_unmanaged_response_ids": True}, UserAPIKeyAuth(user_id="user-2"), None),
|
||||
({"disable_responses_id_security": True}, UserAPIKeyAuth(user_id="user-2"), None),
|
||||
({}, UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), None),
|
||||
({}, UserAPIKeyAuth(user_id="user-1"), "user-1"),
|
||||
],
|
||||
ids=["allow_unmanaged_response_ids", "disable_responses_id_security", "proxy_admin", "issued_to_the_caller"],
|
||||
)
|
||||
async def test_managed_pass_through_forwards_ids_it_may_address(
|
||||
method, target, body, general_settings, user_api_key_dict, owner_user_id
|
||||
):
|
||||
outcome, upstream_calls = await _run_managed_responses_pass_through(
|
||||
method, target, body, general_settings, user_api_key_dict, _passthrough_object_rows(owner_user_id)
|
||||
)
|
||||
|
||||
assert (outcome, upstream_calls) == ("200", 1)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("method", "target", "body"), _RESPONSE_ID_ROUTES)
|
||||
async def test_managed_pass_through_answers_a_retryable_503_when_it_cannot_check_a_response_id(method, target, body):
|
||||
unreachable_database: Final = SimpleNamespace(
|
||||
db=SimpleNamespace(
|
||||
litellm_managedobjecttable=SimpleNamespace(count=AsyncMock(side_effect=httpx.ConnectError("db down")))
|
||||
)
|
||||
)
|
||||
|
||||
outcome, upstream_calls = await _run_managed_responses_pass_through(
|
||||
method, target, body, {}, UserAPIKeyAuth(user_id="user-1"), unreachable_database
|
||||
)
|
||||
|
||||
assert outcome.startswith("503 ")
|
||||
assert "Retry the request" in outcome
|
||||
assert upstream_calls == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_managed_pass_through_still_hides_a_response_issued_to_someone_else():
|
||||
outcome, upstream_calls = await _run_managed_responses_pass_through(
|
||||
"GET",
|
||||
f"https://api.openai.com/v1/responses/{_UNISSUED_RESPONSE_ID}",
|
||||
None,
|
||||
{"allow_unmanaged_response_ids": True},
|
||||
UserAPIKeyAuth(user_id="user-2"),
|
||||
_passthrough_object_rows("user-1"),
|
||||
)
|
||||
|
||||
assert outcome.startswith("404 ")
|
||||
assert upstream_calls == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pass_through_request_follows_redirect_to_final_response(httpx_transport):
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue