mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
Merge remote-tracking branch 'origin/main' into litellm_responses_background_stream_resume
# Conflicts: # litellm/responses/streaming_iterator.py
This commit is contained in:
commit
9d6f615392
80 changed files with 4385 additions and 635 deletions
|
|
@ -19,7 +19,7 @@ jobs:
|
|||
name: interactions-spec-canary
|
||||
if: github.repository == 'BerriAI/litellm'
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
timeout-minutes: 30
|
||||
permissions:
|
||||
contents: read
|
||||
outputs:
|
||||
|
|
@ -39,8 +39,11 @@ jobs:
|
|||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Set up the Rust bridge
|
||||
uses: ./.github/actions/rust-bridge
|
||||
|
||||
- name: Install dependencies
|
||||
run: .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy --no-install-project
|
||||
run: .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy
|
||||
|
||||
- name: Run the compliance tests against Google's live spec
|
||||
id: compliance
|
||||
|
|
|
|||
7
.github/workflows/test-merge-smoke.yml
vendored
7
.github/workflows/test-merge-smoke.yml
vendored
|
|
@ -132,7 +132,7 @@ jobs:
|
|||
windows-sdk:
|
||||
name: windows-sdk
|
||||
runs-on: windows-2025
|
||||
timeout-minutes: 10
|
||||
timeout-minutes: 30
|
||||
env:
|
||||
LITELLM_LOCAL_MODEL_COST_MAP: "True"
|
||||
defaults:
|
||||
|
|
@ -154,8 +154,11 @@ jobs:
|
|||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Set up the Rust bridge
|
||||
uses: ./.github/actions/rust-bridge
|
||||
|
||||
- name: Install dependencies
|
||||
run: .github/scripts/uv_sync_with_retries.sh --frozen --no-install-workspace --group dev --python 3.11
|
||||
run: .github/scripts/uv_sync_with_retries.sh --frozen --group dev --python 3.11
|
||||
|
||||
- name: Run the Windows tests
|
||||
run: uv run --no-sync python -m pytest --tb=short -v tests/windows_tests
|
||||
|
|
|
|||
10
.github/workflows/test-unit.yml
vendored
10
.github/workflows/test-unit.yml
vendored
|
|
@ -1306,6 +1306,8 @@ jobs:
|
|||
done
|
||||
openai-sdk-compat:
|
||||
name: OpenAI ${{ matrix.openai-version }} compatibility
|
||||
needs: rust-bridge
|
||||
if: ${{ !cancelled() }}
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: read
|
||||
|
|
@ -1342,10 +1344,18 @@ jobs:
|
|||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Set up the Rust bridge
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 5
|
||||
uses: ./.github/actions/rust-bridge
|
||||
with:
|
||||
artifact: ${{ needs.rust-bridge.outputs.artifact }}
|
||||
|
||||
- name: Install dependencies
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group dev --group proxy-dev --extra proxy --no-install-workspace
|
||||
cp rust-bridge-dist/litellm/rust_bridge/_native.* litellm/rust_bridge/
|
||||
uv pip install "openai==${OPENAI_VERSION:?}" --exclude-newer-package openai=2026-10-07T00:00:00Z
|
||||
uv run --no-sync python -c 'import os, openai; assert openai.__version__ == os.environ["OPENAI_VERSION"], openai.__version__'
|
||||
|
||||
|
|
|
|||
|
|
@ -28,7 +28,7 @@ import httpx
|
|||
from fastapi import HTTPException
|
||||
from openai import APIConnectionError, AsyncOpenAI
|
||||
from openai.types.file_deleted import FileDeleted
|
||||
from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
|
||||
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError
|
||||
from typing_extensions import ReadOnly, Unpack
|
||||
|
||||
import litellm
|
||||
|
|
@ -53,7 +53,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
|||
extract_file_metadata,
|
||||
)
|
||||
from litellm.llms.base_llm.files.storage_backend_factory import get_storage_backend
|
||||
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
|
||||
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints, FileListModelScope
|
||||
from litellm.llms.base_llm.managed_resources.isolation import (
|
||||
build_list_page,
|
||||
build_owner_filter,
|
||||
|
|
@ -309,6 +309,7 @@ def _is_xai_batch_output(
|
|||
class _ManagedFileRow(Protocol):
|
||||
unified_file_id: str
|
||||
file_object: OpenAIFileObject
|
||||
model_mappings: Mapping[str, str]
|
||||
flat_model_file_ids: Sequence[str]
|
||||
storage_backend: Optional[str]
|
||||
storage_url: Optional[str]
|
||||
|
|
@ -362,6 +363,8 @@ class _ManagedResourceDatabase(Protocol):
|
|||
@property
|
||||
def litellm_managedfiletable(self) -> _ManagedFileTableActions: ...
|
||||
|
||||
async def query_raw(self, query: str, *args: object) -> object: ...
|
||||
|
||||
@property
|
||||
def litellm_managedobjecttable(self) -> _ManagedObjectTableActions: ...
|
||||
|
||||
|
|
@ -385,10 +388,83 @@ class _RouterFileCallKwargs(TypedDict, total=False):
|
|||
custom_llm_provider: ReadOnly[str | None]
|
||||
|
||||
|
||||
def _managed_database(prisma_client: _ManagedResourcePrismaClient) -> _ManagedResourceDatabase:
|
||||
return prisma_client.db
|
||||
|
||||
|
||||
def _managed_file_table(prisma_client: _ManagedResourcePrismaClient) -> _ManagedFileTableActions:
|
||||
return prisma_client.db.litellm_managedfiletable
|
||||
|
||||
|
||||
class _ListedFileRow(Protocol):
|
||||
@property
|
||||
def unified_file_id(self) -> str: ...
|
||||
|
||||
@property
|
||||
def file_object(self) -> object: ...
|
||||
|
||||
@property
|
||||
def model_mappings(self) -> Mapping[str, object]: ...
|
||||
|
||||
@property
|
||||
def storage_backend(self) -> str | None: ...
|
||||
|
||||
|
||||
class _ScopedFileRow(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
unified_file_id: str
|
||||
file_object: JsonValue = None
|
||||
model_mappings: Mapping[str, JsonValue]
|
||||
storage_backend: str | None = None
|
||||
|
||||
|
||||
_SCOPED_FILE_ROWS: Final = TypeAdapter(tuple[_ScopedFileRow, ...])
|
||||
|
||||
|
||||
def _owner_condition(owner_filter: Mapping[str, object]) -> tuple[str, tuple[str, ...]]:
|
||||
match owner_filter:
|
||||
case {"OR": [{"created_by": str() as user_id}, {"team_id": str() as team_id}]}:
|
||||
return '("created_by" = $1 OR "team_id" = $2)', (user_id, team_id)
|
||||
case {"created_by": str() as user_id}:
|
||||
return '"created_by" = $1', (user_id,)
|
||||
case {"team_id": str() as team_id}:
|
||||
return '"team_id" = $1', (team_id,)
|
||||
case _:
|
||||
return ("FALSE", ()) if owner_filter else ("TRUE", ())
|
||||
|
||||
|
||||
async def _scoped_file_rows(
|
||||
database: _ManagedResourceDatabase,
|
||||
owner_filter: Mapping[str, object],
|
||||
model_mapping_keys: frozenset[str],
|
||||
cursor_id: str | None,
|
||||
direction: Literal["asc", "desc"],
|
||||
take: int,
|
||||
) -> tuple[_ScopedFileRow, ...]:
|
||||
owner_condition, owner_params = _owner_condition(owner_filter)
|
||||
keys_param: Final = len(owner_params) + 1
|
||||
comparison: Final = "<" if direction == "desc" else ">"
|
||||
cursor_condition: Final = (
|
||||
f'AND ("created_at", "unified_file_id") {comparison} (SELECT "created_at", "unified_file_id" '
|
||||
f'FROM "LiteLLM_ManagedFileTable" WHERE "unified_file_id" = ${keys_param + 1}) '
|
||||
if cursor_id
|
||||
else ""
|
||||
)
|
||||
order: Final = direction.upper()
|
||||
rows: Final = await database.query_raw(
|
||||
'SELECT "unified_file_id", "file_object", "model_mappings", "storage_backend" FROM "LiteLLM_ManagedFileTable" '
|
||||
f'WHERE {owner_condition} AND ("model_mappings" ?| ${keys_param}::text[] OR "storage_backend" IS NOT NULL) '
|
||||
f'{cursor_condition}ORDER BY "created_at" {order}, "unified_file_id" {order} '
|
||||
f"LIMIT ${keys_param + (2 if cursor_id else 1)}",
|
||||
*owner_params,
|
||||
sorted(model_mapping_keys),
|
||||
*((cursor_id,) if cursor_id else ()),
|
||||
take,
|
||||
)
|
||||
return _SCOPED_FILE_ROWS.validate_python(rows)
|
||||
|
||||
|
||||
def _iter_provider_file_id_pairs(
|
||||
rows: Sequence[_ManagedFileRow],
|
||||
requested_provider_file_ids: frozenset[str],
|
||||
|
|
@ -2265,14 +2341,23 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
user_api_key_dict: UserAPIKeyAuth,
|
||||
limit: Optional[int] = None,
|
||||
after: Optional[str] = None,
|
||||
model_scope: Optional[FileListModelScope] = None,
|
||||
before: Optional[str] = None,
|
||||
**data: Dict,
|
||||
) -> FileListPage:
|
||||
"""List the managed files the caller owns, newest first.
|
||||
|
||||
Pagination is keyset based on ``unified_file_id`` so a key that owns
|
||||
every file on the proxy still reads one bounded page at a time.
|
||||
``purpose`` is applied after parsing, because the managed file table
|
||||
keeps it inside the ``file_object`` blob instead of a column, and rows
|
||||
``after`` pages toward older files and ``before`` toward newer ones,
|
||||
each page still ordered newest first.
|
||||
``model_scope`` is applied in the query, so only rows whose
|
||||
``model_mappings`` name one of its deployment ids or model names, and
|
||||
storage-backed rows, leave the database; a storage-backed row then
|
||||
stays only when its model name is one the router sends to the listed
|
||||
provider. ``purpose`` is applied after parsing, because the managed
|
||||
file table keeps it inside the ``file_object`` blob instead of a
|
||||
column. Rows
|
||||
whose blob will not parse drop out there too, so a chunk of rows can
|
||||
yield fewer matches than the page holds. Successive chunks are read
|
||||
until the page is full or the caller's rows run out, which keeps
|
||||
|
|
@ -2281,9 +2366,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
scan has to continue past it, the chunk widens to
|
||||
``FILE_LIST_CONTINUATION_CHUNK_SIZE``, so the walk costs one query per
|
||||
that many rows instead of one per page. That bound is per query, not
|
||||
per request: the work is still linear in the rows the caller owns, and
|
||||
a filter matching nothing reads every one of them, with no index
|
||||
covering either the owner filter or the sort.
|
||||
per request: a ``purpose`` filter matching nothing still reads every
|
||||
row that passes the query.
|
||||
"""
|
||||
validate_file_list_limit(limit)
|
||||
validate_file_list_purpose(purpose)
|
||||
|
|
@ -2292,36 +2376,60 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
if owner_filter is None:
|
||||
return FileListPage.model_validate(build_list_page([]))
|
||||
|
||||
if after:
|
||||
cursor_row = await _managed_file_table(self.prisma_client).find_first(
|
||||
where={**owner_filter, "unified_file_id": after}
|
||||
if after and before:
|
||||
raise ProxyException(
|
||||
message="Pass either 'after' (or 'after_id') or 'before_id' when listing files, not both.",
|
||||
type="invalid_request_error",
|
||||
param="before_id",
|
||||
code=400,
|
||||
openai_code="invalid_value",
|
||||
)
|
||||
cursor_param: Final = "before_id" if before else "after"
|
||||
cursor: Final = before or after
|
||||
database: Final = _managed_database(self.prisma_client)
|
||||
if cursor:
|
||||
cursor_row = await database.litellm_managedfiletable.find_first(
|
||||
where={**owner_filter, "unified_file_id": cursor}
|
||||
)
|
||||
if cursor_row is None:
|
||||
raise ProxyException(
|
||||
message=f"Invalid 'after' cursor: no file found with id '{after}'.",
|
||||
message=f"Invalid '{cursor_param}' cursor: no file found with id '{cursor}'.",
|
||||
type="invalid_request_error",
|
||||
param="after",
|
||||
param=cursor_param,
|
||||
code=400,
|
||||
openai_code="invalid_value",
|
||||
)
|
||||
|
||||
if model_scope is not None and not model_scope.model_mapping_keys:
|
||||
return FileListPage.model_validate(build_list_page([]))
|
||||
|
||||
page_size: Final = min(limit or MAX_FILE_LIST_LIMIT, MAX_FILE_LIST_LIMIT)
|
||||
matches: Final[List[OpenAIFileObject]] = []
|
||||
cursor_id = after
|
||||
cursor_id = cursor
|
||||
chunk_size = page_size + 1
|
||||
direction: Final[Literal["asc", "desc"]] = "asc" if before else "desc"
|
||||
|
||||
while len(matches) <= page_size:
|
||||
cursor_args: _CursorPageArgs = {"cursor": {"unified_file_id": cursor_id}, "skip": 1} if cursor_id else {}
|
||||
chunk = await _managed_file_table(self.prisma_client).find_many(
|
||||
where=owner_filter,
|
||||
take=chunk_size,
|
||||
order=[{"created_at": "desc"}, {"unified_file_id": "desc"}],
|
||||
**cursor_args,
|
||||
chunk: Sequence[_ListedFileRow] = (
|
||||
await _scoped_file_rows(
|
||||
database, owner_filter, model_scope.model_mapping_keys, cursor_id, direction, chunk_size
|
||||
)
|
||||
if model_scope is not None
|
||||
else await database.litellm_managedfiletable.find_many(
|
||||
where=owner_filter,
|
||||
take=chunk_size,
|
||||
order=[{"created_at": direction}, {"unified_file_id": direction}],
|
||||
**cursor_args,
|
||||
)
|
||||
)
|
||||
matches.extend(
|
||||
_public_file_object(parsed_file_object, row.unified_file_id)
|
||||
for row in chunk
|
||||
if (parsed_file_object := _parse_managed_file_object(row.file_object, row.unified_file_id)) is not None
|
||||
if (
|
||||
model_scope is None or model_scope.holds(row.model_mappings, stored=row.storage_backend is not None)
|
||||
)
|
||||
and (parsed_file_object := _parse_managed_file_object(row.file_object, row.unified_file_id)) is not None
|
||||
and (purpose is None or parsed_file_object.purpose == purpose)
|
||||
)
|
||||
if len(chunk) < chunk_size:
|
||||
|
|
@ -2329,7 +2437,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
cursor_id = chunk[-1].unified_file_id
|
||||
chunk_size = max(chunk_size, FILE_LIST_CONTINUATION_CHUNK_SIZE)
|
||||
|
||||
return FileListPage.model_validate(build_list_page(matches[:page_size], has_more=len(matches) > page_size))
|
||||
page: Final = matches[:page_size]
|
||||
return FileListPage.model_validate(
|
||||
build_list_page(page[::-1] if before else page, has_more=len(matches) > page_size)
|
||||
)
|
||||
|
||||
def _is_batch_polling_enabled(self) -> bool:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -153,7 +153,7 @@ RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
|
|||
find /app/.venv -type d -path "*/tornado/test" -delete && \
|
||||
chmod -R a+rX /opt/prisma && \
|
||||
python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths" && \
|
||||
python -c "import litellm; from litellm.rust_bridge.loader import native_bridge_available; assert litellm.__file__ == '/app/litellm/__init__.py', litellm.__file__; assert native_bridge_available()"
|
||||
python -c "import litellm, litellm.rust_bridge._native; assert litellm.__file__ == '/app/litellm/__init__.py', litellm.__file__"
|
||||
|
||||
USER nonroot
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,2 @@
|
|||
CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_ManagedFileTable_created_by_created_at_idx"
|
||||
ON "LiteLLM_ManagedFileTable" ("created_by", "created_at" DESC);
|
||||
|
|
@ -1163,6 +1163,7 @@ model LiteLLM_ManagedFileTable {
|
|||
@@index([unified_file_id])
|
||||
@@index([flat_model_file_ids], type: Gin)
|
||||
@@index([team_id, created_at(sort: Desc)])
|
||||
@@index([created_by, created_at(sort: Desc)])
|
||||
}
|
||||
|
||||
model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use the
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ const FINISH_REASONS: &[(&str, &str)] = &[
|
|||
("max_tokens", "length"),
|
||||
("refusal", "content_filter"),
|
||||
("compaction", "length"),
|
||||
("model_context_window_exceeded", "length"),
|
||||
("guardrail_intervened", "content_filter"),
|
||||
("content_filtered", "content_filter"),
|
||||
("content_filter", "content_filter"),
|
||||
|
|
@ -74,6 +75,7 @@ mod tests {
|
|||
#[case::end_turn("end_turn", "stop")]
|
||||
#[case::stop_sequence("stop_sequence", "stop")]
|
||||
#[case::max_tokens("max_tokens", "length")]
|
||||
#[case::model_context_window_exceeded("model_context_window_exceeded", "length")]
|
||||
#[case::refusal("refusal", "content_filter")]
|
||||
#[case::guardrail_intervened("guardrail_intervened", "content_filter")]
|
||||
#[case::content_filtered("content_filtered", "content_filter")]
|
||||
|
|
|
|||
|
|
@ -187,6 +187,7 @@ fn native_events_reach_python_with_levels_context_reentry_and_http_deduplication
|
|||
locals
|
||||
.set_item("http_warning", wrap_pyfunction!(http_warning, py).unwrap())
|
||||
.unwrap();
|
||||
locals.set_item("native", crate::native_module(py)).unwrap();
|
||||
let importable = py
|
||||
.eval(
|
||||
c"__import__('importlib.util', fromlist=['util']).find_spec('dotenv') is not None",
|
||||
|
|
@ -205,6 +206,7 @@ import asyncio
|
|||
import logging
|
||||
import sys
|
||||
sys.path.insert(0, repo_root)
|
||||
sys.modules['litellm.rust_bridge._native'] = native
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger, session_id_var, trace_id_var
|
||||
|
||||
|
|
|
|||
|
|
@ -129,6 +129,7 @@ MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES: Final = 1024
|
|||
S3_LOG_PROMPTS_ONLY_ENV_VAR: Final = "S3_LOG_PROMPTS_ONLY"
|
||||
S3_PARTITION_GRANULARITY_ENV_VAR: Final = "S3_PARTITION_GRANULARITY"
|
||||
MAX_FILE_LIST_LIMIT: Final = 10000
|
||||
ANTHROPIC_DEFAULT_FILE_LIST_LIMIT: Final = 20
|
||||
DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = get_env_int("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10)
|
||||
DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = get_env_int("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1)
|
||||
GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS: Final = 10
|
||||
|
|
|
|||
|
|
@ -11,4 +11,4 @@ Core files:
|
|||
- `duration_parser.py`: code for parsing durations - e.g. "1d", "1mo", "10s"
|
||||
- `api_route_to_call_types.py`: mapping of API routes to their corresponding CallTypes (e.g., `/chat/completions` -> [acompletion, completion])
|
||||
|
||||
Tokenizer factories return Python tokenizer objects by default. Set `LITELLM_RUST=1` or call `litellm.rust(True)` before constructing tokenizers to select the Rust backend through `Route.TOKENIZER` in the Rust catalog. Missing native bindings or unsupported native features fall back to Python. Existing tokenizer objects keep their selected backend. Rust-backed tokenizer objects carry the read-only `tiktoken.Encoding` / `tokenizers.Tokenizer` surface and are immutable: `enable_padding`, `enable_truncation` and `add_tokens` stay on the Python tokenizer.
|
||||
Tokenizer factories return Python tokenizer objects by default. Set `LITELLM_RUST=1` or call `litellm.rust(True)` before constructing tokenizers to select the Rust backend through `Route.TOKENIZER` in the Rust catalog. Unsupported native features fall back to Python. Existing tokenizer objects keep their selected backend. Rust-backed tokenizer objects carry the read-only `tiktoken.Encoding` / `tokenizers.Tokenizer` surface and are immutable: `enable_padding`, `enable_truncation` and `add_tokens` stay on the Python tokenizer.
|
||||
|
|
|
|||
|
|
@ -353,9 +353,9 @@ def is_cli_token_fresh(
|
|||
buffer_hours: float = CLI_TOKEN_FRESHNESS_BUFFER_SECONDS / 3600,
|
||||
) -> bool:
|
||||
"""Check whether a cached CLI token is still within its expiration window.
|
||||
Used by `lite auth print-token` to fail fast, without a network round trip,
|
||||
once the cached token is past `LITELLM_CLI_JWT_EXPIRATION_HOURS`. A `--pkce`
|
||||
credential carries its own `expires_at`, which is authoritative when present."""
|
||||
Used by `lite auth print-token` to skip asking the proxy while the cached token
|
||||
is younger than `LITELLM_CLI_JWT_EXPIRATION_HOURS`. A `--pkce` credential
|
||||
carries its own `expires_at`, which is authoritative when present."""
|
||||
from litellm.constants import CLI_JWT_EXPIRATION_HOURS
|
||||
|
||||
expires_at: Final = (
|
||||
|
|
|
|||
|
|
@ -266,6 +266,7 @@ _FINISH_REASON_MAP: Final[dict[str, OpenAIChatCompletionFinishReason]] = {
|
|||
"tool_use": "tool_calls",
|
||||
"refusal": "content_filter",
|
||||
"compaction": "length",
|
||||
"model_context_window_exceeded": "length",
|
||||
# Cohere
|
||||
"COMPLETE": "stop",
|
||||
"ERROR_TOXIC": "content_filter",
|
||||
|
|
|
|||
|
|
@ -842,6 +842,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
self._async_success_enqueued: bool = False
|
||||
self._sync_success_scheduled: bool = False
|
||||
self.on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None
|
||||
self._stream_failure_claimed: bool = False
|
||||
self.shadow_eval_request_snapshot: GuardrailRequestSnapshot | None = None
|
||||
|
||||
def copy_for_new_call(self, litellm_call_id: str, start_time: datetime.datetime) -> "Logging":
|
||||
|
|
@ -2561,6 +2562,12 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
},
|
||||
)
|
||||
|
||||
def claim_stream_failure_report(self) -> bool:
|
||||
if self._stream_failure_claimed:
|
||||
return False
|
||||
self._stream_failure_claimed = True
|
||||
return True
|
||||
|
||||
def record_assembled_response_for_failure(self, assembled: ModelResponse) -> None:
|
||||
"""Bill a fully streamed response on the failure log when a post-call hook rejects it."""
|
||||
usage: Final = getattr(assembled, "usage", None)
|
||||
|
|
|
|||
|
|
@ -1557,6 +1557,12 @@ class CustomStreamWrapper:
|
|||
else: # openai / azure chat model
|
||||
if isinstance(chunk, ChatCompletionChunk):
|
||||
self._openai_handler_receives_chat_chunks = True
|
||||
if any(str(choice.finish_reason) == "abort" for choice in chunk.choices or ()):
|
||||
raise litellm.InternalServerError(
|
||||
message="Upstream generation aborted (finish_reason=abort)",
|
||||
model=chunk.model,
|
||||
llm_provider=self.custom_llm_provider,
|
||||
)
|
||||
if self.custom_llm_provider in [
|
||||
LlmProviders.AZURE.value,
|
||||
LlmProviders.AZURE_AI.value,
|
||||
|
|
|
|||
|
|
@ -82,4 +82,6 @@ class AiohttpOpenAIChatConfig(OpenAILikeChatConfig):
|
|||
model_response.model = _json_response.get("model")
|
||||
model_response.object = _json_response.get("object")
|
||||
model_response.system_fingerprint = _json_response.get("system_fingerprint")
|
||||
if "service_tier" in _json_response:
|
||||
return model_response.model_copy(update={"service_tier": _json_response.get("service_tier")})
|
||||
return model_response
|
||||
|
|
|
|||
|
|
@ -2,15 +2,16 @@
|
|||
## Translates OpenAI call to Anthropic `/v1/messages` format
|
||||
import asyncio
|
||||
from collections import deque
|
||||
from collections.abc import AsyncGenerator, Iterator, Mapping
|
||||
from collections.abc import AsyncGenerator, Awaitable, Callable, Iterator, Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
import anyio
|
||||
from pydantic import ConfigDict, field_validator
|
||||
|
||||
from litellm import verbose_logger
|
||||
from litellm._logging import redact_internal_details_from_client_message
|
||||
from litellm._uuid import uuid
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
from litellm.exceptions import APIError, MidStreamFallbackError
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
encrypted_reasoning_signature,
|
||||
)
|
||||
|
|
@ -22,6 +23,7 @@ from litellm.llms.anthropic.messages.native.utils import (
|
|||
from litellm.llms.anthropic.messages.safeguards import StreamedSafeguardResults
|
||||
from litellm.llms.anthropic.messages.utils import anthropic_sse_frame
|
||||
from litellm.responses.streaming_iterator import (
|
||||
map_stream_error_to_exception,
|
||||
partial_responses_incomplete_event_on_disconnect,
|
||||
stream_error_status_and_message,
|
||||
)
|
||||
|
|
@ -130,10 +132,12 @@ class AnthropicResponsesStreamWrapper:
|
|||
) -> None:
|
||||
self.responses_stream = responses_stream
|
||||
self.model = model
|
||||
self.litellm_logging_obj = litellm_logging_obj
|
||||
self.safeguards_evaluator: Final = safeguards_evaluator
|
||||
self._message_id: str = f"msg_{uuid.uuid4()}"
|
||||
if litellm_logging_obj is not None:
|
||||
litellm_logging_obj.record_streamed_anthropic_message_id(self._message_id)
|
||||
self._unreported_failure: Exception | None = None
|
||||
self._current_block_index: int = -1
|
||||
# Map item_id -> content_block_index so we can stop the right block later
|
||||
self._item_id_to_block_index: dict[str, int] = {}
|
||||
|
|
@ -377,16 +381,24 @@ class AnthropicResponsesStreamWrapper:
|
|||
|
||||
if event_type == "response.failed":
|
||||
failed: Final = _FailedResponseEvent.model_validate(event)
|
||||
status_code, message = stream_error_status_and_message(
|
||||
failed.response.error if failed.response is not None else None
|
||||
)
|
||||
upstream_error: Final = failed.response.error if failed.response is not None else None
|
||||
status_code, message = stream_error_status_and_message(upstream_error)
|
||||
verbose_logger.error(
|
||||
"AnthropicResponsesStreamWrapper: upstream Responses stream for %s failed (%s): %s",
|
||||
self.model,
|
||||
status_code,
|
||||
message,
|
||||
)
|
||||
self._fail_stream(status_code, message)
|
||||
self._fail_stream(
|
||||
status_code,
|
||||
message,
|
||||
failure=map_stream_error_to_exception(
|
||||
upstream_error,
|
||||
self.model,
|
||||
getattr(self.responses_stream, "custom_llm_provider", None) or "",
|
||||
logging_obj=self.litellm_logging_obj,
|
||||
),
|
||||
)
|
||||
return
|
||||
|
||||
# ---- response completed -> message_delta + message_stop ----
|
||||
|
|
@ -439,10 +451,36 @@ class AnthropicResponsesStreamWrapper:
|
|||
self._sent_message_stop = True
|
||||
return
|
||||
|
||||
def _fail_stream(self, status_code: int, message: str) -> None:
|
||||
def _fail_stream(self, status_code: int, message: str, failure: Exception | None = None) -> None:
|
||||
self._stream_failed = True
|
||||
self._unreported_failure = failure or APIError(
|
||||
status_code=status_code,
|
||||
message=message,
|
||||
llm_provider=getattr(self.responses_stream, "custom_llm_provider", None),
|
||||
model=self.model,
|
||||
)
|
||||
self._chunk_queue.append(_anthropic_error_chunk(status_code, message))
|
||||
|
||||
async def _report_unreported_failure(self) -> None:
|
||||
failure: Final = self._unreported_failure
|
||||
if failure is None:
|
||||
return
|
||||
self._unreported_failure = None
|
||||
on_failure: Final[Callable[[Exception], Awaitable[None]] | None] = getattr(
|
||||
self.litellm_logging_obj, "_on_detached_stream_failure", None
|
||||
)
|
||||
if on_failure is None:
|
||||
return
|
||||
try:
|
||||
with anyio.CancelScope(shield=True):
|
||||
await on_failure(failure)
|
||||
except Exception as hook_failure: # noqa: BLE001 # a failing proxy hook must not hide the error event from the client
|
||||
verbose_logger.warning(
|
||||
"AnthropicResponsesStreamWrapper: stream failure hook raised: %s(%s)",
|
||||
type(hook_failure).__name__,
|
||||
hook_failure,
|
||||
)
|
||||
|
||||
def __aiter__(self) -> "AnthropicResponsesStreamWrapper":
|
||||
return self
|
||||
|
||||
|
|
@ -450,6 +488,7 @@ class AnthropicResponsesStreamWrapper:
|
|||
if self._chunk_queue:
|
||||
return self._chunk_queue.popleft()
|
||||
if self._stream_failed:
|
||||
await self._report_unreported_failure()
|
||||
raise StopAsyncIteration
|
||||
|
||||
if not self._sent_message_start:
|
||||
|
|
@ -478,7 +517,7 @@ class AnthropicResponsesStreamWrapper:
|
|||
verbose_logger.exception(
|
||||
"AnthropicResponsesStreamWrapper: upstream Responses stream for %s failed", self.model
|
||||
)
|
||||
self._fail_stream(*_failure_status_and_message(e))
|
||||
self._fail_stream(*_failure_status_and_message(e), failure=_original_failure(e))
|
||||
|
||||
if not self._chunk_queue and not self._sent_message_stop and not self._stream_failed:
|
||||
verbose_logger.error(
|
||||
|
|
@ -498,8 +537,14 @@ class AnthropicResponsesStreamWrapper:
|
|||
|
||||
async def _sse_frames(self) -> AsyncGenerator[bytes, None]:
|
||||
safeguard_results: Final = StreamedSafeguardResults(self.safeguards_evaluator)
|
||||
async for chunk in self:
|
||||
yield anthropic_sse_frame(await safeguard_results.observe(chunk))
|
||||
try:
|
||||
async for chunk in self:
|
||||
yield anthropic_sse_frame(await safeguard_results.observe(chunk))
|
||||
finally:
|
||||
from litellm.llms.anthropic.messages.native.streaming_iterator import closing_superseded_attempt
|
||||
|
||||
if not closing_superseded_attempt():
|
||||
await self._report_unreported_failure()
|
||||
|
||||
def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None":
|
||||
return partial_responses_incomplete_event_on_disconnect(self.responses_stream)
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
from abc import ABC, abstractmethod
|
||||
from collections.abc import AsyncGenerator, Iterator, Mapping
|
||||
from typing import TYPE_CHECKING, Any, Union
|
||||
from collections.abc import AsyncGenerator, Callable, Iterable, Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, Final, Union
|
||||
|
||||
import httpx
|
||||
from openai.types.file_deleted import FileDeleted
|
||||
|
|
@ -38,6 +39,22 @@ else:
|
|||
UserAPIKeyAuth = Any
|
||||
|
||||
|
||||
def _routes_nowhere(_model_name: str) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class FileListModelScope:
|
||||
model_mapping_keys: frozenset[str]
|
||||
routes_stored_model_name: Callable[[str], bool] = _routes_nowhere
|
||||
|
||||
def holds(self, model_mapping_keys: Iterable[str], stored: bool) -> bool:
|
||||
mapped_keys: Final = frozenset(model_mapping_keys)
|
||||
return not self.model_mapping_keys.isdisjoint(mapped_keys) or (
|
||||
stored and any(map(self.routes_stored_model_name, mapped_keys))
|
||||
)
|
||||
|
||||
|
||||
class BaseFileUploadStream(ABC):
|
||||
"""Re-iterable request body that yields an upload's bytes lazily.
|
||||
|
||||
|
|
@ -274,6 +291,8 @@ class BaseFileEndpoints(ABC):
|
|||
user_api_key_dict: UserAPIKeyAuth,
|
||||
limit: int | None = None,
|
||||
after: str | None = None,
|
||||
model_scope: FileListModelScope | None = None,
|
||||
before: str | None = None,
|
||||
**data: dict,
|
||||
) -> FileListPage:
|
||||
pass
|
||||
|
|
|
|||
|
|
@ -76,6 +76,7 @@ else:
|
|||
|
||||
_NO_TOOLS_UPDATE: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
_OBJECT_MAPPING: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
|
||||
_CHOICE_MAPPINGS: Final = TypeAdapter(tuple[Mapping[str, object], ...])
|
||||
_DECODED_JSON: Final = TypeAdapter(object)
|
||||
|
||||
|
||||
|
|
@ -958,8 +959,13 @@ class OpenAIChatCompletionStreamingHandler(BaseModelResponseIterator):
|
|||
status_code=error_status_code,
|
||||
message=error_message,
|
||||
)
|
||||
choices = chunk.get("choices", [])
|
||||
choices = self._map_reasoning_to_reasoning_content(choices)
|
||||
raw_choices: Final = chunk.get("choices") or []
|
||||
if any(choice.get("finish_reason") == "abort" for choice in _CHOICE_MAPPINGS.validate_python(raw_choices)):
|
||||
raise OpenAIError(
|
||||
status_code=500,
|
||||
message="Upstream generation aborted (finish_reason=abort)",
|
||||
)
|
||||
choices: Final = self._map_reasoning_to_reasoning_content(raw_choices)
|
||||
fingerprint: Final = _OBJECT_MAPPING.validate_python(chunk).get("system_fingerprint")
|
||||
|
||||
kwargs: Final[dict[str, Any]] = {
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -1225,7 +1225,7 @@
|
|||
"responses": true,
|
||||
"text_completion": true,
|
||||
"embeddings": true,
|
||||
"image_generations": false,
|
||||
"image_generations": true,
|
||||
"image_edits": false,
|
||||
"audio_transcriptions": false,
|
||||
"audio_speech": false,
|
||||
|
|
|
|||
|
|
@ -342,7 +342,7 @@ The CLI provides these authentication commands:
|
|||
- **`lite login --pkce`** - Sign in through the system browser with OAuth authorization code + PKCE; the key renews itself with a refresh token
|
||||
- **`lite logout`** - Clear stored authentication token (and revoke a `--pkce` refresh token on the proxy)
|
||||
- **`lite whoami`** - Show current authentication status
|
||||
- **`lite auth print-token`** - Print the cached token (used as Claude Code's `apiKeyHelper`); renews a `--pkce` key first and fails once a classic token has expired
|
||||
- **`lite auth print-token`** - Print the cached token (used as Claude Code's `apiKeyHelper`); renews a `--pkce` key first, and fails once a classic token is past 24 hours and the proxy rejects it
|
||||
|
||||
### Authentication Flow Steps
|
||||
|
||||
|
|
@ -382,7 +382,7 @@ Keychain storage needs the `keyring` package, which ships with `pip install 'lit
|
|||
|
||||
`lite logout` clears both stores. If the keychain is locked at that moment it says so, and re-running it once the keychain is unlocked finishes the job.
|
||||
|
||||
The stored credential is a short-lived, per-session agent token, not a managed virtual key. It is scoped to the user and team you logged in as and inherits their models and budgets; spend is tracked against the shared team and user budgets rather than a separate per-session cap, so multiple logins or several concurrent agents all draw down the same allowance. It is short-lived by design (default 24h, configurable via `LITELLM_CLI_JWT_EXPIRATION_HOURS`); re-run `lite login` to refresh it and pick up your latest team and user settings. `lite auth print-token` (usable as Claude Code's `apiKeyHelper`) prints it while fresh and fails once it expires -- there is no silent renewal. It is accepted on a default deployment without `EXPERIMENTAL_UI_LOGIN`, does not appear in the Keys UI, and cannot be rotated or revoked mid-session. A credential from `lite login --pkce` is the exception: it carries a refresh token, so the CLI renews the key shortly before it expires and `lite logout` revokes the refresh token on the proxy (see [Browser sign-in with PKCE](https://docs.litellm.ai/docs/proxy/cli_sso#browser-sign-in-with-pkce)). Only the holder can end a `--pkce` session early, with `lite logout`; an admin has no button for it, but every renewal re-reads the user on the proxy, so deactivating the user or removing them from the team makes the next renewal fail and the key runs out within `LITELLM_CLI_JWT_EXPIRATION_HOURS`. On a proxy with more than one worker or replica, configure Redis (`litellm_settings.cache` with Redis `cache_params`, or `general_settings.coordination_redis`) so a refresh token stays single-use and `lite logout` holds on every worker; without Redis each worker keeps its own record. For a long-lived, rotatable, Keys-UI-visible credential, create a dedicated virtual key in the dashboard and pass it via `--api-key` or `LITELLM_PROXY_API_KEY`.
|
||||
The stored credential is a short-lived, per-session agent token, not a managed virtual key. It is scoped to the user and team you logged in as and inherits their models and budgets; spend is tracked against the shared team and user budgets rather than a separate per-session cap, so multiple logins or several concurrent agents all draw down the same allowance. It is short-lived by design (default 24h, configurable via `LITELLM_CLI_JWT_EXPIRATION_HOURS`); re-run `lite login` to refresh it and pick up your latest team and user settings. `lite auth print-token` (usable as Claude Code's `apiKeyHelper`) prints it, and once it is past 24 hours (or the client's `LITELLM_CLI_JWT_EXPIRATION_HOURS`) asks the proxy first and fails if the proxy rejects it. There is no silent renewal. It is accepted on a default deployment without `EXPERIMENTAL_UI_LOGIN`, does not appear in the Keys UI, and cannot be rotated or revoked mid-session. A credential from `lite login --pkce` is the exception: it carries a refresh token, so the CLI renews the key shortly before it expires and `lite logout` revokes the refresh token on the proxy (see [Browser sign-in with PKCE](https://docs.litellm.ai/docs/proxy/cli_sso#browser-sign-in-with-pkce)). Only the holder can end a `--pkce` session early, with `lite logout`; an admin has no button for it, but every renewal re-reads the user on the proxy, so deactivating the user or removing them from the team makes the next renewal fail and the key runs out within `LITELLM_CLI_JWT_EXPIRATION_HOURS`. On a proxy with more than one worker or replica, configure Redis (`litellm_settings.cache` with Redis `cache_params`, or `general_settings.coordination_redis`) so a refresh token stays single-use and `lite logout` holds on every worker; without Redis each worker keeps its own record. For a long-lived, rotatable, Keys-UI-visible credential, create a dedicated virtual key in the dashboard and pass it via `--api-key` or `LITELLM_PROXY_API_KEY`.
|
||||
|
||||
### Usage
|
||||
|
||||
|
|
|
|||
|
|
@ -506,7 +506,7 @@ To pin the model, pass the agent's own model flag (for example `lite claude --mo
|
|||
|
||||
The token minted by `lite login` is a short-lived, per-session agent credential, not a managed virtual key. It is scoped to the user and team you authenticated as, inherits that user's and team's models and budgets, and is enforced on the proxy exactly like a virtual key on the same team (guardrails, routing, logging, spend). Spend is tracked against the shared team and user budgets, so running several agents (or logging in more than once) does not hand each session its own separate budget; they all draw down the same team/user allowance. There is no separate per-session cap, so sustained agent use is not capped at a small chat-session limit.
|
||||
|
||||
The credential is short-lived by design (default 24h, configurable via `LITELLM_CLI_JWT_EXPIRATION_HOURS`); run `lite login` again to refresh it, which also re-reads your latest team and user settings. It does not appear in the Keys UI and cannot be rotated or revoked mid-session. `lite auth print-token` (usable as Claude Code's `apiKeyHelper`) prints it while it's still fresh and fails once it expires -- there is no silent renewal, so a long-running session needs a fresh `lite login` once a day. `lite claude`, `lite codex`, `lite opencode`, and `lite pi` work with it on a default deployment; `EXPERIMENTAL_UI_LOGIN` is not required. `lite login --pkce` is the exception to the daily re-login: it signs in through your system browser with OAuth authorization code and PKCE and stores a refresh token next to the key, so every `lite` command and `lite auth print-token` renew the key on their own shortly before it expires, `lite whoami` shows when the current key expires, and `lite logout` revokes the refresh token on the proxy (it needs a proxy that serves `/.well-known/litellm-cli-auth`; see [Browser sign-in with PKCE](https://docs.litellm.ai/docs/proxy/cli_sso#browser-sign-in-with-pkce)). When a renewal is refused, for example after a `lite logout` run from another copy of the credential, the command prints why on stderr and, once the key has run out, tells you to run `lite login --pkce` again. Only the holder can end a `--pkce` session early, with `lite logout`; an admin has no button for it, but every renewal re-reads the user on the proxy, so deactivating the user or removing them from the team makes the next renewal fail and the key runs out within `LITELLM_CLI_JWT_EXPIRATION_HOURS`. On a proxy with more than one worker or replica, configure Redis (`litellm_settings.cache` with Redis `cache_params`, or `general_settings.coordination_redis`) so a refresh token stays single-use and `lite logout` holds on every worker; without Redis each worker keeps its own record. If you need a long-lived, rotatable key that shows up in the Keys UI, create a dedicated virtual key in the dashboard and pass it via `--api-key` or `LITELLM_PROXY_API_KEY` instead.
|
||||
The credential is short-lived by design (default 24h, configurable via `LITELLM_CLI_JWT_EXPIRATION_HOURS`); run `lite login` again to refresh it, which also re-reads your latest team and user settings. It does not appear in the Keys UI and cannot be rotated or revoked mid-session. `lite auth print-token` (usable as Claude Code's `apiKeyHelper`) prints it without asking the proxy for its first 24 hours, or `LITELLM_CLI_JWT_EXPIRATION_HOURS` if you set it on the client. After that, `print-token` and `lite whoami` ask the proxy whether it still accepts the token, and `print-token` fails once the proxy rejects it. There is no silent renewal, so a long-running session needs a fresh `lite login` each time the proxy's token lifetime runs out. If the proxy's tokens last less than 24 hours, set the same `LITELLM_CLI_JWT_EXPIRATION_HOURS` on the client so `print-token` stops printing the token once it runs out. `lite up` still starts a new login once the token is older than that same client window. It copies the token into Claude Code's settings as a static key that it never renews, so a fresh login gives Claude Code the proxy's full token lifetime instead of whatever an old token has left. `lite claude`, `lite codex`, `lite opencode`, and `lite pi` work with it on a default deployment; `EXPERIMENTAL_UI_LOGIN` is not required. `lite login --pkce` is the exception to the re-login: it signs in through your system browser with OAuth authorization code and PKCE and stores a refresh token next to the key, so every `lite` command and `lite auth print-token` renew the key on their own shortly before it expires, `lite whoami` shows when the current key expires, and `lite logout` revokes the refresh token on the proxy (it needs a proxy that serves `/.well-known/litellm-cli-auth`; see [Browser sign-in with PKCE](https://docs.litellm.ai/docs/proxy/cli_sso#browser-sign-in-with-pkce)). When a renewal is refused, for example after a `lite logout` run from another copy of the credential, the command prints why on stderr and, once the key has run out, tells you to run `lite login --pkce` again. Only the holder can end a `--pkce` session early, with `lite logout`; an admin has no button for it, but every renewal re-reads the user on the proxy, so deactivating the user or removing them from the team makes the next renewal fail and the key runs out within `LITELLM_CLI_JWT_EXPIRATION_HOURS`. On a proxy with more than one worker or replica, configure Redis (`litellm_settings.cache` with Redis `cache_params`, or `general_settings.coordination_redis`) so a refresh token stays single-use and `lite logout` holds on every worker; without Redis each worker keeps its own record. If you need a long-lived, rotatable key that shows up in the Keys UI, create a dedicated virtual key in the dashboard and pass it via `--api-key` or `LITELLM_PROXY_API_KEY` instead.
|
||||
|
||||
### Route Every Claude Code Session Through the Proxy
|
||||
|
||||
|
|
|
|||
|
|
@ -13,7 +13,6 @@ from rich.console import Console
|
|||
from rich.table import Table
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict, assert_never
|
||||
|
||||
from litellm.constants import CLI_JWT_EXPIRATION_HOURS
|
||||
from litellm.litellm_core_utils.cli_keyring import (
|
||||
DISABLE_KEYRING_ENV_VAR,
|
||||
SYSTEM_KEYRING,
|
||||
|
|
@ -143,6 +142,9 @@ _TeamMapping: Final = TypeVar("_TeamMapping", bound=Mapping[str, object])
|
|||
|
||||
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
|
||||
|
||||
_PROXY_KEY_CHECK_TIMEOUT_SECONDS: Final = 3
|
||||
_PROXY_REJECTED_KEY_STATUSES: Final = frozenset({401, 403})
|
||||
|
||||
KEYRING_INSTALL_HINT: Final = "pip install 'litellm[cli]'"
|
||||
|
||||
KEYRING_ENABLE_HINT: Final = "keyring --enable (or unset PYTHON_KEYRING_BACKEND)"
|
||||
|
|
@ -1033,10 +1035,12 @@ def print_token(ctx: click.Context):
|
|||
|
||||
Designed to be used as Claude Code's `apiKeyHelper`
|
||||
(https://docs.claude.com/en/docs/claude-code/settings): stdout must
|
||||
contain only the token, so all diagnostics go to stderr. The token
|
||||
expires after `LITELLM_CLI_JWT_EXPIRATION_HOURS` (default 24h); a
|
||||
`lite login --pkce` token renews itself here first, and once a token
|
||||
has expired for good, run the same `lite login` command again.
|
||||
contain only the token, so all diagnostics go to stderr. A token younger
|
||||
than `LITELLM_CLI_JWT_EXPIRATION_HOURS` (default 24h) is printed without
|
||||
asking the proxy. An older one is checked against the proxy and refused
|
||||
only if the proxy rejects it. A `lite login --pkce`
|
||||
token renews itself here first, and once a token has expired for good,
|
||||
run the same `lite login` command again.
|
||||
"""
|
||||
vault: Final = context_secret_vault(ctx)
|
||||
token_data: Final = load_token(vault=vault)
|
||||
|
|
@ -1055,7 +1059,7 @@ def print_token(ctx: click.Context):
|
|||
sys.exit(1)
|
||||
|
||||
renews: Final = "refresh_token" in token_data
|
||||
if not is_cli_token_fresh(token_data) and not renews:
|
||||
if not renews and token_expired(token_data):
|
||||
click.echo("Token expired. Run 'lite login' again.", err=True)
|
||||
sys.exit(1)
|
||||
|
||||
|
|
@ -1111,8 +1115,32 @@ def whoami(ctx: click.Context):
|
|||
expires_at: Final = token_data.get("expires_at")
|
||||
if isinstance(expires_at, (int, float)):
|
||||
click.echo(_key_expiry_line(expires_at, renews="refresh_token" in token_data))
|
||||
elif age_hours > CLI_JWT_EXPIRATION_HOURS:
|
||||
click.echo(f"Warning: Token is more than {CLI_JWT_EXPIRATION_HOURS} hours old and may have expired.")
|
||||
elif token_expired(token_data):
|
||||
click.echo("Token expired. Run 'lite login' again.")
|
||||
|
||||
|
||||
def token_expired(token_data: Mapping[str, object]) -> bool:
|
||||
if is_cli_token_fresh(token_data):
|
||||
return False
|
||||
if isinstance(token_data.get("expires_at"), (int, float)):
|
||||
return True
|
||||
return _proxy_rejects_key(token_data)
|
||||
|
||||
|
||||
def _proxy_rejects_key(token_data: Mapping[str, object]) -> bool:
|
||||
key: Final = token_data.get("key")
|
||||
base_url: Final = token_data.get("base_url")
|
||||
if not isinstance(key, str) or not isinstance(base_url, str) or not base_url:
|
||||
return True
|
||||
try:
|
||||
response: Final = requests.get(
|
||||
f"{base_url}/v1/models",
|
||||
headers={"Authorization": f"Bearer {key}"},
|
||||
timeout=_PROXY_KEY_CHECK_TIMEOUT_SECONDS,
|
||||
)
|
||||
except requests.RequestException:
|
||||
return False
|
||||
return response.status_code in _PROXY_REJECTED_KEY_STATUSES
|
||||
|
||||
|
||||
def _key_expiry_line(expires_at: float, renews: bool) -> str:
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import json
|
|||
import logging
|
||||
import math
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from functools import lru_cache
|
||||
from itertools import chain
|
||||
|
|
@ -136,7 +137,7 @@ from litellm.proxy.utils import ( # noqa: F401 # legacy module exports
|
|||
_check_and_merge_model_level_guardrails, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
||||
check_and_merge_model_level_guardrails,
|
||||
)
|
||||
from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect
|
||||
from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect, spawn_logging_task
|
||||
from litellm.router import Router
|
||||
from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
|
||||
from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX
|
||||
|
|
@ -2006,6 +2007,49 @@ def _content_text(fields: Mapping[str, object]) -> str:
|
|||
return str(fields.get("content", ""))
|
||||
|
||||
|
||||
async def _no_stream_failure_to_release() -> None:
|
||||
return None
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _ClientStreamFailureReport:
|
||||
logging_obj: LiteLLMLoggingObj
|
||||
report: Callable[[Exception], Awaitable[None]]
|
||||
held: Exception | None = None
|
||||
|
||||
@classmethod
|
||||
def hold_while_streaming(cls, logging_obj: LiteLLMLoggingObj | None) -> Callable[[], Awaitable[None]]:
|
||||
if logging_obj is None or logging_obj.on_detached_stream_failure is None:
|
||||
return _no_stream_failure_to_release
|
||||
return cls(logging_obj, logging_obj.on_detached_stream_failure)._start_holding()
|
||||
|
||||
def _start_holding(self) -> Callable[[], Awaitable[None]]:
|
||||
self.logging_obj.on_detached_stream_failure = self.hold
|
||||
return self.release
|
||||
|
||||
async def hold(self, failure: Exception) -> None:
|
||||
if self.held is None:
|
||||
self.held = failure
|
||||
|
||||
async def release(self) -> None:
|
||||
self.logging_obj.on_detached_stream_failure = self.report
|
||||
if self.held is None:
|
||||
return
|
||||
await asyncio.shield(
|
||||
spawn_logging_task(
|
||||
asyncio.get_running_loop(), self._report_held(self.held), task_name="stream failure report"
|
||||
)
|
||||
)
|
||||
|
||||
async def _report_held(self, failure: Exception) -> None:
|
||||
try:
|
||||
await self.report(failure)
|
||||
except Exception as hook_failure: # noqa: BLE001 # a failing failure hook must not skip the stream cleanup
|
||||
verbose_proxy_logger.warning(
|
||||
"Stream failure hook raised: %s(%s)", type(hook_failure).__name__, sanitize_for_log(hook_failure)
|
||||
)
|
||||
|
||||
|
||||
class ProxyBaseLLMRequestProcessing:
|
||||
def __init__(self, data: dict, served_model: str | None = None):
|
||||
self.data = data
|
||||
|
|
@ -3498,6 +3542,8 @@ class ProxyBaseLLMRequestProcessing:
|
|||
request_data: Final = self.data
|
||||
|
||||
async def _on_detached_stream_failure(exc: Exception) -> None:
|
||||
if not logging_obj.claim_stream_failure_report():
|
||||
return
|
||||
try:
|
||||
await proxy_logging_obj.post_call_failure_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
|
|
@ -4448,11 +4494,15 @@ class ProxyBaseLLMRequestProcessing:
|
|||
client_disconnected = False
|
||||
delivered_chunk = False
|
||||
recent_tail = SSE_STREAM_START_TAIL # rebind-ok: rolling window over the yielded bytes
|
||||
stream_logging_obj: Final[object] = request_data.get("litellm_logging_obj") # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] # legacy request data is narrowed to Logging below
|
||||
guarded_stream: Final[AsyncGenerator[object, None]] = proxy_logging_obj.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=response,
|
||||
request_data=request_data,
|
||||
)
|
||||
release_stream_failure: Final = _ClientStreamFailureReport.hold_while_streaming(
|
||||
stream_logging_obj if isinstance(stream_logging_obj, LiteLLMLoggingObj) else None
|
||||
)
|
||||
try:
|
||||
str_so_far = ""
|
||||
async for chunk in guarded_stream:
|
||||
|
|
@ -4523,10 +4573,18 @@ class ProxyBaseLLMRequestProcessing:
|
|||
verbose_proxy_logger.exception(
|
||||
"litellm.proxy.proxy_server.async_data_generator(): Exception occured - %s", e
|
||||
)
|
||||
transformed_exception: Final = await proxy_logging_obj.post_call_failure_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
original_exception=e,
|
||||
request_data=request_data,
|
||||
already_reported: Final = (
|
||||
isinstance(stream_logging_obj, LiteLLMLoggingObj)
|
||||
and not stream_logging_obj.claim_stream_failure_report()
|
||||
)
|
||||
transformed_exception: Final = (
|
||||
None
|
||||
if already_reported
|
||||
else await proxy_logging_obj.post_call_failure_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
original_exception=e,
|
||||
request_data=request_data,
|
||||
)
|
||||
)
|
||||
if transformed_exception is not None:
|
||||
e = transformed_exception
|
||||
|
|
@ -4549,15 +4607,18 @@ class ProxyBaseLLMRequestProcessing:
|
|||
seal: Final = "" if seal_open_frame is None else seal_open_frame(recent_tail)
|
||||
yield seal + error_frame if seal else error_frame
|
||||
finally:
|
||||
await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup(
|
||||
request=request,
|
||||
request_data=request_data,
|
||||
response=response,
|
||||
stream_completed=stream_completed,
|
||||
client_disconnected=client_disconnected,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
try:
|
||||
await release_stream_failure()
|
||||
finally:
|
||||
await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup(
|
||||
request=request,
|
||||
request_data=request_data,
|
||||
response=response,
|
||||
stream_completed=stream_completed,
|
||||
client_disconnected=client_disconnected,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def async_sse_data_generator(
|
||||
|
|
|
|||
|
|
@ -6,9 +6,12 @@
|
|||
######################################################################
|
||||
|
||||
import asyncio
|
||||
import functools
|
||||
import itertools
|
||||
import traceback
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Any, BinaryIO, Final, TypedDict, cast, get_args
|
||||
from collections.abc import Iterable, Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, BinaryIO, Final, NoReturn, TypedDict, cast, get_args
|
||||
|
||||
import httpx
|
||||
from fastapi import (
|
||||
|
|
@ -28,12 +31,13 @@ from typing_extensions import ReadOnly
|
|||
import litellm
|
||||
from litellm import CreateFileRequest, get_secret_str
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import ANTHROPIC_DEFAULT_FILE_LIST_LIMIT
|
||||
from litellm.litellm_core_utils.cloud_storage_security import (
|
||||
is_managed_cloud_storage_uri,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket
|
||||
from litellm.llms.base_llm.files.litellm_db_storage_backend import LITELLM_DB_STORAGE_BACKEND_NAME
|
||||
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
|
||||
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints, FileListModelScope
|
||||
from litellm.llms.base_llm.managed_resources.isolation import build_list_page
|
||||
from litellm.proxy._types import *
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
|
@ -112,6 +116,9 @@ from litellm.types.llms.openai import (
|
|||
OpenAIFileObject,
|
||||
OpenAIFilesPurpose,
|
||||
)
|
||||
from litellm.types.proxy.auth.user_api_key_auth import user_api_key_has_admin_view
|
||||
from litellm.types.router import Deployment
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
router: Final = APIRouter()
|
||||
|
||||
|
|
@ -1673,6 +1680,185 @@ def _as_file_list_page(response: object) -> object:
|
|||
return FileListPage.model_validate(page)
|
||||
|
||||
|
||||
async def _list_owned_managed_files(
|
||||
managed_files_obj: BaseFileEndpoints,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
purpose: str | None,
|
||||
limit: int | None,
|
||||
after: str | None,
|
||||
before: str | None,
|
||||
model_scope: FileListModelScope | None,
|
||||
) -> FileListPage:
|
||||
return await managed_files_obj.afile_list(
|
||||
purpose=purpose,
|
||||
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
limit=limit,
|
||||
after=after,
|
||||
before=before,
|
||||
model_scope=model_scope,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _UnscopedFileList:
|
||||
status_code: int
|
||||
detail: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ListedFiles:
|
||||
model_scope: FileListModelScope | None
|
||||
custom_llm_provider: str | None
|
||||
|
||||
|
||||
def _raise_unscoped_file_list(failure: _UnscopedFileList) -> NoReturn:
|
||||
raise HTTPException(status_code=failure.status_code, detail=failure.detail)
|
||||
|
||||
|
||||
def _single_target_model_name(target_model_names: str) -> str | _UnscopedFileList:
|
||||
target_model_names_list: Final = target_model_names.split(",")
|
||||
if len(target_model_names_list) != 1:
|
||||
return _UnscopedFileList(
|
||||
status_code=400,
|
||||
detail="target_model_names on list files must be a list of one model name. Example: ['gpt-4o']",
|
||||
)
|
||||
return target_model_names_list[0]
|
||||
|
||||
|
||||
async def _requested_custom_llm_provider(request: Request, provider: str | None) -> str | None:
|
||||
return (
|
||||
provider
|
||||
or get_custom_llm_provider_from_request_headers(request=request)
|
||||
or get_custom_llm_provider_from_request_query(request=request)
|
||||
or await get_custom_llm_provider_from_request_body(request=request)
|
||||
)
|
||||
|
||||
|
||||
def _default_file_list_limit(custom_llm_provider: str | None) -> int | None:
|
||||
return ANTHROPIC_DEFAULT_FILE_LIST_LIMIT if custom_llm_provider == LlmProviders.ANTHROPIC.value else None
|
||||
|
||||
|
||||
def _deployment_model_mapping_keys(deployment: Deployment) -> frozenset[str]:
|
||||
return frozenset(key for key in (deployment.model_name, deployment.model_info.id) if key is not None)
|
||||
|
||||
|
||||
def _model_mapping_keys(deployments: Iterable[Deployment]) -> frozenset[str]:
|
||||
return frozenset(itertools.chain.from_iterable(map(_deployment_model_mapping_keys, deployments)))
|
||||
|
||||
|
||||
def _deployments_serving_model(llm_router: Router, model: str) -> tuple[Deployment, ...]:
|
||||
deployment: Final = llm_router.get_deployment(model_id=model) if llm_router.has_model_id(model) else None
|
||||
if deployment is not None:
|
||||
return (deployment,)
|
||||
return tuple(Deployment.model_validate(listed) for listed in llm_router.get_model_list(model_name=model) or ())
|
||||
|
||||
|
||||
def _resolved_deployment_provider(deployment: Deployment) -> str | None:
|
||||
try:
|
||||
return litellm.get_llm_provider(
|
||||
model=deployment.litellm_params.model,
|
||||
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
|
||||
)[1]
|
||||
except litellm.BadRequestError:
|
||||
return None
|
||||
|
||||
|
||||
def _deployments_on_provider(llm_router: Router | None, custom_llm_provider: str) -> Iterator[Deployment]:
|
||||
if llm_router is None:
|
||||
return iter(())
|
||||
return (
|
||||
deployment
|
||||
for deployment_id in llm_router.get_model_ids()
|
||||
if (deployment := llm_router.get_deployment(model_id=deployment_id)) is not None
|
||||
and _resolved_deployment_provider(deployment) == custom_llm_provider
|
||||
)
|
||||
|
||||
|
||||
def _routed_model_name_on_provider(
|
||||
llm_router: Router, custom_llm_provider: str, provider_model_names: frozenset[str], model_name: str
|
||||
) -> bool:
|
||||
if model_name in llm_router.model_names:
|
||||
return model_name in provider_model_names
|
||||
return llm_router.pattern_router.routes_to_provider(model_name, custom_llm_provider)
|
||||
|
||||
|
||||
def _stored_model_name_routes_to_provider(
|
||||
llm_router: Router, custom_llm_provider: str, provider_model_names: frozenset[str], model_name: str
|
||||
) -> bool:
|
||||
routing_group: Final = llm_router.get_routing_group(model_name)
|
||||
routed_model_names: Final = (
|
||||
tuple(routing_group.models) if routing_group is not None else (llm_router.routable_model_group(model_name),)
|
||||
)
|
||||
return any(
|
||||
_routed_model_name_on_provider(llm_router, custom_llm_provider, provider_model_names, routed_model_name)
|
||||
for routed_model_name in routed_model_names
|
||||
)
|
||||
|
||||
|
||||
def _files_on_provider(llm_router: Router | None, custom_llm_provider: str) -> _ListedFiles | _UnscopedFileList:
|
||||
from litellm.proxy.pass_through_endpoints.managed_id_rewriter import passthrough_file_model_id
|
||||
|
||||
if custom_llm_provider not in litellm.provider_list:
|
||||
return _UnscopedFileList(status_code=400, detail=f"Unknown provider '{custom_llm_provider}' for listing files.")
|
||||
passthrough_model_id: Final = passthrough_file_model_id(custom_llm_provider)
|
||||
deployments: Final = tuple(_deployments_on_provider(llm_router, custom_llm_provider))
|
||||
model_mapping_keys: Final = _model_mapping_keys(deployments) | frozenset(
|
||||
(passthrough_model_id,) if passthrough_model_id is not None else ()
|
||||
)
|
||||
if llm_router is None:
|
||||
return _ListedFiles(FileListModelScope(model_mapping_keys), custom_llm_provider)
|
||||
return _ListedFiles(
|
||||
FileListModelScope(
|
||||
model_mapping_keys,
|
||||
routes_stored_model_name=functools.partial(
|
||||
_stored_model_name_routes_to_provider,
|
||||
llm_router,
|
||||
custom_llm_provider,
|
||||
frozenset(deployment.model_name for deployment in deployments),
|
||||
),
|
||||
),
|
||||
custom_llm_provider,
|
||||
)
|
||||
|
||||
|
||||
def _files_serving_model(llm_router: Router, model: str) -> _ListedFiles:
|
||||
deployments: Final = _deployments_serving_model(llm_router, model)
|
||||
providers: Final = frozenset(map(_resolved_deployment_provider, deployments))
|
||||
return _ListedFiles(
|
||||
FileListModelScope(_model_mapping_keys(deployments)), next(iter(providers)) if len(providers) == 1 else None
|
||||
)
|
||||
|
||||
|
||||
async def _listed_files(
|
||||
llm_router: Router | None,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
custom_llm_provider: str | None,
|
||||
target_model_names: str | None,
|
||||
routed_model: str | None,
|
||||
) -> _ListedFiles | _UnscopedFileList:
|
||||
if routed_model is not None and llm_router is not None:
|
||||
return _files_serving_model(llm_router, routed_model)
|
||||
if target_model_names:
|
||||
target_model_name: Final = _single_target_model_name(target_model_names)
|
||||
if isinstance(target_model_name, _UnscopedFileList):
|
||||
return target_model_name
|
||||
if llm_router is None:
|
||||
return _UnscopedFileList(
|
||||
status_code=500, detail="LLM Router not initialized. Ensure models added to proxy."
|
||||
)
|
||||
await get_authorized_credentials_for_model(
|
||||
llm_router=llm_router,
|
||||
model_id=target_model_name,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
operation_context="file list",
|
||||
)
|
||||
return _files_serving_model(llm_router, target_model_name)
|
||||
if custom_llm_provider is None:
|
||||
return _ListedFiles(None, None)
|
||||
return _files_on_provider(llm_router, custom_llm_provider)
|
||||
|
||||
|
||||
@router.get(
|
||||
"/{provider}/v1/files",
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
|
|
@ -1697,6 +1883,8 @@ async def list_files(
|
|||
purpose: str | None = None,
|
||||
limit: int | None = None,
|
||||
after: str | None = None,
|
||||
after_id: str | None = None,
|
||||
before_id: str | None = None,
|
||||
):
|
||||
"""
|
||||
Returns information about a specific file. that can be used across - Assistants API, Batch API
|
||||
|
|
@ -1739,6 +1927,10 @@ async def list_files(
|
|||
)
|
||||
|
||||
response: Any | None = None
|
||||
owned_files_hook: Final = proxy_logging_obj.get_proxy_hook("managed_files")
|
||||
list_owned_files_only: Final = litellm.require_managed_files is True and not user_api_key_has_admin_view(
|
||||
user_api_key_dict
|
||||
)
|
||||
|
||||
# Check for model-based credential routing (no file_id encoding check for list)
|
||||
should_route, model_used, _, credentials = await handle_model_based_routing(
|
||||
|
|
@ -1750,7 +1942,36 @@ async def list_files(
|
|||
check_file_id_encoding=False,
|
||||
)
|
||||
|
||||
if should_route and credentials is not None:
|
||||
if list_owned_files_only:
|
||||
if after and after_id and after != after_id:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Pass the list cursor as either 'after' or 'after_id', not both with different ids.",
|
||||
)
|
||||
if not isinstance(owned_files_hook, BaseFileEndpoints):
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail="Managed file listing is unavailable, so require_managed_files cannot scope the list to the caller.",
|
||||
)
|
||||
listed_files: Final = await _listed_files(
|
||||
llm_router=llm_router,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
custom_llm_provider=await _requested_custom_llm_provider(request=request, provider=provider),
|
||||
target_model_names=target_model_names,
|
||||
routed_model=model_used if should_route else None,
|
||||
)
|
||||
if isinstance(listed_files, _UnscopedFileList):
|
||||
_raise_unscoped_file_list(listed_files)
|
||||
response = await _list_owned_managed_files( # rebind-ok: each dispatch branch binds response
|
||||
managed_files_obj=owned_files_hook,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
purpose=purpose,
|
||||
limit=limit if limit is not None else _default_file_list_limit(listed_files.custom_llm_provider),
|
||||
after=after or after_id,
|
||||
before=before_id,
|
||||
model_scope=listed_files.model_scope,
|
||||
)
|
||||
elif should_route and credentials is not None:
|
||||
# Use model-based routing with credentials from config
|
||||
prepare_data_with_credentials(data=data, credentials=credentials, include_internal_credentials=True)
|
||||
response = await litellm.afile_list(
|
||||
|
|
@ -1762,12 +1983,9 @@ async def list_files(
|
|||
verbose_proxy_logger.debug("Listed files using model: %s", model_used)
|
||||
|
||||
elif target_model_names and isinstance(target_model_names, str):
|
||||
target_model_names_list: Final = target_model_names.split(",")
|
||||
if len(target_model_names_list) != 1:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="target_model_names on list files must be a list of one model name. Example: ['gpt-4o']",
|
||||
)
|
||||
target_model_name: Final = _single_target_model_name(target_model_names)
|
||||
if isinstance(target_model_name, _UnscopedFileList):
|
||||
_raise_unscoped_file_list(target_model_name)
|
||||
if llm_router is None:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
|
|
@ -1775,7 +1993,7 @@ async def list_files(
|
|||
)
|
||||
credentials = await get_authorized_credentials_for_model(
|
||||
llm_router=llm_router,
|
||||
model_id=target_model_names_list[0],
|
||||
model_id=target_model_name,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
operation_context="file list",
|
||||
)
|
||||
|
|
@ -1786,20 +2004,17 @@ async def list_files(
|
|||
**data,
|
||||
)
|
||||
else:
|
||||
custom_llm_provider: Final = (
|
||||
provider
|
||||
or get_custom_llm_provider_from_request_headers(request=request)
|
||||
or get_custom_llm_provider_from_request_query(request=request)
|
||||
or await get_custom_llm_provider_from_request_body(request=request)
|
||||
)
|
||||
custom_llm_provider: Final = await _requested_custom_llm_provider(request=request, provider=provider)
|
||||
managed_files_obj: Final = proxy_logging_obj.get_proxy_hook("managed_files")
|
||||
if custom_llm_provider is None and isinstance(managed_files_obj, BaseFileEndpoints):
|
||||
response = await managed_files_obj.afile_list(
|
||||
purpose=purpose,
|
||||
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
|
||||
response = await _list_owned_managed_files( # rebind-ok: each dispatch branch binds response
|
||||
managed_files_obj=managed_files_obj,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
purpose=purpose,
|
||||
limit=limit,
|
||||
after=after,
|
||||
before=None,
|
||||
model_scope=None,
|
||||
)
|
||||
else:
|
||||
resolved_custom_llm_provider: Final = custom_llm_provider or "openai"
|
||||
|
|
|
|||
|
|
@ -242,6 +242,12 @@ def _passthrough_sentinel_model_id(provider: str) -> str:
|
|||
return f"_passthrough_{provider}"
|
||||
|
||||
|
||||
def passthrough_file_model_id(provider: str) -> str | None:
|
||||
if (provider, "/v1/files") not in _LIST_ROUTE_TABLE:
|
||||
return None
|
||||
return _passthrough_sentinel_model_id(provider)
|
||||
|
||||
|
||||
# Key under which the provider marker is stored in a file row's model_mappings.
|
||||
# Its value lands in flat_model_file_ids (built from model_mappings.values()),
|
||||
# giving the file table a DB-queryable provider scope it otherwise lacks.
|
||||
|
|
|
|||
|
|
@ -1163,6 +1163,7 @@ model LiteLLM_ManagedFileTable {
|
|||
@@index([unified_file_id])
|
||||
@@index([flat_model_file_ids], type: Gin)
|
||||
@@index([team_id, created_at(sort: Desc)])
|
||||
@@index([created_by, created_at(sort: Desc)])
|
||||
}
|
||||
|
||||
model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use the
|
||||
|
|
|
|||
|
|
@ -41,7 +41,11 @@ from litellm.exceptions import MidStreamFallbackError, RateLimitError
|
|||
from litellm.litellm_core_utils.asyncify import run_async_function
|
||||
from litellm.litellm_core_utils.core_helpers import process_response_headers
|
||||
from litellm.litellm_core_utils.exception_mapping_utils import with_deployment_api_base
|
||||
from litellm.litellm_core_utils.internal_call_metadata import without_budget_reservation
|
||||
from litellm.litellm_core_utils.internal_call_metadata import (
|
||||
is_background_response,
|
||||
requests_background,
|
||||
without_budget_reservation,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base
|
||||
|
|
@ -65,7 +69,7 @@ from litellm.types.llms.openai import (
|
|||
ResponsesAPIStreamingResponse,
|
||||
)
|
||||
from litellm.types.responses.main import GenericResponseOutputItem, OutputText
|
||||
from litellm.types.utils import CallTypes
|
||||
from litellm.types.utils import CallTypes, ModelResponse
|
||||
from litellm.utils import async_post_call_success_deployment_hook
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -288,7 +292,7 @@ def stream_error_status_and_message(error_obj: object) -> tuple[int, str]:
|
|||
return _status_code_for_error_fields(error_type, error_code), message
|
||||
|
||||
|
||||
def _map_stream_error_to_exception(
|
||||
def map_stream_error_to_exception(
|
||||
error_obj: object,
|
||||
model: str,
|
||||
custom_llm_provider: str,
|
||||
|
|
@ -579,6 +583,7 @@ class BaseResponsesAPIStreamingIterator:
|
|||
self.request_data.get("input"),
|
||||
self.request_data,
|
||||
self._generated_content + self._generated_tool_arguments,
|
||||
use_default_image_token_count=False,
|
||||
)
|
||||
if _estimate_wanted
|
||||
else None
|
||||
|
|
@ -732,7 +737,7 @@ class BaseResponsesAPIStreamingIterator:
|
|||
self._earlier_rounds_spend = earlier
|
||||
|
||||
def _map_error_event_exception(self, error_obj: object) -> Exception:
|
||||
return _map_stream_error_to_exception(
|
||||
return map_stream_error_to_exception(
|
||||
error_obj,
|
||||
self.model or "",
|
||||
self.custom_llm_provider or "",
|
||||
|
|
@ -982,6 +987,40 @@ class BaseResponsesAPIStreamingIterator:
|
|||
except Exception:
|
||||
pass
|
||||
|
||||
def _record_partial_usage_for_failure(self) -> None:
|
||||
generated_text: Final = self._generated_content + self._generated_tool_arguments
|
||||
if (
|
||||
not generated_text
|
||||
or self._replays_stored_stream
|
||||
or requests_background(self.request_data)
|
||||
or is_background_response(self._created_response)
|
||||
or self._completed_response_logged
|
||||
or self.logging_obj.model_call_details.get("combined_usage_object") is not None
|
||||
):
|
||||
return
|
||||
estimate: Final = _estimate_usage_safely(
|
||||
self.model or "",
|
||||
self.request_data.get("input"),
|
||||
self.request_data,
|
||||
generated_text,
|
||||
use_default_image_token_count=True,
|
||||
)
|
||||
if estimate is None:
|
||||
return
|
||||
usage: Final = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( # pyright: ignore[reportPrivateUsage] # same shared transform the success path above uses
|
||||
estimate
|
||||
)
|
||||
served_tier: Final[object] = getattr(self._created_response, "service_tier", None)
|
||||
assembled: Final = ModelResponse(
|
||||
model=self.model, usage=usage, service_tier=served_tier if isinstance(served_tier, str) else None
|
||||
)
|
||||
if self._earlier_rounds_spend is None:
|
||||
self.logging_obj.record_assembled_response_for_failure(assembled)
|
||||
return
|
||||
self.logging_obj.record_failure_after_earlier_rounds(
|
||||
self._earlier_rounds_spend, usage, self.logging_obj.response_cost_calculator(result=assembled) or 0.0
|
||||
)
|
||||
|
||||
def _handle_failure(self, exception: Exception):
|
||||
"""
|
||||
Trigger failure handlers before bubbling the exception.
|
||||
|
|
@ -991,6 +1030,7 @@ class BaseResponsesAPIStreamingIterator:
|
|||
if self._failure_handled:
|
||||
return
|
||||
self._failure_handled = True
|
||||
self._record_partial_usage_for_failure()
|
||||
|
||||
traceback_exception: Final = traceback.format_exc()
|
||||
end_time: Final = datetime.now()
|
||||
|
|
@ -1061,15 +1101,20 @@ class BaseResponsesAPIStreamingIterator:
|
|||
the same shape the success handler assembles for a natural end of stream.
|
||||
Its response is the response.created payload (it carries the served
|
||||
service_tier) plus the text and tool arguments relayed so far, with counted
|
||||
usage since the provider never sent any. None when nothing was relayed yet
|
||||
or the stream already reached its terminal event and billed itself.
|
||||
usage since the provider never sent any. None when nothing was relayed yet,
|
||||
the stream already reached its terminal event and billed itself, or its failure
|
||||
log already billed the relayed output.
|
||||
"""
|
||||
created: Final = self._created_response
|
||||
generated: Final = self._generated_content + self._generated_tool_arguments
|
||||
if self._replays_stored_stream or self.completed_response is not None or created is None or not generated:
|
||||
if self._replays_stored_stream or created is None or not generated or self._billed_without_disconnect():
|
||||
return None
|
||||
usage: Final = _estimate_usage_safely(
|
||||
self.model or "", self.request_data.get("input"), self.request_data, generated
|
||||
self.model or "",
|
||||
self.request_data.get("input"),
|
||||
self.request_data,
|
||||
generated,
|
||||
use_default_image_token_count=False,
|
||||
)
|
||||
if usage is None:
|
||||
return None
|
||||
|
|
@ -1083,6 +1128,9 @@ class BaseResponsesAPIStreamingIterator:
|
|||
stamp_responses_usage_cost(partial, self.logging_obj)
|
||||
return ResponseIncompleteEvent(type=ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, response=partial)
|
||||
|
||||
def _billed_without_disconnect(self) -> bool:
|
||||
return self.completed_response is not None or self._failure_handled
|
||||
|
||||
def _note_yielded_event(self, event: ResponsesAPIStreamingResponse) -> None:
|
||||
self._yielded_first_chunk = True
|
||||
if event.type not in PRE_OUTPUT_LIFECYCLE_EVENT_TYPES:
|
||||
|
|
@ -1226,6 +1274,9 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
raise StopAsyncIteration from e
|
||||
self._raise_for_transport_error(e)
|
||||
except httpx.TimeoutException as e:
|
||||
if self.completed_response is not None:
|
||||
self.finished = True
|
||||
raise StopAsyncIteration from e
|
||||
self._raise_for_transport_timeout(e)
|
||||
except httpx.HTTPError as e:
|
||||
# Handle HTTP errors
|
||||
|
|
@ -1316,6 +1367,9 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
raise StopIteration from e
|
||||
self._raise_for_transport_error(e)
|
||||
except httpx.TimeoutException as e:
|
||||
if self.completed_response is not None:
|
||||
self.finished = True
|
||||
raise StopIteration from e
|
||||
self._raise_for_transport_timeout(e)
|
||||
except httpx.HTTPError as e:
|
||||
# Handle HTTP errors
|
||||
|
|
@ -1740,6 +1794,8 @@ def _estimate_usage_from_text(
|
|||
request_input: object,
|
||||
responses_api_request: Mapping[str, object],
|
||||
generated_text: str,
|
||||
*,
|
||||
use_default_image_token_count: bool = False,
|
||||
) -> ResponseAPIUsage:
|
||||
messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # the transformer's signature is partially untyped
|
||||
input=request_input, # pyright: ignore[reportArgumentType] # the raw Responses API input is a str or ResponseInputParam list, matching the helper's declared union
|
||||
|
|
@ -1747,7 +1803,7 @@ def _estimate_usage_from_text(
|
|||
only_since_compaction=True,
|
||||
)
|
||||
input_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped
|
||||
model=model, messages=messages
|
||||
model=model, messages=messages, use_default_image_token_count=use_default_image_token_count
|
||||
)
|
||||
output_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped
|
||||
model=model, text=generated_text, count_response_tokens=True
|
||||
|
|
@ -1764,6 +1820,8 @@ def _estimate_usage_safely(
|
|||
request_input: object,
|
||||
responses_api_request: Mapping[str, object],
|
||||
generated_text: str,
|
||||
*,
|
||||
use_default_image_token_count: bool,
|
||||
) -> ResponseAPIUsage | None:
|
||||
try:
|
||||
return _estimate_usage_from_text(
|
||||
|
|
@ -1771,6 +1829,7 @@ def _estimate_usage_safely(
|
|||
request_input=request_input,
|
||||
responses_api_request=responses_api_request,
|
||||
generated_text=generated_text,
|
||||
use_default_image_token_count=use_default_image_token_count,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.debug("Could not estimate usage from stream text, billing $0: %s", e)
|
||||
|
|
@ -2289,7 +2348,7 @@ class _ResponsesWsTurn:
|
|||
)
|
||||
if failed_event is None:
|
||||
return None
|
||||
return _map_stream_error_to_exception(
|
||||
return map_stream_error_to_exception(
|
||||
_ws_event_error(failed_event), model, custom_llm_provider, logging_obj=connection_logging_obj
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -6458,6 +6458,9 @@ class Router:
|
|||
hop_kwargs: Final = mid_stream_fallback_hop_kwargs(
|
||||
model=model, original_generic_function=original_generic_function, controls=controls, kwargs=kwargs
|
||||
)
|
||||
shared_logging: Final = kwargs.get("litellm_logging_obj")
|
||||
if isinstance(shared_logging, LiteLLMLogging):
|
||||
shared_logging.carry_failed_attempt_usage()
|
||||
response: Final = await self._ageneric_api_call_with_fallbacks_helper(
|
||||
model=model, original_generic_function=original_generic_function, **kwargs
|
||||
)
|
||||
|
|
|
|||
|
|
@ -46,6 +46,12 @@
|
|||
- A failure after execution starts is terminal and never triggers fallback
|
||||
- Creating an unstarted native coroutine acquires no clients or credentials and captures no execution context
|
||||
|
||||
# Native extension
|
||||
|
||||
- `_native` ships in every wheel and the sdist builds it, so it is always importable
|
||||
- Import it with `from litellm.rust_bridge import _native` and use the attributes `_native.pyi` declares
|
||||
- Never add a `None` branch, `getattr` default or stand-in class for a missing extension
|
||||
|
||||
# Verification
|
||||
|
||||
- Check cross-language behavior against a freshly installed extension, covering unstarted-call release, setup ordering, exception identity, sync and async streams, cancellation and GC
|
||||
|
|
|
|||
|
|
@ -1,10 +1,5 @@
|
|||
"""LiteLLM Rust bridge package."""
|
||||
|
||||
from litellm.rust_bridge.configuration import rust
|
||||
from litellm.rust_bridge.loader import (
|
||||
get_native_bridge,
|
||||
native_bridge_available,
|
||||
reset_native_bridge_cache,
|
||||
)
|
||||
|
||||
__all__ = ["get_native_bridge", "native_bridge_available", "reset_native_bridge_cache", "rust"]
|
||||
__all__ = ["rust"]
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ from __future__ import annotations
|
|||
from collections.abc import Callable
|
||||
from typing import Final, Generic, TypeVar
|
||||
|
||||
from litellm.rust_bridge.loader import get_native_bridge
|
||||
from litellm.rust_bridge import _native
|
||||
|
||||
BindingT = TypeVar("BindingT")
|
||||
|
||||
|
|
@ -18,32 +18,24 @@ _UNSET: Final = _Unset()
|
|||
class NativeBinding(Generic[BindingT]):
|
||||
"""Resolve one native attribute with an explicit, resettable test override."""
|
||||
|
||||
def __init__(self, attribute: str, *, validate: Callable[[object], BindingT | None]) -> None:
|
||||
def __init__(
|
||||
self, attribute: str, *, validate: Callable[[object], BindingT | None], source: object = _native
|
||||
) -> None:
|
||||
self._attribute: Final = attribute
|
||||
self._validate: Final = validate
|
||||
self._source: Final = source
|
||||
self._override: BindingT | None | _Unset = _UNSET
|
||||
|
||||
def load(self) -> BindingT | None:
|
||||
if not isinstance(self._override, _Unset):
|
||||
return self._override
|
||||
native: Final = get_native_bridge()
|
||||
if native is None:
|
||||
return None
|
||||
return self._validate(getattr(native, self._attribute, None))
|
||||
return self._validate(getattr(self._source, self._attribute, None))
|
||||
|
||||
def with_source(self, source: object) -> NativeBinding[BindingT]:
|
||||
return NativeBinding(self._attribute, validate=self._validate, source=source)
|
||||
|
||||
def override(self, value: BindingT | None) -> None:
|
||||
self._override = value
|
||||
|
||||
def reset(self) -> None:
|
||||
self._override = _UNSET
|
||||
|
||||
|
||||
def native_exception_types() -> tuple[type[BaseException], type[BaseException]] | None:
|
||||
native: Final = get_native_bridge()
|
||||
if native is None:
|
||||
return None
|
||||
declined: Final = getattr(native, "RustBridgeDeclined", None)
|
||||
upstream: Final = getattr(native, "RustUpstreamError", None)
|
||||
if not isinstance(declined, type) or not isinstance(upstream, type):
|
||||
return None
|
||||
return declined, upstream
|
||||
|
|
|
|||
|
|
@ -1,12 +1,10 @@
|
|||
import os
|
||||
from collections.abc import Awaitable, Mapping, Sequence
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Final, Protocol, runtime_checkable
|
||||
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
from typing import Final
|
||||
|
||||
from litellm.constants import DEFAULT_AGENT_TRACING_RETENTION_DAYS, DEFAULT_CLICKHOUSE_DATABASE
|
||||
from litellm.rust_bridge.loader import get_native_bridge
|
||||
from litellm.rust_bridge import _native
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True, repr=False)
|
||||
|
|
@ -16,27 +14,6 @@ class SpendStorageConfig:
|
|||
retention_days: int = DEFAULT_AGENT_TRACING_RETENTION_DAYS
|
||||
|
||||
|
||||
class NativeSpendConfig(Protocol):
|
||||
def __init__(self, database: str, url: str, retention_days: int) -> None: ...
|
||||
|
||||
|
||||
class NativeSpendStorage(Protocol):
|
||||
def __init__(self, config: NativeSpendConfig) -> None: ...
|
||||
def ensure_schema(self) -> Awaitable[None]: ...
|
||||
def insert_rows(self, rows: Sequence[Mapping[str, object]]) -> Awaitable[None]: ...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class NativeSpendBridge(Protocol):
|
||||
NativeClickHouseSpendConfig: type[NativeSpendConfig]
|
||||
NativeClickHouseSpendStorage: type[NativeSpendStorage]
|
||||
|
||||
|
||||
_NATIVE: Final[TypeAdapter[NativeSpendBridge]] = TypeAdapter(
|
||||
NativeSpendBridge, config=ConfigDict(arbitrary_types_allowed=True)
|
||||
)
|
||||
|
||||
|
||||
def spend_storage_config(environ: Mapping[str, str] = os.environ) -> SpendStorageConfig:
|
||||
url: Final = environ.get("CLICKHOUSE_URL")
|
||||
if not url:
|
||||
|
|
@ -53,12 +30,8 @@ def spend_storage_config(environ: Mapping[str, str] = os.environ) -> SpendStorag
|
|||
|
||||
class ClickHouseSpendStorage:
|
||||
def __init__(self, config: SpendStorageConfig) -> None:
|
||||
bridge: Final = get_native_bridge()
|
||||
if bridge is None:
|
||||
raise RuntimeError("ClickHouse spend logging requires the Rust extension")
|
||||
native: Final = _NATIVE.validate_python(bridge)
|
||||
self._native: Final = native.NativeClickHouseSpendStorage(
|
||||
native.NativeClickHouseSpendConfig(config.database, config.url, config.retention_days)
|
||||
self._native: Final = _native.NativeClickHouseSpendStorage(
|
||||
_native.NativeClickHouseSpendConfig(config.database, config.url, config.retention_days)
|
||||
)
|
||||
|
||||
async def ensure_schema(self) -> None:
|
||||
|
|
|
|||
|
|
@ -11,32 +11,22 @@ site, so the runtime can never start there. Workers forked from it are unaffecte
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import Final
|
||||
|
||||
from litellm.rust_bridge.loader import get_native_bridge
|
||||
from litellm.rust_bridge import _native
|
||||
|
||||
|
||||
class NativeStateStartedBeforeFork(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class _NeverRaised(RuntimeError):
|
||||
"""Stands in for a native exception when the extension is unavailable or predates it."""
|
||||
ForkedAfterNativeRuntimeStarted: Final = _native.ForkedAfterNativeRuntimeStarted
|
||||
ProcessReservedForForking: Final = _native.ProcessReservedForForking
|
||||
|
||||
|
||||
_native: Final = get_native_bridge()
|
||||
ForkedAfterNativeRuntimeStarted: Final[type[RuntimeError]] = getattr(
|
||||
_native, "ForkedAfterNativeRuntimeStarted", _NeverRaised
|
||||
)
|
||||
ProcessReservedForForking: Final[type[RuntimeError]] = getattr(_native, "ProcessReservedForForking", _NeverRaised)
|
||||
|
||||
|
||||
def reserve_process_for_forking(where: str) -> None:
|
||||
def reserve_process_for_forking(where: str, reserve: Callable[[], None] = _native.reserve_process_for_forking) -> None:
|
||||
"""Forbid native routes in this process. Raises if one already ran here."""
|
||||
native: Final = get_native_bridge()
|
||||
reserve: Final = getattr(native, "reserve_process_for_forking", None)
|
||||
if not callable(reserve):
|
||||
return
|
||||
try:
|
||||
reserve()
|
||||
except RuntimeError as error:
|
||||
|
|
|
|||
|
|
@ -1,35 +0,0 @@
|
|||
"""Loader for the packaged LiteLLM Rust extension."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import ModuleType
|
||||
from typing import Final
|
||||
|
||||
_BRIDGE_SENTINEL: Final = object()
|
||||
_cached_bridge: ModuleType | None | object = _BRIDGE_SENTINEL
|
||||
|
||||
|
||||
def get_native_bridge() -> ModuleType | None:
|
||||
"""Return the packaged Rust extension, or ``None`` when unavailable."""
|
||||
global _cached_bridge
|
||||
if _cached_bridge is not _BRIDGE_SENTINEL:
|
||||
return _cached_bridge if isinstance(_cached_bridge, ModuleType) else None
|
||||
|
||||
try:
|
||||
from litellm.rust_bridge import _native
|
||||
except ImportError:
|
||||
_cached_bridge = None
|
||||
return None
|
||||
_cached_bridge = _native
|
||||
return _native
|
||||
|
||||
|
||||
def reset_native_bridge_cache() -> None:
|
||||
"""Forget the cached extension so the next lookup reimports it from disk."""
|
||||
global _cached_bridge
|
||||
_cached_bridge = _BRIDGE_SENTINEL
|
||||
|
||||
|
||||
def native_bridge_available() -> bool:
|
||||
"""Whether the packaged Rust extension is importable."""
|
||||
return get_native_bridge() is not None
|
||||
|
|
@ -2,32 +2,33 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable
|
||||
from dataclasses import dataclass
|
||||
from typing import Final, Protocol
|
||||
|
||||
import httpx
|
||||
from websockets.exceptions import ConnectionClosedOK
|
||||
|
||||
from litellm.rust_bridge import _native
|
||||
from litellm.rust_bridge.host.timeouts import timeout_to_seconds
|
||||
from litellm.rust_bridge.loader import get_native_bridge
|
||||
|
||||
|
||||
class RustResponsesWebSocket(Protocol):
|
||||
async def send_text(self, text: str) -> None: ...
|
||||
def send_text(self, text: str) -> Awaitable[None]: ...
|
||||
|
||||
async def recv_text(self) -> str | None: ...
|
||||
def recv_text(self) -> Awaitable[str | None]: ...
|
||||
|
||||
async def close(self) -> None: ...
|
||||
def close(self) -> Awaitable[None]: ...
|
||||
|
||||
|
||||
class RustResponsesWebSocketConnection(Protocol):
|
||||
@classmethod
|
||||
async def connect(
|
||||
def connect(
|
||||
cls,
|
||||
url: str,
|
||||
headers: dict[str, str],
|
||||
timeout_seconds: float | None,
|
||||
) -> RustResponsesWebSocket: ...
|
||||
) -> Awaitable[RustResponsesWebSocket]: ...
|
||||
|
||||
|
||||
class _Unset:
|
||||
|
|
@ -53,16 +54,10 @@ def set_rust_responses_websocket(
|
|||
_STATE.connection = connection
|
||||
|
||||
|
||||
def load_rust_responses_websocket() -> RustResponsesWebSocketConnection | None:
|
||||
def load_rust_responses_websocket() -> RustResponsesWebSocketConnection:
|
||||
if _STATE.connection is not None:
|
||||
return _STATE.connection
|
||||
native_bridge: Final = get_native_bridge()
|
||||
if native_bridge is None:
|
||||
return None
|
||||
connection_type: Final[RustResponsesWebSocketConnection | None] = getattr(
|
||||
native_bridge, "ResponsesWebSocketConnection", None
|
||||
)
|
||||
return connection_type
|
||||
return _native.ResponsesWebSocketConnection
|
||||
|
||||
|
||||
class _ConnectionAdapter:
|
||||
|
|
@ -89,8 +84,6 @@ async def connect(
|
|||
timeout: float | httpx.Timeout | None,
|
||||
) -> _ConnectionAdapter | None:
|
||||
connection_type: Final = load_rust_responses_websocket()
|
||||
if connection_type is None:
|
||||
return None
|
||||
try:
|
||||
connection: Final = await connection_type.connect(
|
||||
url=url,
|
||||
|
|
|
|||
|
|
@ -7,7 +7,8 @@ from typing import Final, Generic, NoReturn, TypeAlias, TypeVar
|
|||
from typing_extensions import assert_never
|
||||
|
||||
from litellm.exceptions import APIError
|
||||
from litellm.rust_bridge.bindings import NativeBinding, native_exception_types
|
||||
from litellm.rust_bridge import _native
|
||||
from litellm.rust_bridge.bindings import NativeBinding
|
||||
from litellm.rust_bridge.catalog import RouteContext, Rules, decision
|
||||
from litellm.rust_bridge.configuration import Decision
|
||||
from litellm.rust_bridge.response_metadata import mark_rust_response
|
||||
|
|
@ -153,15 +154,11 @@ def attempt(
|
|||
) -> RustAttempt[ResultT]:
|
||||
if native_call is None:
|
||||
return RustUnavailable()
|
||||
exceptions: Final = native_exception_types()
|
||||
if exceptions is None:
|
||||
return RustHandled(adapt(native_call()))
|
||||
declined, upstream = exceptions
|
||||
try:
|
||||
value: Final = native_call()
|
||||
except declined as error:
|
||||
except _native.RustBridgeDeclined as error:
|
||||
return RustDeclined(reason=_decline_reason(error))
|
||||
except upstream as error:
|
||||
except _native.RustUpstreamError as error:
|
||||
_raise_upstream(error, context)
|
||||
return RustHandled(adapt(value))
|
||||
|
||||
|
|
@ -174,15 +171,11 @@ async def aattempt(
|
|||
) -> RustAttempt[ResultT]:
|
||||
if native_call is None:
|
||||
return RustUnavailable()
|
||||
exceptions: Final = native_exception_types()
|
||||
if exceptions is None:
|
||||
return RustHandled(adapt(await native_call()))
|
||||
declined, upstream = exceptions
|
||||
try:
|
||||
value: Final = await native_call()
|
||||
except declined as error:
|
||||
except _native.RustBridgeDeclined as error:
|
||||
return RustDeclined(reason=_decline_reason(error))
|
||||
except upstream as error:
|
||||
except _native.RustUpstreamError as error:
|
||||
_raise_upstream(error, context)
|
||||
return RustHandled(adapt(value))
|
||||
|
||||
|
|
|
|||
|
|
@ -1,20 +1,9 @@
|
|||
import json
|
||||
from typing import Final, Protocol, runtime_checkable
|
||||
from typing import Final
|
||||
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm.rust_bridge.loader import get_native_bridge
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class NativeOtlpError(Protocol):
|
||||
def trace_encode_error(self, message: str) -> bytes: ...
|
||||
|
||||
|
||||
_NATIVE: Final[TypeAdapter[NativeOtlpError]] = TypeAdapter(
|
||||
NativeOtlpError, config=ConfigDict(arbitrary_types_allowed=True)
|
||||
)
|
||||
from litellm.rust_bridge import _native
|
||||
|
||||
|
||||
class OTLPError(TypedDict):
|
||||
|
|
@ -22,8 +11,7 @@ class OTLPError(TypedDict):
|
|||
|
||||
|
||||
def encode_error(message: str) -> bytes:
|
||||
native: Final = get_native_bridge()
|
||||
return b"" if native is None else _NATIVE.validate_python(native).trace_encode_error(message)
|
||||
return _native.trace_encode_error(message)
|
||||
|
||||
|
||||
def encode_otlp_response(content_type: str | None, error: str | None = None) -> tuple[bytes, str]:
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -1225,7 +1225,7 @@
|
|||
"responses": true,
|
||||
"text_completion": true,
|
||||
"embeddings": true,
|
||||
"image_generations": false,
|
||||
"image_generations": true,
|
||||
"image_edits": false,
|
||||
"audio_transcriptions": false,
|
||||
"audio_speech": false,
|
||||
|
|
|
|||
|
|
@ -401,7 +401,7 @@ markers = [
|
|||
"asyncio: mark test as an asyncio test",
|
||||
"limit_leaks: mark test with memory limit for leak detection (e.g., '40 MB')",
|
||||
"no_parallel: mark test to run sequentially (not in parallel) - typically for memory measurement tests",
|
||||
"requires_rust_extension: public Python contract requiring an enabled, compiled Rust extension",
|
||||
"requires_rust_extension: public Python contract tests that run against the release-built wheel",
|
||||
]
|
||||
filterwarnings = [
|
||||
# Suppress Pydantic serializer warnings from mock server responses (non-critical for memory tests)
|
||||
|
|
|
|||
|
|
@ -1163,6 +1163,7 @@ model LiteLLM_ManagedFileTable {
|
|||
@@index([unified_file_id])
|
||||
@@index([flat_model_file_ids], type: Gin)
|
||||
@@index([team_id, created_at(sort: Desc)])
|
||||
@@index([created_by, created_at(sort: Desc)])
|
||||
}
|
||||
|
||||
model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use the
|
||||
|
|
|
|||
|
|
@ -8,6 +8,8 @@ the very class of undeclared-dependency bug this guards against.
|
|||
|
||||
import argparse
|
||||
import importlib.util
|
||||
from importlib.machinery import EXTENSION_SUFFIXES
|
||||
from pathlib import Path
|
||||
import re
|
||||
import sys
|
||||
import traceback
|
||||
|
|
@ -89,6 +91,18 @@ def check_import() -> str:
|
|||
return f"imported {distribution} {sdk_version}"
|
||||
|
||||
|
||||
def check_native_extension() -> str:
|
||||
import litellm
|
||||
from litellm.rust_bridge import _native
|
||||
|
||||
native_path: Final = Path(_native.__file__).resolve()
|
||||
package_path: Final = Path(litellm.__file__).resolve().parent
|
||||
_require(native_path.parent == package_path / "rust_bridge", "native extension is outside the SDK package")
|
||||
_require(any(str(native_path).endswith(suffix) for suffix in EXTENSION_SUFFIXES), "native module is not compiled")
|
||||
_require(_native.trace_encode_error("ping") == b"\x12\x04ping", "native OTLP error encoding failed")
|
||||
return f"loaded and exercised {native_path}"
|
||||
|
||||
|
||||
def check_completion() -> str:
|
||||
import litellm
|
||||
|
||||
|
|
@ -230,6 +244,7 @@ def check_bedrock_credential_resolution() -> str:
|
|||
CHECKS: tuple[tuple[str, Callable[[], str]], ...] = (
|
||||
("environment is base-only", check_environment_is_base_only),
|
||||
("import litellm", check_import),
|
||||
("native extension", check_native_extension),
|
||||
("optional MCP installation guidance", check_mcp_install_guidance),
|
||||
("chat completion", check_completion),
|
||||
("chat completion with a function tool", check_completion_with_function_tool),
|
||||
|
|
@ -250,8 +265,11 @@ def _run(check: Callable[[], str]) -> tuple[bool, str]:
|
|||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--profile", choices=("legacy", "core", "dependencies"), default="legacy")
|
||||
parser.add_argument("--profile", choices=("legacy", "core", "dependencies", "native"), default="legacy")
|
||||
profile = parser.parse_args().profile
|
||||
if profile == "native":
|
||||
print(check_native_extension())
|
||||
return 0
|
||||
for module in ("boto3", "botocore", "tokenizers", "huggingface_hub"):
|
||||
present = importlib.util.find_spec(module) is not None
|
||||
_require(present == (profile != "core"), f"{profile}: unexpected presence of {module}: {present}")
|
||||
|
|
|
|||
|
|
@ -0,0 +1,122 @@
|
|||
import json
|
||||
import uuid
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.client import Gateway
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
MODEL: Final = "claude-haiku-4-5"
|
||||
|
||||
|
||||
def _sse(event: Mapping[str, JsonValue]) -> bytes:
|
||||
return f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode()
|
||||
|
||||
|
||||
def _message(identity: str, stop_reason: str | None) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": "msg_" + identity,
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": MODEL,
|
||||
"content": [{"type": "text", "text": "alpha bravo"}],
|
||||
"stop_reason": stop_reason,
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 199878, "output_tokens": 12799},
|
||||
}
|
||||
|
||||
|
||||
def _stream(identity: str, stop_reason: str) -> tuple[bytes, ...]:
|
||||
return (
|
||||
_sse({"type": "message_start", "message": {**_message(identity, None), "content": []}}),
|
||||
_sse({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
|
||||
_sse({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "alpha bravo"}}),
|
||||
_sse({"type": "content_block_stop", "index": 0}),
|
||||
_sse(
|
||||
{
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": stop_reason, "stop_sequence": None},
|
||||
"usage": {"output_tokens": 12799},
|
||||
}
|
||||
),
|
||||
_sse({"type": "message_stop"}),
|
||||
)
|
||||
|
||||
|
||||
def _anthropic(identity: str, stop_reason: str):
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST" and request.target == "/v1/messages", request.target
|
||||
if json.loads(request.body).get("stream"):
|
||||
return Reply(content_type="text/event-stream", chunks=_stream(identity, stop_reason))
|
||||
return Reply(body=json.dumps(_message(identity, stop_reason)).encode())
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def _events(gateway: Gateway, path: str, body: Mapping[str, JsonValue]) -> tuple[dict[str, JsonValue], ...]:
|
||||
with gateway.client.stream(
|
||||
"POST", path, json={**body, "stream": True}, headers={"Authorization": f"Bearer {gateway.key}"}
|
||||
) as response:
|
||||
assert response.status_code == 200, response.read().decode()
|
||||
return tuple(
|
||||
json.loads(line.removeprefix("data: "))
|
||||
for line in response.iter_lines()
|
||||
if line.startswith("data: ") and line != "data: [DONE]"
|
||||
)
|
||||
|
||||
|
||||
def _chat_finish_reason(gateway: Gateway, model: str, *, stream: bool) -> JsonValue:
|
||||
body: Final = {"model": model, "max_tokens": 64000, "messages": [{"role": "user", "content": "copy the list"}]}
|
||||
if stream:
|
||||
reasons: Final = tuple(
|
||||
choice["finish_reason"]
|
||||
for chunk in _events(gateway, "/v1/chat/completions", body)
|
||||
for choice in chunk["choices"] # comprehension-ok: chunk -> choices flatten
|
||||
if choice.get("finish_reason")
|
||||
)
|
||||
assert len(reasons) == 1, reasons
|
||||
return reasons[0]
|
||||
response: Final = gateway.request("POST", "/v1/chat/completions", body)
|
||||
assert response.status_code == 200, response.text
|
||||
return response.json()["choices"][0]["finish_reason"]
|
||||
|
||||
|
||||
def _responses_status(gateway: Gateway, model: str, *, stream: bool) -> JsonValue:
|
||||
body: Final = {"model": model, "max_output_tokens": 64000, "input": "copy the list"}
|
||||
if stream:
|
||||
terminal: Final = tuple(
|
||||
event["response"]
|
||||
for event in _events(gateway, "/v1/responses", body)
|
||||
if event["type"] in ("response.completed", "response.incomplete", "response.failed")
|
||||
)
|
||||
assert len(terminal) == 1, terminal
|
||||
return terminal[0]["status"]
|
||||
response: Final = gateway.request("POST", "/v1/responses", body)
|
||||
assert response.status_code == 200, response.text
|
||||
return response.json()["status"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream", [False, True], ids=["json", "stream"])
|
||||
@pytest.mark.parametrize(
|
||||
("stop_reason", "finish_reason"),
|
||||
[("model_context_window_exceeded", "length"), ("max_tokens", "length"), ("end_turn", "stop")],
|
||||
)
|
||||
def test_chat_maps_anthropic_stop_reason(gateway: Gateway, stop_reason: str, finish_reason: str, stream: bool) -> None:
|
||||
identity: Final = uuid.uuid4().hex
|
||||
with wire_server(_anthropic(identity, stop_reason)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"anthropic/{MODEL}", api_base=wire.url, api_key="scripted-anthropic-key")
|
||||
assert _chat_finish_reason(gateway, model, stream=stream) == finish_reason
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream", [False, True], ids=["json", "stream"])
|
||||
@pytest.mark.parametrize(
|
||||
("stop_reason", "status"),
|
||||
[("model_context_window_exceeded", "incomplete"), ("max_tokens", "incomplete"), ("end_turn", "completed")],
|
||||
)
|
||||
def test_responses_maps_anthropic_stop_reason(gateway: Gateway, stop_reason: str, status: str, stream: bool) -> None:
|
||||
identity: Final = uuid.uuid4().hex
|
||||
with wire_server(_anthropic(identity, stop_reason)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"anthropic/{MODEL}", api_base=wire.url, api_key="scripted-anthropic-key")
|
||||
assert _responses_status(gateway, model, stream=stream) == status
|
||||
263
tests/integration/security/test_managed_file_listing.py
Normal file
263
tests/integration/security/test_managed_file_listing.py
Normal file
|
|
@ -0,0 +1,263 @@
|
|||
import json
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass, field
|
||||
from itertools import count, islice
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from urllib.parse import parse_qs, urlsplit
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, gateway_from_environment
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from anthropic import Anthropic
|
||||
from openai import OpenAI
|
||||
from pydantic import BaseModel, ConfigDict, JsonValue
|
||||
|
||||
pytestmark: Final = pytest.mark.shares_worker_state
|
||||
|
||||
UNMANAGED_FILE: Final = "file-uploaded-straight-to-the-provider"
|
||||
UPLOAD: Final = b"isolation check\n"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Deployment:
|
||||
provider: str
|
||||
model: str
|
||||
litellm_model: str
|
||||
api_key: str
|
||||
api_path: str
|
||||
|
||||
|
||||
OPENAI: Final = Deployment("openai", "gpt-4.1-managed-listing", "openai/gpt-4.1", "sk-managed-listing", "/v1")
|
||||
ANTHROPIC: Final = Deployment(
|
||||
"anthropic", "claude-managed-listing", "anthropic/claude-haiku-5-5", "sk-ant-managed-listing", ""
|
||||
)
|
||||
DEPLOYMENTS: Final = (pytest.param(OPENAI, id="openai"), pytest.param(ANTHROPIC, id="anthropic"))
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ListCall:
|
||||
path: str
|
||||
params: Mapping[str, str] = field(default_factory=lambda: MappingProxyType({}))
|
||||
headers: Mapping[str, str] = field(default_factory=lambda: MappingProxyType({}))
|
||||
|
||||
|
||||
LIST_FORMS: Final = (
|
||||
pytest.param(lambda d: ListCall(f"/{d.provider}/v1/files"), id="provider-path"),
|
||||
pytest.param(lambda d: ListCall("/v1/files", params={"provider": d.provider}), id="provider-query"),
|
||||
pytest.param(lambda d: ListCall("/v1/files", headers={"custom-llm-provider": d.provider}), id="provider-header"),
|
||||
pytest.param(lambda d: ListCall("/v1/files", params={"target_model_names": d.model}), id="target-model-names"),
|
||||
pytest.param(lambda d: ListCall("/v1/files", headers={"x-litellm-model": d.model}), id="model-header"),
|
||||
)
|
||||
|
||||
|
||||
class FilePage(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore", frozen=True)
|
||||
|
||||
data: list[dict[str, JsonValue]]
|
||||
has_more: bool
|
||||
|
||||
@property
|
||||
def ids(self) -> list[str]:
|
||||
return [str(entry["id"]) for entry in self.data]
|
||||
|
||||
|
||||
def _openai_file(file_id: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": file_id,
|
||||
"object": "file",
|
||||
"bytes": len(UPLOAD),
|
||||
"created_at": 1700000000,
|
||||
"filename": "notes.txt",
|
||||
"purpose": "user_data",
|
||||
"status": "processed",
|
||||
}
|
||||
|
||||
|
||||
def _anthropic_file(file_id: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": file_id,
|
||||
"type": "file",
|
||||
"filename": "notes.txt",
|
||||
"mime_type": "text/plain",
|
||||
"size_bytes": len(UPLOAD),
|
||||
"created_at": "2026-01-01T00:00:00Z",
|
||||
"downloadable": False,
|
||||
}
|
||||
|
||||
|
||||
def _reply(body: Mapping[str, JsonValue]) -> Reply:
|
||||
return Reply(body=json.dumps(body).encode())
|
||||
|
||||
|
||||
def _list_reply(provider: str, file_ids: tuple[str, ...]) -> Reply:
|
||||
if provider == ANTHROPIC.provider:
|
||||
return _reply(
|
||||
{
|
||||
"data": [_anthropic_file(file_id) for file_id in file_ids],
|
||||
"has_more": False,
|
||||
"first_id": file_ids[0],
|
||||
"last_id": file_ids[-1],
|
||||
}
|
||||
)
|
||||
return _reply({"object": "list", "data": [_openai_file(file_id) for file_id in file_ids]})
|
||||
|
||||
|
||||
def _responder(provider_files: Mapping[str, tuple[str, ...]]) -> Callable[[Request], Reply]:
|
||||
uploads: Final = {provider: count() for provider in provider_files}
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
target: Final = urlsplit(request.target)
|
||||
path: Final = target.path
|
||||
provider: Final = ANTHROPIC.provider if "x-api-key" in request.headers else OPENAI.provider
|
||||
files: Final = provider_files[provider]
|
||||
if (request.method, path) == ("GET", "/v1/models"):
|
||||
return _reply({"object": "list", "data": []})
|
||||
if (request.method, path) == ("POST", "/v1/files"):
|
||||
uploaded: Final = files[next(uploads[provider])]
|
||||
return _reply(_anthropic_file(uploaded) if provider == ANTHROPIC.provider else _openai_file(uploaded))
|
||||
if (request.method, path) == ("GET", "/v1/files"):
|
||||
limit: Final = int(parse_qs(target.query).get("limit", ["10000"])[0])
|
||||
return _list_reply(provider, (UNMANAGED_FILE, *reversed(files))[:limit])
|
||||
return Reply(status=404, body=json.dumps({"error": {"message": f"unscripted {request.target}"}}).encode())
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def _config(wire: Wire) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"model_list": [
|
||||
{
|
||||
"model_name": deployment.model,
|
||||
"litellm_params": {
|
||||
"model": deployment.litellm_model,
|
||||
"api_key": deployment.api_key,
|
||||
"api_base": f"{wire.url}{deployment.api_path}",
|
||||
},
|
||||
}
|
||||
for deployment in (OPENAI, ANTHROPIC)
|
||||
],
|
||||
"litellm_settings": {"require_managed_files": True},
|
||||
"general_settings": {"master_key": "os.environ/LITELLM_MASTER_KEY", "database_url": "os.environ/DATABASE_URL"},
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Rig:
|
||||
proxy: Gateway
|
||||
owner: str
|
||||
intruder: str
|
||||
files_by_provider: Mapping[str, tuple[str, ...]]
|
||||
every_file: tuple[str, ...]
|
||||
|
||||
@property
|
||||
def base_url(self) -> str:
|
||||
return str(self.proxy.client.base_url).rstrip("/")
|
||||
|
||||
def list(self, call: ListCall, key: str | None = None) -> FilePage:
|
||||
response: Final = self.proxy.request("GET", call.path, key=key, params=call.params, headers=call.headers)
|
||||
assert response.status_code == 200, response.text
|
||||
return FilePage.model_validate_json(response.text)
|
||||
|
||||
|
||||
def _upload(proxy: Gateway, key: str, deployments: tuple[Deployment, ...]) -> str:
|
||||
with OpenAI(base_url=f"{str(proxy.client.base_url).rstrip('/')}/openai/v1", api_key=key, max_retries=0) as client:
|
||||
uploaded: Final = client.files.with_raw_response.create(
|
||||
file=("notes.txt", UPLOAD, "text/plain"),
|
||||
purpose="user_data",
|
||||
extra_body={"target_model_names": ",".join(deployment.model for deployment in deployments)},
|
||||
)
|
||||
assert uploaded.status_code == 200, uploaded.text
|
||||
return uploaded.parse().id
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _rig(directory: Path) -> Iterator[Rig]:
|
||||
build: Final = uuid.uuid4().hex
|
||||
provider_files: Final = {
|
||||
OPENAI.provider: (f"file-older-{build}", f"file-newer-{build}"),
|
||||
ANTHROPIC.provider: (f"file_older_{build}", f"file_newer_{build}"),
|
||||
}
|
||||
with gateway_from_environment() as shared, wire_server(_responder(provider_files)) as wire:
|
||||
config: Final = directory / "config.yaml"
|
||||
config.write_text(yaml.safe_dump(_config(wire)))
|
||||
with owned_proxy(shared, directory, {}, config=config) as proxy, proxy.scenario() as scenario:
|
||||
owner: Final = scenario.key(user_id=scenario.user())
|
||||
intruder: Final = scenario.key(user_id=scenario.user())
|
||||
on_both: Final = _upload(proxy, owner, (OPENAI, ANTHROPIC))
|
||||
on_openai: Final = _upload(proxy, owner, (OPENAI,))
|
||||
on_anthropic: Final = _upload(proxy, owner, (ANTHROPIC,))
|
||||
yield Rig(
|
||||
proxy,
|
||||
owner,
|
||||
intruder,
|
||||
MappingProxyType(
|
||||
{OPENAI.provider: (on_openai, on_both), ANTHROPIC.provider: (on_anthropic, on_both)}
|
||||
),
|
||||
(on_anthropic, on_openai, on_both),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
with _rig(tmp_path_factory.mktemp("managed-file-listing")) as built:
|
||||
yield built
|
||||
|
||||
|
||||
@pytest.mark.parametrize("deployment", DEPLOYMENTS)
|
||||
@pytest.mark.parametrize("form", LIST_FORMS)
|
||||
def test_another_key_lists_none_of_the_files_on_the_shared_key(
|
||||
rig: Rig, deployment: Deployment, form: Callable[[Deployment], ListCall]
|
||||
) -> None:
|
||||
page: Final = rig.list(form(deployment), key=rig.intruder)
|
||||
assert (page.ids, page.has_more) == ([], False)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("deployment", DEPLOYMENTS)
|
||||
@pytest.mark.parametrize("form", LIST_FORMS)
|
||||
def test_owner_lists_their_own_managed_files_on_that_provider_or_model_newest_first(
|
||||
rig: Rig, deployment: Deployment, form: Callable[[Deployment], ListCall]
|
||||
) -> None:
|
||||
page: Final = rig.list(form(deployment), key=rig.owner)
|
||||
assert (page.ids, page.has_more) == (list(rig.files_by_provider[deployment.provider]), False)
|
||||
|
||||
|
||||
def test_owner_lists_every_managed_file_when_no_provider_or_model_is_named(rig: Rig) -> None:
|
||||
page: Final = rig.list(ListCall("/v1/files"), key=rig.owner)
|
||||
assert (page.ids, page.has_more) == (list(rig.every_file), False)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("cursor_param", ["after", "after_id"])
|
||||
def test_owner_pages_through_a_provider_list_with_limit_and_a_cursor(rig: Rig, cursor_param: str) -> None:
|
||||
first: Final = rig.list(ListCall("/openai/v1/files", params={"limit": "1"}), key=rig.owner)
|
||||
second: Final = rig.list(
|
||||
ListCall("/openai/v1/files", params={"limit": "1", cursor_param: first.ids[0]}), key=rig.owner
|
||||
)
|
||||
expected: Final = rig.files_by_provider[OPENAI.provider]
|
||||
assert [(page.ids, page.has_more) for page in (first, second)] == [
|
||||
([expected[0]], True),
|
||||
([expected[1]], False),
|
||||
]
|
||||
|
||||
|
||||
def test_anthropic_sdk_auto_pages_the_owners_files_to_the_end(rig: Rig) -> None:
|
||||
with Anthropic(base_url=f"{rig.base_url}/anthropic", api_key=rig.owner, max_retries=0) as client:
|
||||
listed: Final = [file.id for file in islice(client.beta.files.list(limit=1), 5)]
|
||||
assert listed == list(rig.files_by_provider[ANTHROPIC.provider])
|
||||
|
||||
|
||||
def test_anthropic_sdk_pages_backwards_to_the_newest_file_and_stops(rig: Rig) -> None:
|
||||
newest_first: Final = rig.files_by_provider[ANTHROPIC.provider]
|
||||
with Anthropic(base_url=f"{rig.base_url}/anthropic", api_key=rig.owner, max_retries=0) as client:
|
||||
listed: Final = [file.id for file in islice(client.beta.files.list(limit=1, before_id=newest_first[-1]), 5)]
|
||||
assert listed == list(reversed(newest_first[:-1]))
|
||||
|
||||
|
||||
def test_master_key_still_lists_the_files_on_the_shared_provider_key(rig: Rig) -> None:
|
||||
listed: Final = rig.list(ListCall(f"/{ANTHROPIC.provider}/v1/files")).ids
|
||||
assert UNMANAGED_FILE in listed
|
||||
|
|
@ -1,6 +1,4 @@
|
|||
from collections.abc import AsyncIterator, Generator
|
||||
from importlib import import_module
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
|
@ -23,16 +21,6 @@ def recording_server() -> Generator[RecordingServer]:
|
|||
yield server
|
||||
|
||||
|
||||
def pytest_collection_modifyitems(items: list[pytest.Item]) -> None:
|
||||
try:
|
||||
import_module("litellm.rust_bridge._native")
|
||||
except ImportError:
|
||||
skip: Final = pytest.mark.skip(reason="requires a compiled litellm.rust_bridge._native extension")
|
||||
for item in items:
|
||||
if "test_litellm_rust" in item.path.parts:
|
||||
item.add_marker(skip)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def isolated_azure_auth(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
for name in (
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ import pytest
|
|||
from typing import Final, Optional
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from litellm.llms.base_llm.files.transformation import FileListModelScope
|
||||
from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import BATCH_CREATE_HIDDEN_PARAM
|
||||
from litellm.types.llms.openai import FileListPage, OpenAIFileObject
|
||||
|
|
@ -649,6 +650,22 @@ async def test_afile_list_denies_a_caller_with_no_identity_at_all():
|
|||
assert table.find_many_calls == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_afile_list_reads_no_rows_when_no_deployment_can_hold_a_file() -> None:
|
||||
managed_files, table = _make_managed_files_over_rows([_make_managed_file_row("unified-mine")])
|
||||
|
||||
response: Final = await managed_files.afile_list(
|
||||
purpose=None,
|
||||
litellm_parent_otel_span=None,
|
||||
user_api_key_dict=_make_user_api_key_dict(),
|
||||
model_scope=FileListModelScope(frozenset()),
|
||||
)
|
||||
|
||||
assert response.data == []
|
||||
assert response.has_more is False
|
||||
assert table.find_many_calls == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_afile_list_filters_by_purpose():
|
||||
managed_files, _ = _make_managed_files_over_rows(
|
||||
|
|
|
|||
|
|
@ -256,6 +256,7 @@ class TestMapFinishReasonAnthropic:
|
|||
("max_tokens", "length"),
|
||||
("tool_use", "tool_calls"),
|
||||
("compaction", "length"),
|
||||
("model_context_window_exceeded", "length"),
|
||||
("content_filtered", "content_filter"),
|
||||
],
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1516,9 +1516,9 @@ def test_bedrock_mantle_openai_envelope_context_overflow_maps_to_context_window_
|
|||
|
||||
|
||||
def test_bedrock_mantle_streamed_context_overflow_event_maps_to_context_window_exceeded():
|
||||
from litellm.responses.streaming_iterator import _map_stream_error_to_exception
|
||||
from litellm.responses.streaming_iterator import map_stream_error_to_exception
|
||||
|
||||
mapped_exception = _map_stream_error_to_exception(
|
||||
mapped_exception = map_stream_error_to_exception(
|
||||
{
|
||||
"code": "context_length_exceeded",
|
||||
"message": "Your input exceeds the context window of this model. Please adjust your input and try again.",
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import asyncio
|
||||
import datetime
|
||||
import itertools
|
||||
import json
|
||||
import queue
|
||||
import sys
|
||||
|
|
@ -13,6 +14,7 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch
|
|||
|
||||
import httpx
|
||||
import pytest
|
||||
from openai import AsyncOpenAI, AsyncStream, OpenAI, Stream
|
||||
from openai.types.chat import ChatCompletionChunk
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
import litellm
|
||||
|
|
@ -28,6 +30,7 @@ from litellm.litellm_core_utils.streaming_handler import (
|
|||
_ProviderChunkParsed,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler
|
||||
from litellm.responses.mcp.chat_completions_handler import _add_mcp_metadata_to_response
|
||||
from litellm.types.utils import (
|
||||
CompletionTokensDetailsWrapper,
|
||||
|
|
@ -1367,6 +1370,148 @@ async def test_bedrock_validation_error_raises_directly(logging_obj: Logging):
|
|||
assert getattr(excinfo.value, "status_code", None) == 400
|
||||
|
||||
|
||||
@pytest.mark.parametrize("async_stream", [False, True])
|
||||
@pytest.mark.parametrize("transport", ["sdk", "http_handler"])
|
||||
@pytest.mark.parametrize("prefix", ["", "Partial answer"])
|
||||
@pytest.mark.parametrize("abort_index", [0, 1])
|
||||
@pytest.mark.parametrize("abort_content", ["", "unfinished"])
|
||||
async def test_aborted_stream_raises_instead_of_finishing_successfully(
|
||||
async_stream: bool, transport: str, prefix: str, abort_index: int, abort_content: str, logging_obj: Logging
|
||||
) -> None:
|
||||
initial: Final = (
|
||||
(({"index": 0, "delta": {"role": "assistant", "content": prefix}, "finish_reason": None},),) if prefix else ()
|
||||
)
|
||||
final: Final = tuple(
|
||||
{
|
||||
"index": index,
|
||||
"delta": {"content": abort_content},
|
||||
"finish_reason": "abort" if index == abort_index else "stop",
|
||||
}
|
||||
for index in range(abort_index + 1)
|
||||
)
|
||||
body: Final = (
|
||||
"".join(
|
||||
"data: "
|
||||
+ json.dumps(
|
||||
{
|
||||
"id": "chatcmpl-aborted",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": "test-model",
|
||||
"choices": choices,
|
||||
}
|
||||
)
|
||||
+ "\n\n"
|
||||
for choices in (*initial, final)
|
||||
)
|
||||
+ "data: [DONE]\n\n"
|
||||
)
|
||||
raw_response: Final = httpx.Response(
|
||||
200,
|
||||
content=body,
|
||||
headers={"content-type": "text/event-stream"},
|
||||
request=httpx.Request("POST", "https://provider.example/v1/chat/completions"),
|
||||
)
|
||||
client: Final = AsyncOpenAI(api_key="unused") if async_stream else OpenAI(api_key="unused")
|
||||
stream: Final = (
|
||||
(
|
||||
AsyncStream(cast_to=ChatCompletionChunk, response=raw_response, client=client)
|
||||
if isinstance(client, AsyncOpenAI)
|
||||
else Stream(cast_to=ChatCompletionChunk, response=raw_response, client=client)
|
||||
)
|
||||
if transport == "sdk"
|
||||
else OpenAIChatCompletionStreamingHandler(
|
||||
streaming_response=raw_response.aiter_lines() if async_stream else raw_response.iter_lines(),
|
||||
sync_stream=not async_stream,
|
||||
)
|
||||
)
|
||||
response: Final = CustomStreamWrapper(
|
||||
completion_stream=stream,
|
||||
model="test-model",
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider="openai" if transport == "sdk" else "hosted_vllm",
|
||||
)
|
||||
|
||||
try:
|
||||
with pytest.raises(MidStreamFallbackError, match="abort") as excinfo:
|
||||
_chunks: Final = [chunk async for chunk in response] if async_stream else list(response)
|
||||
|
||||
assert excinfo.value.is_pre_first_chunk is (not prefix)
|
||||
assert excinfo.value.generated_content == prefix
|
||||
assert getattr(excinfo.value.original_exception, "status_code", None) == 500
|
||||
finally:
|
||||
if isinstance(client, AsyncOpenAI):
|
||||
await client.close()
|
||||
else:
|
||||
client.close()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("async_stream", [False, True])
|
||||
@pytest.mark.parametrize("transport", ["sdk", "http_handler"])
|
||||
async def test_stream_with_null_choices_chunk_finishes_normally(
|
||||
async_stream: bool, transport: str, logging_obj: Logging
|
||||
) -> None:
|
||||
body: Final = (
|
||||
"".join(
|
||||
"data: "
|
||||
+ json.dumps(
|
||||
{
|
||||
"id": "chatcmpl-null-choices",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": "test-model",
|
||||
"choices": choices,
|
||||
**extra,
|
||||
}
|
||||
)
|
||||
+ "\n\n"
|
||||
for choices, extra in (
|
||||
([{"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": None}], {}),
|
||||
([{"index": 0, "delta": {}, "finish_reason": "stop"}], {}),
|
||||
(None, {"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}}),
|
||||
)
|
||||
)
|
||||
+ "data: [DONE]\n\n"
|
||||
)
|
||||
raw_response: Final = httpx.Response(
|
||||
200,
|
||||
content=body,
|
||||
headers={"content-type": "text/event-stream"},
|
||||
request=httpx.Request("POST", "https://provider.example/v1/chat/completions"),
|
||||
)
|
||||
client: Final = AsyncOpenAI(api_key="unused") if async_stream else OpenAI(api_key="unused")
|
||||
stream: Final = (
|
||||
(
|
||||
AsyncStream(cast_to=ChatCompletionChunk, response=raw_response, client=client)
|
||||
if isinstance(client, AsyncOpenAI)
|
||||
else Stream(cast_to=ChatCompletionChunk, response=raw_response, client=client)
|
||||
)
|
||||
if transport == "sdk"
|
||||
else OpenAIChatCompletionStreamingHandler(
|
||||
streaming_response=raw_response.aiter_lines() if async_stream else raw_response.iter_lines(),
|
||||
sync_stream=not async_stream,
|
||||
)
|
||||
)
|
||||
response: Final = CustomStreamWrapper(
|
||||
completion_stream=stream,
|
||||
model="test-model",
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider="openai" if transport == "sdk" else "hosted_vllm",
|
||||
)
|
||||
|
||||
try:
|
||||
chunks: Final = [chunk async for chunk in response] if async_stream else list(response)
|
||||
finally:
|
||||
if isinstance(client, AsyncOpenAI):
|
||||
await client.close()
|
||||
else:
|
||||
client.close()
|
||||
|
||||
streamed_choices: Final = tuple(itertools.chain.from_iterable(chunk.choices for chunk in chunks))
|
||||
assert "".join(choice.delta.content or "" for choice in streamed_choices) == "hi"
|
||||
assert [choice.finish_reason for choice in streamed_choices if choice.finish_reason] == ["stop"]
|
||||
|
||||
|
||||
def _hosted_vllm_stream_wrapper(logging_obj: Logging, error_payload: dict) -> CustomStreamWrapper:
|
||||
"""A CustomStreamWrapper over the real OpenAI-compatible line iterator,
|
||||
fed an HTTP 200 SSE body that carries an in-body error payload the way
|
||||
|
|
|
|||
|
|
@ -241,20 +241,14 @@ print("compatible")
|
|||
|
||||
|
||||
@pytest.mark.parametrize("rust", (None, "0", "1"))
|
||||
def test_tokenization_without_native_extension_stays_offline(tmp_path: Path, rust: str | None) -> None:
|
||||
def test_tokenization_stays_offline(tmp_path: Path, rust: str | None) -> None:
|
||||
script: Final = """
|
||||
import importlib.abc
|
||||
import sys
|
||||
sys.path.insert(0, sys.argv[1])
|
||||
def reject_network(event, args):
|
||||
if event == "socket.connect":
|
||||
raise AssertionError("tokenizer attempted a network connection")
|
||||
sys.addaudithook(reject_network)
|
||||
class Block(importlib.abc.MetaPathFinder):
|
||||
def find_spec(self, fullname, path=None, target=None):
|
||||
if fullname == "litellm.rust_bridge._native":
|
||||
raise ImportError("native extension is unavailable")
|
||||
sys.meta_path.insert(0, Block())
|
||||
import litellm
|
||||
from litellm.rust_bridge.tokenizer import get_encoding
|
||||
import tiktoken
|
||||
|
|
|
|||
|
|
@ -1,15 +1,16 @@
|
|||
from typing import Final
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
from aiohttp import ClientResponse
|
||||
from pydantic import ValidationError
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm.llms.aiohttp_openai.chat.transformation import AiohttpOpenAIChatConfig
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
|
||||
async def _transform(body: object) -> ModelResponse:
|
||||
raw_response = Mock(spec=ClientResponse)
|
||||
async def _transform(body: object, *, requested_tier: str | None = None) -> ModelResponse:
|
||||
raw_response: Final = Mock(spec=ClientResponse)
|
||||
raw_response.json = AsyncMock(return_value=body)
|
||||
return await AiohttpOpenAIChatConfig().transform_response(
|
||||
model="gpt-4o",
|
||||
|
|
@ -18,7 +19,7 @@ async def _transform(body: object) -> ModelResponse:
|
|||
logging_obj=Mock(),
|
||||
request_data={},
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
optional_params={},
|
||||
optional_params={"service_tier": requested_tier} if requested_tier is not None else {},
|
||||
litellm_params={},
|
||||
encoding=None,
|
||||
)
|
||||
|
|
@ -99,3 +100,20 @@ async def test_transform_response_rejects_choices_that_are_not_a_list_of_objects
|
|||
async def test_transform_response_rejects_a_choice_whose_message_is_not_an_object():
|
||||
with pytest.raises(ValidationError, match="validation error for Choices"):
|
||||
await _transform({"choices": [{"message": "text"}]})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tier", ("flex", "priority", "default", "custom-tier", None))
|
||||
async def test_transform_response_serializes_served_tier_instead_of_requested_tier(tier: str | None) -> None:
|
||||
response: Final = await _transform(
|
||||
{"id": "chatcmpl-1", "choices": [{}], "service_tier": tier}, requested_tier="auto"
|
||||
)
|
||||
serialized: Final = TypeAdapter(dict[str, object]).validate_json(response.model_dump_json())
|
||||
|
||||
assert serialized["service_tier"] == tier
|
||||
|
||||
|
||||
async def test_transform_response_does_not_invent_a_served_tier_when_omitted() -> None:
|
||||
response: Final = await _transform({"id": "chatcmpl-1", "choices": [{}]}, requested_tier="flex")
|
||||
serialized: Final = TypeAdapter(dict[str, object]).validate_json(response.model_dump_json())
|
||||
|
||||
assert "service_tier" not in serialized
|
||||
|
|
|
|||
|
|
@ -4,20 +4,25 @@ Tests for AnthropicResponsesStreamWrapper
|
|||
"""
|
||||
|
||||
import asyncio
|
||||
import datetime
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable
|
||||
from types import SimpleNamespace
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../../..")))
|
||||
|
||||
import litellm
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
encrypted_reasoning_signature,
|
||||
)
|
||||
from litellm.llms.anthropic.messages.native.streaming_iterator import aclose_superseded_attempt
|
||||
from litellm.llms.anthropic.messages.native.utils import INCOMPLETE_STREAM_ERROR_MESSAGE
|
||||
from litellm.llms.anthropic.messages.responses_adapters.streaming_iterator import (
|
||||
AnthropicResponsesStreamWrapper,
|
||||
|
|
@ -622,3 +627,194 @@ class TestUpstreamFailureEndsStreamWithErrorEvent:
|
|||
chunks = _collect(iter(()))
|
||||
assert [chunk["type"] for chunk in chunks] == ["message_start", "error"]
|
||||
assert chunks[1]["error"] == {"type": "api_error", "message": INCOMPLETE_STREAM_ERROR_MESSAGE}
|
||||
|
||||
|
||||
class _FailureHookRecorder:
|
||||
def __init__(self) -> None:
|
||||
self.received: tuple[Exception, ...] = ()
|
||||
|
||||
async def __call__(self, failure: Exception) -> None:
|
||||
self.received = (*self.received, failure)
|
||||
|
||||
|
||||
def _logging_obj_with_failure_hook(on_failure: Callable[[Exception], Awaitable[None]]) -> LitellmLogging:
|
||||
logging_obj = LitellmLogging(
|
||||
model="m",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=True,
|
||||
call_type="anthropic_messages",
|
||||
start_time=datetime.datetime(2026, 1, 1),
|
||||
litellm_call_id="lit-6871-test",
|
||||
function_id="lit-6871-test",
|
||||
)
|
||||
logging_obj._on_detached_stream_failure = on_failure
|
||||
return logging_obj
|
||||
|
||||
|
||||
def _collect_with_logging(
|
||||
stream: AsyncIterator[dict[str, object]], logging_obj: LitellmLogging
|
||||
) -> list[dict[str, object]]:
|
||||
async def _run() -> list[dict[str, object]]:
|
||||
wrapper = AnthropicResponsesStreamWrapper(responses_stream=stream, model="m", litellm_logging_obj=logging_obj)
|
||||
return [chunk async for chunk in wrapper]
|
||||
|
||||
return asyncio.run(_run())
|
||||
|
||||
|
||||
class _DeliveryCountingConsumer:
|
||||
def __init__(self) -> None:
|
||||
self.delivered: tuple[dict[str, object], ...] = ()
|
||||
self.delivered_when_hook_fired: tuple[int, ...] = ()
|
||||
|
||||
async def hook(self, failure: Exception) -> None:
|
||||
self.delivered_when_hook_fired = (*self.delivered_when_hook_fired, len(self.delivered))
|
||||
|
||||
async def drain(self, wrapper: AnthropicResponsesStreamWrapper) -> None:
|
||||
async for chunk in wrapper:
|
||||
self.delivered = (*self.delivered, chunk)
|
||||
|
||||
|
||||
def _failing_after_one_delta() -> AsyncIterator[dict[str, object]]:
|
||||
async def _gen() -> AsyncIterator[dict[str, object]]:
|
||||
yield {"type": "response.created"}
|
||||
yield {"type": "response.output_item.added", "item": {"type": "message", "id": "msg_1"}}
|
||||
yield {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "Hi"}
|
||||
raise litellm.RateLimitError(message="You have no credits remaining.", llm_provider="openai", model="m")
|
||||
|
||||
return _gen()
|
||||
|
||||
|
||||
class TestUpstreamFailureReachesTheProxyFailureHook:
|
||||
def test_raised_upstream_failure_reaches_the_hook_unwrapped_once_the_error_event_was_consumed(self):
|
||||
rate_limit = litellm.RateLimitError(message="You have no credits remaining.", llm_provider="openai", model="m")
|
||||
wrapped = MidStreamFallbackError(
|
||||
message=str(rate_limit),
|
||||
model="m",
|
||||
llm_provider="openai",
|
||||
original_exception=rate_limit,
|
||||
is_pre_first_chunk=False,
|
||||
)
|
||||
recorder = _FailureHookRecorder()
|
||||
|
||||
async def _gen() -> AsyncIterator[dict[str, object]]:
|
||||
yield {"type": "response.created"}
|
||||
yield {"type": "response.output_item.added", "item": {"type": "message", "id": "msg_1"}}
|
||||
yield {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "Hi"}
|
||||
raise wrapped
|
||||
|
||||
chunks = _collect_with_logging(_gen(), _logging_obj_with_failure_hook(recorder))
|
||||
assert [chunk["type"] for chunk in chunks] == [
|
||||
"message_start",
|
||||
"content_block_start",
|
||||
"content_block_delta",
|
||||
"error",
|
||||
]
|
||||
assert recorder.received == (rate_limit,)
|
||||
|
||||
def test_response_failed_event_reaches_the_hook_as_an_exception_carrying_the_upstream_status(self):
|
||||
recorder = _FailureHookRecorder()
|
||||
failed = SimpleNamespace(
|
||||
status="failed",
|
||||
output=[],
|
||||
usage=None,
|
||||
error={"code": "rate_limit_exceeded", "message": "Rate limit reached for gpt-5.5, try again in 20s."},
|
||||
)
|
||||
|
||||
async def _gen() -> AsyncIterator[dict[str, object]]:
|
||||
yield {"type": "response.created"}
|
||||
yield {"type": "response.failed", "response": failed}
|
||||
|
||||
chunks = _collect_with_logging(_gen(), _logging_obj_with_failure_hook(recorder))
|
||||
assert chunks[-1]["type"] == "error"
|
||||
(failure,) = recorder.received
|
||||
assert isinstance(failure, litellm.RateLimitError)
|
||||
assert failure.status_code == 429
|
||||
assert "Rate limit reached for gpt-5.5, try again in 20s." in failure.message
|
||||
|
||||
def test_a_raising_hook_never_hides_the_error_event_from_the_client(self):
|
||||
async def _exploding_hook(failure: Exception) -> None:
|
||||
raise RuntimeError("spend writer is down")
|
||||
|
||||
logging_obj = _logging_obj_with_failure_hook(_exploding_hook)
|
||||
|
||||
async def _gen() -> AsyncIterator[dict[str, object]]:
|
||||
yield {"type": "response.created"}
|
||||
raise ConnectionResetError("Response payload is not completed")
|
||||
|
||||
chunks = _collect_with_logging(_gen(), logging_obj)
|
||||
assert [chunk["type"] for chunk in chunks] == ["message_start", "error"]
|
||||
assert chunks[-1]["error"] == {"type": "api_error", "message": "Response payload is not completed"}
|
||||
|
||||
def test_the_hook_fires_only_after_the_consumer_took_the_error_event(self):
|
||||
consumer = _DeliveryCountingConsumer()
|
||||
|
||||
async def _run() -> None:
|
||||
wrapper = AnthropicResponsesStreamWrapper(
|
||||
responses_stream=_failing_after_one_delta(),
|
||||
model="m",
|
||||
litellm_logging_obj=_logging_obj_with_failure_hook(consumer.hook),
|
||||
)
|
||||
await consumer.drain(wrapper)
|
||||
with pytest.raises(StopAsyncIteration):
|
||||
await wrapper.__anext__()
|
||||
|
||||
asyncio.run(_run())
|
||||
assert [chunk["type"] for chunk in consumer.delivered] == [
|
||||
"message_start",
|
||||
"content_block_start",
|
||||
"content_block_delta",
|
||||
"error",
|
||||
]
|
||||
assert consumer.delivered_when_hook_fired == (len(consumer.delivered),)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("close", "expected_reports"),
|
||||
[
|
||||
pytest.param(lambda sse: sse.aclose(), 1, id="client_closes"),
|
||||
pytest.param(aclose_superseded_attempt, 0, id="fallback_takes_over"),
|
||||
],
|
||||
)
|
||||
def test_closing_the_stream_on_the_error_event_reports_unless_a_fallback_takes_over(self, close, expected_reports):
|
||||
recorder = _FailureHookRecorder()
|
||||
|
||||
async def _close_on_the_error_event() -> None:
|
||||
wrapper = AnthropicResponsesStreamWrapper(
|
||||
responses_stream=_failing_after_one_delta(), model="m", litellm_logging_obj=_logging_obj_with_failure_hook(recorder)
|
||||
)
|
||||
sse = wrapper.async_anthropic_sse_wrapper()
|
||||
async for frame in sse:
|
||||
if frame.startswith(b"event: error\n"):
|
||||
await close(sse)
|
||||
return
|
||||
pytest.fail("the stream ended without an error event")
|
||||
|
||||
asyncio.run(_close_on_the_error_event())
|
||||
assert len(recorder.received) == expected_reports
|
||||
|
||||
def test_a_client_disconnect_during_the_failure_hook_lets_the_hook_finish(self):
|
||||
hook_started = anyio.Event()
|
||||
hook_released = anyio.Event()
|
||||
finished: list[Exception] = []
|
||||
|
||||
async def _held_hook(failure: Exception) -> None:
|
||||
hook_started.set()
|
||||
await hook_released.wait()
|
||||
finished.append(failure)
|
||||
|
||||
async def _disconnect_mid_hook() -> None:
|
||||
wrapper = AnthropicResponsesStreamWrapper(
|
||||
responses_stream=_failing_after_one_delta(), model="m", litellm_logging_obj=_logging_obj_with_failure_hook(_held_hook)
|
||||
)
|
||||
|
||||
async def _consume() -> None:
|
||||
async for _ in wrapper.async_anthropic_sse_wrapper():
|
||||
pass
|
||||
|
||||
async with anyio.create_task_group() as task_group:
|
||||
task_group.start_soon(_consume)
|
||||
await hook_started.wait()
|
||||
task_group.cancel_scope.cancel()
|
||||
hook_released.set()
|
||||
|
||||
asyncio.run(_disconnect_mid_hook())
|
||||
assert len(finished) == 1
|
||||
|
|
|
|||
|
|
@ -33,6 +33,13 @@ PRICE_FIELDS: Final = (
|
|||
"cache_read_input_token_cost",
|
||||
)
|
||||
PUBLISHED_DBU_PER_MILLION: Final = {
|
||||
"databricks/databricks-claude-haiku-5-5": ("1.429", "7.143", "1.786", "0.143"),
|
||||
"databricks/databricks-claude-sonnet-5-5": ("28.571", "142.857", "35.714", "1.429"),
|
||||
"databricks/databricks-gpt-6-1-sol": ("28.571", "142.857", "35.714", "1.429"),
|
||||
"databricks/databricks-gpt-6-1-terra": ("14.286", "71.429", "17.857", "0.714"),
|
||||
"databricks/databricks-gpt-6-astra": ("142.857", "714.286", "178.571", "14.286"),
|
||||
"databricks/databricks-gpt-6-luna": ("1.429", "7.143", "1.786", "0.143"),
|
||||
"databricks/databricks-gpt-6-sol": ("28.571", "142.857", "35.714", "2.857"),
|
||||
"databricks/databricks-claude-opus-5-5": ("57.143", "285.714", "71.429", "2.857"),
|
||||
"databricks/databricks-claude-fable-5-1": ("142.858", "714.286", "178.572", "3.572"),
|
||||
"databricks/databricks-claude-fable-5": ("142.858", "714.286", "178.572", "14.286"),
|
||||
|
|
@ -43,7 +50,7 @@ PUBLISHED_DBU_PER_MILLION: Final = {
|
|||
"databricks/databricks-claude-opus-4-5": ("71.429", "357.143", "89.286", "7.143"),
|
||||
"databricks/databricks-claude-opus-4-1": ("214.286", "1071.429", "267.857", "21.429"),
|
||||
"databricks/databricks-claude-opus-4": ("214.286", "1071.429", "267.857", "21.429"),
|
||||
"databricks/databricks-claude-sonnet-5": ("42.857", "214.286", "53.571", "4.286"),
|
||||
"databricks/databricks-claude-sonnet-5": ("28.571", "142.857", "35.714", "2.857"),
|
||||
"databricks/databricks-claude-sonnet-4-6": ("42.857", "214.286", "53.571", "4.286"),
|
||||
"databricks/databricks-claude-sonnet-4-5": ("42.857", "214.286", "53.571", "4.286"),
|
||||
"databricks/databricks-claude-sonnet-4-1": ("42.857", "214.286", "53.571", "4.286"),
|
||||
|
|
@ -197,7 +204,7 @@ def test_models_without_a_cache_discount_bill_cache_tokens_at_the_input_rate(
|
|||
assert prompt_cost > 8000 * info["input_cost_per_token"]
|
||||
|
||||
|
||||
def test_every_model_without_published_cache_dbu_bills_cache_at_its_own_input_rate(
|
||||
def test_every_model_without_published_cache_dbu_bills_cache_writes_at_its_own_input_rate(
|
||||
local_model_cost_map: None,
|
||||
) -> None:
|
||||
without_published_rates: Final = [
|
||||
|
|
@ -210,8 +217,7 @@ def test_every_model_without_published_cache_dbu_bills_cache_at_its_own_input_ra
|
|||
|
||||
for model in without_published_rates:
|
||||
info = _model_info(model)
|
||||
for field in CACHE_FIELDS:
|
||||
assert info[field] == pytest.approx(info["input_cost_per_token"]), (model, field)
|
||||
assert info["cache_creation_input_token_cost"] == pytest.approx(info["input_cost_per_token"]), model
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", NEW_MODELS)
|
||||
|
|
@ -224,14 +230,6 @@ def test_backup_price_map_matches_main(model: str) -> None:
|
|||
assert backup_cost[model] == main_cost[model]
|
||||
|
||||
|
||||
def test_sonnet_5_ships_standard_rates_not_introductory(local_model_cost_map: None) -> None:
|
||||
sonnet_5: Final = _model_info("databricks/databricks-claude-sonnet-5")
|
||||
sonnet_4_6: Final = _model_info("databricks/databricks-claude-sonnet-4-6")
|
||||
|
||||
for field in PRICE_FIELDS:
|
||||
assert sonnet_5[field] == pytest.approx(sonnet_4_6[field]), field
|
||||
|
||||
|
||||
def test_cost_per_token_bills_the_served_priority_tier(
|
||||
local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
|
|
|
|||
|
|
@ -3,13 +3,15 @@ import os
|
|||
import stat
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
import responses
|
||||
from click.testing import CliRunner
|
||||
from responses import matchers
|
||||
|
||||
from litellm.constants import CLI_JWT_EXPIRATION_HOURS
|
||||
from litellm.litellm_core_utils.cli_keyring import (
|
||||
DISABLE_KEYRING_ENV_VAR,
|
||||
KeyringDisabled,
|
||||
|
|
@ -56,6 +58,20 @@ def _write_token_file(home: Path, *, key: str | None) -> None:
|
|||
_write_home_json(home, "token.json", payload if key is None else {**payload, "key": key})
|
||||
|
||||
|
||||
STALE_TOKEN_BASE_URL: Final = "https://stale-token.example.com"
|
||||
STALE_TOKEN_TIMESTAMP: Final = 1_700_000_000.0
|
||||
|
||||
|
||||
def _register_models_answer(key: str, answer: int | Exception) -> None:
|
||||
models_url: Final = f"{STALE_TOKEN_BASE_URL}/v1/models"
|
||||
authorized: Final = [matchers.header_matcher({"Authorization": f"Bearer {key}"})]
|
||||
if isinstance(answer, Exception):
|
||||
responses.get(models_url, body=answer, match=authorized)
|
||||
else:
|
||||
responses.get(models_url, json={"data": []}, status=answer, match=authorized)
|
||||
responses.get(models_url, status=401)
|
||||
|
||||
|
||||
def _secret_blob(base_url: str, key: str) -> str:
|
||||
return json.dumps({"base_url": base_url, "key": key, "jwt_token": ""})
|
||||
|
||||
|
|
@ -696,22 +712,32 @@ class TestWhoamiCommand:
|
|||
assert "Not authenticated" in result.output
|
||||
assert "Run 'lite login'" in result.output
|
||||
|
||||
def test_whoami_old_token(self):
|
||||
"""Test whoami with old token showing warning"""
|
||||
token_data = CliTokenRecord(
|
||||
@pytest.mark.parametrize(
|
||||
("proxy_status", "expected_last_line_start"),
|
||||
[
|
||||
(200, "Token age: "),
|
||||
(401, "Token expired. Run 'lite login' again."),
|
||||
],
|
||||
)
|
||||
@responses.activate
|
||||
def test_whoami_lets_the_proxy_judge_a_token_older_than_the_client_window(
|
||||
self, proxy_status: int, expected_last_line_start: str
|
||||
):
|
||||
_register_models_answer("sk-live", proxy_status)
|
||||
token_data: Final = CliTokenRecord(
|
||||
base_url=STALE_TOKEN_BASE_URL,
|
||||
user_email="test@example.com",
|
||||
user_id="test-user-123",
|
||||
user_role="admin",
|
||||
key="sk-live",
|
||||
timestamp=time.time() - (25 * 3600),
|
||||
timestamp=STALE_TOKEN_TIMESTAMP,
|
||||
)
|
||||
|
||||
with patch("litellm.proxy.client.cli.commands.auth.load_cli_token", return_value=token_data):
|
||||
result = self.runner.invoke(whoami)
|
||||
result: Final = self.runner.invoke(whoami)
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert "Authenticated" in result.output
|
||||
assert "Warning: Token is more than 24 hours old" in result.output
|
||||
assert result.exit_code == 0
|
||||
assert result.output.splitlines()[-1].startswith(expected_last_line_start)
|
||||
assert "Warning" not in result.output
|
||||
assert len(responses.calls) == 1
|
||||
|
||||
def test_whoami_missing_fields(self):
|
||||
"""Test whoami with token missing some fields"""
|
||||
|
|
@ -1024,51 +1050,93 @@ class TestPrintTokenCommand:
|
|||
assert result.exit_code == 0
|
||||
assert result.output.strip() == "sk-matches"
|
||||
|
||||
@responses.activate
|
||||
def test_fresh_cached_key_printed_without_network_call(self):
|
||||
"""A recently-issued key should be printed straight from cache -- no
|
||||
refresh call on every single invocation (apiKeyHelper gets called
|
||||
frequently)."""
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.client.cli.commands.auth.load_cli_token",
|
||||
return_value=CliTokenRecord(
|
||||
base_url="http://localhost:4000",
|
||||
key="sk-cached-fresh",
|
||||
timestamp=time.time(),
|
||||
),
|
||||
with patch(
|
||||
"litellm.proxy.client.cli.commands.auth.load_cli_token",
|
||||
return_value=CliTokenRecord(
|
||||
base_url="http://localhost:4000",
|
||||
key="sk-cached-fresh",
|
||||
timestamp=time.time(),
|
||||
),
|
||||
patch("requests.post") as mock_post,
|
||||
):
|
||||
result = self.runner.invoke(print_token, obj={})
|
||||
result: Final = self.runner.invoke(print_token, obj={})
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert result.output.strip() == "sk-cached-fresh"
|
||||
mock_post.assert_not_called()
|
||||
assert len(responses.calls) == 0
|
||||
|
||||
def test_stale_key_fails_fast_without_network_call(self):
|
||||
"""There is no silent refresh: an expired cached key must fail
|
||||
loudly (stderr, nonzero exit) telling the user to `lite login`
|
||||
again, rather than making a network call or printing a dead key
|
||||
that will just 401 Claude Code."""
|
||||
old_timestamp = time.time() - (CLI_JWT_EXPIRATION_HOURS + 1) * 3600
|
||||
@pytest.mark.parametrize(
|
||||
("proxy_answer", "expected_exit_code", "expected_stdout", "expected_stderr"),
|
||||
[
|
||||
(200, 0, "sk-stale-key\n", ""),
|
||||
(401, 1, "", "Token expired. Run 'lite login' again.\n"),
|
||||
(403, 1, "", "Token expired. Run 'lite login' again.\n"),
|
||||
(404, 0, "sk-stale-key\n", ""),
|
||||
(500, 0, "sk-stale-key\n", ""),
|
||||
(requests.ConnectionError("proxy unreachable"), 0, "sk-stale-key\n", ""),
|
||||
],
|
||||
)
|
||||
@responses.activate
|
||||
def test_key_older_than_the_client_window_is_judged_by_the_proxy(
|
||||
self,
|
||||
proxy_answer: int | Exception,
|
||||
expected_exit_code: int,
|
||||
expected_stdout: str,
|
||||
expected_stderr: str,
|
||||
):
|
||||
_register_models_answer("sk-stale-key", proxy_answer)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.client.cli.commands.auth.load_cli_token",
|
||||
return_value=CliTokenRecord(
|
||||
base_url="http://localhost:4000",
|
||||
key="sk-stale-key",
|
||||
timestamp=old_timestamp,
|
||||
),
|
||||
with patch(
|
||||
"litellm.proxy.client.cli.commands.auth.load_cli_token",
|
||||
return_value=CliTokenRecord(
|
||||
base_url=STALE_TOKEN_BASE_URL,
|
||||
key="sk-stale-key",
|
||||
timestamp=STALE_TOKEN_TIMESTAMP,
|
||||
),
|
||||
patch("requests.post") as mock_post,
|
||||
):
|
||||
result = self.runner.invoke(print_token, obj={})
|
||||
result: Final = self.runner.invoke(print_token, obj={})
|
||||
|
||||
assert result.exit_code != 0
|
||||
assert "sk-stale-key" not in result.output
|
||||
assert "lite login" in result.output
|
||||
mock_post.assert_not_called()
|
||||
assert result.exit_code == expected_exit_code
|
||||
assert result.stdout == expected_stdout
|
||||
assert result.stderr == expected_stderr
|
||||
assert len(responses.calls) == 1
|
||||
|
||||
@responses.activate
|
||||
def test_proxy_check_waits_only_a_few_seconds_for_a_silent_proxy(self):
|
||||
_register_models_answer("sk-stale-key", requests.ReadTimeout("proxy sent nothing"))
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.client.cli.commands.auth.load_cli_token",
|
||||
return_value=CliTokenRecord(
|
||||
base_url=STALE_TOKEN_BASE_URL,
|
||||
key="sk-stale-key",
|
||||
timestamp=STALE_TOKEN_TIMESTAMP,
|
||||
),
|
||||
):
|
||||
result: Final = self.runner.invoke(print_token, obj={})
|
||||
|
||||
timeout: Final = responses.calls[0].request.req_kwargs["timeout"]
|
||||
assert result.exit_code == 0
|
||||
assert result.stdout == "sk-stale-key\n"
|
||||
assert isinstance(timeout, int | float)
|
||||
assert 0 < timeout <= 5
|
||||
|
||||
@responses.activate
|
||||
def test_key_without_a_recorded_proxy_still_expires_after_the_client_window(self):
|
||||
with patch(
|
||||
"litellm.proxy.client.cli.commands.auth.load_cli_token",
|
||||
return_value=CliTokenRecord(key="sk-old-format", timestamp=STALE_TOKEN_TIMESTAMP),
|
||||
):
|
||||
result: Final = self.runner.invoke(print_token, obj={})
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert result.stdout == ""
|
||||
assert result.stderr == "Token expired. Run 'lite login' again.\n"
|
||||
assert len(responses.calls) == 0
|
||||
|
||||
|
||||
class TestPrintTokenWithConfigFile:
|
||||
|
|
@ -2019,18 +2087,27 @@ class TestPkcePrintToken:
|
|||
assert sum(len(session.posts) for session in _FakeSession.instances) == 1
|
||||
assert save.call_count == 1
|
||||
|
||||
def test_print_token_for_an_expired_classic_token_makes_no_request(self):
|
||||
@responses.activate
|
||||
def test_print_token_trusts_a_recorded_expiry_without_asking_the_proxy(self):
|
||||
_register_models_answer("sk-classic", 200)
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.client.cli.commands.auth.load_token",
|
||||
return_value={"key": "sk-classic", "timestamp": time.time() - (CLI_JWT_EXPIRATION_HOURS + 1) * 3600},
|
||||
return_value={
|
||||
"key": "sk-classic",
|
||||
"base_url": STALE_TOKEN_BASE_URL,
|
||||
"timestamp": STALE_TOKEN_TIMESTAMP,
|
||||
"expires_at": STALE_TOKEN_TIMESTAMP + 3600,
|
||||
},
|
||||
),
|
||||
patch("litellm.proxy.client.cli.commands.auth.requests.Session", _FakeSession),
|
||||
):
|
||||
result = self.runner.invoke(print_token, obj={})
|
||||
result: Final = self.runner.invoke(print_token, obj={})
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert "Token expired" in result.output
|
||||
assert result.stdout == ""
|
||||
assert result.stderr == "Token expired. Run 'lite login' again.\n"
|
||||
assert len(responses.calls) == 0
|
||||
assert _FakeSession.instances == []
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import itertools
|
||||
import json
|
||||
from typing import Final, List, Optional
|
||||
from urllib.parse import quote
|
||||
|
|
@ -12,6 +13,7 @@ from pytest_mock import MockerFixture
|
|||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.constants import ANTHROPIC_DEFAULT_FILE_LIST_LIMIT
|
||||
from litellm.files.types import FileContentStreamingResult
|
||||
from litellm.proxy._types import LiteLLM_UserTableFiltered, LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.hooks import get_proxy_hook
|
||||
|
|
@ -4969,6 +4971,8 @@ def _managed_file_row(
|
|||
stored_file_id: Optional[str] = None,
|
||||
created_by: str = "test-user",
|
||||
team_id: Optional[str] = None,
|
||||
model_mappings: Optional[dict[str, str]] = None,
|
||||
storage_backend: Optional[str] = None,
|
||||
) -> MagicMock:
|
||||
file_object = OpenAIFileObject(
|
||||
id=stored_file_id or unified_file_id,
|
||||
|
|
@ -4984,6 +4988,8 @@ def _managed_file_row(
|
|||
file_object=file_object.model_dump(),
|
||||
created_by=created_by,
|
||||
team_id=team_id,
|
||||
model_mappings=model_mappings or {},
|
||||
storage_backend=storage_backend,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -5010,16 +5016,49 @@ class _ManagedFileTableOverRows:
|
|||
return next(iter(self._owned_rows(where)), None)
|
||||
|
||||
async def find_many(self, where, take=None, order=None, cursor=None, skip=0):
|
||||
rows = self._owned_rows(where)
|
||||
if cursor is not None:
|
||||
start = next(
|
||||
index
|
||||
for index, row in enumerate(rows)
|
||||
if row.unified_file_id == cursor["unified_file_id"]
|
||||
)
|
||||
rows = rows[start + skip :]
|
||||
owned_rows: Final = self._owned_rows(where)
|
||||
oldest_first: Final = bool(order) and order[0].get("created_at") == "asc"
|
||||
ordered_rows: Final = owned_rows[::-1] if oldest_first else owned_rows
|
||||
start: Final = (
|
||||
0
|
||||
if cursor is None
|
||||
else skip
|
||||
+ next(index for index, row in enumerate(ordered_rows) if row.unified_file_id == cursor["unified_file_id"])
|
||||
)
|
||||
rows: Final = ordered_rows[start:]
|
||||
return rows if take is None else rows[:take]
|
||||
|
||||
async def query_raw(self, query: str, *params: object) -> list[dict[str, object]]:
|
||||
keys_index: Final = next(index for index, param in enumerate(params) if isinstance(param, list))
|
||||
owner_ids: Final = frozenset(params[:keys_index])
|
||||
model_mapping_keys: Final = frozenset(params[keys_index])
|
||||
cursor_id: Final = params[keys_index + 1] if len(params) == keys_index + 3 else None
|
||||
ordered_rows: Final = [
|
||||
row
|
||||
for row in (self.rows[::-1] if '"created_at" ASC' in query else self.rows)
|
||||
if not owner_ids or {row.created_by, row.team_id} & owner_ids
|
||||
]
|
||||
start: Final = (
|
||||
0
|
||||
if cursor_id is None
|
||||
else 1 + next(index for index, row in enumerate(ordered_rows) if row.unified_file_id == cursor_id)
|
||||
)
|
||||
listed_rows: Final = [
|
||||
row
|
||||
for row in ordered_rows[start:]
|
||||
if not model_mapping_keys.isdisjoint(row.model_mappings)
|
||||
or ('"storage_backend" IS NOT NULL' in query and row.storage_backend is not None)
|
||||
]
|
||||
return [
|
||||
{
|
||||
"unified_file_id": row.unified_file_id,
|
||||
"file_object": row.file_object,
|
||||
"model_mappings": row.model_mappings,
|
||||
"storage_backend": row.storage_backend,
|
||||
}
|
||||
for row in listed_rows[: params[-1]]
|
||||
]
|
||||
|
||||
|
||||
def _setup_unscoped_list_files_route_over_real_hook(
|
||||
mocker, monkeypatch, llm_router: Router, rows
|
||||
|
|
@ -5041,7 +5080,9 @@ def _setup_unscoped_list_files_route_over_real_hook(
|
|||
managed_files = PROXY_LiteLLMManagedFiles(
|
||||
internal_usage_cache=MagicMock(), prisma_client=MagicMock()
|
||||
)
|
||||
managed_files.prisma_client.db.litellm_managedfiletable = _ManagedFileTableOverRows(rows)
|
||||
managed_file_table = _ManagedFileTableOverRows(rows)
|
||||
managed_files.prisma_client.db.litellm_managedfiletable = managed_file_table
|
||||
managed_files.prisma_client.db.query_raw = managed_file_table.query_raw
|
||||
|
||||
proxy_logging_obj = setup_proxy_logging_object(monkeypatch, llm_router)
|
||||
proxy_logging_obj.proxy_hook_mapping["managed_files"] = managed_files
|
||||
|
|
@ -6263,3 +6304,513 @@ def test_create_file_passthrough_fails_closed_when_guardrails_would_scan_the_bat
|
|||
assert error["param"] == "passthrough"
|
||||
assert "guardrails" in error["message"]
|
||||
assert forwarded_calls == []
|
||||
|
||||
|
||||
def _callers_files_on_every_deployment() -> list[MagicMock]:
|
||||
return [
|
||||
_managed_file_row("unified-gpt-newer", model_mappings={"gpt-3.5-turbo-id": "file-gpt-newer"}),
|
||||
_managed_file_row("unified-azure", model_mappings={"azure-gpt-3-5-turbo-id": "file-azure"}),
|
||||
_managed_file_row("unified-gemini", model_mappings={"gemini-2.0-flash-id": "file-gemini"}),
|
||||
_managed_file_row(
|
||||
"unified-theirs", created_by="other-user", model_mappings={"gpt-3.5-turbo-id": "file-theirs"}
|
||||
),
|
||||
_managed_file_row("unified-gpt-older", model_mappings={"gpt-3.5-turbo-id": "file-gpt-older"}),
|
||||
]
|
||||
|
||||
|
||||
def _list_files_as_test_user(path: str, params: dict[str, str], headers: dict[str, str]) -> httpx.Response:
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
||||
try:
|
||||
return client.get(path, params=params, headers={"Authorization": "Bearer test-key", **headers})
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("path", "params", "headers", "expected_ids"),
|
||||
[
|
||||
("/openai/v1/files", {}, {}, ["unified-gpt-newer", "unified-gpt-older"]),
|
||||
("/v1/files", {"provider": "azure"}, {}, ["unified-azure"]),
|
||||
("/v1/files", {}, {"custom-llm-provider": "gemini"}, ["unified-gemini"]),
|
||||
("/v1/files", {"target_model_names": "gemini-2.0-flash"}, {}, ["unified-gemini"]),
|
||||
("/v1/files", {}, {"x-litellm-model": "gpt-3.5-turbo"}, ["unified-gpt-newer", "unified-gpt-older"]),
|
||||
("/v1/files", {}, {}, ["unified-gpt-newer", "unified-azure", "unified-gemini", "unified-gpt-older"]),
|
||||
],
|
||||
ids=["provider-path", "provider-query", "provider-header", "target-model-names", "model-header", "no-provider"],
|
||||
)
|
||||
def test_require_managed_files_lists_the_callers_files_on_the_named_provider_or_model(
|
||||
mocker: MockerFixture,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
llm_router: Router,
|
||||
path: str,
|
||||
params: dict[str, str],
|
||||
headers: dict[str, str],
|
||||
expected_ids: list[str],
|
||||
) -> None:
|
||||
_, provider_list = _setup_unscoped_list_files_route_over_real_hook(
|
||||
mocker, monkeypatch, llm_router, _callers_files_on_every_deployment()
|
||||
)
|
||||
monkeypatch.setattr("litellm.require_managed_files", True)
|
||||
|
||||
response: Final = _list_files_as_test_user(path, params, headers)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert [file["id"] for file in response.json()["data"]] == expected_ids
|
||||
provider_list.assert_not_awaited()
|
||||
|
||||
|
||||
def test_require_managed_files_lists_a_teammates_files_on_the_provider_and_no_other_teams(
|
||||
mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, llm_router: Router
|
||||
) -> None:
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
_setup_unscoped_list_files_route_over_real_hook(
|
||||
mocker,
|
||||
monkeypatch,
|
||||
llm_router,
|
||||
[
|
||||
_managed_file_row(
|
||||
"unified-teammate", created_by="teammate", team_id="team-a", model_mappings={"gpt-3.5-turbo-id": "f-1"}
|
||||
),
|
||||
_managed_file_row("unified-own", model_mappings={"gpt-3.5-turbo-id": "f-2"}),
|
||||
_managed_file_row(
|
||||
"unified-other-team", created_by="outsider", team_id="team-b", model_mappings={"gpt-3.5-turbo-id": "f-3"}
|
||||
),
|
||||
],
|
||||
)
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
api_key="test-key", user_role=LitellmUserRoles.INTERNAL_USER, user_id="test-user", team_id="team-a"
|
||||
)
|
||||
monkeypatch.setattr("litellm.require_managed_files", True)
|
||||
|
||||
response: Final = _list_files_as_test_user("/openai/v1/files", {}, {})
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert [file["id"] for file in response.json()["data"]] == ["unified-teammate", "unified-own"]
|
||||
|
||||
|
||||
def _router_with_wildcard_and_unprefixed_anthropic_deployments() -> Router:
|
||||
unprefixed_anthropic_model: Final = next(model for model in sorted(litellm.anthropic_models) if "/" not in model)
|
||||
return Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {"model": "openai/gpt-3.5-turbo", "api_key": "openai_api_key"},
|
||||
"model_info": {"id": "gpt-3.5-turbo-id"},
|
||||
},
|
||||
{
|
||||
"model_name": "anthropic/*",
|
||||
"litellm_params": {"model": "anthropic/*", "api_key": "anthropic_api_key"},
|
||||
"model_info": {"id": "anthropic-wildcard-id"},
|
||||
},
|
||||
{
|
||||
"model_name": "claude-unprefixed",
|
||||
"litellm_params": {"model": unprefixed_anthropic_model, "api_key": "anthropic_api_key"},
|
||||
"model_info": {"id": "claude-unprefixed-id"},
|
||||
},
|
||||
],
|
||||
model_group_alias={"claude-alias": "claude-unprefixed"},
|
||||
routing_groups=[
|
||||
{"group_name": "claude-group", "models": ["claude-unprefixed"], "routing_strategy": "simple-shuffle"}
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def _callers_files_on_wildcard_unprefixed_and_passthrough_deployments() -> list[MagicMock]:
|
||||
return [
|
||||
_managed_file_row("unified-wildcard", model_mappings={"anthropic-wildcard-id": "file-wildcard"}),
|
||||
_managed_file_row("unified-unprefixed", model_mappings={"claude-unprefixed-id": "file-unprefixed"}),
|
||||
_managed_file_row("unified-gpt", model_mappings={"gpt-3.5-turbo-id": "file-gpt"}),
|
||||
_managed_file_row(
|
||||
"unified-passthrough",
|
||||
model_mappings={
|
||||
"_passthrough_openai": "file-passthrough",
|
||||
"_passthrough_provider_marker": "_passthrough_provider:openai",
|
||||
},
|
||||
),
|
||||
_managed_file_row(
|
||||
"unified-storage",
|
||||
model_mappings={"claude-unprefixed": "litellm_db://unified-storage"},
|
||||
storage_backend="litellm_db",
|
||||
),
|
||||
_managed_file_row(
|
||||
"unified-wildcard-storage",
|
||||
model_mappings={"anthropic/claude-haiku-5-5": "litellm_db://unified-wildcard-storage"},
|
||||
storage_backend="litellm_db",
|
||||
),
|
||||
_managed_file_row(
|
||||
"unified-alias-storage",
|
||||
model_mappings={"claude-alias": "s3://bucket/unified-alias-storage"},
|
||||
storage_backend="s3",
|
||||
),
|
||||
_managed_file_row(
|
||||
"unified-group-storage",
|
||||
model_mappings={"claude-group": "s3://bucket/unified-group-storage"},
|
||||
storage_backend="s3",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("path", "params", "headers", "expected_ids"),
|
||||
[
|
||||
(
|
||||
"/v1/files",
|
||||
{"target_model_names": "anthropic/claude-haiku-5-5"},
|
||||
{},
|
||||
["unified-wildcard", "unified-wildcard-storage"],
|
||||
),
|
||||
(
|
||||
"/v1/files",
|
||||
{},
|
||||
{"x-litellm-model": "anthropic/claude-haiku-5-5"},
|
||||
["unified-wildcard", "unified-wildcard-storage"],
|
||||
),
|
||||
(
|
||||
"/anthropic/v1/files",
|
||||
{},
|
||||
{},
|
||||
[
|
||||
"unified-wildcard",
|
||||
"unified-unprefixed",
|
||||
"unified-storage",
|
||||
"unified-wildcard-storage",
|
||||
"unified-alias-storage",
|
||||
"unified-group-storage",
|
||||
],
|
||||
),
|
||||
("/openai/v1/files", {}, {}, ["unified-gpt", "unified-passthrough"]),
|
||||
("/v1/files", {"target_model_names": "claude-unprefixed"}, {}, ["unified-unprefixed", "unified-storage"]),
|
||||
],
|
||||
ids=[
|
||||
"wildcard-target-model-names",
|
||||
"wildcard-model-header",
|
||||
"unprefixed-provider",
|
||||
"passthrough-provider",
|
||||
"storage-backed-model",
|
||||
],
|
||||
)
|
||||
def test_require_managed_files_lists_the_callers_files_wherever_the_router_put_them(
|
||||
mocker: MockerFixture,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
path: str,
|
||||
params: dict[str, str],
|
||||
headers: dict[str, str],
|
||||
expected_ids: list[str],
|
||||
) -> None:
|
||||
_, provider_list = _setup_unscoped_list_files_route_over_real_hook(
|
||||
mocker,
|
||||
monkeypatch,
|
||||
_router_with_wildcard_and_unprefixed_anthropic_deployments(),
|
||||
_callers_files_on_wildcard_unprefixed_and_passthrough_deployments(),
|
||||
)
|
||||
monkeypatch.setattr("litellm.require_managed_files", True)
|
||||
|
||||
response: Final = _list_files_as_test_user(path, params, headers)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert [file["id"] for file in response.json()["data"]] == expected_ids
|
||||
provider_list.assert_not_awaited()
|
||||
|
||||
|
||||
def _router_with_a_catch_all_openai_deployment() -> Router:
|
||||
unprefixed_anthropic_model: Final = next(model for model in sorted(litellm.anthropic_models) if "/" not in model)
|
||||
return Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {"model": "openai/gpt-3.5-turbo", "api_key": "openai_api_key"},
|
||||
"model_info": {"id": "gpt-3.5-turbo-id"},
|
||||
},
|
||||
{
|
||||
"model_name": "*",
|
||||
"litellm_params": {"model": "openai/*", "api_key": "openai_api_key"},
|
||||
"model_info": {"id": "catch-all-id"},
|
||||
},
|
||||
{
|
||||
"model_name": "claude-unprefixed",
|
||||
"litellm_params": {"model": unprefixed_anthropic_model, "api_key": "anthropic_api_key"},
|
||||
"model_info": {"id": "claude-unprefixed-id"},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("path", "expected_ids"),
|
||||
[
|
||||
("/openai/v1/files", ["unified-gpt", "unified-catch-all-storage"]),
|
||||
("/anthropic/v1/files", ["unified-unprefixed", "unified-storage"]),
|
||||
],
|
||||
ids=["catch-all-provider", "explicit-provider"],
|
||||
)
|
||||
def test_require_managed_files_keeps_a_catch_all_route_from_claiming_files_on_other_providers(
|
||||
mocker: MockerFixture,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
path: str,
|
||||
expected_ids: list[str],
|
||||
) -> None:
|
||||
_, provider_list = _setup_unscoped_list_files_route_over_real_hook(
|
||||
mocker,
|
||||
monkeypatch,
|
||||
_router_with_a_catch_all_openai_deployment(),
|
||||
[
|
||||
_managed_file_row("unified-unprefixed", model_mappings={"claude-unprefixed-id": "file-unprefixed"}),
|
||||
_managed_file_row("unified-gpt", model_mappings={"gpt-3.5-turbo-id": "file-gpt"}),
|
||||
_managed_file_row(
|
||||
"unified-storage",
|
||||
model_mappings={"claude-unprefixed": "litellm_db://unified-storage"},
|
||||
storage_backend="litellm_db",
|
||||
),
|
||||
_managed_file_row(
|
||||
"unified-catch-all-storage",
|
||||
model_mappings={"model-only-the-catch-all-serves": "litellm_db://unified-catch-all-storage"},
|
||||
storage_backend="litellm_db",
|
||||
),
|
||||
],
|
||||
)
|
||||
monkeypatch.setattr("litellm.require_managed_files", True)
|
||||
|
||||
response: Final = _list_files_as_test_user(path, {}, {})
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert [file["id"] for file in response.json()["data"]] == expected_ids
|
||||
provider_list.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("path", "params"),
|
||||
[
|
||||
("/antropic/v1/files", {}),
|
||||
("/openai/v1/files", {"after": "unified-gpt", "after_id": "unified-wildcard"}),
|
||||
("/v1/files", {"target_model_names": "gpt-3.5-turbo,claude-unprefixed"}),
|
||||
],
|
||||
ids=["unknown-provider", "conflicting-after-cursors", "two-target-model-names"],
|
||||
)
|
||||
def test_require_managed_files_rejects_a_list_it_cannot_scope(
|
||||
mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, path: str, params: dict[str, str]
|
||||
) -> None:
|
||||
_, provider_list = _setup_unscoped_list_files_route_over_real_hook(
|
||||
mocker,
|
||||
monkeypatch,
|
||||
_router_with_wildcard_and_unprefixed_anthropic_deployments(),
|
||||
_callers_files_on_wildcard_unprefixed_and_passthrough_deployments(),
|
||||
)
|
||||
monkeypatch.setattr("litellm.require_managed_files", True)
|
||||
|
||||
response: Final = _list_files_as_test_user(path, params, {})
|
||||
|
||||
assert response.status_code == 400, response.text
|
||||
provider_list.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("cursor_param", ["after", "after_id"])
|
||||
def test_require_managed_files_pages_a_provider_list_with_either_cursor_param(
|
||||
mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, llm_router: Router, cursor_param: str
|
||||
) -> None:
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
||||
_setup_unscoped_list_files_route_over_real_hook(
|
||||
mocker, monkeypatch, llm_router, _callers_files_on_every_deployment()
|
||||
)
|
||||
monkeypatch.setattr("litellm.require_managed_files", True)
|
||||
|
||||
try:
|
||||
first_page: Final = client.get(
|
||||
"/openai/v1/files", params={"limit": "1"}, headers={"Authorization": "Bearer test-key"}
|
||||
)
|
||||
second_page: Final = client.get(
|
||||
"/openai/v1/files",
|
||||
params={"limit": "1", cursor_param: first_page.json()["last_id"]},
|
||||
headers={"Authorization": "Bearer test-key"},
|
||||
)
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
assert (first_page.status_code, second_page.status_code) == (200, 200), (first_page.text, second_page.text)
|
||||
assert [file["id"] for file in first_page.json()["data"]] == ["unified-gpt-newer"]
|
||||
assert first_page.json()["has_more"] is True
|
||||
assert [file["id"] for file in second_page.json()["data"]] == ["unified-gpt-older"]
|
||||
assert second_page.json()["has_more"] is False
|
||||
|
||||
|
||||
def test_require_managed_files_pages_toward_newer_files_with_before_id(
|
||||
mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, llm_router: Router
|
||||
) -> None:
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
||||
_setup_unscoped_list_files_route_over_real_hook(
|
||||
mocker, monkeypatch, llm_router, _callers_files_on_every_deployment()
|
||||
)
|
||||
monkeypatch.setattr("litellm.require_managed_files", True)
|
||||
|
||||
try:
|
||||
page_before_oldest: Final = client.get(
|
||||
"/v1/files",
|
||||
params={"limit": "2", "before_id": "unified-gpt-older"},
|
||||
headers={"Authorization": "Bearer test-key"},
|
||||
)
|
||||
newest_page: Final = client.get(
|
||||
"/v1/files",
|
||||
params={"limit": "2", "before_id": page_before_oldest.json()["first_id"]},
|
||||
headers={"Authorization": "Bearer test-key"},
|
||||
)
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
assert (page_before_oldest.status_code, newest_page.status_code) == (200, 200), (
|
||||
page_before_oldest.text,
|
||||
newest_page.text,
|
||||
)
|
||||
assert [file["id"] for file in page_before_oldest.json()["data"]] == ["unified-azure", "unified-gemini"]
|
||||
assert page_before_oldest.json()["has_more"] is True
|
||||
assert [file["id"] for file in newest_page.json()["data"]] == ["unified-gpt-newer"]
|
||||
assert newest_page.json()["has_more"] is False
|
||||
|
||||
|
||||
def test_managed_file_list_rejects_both_cursor_directions(
|
||||
mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, llm_router: Router
|
||||
) -> None:
|
||||
_setup_unscoped_list_files_route_over_real_hook(
|
||||
mocker, monkeypatch, llm_router, _callers_files_on_every_deployment()
|
||||
)
|
||||
monkeypatch.setattr("litellm.require_managed_files", True)
|
||||
|
||||
response: Final = _list_files_as_test_user(
|
||||
"/v1/files", {"after": "unified-gpt-newer", "before_id": "unified-gpt-older"}, {}
|
||||
)
|
||||
|
||||
assert response.status_code == 400, response.text
|
||||
assert "not both" in response.json()["error"]["message"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("params", "headers"),
|
||||
[({"target_model_names": "gemini-2.0-flash"}, {}), ({}, {"x-litellm-model": "gemini-2.0-flash"})],
|
||||
ids=["target-model-names", "model-header"],
|
||||
)
|
||||
def test_require_managed_files_list_rejects_a_model_the_key_cannot_use(
|
||||
mocker: MockerFixture,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
llm_router: Router,
|
||||
params: dict[str, str],
|
||||
headers: dict[str, str],
|
||||
) -> None:
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
||||
_setup_unscoped_list_files_route_over_real_hook(
|
||||
mocker, monkeypatch, llm_router, _callers_files_on_every_deployment()
|
||||
)
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
user_id="test-user",
|
||||
models=["gpt-3.5-turbo"],
|
||||
)
|
||||
monkeypatch.setattr("litellm.require_managed_files", True)
|
||||
|
||||
response: Final = _list_files_as_test_user("/v1/files", params, headers)
|
||||
|
||||
assert response.status_code == 403, response.text
|
||||
assert "unified-gemini" not in response.text
|
||||
|
||||
|
||||
def test_require_managed_files_list_fails_closed_without_the_managed_files_hook(
|
||||
mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, llm_router: Router
|
||||
) -> None:
|
||||
proxy_logging_obj, provider_list = _setup_unscoped_list_files_route_over_real_hook(
|
||||
mocker, monkeypatch, llm_router, _callers_files_on_every_deployment()
|
||||
)
|
||||
proxy_logging_obj.proxy_hook_mapping.pop("managed_files")
|
||||
provider_list.return_value = [_managed_file("file-on-the-shared-key")]
|
||||
monkeypatch.setattr("litellm.require_managed_files", True)
|
||||
|
||||
response: Final = _list_files_as_test_user("/openai/v1/files", {}, {})
|
||||
|
||||
assert response.status_code == 500, response.text
|
||||
assert "file-on-the-shared-key" not in response.text
|
||||
assert "require_managed_files cannot scope the list to the caller" in response.json()["error"]["message"]
|
||||
provider_list.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("hook_loaded", [True, False], ids=["hook-loaded", "no-hook"])
|
||||
def test_require_managed_files_keeps_the_provider_list_for_an_admin(
|
||||
mocker: MockerFixture,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
llm_router: Router,
|
||||
hook_loaded: bool,
|
||||
) -> None:
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
||||
proxy_logging_obj, provider_list = _setup_unscoped_list_files_route_over_real_hook(
|
||||
mocker, monkeypatch, llm_router, _callers_files_on_every_deployment()
|
||||
)
|
||||
if not hook_loaded:
|
||||
proxy_logging_obj.proxy_hook_mapping.pop("managed_files")
|
||||
provider_list.return_value = [_managed_file("file-on-the-shared-key")]
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
api_key="test-key", user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user"
|
||||
)
|
||||
monkeypatch.setattr("litellm.require_managed_files", True)
|
||||
|
||||
response: Final = _list_files_as_test_user("/anthropic/v1/files", {}, {})
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert [file["id"] for file in response.json()["data"]] == ["file-on-the-shared-key"]
|
||||
|
||||
|
||||
def test_managed_file_list_without_require_managed_files_ignores_cursor_params_openai_does_not_define(
|
||||
mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, llm_router: Router
|
||||
) -> None:
|
||||
_setup_unscoped_list_files_route_over_real_hook(
|
||||
mocker, monkeypatch, llm_router, _callers_files_on_every_deployment()
|
||||
)
|
||||
monkeypatch.setattr("litellm.require_managed_files", False)
|
||||
|
||||
response: Final = _list_files_as_test_user(
|
||||
"/v1/files", {"after": "unified-gpt-newer", "after_id": "unified-gemini", "before_id": "file-raw"}, {}
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert [file["id"] for file in response.json()["data"]] == ["unified-azure", "unified-gemini", "unified-gpt-older"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("path", "headers", "expected_count", "has_more"),
|
||||
[
|
||||
("/anthropic/v1/files", {}, ANTHROPIC_DEFAULT_FILE_LIST_LIMIT, True),
|
||||
("/v1/files", {"x-litellm-model": "claude-unprefixed"}, ANTHROPIC_DEFAULT_FILE_LIST_LIMIT, True),
|
||||
("/anthropic/v1/files", {"x-litellm-model": "gpt-3.5-turbo"}, ANTHROPIC_DEFAULT_FILE_LIST_LIMIT + 1, False),
|
||||
("/v1/files", {}, 2 * (ANTHROPIC_DEFAULT_FILE_LIST_LIMIT + 1), False),
|
||||
],
|
||||
ids=["anthropic-provider", "anthropic-model-header", "openai-model-header", "every-provider"],
|
||||
)
|
||||
def test_require_managed_files_pages_an_anthropic_list_at_the_anthropic_default_size(
|
||||
mocker: MockerFixture,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
path: str,
|
||||
headers: dict[str, str],
|
||||
expected_count: int,
|
||||
has_more: bool,
|
||||
) -> None:
|
||||
_setup_unscoped_list_files_route_over_real_hook(
|
||||
mocker,
|
||||
monkeypatch,
|
||||
_router_with_wildcard_and_unprefixed_anthropic_deployments(),
|
||||
[
|
||||
_managed_file_row(
|
||||
f"unified-{deployment_id}-{index:03d}", model_mappings={deployment_id: f"file-{index:03d}"}
|
||||
)
|
||||
for deployment_id, index in itertools.product(
|
||||
("claude-unprefixed-id", "gpt-3.5-turbo-id"), range(ANTHROPIC_DEFAULT_FILE_LIST_LIMIT + 1)
|
||||
)
|
||||
],
|
||||
)
|
||||
monkeypatch.setattr("litellm.require_managed_files", True)
|
||||
|
||||
response: Final = _list_files_as_test_user(path, {}, headers)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert (len(response.json()["data"]), response.json()["has_more"]) == (expected_count, has_more)
|
||||
|
|
|
|||
|
|
@ -488,22 +488,17 @@ async def test_otel_unhandled_exception_handler_reraises_http_exception_invalid(
|
|||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("media_type", ["application/json", "application/x-protobuf"])
|
||||
@pytest.mark.parametrize("root_path", ["", "/tenant-a"])
|
||||
@pytest.mark.parametrize("native_available", [True, False])
|
||||
@pytest.mark.parametrize("error", [
|
||||
ProxyException("database credentials: secret", "auth_error", None, 401),
|
||||
HTTPException(403, "database credentials: secret"),
|
||||
])
|
||||
async def test_otlp_auth_errors_hide_internal_details_and_survive_missing_native(
|
||||
media_type: str, root_path: str, native_available: bool,
|
||||
error: ProxyException | HTTPException, monkeypatch: pytest.MonkeyPatch,
|
||||
async def test_otlp_auth_errors_hide_internal_details(
|
||||
media_type: str, root_path: str, error: ProxyException | HTTPException,
|
||||
) -> None:
|
||||
from google.rpc.status_pb2 import Status
|
||||
|
||||
from litellm.proxy.proxy_server import otlp_http_exception_handler
|
||||
from litellm.rust_bridge import loader
|
||||
|
||||
if not native_available:
|
||||
monkeypatch.setattr(loader, "_cached_bridge", None)
|
||||
request: Final = Request({
|
||||
"type": "http", "method": "POST", "path": root_path + "/v1/traces", "root_path": root_path,
|
||||
"headers": [(b"content-type", media_type.encode())],
|
||||
|
|
@ -522,7 +517,7 @@ async def test_otlp_auth_errors_hide_internal_details_and_survive_missing_native
|
|||
else Status.FromString(response.body).message
|
||||
)
|
||||
expected: Final = "Unauthorized" if isinstance(error, ProxyException) else "Forbidden"
|
||||
assert message == (expected if native_available or media_type == "application/json" else "")
|
||||
assert message == expected
|
||||
|
||||
|
||||
def _websocket(sent: list[Message]) -> WebSocket:
|
||||
|
|
|
|||
|
|
@ -10826,12 +10826,19 @@ class TestResolveLitellmCallId:
|
|||
class _FailureHookRecorder:
|
||||
"""Stands in for ProxyLogging.post_call_failure_hook, recording what the detached-failure closure hands it."""
|
||||
|
||||
def __init__(self, raises: Optional[Exception] = None):
|
||||
def __init__(self, raises: Optional[Exception] = None, gate: Optional[asyncio.Event] = None):
|
||||
self.calls = []
|
||||
self.entered = asyncio.Event()
|
||||
self.reported = asyncio.Event()
|
||||
self._raises = raises
|
||||
self._gate = gate
|
||||
|
||||
async def post_call_failure_hook(self, **kwargs):
|
||||
self.entered.set()
|
||||
if self._gate is not None:
|
||||
await self._gate.wait()
|
||||
self.calls.append(kwargs)
|
||||
self.reported.set()
|
||||
if self._raises is not None:
|
||||
raise self._raises
|
||||
|
||||
|
|
@ -10932,6 +10939,171 @@ class TestDetachedStreamFailureHook:
|
|||
|
||||
assert [call["original_exception"] for call in recorder.calls] == [failure]
|
||||
|
||||
@staticmethod
|
||||
def _armed_stream(recorder: _FailureHookRecorder, iterator_hook_for, response: Optional[MagicMock] = None):
|
||||
from litellm.proxy._types import UserAPIKeyAuth as RealUserAPIKeyAuth
|
||||
|
||||
logging_obj = LiteLLMLoggingObj(
|
||||
model="gpt-6.1-sol",
|
||||
messages=[{"role": "user", "content": "Write a story."}],
|
||||
stream=True,
|
||||
call_type="anthropic_messages",
|
||||
start_time=None,
|
||||
litellm_call_id="call-lit6871",
|
||||
function_id="function-lit6871",
|
||||
)
|
||||
user_api_key_dict = RealUserAPIKeyAuth(api_key="sk-test")
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(data={"litellm_logging_obj": logging_obj})
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache=MagicMock())
|
||||
proxy_logging_obj.post_call_failure_hook = recorder.post_call_failure_hook
|
||||
proxy_logging_obj.async_post_call_streaming_iterator_hook = iterator_hook_for(logging_obj)
|
||||
processing_obj._arm_detached_stream_failure_hook(
|
||||
logging_obj=logging_obj,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
stream = ProxyBaseLLMRequestProcessing.async_streaming_data_generator(
|
||||
response=MagicMock() if response is None else response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_data=processing_obj.data,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
serialize_chunk=lambda chunk: "data: chunk\n\n",
|
||||
serialize_error=lambda error: "data: error\n\n",
|
||||
)
|
||||
return logging_obj, stream
|
||||
|
||||
@staticmethod
|
||||
def _hook_reporting(upstream_failure: Exception, after_report=None):
|
||||
def iterator_hook_for(logging_obj):
|
||||
async def iterator_hook(**_kwargs):
|
||||
yield b"event: message_start\n\n"
|
||||
await logging_obj.on_detached_stream_failure(upstream_failure)
|
||||
yield b"event: error\n\n"
|
||||
yield b"event: message_stop\n\n"
|
||||
if after_report is not None:
|
||||
raise after_report
|
||||
|
||||
return iterator_hook
|
||||
|
||||
return iterator_hook_for
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_iterator_hook_failure_after_a_bridged_failure_is_the_one_the_failure_hook_gets(self):
|
||||
recorder = _FailureHookRecorder()
|
||||
upstream_failure = RuntimeError("upstream Responses stream failed")
|
||||
hook_failure = RuntimeError("iterator hook rejected the stream")
|
||||
|
||||
_, stream = self._armed_stream(recorder, self._hook_reporting(upstream_failure, after_report=hook_failure))
|
||||
|
||||
frames = [frame async for frame in stream]
|
||||
|
||||
assert frames == ["data: chunk\n\n"] * 3 + ["data: error\n\n"]
|
||||
assert [call["original_exception"] for call in recorder.calls] == [hook_failure]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failure_hook_replacing_the_iterator_hook_failure_aborts_the_stream(self):
|
||||
replacement = HTTPException(status_code=418, detail="rejected by a custom failure hook")
|
||||
recorder = _FailureHookRecorder(raises=replacement)
|
||||
upstream_failure = RuntimeError("upstream Responses stream failed")
|
||||
hook_failure = RuntimeError("iterator hook rejected the stream")
|
||||
|
||||
_, stream = self._armed_stream(recorder, self._hook_reporting(upstream_failure, after_report=hook_failure))
|
||||
|
||||
frames = [await stream.__anext__() for _ in range(3)]
|
||||
with pytest.raises(HTTPException) as raised:
|
||||
await stream.__anext__()
|
||||
|
||||
assert raised.value is replacement
|
||||
assert frames == ["data: chunk\n\n"] * 3
|
||||
assert [call["original_exception"] for call in recorder.calls] == [hook_failure]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failure_reported_by_the_stream_reaches_the_failure_hook_once_after_the_stream_ends(self):
|
||||
recorder = _FailureHookRecorder()
|
||||
upstream_failure = RuntimeError("upstream Responses stream failed")
|
||||
|
||||
_, stream = self._armed_stream(recorder, self._hook_reporting(upstream_failure))
|
||||
|
||||
frames = []
|
||||
async for frame in stream:
|
||||
frames.append(frame)
|
||||
assert recorder.calls == []
|
||||
|
||||
assert frames == ["data: chunk\n\n"] * 3
|
||||
assert [call["original_exception"] for call in recorder.calls] == [upstream_failure]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failure_report_and_stream_cleanup_survive_the_client_task_being_cancelled_mid_report(self):
|
||||
gate = asyncio.Event()
|
||||
recorder = _FailureHookRecorder(gate=gate)
|
||||
upstream_failure = RuntimeError("upstream Responses stream failed")
|
||||
response = MagicMock()
|
||||
response.aclose = AsyncMock()
|
||||
|
||||
_, stream = self._armed_stream(recorder, self._hook_reporting(upstream_failure), response=response)
|
||||
|
||||
async def client_reads_the_whole_stream():
|
||||
return [frame async for frame in stream]
|
||||
|
||||
client = asyncio.create_task(client_reads_the_whole_stream())
|
||||
await asyncio.wait_for(recorder.entered.wait(), timeout=5)
|
||||
client.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await client
|
||||
|
||||
response.aclose.assert_awaited_once()
|
||||
gate.set()
|
||||
await asyncio.wait_for(recorder.reported.wait(), timeout=5)
|
||||
assert [call["original_exception"] for call in recorder.calls] == [upstream_failure]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failure_reported_by_the_stream_reaches_the_failure_hook_once_when_the_client_leaves_early(self):
|
||||
recorder = _FailureHookRecorder()
|
||||
upstream_failure = RuntimeError("upstream Responses stream failed")
|
||||
|
||||
_, stream = self._armed_stream(recorder, self._hook_reporting(upstream_failure))
|
||||
|
||||
await stream.__anext__()
|
||||
await stream.__anext__()
|
||||
await stream.aclose()
|
||||
|
||||
assert [call["original_exception"] for call in recorder.calls] == [upstream_failure]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failure_reported_after_the_stream_ends_still_reaches_the_failure_hook(self):
|
||||
recorder = _FailureHookRecorder()
|
||||
upstream_failure = RuntimeError("upstream Responses stream failed after the client stream closed")
|
||||
|
||||
async def iterator_hook(**_kwargs):
|
||||
yield b"event: message_start\n\n"
|
||||
|
||||
logging_obj, stream = self._armed_stream(recorder, lambda _logging_obj: iterator_hook)
|
||||
|
||||
frames = [frame async for frame in stream]
|
||||
await logging_obj.on_detached_stream_failure(upstream_failure)
|
||||
|
||||
assert frames == ["data: chunk\n\n"]
|
||||
assert [call["original_exception"] for call in recorder.calls] == [upstream_failure]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failure_reported_by_an_iterator_hook_is_not_reported_again_when_the_stream_closes(self):
|
||||
recorder = _FailureHookRecorder()
|
||||
upstream_failure = RuntimeError("upstream Responses stream failed")
|
||||
hook_failure = RuntimeError("iterator hook rejected the error frame")
|
||||
|
||||
async def iterator_hook(**_kwargs):
|
||||
yield b"event: message_start\n\n"
|
||||
raise hook_failure
|
||||
|
||||
logging_obj, stream = self._armed_stream(recorder, lambda _logging_obj: iterator_hook)
|
||||
|
||||
frames = [frame async for frame in stream]
|
||||
await logging_obj.on_detached_stream_failure(upstream_failure)
|
||||
|
||||
assert frames == ["data: chunk\n\n", "data: error\n\n"]
|
||||
assert [call["original_exception"] for call in recorder.calls] == [hook_failure]
|
||||
|
||||
|
||||
class TestPostCallMaskedOutputReachesDeferredLogging:
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -30,7 +30,6 @@ from litellm.proxy.auth.authorization import OwnedRows, ReadScope
|
|||
from litellm.proxy.auth.authorization_dependencies import get_log_team_lookup
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.tracing_runtime import manage_tracing, provide_storage
|
||||
from litellm.rust_bridge import loader
|
||||
from litellm.tracing.errors import TraceChanged
|
||||
from litellm.tracing.generated.models import TraceQueryHelp
|
||||
from litellm.tracing.generated.responses import TraceSQLResponse
|
||||
|
|
|
|||
|
|
@ -2668,6 +2668,65 @@ async def test_stream_auto_execute_bills_the_paid_first_round_on_its_success_row
|
|||
assert _request_failure_row(logging_obj, recorder.logged[0]).get("response_cost", 0.0) == 0.0
|
||||
|
||||
|
||||
class _CutAfterStreamedText(httpx.AsyncByteStream):
|
||||
async def __aiter__(self):
|
||||
opened: Final = {**json.loads(json.dumps(_ROUND_TWO_BODY, default=dict)), "status": "in_progress"}
|
||||
yield _sse(
|
||||
(
|
||||
{"type": "response.created", "response": opened, "sequence_number": 0},
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"item_id": "msg_round_two",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": "PING " * 40,
|
||||
"sequence_number": 1,
|
||||
},
|
||||
)
|
||||
)
|
||||
raise httpx.ReadError("connection reset mid-stream")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_auto_execute_bills_a_follow_up_cut_after_streamed_text_with_no_second_fixed_margin(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}})
|
||||
_serve_responses(
|
||||
respx_mock,
|
||||
monkeypatch,
|
||||
(_streamed(_ROUND_ONE_BODY), httpx.Response(200, stream=_CutAfterStreamedText(), headers=SSE_HEADERS)),
|
||||
)
|
||||
recorder: Final = _RoundRecorder()
|
||||
logging_obj: Final = _rounds_logging_obj(recorder, "mcp-stream-follow-up-cut-after-text")
|
||||
|
||||
with pytest.raises(httpx.ReadError):
|
||||
await _stream_auto_execute(logging_obj)
|
||||
ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj})
|
||||
await _settle(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0)
|
||||
|
||||
_assert_one_row_bills(recorder.logged, _ROUND_ONE_BODY)
|
||||
assert len(recorder.failed) == 1
|
||||
failed: Final = recorder.failed[0]
|
||||
assert failed["completion_tokens"] > 0
|
||||
priced_with_its_own_margin: Final = litellm.completion_cost(
|
||||
completion_response=ModelResponse(
|
||||
model="gpt-5.6",
|
||||
usage=Usage(
|
||||
prompt_tokens=failed["prompt_tokens"],
|
||||
completion_tokens=failed["completion_tokens"],
|
||||
total_tokens=failed["prompt_tokens"] + failed["completion_tokens"],
|
||||
),
|
||||
),
|
||||
model="gpt-5.6",
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
assert failed["response_cost"] == pytest.approx(priced_with_its_own_margin - _FIXED_MARGIN)
|
||||
_assert_breakdown_agrees_with_the_billed_cost(failed)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_auto_execute_blocked_by_a_post_call_guardrail_keeps_both_rounds_on_the_success_row(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
|
|
|
|||
|
|
@ -9,20 +9,23 @@ crashed with exit 139 whenever any CustomLogger was registered).
|
|||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
from collections.abc import AsyncIterator, Mapping
|
||||
from datetime import datetime
|
||||
from typing import Final, cast
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils import thread_pool_executor as thread_pool_executor_module
|
||||
from litellm.responses import streaming_iterator as responses_streaming_iterator_module
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging
|
||||
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
|
||||
from litellm.responses import streaming_iterator as responses_streaming_iterator_module
|
||||
from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse, ResponsesAPIStreamingResponse
|
||||
from litellm.types.utils import StandardLoggingPayload
|
||||
|
||||
|
||||
class RecordingCustomLogger(CustomLogger):
|
||||
|
|
@ -169,6 +172,253 @@ async def test_sync_callbacks_run_only_after_async_handler_completes(recording_e
|
|||
assert submit_times[0] >= recorder.async_hook_finished
|
||||
|
||||
|
||||
class FailureRecorder(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.failure_payloads: tuple[StandardLoggingPayload, ...] = ()
|
||||
self.success_payloads: tuple[StandardLoggingPayload, ...] = ()
|
||||
|
||||
async def async_log_failure_event(
|
||||
self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
self.failure_payloads = (
|
||||
*self.failure_payloads,
|
||||
cast(StandardLoggingPayload, kwargs["standard_logging_object"]),
|
||||
)
|
||||
|
||||
async def async_log_success_event(
|
||||
self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
self.success_payloads = (
|
||||
*self.success_payloads,
|
||||
cast(StandardLoggingPayload, kwargs["standard_logging_object"]),
|
||||
)
|
||||
|
||||
|
||||
class _ChunkRecorder:
|
||||
def __init__(self) -> None:
|
||||
self.chunks: tuple[ResponsesAPIStreamingResponse, ...] = ()
|
||||
|
||||
async def drain(self, iterator: ResponsesAPIStreamingIterator) -> None:
|
||||
async for chunk in iterator:
|
||||
self.chunks = (*self.chunks, chunk)
|
||||
|
||||
|
||||
_SEA_DELTAS = ("The sea ", "is wide ", "and deep.")
|
||||
_SEA_PROMPT = "Write a 300 word story about the sea."
|
||||
_UNREACHABLE_IMAGE_INPUT: list[dict[str, object]] = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "input_text", "text": "Describe this picture."},
|
||||
{"type": "input_image", "image_url": "http://images.invalid/photo.png", "detail": "high"},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
_TOOL_ARGUMENT_DELTAS = ('{"city": ', '"Lisbon", ', '"unit": "celsius"}')
|
||||
|
||||
|
||||
class _UpstreamThatTimesOutAfterThreeDeltas:
|
||||
headers: dict[str, str] = {}
|
||||
|
||||
def __init__(self, event_type: str = "response.output_text.delta", deltas: tuple[str, ...] = _SEA_DELTAS) -> None:
|
||||
self.event_type = event_type
|
||||
self.deltas = deltas
|
||||
|
||||
async def aiter_bytes(self) -> AsyncIterator[bytes]:
|
||||
for sequence_number, delta in enumerate(self.deltas):
|
||||
event = {
|
||||
"type": self.event_type,
|
||||
"item_id": "item_1",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": delta,
|
||||
"sequence_number": sequence_number,
|
||||
}
|
||||
yield f"data: {json.dumps(event)}\n\n".encode()
|
||||
raise httpx.ReadTimeout("Timeout on reading data from socket")
|
||||
|
||||
|
||||
class _UpstreamThatFailsWithUsageAfterThreeDeltas:
|
||||
headers: dict[str, str] = {}
|
||||
|
||||
async def aiter_bytes(self) -> AsyncIterator[bytes]:
|
||||
for sequence_number, delta in enumerate(_SEA_DELTAS):
|
||||
event = {
|
||||
"type": "response.output_text.delta",
|
||||
"item_id": "msg_1",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": delta,
|
||||
"sequence_number": sequence_number,
|
||||
}
|
||||
yield f"data: {json.dumps(event)}\n\n".encode()
|
||||
failed = {
|
||||
"type": "response.failed",
|
||||
"sequence_number": len(_SEA_DELTAS),
|
||||
"response": {
|
||||
"id": "resp_1",
|
||||
"object": "response",
|
||||
"created_at": 1700000000,
|
||||
"status": "failed",
|
||||
"model": "gpt-5.4-nano",
|
||||
"output": [],
|
||||
"error": {"code": "server_error", "message": "The model had an internal error"},
|
||||
"usage": {"input_tokens": 400, "output_tokens": 9, "total_tokens": 409},
|
||||
},
|
||||
}
|
||||
yield f"data: {json.dumps(failed)}\n\n".encode()
|
||||
|
||||
|
||||
async def _failure_payload_after_mid_stream_timeout(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
request_input: str | list[dict[str, object]],
|
||||
upstream: _UpstreamThatTimesOutAfterThreeDeltas | None = None,
|
||||
background_fields: Mapping[str, object] | None = None,
|
||||
replays_stored_stream: bool = False,
|
||||
) -> tuple[FailureRecorder, StandardLoggingPayload]:
|
||||
upstream_response: Final = upstream or _UpstreamThatTimesOutAfterThreeDeltas()
|
||||
recorder = FailureRecorder()
|
||||
monkeypatch.setattr(litellm, "failure_callback", [recorder])
|
||||
monkeypatch.setattr(litellm, "_async_failure_callback", [recorder])
|
||||
monkeypatch.setattr(litellm, "success_callback", [recorder])
|
||||
monkeypatch.setattr(litellm, "_async_success_callback", [recorder])
|
||||
logging_obj = _make_logging_obj()
|
||||
logging_obj.update_environment_variables(
|
||||
model="gpt-5.4-nano",
|
||||
optional_params={"stream": True},
|
||||
litellm_params={"custom_llm_provider": "openai", "aresponses": True},
|
||||
)
|
||||
iterator = ResponsesAPIStreamingIterator(
|
||||
response=cast(httpx.Response, upstream_response),
|
||||
model="gpt-5.4-nano",
|
||||
responses_api_provider_config=OpenAIResponsesAPIConfig(),
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider="openai",
|
||||
request_data={"model": "gpt-5.4-nano", "input": request_input, "stream": True, **(background_fields or {})},
|
||||
replays_stored_stream=replays_stored_stream,
|
||||
)
|
||||
received = _ChunkRecorder()
|
||||
|
||||
with pytest.raises(litellm.Timeout, match="Timeout on reading data from socket"):
|
||||
await received.drain(iterator)
|
||||
assert [getattr(chunk, "delta", None) for chunk in received.chunks] == list(upstream_response.deltas)
|
||||
|
||||
await iterator._await_pending_logging()
|
||||
(payload,) = recorder.failure_payloads
|
||||
return recorder, payload
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mid_stream_read_timeout_logs_failure_with_partial_usage_and_no_success(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
recorder, payload = await _failure_payload_after_mid_stream_timeout(monkeypatch, _SEA_PROMPT)
|
||||
|
||||
assert payload["status"] == "failure"
|
||||
assert payload["prompt_tokens"] == litellm.token_counter(
|
||||
model="gpt-5.4-nano", messages=[{"role": "user", "content": _SEA_PROMPT}]
|
||||
)
|
||||
assert payload["completion_tokens"] == litellm.token_counter(
|
||||
model="gpt-5.4-nano", text="".join(_SEA_DELTAS), count_response_tokens=True
|
||||
)
|
||||
assert payload["response_cost"] > 0
|
||||
assert recorder.success_payloads == ()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mid_stream_read_timeout_bills_the_tool_call_arguments_already_streamed(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
upstream: Final = _UpstreamThatTimesOutAfterThreeDeltas(
|
||||
event_type="response.function_call_arguments.delta", deltas=_TOOL_ARGUMENT_DELTAS
|
||||
)
|
||||
_, payload = await _failure_payload_after_mid_stream_timeout(monkeypatch, _SEA_PROMPT, upstream)
|
||||
|
||||
assert payload["status"] == "failure"
|
||||
assert payload["completion_tokens"] == litellm.token_counter(
|
||||
model="gpt-5.4-nano", text="".join(_TOOL_ARGUMENT_DELTAS), count_response_tokens=True
|
||||
)
|
||||
assert payload["response_cost"] > 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mid_stream_failure_of_a_background_response_bills_nothing_since_retrieving_it_bills_the_whole_run(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
_, payload = await _failure_payload_after_mid_stream_timeout(
|
||||
monkeypatch, _SEA_PROMPT, background_fields={"background": True}
|
||||
)
|
||||
|
||||
assert payload["status"] == "failure"
|
||||
assert payload["completion_tokens"] == 0
|
||||
assert payload["response_cost"] == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mid_stream_failure_of_a_resumed_background_stream_bills_nothing_since_its_create_bills_the_run(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
_, payload = await _failure_payload_after_mid_stream_timeout(monkeypatch, _SEA_PROMPT, replays_stored_stream=True)
|
||||
|
||||
assert payload["status"] == "failure"
|
||||
assert payload["completion_tokens"] == 0
|
||||
assert payload["response_cost"] == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mid_stream_failure_usage_estimate_never_fetches_image_dimensions(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
_, payload = await _failure_payload_after_mid_stream_timeout(monkeypatch, _UNREACHABLE_IMAGE_INPUT)
|
||||
|
||||
assert payload["status"] == "failure"
|
||||
assert payload["prompt_tokens"] == litellm.token_counter(
|
||||
model="gpt-5.4-nano",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Describe this picture."},
|
||||
{"type": "image_url", "image_url": {"url": "http://images.invalid/photo.png", "detail": "high"}},
|
||||
],
|
||||
}
|
||||
],
|
||||
use_default_image_token_count=True,
|
||||
)
|
||||
assert payload["completion_tokens"] > 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mid_stream_failure_keeps_the_usage_the_provider_reported(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
recorder = FailureRecorder()
|
||||
monkeypatch.setattr(litellm, "failure_callback", [recorder])
|
||||
monkeypatch.setattr(litellm, "_async_failure_callback", [recorder])
|
||||
logging_obj = _make_logging_obj()
|
||||
logging_obj.update_environment_variables(
|
||||
model="gpt-5.4-nano",
|
||||
optional_params={"stream": True},
|
||||
litellm_params={"custom_llm_provider": "openai", "aresponses": True},
|
||||
)
|
||||
iterator = ResponsesAPIStreamingIterator(
|
||||
response=cast(httpx.Response, _UpstreamThatFailsWithUsageAfterThreeDeltas()),
|
||||
model="gpt-5.4-nano",
|
||||
responses_api_provider_config=OpenAIResponsesAPIConfig(),
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider="openai",
|
||||
request_data={"model": "gpt-5.4-nano", "input": _SEA_PROMPT, "stream": True},
|
||||
)
|
||||
|
||||
with pytest.raises(MidStreamFallbackError, match="The model had an internal error"):
|
||||
await _ChunkRecorder().drain(iterator)
|
||||
await iterator._await_pending_logging()
|
||||
|
||||
(payload,) = recorder.failure_payloads
|
||||
assert (payload["prompt_tokens"], payload["completion_tokens"]) == (400, 9)
|
||||
|
||||
|
||||
_CREATED_RESPONSE_WITH_PRIORITY_TIER: Final[Mapping[str, object]] = {
|
||||
"id": "resp_partial",
|
||||
"object": "response",
|
||||
|
|
@ -211,17 +461,135 @@ def _completed_after(text: str) -> Mapping[str, object]:
|
|||
}
|
||||
|
||||
|
||||
def _live_iterator(logging_obj: LitellmLogging) -> ResponsesAPIStreamingIterator:
|
||||
def _live_iterator(
|
||||
logging_obj: LitellmLogging, upstream: object = None, background: bool = False
|
||||
) -> ResponsesAPIStreamingIterator:
|
||||
return ResponsesAPIStreamingIterator(
|
||||
response=httpx.Response(200),
|
||||
response=cast(httpx.Response, upstream or httpx.Response(200)),
|
||||
model="gpt-5.4-nano",
|
||||
responses_api_provider_config=OpenAIResponsesAPIConfig(),
|
||||
logging_obj=logging_obj,
|
||||
request_data={"model": "gpt-5.4-nano", "input": "tell me a story"},
|
||||
request_data={"model": "gpt-5.4-nano", "input": "tell me a story", "background": background},
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
|
||||
|
||||
class _UpstreamThatTimesOutAfterRelayingText:
|
||||
headers: dict[str, str] = {}
|
||||
|
||||
async def aiter_bytes(self) -> AsyncIterator[bytes]:
|
||||
for event in (_CREATED_WITH_PRIORITY_TIER, _text_delta("Once upon a time", 1)):
|
||||
yield f"data: {json.dumps(event)}\n\n".encode()
|
||||
raise httpx.ReadTimeout("Timeout on reading data from socket")
|
||||
|
||||
|
||||
class _UpstreamThatCompletesThenTimesOut:
|
||||
headers: dict[str, str] = {}
|
||||
|
||||
async def aiter_bytes(self) -> AsyncIterator[bytes]:
|
||||
events: Final = (
|
||||
_CREATED_WITH_PRIORITY_TIER,
|
||||
*(_text_delta(delta, sequence_number) for sequence_number, delta in enumerate(_SEA_DELTAS, start=1)),
|
||||
_completed_after("".join(_SEA_DELTAS)),
|
||||
)
|
||||
for event in events:
|
||||
yield f"data: {json.dumps(event)}\n\n".encode()
|
||||
raise httpx.ReadTimeout("Timeout on reading data from socket")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_timeout_after_completed_response_ends_the_stream_and_bills_it_once(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
recorder: Final = FailureRecorder()
|
||||
monkeypatch.setattr(litellm, "failure_callback", [recorder])
|
||||
monkeypatch.setattr(litellm, "_async_failure_callback", [recorder])
|
||||
monkeypatch.setattr(litellm, "success_callback", [recorder])
|
||||
monkeypatch.setattr(litellm, "_async_success_callback", [recorder])
|
||||
logging_obj: Final = _make_logging_obj()
|
||||
logging_obj.update_environment_variables(
|
||||
model="gpt-5.4-nano",
|
||||
optional_params={"stream": True},
|
||||
litellm_params={"custom_llm_provider": "openai", "aresponses": True},
|
||||
)
|
||||
iterator: Final = _live_iterator(logging_obj, upstream=_UpstreamThatCompletesThenTimesOut())
|
||||
|
||||
chunks: Final = _ChunkRecorder()
|
||||
await chunks.drain(iterator)
|
||||
await iterator._await_pending_logging()
|
||||
|
||||
assert chunks.chunks[-1].type == "response.completed"
|
||||
assert len(recorder.success_payloads) == 1
|
||||
success_payload: Final = recorder.success_payloads[0]
|
||||
response_payload: Final = success_payload["response"]
|
||||
assert isinstance(response_payload, dict)
|
||||
usage: Final = cast(dict[str, object], response_payload["usage"])
|
||||
assert (usage["prompt_tokens"], usage["completion_tokens"], usage["total_tokens"]) == (5, 7, 12)
|
||||
assert recorder.failure_payloads == ()
|
||||
assert logging_obj.model_call_details.get("combined_usage_object") is None
|
||||
|
||||
|
||||
class _UpstreamThatTimesOutAfterCreated:
|
||||
headers: dict[str, str] = {}
|
||||
|
||||
def __init__(self, **created_fields: object) -> None:
|
||||
self.created_fields = created_fields
|
||||
|
||||
async def aiter_bytes(self) -> AsyncIterator[bytes]:
|
||||
created: Final = {
|
||||
**_CREATED_WITH_PRIORITY_TIER,
|
||||
"response": {**_CREATED_RESPONSE_WITH_PRIORITY_TIER, **self.created_fields},
|
||||
}
|
||||
for event in (created, _text_delta("Once upon a time", 1)):
|
||||
yield f"data: {json.dumps(event)}\n\n".encode()
|
||||
raise httpx.ReadTimeout("Timeout on reading data from socket")
|
||||
|
||||
|
||||
async def _mid_stream_failure_payload(
|
||||
monkeypatch: pytest.MonkeyPatch, requested_tier: str | None, **created_fields: object
|
||||
) -> StandardLoggingPayload:
|
||||
recorder: Final = FailureRecorder()
|
||||
monkeypatch.setattr(litellm, "failure_callback", [recorder])
|
||||
monkeypatch.setattr(litellm, "_async_failure_callback", [recorder])
|
||||
logging_obj: Final = _make_logging_obj()
|
||||
logging_obj.update_environment_variables(
|
||||
model="gpt-5.4-nano",
|
||||
optional_params={"stream": True}
|
||||
if requested_tier is None
|
||||
else {"stream": True, "service_tier": requested_tier},
|
||||
litellm_params={"custom_llm_provider": "openai", "aresponses": True},
|
||||
)
|
||||
iterator: Final = _live_iterator(logging_obj, upstream=_UpstreamThatTimesOutAfterCreated(**created_fields))
|
||||
|
||||
with pytest.raises(litellm.Timeout):
|
||||
await _ChunkRecorder().drain(iterator)
|
||||
await iterator._await_pending_logging()
|
||||
|
||||
(payload,) = recorder.failure_payloads
|
||||
return payload
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mid_stream_failure_bills_the_relayed_text_at_the_tier_the_provider_served(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
flex_requested: Final = await _mid_stream_failure_payload(monkeypatch, "flex", service_tier="default")
|
||||
nothing_requested: Final = await _mid_stream_failure_payload(monkeypatch, None, service_tier="default")
|
||||
|
||||
assert nothing_requested["response_cost"] > 0
|
||||
assert flex_requested["response_cost"] == nothing_requested["response_cost"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mid_stream_failure_bills_nothing_when_the_provider_runs_the_response_in_background(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
payload: Final = await _mid_stream_failure_payload(monkeypatch, None, background=True)
|
||||
|
||||
assert payload["status"] == "failure"
|
||||
assert (payload["completion_tokens"], payload["response_cost"]) == (0, 0)
|
||||
|
||||
|
||||
def test_partial_response_on_disconnect_bills_the_relayed_text_at_the_served_service_tier():
|
||||
iterator: Final = _live_iterator(_make_logging_obj())
|
||||
for event in (_CREATED_WITH_PRIORITY_TIER, _text_delta("Once upon a time ", 1), _text_delta("a fox slept.", 2)):
|
||||
|
|
@ -278,3 +646,25 @@ async def test_partial_response_on_disconnect_is_none_when_there_is_nothing_to_b
|
|||
|
||||
assert iterator.partial_response_on_disconnect() is None
|
||||
await iterator._await_pending_logging()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_partial_response_on_disconnect_is_none_once_the_failure_log_billed_the_relayed_text():
|
||||
iterator: Final = _live_iterator(_make_logging_obj(), upstream=_UpstreamThatTimesOutAfterRelayingText())
|
||||
|
||||
with pytest.raises(litellm.Timeout):
|
||||
await _ChunkRecorder().drain(iterator)
|
||||
|
||||
assert iterator.partial_response_on_disconnect() is None
|
||||
await iterator._await_pending_logging()
|
||||
|
||||
|
||||
def test_partial_response_on_disconnect_reports_the_relayed_text_of_a_background_response():
|
||||
iterator: Final = _live_iterator(_make_logging_obj(), background=True)
|
||||
for event in (_CREATED_WITH_PRIORITY_TIER, _text_delta("Once upon a time", 1)):
|
||||
iterator._process_chunk(json.dumps(event))
|
||||
|
||||
incomplete_event: Final = iterator.partial_response_on_disconnect()
|
||||
|
||||
assert incomplete_event is not None
|
||||
assert incomplete_event.response.output[0].content[0].text == "Once upon a time"
|
||||
|
|
|
|||
|
|
@ -55,21 +55,6 @@ async def test_adapter_raises_clean_close_when_rust_connection_ends() -> None:
|
|||
await adapter.recv()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bridge_unavailable_returns_none(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(responses_websocket, "_STATE", responses_websocket._RustResponsesWebSocketState())
|
||||
monkeypatch.setattr(responses_websocket, "get_native_bridge", lambda: None)
|
||||
|
||||
assert (
|
||||
await responses_websocket.connect(
|
||||
url="wss://example.test/responses",
|
||||
headers={},
|
||||
timeout=None,
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_enabled_bridge_connects_and_adapts_socket(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
|
|
|
|||
|
|
@ -210,9 +210,11 @@ _TRAILING_ERRORS = [
|
|||
httpx.RemoteProtocolError("peer closed connection without sending complete message body"),
|
||||
]
|
||||
|
||||
_TRAILING_ERRORS_AFTER_COMPLETED = [*_TRAILING_ERRORS, httpx.ReadTimeout("Timeout on reading data from socket")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("trailing_error", _TRAILING_ERRORS, ids=type)
|
||||
@pytest.mark.parametrize("trailing_error", _TRAILING_ERRORS_AFTER_COMPLETED, ids=type)
|
||||
async def test_transport_error_after_completed_event_ends_stream_cleanly(trailing_error):
|
||||
"""A sloppy connection close after `response.completed` must not turn a
|
||||
complete stream into an error (regression guard for the transport no longer
|
||||
|
|
@ -243,7 +245,7 @@ async def test_transport_error_before_completed_event_raises():
|
|||
pass
|
||||
|
||||
|
||||
@pytest.mark.parametrize("trailing_error", _TRAILING_ERRORS, ids=type)
|
||||
@pytest.mark.parametrize("trailing_error", _TRAILING_ERRORS_AFTER_COMPLETED, ids=type)
|
||||
def test_sync_transport_error_after_completed_event_ends_stream_cleanly(trailing_error):
|
||||
iterator = _make_sync_iterator(
|
||||
sse_events=_COMPLETE_STREAM_EVENTS,
|
||||
|
|
@ -1205,7 +1207,9 @@ async def test_completed_event_survives_a_failing_usage_estimate():
|
|||
falls back to usage None."""
|
||||
malformed_input: Final = [{"type": "message", "role": "user", "content": 42}]
|
||||
with pytest.raises(ValueError, match="Invalid content type"):
|
||||
_estimate_usage_from_text("gpt-4o-mini", malformed_input, {"input": malformed_input}, "hello world")
|
||||
_estimate_usage_from_text(
|
||||
"gpt-4o-mini", malformed_input, {"input": malformed_input}, "hello world"
|
||||
)
|
||||
|
||||
response = _responses_api_response_without_usage()
|
||||
iterator = _make_iterator(
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ from botocore.credentials import Credentials
|
|||
from pydantic import JsonValue
|
||||
|
||||
import litellm
|
||||
from litellm.rust_bridge import _native as native
|
||||
from litellm.rust_bridge.host.secret_manager import capture_secret_manager, native_secret_manager_config
|
||||
from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2
|
||||
from litellm.secret_managers.cyberark_secret_manager import CyberArkSecretManager
|
||||
|
|
@ -67,7 +68,6 @@ def _vault_body(value: str) -> dict[str, object]:
|
|||
|
||||
|
||||
def _native_handle(client: object) -> "_SecretManagerRuntime":
|
||||
native: Final = pytest.importorskip("litellm.rust_bridge._native")
|
||||
handle: Final = native._SecretManagerRuntime.from_client(client)
|
||||
assert handle is not None
|
||||
return handle
|
||||
|
|
@ -75,7 +75,6 @@ def _native_handle(client: object) -> "_SecretManagerRuntime":
|
|||
|
||||
@pytest.mark.parametrize("system", ("aws_secret_manager", "hashicorp_vault", "cyberark"))
|
||||
async def test_python_native_handle_reuses_backend_across_sync_and_async_reads(system: str) -> None:
|
||||
native: Final = pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
environment: Final = {
|
||||
"AWS_REGION_NAME": "us-east-1",
|
||||
|
|
@ -109,7 +108,6 @@ async def test_python_native_handle_reuses_backend_across_sync_and_async_reads(s
|
|||
|
||||
|
||||
def test_shared_initializer_captures_credentials_and_tracks_instance_settings(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
native: Final = pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
server.default_response = ResponseSpec(body={"SecretString": "native-value"})
|
||||
server.expected_requests = 2
|
||||
|
|
@ -141,7 +139,6 @@ def test_shared_initializer_captures_credentials_and_tracks_instance_settings(mo
|
|||
def test_configuration_replacement_rebuilds_without_invalidating_existing_handles(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
native: Final = pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as first_server, recording_service() as second_server:
|
||||
first_server.default_response = ResponseSpec(body=_vault_body("first"))
|
||||
second_server.default_response = ResponseSpec(body=_vault_body("second"))
|
||||
|
|
@ -158,7 +155,6 @@ def test_configuration_replacement_rebuilds_without_invalidating_existing_handle
|
|||
|
||||
|
||||
def test_custom_subclass_is_not_captured_for_a_native_handle() -> None:
|
||||
native: Final = pytest.importorskip("litellm.rust_bridge._native")
|
||||
|
||||
class CustomManager(AWSSecretsManagerV2):
|
||||
def sync_read_secret(self, secret_name: str, primary_secret_name: str | None = None) -> str:
|
||||
|
|
@ -171,19 +167,16 @@ def test_custom_subclass_is_not_captured_for_a_native_handle() -> None:
|
|||
|
||||
@pytest.mark.parametrize("system", ("google_secret_manager", "hashicorp_vault", "cyberark"))
|
||||
def test_enterprise_backends_cannot_initialize_without_entitlement(system: str) -> None:
|
||||
native: Final = pytest.importorskip("litellm.rust_bridge._native")
|
||||
with pytest.raises(ValueError, match=r"[Ee]nterprise|[Pp]remium"):
|
||||
native._SecretManagerRuntime.from_config(system, {"CYBERARK_API_KEY": "key"})
|
||||
|
||||
|
||||
def test_azure_factory_rejects_unencrypted_vault_endpoints() -> None:
|
||||
native: Final = pytest.importorskip("litellm.rust_bridge._native")
|
||||
with pytest.raises(ValueError, match="https"):
|
||||
native._SecretManagerRuntime.from_config("azure_key_vault", {"AZURE_KEY_VAULT_URI": "http://127.0.0.1:1"})
|
||||
|
||||
|
||||
def test_direct_builtin_constructor_can_use_native_without_registration(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
native: Final = pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
server.default_response = ResponseSpec(body=_vault_body("native-value"))
|
||||
manager: Final = _vault(monkeypatch, server.base_url)
|
||||
|
|
@ -193,12 +186,10 @@ def test_direct_builtin_constructor_can_use_native_without_registration(monkeypa
|
|||
|
||||
|
||||
def test_uncaptured_clients_have_no_native_handle() -> None:
|
||||
native: Final = pytest.importorskip("litellm.rust_bridge._native")
|
||||
assert native._SecretManagerRuntime.from_client(SimpleNamespace()) is None
|
||||
|
||||
|
||||
def test_native_failure_is_not_replayed_in_python(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
server.default_response = ResponseSpec(status=403, body={"errors": ["denied"]})
|
||||
server.expected_requests = 1
|
||||
|
|
@ -267,7 +258,6 @@ def test_same_named_custom_client_is_not_captured() -> None:
|
|||
|
||||
@pytest.mark.parametrize("value", ("text", "", True, False, 42, 2**100, [1, "two"], {"nested": True}, None))
|
||||
async def test_aws_primary_values_match_python_handler(monkeypatch: pytest.MonkeyPatch, value: JsonValue) -> None:
|
||||
native: Final = pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
server.default_response = ResponseSpec(body={"SecretString": json.dumps({"KEY": value})})
|
||||
server.expected_requests = 3
|
||||
|
|
@ -290,7 +280,6 @@ async def test_aws_primary_values_match_python_handler(monkeypatch: pytest.Monke
|
|||
|
||||
@pytest.mark.parametrize("primary", (None, "primary"))
|
||||
async def test_aws_primary_values_match_python_reads(monkeypatch: pytest.MonkeyPatch, primary: str | None) -> None:
|
||||
native: Final = pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
server.default_response = ResponseSpec(body={"SecretString": '{"KEY":"value"}'})
|
||||
server.expected_requests = 3
|
||||
|
|
@ -343,7 +332,6 @@ def test_aws_absence_and_failed_reads_match_python_without_environment_fallback(
|
|||
async def test_aws_primary_json_errors_preserve_python_exception_details(
|
||||
monkeypatch: pytest.MonkeyPatch, document: str
|
||||
) -> None:
|
||||
native: Final = pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
server.default_response = ResponseSpec(body={"SecretString": document})
|
||||
server.expected_requests = 3
|
||||
|
|
@ -370,7 +358,6 @@ async def test_aws_primary_json_errors_preserve_python_exception_details(
|
|||
async def test_public_aws_reads_preserve_coroutines_and_per_call_credentials(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as initial, recording_service() as selected:
|
||||
initial.expected_requests = 0
|
||||
selected.expected_requests = 2
|
||||
|
|
@ -421,7 +408,6 @@ async def test_public_aws_reads_preserve_coroutines_and_per_call_credentials(
|
|||
|
||||
|
||||
async def test_public_aws_primary_reads_ignore_operation_overrides(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server, recording_service() as unused:
|
||||
server.default_response = ResponseSpec(body={"SecretString": '{"KEY":true}'})
|
||||
server.expected_requests = 2
|
||||
|
|
@ -440,7 +426,6 @@ async def test_public_aws_primary_reads_ignore_operation_overrides(monkeypatch:
|
|||
|
||||
|
||||
async def test_public_aws_bootstrap_names_only_bypass_sync_reads(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
server.default_response = ResponseSpec(body={"SecretString": "remote-access"})
|
||||
server.expected_requests = 1
|
||||
|
|
@ -457,7 +442,6 @@ async def test_public_aws_bootstrap_names_only_bypass_sync_reads(monkeypatch: py
|
|||
async def test_public_aws_read_timeouts_follow_the_python_http_handler(
|
||||
monkeypatch: pytest.MonkeyPatch, timeout: float | httpx.Timeout
|
||||
) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
server.default_response = ResponseSpec(body={"SecretString": "too-late"}, delay=0.25)
|
||||
server.expected_requests = 2
|
||||
|
|
@ -470,7 +454,6 @@ async def test_public_aws_read_timeouts_follow_the_python_http_handler(
|
|||
|
||||
|
||||
async def test_public_vault_reads_keep_overrides_cache_and_coroutines(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
server.default_response = ResponseSpec(body=_vault_body("value"))
|
||||
server.expected_requests = 2
|
||||
|
|
@ -490,7 +473,6 @@ async def test_public_vault_reads_keep_overrides_cache_and_coroutines(monkeypatc
|
|||
async def test_public_vault_failed_reads_return_none_without_replay(
|
||||
monkeypatch: pytest.MonkeyPatch, status: int
|
||||
) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
server.default_response = ResponseSpec(status=status, body={"errors": ["unavailable"]})
|
||||
server.expected_requests = 2
|
||||
|
|
@ -502,7 +484,6 @@ async def test_public_vault_failed_reads_return_none_without_replay(
|
|||
async def test_public_cyberark_reads_reuse_authentication_and_cached_values(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.secret_managers.cyberark_secret_manager import CyberArkSecretManager
|
||||
|
||||
|
|
@ -536,7 +517,6 @@ def test_public_aws_bootstrap_read_does_not_hit_the_provider(monkeypatch: pytest
|
|||
def test_public_vault_prefix_overrides_match_python_string_conversion(
|
||||
monkeypatch: pytest.MonkeyPatch, override: str | int
|
||||
) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
monkeypatch.setenv("HCP_VAULT_PATH_PREFIX", "default-prefix")
|
||||
with recording_service() as server:
|
||||
server.default_response = ResponseSpec(body=_vault_body("value"))
|
||||
|
|
@ -562,7 +542,6 @@ def _cyberark(monkeypatch: pytest.MonkeyPatch, address: str) -> CyberArkSecretMa
|
|||
async def test_public_cyberark_writes_and_deletes_share_the_read_cache(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
for body in (b"token", b"old", {}, {}, b"provider-after-delete"):
|
||||
server.enqueue(ResponseSpec(body=body))
|
||||
|
|
@ -609,7 +588,6 @@ async def test_public_cyberark_write_errors_match_python_without_http_retries(
|
|||
status: int,
|
||||
authentication: bool,
|
||||
) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
responses: Final = (
|
||||
(ResponseSpec(status=status, body={}),) * 2
|
||||
|
|
@ -635,7 +613,6 @@ async def test_public_cyberark_write_errors_match_python_without_http_retries(
|
|||
async def test_public_cyberark_write_recovers_from_initial_policy_authentication_failure(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
server.enqueue(ResponseSpec(status=401, body={}))
|
||||
server.enqueue(ResponseSpec(body=b"token"))
|
||||
|
|
@ -658,7 +635,6 @@ async def test_public_cyberark_write_rejects_unsafe_names_before_authentication(
|
|||
monkeypatch: pytest.MonkeyPatch,
|
||||
name: str,
|
||||
) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
server.expected_requests = 0
|
||||
manager: Final = _cyberark(monkeypatch, server.base_url)
|
||||
|
|
@ -673,7 +649,6 @@ async def test_public_cyberark_rotation_returns_the_write_response_and_retains_o
|
|||
monkeypatch: pytest.MonkeyPatch,
|
||||
same_name: bool,
|
||||
) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
for body in (b"token", b"old-value", {}, {}):
|
||||
server.enqueue(ResponseSpec(body=body))
|
||||
|
|
@ -696,7 +671,6 @@ async def test_public_native_cyberark_rotation_requires_a_fresh_matching_replace
|
|||
monkeypatch: pytest.MonkeyPatch,
|
||||
replacement: bytes | None,
|
||||
) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
for body in (b"token", b"old-value", {}, {}):
|
||||
server.enqueue(ResponseSpec(body=body))
|
||||
|
|
@ -720,7 +694,6 @@ async def test_public_native_cyberark_rotation_requires_a_fresh_matching_replace
|
|||
async def test_public_cyberark_cached_authentication_does_not_retry_denied_reads(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
server.enqueue(ResponseSpec(body=b"token"))
|
||||
server.enqueue(ResponseSpec(body=b"value"))
|
||||
|
|
@ -734,7 +707,6 @@ async def test_public_cyberark_cached_authentication_does_not_retry_denied_reads
|
|||
async def test_cyberark_handler_errors_match_python_after_cached_authentication_is_denied(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
for response in (
|
||||
ResponseSpec(body=b"token"),
|
||||
|
|
@ -757,7 +729,6 @@ async def test_cyberark_handler_errors_match_python_after_cached_authentication_
|
|||
async def test_public_cyberark_connection_errors_match_python(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
server.expected_requests = 0
|
||||
address: Final = server.base_url
|
||||
|
|
@ -769,7 +740,6 @@ async def test_public_cyberark_connection_errors_match_python(
|
|||
|
||||
|
||||
async def test_public_native_cyberark_rotation_stops_after_a_failed_write(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
for body in (b"token", b"old-value", {}):
|
||||
server.enqueue(ResponseSpec(body=body))
|
||||
|
|
@ -788,7 +758,6 @@ async def test_public_vault_writes_preserve_complete_responses_and_request_field
|
|||
monkeypatch: pytest.MonkeyPatch,
|
||||
description: str | None,
|
||||
) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
response: Final = {
|
||||
"request_id": "test-request",
|
||||
|
|
@ -829,7 +798,6 @@ async def test_public_vault_mutation_http_errors_match_python_without_retry(
|
|||
operation: str,
|
||||
status: int,
|
||||
) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
server.default_response = ResponseSpec(status=status, body={"errors": ["denied"]})
|
||||
server.expected_requests = 2
|
||||
|
|
@ -856,7 +824,6 @@ async def test_public_vault_write_response_conversion_matches_python(
|
|||
monkeypatch: pytest.MonkeyPatch,
|
||||
body: bytes,
|
||||
) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
server.default_response = ResponseSpec(body=body)
|
||||
server.expected_requests = 2
|
||||
|
|
@ -870,7 +837,6 @@ async def test_public_vault_write_response_conversion_matches_python(
|
|||
async def test_public_vault_deletion_invalidates_cached_fields(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
server.enqueue(ResponseSpec(body=_vault_body("old")))
|
||||
server.enqueue(ResponseSpec(status=204, body=b""))
|
||||
|
|
@ -893,7 +859,6 @@ async def test_public_vault_rotation_preserves_response_and_best_effort_deletion
|
|||
same_name: bool,
|
||||
delete_status: int,
|
||||
) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
response: Final = {"request_id": "write-id", "data": {"version": 3}, "extra": [1, 2]}
|
||||
server.enqueue(ResponseSpec(body=b"current-existence-is-status-only"))
|
||||
|
|
@ -924,7 +889,6 @@ async def test_public_vault_rotation_failure_messages_and_request_counts_match_p
|
|||
stage: str,
|
||||
status: int,
|
||||
) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
before: Final = (
|
||||
()
|
||||
|
|
@ -960,7 +924,6 @@ async def test_public_vault_rotation_mismatches_do_not_delete_the_old_alias(
|
|||
monkeypatch: pytest.MonkeyPatch,
|
||||
value: JsonValue,
|
||||
) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
responses: Final = (
|
||||
ResponseSpec(body=_vault_body("old")),
|
||||
|
|
@ -983,7 +946,6 @@ async def test_public_vault_mutation_timeouts_match_python(
|
|||
monkeypatch: pytest.MonkeyPatch,
|
||||
operation: str,
|
||||
) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
server.default_response = ResponseSpec(body=_vault_body("value"), delay=0.25)
|
||||
server.expected_requests = 2
|
||||
|
|
@ -1021,7 +983,6 @@ async def test_public_vault_unsafe_names_fail_before_authentication(
|
|||
monkeypatch: pytest.MonkeyPatch,
|
||||
operation: str,
|
||||
) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
server.expected_requests = 0
|
||||
manager: Final = _vault(monkeypatch, server.base_url)
|
||||
|
|
@ -1038,7 +999,6 @@ async def test_public_vault_authentication_errors_match_python(
|
|||
monkeypatch: pytest.MonkeyPatch,
|
||||
operation: str,
|
||||
) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
monkeypatch.setenv("HCP_VAULT_APPROLE_ROLE_ID", "role")
|
||||
monkeypatch.setenv("HCP_VAULT_APPROLE_SECRET_ID", "secret-id")
|
||||
with recording_service() as server:
|
||||
|
|
@ -1064,7 +1024,6 @@ async def test_public_vault_authentication_errors_match_python(
|
|||
async def test_public_vault_rotation_stops_on_a_success_response_containing_an_error(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
result: Final = {"status": "error", "message": "write rejected", "extra": 2**100}
|
||||
for response in (ResponseSpec(body=_vault_body("old")), ResponseSpec(body=result)) * 2:
|
||||
|
|
@ -1078,7 +1037,6 @@ async def test_public_vault_rotation_stops_on_a_success_response_containing_an_e
|
|||
|
||||
|
||||
async def test_public_native_vault_write_invalidates_stale_cached_values(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
server.enqueue(ResponseSpec(body=_vault_body("old")))
|
||||
server.enqueue(ResponseSpec(body={"data": {"version": 2}}))
|
||||
|
|
@ -1094,7 +1052,6 @@ async def test_public_native_vault_write_invalidates_stale_cached_values(monkeyp
|
|||
async def test_public_native_vault_write_rejects_description_overwriting_the_secret(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
server.expected_requests = 0
|
||||
handle: Final = _native_handle(_vault(monkeypatch, server.base_url))
|
||||
|
|
@ -1111,7 +1068,6 @@ async def test_public_vault_rotation_preserves_malformed_verification_errors(
|
|||
monkeypatch: pytest.MonkeyPatch,
|
||||
body: JsonValue,
|
||||
) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
with recording_service() as server:
|
||||
responses: Final = (
|
||||
ResponseSpec(body=_vault_body("old")),
|
||||
|
|
|
|||
|
|
@ -17,8 +17,6 @@ from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagem
|
|||
from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_MODEL, MESSAGES_RESPONSE
|
||||
from tests.test_litellm_rust.support.recording_server import ResponseSpec, recording_service
|
||||
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
|
||||
pytestmark = pytest.mark.usefixtures("local_model_cost_map")
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ import pytest
|
|||
import litellm
|
||||
from litellm.integrations.custom_secret_manager import CustomSecretManager
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRResponse
|
||||
from litellm.rust_bridge import _native as native
|
||||
from litellm.rust_bridge.host import settings
|
||||
from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR
|
||||
from litellm.rust_bridge.public_call import NativeCall
|
||||
|
|
@ -19,8 +20,6 @@ from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagem
|
|||
from tests.test_litellm_rust.ocr.data import OCR_DOCUMENT, OCR_MODEL, OCR_RESPONSE
|
||||
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, recording_service
|
||||
|
||||
native: Final = pytest.importorskip("litellm.rust_bridge._native")
|
||||
|
||||
AccessMode: TypeAlias = Literal["read_only", "write_only", "read_and_write"]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ from typing import Final, cast # noqa: TID251 # narrows the async binding's un
|
|||
|
||||
import pytest
|
||||
|
||||
from litellm.rust_bridge import bindings
|
||||
from litellm.rust_bridge import _native, bindings
|
||||
from litellm.rust_bridge.chat_completions import entrypoints as chat_completions
|
||||
from litellm.rust_bridge.embeddings import entrypoints as embeddings
|
||||
from litellm.rust_bridge.execution import Complete, Step
|
||||
|
|
@ -17,12 +17,15 @@ from litellm.rust_bridge.responses import entrypoints as responses
|
|||
from litellm.rust_bridge.transcription import native as transcription
|
||||
|
||||
|
||||
def test_binding_distinguishes_disable_from_reset(monkeypatch) -> None:
|
||||
native = SimpleNamespace(route=lambda: "native")
|
||||
monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
|
||||
binding: bindings.NativeBinding[object] = bindings.NativeBinding("route", validate=lambda value: value)
|
||||
def test_binding_distinguishes_disable_from_reset() -> None:
|
||||
def route() -> str:
|
||||
return "native"
|
||||
|
||||
assert binding.load() is native.route
|
||||
binding: bindings.NativeBinding[object] = bindings.NativeBinding(
|
||||
"route", validate=lambda value: value, source=SimpleNamespace(route=route)
|
||||
)
|
||||
|
||||
assert binding.load() is route
|
||||
|
||||
binding.override(None)
|
||||
assert binding.load() is None
|
||||
|
|
@ -32,16 +35,15 @@ def test_binding_distinguishes_disable_from_reset(monkeypatch) -> None:
|
|||
assert binding.load() is replacement
|
||||
|
||||
binding.reset()
|
||||
assert binding.load() is native.route
|
||||
assert binding.load() is route
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("value", "expected"), ((3, 3), ("invalid", None), (None, None)))
|
||||
def test_binding_validates_native_attribute(
|
||||
monkeypatch: pytest.MonkeyPatch, value: object, expected: int | None
|
||||
) -> None:
|
||||
def test_binding_validates_native_attribute(value: object, expected: int | None) -> None:
|
||||
native: Final = SimpleNamespace(route=value)
|
||||
monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
|
||||
binding: Final = bindings.NativeBinding("route", validate=lambda item: item if isinstance(item, int) else None)
|
||||
binding: Final = bindings.NativeBinding(
|
||||
"route", validate=lambda item: item if isinstance(item, int) else None, source=native
|
||||
)
|
||||
|
||||
assert binding.load() == expected
|
||||
|
||||
|
|
@ -72,28 +74,33 @@ ROUTE_BINDINGS: Final = tuple((attribute, binding) for attribute, binding, _ in
|
|||
@pytest.mark.parametrize(
|
||||
("attribute", "route_binding"), ROUTE_BINDINGS, ids=[attribute for attribute, _ in ROUTE_BINDINGS]
|
||||
)
|
||||
def test_route_bindings_reject_non_callable_native_attributes(
|
||||
monkeypatch: pytest.MonkeyPatch, attribute: str, route_binding: bindings.NativeBinding[object]
|
||||
def test_route_bindings_resolve_the_packaged_native_route(
|
||||
attribute: str, route_binding: bindings.NativeBinding[object]
|
||||
) -> None:
|
||||
monkeypatch.setattr(bindings, "get_native_bridge", lambda: SimpleNamespace(**{attribute: "not callable"}))
|
||||
route_binding.reset()
|
||||
assert route_binding.load() is None
|
||||
route_binding.reset()
|
||||
assert attribute in dir(_native)
|
||||
assert callable(route_binding.load())
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("attribute", "route_binding"), ROUTE_BINDINGS, ids=[attribute for attribute, _ in ROUTE_BINDINGS]
|
||||
)
|
||||
def test_route_bindings_reject_non_callable_native_attributes(
|
||||
attribute: str, route_binding: bindings.NativeBinding[object]
|
||||
) -> None:
|
||||
assert route_binding.with_source(SimpleNamespace(**{attribute: "not callable"})).load() is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("attribute", "route_binding"), TRANSCRIPTION_BINDINGS, ids=[attribute for attribute, _ in TRANSCRIPTION_BINDINGS]
|
||||
)
|
||||
def test_transcription_bindings_expose_the_native_callable_itself(
|
||||
monkeypatch: pytest.MonkeyPatch, attribute: str, route_binding: bindings.NativeBinding[object]
|
||||
attribute: str, route_binding: bindings.NativeBinding[object]
|
||||
) -> None:
|
||||
def native_route() -> None:
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(bindings, "get_native_bridge", lambda: SimpleNamespace(**{attribute: native_route}))
|
||||
route_binding.reset()
|
||||
assert route_binding.load() is native_route
|
||||
route_binding.reset()
|
||||
assert route_binding.with_source(SimpleNamespace(**{attribute: native_route})).load() is native_route
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -116,7 +123,6 @@ class _EchoingExecution:
|
|||
ids=[attribute for attribute, _, _ in EXECUTION_BINDINGS],
|
||||
)
|
||||
def test_execution_bindings_start_the_native_execution_and_return_its_result(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
attribute: str,
|
||||
route_binding: bindings.NativeBinding[object],
|
||||
asynchronous: bool,
|
||||
|
|
@ -128,9 +134,7 @@ def test_execution_bindings_start_the_native_execution_and_return_its_result(
|
|||
received.append(incoming)
|
||||
return _EchoingExecution()
|
||||
|
||||
monkeypatch.setattr(bindings, "get_native_bridge", lambda: SimpleNamespace(**{attribute: native_route}))
|
||||
route_binding.reset()
|
||||
bound: Final = route_binding.load()
|
||||
bound: Final = route_binding.with_source(SimpleNamespace(**{attribute: native_route})).load()
|
||||
assert callable(bound)
|
||||
|
||||
outcome: Final = bound(call)
|
||||
|
|
@ -139,7 +143,6 @@ def test_execution_bindings_start_the_native_execution_and_return_its_result(
|
|||
if asynchronous
|
||||
else outcome
|
||||
)
|
||||
route_binding.reset()
|
||||
|
||||
assert received == [call]
|
||||
assert result == ("completed with", None)
|
||||
|
|
|
|||
|
|
@ -1,36 +1,21 @@
|
|||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.rust_bridge import fork_guard
|
||||
|
||||
|
||||
def _reserve_with(monkeypatch: pytest.MonkeyPatch, native: object) -> None:
|
||||
monkeypatch.setattr(fork_guard, "get_native_bridge", lambda: native)
|
||||
fork_guard.reserve_process_for_forking("the gunicorn master")
|
||||
|
||||
|
||||
def test_missing_extension_has_nothing_to_reserve(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
assert _reserve_with(monkeypatch, None) is None
|
||||
|
||||
|
||||
def test_extension_built_before_reservation_existed_passes(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
assert _reserve_with(monkeypatch, SimpleNamespace()) is None
|
||||
|
||||
|
||||
def test_unused_extension_is_reserved(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def test_unused_extension_is_reserved() -> None:
|
||||
calls: list[None] = []
|
||||
|
||||
_reserve_with(monkeypatch, SimpleNamespace(reserve_process_for_forking=lambda: calls.append(None)))
|
||||
fork_guard.reserve_process_for_forking("the gunicorn master", reserve=lambda: calls.append(None))
|
||||
|
||||
assert calls == [None]
|
||||
|
||||
|
||||
def test_used_extension_refuses_and_names_the_place(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def test_used_extension_refuses_and_names_the_place() -> None:
|
||||
def reserve() -> None:
|
||||
raise RuntimeError("the native runtime already started in this process")
|
||||
|
||||
with pytest.raises(fork_guard.NativeStateStartedBeforeFork, match="the gunicorn master") as raised:
|
||||
_reserve_with(monkeypatch, SimpleNamespace(reserve_process_for_forking=reserve))
|
||||
fork_guard.reserve_process_for_forking("the gunicorn master", reserve=reserve)
|
||||
|
||||
assert isinstance(raised.value.__cause__, RuntimeError)
|
||||
|
|
|
|||
|
|
@ -80,7 +80,6 @@ def test_native_logging_observes_level_changes(caplog: pytest.LogCaptureFixture)
|
|||
),
|
||||
)
|
||||
def test_native_credential_patterns_match_python(text: str) -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
from litellm.rust_bridge._native import NativeDiagnosticProcessor
|
||||
|
||||
processor: Final = NativeDiagnosticProcessor(MINIMUM_CUSTOM_KEY_LENGTH)
|
||||
|
|
@ -91,7 +90,6 @@ def test_native_credential_patterns_match_python(text: str) -> None:
|
|||
|
||||
|
||||
def test_native_client_redaction_matches_python() -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
from litellm.rust_bridge._native import NativeDiagnosticProcessor
|
||||
|
||||
text: Final = "error at /etc/secrets/config on db.internal\nTraceback (most recent call last):\nsecret"
|
||||
|
|
@ -100,7 +98,6 @@ def test_native_client_redaction_matches_python() -> None:
|
|||
|
||||
|
||||
def test_native_diagnostic_batch_matches_python() -> None:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
from litellm.rust_bridge._native import NativeDiagnosticProcessor
|
||||
|
||||
message: Final = "é" * 110 + "sk-" + "q" * 48 + "界" * 1000
|
||||
|
|
|
|||
|
|
@ -1,7 +1,6 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable, Generator
|
||||
from types import SimpleNamespace
|
||||
from typing import Final, Protocol
|
||||
|
||||
import pytest
|
||||
|
|
@ -10,26 +9,14 @@ from litellm.exceptions import APIError
|
|||
from litellm.llms.base_llm.ocr.transformation import OCRResponse
|
||||
from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
|
||||
from litellm.rust_bridge import bindings, configuration, runtime
|
||||
from litellm.rust_bridge._native import RustBridgeDeclined, RustUpstreamError
|
||||
from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule
|
||||
from litellm.rust_bridge.configuration import Rollout
|
||||
from litellm.rust_bridge.execution import AsyncStream, Complete, StreamChunk, SyncStream
|
||||
|
||||
|
||||
class RustBridgeDeclined(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class RustUpstreamError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def native_exceptions(monkeypatch: pytest.MonkeyPatch) -> Generator[None]:
|
||||
native: Final = SimpleNamespace(
|
||||
RustBridgeDeclined=RustBridgeDeclined,
|
||||
RustUpstreamError=RustUpstreamError,
|
||||
)
|
||||
monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
|
||||
def rust_configuration(monkeypatch: pytest.MonkeyPatch) -> Generator[None]:
|
||||
monkeypatch.delenv("LITELLM_RUST", raising=False)
|
||||
configuration.reset_rust_configuration()
|
||||
yield
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ import litellm
|
|||
from litellm.constants import TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS
|
||||
from litellm.litellm_core_utils.token_counter import openai_tokenizer_encoding
|
||||
from litellm.proxy.spend_tracking.input_tokens import count_input_tokens_for_model
|
||||
from litellm.rust_bridge import _native as native
|
||||
from litellm.rust_bridge import token_counter as bridge
|
||||
from litellm.rust_bridge import tokenizer as tokenizer_dispatch
|
||||
from litellm.rust_bridge._native import Tokenizer
|
||||
|
|
@ -313,7 +314,6 @@ PARITY_MODELS: Final[tuple[tuple[str, bridge.RustTokenizer], ...]] = (
|
|||
async def test_native_count_matches_python_budget_counter(
|
||||
request_body: dict[str, object], model: str, tokenizer: bridge.RustTokenizer
|
||||
) -> None:
|
||||
native: Final = pytest.importorskip("litellm.rust_bridge._native")
|
||||
body: Final = json.dumps(request_body).replace(MODEL, model)
|
||||
parsed: Final = json.loads(body)
|
||||
|
||||
|
|
@ -328,7 +328,6 @@ async def test_tiktoken_counts_long_text_exactly_where_python_chunks(
|
|||
model: str, tokenizer: bridge.RustTokenizer
|
||||
) -> None:
|
||||
"""Python encodes tiktoken text in fixed-size chunks (drift of up to one token per chunk boundary); Rust does not."""
|
||||
native: Final = pytest.importorskip("litellm.rust_bridge._native")
|
||||
text: Final = "x " * 20_000
|
||||
body: Final = {"model": model, "messages": [{"role": "user", "content": text}]}
|
||||
encoding: Final = Tokenizer.from_tiktoken(tokenizer)
|
||||
|
|
@ -362,7 +361,6 @@ DECLINED_REQUESTS: Final[tuple[dict[str, object], ...]] = (
|
|||
async def test_native_declines_shapes_python_prices_differently(
|
||||
request_body: dict[str, object], tokenizer: bridge.RustTokenizer
|
||||
) -> None:
|
||||
native: Final = pytest.importorskip("litellm.rust_bridge._native")
|
||||
raw, _ = _counted(request_body, MODEL_BY_TOKENIZER[tokenizer])
|
||||
|
||||
with pytest.raises(native.RustBridgeDeclined):
|
||||
|
|
|
|||
|
|
@ -1,14 +1,14 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Generator
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.llms.bedrock.audio_transcription import BedrockAudioTranscriptionRustDispatch
|
||||
from litellm.rust_bridge import bindings, configuration
|
||||
from litellm.rust_bridge import configuration
|
||||
from litellm.rust_bridge._native import RustBridgeDeclined, RustUpstreamError
|
||||
from litellm.rust_bridge.public_call import NativeCall
|
||||
from litellm.rust_bridge.transcription.native import NATIVE_ATRANSCRIPTION, NATIVE_TRANSCRIPTION
|
||||
|
||||
|
|
@ -28,18 +28,8 @@ TRANSCRIPTION_FIELDS: Final = frozenset(
|
|||
)
|
||||
|
||||
|
||||
class RustBridgeDeclined(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class RustUpstreamError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def isolated_bridge(monkeypatch: pytest.MonkeyPatch) -> Generator[None]:
|
||||
native: Final = SimpleNamespace(RustBridgeDeclined=RustBridgeDeclined, RustUpstreamError=RustUpstreamError)
|
||||
monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
|
||||
monkeypatch.delenv("LITELLM_RUST", raising=False)
|
||||
configuration.reset_rust_configuration()
|
||||
yield
|
||||
|
|
|
|||
|
|
@ -862,8 +862,6 @@ def test_secret_filter_keeps_truncated_traceback(monkeypatch):
|
|||
|
||||
@pytest.mark.parametrize("native", (False, True), ids=("python", "rust"))
|
||||
def test_diagnostic_redaction_precedes_a_credential_cut(monkeypatch, native):
|
||||
if native:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
monkeypatch.setenv("LITELLM_RUST", "1" if native else "0")
|
||||
monkeypatch.setenv("MAX_STRING_LENGTH_STDOUT_LOG", "500")
|
||||
monkeypatch.setenv("MAX_BASE64_LENGTH_STDOUT_LOG", "0")
|
||||
|
|
@ -890,8 +888,6 @@ def test_correlation_id_redacts_before_its_length_bound(monkeypatch):
|
|||
|
||||
@pytest.mark.parametrize("native", (False, True), ids=("python", "rust"))
|
||||
def test_malformed_interpolation_still_scrubs_a_record(monkeypatch, native):
|
||||
if native:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
monkeypatch.setenv("LITELLM_RUST", "1" if native else "0")
|
||||
monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True)
|
||||
record = _make_record(logging.WARNING, "bad % api_key=secret123", ("value",))
|
||||
|
|
@ -905,8 +901,6 @@ def test_malformed_interpolation_still_scrubs_a_record(monkeypatch, native):
|
|||
|
||||
@pytest.mark.parametrize("native", (False, True), ids=("python", "rust"))
|
||||
def test_key_pattern_template_keeps_the_rendered_redacted_line(monkeypatch, native):
|
||||
if native:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
monkeypatch.setenv("LITELLM_RUST", "1" if native else "0")
|
||||
monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True)
|
||||
record = _make_record(logging.INFO, "password=%s ok", ("hunter2",))
|
||||
|
|
@ -1733,8 +1727,6 @@ def test_access_log_path_filter_keeps_a_record_without_a_string_path_arg(monkeyp
|
|||
|
||||
@pytest.mark.parametrize("native", (False, True), ids=("python", "rust"))
|
||||
def test_diagnostic_filter_scrubs_exc_stack_and_nested_extras(monkeypatch, native):
|
||||
if native:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
monkeypatch.setenv("LITELLM_RUST", "1" if native else "0")
|
||||
monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True)
|
||||
secret = "sk-" + "q" * 48
|
||||
|
|
@ -1761,8 +1753,6 @@ def test_diagnostic_filter_scrubs_exc_stack_and_nested_extras(monkeypatch, nativ
|
|||
|
||||
@pytest.mark.parametrize("native", (False, True), ids=("python", "rust"))
|
||||
def test_diagnostic_filter_stamps_records_so_a_second_pass_is_free(monkeypatch, native):
|
||||
if native:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
monkeypatch.setenv("LITELLM_RUST", "1" if native else "0")
|
||||
monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True)
|
||||
record = _make_record(logging.WARNING, "api_key=secret123")
|
||||
|
|
@ -1775,8 +1765,6 @@ def test_diagnostic_filter_stamps_records_so_a_second_pass_is_free(monkeypatch,
|
|||
|
||||
@pytest.mark.parametrize("native", (False, True), ids=("python", "rust"))
|
||||
def test_json_formatter_scrubs_unfiltered_extras(monkeypatch, native):
|
||||
if native:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
monkeypatch.setenv("LITELLM_RUST", "1" if native else "0")
|
||||
monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True)
|
||||
secret = "sk-" + "q" * 48
|
||||
|
|
@ -1791,8 +1779,6 @@ def test_json_formatter_scrubs_unfiltered_extras(monkeypatch, native):
|
|||
|
||||
@pytest.mark.parametrize("native", (False, True), ids=("python", "rust"))
|
||||
def test_diagnostic_filter_redacts_a_non_string_message_object(monkeypatch, native):
|
||||
if native:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
monkeypatch.setenv("LITELLM_RUST", "1" if native else "0")
|
||||
monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True)
|
||||
secret = "sk-" + "q" * 48
|
||||
|
|
|
|||
|
|
@ -3774,6 +3774,132 @@ async def test_continued_stream_sharing_the_proxy_logging_object_bills_the_cut_t
|
|||
)
|
||||
|
||||
|
||||
def _responses_sse(*events: Mapping[str, object]) -> str:
|
||||
return "".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events)
|
||||
|
||||
|
||||
def _responses_envelope(response_id: str, model: str, status: str, **fields: object) -> dict[str, object]:
|
||||
return {
|
||||
"id": response_id,
|
||||
"object": "response",
|
||||
"created_at": 1700000000,
|
||||
"status": status,
|
||||
"model": model,
|
||||
"output": [],
|
||||
**fields,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_responses_fallback_sharing_the_proxy_logging_object_logs_the_fallback_usage(
|
||||
monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
recorder: Final = _SpendRecorder()
|
||||
monkeypatch.setattr(litellm, "callbacks", [recorder])
|
||||
respx_mock.post("https://primary.example/v1/responses").mock(
|
||||
return_value=httpx.Response(
|
||||
200,
|
||||
headers={"content-type": "text/event-stream"},
|
||||
text=_responses_sse(
|
||||
{"type": "response.created", "response": _responses_envelope("resp_primary", "gpt-4o", "in_progress")},
|
||||
*(
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"item_id": "msg_primary",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": word,
|
||||
}
|
||||
for word in ("Once ", "upon ", "a ", "time ", "there")
|
||||
),
|
||||
{
|
||||
"type": "response.failed",
|
||||
"response": _responses_envelope(
|
||||
"resp_primary",
|
||||
"gpt-4o",
|
||||
"failed",
|
||||
error={"code": "server_error", "message": "The model had an internal error"},
|
||||
),
|
||||
},
|
||||
),
|
||||
)
|
||||
)
|
||||
respx_mock.post("https://backup.example/v1/responses").mock(
|
||||
return_value=httpx.Response(
|
||||
200,
|
||||
headers={"content-type": "text/event-stream"},
|
||||
text=_responses_sse(
|
||||
{"type": "response.created", "response": _responses_envelope("resp_backup", "gpt-4o-mini", "in_progress")},
|
||||
{
|
||||
"type": "response.completed",
|
||||
"response": _responses_envelope(
|
||||
"resp_backup",
|
||||
"gpt-4o-mini",
|
||||
"completed",
|
||||
output=[
|
||||
{
|
||||
"type": "message",
|
||||
"id": "msg_backup",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "Once upon a time", "annotations": []}],
|
||||
}
|
||||
],
|
||||
usage={"input_tokens": 1000, "output_tokens": 700, "total_tokens": 1700},
|
||||
),
|
||||
},
|
||||
),
|
||||
)
|
||||
)
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "primary",
|
||||
"litellm_params": {"model": "openai/gpt-4o", "api_key": "k1", "api_base": "https://primary.example/v1"},
|
||||
},
|
||||
{
|
||||
"model_name": "backup",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"api_key": "k2",
|
||||
"api_base": "https://backup.example/v1",
|
||||
},
|
||||
},
|
||||
],
|
||||
fallbacks=[{"primary": ["backup"]}],
|
||||
num_retries=0,
|
||||
)
|
||||
shared_logging, _ = litellm.utils.function_setup(
|
||||
original_function="aresponses",
|
||||
rules_obj=litellm.utils.Rules(),
|
||||
start_time=datetime.now(),
|
||||
model="primary",
|
||||
input="Tell me a story",
|
||||
stream=True,
|
||||
litellm_call_id="shared-call",
|
||||
)
|
||||
|
||||
response: Final = await router.aresponses(
|
||||
model="primary",
|
||||
input="Tell me a story",
|
||||
stream=True,
|
||||
litellm_call_id="shared-call",
|
||||
litellm_logging_obj=shared_logging,
|
||||
)
|
||||
_ = [event async for event in response]
|
||||
await asyncio.wait_for(recorder.succeeded.wait(), timeout=10)
|
||||
|
||||
success: Final = recorder.successes[-1]
|
||||
assert (success["prompt_tokens"], success["completion_tokens"]) == (1000, 700)
|
||||
spend_row: Final = recorder.spend_rows[-1]
|
||||
assert (spend_row["prompt_tokens"], spend_row["completion_tokens"], spend_row["spend"]) == (
|
||||
1000,
|
||||
700,
|
||||
success["response_cost"],
|
||||
)
|
||||
|
||||
|
||||
def test_carry_partial_spend_to_fallback_adds_the_partial_the_cut_stream_inherited():
|
||||
from litellm.types.utils import Usage
|
||||
|
||||
|
|
|
|||
|
|
@ -76,8 +76,6 @@ def test_redact_string_catches_secret_patterns():
|
|||
|
||||
@pytest.mark.parametrize("native", (False, True), ids=("python", "rust"))
|
||||
def test_diagnostic_redaction_policy_matches_across_backends(monkeypatch: pytest.MonkeyPatch, native: bool) -> None:
|
||||
if native:
|
||||
pytest.importorskip("litellm.rust_bridge._native")
|
||||
monkeypatch.setenv("LITELLM_RUST", "1" if native else "0")
|
||||
|
||||
assert redact_string("GET /v1?api_key=abcdefgh12345&page=2") == "GET /v1?REDACTED&page=2"
|
||||
|
|
|
|||
|
|
@ -87,10 +87,20 @@ def main(argv):
|
|||
f"::error::installing {os.path.basename(wheel)} into a deep prefix failed"
|
||||
)
|
||||
return 1
|
||||
if _run([python, "-c", "import litellm; import litellm.types.utils"]) != 0:
|
||||
if _run([python, "-I", "-c", "import litellm; import litellm.types.utils"]) != 0:
|
||||
print("::error::litellm did not import after install (half-unpacked package)")
|
||||
return 1
|
||||
|
||||
checker = os.path.abspath(os.path.join("tests", "base_sdk_tests", "check_base_sdk_install.py"))
|
||||
for rust in ("0", "1"):
|
||||
result = subprocess.run(
|
||||
[python, "-I", checker, "--profile", "native"],
|
||||
env={**os.environ, "LITELLM_RUST": rust, "LITELLM_LOCAL_MODEL_COST_MAP": "True"},
|
||||
check=False,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
return result.returncode
|
||||
|
||||
print(
|
||||
f"ok: {os.path.basename(wheel)} installs into a worst-case prefix and imports"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,21 @@
|
|||
import { render, screen } from "@testing-library/react";
|
||||
import { afterEach, describe, expect, it } from "vitest";
|
||||
import { setServerRootPath } from "@/lib/serverRootPath";
|
||||
import { AriaAvatar } from "./AriaAvatar";
|
||||
|
||||
describe("AriaAvatar", () => {
|
||||
afterEach(() => {
|
||||
setServerRootPath("/");
|
||||
});
|
||||
|
||||
it.each(["/", "/litellm"])("loads the bundled image when the server root is %s", (root) => {
|
||||
setServerRootPath(root);
|
||||
|
||||
render(<AriaAvatar />);
|
||||
|
||||
expect(screen.getByRole("presentation", { hidden: true })).toHaveAttribute(
|
||||
"src",
|
||||
"/_next/static/media/aria-agent.png",
|
||||
);
|
||||
});
|
||||
});
|
||||
|
|
@ -1,10 +1,10 @@
|
|||
import { resolveLogoSrc } from "@/lib/assetPaths";
|
||||
import ariaAgentImage from "../../../public/assets/aria-agent.png";
|
||||
import { cn } from "@/lib/cva.config";
|
||||
|
||||
export function AriaAvatar({ size = 28, className }: { size?: number; className?: string }) {
|
||||
return (
|
||||
<img
|
||||
src={resolveLogoSrc("/ui/assets/aria-agent.png")}
|
||||
src={ariaAgentImage.src}
|
||||
width={size}
|
||||
height={size}
|
||||
alt=""
|
||||
|
|
|
|||
6
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
6
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -59866,6 +59866,8 @@ export interface operations {
|
|||
purpose?: string | null;
|
||||
limit?: number | null;
|
||||
after?: string | null;
|
||||
after_id?: string | null;
|
||||
before_id?: string | null;
|
||||
};
|
||||
header?: never;
|
||||
path?: never;
|
||||
|
|
@ -79082,6 +79084,8 @@ export interface operations {
|
|||
purpose?: string | null;
|
||||
limit?: number | null;
|
||||
after?: string | null;
|
||||
after_id?: string | null;
|
||||
before_id?: string | null;
|
||||
};
|
||||
header?: never;
|
||||
path?: never;
|
||||
|
|
@ -86427,6 +86431,8 @@ export interface operations {
|
|||
purpose?: string | null;
|
||||
limit?: number | null;
|
||||
after?: string | null;
|
||||
after_id?: string | null;
|
||||
before_id?: string | null;
|
||||
};
|
||||
header?: never;
|
||||
path: {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue