mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-31 06:56:59 +00:00
Merge remote-tracking branch 'origin/main' into litellm_hosted_vllm_google_oidc_api_key
# Conflicts: # litellm/secret_managers/main.py # tests/unit/secret_managers/test_secret_managers_main.py
This commit is contained in:
commit
2087d8bd89
150 changed files with 10084 additions and 953 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__'
|
||||
|
||||
|
|
|
|||
|
|
@ -427,6 +427,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th
|
|||
| [Watsonx Text (`watsonx_text`)](https://docs.litellm.ai/docs/providers/watsonx) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [xAI (`xai`)](https://docs.litellm.ai/docs/providers/xai) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Xinference (`xinference`)](https://docs.litellm.ai/docs/providers/xinference) | | | | ✅ | | | | | | |
|
||||
| [ZeroGPU (`zerogpu`)](https://docs.litellm.ai/docs/providers/zerogpu) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
|
||||
[**Read the Docs**](https://docs.litellm.ai/docs/)
|
||||
|
||||
|
|
|
|||
|
|
@ -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")]
|
||||
|
|
|
|||
|
|
@ -113,6 +113,7 @@ const OPENAI_COMPATIBLE_PROVIDERS: &[&str] = &[
|
|||
"meta",
|
||||
"cognition",
|
||||
"scx-ai",
|
||||
"zerogpu",
|
||||
];
|
||||
|
||||
impl ExceptionFamily {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -681,6 +681,7 @@ moonshot_models: Set = set()
|
|||
publicai_models: Set = set()
|
||||
darkbloom_models: Set = set()
|
||||
openinfer_models: Final[Set[str]] = set()
|
||||
zerogpu_models: Final[Set[str]] = set()
|
||||
v0_models: Set = set()
|
||||
morph_models: Set = set()
|
||||
opencode_go_models: Final[set[str]] = set() # mutable-ok: populated when the shared model catalog loads
|
||||
|
|
@ -969,6 +970,8 @@ def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, obje
|
|||
darkbloom_models.add(key)
|
||||
elif value.get("litellm_provider") == "openinfer":
|
||||
openinfer_models.add(key)
|
||||
elif value.get("litellm_provider") == "zerogpu":
|
||||
zerogpu_models.add(key)
|
||||
elif value.get("litellm_provider") == "v0":
|
||||
v0_models.add(key)
|
||||
elif value.get("litellm_provider") == "morph":
|
||||
|
|
@ -1142,6 +1145,7 @@ model_list = list(
|
|||
| publicai_models
|
||||
| darkbloom_models
|
||||
| openinfer_models
|
||||
| zerogpu_models
|
||||
| v0_models
|
||||
| morph_models
|
||||
| opencode_go_models
|
||||
|
|
@ -1261,6 +1265,7 @@ def _build_models_by_provider() -> dict:
|
|||
"publicai": publicai_models,
|
||||
"darkbloom": darkbloom_models,
|
||||
"openinfer": openinfer_models,
|
||||
"zerogpu": zerogpu_models,
|
||||
"v0": v0_models,
|
||||
"morph": morph_models,
|
||||
"opencode_go": opencode_go_models,
|
||||
|
|
@ -1739,6 +1744,9 @@ if TYPE_CHECKING:
|
|||
from .llms.lm_studio.decisions.transformation import (
|
||||
LMStudioDecisionsConfig as LMStudioDecisionsConfig,
|
||||
)
|
||||
from .llms.ollama.decisions.transformation import (
|
||||
OllamaDecisionsConfig as OllamaDecisionsConfig,
|
||||
)
|
||||
from .llms.laya.decisions.transformation import (
|
||||
LayaDecisionsConfig as LayaDecisionsConfig,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -171,6 +171,7 @@ LLM_CONFIG_NAMES: Final = (
|
|||
"DatabricksDecisionsConfig",
|
||||
"HostedVLLMDecisionsConfig",
|
||||
"LMStudioDecisionsConfig",
|
||||
"OllamaDecisionsConfig",
|
||||
"LayaDecisionsConfig",
|
||||
"BespokeDecisionsConfig",
|
||||
"OpenAIDecisionsConfig",
|
||||
|
|
@ -749,6 +750,7 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
|
|||
"HostedVLLMDecisionsConfig",
|
||||
),
|
||||
"LMStudioDecisionsConfig": (".llms.lm_studio.decisions.transformation", "LMStudioDecisionsConfig"),
|
||||
"OllamaDecisionsConfig": (".llms.ollama.decisions.transformation", "OllamaDecisionsConfig"),
|
||||
"LayaDecisionsConfig": (".llms.laya.decisions.transformation", "LayaDecisionsConfig"),
|
||||
"BespokeDecisionsConfig": (".llms.bespoke.decisions.transformation", "BespokeDecisionsConfig"),
|
||||
"OpenAIDecisionsConfig": (".llms.openai.decisions.transformation", "OpenAIDecisionsConfig"),
|
||||
|
|
|
|||
|
|
@ -3,7 +3,6 @@ Handler for transforming /chat/completions api requests to litellm.responses req
|
|||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence
|
||||
from itertools import accumulate, chain
|
||||
from types import MappingProxyType
|
||||
|
|
@ -43,8 +42,8 @@ from litellm.responses.sse_output_recovery import (
|
|||
)
|
||||
from litellm.responses.utils import (
|
||||
ResponsesAPIRequestUtils,
|
||||
auto_reasoning_summary,
|
||||
normalize_responses_api_stream_options,
|
||||
opt_in_auto_reasoning_summary,
|
||||
)
|
||||
from litellm.types.llms.openai import (
|
||||
ChatCompletionAnnotation,
|
||||
|
|
@ -1745,19 +1744,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
**cast(Reasoning, reasoning_effort) # cast-ok: dict is forwarded verbatim to the provider
|
||||
)
|
||||
|
||||
# Check if auto-summary is enabled via flag or environment variable
|
||||
# Priority: litellm.reasoning_auto_summary flag > LITELLM_REASONING_AUTO_SUMMARY env var
|
||||
auto_summary_enabled: Final = (
|
||||
litellm.reasoning_auto_summary or os.getenv("LITELLM_REASONING_AUTO_SUMMARY", "false").lower() == "true"
|
||||
)
|
||||
|
||||
summary: Final = opt_in_auto_reasoning_summary(model, custom_llm_provider)
|
||||
return (
|
||||
Reasoning(
|
||||
effort=reasoning_effort,
|
||||
summary="detailed" if model is None else auto_reasoning_summary(model, custom_llm_provider),
|
||||
)
|
||||
if auto_summary_enabled
|
||||
else Reasoning(effort=reasoning_effort)
|
||||
Reasoning(effort=reasoning_effort)
|
||||
if summary is None
|
||||
else Reasoning(effort=reasoning_effort, summary=summary)
|
||||
)
|
||||
|
||||
_map_reasoning_effort = map_reasoning_effort
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -1096,6 +1097,7 @@ openai_compatible_providers: Final[list[str]] = [
|
|||
"llamafile",
|
||||
"lm_studio",
|
||||
"galadriel",
|
||||
"zerogpu",
|
||||
"github_copilot", # GitHub Copilot Chat API
|
||||
"chatgpt", # ChatGPT subscription API
|
||||
"novita",
|
||||
|
|
|
|||
|
|
@ -32,11 +32,11 @@ Usage:
|
|||
|
||||
# Create a managed agent on the provider side
|
||||
result = litellm.interactions.agents.create(
|
||||
name="waverunner",
|
||||
id="waverunner",
|
||||
custom_llm_provider="gemini",
|
||||
api_key="...",
|
||||
base_agent="gemini-2.5-flash",
|
||||
instructions="You are a helpful assistant.",
|
||||
base_agent="antigravity-preview-09-2026",
|
||||
system_instruction="You are a helpful assistant.",
|
||||
)
|
||||
|
||||
Methods:
|
||||
|
|
|
|||
|
|
@ -3,11 +3,11 @@ litellm.interactions.agents
|
|||
|
||||
Full CRUD SDK for provider-side managed agents (e.g. Gemini v1beta/agents).
|
||||
|
||||
litellm.interactions.agents.create(name=..., ...)
|
||||
litellm.interactions.agents.create(id=..., ...)
|
||||
litellm.interactions.agents.list(api_key=...)
|
||||
litellm.interactions.agents.get(name=..., ...)
|
||||
litellm.interactions.agents.delete(name=..., ...)
|
||||
litellm.interactions.agents.list_versions(name=..., ...)
|
||||
litellm.interactions.agents.get(id=..., ...)
|
||||
litellm.interactions.agents.delete(id=..., ...)
|
||||
litellm.interactions.agents.list_versions(id=..., ...)
|
||||
|
||||
Async counterparts: acreate, alist, aget, adelete, alist_versions
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
|||
from litellm.llms.base_llm.agents.transformation import BaseAgentsAPIConfig
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.types.agents import (
|
||||
AgentCreateRequest,
|
||||
AgentCreateResponse,
|
||||
AgentDeleteResult,
|
||||
AgentListResponse,
|
||||
|
|
@ -35,7 +36,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler):
|
|||
def create_agent(
|
||||
self,
|
||||
agents_api_config: BaseAgentsAPIConfig,
|
||||
name: str,
|
||||
create_request: AgentCreateRequest,
|
||||
litellm_params: GenericLiteLLMParams,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
extra_headers: dict[str, str] | None = None,
|
||||
|
|
@ -47,7 +48,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler):
|
|||
if _is_async:
|
||||
return self.async_create_agent(
|
||||
agents_api_config=agents_api_config,
|
||||
name=name,
|
||||
create_request=create_request,
|
||||
litellm_params=litellm_params,
|
||||
logging_obj=logging_obj,
|
||||
extra_headers=extra_headers,
|
||||
|
|
@ -63,12 +64,13 @@ class AgentsHTTPHandler(InteractionsHTTPHandler):
|
|||
api_base=litellm_params.get("api_base"),
|
||||
litellm_params=dict(litellm_params),
|
||||
)
|
||||
data: Final = agents_api_config.transform_create_request(name=name, litellm_params=dict(litellm_params))
|
||||
if extra_body:
|
||||
data.update(extra_body)
|
||||
data: Final = {
|
||||
**agents_api_config.transform_create_request(create_request),
|
||||
**(extra_body or {}),
|
||||
}
|
||||
|
||||
logging_obj.pre_call(
|
||||
input=name,
|
||||
input=create_request.id,
|
||||
api_key="",
|
||||
additional_args={
|
||||
"complete_input_dict": data,
|
||||
|
|
@ -85,12 +87,12 @@ class AgentsHTTPHandler(InteractionsHTTPHandler):
|
|||
original_response=response.text,
|
||||
additional_args={"complete_input_dict": data},
|
||||
)
|
||||
return agents_api_config.transform_create_response(raw_response=response, name=name)
|
||||
return agents_api_config.transform_create_response(raw_response=response, name=create_request.id)
|
||||
|
||||
async def async_create_agent(
|
||||
self,
|
||||
agents_api_config: BaseAgentsAPIConfig,
|
||||
name: str,
|
||||
create_request: AgentCreateRequest,
|
||||
litellm_params: GenericLiteLLMParams,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
extra_headers: dict[str, str] | None = None,
|
||||
|
|
@ -106,12 +108,13 @@ class AgentsHTTPHandler(InteractionsHTTPHandler):
|
|||
api_base=litellm_params.get("api_base"),
|
||||
litellm_params=dict(litellm_params),
|
||||
)
|
||||
data: Final = agents_api_config.transform_create_request(name=name, litellm_params=dict(litellm_params))
|
||||
if extra_body:
|
||||
data.update(extra_body)
|
||||
data: Final = {
|
||||
**agents_api_config.transform_create_request(create_request),
|
||||
**(extra_body or {}),
|
||||
}
|
||||
|
||||
logging_obj.pre_call(
|
||||
input=name,
|
||||
input=create_request.id,
|
||||
api_key="",
|
||||
additional_args={
|
||||
"complete_input_dict": data,
|
||||
|
|
@ -130,7 +133,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler):
|
|||
original_response=response.text,
|
||||
additional_args={"complete_input_dict": data},
|
||||
)
|
||||
return agents_api_config.transform_create_response(raw_response=response, name=name)
|
||||
return agents_api_config.transform_create_response(raw_response=response, name=create_request.id)
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# LIST #
|
||||
|
|
|
|||
|
|
@ -6,41 +6,47 @@ Usage:
|
|||
|
||||
# Create
|
||||
response = litellm.interactions.agents.create(
|
||||
name="waverunner",
|
||||
id="waverunner",
|
||||
custom_llm_provider="gemini",
|
||||
api_key="...",
|
||||
base_agent="gemini-2.5-flash",
|
||||
instructions="You are a helpful assistant.",
|
||||
base_agent="antigravity-preview-09-2026",
|
||||
system_instruction="You are a helpful assistant.",
|
||||
agent_config={"type": "antigravity", "model": "gemini-3.8-flash"},
|
||||
)
|
||||
|
||||
# List
|
||||
response = litellm.interactions.agents.list(api_key="...", custom_llm_provider="gemini")
|
||||
|
||||
# Get
|
||||
response = litellm.interactions.agents.get(name="waverunner", api_key="...")
|
||||
response = litellm.interactions.agents.get(id="waverunner", api_key="...")
|
||||
|
||||
# Delete
|
||||
result = litellm.interactions.agents.delete(name="waverunner", api_key="...")
|
||||
result = litellm.interactions.agents.delete(id="waverunner", api_key="...")
|
||||
|
||||
# List versions
|
||||
result = litellm.interactions.agents.list_versions(name="waverunner", api_key="...")
|
||||
result = litellm.interactions.agents.list_versions(id="waverunner", api_key="...")
|
||||
|
||||
# Async versions: acreate, alist, aget, adelete, alist_versions
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import contextvars
|
||||
from collections.abc import Coroutine, Mapping
|
||||
from collections.abc import Coroutine, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final
|
||||
|
||||
import httpx
|
||||
from pydantic import ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm.interactions.agents.http_handler import agents_http_handler
|
||||
from litellm.interactions.agents.utils import get_provider_agents_api_config
|
||||
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.types.agents import (
|
||||
AgentCreateRequest,
|
||||
AgentCreateResponse,
|
||||
AgentDeleteResult,
|
||||
AgentListResponse,
|
||||
|
|
@ -89,14 +95,120 @@ def _make_logging_obj(
|
|||
return litellm_logging_obj
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _InvalidAgentRequest:
|
||||
message: str
|
||||
|
||||
|
||||
def _alias_conflict(field: str, value: str | None, alias: str, alias_value: str | None) -> str | None:
|
||||
if value is None or alias_value is None or value == alias_value:
|
||||
return None
|
||||
return f"'{alias}' is an alias of '{field}', but the request gives them different values. Pass only '{field}'."
|
||||
|
||||
|
||||
def _agent_id_or_error(id: str | None, name: str | None) -> str | _InvalidAgentRequest:
|
||||
conflict: Final = _alias_conflict("id", id, "name", name)
|
||||
if conflict is not None:
|
||||
return _InvalidAgentRequest(conflict)
|
||||
agent_id: Final = id if id is not None else name
|
||||
if agent_id is None:
|
||||
return _InvalidAgentRequest("Missing required parameter: 'id' (or its alias 'name').")
|
||||
return agent_id
|
||||
|
||||
|
||||
def _required_agent_id(id: str | None, name: str | None, llm_provider: str) -> str:
|
||||
agent_id: Final = _agent_id_or_error(id, name)
|
||||
if isinstance(agent_id, _InvalidAgentRequest):
|
||||
raise litellm.BadRequestError(message=agent_id.message, model=id or name or "", llm_provider=llm_provider)
|
||||
return agent_id
|
||||
|
||||
|
||||
# ================================================================== #
|
||||
# CREATE #
|
||||
# ================================================================== #
|
||||
|
||||
|
||||
_CREATE_FIELD_SHAPES: Final = MappingProxyType(
|
||||
{
|
||||
"id": "a string",
|
||||
"base_agent": "a string",
|
||||
"agent_config": "an object",
|
||||
"system_instruction": "a string",
|
||||
"base_environment": "an environment id string or an object",
|
||||
"description": "a string",
|
||||
"tools": "a list of objects",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _invalid_create_fields_message(invalid_fields: frozenset[str], caller_names: Mapping[str, str]) -> str:
|
||||
problems: Final = "; ".join(
|
||||
f"'{caller_names.get(field, field)}' must be {_CREATE_FIELD_SHAPES[field]}" for field in sorted(invalid_fields)
|
||||
)
|
||||
drop_hint: Final = "" if "id" in invalid_fields else " Set drop_params to drop invalid optional fields."
|
||||
return f"Invalid agent create request: {problems}.{drop_hint}"
|
||||
|
||||
|
||||
def _validated_create_request(
|
||||
fields: Mapping[str, object], drop_invalid_fields: bool, caller_names: Mapping[str, str]
|
||||
) -> AgentCreateRequest | _InvalidAgentRequest:
|
||||
try:
|
||||
return AgentCreateRequest.model_validate(fields)
|
||||
except ValidationError as e:
|
||||
invalid_fields: Final = frozenset(str(error["loc"][0]) for error in e.errors())
|
||||
if drop_invalid_fields and "id" not in invalid_fields:
|
||||
return AgentCreateRequest.model_validate(
|
||||
{field: value for field, value in fields.items() if field not in invalid_fields}
|
||||
)
|
||||
return _InvalidAgentRequest(_invalid_create_fields_message(invalid_fields, caller_names))
|
||||
|
||||
|
||||
def _create_request_or_error(
|
||||
id: str | None,
|
||||
name: str | None,
|
||||
base_agent: str | None,
|
||||
agent_config: Mapping[str, object] | None,
|
||||
system_instruction: str | None,
|
||||
instructions: str | None,
|
||||
base_environment: InteractionEnvironment | None,
|
||||
description: str | None,
|
||||
tools: Sequence[Mapping[str, object]] | None,
|
||||
drop_invalid_fields: bool,
|
||||
) -> AgentCreateRequest | _InvalidAgentRequest:
|
||||
agent_id: Final = _agent_id_or_error(id, name)
|
||||
if isinstance(agent_id, _InvalidAgentRequest):
|
||||
return agent_id
|
||||
instruction_conflict: Final = _alias_conflict(
|
||||
"system_instruction", system_instruction, "instructions", instructions
|
||||
)
|
||||
if instruction_conflict is not None and not drop_invalid_fields:
|
||||
return _InvalidAgentRequest(instruction_conflict)
|
||||
sent_only_instructions: Final = system_instruction is None and instructions is not None
|
||||
return _validated_create_request(
|
||||
MappingProxyType(
|
||||
{
|
||||
"id": agent_id,
|
||||
"base_agent": base_agent,
|
||||
"agent_config": agent_config,
|
||||
"system_instruction": instructions if sent_only_instructions else system_instruction,
|
||||
"base_environment": base_environment,
|
||||
"description": description,
|
||||
"tools": tools,
|
||||
}
|
||||
),
|
||||
drop_invalid_fields,
|
||||
MappingProxyType(
|
||||
{
|
||||
"id": "name" if id is None else "id",
|
||||
"system_instruction": "instructions" if sent_only_instructions else "system_instruction",
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@client
|
||||
async def acreate(
|
||||
name: str,
|
||||
id: str | None = None,
|
||||
base_agent: str | None = None,
|
||||
instructions: str | None = None,
|
||||
base_environment: InteractionEnvironment | None = None,
|
||||
|
|
@ -104,6 +216,12 @@ async def acreate(
|
|||
extra_headers: dict[str, str] | None = None,
|
||||
extra_body: Mapping[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
*,
|
||||
name: str | None = None,
|
||||
agent_config: Mapping[str, object] | None = None,
|
||||
system_instruction: str | None = None,
|
||||
description: str | None = None,
|
||||
tools: Sequence[Mapping[str, object]] | None = None,
|
||||
**kwargs,
|
||||
) -> AgentCreateResponse:
|
||||
"""Async: Create a managed agent on the provider side."""
|
||||
|
|
@ -113,7 +231,7 @@ async def acreate(
|
|||
kwargs["acreate_agent"] = True
|
||||
func: Final = partial(
|
||||
create,
|
||||
name=name,
|
||||
id=id,
|
||||
base_agent=base_agent,
|
||||
instructions=instructions,
|
||||
base_environment=base_environment,
|
||||
|
|
@ -121,6 +239,11 @@ async def acreate(
|
|||
extra_headers=extra_headers,
|
||||
extra_body=extra_body,
|
||||
timeout=timeout,
|
||||
name=name,
|
||||
agent_config=agent_config,
|
||||
system_instruction=system_instruction,
|
||||
description=description,
|
||||
tools=tools,
|
||||
**kwargs,
|
||||
)
|
||||
ctx: Final = contextvars.copy_context()
|
||||
|
|
@ -130,7 +253,7 @@ async def acreate(
|
|||
return init_response
|
||||
except Exception as e:
|
||||
raise litellm.exception_type(
|
||||
model=name,
|
||||
model=id or name or "",
|
||||
custom_llm_provider=custom_llm_provider or "gemini",
|
||||
original_exception=e,
|
||||
completion_kwargs=local_vars,
|
||||
|
|
@ -140,7 +263,7 @@ async def acreate(
|
|||
|
||||
@client
|
||||
def create(
|
||||
name: str,
|
||||
id: str | None = None,
|
||||
base_agent: str | None = None,
|
||||
instructions: str | None = None,
|
||||
base_environment: InteractionEnvironment | None = None,
|
||||
|
|
@ -148,40 +271,61 @@ def create(
|
|||
extra_headers: dict[str, str] | None = None,
|
||||
extra_body: Mapping[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
*,
|
||||
name: str | None = None,
|
||||
agent_config: Mapping[str, object] | None = None,
|
||||
system_instruction: str | None = None,
|
||||
description: str | None = None,
|
||||
tools: Sequence[Mapping[str, object]] | None = None,
|
||||
**kwargs,
|
||||
) -> AgentCreateResponse | Coroutine[object, object, AgentCreateResponse]:
|
||||
"""
|
||||
Sync: Create a managed agent on the provider side.
|
||||
|
||||
Args:
|
||||
name: Name for the agent (required).
|
||||
base_agent: Base agent to derive from (e.g. "waverunner").
|
||||
instructions: System instructions for the agent.
|
||||
id: Identifier for the agent. Exactly one of ``id`` or its alias ``name`` is required.
|
||||
base_agent: Base agent to derive from (e.g. "antigravity-preview-09-2026").
|
||||
instructions: Alias of ``system_instruction``.
|
||||
base_environment: Environment to fork from — an env_id string or a
|
||||
dict like ``{"type": "remote", "sources": [...]}``.
|
||||
custom_llm_provider: Provider to use, e.g. "gemini".
|
||||
extra_headers: Additional HTTP headers.
|
||||
extra_body: Additional request body fields.
|
||||
timeout: Request timeout.
|
||||
name: Alias of ``id``.
|
||||
agent_config: Agent runtime config, e.g. ``{"type": "antigravity", "model": "gemini-3.8-flash"}``.
|
||||
system_instruction: System instruction for the agent.
|
||||
description: Description of the agent for developers.
|
||||
tools: Tools available to the agent, e.g. ``[{"type": "code_execution"}]``.
|
||||
**kwargs: Forwarded to GenericLiteLLMParams (api_key, api_base, etc.).
|
||||
"""
|
||||
local_vars: Final = locals()
|
||||
custom_llm_provider = custom_llm_provider or kwargs.get("custom_llm_provider") or "gemini"
|
||||
custom_llm_provider = custom_llm_provider or "gemini"
|
||||
try:
|
||||
_is_async: Final = kwargs.pop("acreate_agent", False) is True
|
||||
if base_agent is not None:
|
||||
kwargs["base_agent"] = base_agent
|
||||
if instructions is not None:
|
||||
kwargs["instructions"] = instructions
|
||||
if base_environment is not None:
|
||||
kwargs["base_environment"] = base_environment
|
||||
create_request: Final = _create_request_or_error(
|
||||
id=id,
|
||||
name=name,
|
||||
base_agent=base_agent,
|
||||
agent_config=agent_config,
|
||||
system_instruction=system_instruction,
|
||||
instructions=instructions,
|
||||
base_environment=base_environment,
|
||||
description=description,
|
||||
tools=tools,
|
||||
drop_invalid_fields=drops_unsupported_params(kwargs.get("drop_params")),
|
||||
)
|
||||
if isinstance(create_request, _InvalidAgentRequest):
|
||||
raise litellm.BadRequestError(
|
||||
message=create_request.message, model=id or name or "", llm_provider=custom_llm_provider
|
||||
)
|
||||
kwargs.setdefault("custom_llm_provider", custom_llm_provider)
|
||||
litellm_params: Final = GenericLiteLLMParams(**kwargs)
|
||||
logging_obj: Final = _make_logging_obj(kwargs, name, custom_llm_provider, "create_agent", {})
|
||||
logging_obj: Final = _make_logging_obj(kwargs, create_request.id, custom_llm_provider, "create_agent", {})
|
||||
config: Final = _get_agents_api_config(custom_llm_provider)
|
||||
return agents_http_handler.create_agent(
|
||||
agents_api_config=config,
|
||||
name=name,
|
||||
create_request=create_request,
|
||||
litellm_params=litellm_params,
|
||||
logging_obj=logging_obj,
|
||||
extra_headers=extra_headers,
|
||||
|
|
@ -191,7 +335,7 @@ def create(
|
|||
)
|
||||
except Exception as e:
|
||||
raise litellm.exception_type(
|
||||
model=name,
|
||||
model=id or name or "",
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
original_exception=e,
|
||||
completion_kwargs=local_vars,
|
||||
|
|
@ -247,7 +391,7 @@ def list(
|
|||
) -> AgentListResponse | Coroutine[object, object, AgentListResponse]:
|
||||
"""Sync: List all agents on the provider side."""
|
||||
local_vars: Final = locals()
|
||||
custom_llm_provider = custom_llm_provider or kwargs.get("custom_llm_provider") or "gemini"
|
||||
custom_llm_provider = custom_llm_provider or "gemini"
|
||||
try:
|
||||
_is_async: Final = kwargs.pop("alist_agents", False) is True
|
||||
kwargs.setdefault("custom_llm_provider", custom_llm_provider)
|
||||
|
|
@ -279,19 +423,22 @@ def list(
|
|||
|
||||
@client
|
||||
async def aget(
|
||||
name: str,
|
||||
id: str | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
extra_headers: dict[str, str] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
*,
|
||||
name: str | None = None,
|
||||
**kwargs,
|
||||
) -> AgentCreateResponse:
|
||||
"""Async: Get a specific agent by name."""
|
||||
"""Async: Get a specific agent by id (``name`` is an alias of ``id``)."""
|
||||
local_vars: Final = locals()
|
||||
try:
|
||||
loop: Final = asyncio.get_event_loop()
|
||||
kwargs["aget_agent"] = True
|
||||
func: Final = partial(
|
||||
get,
|
||||
id=id,
|
||||
name=name,
|
||||
custom_llm_provider=custom_llm_provider or "gemini",
|
||||
extra_headers=extra_headers,
|
||||
|
|
@ -305,7 +452,7 @@ async def aget(
|
|||
return init_response
|
||||
except Exception as e:
|
||||
raise litellm.exception_type(
|
||||
model=name,
|
||||
model=id or name or "",
|
||||
custom_llm_provider=custom_llm_provider or "gemini",
|
||||
original_exception=e,
|
||||
completion_kwargs=local_vars,
|
||||
|
|
@ -315,24 +462,27 @@ async def aget(
|
|||
|
||||
@client
|
||||
def get(
|
||||
name: str,
|
||||
id: str | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
extra_headers: dict[str, str] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
*,
|
||||
name: str | None = None,
|
||||
**kwargs,
|
||||
) -> AgentCreateResponse | Coroutine[object, object, AgentCreateResponse]:
|
||||
"""Sync: Get a specific agent by name."""
|
||||
"""Sync: Get a specific agent by id (``name`` is an alias of ``id``)."""
|
||||
local_vars: Final = locals()
|
||||
custom_llm_provider = custom_llm_provider or kwargs.get("custom_llm_provider") or "gemini"
|
||||
custom_llm_provider = custom_llm_provider or "gemini"
|
||||
try:
|
||||
_is_async: Final = kwargs.pop("aget_agent", False) is True
|
||||
agent_id: Final = _required_agent_id(id, name, custom_llm_provider)
|
||||
kwargs.setdefault("custom_llm_provider", custom_llm_provider)
|
||||
litellm_params: Final = GenericLiteLLMParams(**kwargs)
|
||||
logging_obj: Final = _make_logging_obj(kwargs, name, custom_llm_provider, "get_agent", {"name": name})
|
||||
logging_obj: Final = _make_logging_obj(kwargs, agent_id, custom_llm_provider, "get_agent", {"name": agent_id})
|
||||
config: Final = _get_agents_api_config(custom_llm_provider)
|
||||
return agents_http_handler.get_agent(
|
||||
agents_api_config=config,
|
||||
name=name,
|
||||
name=agent_id,
|
||||
litellm_params=litellm_params,
|
||||
logging_obj=logging_obj,
|
||||
extra_headers=extra_headers,
|
||||
|
|
@ -341,7 +491,7 @@ def get(
|
|||
)
|
||||
except Exception as e:
|
||||
raise litellm.exception_type(
|
||||
model=name,
|
||||
model=id or name or "",
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
original_exception=e,
|
||||
completion_kwargs=local_vars,
|
||||
|
|
@ -356,19 +506,22 @@ def get(
|
|||
|
||||
@client
|
||||
async def adelete(
|
||||
name: str,
|
||||
id: str | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
extra_headers: dict[str, str] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
*,
|
||||
name: str | None = None,
|
||||
**kwargs,
|
||||
) -> AgentDeleteResult:
|
||||
"""Async: Delete a specific agent by name."""
|
||||
"""Async: Delete a specific agent by id (``name`` is an alias of ``id``)."""
|
||||
local_vars: Final = locals()
|
||||
try:
|
||||
loop: Final = asyncio.get_event_loop()
|
||||
kwargs["adelete_agent"] = True
|
||||
func: Final = partial(
|
||||
delete,
|
||||
id=id,
|
||||
name=name,
|
||||
custom_llm_provider=custom_llm_provider or "gemini",
|
||||
extra_headers=extra_headers,
|
||||
|
|
@ -382,7 +535,7 @@ async def adelete(
|
|||
return init_response
|
||||
except Exception as e:
|
||||
raise litellm.exception_type(
|
||||
model=name,
|
||||
model=id or name or "",
|
||||
custom_llm_provider=custom_llm_provider or "gemini",
|
||||
original_exception=e,
|
||||
completion_kwargs=local_vars,
|
||||
|
|
@ -392,24 +545,29 @@ async def adelete(
|
|||
|
||||
@client
|
||||
def delete(
|
||||
name: str,
|
||||
id: str | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
extra_headers: dict[str, str] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
*,
|
||||
name: str | None = None,
|
||||
**kwargs,
|
||||
) -> AgentDeleteResult | Coroutine[object, object, AgentDeleteResult]:
|
||||
"""Sync: Delete a specific agent by name."""
|
||||
"""Sync: Delete a specific agent by id (``name`` is an alias of ``id``)."""
|
||||
local_vars: Final = locals()
|
||||
custom_llm_provider = custom_llm_provider or kwargs.get("custom_llm_provider") or "gemini"
|
||||
custom_llm_provider = custom_llm_provider or "gemini"
|
||||
try:
|
||||
_is_async: Final = kwargs.pop("adelete_agent", False) is True
|
||||
agent_id: Final = _required_agent_id(id, name, custom_llm_provider)
|
||||
kwargs.setdefault("custom_llm_provider", custom_llm_provider)
|
||||
litellm_params: Final = GenericLiteLLMParams(**kwargs)
|
||||
logging_obj: Final = _make_logging_obj(kwargs, name, custom_llm_provider, "delete_agent", {"name": name})
|
||||
logging_obj: Final = _make_logging_obj(
|
||||
kwargs, agent_id, custom_llm_provider, "delete_agent", {"name": agent_id}
|
||||
)
|
||||
config: Final = _get_agents_api_config(custom_llm_provider)
|
||||
return agents_http_handler.delete_agent(
|
||||
agents_api_config=config,
|
||||
name=name,
|
||||
name=agent_id,
|
||||
litellm_params=litellm_params,
|
||||
logging_obj=logging_obj,
|
||||
extra_headers=extra_headers,
|
||||
|
|
@ -418,7 +576,7 @@ def delete(
|
|||
)
|
||||
except Exception as e:
|
||||
raise litellm.exception_type(
|
||||
model=name,
|
||||
model=id or name or "",
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
original_exception=e,
|
||||
completion_kwargs=local_vars,
|
||||
|
|
@ -433,19 +591,22 @@ def delete(
|
|||
|
||||
@client
|
||||
async def alist_versions(
|
||||
name: str,
|
||||
id: str | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
extra_headers: dict[str, str] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
*,
|
||||
name: str | None = None,
|
||||
**kwargs,
|
||||
) -> AgentVersionsResponse:
|
||||
"""Async: List versions of a specific agent."""
|
||||
"""Async: List versions of a specific agent by id (``name`` is an alias of ``id``)."""
|
||||
local_vars: Final = locals()
|
||||
try:
|
||||
loop: Final = asyncio.get_event_loop()
|
||||
kwargs["alist_agent_versions"] = True
|
||||
func: Final = partial(
|
||||
list_versions,
|
||||
id=id,
|
||||
name=name,
|
||||
custom_llm_provider=custom_llm_provider or "gemini",
|
||||
extra_headers=extra_headers,
|
||||
|
|
@ -459,7 +620,7 @@ async def alist_versions(
|
|||
return init_response
|
||||
except Exception as e:
|
||||
raise litellm.exception_type(
|
||||
model=name,
|
||||
model=id or name or "",
|
||||
custom_llm_provider=custom_llm_provider or "gemini",
|
||||
original_exception=e,
|
||||
completion_kwargs=local_vars,
|
||||
|
|
@ -469,24 +630,29 @@ async def alist_versions(
|
|||
|
||||
@client
|
||||
def list_versions(
|
||||
name: str,
|
||||
id: str | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
extra_headers: dict[str, str] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
*,
|
||||
name: str | None = None,
|
||||
**kwargs,
|
||||
) -> AgentVersionsResponse | Coroutine[object, object, AgentVersionsResponse]:
|
||||
"""Sync: List versions of a specific agent."""
|
||||
"""Sync: List versions of a specific agent by id (``name`` is an alias of ``id``)."""
|
||||
local_vars: Final = locals()
|
||||
custom_llm_provider = custom_llm_provider or kwargs.get("custom_llm_provider") or "gemini"
|
||||
custom_llm_provider = custom_llm_provider or "gemini"
|
||||
try:
|
||||
_is_async: Final = kwargs.pop("alist_agent_versions", False) is True
|
||||
agent_id: Final = _required_agent_id(id, name, custom_llm_provider)
|
||||
kwargs.setdefault("custom_llm_provider", custom_llm_provider)
|
||||
litellm_params: Final = GenericLiteLLMParams(**kwargs)
|
||||
logging_obj: Final = _make_logging_obj(kwargs, name, custom_llm_provider, "list_agent_versions", {"name": name})
|
||||
logging_obj: Final = _make_logging_obj(
|
||||
kwargs, agent_id, custom_llm_provider, "list_agent_versions", {"name": agent_id}
|
||||
)
|
||||
config: Final = _get_agents_api_config(custom_llm_provider)
|
||||
return agents_http_handler.list_agent_versions(
|
||||
agents_api_config=config,
|
||||
name=name,
|
||||
name=agent_id,
|
||||
litellm_params=litellm_params,
|
||||
logging_obj=logging_obj,
|
||||
extra_headers=extra_headers,
|
||||
|
|
@ -495,7 +661,7 @@ def list_versions(
|
|||
)
|
||||
except Exception as e:
|
||||
raise litellm.exception_type(
|
||||
model=name,
|
||||
model=id or name or "",
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
original_exception=e,
|
||||
completion_kwargs=local_vars,
|
||||
|
|
|
|||
|
|
@ -49,6 +49,7 @@ from litellm.interactions.utils import (
|
|||
InteractionsAPIRequestUtils,
|
||||
get_provider_interactions_api_config,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.types.interactions import (
|
||||
CancelInteractionResult,
|
||||
|
|
@ -100,6 +101,8 @@ async def acreate(
|
|||
timeout: float | httpx.Timeout | None = None,
|
||||
# LiteLLM params
|
||||
custom_llm_provider: str | None = None,
|
||||
*,
|
||||
agent_config: Mapping[str, object] | None = None,
|
||||
**kwargs,
|
||||
) -> InteractionsAPIResponse | AsyncIterator[InteractionsAPIStreamingResponse]:
|
||||
"""
|
||||
|
|
@ -121,6 +124,7 @@ async def acreate(
|
|||
string, or a config object such as
|
||||
``{"type": "remote", "sources": [...]}`` /
|
||||
``{"type": "remote", "network": {...}}``
|
||||
agent_config: Agent runtime config, e.g. ``{"type": "antigravity", "model": "gemini-3.8-flash"}``
|
||||
response_modalities: Requested response modalities (TEXT, IMAGE, AUDIO)
|
||||
response_format: JSON schema for response format
|
||||
response_mime_type: MIME type of the response
|
||||
|
|
@ -155,6 +159,7 @@ async def acreate(
|
|||
store=store,
|
||||
background=background,
|
||||
environment=environment,
|
||||
agent_config=agent_config,
|
||||
response_modalities=response_modalities,
|
||||
response_format=response_format,
|
||||
response_mime_type=response_mime_type,
|
||||
|
|
@ -225,6 +230,8 @@ def create(
|
|||
timeout: float | httpx.Timeout | None = None,
|
||||
# LiteLLM params
|
||||
custom_llm_provider: str | None = None,
|
||||
*,
|
||||
agent_config: Mapping[str, object] | None = None,
|
||||
**kwargs,
|
||||
) -> (
|
||||
InteractionsAPIResponse
|
||||
|
|
@ -250,6 +257,7 @@ def create(
|
|||
string, or a config object such as
|
||||
``{"type": "remote", "sources": [...]}`` /
|
||||
``{"type": "remote", "network": {...}}``
|
||||
agent_config: Agent runtime config, e.g. ``{"type": "antigravity", "model": "gemini-3.8-flash"}``
|
||||
response_modalities: Requested response modalities (TEXT, IMAGE, AUDIO)
|
||||
response_format: JSON schema for response format
|
||||
response_mime_type: MIME type of the response
|
||||
|
|
@ -300,6 +308,15 @@ def create(
|
|||
# Check if this is a bridge provider (litellm_responses) - similar to responses API
|
||||
# Either provider is explicitly "litellm_responses" or no config found (bridge to responses)
|
||||
if custom_llm_provider == "litellm_responses" or interactions_api_config is None:
|
||||
if optional_params.get("agent_config") is not None and not drops_unsupported_params(
|
||||
kwargs.get("drop_params")
|
||||
):
|
||||
raise litellm.UnsupportedParamsError(
|
||||
message=f"'agent_config' is not supported for {custom_llm_provider} interactions. "
|
||||
"Set drop_params to drop it.",
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
)
|
||||
# Bridge to litellm.responses() for non-native providers
|
||||
from litellm.interactions.litellm_responses_transformation.handler import (
|
||||
LiteLLMResponsesInteractionsHandler,
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -2,8 +2,12 @@
|
|||
Utils used for litellm.transcription() and litellm.atranscription()
|
||||
"""
|
||||
|
||||
import base64
|
||||
import binascii
|
||||
import hashlib
|
||||
import math
|
||||
import os
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
|
|
@ -331,6 +335,45 @@ def calculate_request_duration(file: FileTypes) -> float | None:
|
|||
return None
|
||||
|
||||
|
||||
DEFAULT_AUDIO_TOKEN_ESTIMATE: Final = 300
|
||||
AUDIO_TOKENS_PER_SECOND: Final = 10
|
||||
AUDIO_BYTES_PER_TOKEN: Final = 1600
|
||||
_MAX_AUDIO_ESTIMATE_BYTES: Final = 16 * 1024 * 1024
|
||||
|
||||
|
||||
def _audio_token_rate(model: str) -> int | None:
|
||||
if model.startswith(
|
||||
("gemini-", "gemini/gemini-", "vertex_ai/gemini-", "vertex_ai_beta/gemini-", "openrouter/google/gemini-")
|
||||
):
|
||||
family: Final = model.rsplit("/", 1)[-1]
|
||||
return 25 if family.startswith(("gemini-3-", "gemini-3.")) else 32
|
||||
if model.startswith(("gpt-", "openai/gpt-", "azure/gpt-", "openrouter/openai/gpt-")):
|
||||
return AUDIO_TOKENS_PER_SECOND
|
||||
return None
|
||||
|
||||
|
||||
def estimate_audio_tokens(data: str | None, models: Sequence[str] = (), fallback_model: str | None = None) -> int:
|
||||
"""Estimate audio tokens from container duration within a bounded decoding budget."""
|
||||
if not data:
|
||||
return DEFAULT_AUDIO_TOKEN_ESTIMATE
|
||||
fallback_estimate: Final = max(len(data) * 3 // 4 // AUDIO_BYTES_PER_TOKEN, DEFAULT_AUDIO_TOKEN_ESTIMATE)
|
||||
padding: Final = 2 if data.endswith("==") else 1 if data.endswith("=") else 0
|
||||
if len(data) * 3 // 4 - padding > _MAX_AUDIO_ESTIMATE_BYTES:
|
||||
return fallback_estimate
|
||||
try:
|
||||
duration: Final = calculate_request_duration(base64.b64decode(data, validate=True))
|
||||
except (ValueError, binascii.Error):
|
||||
return fallback_estimate
|
||||
if duration is None or not math.isfinite(duration) or duration <= 0:
|
||||
return fallback_estimate
|
||||
fallback_rate: Final = _audio_token_rate(fallback_model) if fallback_model is not None else None
|
||||
rate: Final = max(
|
||||
(_audio_token_rate(model) or fallback_rate or AUDIO_TOKENS_PER_SECOND for model in models),
|
||||
default=fallback_rate or AUDIO_TOKENS_PER_SECOND,
|
||||
)
|
||||
return max(1, math.ceil(duration * rate))
|
||||
|
||||
|
||||
DEFAULT_SPEECH_MEDIA_TYPE: Final = "audio/mpeg"
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ import json
|
|||
import mimetypes
|
||||
import re
|
||||
from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from graphlib import CycleError, TopologicalSorter
|
||||
from itertools import chain, groupby, islice
|
||||
from os import PathLike
|
||||
|
|
@ -236,6 +237,160 @@ _IMAGE_CONTENT_PART_TYPES: Final = frozenset({"image_url", "input_image", "image
|
|||
_IMAGE_SCAN_MAX_DEPTH: Final = 4
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _TPMMediaProjection:
|
||||
content: object
|
||||
media: tuple[Mapping[str, object], ...] = ()
|
||||
|
||||
|
||||
def _tpm_image_part(url: str, detail: object = "auto") -> Mapping[str, object]:
|
||||
return {"type": "image_url", "image_url": {"url": url, "detail": detail}}
|
||||
|
||||
|
||||
def _tpm_blank_fields(fields: Mapping[str, object], names: tuple[str, ...]) -> Mapping[str, object]:
|
||||
return {key: "" if key in names and isinstance(value, str) else value for key, value in fields.items()}
|
||||
|
||||
|
||||
def _tpm_file_media(part: Mapping[str, object]) -> _TPMMediaProjection:
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import parse_mime_type
|
||||
|
||||
nested: Final = part.get("type") == "file"
|
||||
fields: Final = as_str_mapping(part.get("file")) if nested else part
|
||||
if fields is None:
|
||||
return _TPMMediaProjection(part)
|
||||
raw_data: Final = fields.get("file_data")
|
||||
data: Final = raw_data if isinstance(raw_data, str) else ""
|
||||
declared: Final = next(
|
||||
(
|
||||
value
|
||||
for key in ("format", "mime_type", "content_type")
|
||||
if isinstance(value := fields.get(key), str) and value
|
||||
),
|
||||
"",
|
||||
)
|
||||
mime: Final = (declared or parse_mime_type(data[:256]) or "").split(";", 1)[0].strip().lower()
|
||||
if not mime.startswith(("image/", "audio/")):
|
||||
return _TPMMediaProjection(part)
|
||||
raw_url: Final = fields.get("file_url") or fields.get("file_id")
|
||||
url: Final = data or (raw_url if isinstance(raw_url, str) else "")
|
||||
header_end: Final = data.find(",", 0, 256)
|
||||
encoded: Final = (
|
||||
data[header_end + 1 :]
|
||||
if header_end >= 0 and data[:header_end].endswith(";base64")
|
||||
else ""
|
||||
if data.startswith(("data:", "https://", "http://", "gs://"))
|
||||
else data
|
||||
)
|
||||
media: Final = (
|
||||
_tpm_image_part(
|
||||
url
|
||||
if not data or url.startswith(("data:", "https://", "http://", "gs://"))
|
||||
else f"data:{mime};base64,{url}",
|
||||
fields.get("detail", "auto"),
|
||||
)
|
||||
if mime.startswith("image/")
|
||||
else {"type": "input_audio", "input_audio": {"data": encoded, "format": mime.removeprefix("audio/")}}
|
||||
)
|
||||
projected: Final = _tpm_blank_fields(fields, ("file_data", "file_url"))
|
||||
return _TPMMediaProjection({**part, "file": projected} if nested else projected, (media,))
|
||||
|
||||
|
||||
def _tpm_media_leaf(value: object) -> _TPMMediaProjection:
|
||||
part: Final = as_str_mapping(value)
|
||||
if part is None:
|
||||
return _TPMMediaProjection(value)
|
||||
kind: Final = part.get("type")
|
||||
if kind in ("file", "input_file"):
|
||||
return _tpm_file_media(part)
|
||||
if kind == "input_audio":
|
||||
audio: Final = as_str_mapping(part.get("input_audio"))
|
||||
return _TPMMediaProjection(
|
||||
{**part, "input_audio": _tpm_blank_fields(audio, ("data",))} if audio is not None else value,
|
||||
({"type": "input_audio", "input_audio": part.get("input_audio")},),
|
||||
)
|
||||
if kind in ("image_url", "input_image", "image", "computer_screenshot"):
|
||||
image: Final = as_str_mapping(part.get("image_url"))
|
||||
source: Final = as_str_mapping(part.get("source"))
|
||||
canonical: Final = _tpm_image_part(
|
||||
_image_part_url({**part, "type": "input_image"} if kind == "computer_screenshot" else part) or "",
|
||||
(image or part).get("detail", part.get("detail", "auto")),
|
||||
)
|
||||
projected: Final = (
|
||||
{**part, "image_url": _tpm_blank_fields(image, ("url",))}
|
||||
if image is not None and kind in ("image_url", "input_image", "computer_screenshot")
|
||||
else {**part, "source": _tpm_blank_fields(source, ("data", "url"))}
|
||||
if source is not None and kind == "image" and source.get("type") in ("base64", "url")
|
||||
else _tpm_blank_fields(part, ("image_url", "url"))
|
||||
)
|
||||
return _TPMMediaProjection(projected, (canonical,))
|
||||
return _TPMMediaProjection(value)
|
||||
|
||||
|
||||
def _tpm_media_list(value: object, split: Callable[[object], _TPMMediaProjection]) -> _TPMMediaProjection:
|
||||
if not isinstance(value, (list, tuple)):
|
||||
return _TPMMediaProjection(value)
|
||||
parts: Final = cast(Sequence[object], value) # cast-ok: finite JSON content arrays
|
||||
projected: Final = tuple(map(split, parts))
|
||||
media: Final = tuple(chain.from_iterable(part.media for part in projected))
|
||||
content: Final = tuple(part.content for part in projected)
|
||||
return _TPMMediaProjection(parts if not media else content if isinstance(value, tuple) else list(content), media)
|
||||
|
||||
|
||||
def _tpm_replace_field(part: Mapping[str, object], key: str, projection: _TPMMediaProjection) -> _TPMMediaProjection:
|
||||
return _TPMMediaProjection({**part, key: projection.content} if projection.media else part, projection.media)
|
||||
|
||||
|
||||
def _tpm_document_media(value: object) -> _TPMMediaProjection:
|
||||
part: Final = as_str_mapping(value)
|
||||
source: Final = as_str_mapping(part.get("source")) if part is not None else None
|
||||
if part is None or part.get("type") != "document" or source is None or source.get("type") != "content":
|
||||
return _tpm_media_leaf(value)
|
||||
content: Final = _tpm_media_list(source.get("content"), _tpm_media_leaf)
|
||||
return _tpm_replace_field(part, "source", _tpm_replace_field(source, "content", content))
|
||||
|
||||
|
||||
def _tpm_tool_result_media(value: object) -> _TPMMediaProjection:
|
||||
part: Final = as_str_mapping(value)
|
||||
if part is None or part.get("type") != "tool_result":
|
||||
return _tpm_document_media(value)
|
||||
return _tpm_replace_field(part, "content", _tpm_media_list(part.get("content"), _tpm_document_media))
|
||||
|
||||
|
||||
def _tpm_code_output_media(value: object) -> _TPMMediaProjection:
|
||||
part: Final = as_str_mapping(value)
|
||||
url: Final = part.get("url") if part is not None else None
|
||||
if part is None or part.get("type") != "image" or not isinstance(url, str) or not url:
|
||||
return _TPMMediaProjection(value)
|
||||
return _TPMMediaProjection(_tpm_blank_fields(part, ("url",)), (_tpm_image_part(url),))
|
||||
|
||||
|
||||
def _tpm_input_media(value: object) -> _TPMMediaProjection:
|
||||
part: Final = as_str_mapping(value)
|
||||
if part is None:
|
||||
return _TPMMediaProjection(value)
|
||||
kind: Final = part.get("type")
|
||||
if kind in ("function_call_output", "custom_tool_call_output"):
|
||||
return _tpm_replace_field(part, "output", _tpm_media_list(part.get("output"), _tpm_tool_result_media))
|
||||
if kind == "computer_call_output":
|
||||
return _tpm_replace_field(part, "output", _tpm_media_leaf(part.get("output")))
|
||||
if kind == "code_interpreter_call":
|
||||
return _tpm_replace_field(part, "outputs", _tpm_media_list(part.get("outputs"), _tpm_code_output_media))
|
||||
result: Final = part.get("result")
|
||||
if kind == "image_generation_call" and isinstance(result, str) and result:
|
||||
return _TPMMediaProjection(
|
||||
_tpm_blank_fields(part, ("result",)),
|
||||
(_tpm_image_part("data:application/octet-stream;base64," + result),),
|
||||
)
|
||||
if kind == "message" or (kind is None and isinstance(part.get("role"), str)):
|
||||
return _tpm_replace_field(part, "content", _tpm_media_list(part.get("content"), _tpm_tool_result_media))
|
||||
return _tpm_tool_result_media(value)
|
||||
|
||||
|
||||
def split_media_for_tpm(value: object) -> tuple[object, tuple[Mapping[str, object], ...]]:
|
||||
projection: Final = _tpm_media_list(value, _tpm_input_media)
|
||||
return projection.content, projection.media
|
||||
|
||||
|
||||
def _content_parts_contain_image(parts: Sequence[object]) -> bool:
|
||||
"""Depth-bounded frontier walk over nested content lists, iterative because the repo bans
|
||||
recursion; an Anthropic tool_result nests its image parts exactly one level down."""
|
||||
|
|
|
|||
92
litellm/litellm_core_utils/single_flight.py
Normal file
92
litellm/litellm_core_utils/single_flight.py
Normal file
|
|
@ -0,0 +1,92 @@
|
|||
import threading
|
||||
from collections.abc import Callable
|
||||
from concurrent.futures import Future
|
||||
from dataclasses import dataclass, field
|
||||
from functools import lru_cache
|
||||
from typing import Final, Generic, ParamSpec, TypeVar
|
||||
|
||||
_P: Final = ParamSpec("_P")
|
||||
_R: Final = TypeVar("_R")
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class SingleFlight(Generic[_R]):
|
||||
lock: threading.Lock = field(default_factory=threading.Lock)
|
||||
in_flight: Future[_R] | None = None
|
||||
flights_started: int = 0
|
||||
last_flight_failed: bool = False
|
||||
|
||||
def run(self, fetch: Callable[[], _R], while_busy: _R | None = None) -> _R:
|
||||
leader: Final[Future[_R]] = Future()
|
||||
with self.lock:
|
||||
in_flight: Final = self.in_flight
|
||||
joined_flight: Final = self.flights_started
|
||||
retry_if_it_fails: Final = not self.last_flight_failed
|
||||
if in_flight is None:
|
||||
self._take_off(leader)
|
||||
if in_flight is None:
|
||||
return self._lead(leader, fetch)
|
||||
if while_busy is not None:
|
||||
return while_busy
|
||||
try:
|
||||
return in_flight.result()
|
||||
except Exception:
|
||||
retry: Final = self._claim_retry(joined_flight) if retry_if_it_fails else None
|
||||
if retry is None:
|
||||
raise
|
||||
retry_flight, leads_retry = retry
|
||||
return self._lead(retry_flight, fetch) if leads_retry else retry_flight.result()
|
||||
|
||||
def _take_off(self, leader: Future[_R]) -> None:
|
||||
self.in_flight = leader
|
||||
self.flights_started += 1
|
||||
|
||||
def _claim_retry(self, failed_flight: int) -> tuple[Future[_R], bool] | None:
|
||||
leader: Final[Future[_R]] = Future()
|
||||
with self.lock:
|
||||
in_flight: Final = self.in_flight
|
||||
if in_flight is not None:
|
||||
return in_flight, False
|
||||
if self.flights_started > failed_flight and self.last_flight_failed:
|
||||
return None
|
||||
self._take_off(leader)
|
||||
return leader, True
|
||||
|
||||
def _lead(self, flight: Future[_R], fetch: Callable[[], _R]) -> _R:
|
||||
try:
|
||||
result: Final = fetch()
|
||||
except BaseException as error:
|
||||
self._land(failed=True)
|
||||
flight.set_exception(error)
|
||||
raise
|
||||
self._land(failed=False)
|
||||
flight.set_result(result)
|
||||
return result
|
||||
|
||||
def _land(self, failed: bool) -> None:
|
||||
with self.lock:
|
||||
self.in_flight = None
|
||||
self.last_flight_failed = failed
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SingleFlightCache(Generic[_P, _R]):
|
||||
cached: Callable[..., _R]
|
||||
clear: Callable[[], None]
|
||||
lock: threading.Lock = field(default_factory=threading.Lock)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
*args: _P.args,
|
||||
**kwargs: _P.kwargs, # kwargs-ok: forwarded verbatim to the cached builder, typed by its own signature
|
||||
) -> _R:
|
||||
with self.lock:
|
||||
return self.cached(*args, **kwargs)
|
||||
|
||||
def cache_clear(self) -> None:
|
||||
self.clear()
|
||||
|
||||
|
||||
def single_flight_lru_cache(build: Callable[_P, _R]) -> SingleFlightCache[_P, _R]:
|
||||
cached: Final = lru_cache(maxsize=128)(build)
|
||||
return SingleFlightCache(cached, cached.cache_clear)
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -3,9 +3,12 @@
|
|||
import base64
|
||||
import io
|
||||
import math
|
||||
import re
|
||||
import struct
|
||||
from bisect import bisect_right
|
||||
from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
|
||||
from itertools import accumulate
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Literal, cast
|
||||
|
||||
import anyio
|
||||
|
|
@ -35,6 +38,7 @@ from litellm.constants import (
|
|||
TOKEN_COUNTER_MAX_EXACT_CHARS,
|
||||
)
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.litellm_core_utils.prompt_templates.chat_template_tools import (
|
||||
ChatTemplateToolSection,
|
||||
chat_template_tool_section,
|
||||
|
|
@ -188,22 +192,23 @@ def resize_image_high_res(
|
|||
|
||||
# Test the function with the given example
|
||||
def calculate_tiles_needed(
|
||||
resized_width,
|
||||
resized_height,
|
||||
tile_width=MAX_TILE_WIDTH,
|
||||
tile_height=MAX_TILE_HEIGHT,
|
||||
):
|
||||
resized_width: int,
|
||||
resized_height: int,
|
||||
tile_width: int = MAX_TILE_WIDTH,
|
||||
tile_height: int = MAX_TILE_HEIGHT,
|
||||
) -> int:
|
||||
tiles_across: Final = (resized_width + tile_width - 1) // tile_width
|
||||
tiles_down: Final = (resized_height + tile_height - 1) // tile_height
|
||||
total_tiles: Final = tiles_across * tiles_down
|
||||
return total_tiles
|
||||
|
||||
|
||||
def high_detail_image_token_upper_bound(base_tokens: int = 85) -> int:
|
||||
def high_detail_image_token_upper_bound(base_tokens: int = 85, tile_tokens: int | None = None) -> int:
|
||||
largest_tile_count: Final = calculate_tiles_needed(
|
||||
MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES, MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES
|
||||
)
|
||||
return base_tokens + (base_tokens * 2) * largest_tile_count
|
||||
per_tile: Final = base_tokens * 2 if tile_tokens is None else tile_tokens
|
||||
return base_tokens + per_tile * largest_tile_count
|
||||
|
||||
|
||||
def _unpack_ints(fmt: str, buffer: bytes) -> tuple[int, ...]:
|
||||
|
|
@ -381,6 +386,219 @@ def calculate_img_tokens(
|
|||
return total_tokens
|
||||
|
||||
|
||||
_IMAGE_HEADER_BASE64_CHARS: Final = 65_536
|
||||
_MAX_IMAGE_PATCHES: Final = 30_000
|
||||
_ANTHROPIC_IMAGE_LIMITS: Final = MappingProxyType(
|
||||
{"claude-standard": (ANTHROPIC_IMAGE_MAX_LONG_EDGE_PX, 1568), "claude-high": (2576, 4784)}
|
||||
)
|
||||
_TILE_IMAGE_TOKEN_COSTS: Final = MappingProxyType(
|
||||
{
|
||||
"gpt-4o": (85, 170),
|
||||
"gpt-4.1": (85, 170),
|
||||
"gpt-4o-mini": (2833, 5667),
|
||||
"gpt-5": (70, 140),
|
||||
"gpt-5.1": (70, 140),
|
||||
"o1": (75, 150),
|
||||
"o1-pro": (75, 150),
|
||||
"o3": (75, 150),
|
||||
}
|
||||
)
|
||||
_PATCH_IMAGE_MULTIPLIERS: Final = MappingProxyType(
|
||||
{
|
||||
"gpt-6-astra": 1.2,
|
||||
"gpt-5.6-sol": 1.2,
|
||||
"gpt-5.6-terra": 1.2,
|
||||
"gpt-5.6-luna": 1.2,
|
||||
"gpt-5.5": 1.2,
|
||||
"gpt-5.4": 1.2,
|
||||
"gpt-5.4-mini": 1.2,
|
||||
"gpt-5.4-nano": 1.2,
|
||||
"gpt-5.2": 1.2,
|
||||
"gpt-4.1-mini": 1.62,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _inline_image_dimensions(data: str) -> tuple[int, int] | None:
|
||||
separator: Final = data.find(",", 0, 256)
|
||||
if not data.startswith("data:") or separator < 0 or not data[:separator].endswith(";base64"):
|
||||
return None
|
||||
encoded_header: Final = data[separator + 1 : separator + 1 + _IMAGE_HEADER_BASE64_CHARS]
|
||||
try:
|
||||
header: Final = base64.b64decode(encoded_header, validate=True)
|
||||
image_type: Final = get_image_type(header)
|
||||
if image_type not in ("png", "gif", "jpeg", "webp"):
|
||||
return None
|
||||
if image_type == "webp" and header[12:16] not in (b"VP8X", b"VP8 ", b"VP8L"):
|
||||
return None
|
||||
dimensions: Final = _image_dimensions_from_bytes(header)
|
||||
except (ValueError, TypeError, struct.error):
|
||||
return None
|
||||
width, height = dimensions
|
||||
return (width, height) if width > 0 and height > 0 else None
|
||||
|
||||
|
||||
def _patch_image_limits(model: str, detail: str) -> tuple[int, int | None] | None:
|
||||
if model in ("gpt-5.2", "gpt-4.1-mini"):
|
||||
return None if detail == "original" else (2048, 6144)
|
||||
original_default: Final = model in ("gpt-6-astra", "gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna", "gpt-5.5")
|
||||
effective_detail: Final = ("original" if original_default else "high") if detail == "auto" else detail
|
||||
if effective_detail == "low":
|
||||
return (2048, 6144) if model in ("gpt-5.4", "gpt-5.4-mini", "gpt-5.4-nano") else (512, None)
|
||||
if effective_detail == "high":
|
||||
return (65_535 if model == "gpt-6-astra" else 2048, 2500)
|
||||
return (
|
||||
(65_535, None) if model in ("gpt-6-astra", "gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna") else (6000, 10_000)
|
||||
)
|
||||
|
||||
|
||||
def _image_patch_count(width: int, height: int, max_dimension: int, budget: int | None) -> int:
|
||||
dimension_scale: Final = min(1.0, max_dimension / max(width, height))
|
||||
fitted_width: Final = max(1, round(width * dimension_scale))
|
||||
fitted_height: Final = max(1, round(height * dimension_scale))
|
||||
patches: Final = math.ceil(fitted_width / 32) * math.ceil(fitted_height / 32)
|
||||
if budget is None or patches <= budget:
|
||||
return patches
|
||||
shrink: Final = math.sqrt(32**2 * budget / (fitted_width * fitted_height))
|
||||
width_patches: Final = fitted_width * shrink / 32
|
||||
height_patches: Final = fitted_height * shrink / 32
|
||||
aligned_shrink: Final = shrink * min(
|
||||
math.floor(width_patches) / width_patches, math.floor(height_patches) / height_patches
|
||||
)
|
||||
resized_width: Final = max(1, math.floor(fitted_width * aligned_shrink))
|
||||
resized_height: Final = max(1, math.floor(fitted_height * aligned_shrink))
|
||||
return math.ceil(resized_width / 32) * math.ceil(resized_height / 32)
|
||||
|
||||
|
||||
def _image_policy_model_name(model: str) -> str | None:
|
||||
provider, separator, provider_model = model.partition("/")
|
||||
if not separator:
|
||||
return model
|
||||
if provider == "openai":
|
||||
return provider_model
|
||||
if provider == "openrouter":
|
||||
return provider_model.removeprefix("openai/") if provider_model.startswith("openai/") else None
|
||||
if provider not in ("azure", "azure_ai"):
|
||||
return None
|
||||
catalog_entry: Final = as_str_mapping(
|
||||
cast( # cast-ok: read-only view of the legacy string-keyed registry; as_str_mapping checks the entry
|
||||
Mapping[str, object], litellm.model_cost
|
||||
).get(model)
|
||||
)
|
||||
if catalog_entry is not None and catalog_entry.get("litellm_provider") in ("azure", "azure_ai"):
|
||||
return provider_model.rsplit("/", 1)[-1]
|
||||
return provider_model.removeprefix("responses/").removeprefix("o_series/").removeprefix("openai/")
|
||||
|
||||
|
||||
def _anthropic_image_family(model: str) -> str | None:
|
||||
from litellm.utils import get_potential_model_names
|
||||
|
||||
provider, separator, _ = model.partition("/")
|
||||
if separator and provider not in (
|
||||
"anthropic",
|
||||
"bedrock",
|
||||
"bedrock_converse",
|
||||
"bedrock_mantle",
|
||||
"vertex_ai",
|
||||
"vertex_ai_beta",
|
||||
"azure_ai",
|
||||
"openrouter",
|
||||
):
|
||||
return None
|
||||
candidates: Final = get_potential_model_names(model, provider if separator else "anthropic")
|
||||
catalog: Final = cast(Mapping[str, object], litellm.model_cost) # cast-ok: read-only view of the legacy registry
|
||||
for candidate in dict.fromkeys(candidates.values()):
|
||||
if not isinstance(candidate, str):
|
||||
continue
|
||||
if (entry := as_str_mapping(catalog.get(candidate))) is None or entry.get("supports_vision") is not True:
|
||||
continue
|
||||
if (
|
||||
version := re.search(
|
||||
r"(?:^|[./])claude-(?:(?:haiku|sonnet|opus)-)?(\d+)(?:[-.](\d{1,2})(?!\d))?", candidate
|
||||
)
|
||||
) is None:
|
||||
continue
|
||||
major, minor = int(version[1]), int(version[2] or 0)
|
||||
return "claude-high" if (major, minor) >= (4, 7) else "claude-standard"
|
||||
return None
|
||||
|
||||
|
||||
def _anthropic_image_tokens(dimensions: tuple[int, int] | None, limits: tuple[int, int]) -> int:
|
||||
max_edge, token_budget = limits
|
||||
if dimensions is None:
|
||||
return token_budget
|
||||
long_edge, short_edge = max(dimensions), min(dimensions)
|
||||
aspect_ratio: Final = long_edge / short_edge
|
||||
|
||||
def tokens_at_edge(edge: int) -> int:
|
||||
return math.ceil(edge / 28) * math.ceil(max(round(edge / aspect_ratio), 1) / 28)
|
||||
|
||||
fitted_edge: Final = bisect_right(range(1, min(long_edge, max_edge) + 1), token_budget, key=tokens_at_edge)
|
||||
return tokens_at_edge(max(fitted_edge, 1))
|
||||
|
||||
|
||||
def _image_model_family(model: str) -> str | None:
|
||||
provider_model: Final = _image_policy_model_name(model)
|
||||
if provider_model is None:
|
||||
return _anthropic_image_family(model)
|
||||
family: Final = re.sub(r"-\d{4}-\d{2}-\d{2}$", "", provider_model)
|
||||
canonical: Final = litellm.azure_llms.get(f"azure/{family}", family) if model.startswith("azure/") else family
|
||||
return (
|
||||
canonical
|
||||
if canonical in _TILE_IMAGE_TOKEN_COSTS or canonical in _PATCH_IMAGE_MULTIPLIERS
|
||||
else _anthropic_image_family(model)
|
||||
)
|
||||
|
||||
|
||||
def _image_policy_tokens(family: str | None, dimensions: tuple[int, int] | None, detail: str) -> int:
|
||||
if family is None:
|
||||
return DEFAULT_IMAGE_TOKEN_COUNT
|
||||
anthropic_limits: Final = _ANTHROPIC_IMAGE_LIMITS.get(family)
|
||||
if anthropic_limits is not None:
|
||||
return _anthropic_image_tokens(dimensions, anthropic_limits)
|
||||
tile_costs: Final = _TILE_IMAGE_TOKEN_COSTS.get(family)
|
||||
multiplier: Final = _PATCH_IMAGE_MULTIPLIERS.get(family)
|
||||
if tile_costs is not None:
|
||||
base_tokens, tile_tokens = tile_costs
|
||||
if detail == "original":
|
||||
return DEFAULT_IMAGE_TOKEN_COUNT
|
||||
if detail == "low":
|
||||
return base_tokens
|
||||
if dimensions is None:
|
||||
return high_detail_image_token_upper_bound(base_tokens, tile_tokens)
|
||||
width, height = dimensions
|
||||
scale: Final = min(1.0, 2048 / max(width, height), 768 / min(width, height))
|
||||
tiles: Final = calculate_tiles_needed(max(1, math.floor(width * scale)), max(1, math.floor(height * scale)))
|
||||
return base_tokens + tiles * tile_tokens
|
||||
limits: Final = _patch_image_limits(family, detail)
|
||||
if limits is None or multiplier is None:
|
||||
return DEFAULT_IMAGE_TOKEN_COUNT
|
||||
max_dimension, patch_budget = limits
|
||||
patches: Final = (
|
||||
min(math.ceil(max_dimension / 32) ** 2, patch_budget if patch_budget is not None else _MAX_IMAGE_PATCHES)
|
||||
if dimensions is None
|
||||
else _image_patch_count(*dimensions, max_dimension, patch_budget)
|
||||
)
|
||||
return math.ceil(patches * multiplier)
|
||||
|
||||
|
||||
def estimate_image_tokens(
|
||||
model: str | Sequence[str], data: str, detail: str = "auto", *, fallback_model: str | None = None
|
||||
) -> int:
|
||||
"""Estimate image policies from at most 48 KiB of inline headers, without fetching media."""
|
||||
if detail not in ("low", "high", "auto", "original"):
|
||||
return DEFAULT_IMAGE_TOKEN_COUNT
|
||||
models: Final = (model,) if isinstance(model, str) else model
|
||||
fallback_family: Final = _image_model_family(fallback_model) if fallback_model is not None else None
|
||||
families: Final = tuple(
|
||||
dict.fromkeys(_image_model_family(candidate) or fallback_family for candidate in models or ("",))
|
||||
)
|
||||
if all(family is None for family in families):
|
||||
return DEFAULT_IMAGE_TOKEN_COUNT
|
||||
dimensions: Final = _inline_image_dimensions(data)
|
||||
return max(_image_policy_tokens(family, dimensions, detail) for family in families)
|
||||
|
||||
|
||||
TokenCounterFunction = Callable[[str], int]
|
||||
"""
|
||||
Type for a function that counts tokens in a string.
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -5,7 +5,6 @@ import os
|
|||
import re
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from contextlib import closing
|
||||
from functools import lru_cache
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal, NamedTuple, cast
|
||||
|
||||
|
|
@ -18,15 +17,18 @@ import litellm
|
|||
from litellm._logging import verbose_logger
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.constants import DEFAULT_MAX_RETRIES
|
||||
from litellm.litellm_core_utils.single_flight import SingleFlight, single_flight_lru_cache
|
||||
from litellm.litellm_core_utils.url_utils import ensure_api_route
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm.llms.openai.common_utils import BaseOpenAILLM, OpenAIAsyncHTTPClient, OpenAIHTTPClient
|
||||
from litellm.secret_managers.get_azure_ad_token_provider import (
|
||||
get_azure_ad_token_provider,
|
||||
single_flight_bearer_token_provider,
|
||||
)
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.secret_managers.get_azure_ad_token_provider import AzureCredentialType
|
||||
from litellm.utils import add_path_to_api_base
|
||||
|
||||
azure_ad_cache: Final = DualCache()
|
||||
|
|
@ -94,7 +96,7 @@ def process_azure_headers(headers: httpx.Headers | dict) -> dict:
|
|||
return {**llm_response_headers, **openai_headers}
|
||||
|
||||
|
||||
@lru_cache(maxsize=128)
|
||||
@single_flight_lru_cache
|
||||
def _cached_entra_id_token_provider(
|
||||
tenant_id: str,
|
||||
client_id: str,
|
||||
|
|
@ -107,33 +109,38 @@ def _cached_entra_id_token_provider(
|
|||
to expiry, so reusing the provider keeps one AAD round trip per token lifetime instead of one
|
||||
per request.
|
||||
"""
|
||||
from azure.identity import ClientSecretCredential, get_bearer_token_provider
|
||||
from azure.identity import ClientSecretCredential
|
||||
|
||||
return get_bearer_token_provider(ClientSecretCredential(tenant_id, client_id, client_secret), scope)
|
||||
return single_flight_bearer_token_provider(ClientSecretCredential(tenant_id, client_id, client_secret), scope)
|
||||
|
||||
|
||||
@lru_cache(maxsize=128)
|
||||
@single_flight_lru_cache
|
||||
def _cached_username_password_token_provider(
|
||||
client_id: str,
|
||||
azure_username: str,
|
||||
azure_password: str,
|
||||
scope: str,
|
||||
) -> Callable[[], str]:
|
||||
from azure.identity import UsernamePasswordCredential, get_bearer_token_provider
|
||||
from azure.identity import UsernamePasswordCredential
|
||||
|
||||
credential: Final = UsernamePasswordCredential(
|
||||
client_id=client_id,
|
||||
username=azure_username,
|
||||
password=azure_password,
|
||||
)
|
||||
return get_bearer_token_provider(credential, scope)
|
||||
return single_flight_bearer_token_provider(credential, scope)
|
||||
|
||||
|
||||
@lru_cache(maxsize=128)
|
||||
@single_flight_lru_cache
|
||||
def _cached_azure_ad_token_refresh_provider(scope: str) -> Callable[[], str]:
|
||||
return get_azure_ad_token_provider(azure_scope=scope)
|
||||
|
||||
|
||||
@single_flight_lru_cache
|
||||
def _cached_default_azure_credential_provider(scope: str) -> Callable[[], str]:
|
||||
return get_azure_ad_token_provider(azure_scope=scope, azure_credential=AzureCredentialType.DefaultAzureCredential)
|
||||
|
||||
|
||||
def cached_azure_ad_token_refresh_provider(scope: str) -> Callable[[], str]:
|
||||
return _cached_azure_ad_token_refresh_provider(scope)
|
||||
|
||||
|
|
@ -240,6 +247,58 @@ def get_azure_ad_token_from_username_password(
|
|||
return token_provider
|
||||
|
||||
|
||||
@single_flight_lru_cache
|
||||
def _oidc_token_exchange(cache_key: str) -> SingleFlight[str]:
|
||||
return SingleFlight()
|
||||
|
||||
|
||||
def _cached_oidc_token(cache_key: str) -> str | None:
|
||||
cached: Final = azure_ad_cache.get_cache(cache_key)
|
||||
return cached if isinstance(cached, str) else None
|
||||
|
||||
|
||||
def _exchange_oidc_token(cache_key: str, token_url: str, client_id: str, scope: str, oidc_token: str) -> str:
|
||||
cached_access_token: Final = _cached_oidc_token(cache_key)
|
||||
if cached_access_token is not None:
|
||||
return cached_access_token
|
||||
|
||||
with closing(HTTPHandler(timeout=litellm.request_timeout)) as client:
|
||||
req_token: Final = client.post(
|
||||
token_url,
|
||||
data={
|
||||
"client_id": client_id,
|
||||
"grant_type": "client_credentials",
|
||||
"scope": scope,
|
||||
"client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer",
|
||||
"client_assertion": oidc_token,
|
||||
},
|
||||
)
|
||||
|
||||
if req_token.status_code != 200:
|
||||
raise AzureOpenAIError(
|
||||
status_code=req_token.status_code,
|
||||
message=req_token.text,
|
||||
)
|
||||
|
||||
azure_ad_token_json: Final[_AzureAdTokenJson] = req_token.json()
|
||||
azure_ad_token_access_token: Final = azure_ad_token_json.get("access_token", None)
|
||||
azure_ad_token_expires_in: Final = azure_ad_token_json.get("expires_in", None)
|
||||
|
||||
if azure_ad_token_access_token is None:
|
||||
raise AzureOpenAIError(status_code=422, message="Azure AD Token access_token not returned")
|
||||
|
||||
if azure_ad_token_expires_in is None:
|
||||
raise AzureOpenAIError(status_code=422, message="Azure AD Token expires_in not returned")
|
||||
|
||||
azure_ad_cache.set_cache(
|
||||
key=cache_key,
|
||||
value=azure_ad_token_access_token,
|
||||
ttl=azure_ad_token_expires_in,
|
||||
)
|
||||
|
||||
return azure_ad_token_access_token
|
||||
|
||||
|
||||
def get_azure_ad_token_from_oidc(
|
||||
azure_ad_token: str,
|
||||
azure_client_id: str | None = None,
|
||||
|
|
@ -282,50 +341,25 @@ def get_azure_ad_token_from_oidc(
|
|||
"azure_client_id": azure_client_id,
|
||||
"azure_tenant_id": azure_tenant_id,
|
||||
"azure_authority_host": azure_authority_host,
|
||||
"scope": scope,
|
||||
"oidc_token": oidc_token,
|
||||
}
|
||||
)
|
||||
|
||||
azure_ad_token_access_token = azure_ad_cache.get_cache(azure_ad_token_cache_key)
|
||||
if azure_ad_token_access_token is not None:
|
||||
return azure_ad_token_access_token
|
||||
cached_access_token: Final = _cached_oidc_token(azure_ad_token_cache_key)
|
||||
if cached_access_token is not None:
|
||||
return cached_access_token
|
||||
|
||||
with closing(HTTPHandler(timeout=litellm.request_timeout)) as client:
|
||||
req_token: Final = client.post(
|
||||
f"{azure_authority_host}/{azure_tenant_id}/oauth2/v2.0/token",
|
||||
data={
|
||||
"client_id": azure_client_id,
|
||||
"grant_type": "client_credentials",
|
||||
"scope": scope,
|
||||
"client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer",
|
||||
"client_assertion": oidc_token,
|
||||
},
|
||||
return _oidc_token_exchange(azure_ad_token_cache_key).run(
|
||||
lambda: _exchange_oidc_token(
|
||||
azure_ad_token_cache_key,
|
||||
token_url=f"{azure_authority_host}/{azure_tenant_id}/oauth2/v2.0/token",
|
||||
client_id=azure_client_id,
|
||||
scope=scope,
|
||||
oidc_token=oidc_token,
|
||||
)
|
||||
|
||||
if req_token.status_code != 200:
|
||||
raise AzureOpenAIError(
|
||||
status_code=req_token.status_code,
|
||||
message=req_token.text,
|
||||
)
|
||||
|
||||
azure_ad_token_json: Final[_AzureAdTokenJson] = req_token.json()
|
||||
azure_ad_token_access_token = azure_ad_token_json.get("access_token", None)
|
||||
azure_ad_token_expires_in: Final = azure_ad_token_json.get("expires_in", None)
|
||||
|
||||
if azure_ad_token_access_token is None:
|
||||
raise AzureOpenAIError(status_code=422, message="Azure AD Token access_token not returned")
|
||||
|
||||
if azure_ad_token_expires_in is None:
|
||||
raise AzureOpenAIError(status_code=422, message="Azure AD Token expires_in not returned")
|
||||
|
||||
azure_ad_cache.set_cache(
|
||||
key=azure_ad_token_cache_key,
|
||||
value=azure_ad_token_access_token,
|
||||
ttl=azure_ad_token_expires_in,
|
||||
)
|
||||
|
||||
return azure_ad_token_access_token
|
||||
|
||||
|
||||
def select_azure_base_url_or_endpoint(azure_client_params: dict):
|
||||
azure_endpoint: Final = azure_client_params.get("azure_endpoint", None)
|
||||
|
|
@ -503,17 +537,10 @@ class BaseAzureLLM(BaseOpenAILLM):
|
|||
Returns:
|
||||
Token provider callable if DefaultAzureCredential is enabled and available, None otherwise
|
||||
"""
|
||||
from litellm.types.secret_managers.get_azure_ad_token_provider import (
|
||||
AzureCredentialType,
|
||||
)
|
||||
|
||||
verbose_logger.debug("Attempting to use DefaultAzureCredential for Azure Auth")
|
||||
|
||||
try:
|
||||
azure_ad_token_provider: Final = get_azure_ad_token_provider(
|
||||
azure_scope=scope,
|
||||
azure_credential=AzureCredentialType.DefaultAzureCredential,
|
||||
)
|
||||
azure_ad_token_provider: Final = _cached_default_azure_credential_provider(scope)
|
||||
verbose_logger.debug("Successfully obtained Azure AD token provider using DefaultAzureCredential")
|
||||
return azure_ad_token_provider
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ from collections.abc import Mapping
|
|||
import httpx
|
||||
|
||||
from litellm.types.agents import (
|
||||
AgentCreateRequest,
|
||||
AgentCreateResponse,
|
||||
AgentDeleteResult,
|
||||
AgentListResponse,
|
||||
|
|
@ -50,10 +51,9 @@ class BaseAgentsAPIConfig(ABC):
|
|||
@abstractmethod
|
||||
def transform_create_request(
|
||||
self,
|
||||
name: str,
|
||||
litellm_params: Mapping[str, object],
|
||||
) -> dict[str, object]:
|
||||
"""Map name + litellm_params to the provider's create-agent body."""
|
||||
create_request: AgentCreateRequest,
|
||||
) -> Mapping[str, object]:
|
||||
"""Map the create request to the provider's create-agent body."""
|
||||
|
||||
@abstractmethod
|
||||
def transform_create_response(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -20,16 +20,13 @@ from litellm.litellm_core_utils.url_utils import join_api_base
|
|||
from litellm.llms.base_llm.agents.transformation import BaseAgentsAPIConfig
|
||||
from litellm.llms.gemini.common_utils import GeminiError, GeminiModelInfo
|
||||
from litellm.types.agents import (
|
||||
AgentCreateRequest,
|
||||
AgentCreateResponse,
|
||||
AgentDeleteResult,
|
||||
AgentListResponse,
|
||||
AgentVersionsResponse,
|
||||
)
|
||||
|
||||
# Keys inside litellm_params that should be forwarded to the Gemini
|
||||
# create-agent body verbatim.
|
||||
_GEMINI_AGENT_BODY_KEYS: Final = ("base_agent", "instructions", "base_environment")
|
||||
|
||||
|
||||
class _JsonResponse(Protocol):
|
||||
def json(self) -> object: ...
|
||||
|
|
@ -39,30 +36,6 @@ def _json_body(response: _JsonResponse) -> object:
|
|||
return response.json()
|
||||
|
||||
|
||||
# LiteLLM-internal keys that must never be forwarded to Gemini.
|
||||
_LITELLM_INTERNAL_KEYS: Final = frozenset(
|
||||
{
|
||||
"custom_llm_provider",
|
||||
"api_key",
|
||||
"api_base",
|
||||
"make_public",
|
||||
"cost_per_query",
|
||||
"input_cost_per_token",
|
||||
"output_cost_per_token",
|
||||
"require_trace_id_on_calls_to_agent",
|
||||
"require_trace_id_on_calls_by_agent",
|
||||
"max_iterations",
|
||||
"max_budget_per_session",
|
||||
"guardrails",
|
||||
"is_public",
|
||||
"agent_name",
|
||||
"agent_id",
|
||||
"agent_card_params",
|
||||
"provider_agent_response",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class GeminiAgentsConfig(BaseAgentsAPIConfig):
|
||||
"""
|
||||
Configuration for the Google AI Studio (Gemini) native Agents API.
|
||||
|
|
@ -143,14 +116,9 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig):
|
|||
|
||||
def transform_create_request(
|
||||
self,
|
||||
name: str,
|
||||
litellm_params: Mapping[str, object],
|
||||
) -> dict[str, object]:
|
||||
body: Final[dict[str, object]] = {"name": name}
|
||||
for key in _GEMINI_AGENT_BODY_KEYS:
|
||||
value = litellm_params.get(key)
|
||||
if value is not None:
|
||||
body[key] = value
|
||||
create_request: AgentCreateRequest,
|
||||
) -> Mapping[str, object]:
|
||||
body: Final[Mapping[str, object]] = create_request.model_dump(mode="json", exclude_none=True)
|
||||
verbose_logger.debug("GeminiAgentsConfig create body: %s", body)
|
||||
return body
|
||||
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ import httpx
|
|||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.core_helpers import process_response_headers
|
||||
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, process_response_headers
|
||||
from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base
|
||||
from litellm.llms.base_llm.interactions.transformation import BaseInteractionsAPIConfig
|
||||
from litellm.llms.gemini.common_utils import GeminiError, GeminiModelInfo
|
||||
|
|
@ -41,6 +41,14 @@ else:
|
|||
_JsonObject: TypeAlias = dict[str, object]
|
||||
|
||||
|
||||
def _has_agent_config_type(agent_config: object) -> bool:
|
||||
match agent_config:
|
||||
case {"type": str()}:
|
||||
return True
|
||||
case _:
|
||||
return False
|
||||
|
||||
|
||||
class _InteractionPayload(TypedDict, total=False):
|
||||
"""JSON body of an Interactions API interaction, keyed as ``InteractionsAPIResponse`` fields."""
|
||||
|
||||
|
|
@ -105,6 +113,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig):
|
|||
return [
|
||||
"model",
|
||||
"agent",
|
||||
"agent_config",
|
||||
"input",
|
||||
"tools",
|
||||
"system_instruction",
|
||||
|
|
@ -188,6 +197,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig):
|
|||
request_body["input"] = input
|
||||
|
||||
optional_keys: Final = [
|
||||
"agent_config",
|
||||
"tools",
|
||||
"system_instruction",
|
||||
"stream",
|
||||
|
|
@ -197,8 +207,14 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig):
|
|||
"response_modalities",
|
||||
"previous_interaction_id",
|
||||
]
|
||||
dropped_keys: Final = (
|
||||
frozenset({"agent_config"})
|
||||
if drops_unsupported_params(litellm_params.get("drop_params"))
|
||||
and (model or not _has_agent_config_type(optional_params.get("agent_config")))
|
||||
else frozenset()
|
||||
)
|
||||
for key in optional_keys:
|
||||
if optional_params.get(key) is not None:
|
||||
if key not in dropped_keys and optional_params.get(key) is not None:
|
||||
request_body[key] = optional_params[key]
|
||||
|
||||
# response_mime_type is removed — fold it into response_format.
|
||||
|
|
|
|||
17
litellm/llms/ollama/decisions/transformation.py
Normal file
17
litellm/llms/ollama/decisions/transformation.py
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
from typing import Final
|
||||
|
||||
from litellm.constants import PLACEHOLDER_API_KEY
|
||||
from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig
|
||||
|
||||
|
||||
class OllamaDecisionsConfig(BaseDecisionsConfig):
|
||||
api_key_env = ("OLLAMA_API_KEY",)
|
||||
api_base_env = ("OLLAMA_API_BASE",)
|
||||
api_key_required = False
|
||||
|
||||
def get_default_api_base(self) -> str | None:
|
||||
return "http://localhost:11434"
|
||||
|
||||
def resolve_api_key(self, api_key: str | None) -> str | None:
|
||||
resolved: Final = super().resolve_api_key(api_key)
|
||||
return None if resolved == PLACEHOLDER_API_KEY else resolved
|
||||
|
|
@ -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]] = {
|
||||
|
|
|
|||
|
|
@ -58,6 +58,12 @@
|
|||
"max_completion_tokens": "max_tokens"
|
||||
}
|
||||
},
|
||||
"zerogpu": {
|
||||
"base_url": "https://api.zerogpu.ai/v1",
|
||||
"api_key_env": "ZEROGPU_API_KEY",
|
||||
"api_base_env": "ZEROGPU_API_BASE",
|
||||
"supported_endpoints": ["/v1/chat/completions", "/v1/responses"]
|
||||
},
|
||||
"chutes": {
|
||||
"base_url": "https://llm.chutes.ai/v1/",
|
||||
"api_key_env": "CHUTES_API_KEY",
|
||||
|
|
|
|||
|
|
@ -6124,7 +6124,18 @@ def completion(
|
|||
|
||||
optional_params, rs_val = strip_reasoning_summary_aliases_from_optional_params(optional_params)
|
||||
|
||||
if isinstance(resolved_reasoning_effort, dict) and "summary" in resolved_reasoning_effort:
|
||||
if custom_llm_provider == "openai" and isinstance(resolved_reasoning_effort, dict):
|
||||
from litellm.responses.utils import openai_bridge_reasoning
|
||||
|
||||
optional_params["reasoning_effort"] = openai_bridge_reasoning(
|
||||
requested=resolved_reasoning_effort,
|
||||
effort_kept="reasoning_effort" in optional_params,
|
||||
summary_alias=rs_val,
|
||||
model=model,
|
||||
drop_params=drop_unsupported_reasoning,
|
||||
additional_drop_params=routing_kwargs.get("additional_drop_params"),
|
||||
)
|
||||
elif isinstance(resolved_reasoning_effort, dict) and "summary" in resolved_reasoning_effort:
|
||||
optional_params["reasoning_effort"] = resolved_reasoning_effort
|
||||
elif rs_val is not None:
|
||||
eff: Final = optional_params.get("reasoning_effort", resolved_reasoning_effort)
|
||||
|
|
|
|||
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,
|
||||
|
|
@ -4437,6 +4437,32 @@
|
|||
"generateContent": true
|
||||
}
|
||||
},
|
||||
"zerogpu": {
|
||||
"display_name": "ZeroGPU (`zerogpu`)",
|
||||
"url": "https://docs.litellm.ai/docs/providers/zerogpu",
|
||||
"endpoints": {
|
||||
"chat_completions": true,
|
||||
"responses": true,
|
||||
"text_completion": true,
|
||||
"embeddings": false,
|
||||
"image_generations": false,
|
||||
"image_edits": false,
|
||||
"audio_transcriptions": false,
|
||||
"audio_speech": false,
|
||||
"rerank": false,
|
||||
"ocr": false,
|
||||
"search": false,
|
||||
"moderations": false,
|
||||
"realtime": false,
|
||||
"video_generations": false,
|
||||
"batches": false,
|
||||
"messages": true,
|
||||
"interactions": true,
|
||||
"systemone": false,
|
||||
"a2a": true,
|
||||
"generateContent": true
|
||||
}
|
||||
},
|
||||
"ragflow": {
|
||||
"display_name": "RAGFlow (`ragflow`)",
|
||||
"url": "https://docs.litellm.ai/docs/providers/ragflow",
|
||||
|
|
|
|||
|
|
@ -15015,7 +15015,7 @@
|
|||
]
|
||||
},
|
||||
"post": {
|
||||
"description": "Create a named custom agent on the Gemini side.\n\nExample:\n```bash\ncurl -X POST \"http://localhost:4000/v1beta/agents\" \\\n -H \"Authorization: Bearer sk-...\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"name\": \"my-custom-slides-agent\",\n \"base_agent\": \"waverunner\",\n \"instructions\": \"You are a helpful assistant that creates slides.\",\n \"base_environment\": {\n \"type\": \"remote\",\n \"sources\": [\n {\"type\": \"gcs\", \"source\": \"gs://eap-templates/slides-skill\",\n \"target\": \"/.agents/skills/slides-skill\"}\n ]\n }\n }'\n```",
|
||||
"description": "Create a named custom agent on the Gemini side.\n\nExample:\n```bash\ncurl -X POST \"http://localhost:4000/v1beta/agents\" \\\n -H \"Authorization: Bearer sk-...\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"id\": \"my-custom-slides-agent\",\n \"base_agent\": \"antigravity-preview-09-2026\",\n \"system_instruction\": \"You are a helpful assistant that creates slides.\",\n \"agent_config\": {\"type\": \"antigravity\", \"model\": \"gemini-3.8-flash\"},\n \"base_environment\": {\n \"type\": \"remote\",\n \"sources\": [\n {\"type\": \"gcs\", \"source\": \"gs://eap-templates/slides-skill\",\n \"target\": \"/.agents/skills/slides-skill\"}\n ]\n }\n }'\n```",
|
||||
"operationId": "create_gemini_agent_v1beta_agents_post",
|
||||
"responses": {
|
||||
"200": {
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -160,9 +160,10 @@ async def create_gemini_agent(
|
|||
-H "Authorization: Bearer sk-..." \\
|
||||
-H "Content-Type: application/json" \\
|
||||
-d '{
|
||||
"name": "my-custom-slides-agent",
|
||||
"base_agent": "waverunner",
|
||||
"instructions": "You are a helpful assistant that creates slides.",
|
||||
"id": "my-custom-slides-agent",
|
||||
"base_agent": "antigravity-preview-09-2026",
|
||||
"system_instruction": "You are a helpful assistant that creates slides.",
|
||||
"agent_config": {"type": "antigravity", "model": "gemini-3.8-flash"},
|
||||
"base_environment": {
|
||||
"type": "remote",
|
||||
"sources": [
|
||||
|
|
|
|||
|
|
@ -45,11 +45,14 @@ from litellm.caching.redis_cache import log_redis_failure
|
|||
from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE, INTERNAL_CALL_ORIGIN_METADATA_KEY
|
||||
from litellm.cost_calculator import get_usage_object
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.audio_utils.utils import estimate_audio_tokens
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
get_str_from_messages,
|
||||
split_media_for_tpm,
|
||||
)
|
||||
from litellm.litellm_core_utils.ptu_pricing import is_ptu_cost_attribution_enabled
|
||||
from litellm.litellm_core_utils.token_counter import offload_token_count
|
||||
from litellm.litellm_core_utils.token_counter import estimate_image_tokens, offload_token_count
|
||||
from litellm.llms.azure.ptu_capacity import normalized_tokens
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.auth_utils import (
|
||||
|
|
@ -98,6 +101,7 @@ if TYPE_CHECKING:
|
|||
from opentelemetry.trace import Span as _Span
|
||||
|
||||
from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache
|
||||
from litellm.router import Router
|
||||
from litellm.types.agents import AgentResponse
|
||||
from litellm.types.caching import RedisPipelineIncrementOperation
|
||||
|
||||
|
|
@ -135,6 +139,18 @@ class RateLimitedModel:
|
|||
return requested_limit if requested_limit is not None else limits.get(self.group)
|
||||
|
||||
|
||||
def _resolve_media_models(router: "Router | None", model: str, team_id: str | None) -> tuple[str, ...]:
|
||||
if router is None:
|
||||
return (model,)
|
||||
return tuple(dict.fromkeys(router.resolved_litellm_models(model, team_id, use_base_model=True))) or (model,)
|
||||
|
||||
|
||||
def _resolve_media_models_via_proxy_router(model: str, team_id: str | None) -> tuple[str, ...]:
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
|
||||
return _resolve_media_models(llm_router, model, team_id)
|
||||
|
||||
|
||||
def _resolve_model_group_alias_via_proxy_router(model: str) -> str | None:
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
|
||||
|
|
@ -862,6 +878,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
time_provider: Callable[[], datetime] | None = None,
|
||||
tag_rate_limit_resolver: TagRateLimitResolver = resolve_tag_rate_limits_from_db,
|
||||
model_group_resolver: Callable[[str], str | None] = _resolve_model_group_alias_via_proxy_router,
|
||||
media_model_resolver: Callable[[str, str | None], tuple[str, ...]] = _resolve_media_models_via_proxy_router,
|
||||
ptu_team_ceiling_resolver: Callable[
|
||||
[str, str], PTUTeamCeiling | None
|
||||
] = _resolve_ptu_team_ceiling_via_proxy_router,
|
||||
|
|
@ -872,6 +889,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
self._time_provider = time_provider or datetime.now
|
||||
self._tag_rate_limit_resolver = tag_rate_limit_resolver
|
||||
self._model_group_resolver = model_group_resolver
|
||||
self._media_model_resolver = media_model_resolver
|
||||
self._ptu_team_ceiling_resolver = ptu_team_ceiling_resolver
|
||||
self._fail_closed_resolver = fail_closed_resolver
|
||||
self._force_hash_tag_grouping_resolver = force_hash_tag_grouping_resolver
|
||||
|
|
@ -1133,11 +1151,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
|
||||
def _estimate_tokens_for_request(
|
||||
self,
|
||||
data: dict,
|
||||
data: object,
|
||||
model: str | None = None,
|
||||
min_configured_tpm_limit: int | None = None,
|
||||
call_type: str | None = None,
|
||||
configured_output_tokens: int | None = None,
|
||||
team_id: str | None = None,
|
||||
) -> int:
|
||||
"""
|
||||
Estimate total tokens this request will consume so we can reserve them
|
||||
|
|
@ -1164,6 +1183,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
min_configured_tpm_limit=min_configured_tpm_limit,
|
||||
call_type=call_type,
|
||||
configured_output_tokens=configured_output_tokens,
|
||||
team_id=team_id,
|
||||
media_model_resolver=self._media_model_resolver,
|
||||
)
|
||||
total_estimated: Final = estimated_input_tokens + max_tokens_estimate
|
||||
|
||||
|
|
@ -1184,6 +1205,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
call_type: str | None,
|
||||
configured_output_tokens: int | None,
|
||||
raw_estimate: int,
|
||||
team_id: str | None = None,
|
||||
) -> int:
|
||||
"""The team PTU ceiling counts Azure normalized tokens, so its reservation weighs the
|
||||
output budget the way the ceiling does instead of the raw sum the other scopes reserve."""
|
||||
|
|
@ -1194,6 +1216,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
min_configured_tpm_limit=min_configured_tpm_limit,
|
||||
call_type=call_type,
|
||||
configured_output_tokens=configured_output_tokens,
|
||||
team_id=team_id,
|
||||
media_model_resolver=self._media_model_resolver,
|
||||
)
|
||||
normalized: Final = normalized_tokens(
|
||||
ceiling, prompt_tokens=estimated_input_tokens, completion_tokens=max_tokens_estimate
|
||||
|
|
@ -1206,6 +1230,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
min_configured_tpm_limit: int | None = None,
|
||||
call_type: str | None = None,
|
||||
configured_output_tokens: int | None = None,
|
||||
include_media: bool = True,
|
||||
team_id: str | None = None,
|
||||
media_model_resolver: Callable[[str, str | None], tuple[str, ...]] | None = None,
|
||||
) -> tuple[int, int]:
|
||||
"""
|
||||
Estimate input tokens and output (max_tokens) budget separately, so
|
||||
|
|
@ -1229,10 +1256,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
floor entirely, so the reservation reflects what this tenant's model
|
||||
actually emits rather than one constant shared by every tenant.
|
||||
"""
|
||||
if not isinstance(data, dict):
|
||||
request_data: Final = as_str_mapping(data)
|
||||
if not isinstance(data, dict) or request_data is None:
|
||||
return 0, 0
|
||||
request_model: Final = request_data.get("model")
|
||||
translated_data: Final = self._translate_google_genai_native_request(data, call_type)
|
||||
estimable_data: Final = translated_data if translated_data is not None else data
|
||||
estimable_data: Final = translated_data if translated_data is not None else request_data
|
||||
selected_fields: Final[tuple[object | None, object | None, object | None]] = (
|
||||
(None, None, estimable_data.get("input"))
|
||||
if call_type in RESPONSES_API_CALL_TYPES or call_type in EMBEDDING_API_CALL_TYPES
|
||||
|
|
@ -1249,10 +1278,27 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
)
|
||||
)
|
||||
messages, prompt, raw_input = selected_fields
|
||||
input_text: Final = (
|
||||
original_input: Final = (
|
||||
billable_input_since_newest_compaction(raw_input) if call_type in RESPONSES_API_CALL_TYPES else raw_input
|
||||
)
|
||||
|
||||
media_source: Final = messages if messages is not None else original_input
|
||||
projected, media = (
|
||||
split_media_for_tpm(media_source)
|
||||
if include_media and (isinstance(messages, list) or call_type in RESPONSES_API_CALL_TYPES)
|
||||
else (original_input, ())
|
||||
)
|
||||
input_text: Final = projected if call_type in RESPONSES_API_CALL_TYPES else original_input
|
||||
fallback_model: Final = request_model if isinstance(request_model, str) else None
|
||||
resolved_models: Final = (
|
||||
media_model_resolver(fallback_model or "", team_id)
|
||||
if media and media_model_resolver is not None
|
||||
else (fallback_model or "",)
|
||||
)
|
||||
media_tokens: Final = sum(
|
||||
self._estimate_media_block_tokens(block, resolved_models, fallback_model) for block in media
|
||||
)
|
||||
|
||||
total_chars: Final = (
|
||||
len(get_str_from_messages(messages))
|
||||
if isinstance(messages, list) and messages
|
||||
|
|
@ -1267,7 +1313,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
else 0
|
||||
)
|
||||
|
||||
estimated_input_tokens: Final = max(1, total_chars // DEFAULT_CHARS_PER_TOKEN) if total_chars > 0 else 0
|
||||
estimated_input_tokens: Final = (
|
||||
max(1, total_chars // DEFAULT_CHARS_PER_TOKEN) if total_chars > 0 else 0
|
||||
) + media_tokens
|
||||
|
||||
explicit_max_tokens: Final = self._get_explicit_output_cap(data, call_type)
|
||||
generates_no_output: Final = (
|
||||
|
|
@ -1283,7 +1331,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
max_tokens_estimate: Final = (
|
||||
0
|
||||
if generates_no_output
|
||||
or (explicit_max_tokens is None and total_chars == 0 and configured_output_tokens is None)
|
||||
or (explicit_max_tokens is None and estimated_input_tokens == 0 and configured_output_tokens is None)
|
||||
else explicit_max_tokens
|
||||
if explicit_max_tokens is not None
|
||||
else configured_output_tokens
|
||||
|
|
@ -1293,6 +1341,24 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
|
||||
return estimated_input_tokens, max_tokens_estimate * self.get_output_candidate_count(data, call_type)
|
||||
|
||||
@staticmethod
|
||||
def _estimate_media_block_tokens(
|
||||
block: Mapping[str, object], models: tuple[str, ...], fallback_model: str | None
|
||||
) -> int:
|
||||
if block.get("type") == "input_audio":
|
||||
audio: Final = _callback_kwargs_mapping(block.get("input_audio"))
|
||||
data: Final = audio.get("data")
|
||||
return estimate_audio_tokens(data if isinstance(data, str) else None, models, fallback_model=fallback_model)
|
||||
image: Final = _callback_kwargs_mapping(block.get("image_url"))
|
||||
url: Final = image.get("url")
|
||||
detail: Final = image.get("detail")
|
||||
return estimate_image_tokens(
|
||||
models,
|
||||
url if isinstance(url, str) else "",
|
||||
detail if isinstance(detail, str) else "auto",
|
||||
fallback_model=fallback_model,
|
||||
)
|
||||
|
||||
def _is_redis_cluster(self) -> bool:
|
||||
"""
|
||||
Check if the dual cache is using Redis cluster.
|
||||
|
|
@ -3958,7 +4024,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
except Exception: # noqa: BLE001 # tokenizer failures degrade to the cheap estimate
|
||||
if call_type in RERANK_API_CALL_TYPES and isinstance(selected_text, str):
|
||||
return max(0, len(selected_text) // DEFAULT_CHARS_PER_TOKEN)
|
||||
estimated_input_tokens, _ = self._estimate_input_and_output_tokens(data=data, call_type=call_type)
|
||||
estimated_input_tokens, _ = self._estimate_input_and_output_tokens(
|
||||
data=data, call_type=call_type, include_media=False
|
||||
)
|
||||
return estimated_input_tokens + audio_token_estimate
|
||||
|
||||
async def _reserve_project_io_tokens_or_raise(
|
||||
|
|
@ -4003,6 +4071,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
data=data,
|
||||
min_configured_tpm_limit=min_configured_otpm_limit,
|
||||
call_type=call_type,
|
||||
include_media=False,
|
||||
)
|
||||
raw_estimated_input_tokens: Final = await offload_token_count(self._estimate_precise_input_tokens)(
|
||||
data=data, model=requested_model, call_type=call_type
|
||||
|
|
@ -4360,12 +4429,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
# Post-call reconciliation refunds the over-reservation
|
||||
# delta when actual usage comes in below the floor.
|
||||
estimated_tokens: Final = max(
|
||||
self._estimate_tokens_for_request(
|
||||
await offload_token_count(self._estimate_tokens_for_request)(
|
||||
data=data,
|
||||
model=requested_model,
|
||||
min_configured_tpm_limit=min_configured_tpm_limit,
|
||||
call_type=call_type,
|
||||
configured_output_tokens=configured_output_tokens,
|
||||
team_id=user_api_key_dict.team_id,
|
||||
),
|
||||
1,
|
||||
)
|
||||
|
|
@ -4383,13 +4453,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
)
|
||||
|
||||
capped_request: Final = _REQUEST_RATE_LIMIT_DATA.validate_python(data)
|
||||
ptu_estimated_tokens: Final = self._estimate_ptu_tokens_for_request(
|
||||
ptu_estimated_tokens: Final = await offload_token_count(self._estimate_ptu_tokens_for_request)(
|
||||
ceiling=stash.ptu_ceiling,
|
||||
data=capped_request,
|
||||
min_configured_tpm_limit=min_configured_tpm_limit,
|
||||
call_type=call_type,
|
||||
configured_output_tokens=configured_output_tokens,
|
||||
raw_estimate=estimated_tokens,
|
||||
team_id=user_api_key_dict.team_id,
|
||||
)
|
||||
if stash.ptu_ceiling is not None and ptu_estimated_tokens > stash.ptu_ceiling.tpm_limit:
|
||||
if not reuses_request_slot:
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -4145,6 +4145,34 @@
|
|||
],
|
||||
"default_model_placeholder": "gpt-3.5-turbo"
|
||||
},
|
||||
{
|
||||
"provider": "ZeroGPU",
|
||||
"provider_display_name": "ZeroGPU",
|
||||
"litellm_provider": "zerogpu",
|
||||
"credential_fields": [
|
||||
{
|
||||
"key": "api_base",
|
||||
"label": "API Base",
|
||||
"placeholder": "https://api.zerogpu.ai/v1",
|
||||
"tooltip": null,
|
||||
"required": false,
|
||||
"field_type": "text",
|
||||
"options": null,
|
||||
"default_value": null
|
||||
},
|
||||
{
|
||||
"key": "api_key",
|
||||
"label": "API Key",
|
||||
"placeholder": null,
|
||||
"tooltip": null,
|
||||
"required": true,
|
||||
"field_type": "password",
|
||||
"options": null,
|
||||
"default_value": null
|
||||
}
|
||||
],
|
||||
"default_model_placeholder": "zerogpu/gpt-oss-120b"
|
||||
},
|
||||
{
|
||||
"provider": "CURSOR",
|
||||
"provider_display_name": "Cursor",
|
||||
|
|
|
|||
|
|
@ -256,7 +256,6 @@ REQUIRED_PRESENT_BODY_PARAMS_BY_ROUTE: Final[Mapping[str, tuple[str, ...]]] = Ma
|
|||
"avideo_create_character": ("name", "video"),
|
||||
"acreate_container": ("name",),
|
||||
"aupload_container_file": ("file",),
|
||||
"acreate_agent": ("name",),
|
||||
"acreate_interaction": ("input",),
|
||||
"acreate_eval": ("data_source_config", "testing_criteria"),
|
||||
"acreate_run": ("data_source",),
|
||||
|
|
@ -264,7 +263,7 @@ REQUIRED_PRESENT_BODY_PARAMS_BY_ROUTE: Final[Mapping[str, tuple[str, ...]]] = Ma
|
|||
)
|
||||
|
||||
REQUIRED_ONE_OF_BODY_PARAMS_BY_ROUTE: Final[Mapping[str, tuple[str, str]]] = MappingProxyType(
|
||||
{"acreate_interaction": ("model", "agent")}
|
||||
{"acreate_interaction": ("model", "agent"), "acreate_agent": ("id", "name")}
|
||||
)
|
||||
|
||||
JSON_OBJECT_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object])
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
@ -578,6 +582,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
|
||||
|
|
@ -731,7 +736,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 "",
|
||||
|
|
@ -981,6 +986,39 @@ 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 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.
|
||||
|
|
@ -990,6 +1028,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()
|
||||
|
|
@ -1060,15 +1099,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.completed_response is not None or created is None or not generated:
|
||||
if 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
|
||||
|
|
@ -1082,6 +1126,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:
|
||||
|
|
@ -1218,6 +1265,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
|
||||
|
|
@ -1303,6 +1353,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
|
||||
|
|
@ -1727,6 +1780,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
|
||||
|
|
@ -1734,7 +1789,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
|
||||
|
|
@ -1751,6 +1806,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(
|
||||
|
|
@ -1758,6 +1815,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)
|
||||
|
|
@ -2276,7 +2334,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
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,8 +1,10 @@
|
|||
import base64
|
||||
import os
|
||||
import re
|
||||
from collections.abc import Callable, Collection, Iterable, Mapping, Sequence
|
||||
from functools import reduce
|
||||
from typing import Any, Final, Optional, TypeVar, Union, cast, get_type_hints, overload
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, Optional, TypeVar, Union, cast, get_args, get_type_hints, overload
|
||||
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
from typing_extensions import TypedDict as ExtraItemsTypedDict
|
||||
|
|
@ -14,6 +16,7 @@ from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or
|
|||
from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value, is_nested_path
|
||||
from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils
|
||||
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig, ReasoningSummaryLevel
|
||||
from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
InputTokensDetails,
|
||||
|
|
@ -29,6 +32,7 @@ from litellm.types.responses.main import DecodedResponseId
|
|||
from litellm.types.utils import (
|
||||
CacheCreationTokenDetails,
|
||||
CompletionTokensDetailsWrapper,
|
||||
ModelInfo,
|
||||
PromptTokensDetailsWrapper,
|
||||
ProviderSpecificHeader,
|
||||
SpecialEnums,
|
||||
|
|
@ -1382,3 +1386,98 @@ def auto_reasoning_summary(model: str, custom_llm_provider: object) -> Reasoning
|
|||
provider, model.removeprefix(f"{provider}/")
|
||||
)
|
||||
return "detailed" if responses_config is None else responses_config.auto_reasoning_summary()
|
||||
|
||||
|
||||
def opt_in_auto_reasoning_summary(model: str | None, custom_llm_provider: object) -> ReasoningSummaryLevel | None:
|
||||
if not (litellm.reasoning_auto_summary or os.getenv("LITELLM_REASONING_AUTO_SUMMARY", "false").lower() == "true"):
|
||||
return None
|
||||
return "detailed" if model is None else auto_reasoning_summary(model, custom_llm_provider)
|
||||
|
||||
|
||||
def reasoning_with_summary_default(
|
||||
reasoning: Mapping[str, object], summary_alias: object, model: str, custom_llm_provider: object
|
||||
) -> Mapping[str, object]:
|
||||
if "summary" in reasoning:
|
||||
return {**reasoning}
|
||||
summary: Final = (
|
||||
summary_alias if summary_alias is not None else opt_in_auto_reasoning_summary(model, custom_llm_provider)
|
||||
)
|
||||
return {**reasoning} if summary is None else {**reasoning, "summary": summary}
|
||||
|
||||
|
||||
_SUMMARY_VALUES: Final = frozenset((*get_args(ReasoningSummaryLevel), None))
|
||||
|
||||
|
||||
def _openai_model_info(model: str) -> ModelInfo | None:
|
||||
from litellm.utils import ModelInfoLookupMiss, lookup_model_info_for_registration
|
||||
|
||||
model_info: Final = lookup_model_info_for_registration(model.removeprefix("responses/"), "openai")
|
||||
return None if isinstance(model_info, ModelInfoLookupMiss) else model_info
|
||||
|
||||
|
||||
def _declared_values(values: Sequence[str] | None) -> frozenset[str | None] | None:
|
||||
return None if values is None else frozenset(values)
|
||||
|
||||
|
||||
def _openai_reasoning_field_values(model: str) -> Mapping[str, frozenset[str | None] | None]:
|
||||
model_info: Final = _openai_model_info(model)
|
||||
if model_info is None:
|
||||
return MappingProxyType({"summary": _SUMMARY_VALUES, "mode": None, "context": None})
|
||||
if not model_info.get("supports_reasoning"):
|
||||
return MappingProxyType({})
|
||||
return MappingProxyType(
|
||||
{
|
||||
"summary": _SUMMARY_VALUES,
|
||||
"mode": _declared_values(model_info.get("reasoning_modes")),
|
||||
"context": _declared_values(model_info.get("reasoning_contexts")),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _openai_accepts_reasoning_field(
|
||||
key: str, value: object, values_by_field: Mapping[str, frozenset[str | None] | None]
|
||||
) -> bool:
|
||||
if key == "effort":
|
||||
return True
|
||||
if key not in values_by_field:
|
||||
return False
|
||||
allowed: Final = values_by_field[key]
|
||||
return allowed is None or (isinstance(value, str | None) and value in allowed)
|
||||
|
||||
|
||||
def _openai_accepted_reasoning(reasoning: Mapping[str, object], model: str) -> Mapping[str, object]:
|
||||
values_by_field: Final = _openai_reasoning_field_values(model)
|
||||
return {
|
||||
key: value for key, value in reasoning.items() if _openai_accepts_reasoning_field(key, value, values_by_field)
|
||||
}
|
||||
|
||||
|
||||
def _additional_drop_paths(additional_drop_params: object) -> tuple[str, ...]:
|
||||
if additional_drop_params.__class__ is not list:
|
||||
return ()
|
||||
return tuple(path for path in iter_items(additional_drop_params) if isinstance(path, str))
|
||||
|
||||
|
||||
def openai_bridge_reasoning(
|
||||
*,
|
||||
requested: Mapping[str, object],
|
||||
effort_kept: bool,
|
||||
summary_alias: object,
|
||||
model: str,
|
||||
drop_params: bool,
|
||||
additional_drop_params: object,
|
||||
) -> Mapping[str, object] | None:
|
||||
drop_paths: Final = _additional_drop_paths(additional_drop_params)
|
||||
if "reasoning_effort" in drop_paths:
|
||||
return None
|
||||
kept: Final = requested if effort_kept else {key: value for key, value in requested.items() if key != "effort"}
|
||||
with_summary: Final = reasoning_with_summary_default(kept, summary_alias, model, "openai")
|
||||
accepted: Final = (
|
||||
_openai_accepted_reasoning(with_summary, model)
|
||||
if drop_params and OpenAIGPTConfig.is_openai_catalog_model(model)
|
||||
else {**with_summary}
|
||||
)
|
||||
remaining: Final = as_str_mapping(
|
||||
_apply_nested_drop_params({"reasoning_effort": accepted}, list(drop_paths)).get("reasoning_effort")
|
||||
)
|
||||
return remaining or None
|
||||
|
|
|
|||
|
|
@ -6456,6 +6456,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
|
||||
)
|
||||
|
|
@ -10364,28 +10367,38 @@ class Router:
|
|||
model_info[field] = backend_value
|
||||
|
||||
@staticmethod
|
||||
def _cost_map_backend_model(deployment: Deployment) -> str:
|
||||
model_info_base_model: Final = deployment.model_info.base_model
|
||||
def _cost_map_backend_model(deployment: Deployment | DeploymentTypedDict) -> str:
|
||||
model_info_base_model: Final = (
|
||||
deployment.model_info.base_model
|
||||
if isinstance(deployment, Deployment)
|
||||
else _configured_model_info(deployment).get("base_model")
|
||||
)
|
||||
if isinstance(model_info_base_model, str) and model_info_base_model:
|
||||
return model_info_base_model
|
||||
params_base_model: Final = deployment.litellm_params.get("base_model")
|
||||
params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(
|
||||
deployment.litellm_params.model_dump(exclude_none=True)
|
||||
if isinstance(deployment, Deployment)
|
||||
else deployment["litellm_params"]
|
||||
)
|
||||
params_base_model: Final = params.get("base_model")
|
||||
if isinstance(params_base_model, str) and params_base_model:
|
||||
return params_base_model
|
||||
params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(deployment.litellm_params.model_dump(exclude_none=True))
|
||||
model: Final = (
|
||||
deployment.litellm_params.model
|
||||
if isinstance(deployment, Deployment)
|
||||
else deployment["litellm_params"].get("model", "")
|
||||
)
|
||||
invoked_model: Final = params.get("model_id")
|
||||
if (
|
||||
isinstance(invoked_model, str)
|
||||
and invoked_model
|
||||
and (
|
||||
deployment.litellm_params.custom_llm_provider == "bedrock"
|
||||
or deployment.litellm_params.model.startswith("bedrock/")
|
||||
)
|
||||
and (params.get("custom_llm_provider") == "bedrock" or model.startswith("bedrock/"))
|
||||
):
|
||||
profile_model: Final = invoked_model if invoked_model.startswith("bedrock/") else f"bedrock/{invoked_model}"
|
||||
profile_info: Final = lookup_model_info_for_registration(model=profile_model, custom_llm_provider="bedrock")
|
||||
if not isinstance(profile_info, ModelInfoLookupMiss):
|
||||
return profile_model
|
||||
return deployment.litellm_params.model
|
||||
return model
|
||||
|
||||
@staticmethod
|
||||
def _inherit_builtin_service_tier_pricing(
|
||||
|
|
@ -13945,7 +13958,9 @@ class Router:
|
|||
]
|
||||
return named or self._get_wildcard_deployments(model_name=model_group)
|
||||
|
||||
def resolved_litellm_models(self, model_name: str, team_id: str | None = None) -> tuple[str, ...]:
|
||||
def resolved_litellm_models(
|
||||
self, model_name: str, team_id: str | None = None, *, use_base_model: bool = False
|
||||
) -> tuple[str, ...]:
|
||||
"""The provider model strings `model_name` can actually be served by on this proxy.
|
||||
|
||||
`get_model_list` composes every channel the request path itself uses (exact name,
|
||||
|
|
@ -13958,6 +13973,12 @@ class Router:
|
|||
deployment behind it - so the fallback for an empty result is the caller's policy,
|
||||
never this function's.
|
||||
"""
|
||||
if use_base_model:
|
||||
return tuple(
|
||||
backend_model
|
||||
for deployment in self.get_model_list(model_name=model_name, team_id=team_id) or ()
|
||||
if (backend_model := self._cost_map_backend_model(deployment))
|
||||
)
|
||||
return tuple(
|
||||
litellm_model
|
||||
for deployment in self.get_model_list(model_name=model_name, team_id=team_id) or ()
|
||||
|
|
|
|||
|
|
@ -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,14 +1,27 @@
|
|||
import os
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from types import TracebackType
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.single_flight import SingleFlight
|
||||
from litellm.types.secret_managers.get_azure_ad_token_provider import (
|
||||
AzureCredentialType,
|
||||
)
|
||||
|
||||
_HELD_TOKEN_EXPIRY_MARGIN_SECONDS: Final = 60
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from azure.core.credentials import TokenProvider
|
||||
from azure.core.credentials import (
|
||||
AccessToken,
|
||||
AccessTokenInfo,
|
||||
SupportsTokenInfo,
|
||||
TokenCredential,
|
||||
TokenRequestOptions,
|
||||
)
|
||||
from typing_extensions import TypeIs # noqa: TID251 # tells the two azure-core credential protocols apart
|
||||
|
||||
|
||||
def infer_credential_type_from_environment() -> AzureCredentialType:
|
||||
|
|
@ -37,6 +50,86 @@ def infer_credential_type_from_environment() -> AzureCredentialType:
|
|||
return AzureCredentialType.DefaultAzureCredential
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _HeldToken:
|
||||
token: "AccessToken | AccessTokenInfo | None" = None
|
||||
|
||||
def still_valid(self) -> str | None:
|
||||
held: Final = self.token
|
||||
if held is None or held.expires_on - time.time() <= _HELD_TOKEN_EXPIRY_MARGIN_SECONDS:
|
||||
return None
|
||||
return held.token
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _TokenInfoWatcher:
|
||||
credential: "SupportsTokenInfo"
|
||||
held: _HeldToken
|
||||
|
||||
def get_token_info(self, *scopes: str, options: "TokenRequestOptions | None" = None) -> "AccessTokenInfo":
|
||||
token: Final = self.credential.get_token_info(*scopes, options=options)
|
||||
self.held.token = token
|
||||
return token
|
||||
|
||||
def close(self) -> None:
|
||||
self.credential.close()
|
||||
|
||||
def __enter__(self) -> "SupportsTokenInfo":
|
||||
return self.credential.__enter__()
|
||||
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_value: BaseException | None,
|
||||
traceback: TracebackType | None,
|
||||
) -> None:
|
||||
self.credential.__exit__(exc_type, exc_value, traceback)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _TokenWatcher:
|
||||
credential: "TokenCredential"
|
||||
held: _HeldToken
|
||||
|
||||
def get_token(
|
||||
self,
|
||||
*scopes: str,
|
||||
claims: str | None = None,
|
||||
tenant_id: str | None = None,
|
||||
enable_cae: bool = False,
|
||||
**kwargs: object, # kwargs-ok: azure-core forwards its request options verbatim to the wrapped credential
|
||||
) -> "AccessToken":
|
||||
token: Final = self.credential.get_token(
|
||||
*scopes, claims=claims, tenant_id=tenant_id, enable_cae=enable_cae, **kwargs
|
||||
)
|
||||
self.held.token = token
|
||||
return token
|
||||
|
||||
|
||||
def _supports_token_info(
|
||||
credential: "TokenCredential | SupportsTokenInfo",
|
||||
) -> "TypeIs[SupportsTokenInfo]":
|
||||
return hasattr(credential, "get_token_info")
|
||||
|
||||
|
||||
def single_flight_bearer_token_provider(
|
||||
credential: "TokenCredential | SupportsTokenInfo", scope: str
|
||||
) -> Callable[[], str]:
|
||||
from azure.identity import get_bearer_token_provider
|
||||
|
||||
held: Final = _HeldToken()
|
||||
watched: Final = (
|
||||
_TokenInfoWatcher(credential, held) if _supports_token_info(credential) else _TokenWatcher(credential, held)
|
||||
)
|
||||
provider: Final = get_bearer_token_provider(watched, scope)
|
||||
token_request: Final[SingleFlight[str]] = SingleFlight()
|
||||
|
||||
def bearer_token() -> str:
|
||||
return token_request.run(provider, while_busy=held.still_valid())
|
||||
|
||||
return bearer_token
|
||||
|
||||
|
||||
def get_azure_ad_token_provider(
|
||||
azure_scope: str | None = None,
|
||||
azure_credential: AzureCredentialType | None = None,
|
||||
|
|
@ -65,7 +158,6 @@ def get_azure_ad_token_provider(
|
|||
DefaultAzureCredential,
|
||||
ManagedIdentityCredential,
|
||||
WorkloadIdentityCredential,
|
||||
get_bearer_token_provider,
|
||||
)
|
||||
|
||||
if azure_scope is None:
|
||||
|
|
@ -77,7 +169,7 @@ def get_azure_ad_token_provider(
|
|||
else None or os.environ.get("AZURE_CREDENTIAL") or infer_credential_type_from_environment()
|
||||
)
|
||||
verbose_logger.info("For Azure AD Token Provider, choosing credential type: %s", cred)
|
||||
credential: TokenProvider | None = None
|
||||
credential: TokenCredential | SupportsTokenInfo | None = None
|
||||
if cred == AzureCredentialType.ClientSecretCredential:
|
||||
credential = ClientSecretCredential(
|
||||
client_id=os.environ["AZURE_CLIENT_ID"],
|
||||
|
|
@ -133,4 +225,4 @@ def get_azure_ad_token_provider(
|
|||
if credential is None:
|
||||
raise ValueError("No credential provided")
|
||||
|
||||
return get_bearer_token_provider(credential, azure_scope)
|
||||
return single_flight_bearer_token_provider(credential, azure_scope)
|
||||
|
|
|
|||
|
|
@ -5,7 +5,6 @@ import time
|
|||
import traceback
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from functools import lru_cache
|
||||
from typing import Final, TypeAlias
|
||||
|
||||
import httpx
|
||||
|
|
@ -15,6 +14,7 @@ import litellm
|
|||
from litellm._logging import verbose_logger
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.litellm_core_utils.env_utils import get_env_float
|
||||
from litellm.litellm_core_utils.single_flight import single_flight_lru_cache
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm.secret_managers.get_azure_ad_token_provider import (
|
||||
get_azure_ad_token_provider,
|
||||
|
|
@ -28,7 +28,7 @@ _PARSED_LITERAL: Final = TypeAdapter(object)
|
|||
_SECRET_MANAGER_CLIENT: Final = TypeAdapter(object)
|
||||
|
||||
|
||||
@lru_cache(maxsize=128)
|
||||
@single_flight_lru_cache
|
||||
def _cached_azure_oidc_token_provider(audience: str) -> Callable[[], str]:
|
||||
return get_azure_ad_token_provider(azure_scope=audience)
|
||||
|
||||
|
|
|
|||
|
|
@ -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]:
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ from datetime import datetime
|
|||
from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, TypeAlias
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
from pydantic import ConfigDict, Field, PrivateAttr, StrictInt, field_validator
|
||||
from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, StrictInt, field_validator
|
||||
from typing_extensions import ReadOnly, Required, TypedDict
|
||||
|
||||
from litellm.types.llms.base import LiteLLMBaseModel, LiteLLMPydanticObjectBase
|
||||
|
|
@ -343,6 +343,18 @@ class ListAgentsResponse(LiteLLMBaseModel):
|
|||
agents: list[AgentResponse]
|
||||
|
||||
|
||||
class AgentCreateRequest(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
id: str
|
||||
base_agent: str | None = None
|
||||
agent_config: Mapping[str, object] | None = None
|
||||
system_instruction: str | None = None
|
||||
base_environment: str | Mapping[str, object] | None = None
|
||||
description: str | None = None
|
||||
tools: Sequence[Mapping[str, object]] | None = None
|
||||
|
||||
|
||||
class AgentCreateResponse(LiteLLMPydanticObjectBase):
|
||||
"""
|
||||
Response from a provider-side agent creation or get call (e.g. Gemini v1beta/agents).
|
||||
|
|
|
|||
|
|
@ -225,6 +225,8 @@ class ProviderSpecificModelInfo(TypedDict, total=False):
|
|||
supports_xhigh_reasoning_effort: bool | None
|
||||
supports_max_reasoning_effort: bool | None
|
||||
reasoning_effort_levels: ReadOnly[Sequence[str] | None]
|
||||
reasoning_modes: ReadOnly[Sequence[str] | None]
|
||||
reasoning_contexts: ReadOnly[Sequence[str] | None]
|
||||
default_reasoning_effort: ReadOnly[Literal["none", "minimal", "low", "medium", "high", "xhigh"] | None]
|
||||
supports_output_config: bool | None
|
||||
supports_image_size: bool | None
|
||||
|
|
@ -4586,6 +4588,7 @@ class LlmProviders(str, Enum):
|
|||
PRISM = "prism"
|
||||
OPENINFER = "openinfer"
|
||||
INFERENCEAPIS = "inferenceapis"
|
||||
ZEROGPU = "zerogpu"
|
||||
DARKBLOOM = "darkbloom"
|
||||
META = "meta"
|
||||
SAIL = "sail"
|
||||
|
|
|
|||
|
|
@ -7056,6 +7056,8 @@ def get_model_info_helper(
|
|||
supports_xhigh_reasoning_effort=_model_info.get("supports_xhigh_reasoning_effort", None),
|
||||
supports_max_reasoning_effort=_model_info.get("supports_max_reasoning_effort", None),
|
||||
reasoning_effort_levels=_model_info.get("reasoning_effort_levels", None),
|
||||
reasoning_modes=_model_info.get("reasoning_modes", None),
|
||||
reasoning_contexts=_model_info.get("reasoning_contexts", None),
|
||||
default_reasoning_effort=_model_info.get("default_reasoning_effort", None),
|
||||
bedrock_output_config_effort_ceiling=_model_info.get("bedrock_output_config_effort_ceiling", None),
|
||||
bedrock_converse_supports_strict_tools=_model_info.get("bedrock_converse_supports_strict_tools", None),
|
||||
|
|
@ -9668,18 +9670,19 @@ class ProviderConfigManager:
|
|||
return litellm.HostedVLLMDecisionsConfig()
|
||||
if provider == LlmProviders.LM_STUDIO:
|
||||
return litellm.LMStudioDecisionsConfig()
|
||||
if provider in (LlmProviders.OLLAMA, LlmProviders.OLLAMA_CHAT):
|
||||
return litellm.OllamaDecisionsConfig()
|
||||
if provider == LlmProviders.LAYA:
|
||||
return litellm.LayaDecisionsConfig()
|
||||
if provider == LlmProviders.BESPOKE:
|
||||
return litellm.BespokeDecisionsConfig()
|
||||
if provider in (LlmProviders.OPENAI, LlmProviders.AZURE):
|
||||
if provider in (LlmProviders.OPENAI, LlmProviders.AZURE, LlmProviders.AZURE_AI):
|
||||
configs: Final = {
|
||||
LlmProviders.OPENAI: litellm.OpenAIDecisionsConfig,
|
||||
LlmProviders.AZURE: litellm.AzureDecisionsConfig,
|
||||
LlmProviders.AZURE_AI: litellm.AzureAIDecisionsConfig,
|
||||
}
|
||||
return configs[provider]()
|
||||
if provider == LlmProviders.AZURE_AI:
|
||||
return litellm.AzureAIDecisionsConfig()
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
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,
|
||||
|
|
@ -4437,6 +4437,32 @@
|
|||
"generateContent": true
|
||||
}
|
||||
},
|
||||
"zerogpu": {
|
||||
"display_name": "ZeroGPU (`zerogpu`)",
|
||||
"url": "https://docs.litellm.ai/docs/providers/zerogpu",
|
||||
"endpoints": {
|
||||
"chat_completions": true,
|
||||
"responses": true,
|
||||
"text_completion": true,
|
||||
"embeddings": false,
|
||||
"image_generations": false,
|
||||
"image_edits": false,
|
||||
"audio_transcriptions": false,
|
||||
"audio_speech": false,
|
||||
"rerank": false,
|
||||
"ocr": false,
|
||||
"search": false,
|
||||
"moderations": false,
|
||||
"realtime": false,
|
||||
"video_generations": false,
|
||||
"batches": false,
|
||||
"messages": true,
|
||||
"interactions": true,
|
||||
"systemone": false,
|
||||
"a2a": true,
|
||||
"generateContent": true
|
||||
}
|
||||
},
|
||||
"ragflow": {
|
||||
"display_name": "RAGFlow (`ragflow`)",
|
||||
"url": "https://docs.litellm.ai/docs/providers/ragflow",
|
||||
|
|
|
|||
|
|
@ -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}")
|
||||
|
|
|
|||
|
|
@ -254,8 +254,13 @@ _ROUTES: Final[dict[str, _Route]] = {
|
|||
"acreate_agent": (
|
||||
"/v1beta/agents",
|
||||
"/v1beta/agents",
|
||||
("name",),
|
||||
{"name": "agent", "base_agent": "waverunner", "instructions": "You are a helpful assistant."},
|
||||
("id",),
|
||||
{
|
||||
"id": "agent",
|
||||
"base_agent": "antigravity-preview-09-2026",
|
||||
"system_instruction": "You are a helpful assistant.",
|
||||
"agent_config": {"type": "antigravity", "model": "gemini-3.8-flash"},
|
||||
},
|
||||
"gemini/gemini-2.5-flash",
|
||||
),
|
||||
"acreate_interaction": (
|
||||
|
|
@ -1250,11 +1255,7 @@ def test_valid_agent_creation_reaches_upstream(gateway: Gateway) -> None:
|
|||
_assert_scripted_response("acreate_agent", caller)
|
||||
matches: Final = _observed(_Observations(gateway.upstream_url), identity)
|
||||
outbound: Final = object_value(matches[0]["body"])
|
||||
assert outbound == {
|
||||
"name": "agent",
|
||||
"base_agent": "waverunner",
|
||||
"instructions": "You are a helpful assistant.",
|
||||
}, matches
|
||||
assert outbound == body, matches
|
||||
|
||||
|
||||
def test_valid_speech_returns_binary_audio_and_reaches_upstream(gateway: Gateway) -> None:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
625
tests/integration/providers/test_azure_entra_token_reuse_wire.py
Normal file
625
tests/integration/providers/test_azure_entra_token_reuse_wire.py
Normal file
|
|
@ -0,0 +1,625 @@
|
|||
import asyncio
|
||||
import itertools
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass, field
|
||||
from enum import StrEnum
|
||||
from pathlib import Path
|
||||
from queue import SimpleQueue
|
||||
from typing import Final
|
||||
from urllib.parse import parse_qs, urlsplit
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import websockets
|
||||
import yaml
|
||||
from integration._support.client import (
|
||||
JSON_OBJECT,
|
||||
Gateway,
|
||||
Scenario,
|
||||
eventually,
|
||||
gateway_from_environment,
|
||||
object_value,
|
||||
string_value,
|
||||
)
|
||||
from integration._support.openai_wire import chat_reply, responses_reply
|
||||
from integration._support.ports import allocate_port
|
||||
from integration._support.process import graceful_stop_seconds, owned_proxy
|
||||
from integration._support.responses_ws_peer import Peer, completed, responses_peer, session
|
||||
from integration._support.tls import write_self_signed_cert
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import BaseModel, ConfigDict, JsonValue
|
||||
|
||||
pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state]
|
||||
|
||||
_DEPLOYMENT: Final = "gpt-5.4-nano"
|
||||
_DATED_API_VERSION: Final = "2024-10-21"
|
||||
_REALTIME_API_VERSION: Final = "2024-10-01-preview"
|
||||
_IDENTITY_PATH: Final = "/msi/token"
|
||||
_EXCHANGE_SUFFIX: Final = "/oauth2/v2.0/token"
|
||||
_FOUNDRY_PROJECT: Final = "/api/projects/audit"
|
||||
_BING_TOOL: Final = "audit-bing-entra"
|
||||
_BING_RESOURCE: Final = "https://ai.azure.com"
|
||||
_BING_CITATION: Final = "https://result.invalid/bing"
|
||||
_TENANT: Final = "audit-tenant"
|
||||
_APP_CLIENT: Final = "audit-app-client"
|
||||
_HOLD_SECONDS: Final = 60
|
||||
_LIVELINESS_SECONDS: Final = 3
|
||||
_HERD: Final = 6
|
||||
_HERD_WINDOW_SECONDS: Final = 3
|
||||
_TOKEN_LIFETIME_SECONDS: Final = 3600
|
||||
_INSIDE_REFRESH_WINDOW_SECONDS: Final = 120
|
||||
_REMOVED_ENVIRONMENT: Final = (
|
||||
"AZURE_API_KEY",
|
||||
"AZURE_OPENAI_API_KEY",
|
||||
"AZURE_AD_TOKEN",
|
||||
"AZURE_TENANT_ID",
|
||||
"AZURE_CLIENT_SECRET",
|
||||
"AZURE_FEDERATED_TOKEN_FILE",
|
||||
"AZURE_USERNAME",
|
||||
"AZURE_PASSWORD",
|
||||
"AZURE_CREDENTIAL",
|
||||
"AZURE_SCOPE",
|
||||
"MSI_ENDPOINT",
|
||||
"MSI_SECRET",
|
||||
"IMDS_ENDPOINT",
|
||||
"IDENTITY_SERVER_THUMBPRINT",
|
||||
"BING_GROUNDING_TOKEN",
|
||||
"HTTP_PROXY",
|
||||
"HTTPS_PROXY",
|
||||
"http_proxy",
|
||||
"https_proxy",
|
||||
)
|
||||
|
||||
|
||||
class _Surface(StrEnum):
|
||||
LEGACY_CHAT = "legacy-chat"
|
||||
V1_CHAT = "v1-chat"
|
||||
RESPONSES = "responses"
|
||||
RESPONSES_WEBSOCKET = "responses-websocket"
|
||||
REALTIME = "realtime"
|
||||
REALTIME_HEALTH_CHECK = "realtime-health-check"
|
||||
|
||||
|
||||
_HTTP_SURFACES: Final = frozenset({_Surface.LEGACY_CHAT, _Surface.V1_CHAT, _Surface.RESPONSES})
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Rig:
|
||||
gateway: Gateway
|
||||
port: int
|
||||
cert: tuple[Path, Path]
|
||||
|
||||
@property
|
||||
def edge(self) -> str:
|
||||
return f"http://127.0.0.1:{self.port}"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Hold:
|
||||
resource: str
|
||||
arrivals: SimpleQueue[str] = field(default_factory=SimpleQueue)
|
||||
release: threading.Event = field(default_factory=threading.Event)
|
||||
|
||||
|
||||
def _config(directory: Path) -> Path:
|
||||
stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
|
||||
path: Final = directory / "azure-entra-token-reuse.yaml"
|
||||
path.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
**stock,
|
||||
"litellm_settings": {**object_value(stock["litellm_settings"]), "enable_azure_ad_token_refresh": True},
|
||||
"router_settings": {**object_value(stock["router_settings"]), "num_retries": 0},
|
||||
"search_tools": [
|
||||
{"search_tool_name": _BING_TOOL, "litellm_params": {"search_provider": "bing_grounding"}}
|
||||
],
|
||||
}
|
||||
)
|
||||
)
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
|
||||
directory: Final = tmp_path_factory.mktemp("azure-entra-token-reuse")
|
||||
port: Final = allocate_port()
|
||||
cert: Final = write_self_signed_cert(directory)
|
||||
edge: Final = f"http://127.0.0.1:{port}"
|
||||
environment: Final = {
|
||||
"AZURE_CLIENT_ID": "audit-managed-identity",
|
||||
"IDENTITY_ENDPOINT": f"{edge}{_IDENTITY_PATH}",
|
||||
"IDENTITY_HEADER": "audit-identity-header",
|
||||
"AZURE_AUTHORITY_HOST": edge,
|
||||
"BING_GROUNDING_PROJECT_ENDPOINT": f"{edge}{_FOUNDRY_PROJECT}",
|
||||
"BING_GROUNDING_MODEL": _DEPLOYMENT,
|
||||
"BING_GROUNDING_CONNECTION_ID": "audit-bing-connection",
|
||||
"SSL_CERT_FILE": str(cert[0]),
|
||||
}
|
||||
with (
|
||||
gateway_from_environment() as shared,
|
||||
owned_proxy(
|
||||
shared, directory, environment, config=_config(directory), remove_environment=_REMOVED_ENVIRONMENT
|
||||
) as proxy,
|
||||
):
|
||||
yield _Rig(proxy, port, cert)
|
||||
|
||||
|
||||
def _json(payload: Mapping[str, JsonValue]) -> Reply:
|
||||
return Reply(body=json.dumps(payload).encode())
|
||||
|
||||
|
||||
def _identity_token(resource: str) -> str:
|
||||
return f"mi-token:{resource}"
|
||||
|
||||
|
||||
def _identity_resource(request: Request) -> str | None:
|
||||
target: Final = urlsplit(request.target)
|
||||
if target.path != _IDENTITY_PATH:
|
||||
return None
|
||||
return parse_qs(target.query)["resource"][-1]
|
||||
|
||||
|
||||
def _is_token_endpoint(request: Request) -> bool:
|
||||
return _identity_resource(request) is not None or request.target.endswith(_EXCHANGE_SUFFIX)
|
||||
|
||||
|
||||
def _managed_identity(
|
||||
resource: str,
|
||||
hold: _Hold | None,
|
||||
lifetime_seconds: int = _TOKEN_LIFETIME_SECONDS,
|
||||
access_token: str | None = None,
|
||||
) -> Reply:
|
||||
if hold is not None and resource == hold.resource:
|
||||
hold.arrivals.put(resource)
|
||||
hold.release.wait(_HOLD_SECONDS)
|
||||
return _json(
|
||||
{
|
||||
"access_token": access_token or _identity_token(resource),
|
||||
"expires_on": str(int(time.time()) + lifetime_seconds),
|
||||
"resource": resource,
|
||||
"token_type": "Bearer",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _entra_exchange(request: Request) -> Reply:
|
||||
assertion: Final = parse_qs(request.body.decode())["client_assertion"][-1]
|
||||
return _json({"access_token": f"aad:{assertion}", "expires_in": 3600, "token_type": "Bearer"})
|
||||
|
||||
|
||||
def _edge(upstream: Callable[[Request], Reply], hold: _Hold | None = None) -> Callable[[Request], Reply]:
|
||||
def respond(request: Request) -> Reply:
|
||||
resource: Final = _identity_resource(request)
|
||||
if resource is not None:
|
||||
return _managed_identity(resource, hold)
|
||||
if request.target.endswith(_EXCHANGE_SUFFIX):
|
||||
return _entra_exchange(request)
|
||||
return upstream(request)
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def _azure_openai(request: Request) -> Reply:
|
||||
if urlsplit(request.target).path.endswith("/chat/completions"):
|
||||
return chat_reply(f"chatcmpl-{uuid.uuid4().hex}", _DEPLOYMENT, "azure answered", stream=False)
|
||||
return responses_reply(f"resp_{uuid.uuid4().hex}", _DEPLOYMENT, "azure answered", stream=False)
|
||||
|
||||
|
||||
def _identity_fetches(requests: tuple[Request, ...], resource: str) -> int:
|
||||
return sum(1 for request in requests if _identity_resource(request) == resource)
|
||||
|
||||
|
||||
def _provider_calls(requests: tuple[Request, ...]) -> tuple[Request, ...]:
|
||||
return tuple(request for request in requests if not _is_token_endpoint(request))
|
||||
|
||||
|
||||
def _marker() -> str:
|
||||
return f"entra audit {uuid.uuid4().hex}"
|
||||
|
||||
|
||||
def _model(scenario: Scenario, surface: _Surface, api_base: str, scope: str) -> str:
|
||||
match surface:
|
||||
case _Surface.LEGACY_CHAT | _Surface.RESPONSES:
|
||||
return scenario.model(
|
||||
model=f"azure/{_DEPLOYMENT}",
|
||||
api_base=api_base,
|
||||
api_key=None,
|
||||
api_version=_DATED_API_VERSION,
|
||||
azure_scope=scope,
|
||||
)
|
||||
case _Surface.V1_CHAT:
|
||||
return scenario.model(
|
||||
model=f"azure/{_DEPLOYMENT}", api_base=api_base, api_key=None, api_version="v1", azure_scope=scope
|
||||
)
|
||||
case _Surface.RESPONSES_WEBSOCKET:
|
||||
return scenario.model(model=f"azure/{_DEPLOYMENT}", api_base=api_base, api_key=None, azure_scope=scope)
|
||||
case _Surface.REALTIME | _Surface.REALTIME_HEALTH_CHECK:
|
||||
return scenario.model(
|
||||
model_info={"mode": "realtime"},
|
||||
model=f"azure/{_DEPLOYMENT}",
|
||||
api_base=api_base,
|
||||
api_key=None,
|
||||
api_version=_REALTIME_API_VERSION,
|
||||
azure_scope=scope,
|
||||
)
|
||||
|
||||
|
||||
def _post(surface: _Surface, gateway: Gateway, key: str, model: str) -> httpx.Response:
|
||||
if surface is _Surface.RESPONSES:
|
||||
return gateway.request("POST", "/v1/responses", {"model": model, "input": _marker()}, key=key)
|
||||
return gateway.request(
|
||||
"POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": _marker()}]}, key=key
|
||||
)
|
||||
|
||||
|
||||
async def _realtime_handshake(gateway: Gateway, key: str, model: str, peer: Peer) -> str | None:
|
||||
proxy: Final = str(gateway.client.base_url).rstrip("/").replace("http://", "ws://")
|
||||
async with websockets.connect(
|
||||
f"{proxy}/v1/realtime?model={model}",
|
||||
additional_headers={"Authorization": f"Bearer {key}"},
|
||||
open_timeout=_HOLD_SECONDS,
|
||||
):
|
||||
credentials: Final = await asyncio.to_thread(peer.credentials.get, True, _HOLD_SECONDS)
|
||||
return credentials.authorization
|
||||
|
||||
|
||||
def _upstream_authorization(surface: _Surface, gateway: Gateway, key: str, model: str, wire: Wire, peer: Peer) -> str:
|
||||
match surface:
|
||||
case _Surface.LEGACY_CHAT | _Surface.V1_CHAT | _Surface.RESPONSES:
|
||||
response: Final = _post(surface, gateway, key, model)
|
||||
assert response.status_code == 200, response.text
|
||||
(call,) = _provider_calls(wire.drain())
|
||||
return call.headers.get("authorization", "")
|
||||
case _Surface.RESPONSES_WEBSOCKET:
|
||||
(events,) = asyncio.run(session(str(gateway.client.base_url), key, model, (_marker(),)))
|
||||
assert completed(events)["status"] == "completed", events
|
||||
return peer.credentials.get(timeout=_HOLD_SECONDS).authorization or ""
|
||||
case _Surface.REALTIME:
|
||||
return asyncio.run(_realtime_handshake(gateway, key, model, peer)) or ""
|
||||
case _Surface.REALTIME_HEALTH_CHECK:
|
||||
gateway.request("GET", "/health", params={"model": model})
|
||||
return peer.credentials.get(timeout=_HOLD_SECONDS).authorization or ""
|
||||
|
||||
|
||||
def _liveliness(gateway: Gateway) -> int | None:
|
||||
try:
|
||||
return httpx.get(
|
||||
f"{str(gateway.client.base_url).rstrip('/')}/health/liveliness", timeout=_LIVELINESS_SECONDS
|
||||
).status_code
|
||||
except httpx.TimeoutException:
|
||||
return None
|
||||
|
||||
|
||||
def _liveliness_while_held(gateway: Gateway, hold: _Hold) -> int | None:
|
||||
try:
|
||||
assert hold.arrivals.get(timeout=_HOLD_SECONDS) == hold.resource
|
||||
return _liveliness(gateway)
|
||||
finally:
|
||||
hold.release.set()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("surface", tuple(_Surface))
|
||||
def test_a_slow_managed_identity_fetch_leaves_the_proxy_answering_other_requests(rig: _Rig, surface: _Surface) -> None:
|
||||
resource: Final = f"api://audit-hold-{uuid.uuid4().hex}"
|
||||
hold: Final = _Hold(resource)
|
||||
with (
|
||||
wire_server(_edge(_azure_openai, hold), port=rig.port) as wire,
|
||||
responses_peer(rig.cert) as peer,
|
||||
rig.gateway.scenario() as scenario,
|
||||
ThreadPoolExecutor(max_workers=1) as pool,
|
||||
):
|
||||
api_base: Final = rig.edge if surface in _HTTP_SURFACES else peer.url.removesuffix("/v1")
|
||||
model: Final = _model(scenario, surface, api_base, f"{resource}/.default")
|
||||
key: Final = scenario.key()
|
||||
call: Final = pool.submit(_upstream_authorization, surface, rig.gateway, key, model, wire, peer)
|
||||
liveliness: Final = _liveliness_while_held(rig.gateway, hold)
|
||||
authorization: Final = call.result(timeout=_HOLD_SECONDS)
|
||||
|
||||
assert liveliness == 200, f"/health/liveliness did not answer within {_LIVELINESS_SECONDS}s during the fetch"
|
||||
assert authorization == f"Bearer {_identity_token(resource)}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("surface", (_Surface.LEGACY_CHAT, _Surface.RESPONSES))
|
||||
def test_concurrent_cold_requests_share_one_managed_identity_fetch(rig: _Rig, surface: _Surface) -> None:
|
||||
resource: Final = f"api://audit-herd-{uuid.uuid4().hex}"
|
||||
hold: Final = _Hold(resource)
|
||||
with (
|
||||
wire_server(_edge(_azure_openai, hold), port=rig.port) as wire,
|
||||
rig.gateway.scenario() as scenario,
|
||||
ThreadPoolExecutor(max_workers=_HERD) as pool,
|
||||
):
|
||||
model: Final = _model(scenario, surface, rig.edge, f"{resource}/.default")
|
||||
key: Final = scenario.key()
|
||||
calls: Final = tuple(pool.submit(_post, surface, rig.gateway, key, model) for _ in range(_HERD))
|
||||
try:
|
||||
eventually(
|
||||
hold.arrivals.qsize,
|
||||
lambda arrived: arrived >= _HERD,
|
||||
seconds=_HERD_WINDOW_SECONDS,
|
||||
return_last_on_timeout=True,
|
||||
)
|
||||
finally:
|
||||
hold.release.set()
|
||||
responses: Final = tuple(call.result(timeout=_HOLD_SECONDS) for call in calls)
|
||||
requests: Final = wire.drain()
|
||||
|
||||
assert [response.status_code for response in responses] == [200] * _HERD, [r.text for r in responses]
|
||||
assert _identity_fetches(requests, resource) == 1
|
||||
assert [call.headers.get("authorization") for call in _provider_calls(requests)] == [
|
||||
f"Bearer {_identity_token(resource)}"
|
||||
] * _HERD
|
||||
|
||||
|
||||
def _scripted_identity(resource: str, fetches: tuple[Callable[[], Reply], ...]) -> Callable[[Request], Reply]:
|
||||
attempts: Final = itertools.count()
|
||||
edge: Final = _edge(_azure_openai)
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
if _identity_resource(request) != resource:
|
||||
return edge(request)
|
||||
return fetches[min(next(attempts), len(fetches) - 1)]()
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def _held_failure(hold: _Hold) -> Reply:
|
||||
hold.arrivals.put(hold.resource)
|
||||
hold.release.wait(_HOLD_SECONDS)
|
||||
return Reply(body=b"not json")
|
||||
|
||||
|
||||
def _herd(
|
||||
rig: _Rig, edge: Callable[[Request], Reply], resource: str, hold: _Hold, warm_up: bool
|
||||
) -> tuple[tuple[httpx.Response, ...], tuple[Request, ...]]:
|
||||
with (
|
||||
wire_server(edge, port=rig.port) as wire,
|
||||
rig.gateway.scenario() as scenario,
|
||||
ThreadPoolExecutor(max_workers=_HERD) as pool,
|
||||
):
|
||||
model: Final = _model(scenario, _Surface.LEGACY_CHAT, rig.edge, f"{resource}/.default")
|
||||
key: Final = scenario.key()
|
||||
warm_up_responses: Final = (_post(_Surface.LEGACY_CHAT, rig.gateway, key, model),) if warm_up else ()
|
||||
calls: Final = tuple(pool.submit(_post, _Surface.LEGACY_CHAT, rig.gateway, key, model) for _ in range(_HERD))
|
||||
try:
|
||||
eventually(
|
||||
hold.arrivals.qsize,
|
||||
lambda arrived: arrived >= _HERD,
|
||||
seconds=_HERD_WINDOW_SECONDS,
|
||||
return_last_on_timeout=True,
|
||||
)
|
||||
finally:
|
||||
hold.release.set()
|
||||
return warm_up_responses + tuple(call.result(timeout=_HOLD_SECONDS) for call in calls), wire.drain()
|
||||
|
||||
|
||||
def test_concurrent_requests_inside_the_refresh_window_share_one_refresh_and_keep_the_valid_token(rig: _Rig) -> None:
|
||||
resource: Final = f"api://audit-refresh-{uuid.uuid4().hex}"
|
||||
refreshed: Final = f"mi-token-refreshed:{resource}"
|
||||
hold: Final = _Hold(resource)
|
||||
edge: Final = _scripted_identity(
|
||||
resource,
|
||||
(
|
||||
lambda: _managed_identity(resource, None, _INSIDE_REFRESH_WINDOW_SECONDS),
|
||||
lambda: _managed_identity(resource, hold, access_token=refreshed),
|
||||
),
|
||||
)
|
||||
|
||||
responses, requests = _herd(rig, edge, resource, hold, warm_up=True)
|
||||
|
||||
assert [response.status_code for response in responses] == [200] * (_HERD + 1), [r.text for r in responses]
|
||||
assert _identity_fetches(requests, resource) == 2
|
||||
assert sorted(call.headers.get("authorization", "") for call in _provider_calls(requests)) == sorted(
|
||||
[f"Bearer {_identity_token(resource)}"] * _HERD + [f"Bearer {refreshed}"]
|
||||
)
|
||||
|
||||
|
||||
def test_a_failed_cold_fetch_fails_only_the_request_that_made_it(rig: _Rig) -> None:
|
||||
resource: Final = f"api://audit-failure-{uuid.uuid4().hex}"
|
||||
hold: Final = _Hold(resource)
|
||||
edge: Final = _scripted_identity(resource, (lambda: _held_failure(hold), lambda: _managed_identity(resource, None)))
|
||||
|
||||
responses, _ = _herd(rig, edge, resource, hold, warm_up=False)
|
||||
|
||||
assert sorted(response.status_code == 200 for response in responses) == [False] + [True] * (_HERD - 1), [
|
||||
r.text for r in responses
|
||||
]
|
||||
|
||||
|
||||
def _oidc_model(scenario: Scenario, rig: _Rig, audience: str) -> str:
|
||||
return scenario.model(
|
||||
model=f"azure/{_DEPLOYMENT}",
|
||||
api_base=rig.edge,
|
||||
api_key=None,
|
||||
api_version=_DATED_API_VERSION,
|
||||
azure_ad_token=f"oidc/azure/{audience}",
|
||||
tenant_id=_TENANT,
|
||||
client_id=_APP_CLIENT,
|
||||
)
|
||||
|
||||
|
||||
def _exchanges(requests: tuple[Request, ...]) -> tuple[Request, ...]:
|
||||
return tuple(request for request in requests if request.target.endswith(_EXCHANGE_SUFFIX))
|
||||
|
||||
|
||||
def test_an_oidc_federated_deployment_fetches_its_assertion_once_across_responses_calls(rig: _Rig) -> None:
|
||||
audience: Final = f"api://audit-oidc-{uuid.uuid4().hex}"
|
||||
with wire_server(_edge(_azure_openai), port=rig.port) as wire, rig.gateway.scenario() as scenario:
|
||||
model: Final = _oidc_model(scenario, rig, audience)
|
||||
key: Final = scenario.key()
|
||||
responses: Final = tuple(_post(_Surface.RESPONSES, rig.gateway, key, model) for _ in range(3))
|
||||
requests: Final = wire.drain()
|
||||
|
||||
assert [response.status_code for response in responses] == [200] * 3, [r.text for r in responses]
|
||||
assert _identity_fetches(requests, audience) == 1
|
||||
exchanges: Final = _exchanges(requests)
|
||||
assert [request.target for request in exchanges] == [f"/{_TENANT}{_EXCHANGE_SUFFIX}"]
|
||||
assert {name: values[-1] for name, values in parse_qs(exchanges[0].body.decode()).items()} == {
|
||||
"client_id": _APP_CLIENT,
|
||||
"grant_type": "client_credentials",
|
||||
"scope": "https://cognitiveservices.azure.com/.default",
|
||||
"client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer",
|
||||
"client_assertion": _identity_token(audience),
|
||||
}
|
||||
assert [call.headers.get("authorization") for call in _provider_calls(requests)] == [
|
||||
f"Bearer aad:{_identity_token(audience)}"
|
||||
] * 3
|
||||
|
||||
|
||||
def test_concurrent_cold_oidc_requests_share_one_assertion_fetch_and_one_exchange(rig: _Rig) -> None:
|
||||
audience: Final = f"api://audit-oidc-herd-{uuid.uuid4().hex}"
|
||||
hold: Final = _Hold(audience)
|
||||
with (
|
||||
wire_server(_edge(_azure_openai, hold), port=rig.port) as wire,
|
||||
rig.gateway.scenario() as scenario,
|
||||
ThreadPoolExecutor(max_workers=_HERD) as pool,
|
||||
):
|
||||
model: Final = _oidc_model(scenario, rig, audience)
|
||||
key: Final = scenario.key()
|
||||
calls: Final = tuple(pool.submit(_post, _Surface.RESPONSES, rig.gateway, key, model) for _ in range(_HERD))
|
||||
try:
|
||||
eventually(
|
||||
hold.arrivals.qsize,
|
||||
lambda arrived: arrived >= _HERD,
|
||||
seconds=_HERD_WINDOW_SECONDS,
|
||||
return_last_on_timeout=True,
|
||||
)
|
||||
finally:
|
||||
hold.release.set()
|
||||
responses: Final = tuple(call.result(timeout=_HOLD_SECONDS) for call in calls)
|
||||
requests: Final = wire.drain()
|
||||
|
||||
assert [response.status_code for response in responses] == [200] * _HERD, [r.text for r in responses]
|
||||
assert _identity_fetches(requests, audience) == 1
|
||||
assert len(_exchanges(requests)) == 1
|
||||
assert [call.headers.get("authorization") for call in _provider_calls(requests)] == [
|
||||
f"Bearer aad:{_identity_token(audience)}"
|
||||
] * _HERD
|
||||
|
||||
|
||||
def _foundry(request: Request) -> Reply:
|
||||
assert urlsplit(request.target).path == f"{_FOUNDRY_PROJECT}/openai/v1/responses", request.target
|
||||
text: Final = "Audit result."
|
||||
return _json(
|
||||
{
|
||||
"status": "completed",
|
||||
"error": None,
|
||||
"incomplete_details": None,
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
"content": [
|
||||
{
|
||||
"type": "output_text",
|
||||
"text": text,
|
||||
"annotations": [
|
||||
{
|
||||
"type": "url_citation",
|
||||
"url": _BING_CITATION,
|
||||
"title": "Bing result",
|
||||
"start_index": len(text),
|
||||
"end_index": len(text),
|
||||
}
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class _SearchResult(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
url: str
|
||||
|
||||
|
||||
class _SearchResponse(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
results: tuple[_SearchResult, ...]
|
||||
|
||||
|
||||
def test_bing_grounding_search_mints_its_entra_token_once_across_searches(rig: _Rig) -> None:
|
||||
with wire_server(_edge(_foundry), port=rig.port) as wire:
|
||||
responses: Final = tuple(
|
||||
rig.gateway.request("POST", f"/v1/search/{_BING_TOOL}", {"query": _marker(), "max_results": 1})
|
||||
for _ in range(3)
|
||||
)
|
||||
requests: Final = wire.drain()
|
||||
|
||||
assert [response.status_code for response in responses] == [200] * 3, [r.text for r in responses]
|
||||
assert [
|
||||
tuple(result.url for result in _SearchResponse.model_validate_json(response.content).results)
|
||||
for response in responses
|
||||
] == [(_BING_CITATION,)] * 3
|
||||
assert _identity_fetches(requests, _BING_RESOURCE) == 1
|
||||
assert [call.headers.get("authorization") for call in _provider_calls(requests)] == [
|
||||
f"Bearer {_identity_token(_BING_RESOURCE)}"
|
||||
] * 3
|
||||
|
||||
|
||||
def _lifecycle(upstream_id: str) -> Callable[[Request], Reply]:
|
||||
response: Final[dict[str, JsonValue]] = {
|
||||
"id": upstream_id,
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"model": _DEPLOYMENT,
|
||||
"output": [],
|
||||
}
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
path: Final = urlsplit(request.target).path
|
||||
match request.method:
|
||||
case "DELETE":
|
||||
return _json({"id": upstream_id, "object": "response", "deleted": True})
|
||||
case "GET" if path.endswith("/input_items"):
|
||||
return _json({"object": "list", "data": [], "first_id": None, "last_id": None, "has_more": False})
|
||||
case "POST" if path.endswith("/cancel"):
|
||||
return _json({**response, "status": "cancelled"})
|
||||
case _:
|
||||
return _json({**response, "status": "completed"})
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
_LIFECYCLE_ROUTES: Final = (
|
||||
("GET", "/v1/responses/{}", ""),
|
||||
("GET", "/v1/responses/{}/input_items", "/input_items"),
|
||||
("POST", "/v1/responses/{}/cancel", "/cancel"),
|
||||
("DELETE", "/v1/responses/{}", ""),
|
||||
)
|
||||
|
||||
|
||||
def test_every_responses_route_on_an_entra_deployment_reaches_azure_with_the_managed_identity_token(
|
||||
rig: _Rig,
|
||||
) -> None:
|
||||
resource: Final = f"api://audit-routes-{uuid.uuid4().hex}"
|
||||
upstream_id: Final = f"resp_{uuid.uuid4().hex}"
|
||||
with wire_server(_edge(_lifecycle(upstream_id)), port=rig.port) as wire, rig.gateway.scenario() as scenario:
|
||||
model: Final = _model(scenario, _Surface.RESPONSES, rig.edge, f"{resource}/.default")
|
||||
key: Final = scenario.key()
|
||||
created: Final = _post(_Surface.RESPONSES, rig.gateway, key, model)
|
||||
assert created.status_code == 200, created.text
|
||||
response_id: Final = string_value(JSON_OBJECT.validate_json(created.content)["id"])
|
||||
followups: Final = tuple(
|
||||
rig.gateway.request(method, route.format(response_id), key=key) for method, route, _ in _LIFECYCLE_ROUTES
|
||||
)
|
||||
requests: Final = wire.drain()
|
||||
|
||||
assert [response.status_code for response in followups] == [200] * len(_LIFECYCLE_ROUTES), [
|
||||
response.text for response in followups
|
||||
]
|
||||
calls: Final = _provider_calls(requests)
|
||||
assert (calls[0].method, urlsplit(calls[0].target).path.endswith("/responses")) == ("POST", True)
|
||||
assert [
|
||||
(call.method, urlsplit(call.target).path.endswith(f"/{upstream_id}{suffix}"))
|
||||
for call, (_, _, suffix) in zip(calls[1:], _LIFECYCLE_ROUTES, strict=True)
|
||||
] == [(method, True) for method, _, _ in _LIFECYCLE_ROUTES]
|
||||
assert {call.headers.get("authorization") for call in calls} == {f"Bearer {_identity_token(resource)}"}
|
||||
assert _identity_fetches(requests, resource) == 1
|
||||
|
|
@ -1,7 +1,7 @@
|
|||
import json
|
||||
import math
|
||||
import uuid
|
||||
from collections.abc import Sequence
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
|
@ -140,6 +140,38 @@ _PROVIDERS: Final = (
|
|||
"databricks/databricks-openjev-qwen35-4b",
|
||||
),
|
||||
)
|
||||
_LOCAL_PROVIDERS: Final = (
|
||||
pytest.param(
|
||||
_Provider("ollama", "ollama/d3:4b", "/v1/systemone", "d3:4b", None, False, None), id="regression-ollama"
|
||||
),
|
||||
pytest.param(
|
||||
_Provider("ollama_chat", "ollama_chat/d3:4b", "/v1/systemone", "d3:4b", None, False, None),
|
||||
id="regression-ollama_chat",
|
||||
),
|
||||
pytest.param(
|
||||
_Provider("hosted_vllm", "hosted_vllm/d3:4b", "/v1/systemone", "d3:4b", None, False, None),
|
||||
id="guard-hosted_vllm",
|
||||
),
|
||||
pytest.param(
|
||||
_Provider("lm_studio", "lm_studio/d3:4b", "/v1/systemone", "d3:4b", None, False, None),
|
||||
id="guard-lm_studio",
|
||||
),
|
||||
)
|
||||
_LOCAL_STATE: Final = "I was charged twice for order 4417."
|
||||
_LOCAL_QUESTIONS: Final[Mapping[str, JsonValue]] = {
|
||||
"billing": {"type": "noul", "instructions": "Is this a billing issue?"}
|
||||
}
|
||||
_LOCAL_ANSWERS: Final[Mapping[str, JsonValue]] = {"billing": {"type": "noul", "noul": 0.91}}
|
||||
_LOCAL_OPENAI_QUESTIONS: Final[Sequence[JsonValue]] = (
|
||||
{"type": "predicate", "name": "billing", "instructions": "Is this a billing issue?"},
|
||||
)
|
||||
_LOCAL_OPENAI_ANSWERS: Final[Sequence[JsonValue]] = ({"type": "predicate", "name": "billing", "probability": 0.91},)
|
||||
_LOCAL_OPENAI_USAGE: Final[Mapping[str, JsonValue]] = {
|
||||
**_USAGE,
|
||||
"input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0},
|
||||
"output_tokens_details": {"reasoning_tokens": 0},
|
||||
"total_tokens": 370,
|
||||
}
|
||||
_PERPLEXITY: Final = _PROVIDERS[0]
|
||||
_OPENROUTER: Final = next(provider for provider in _PROVIDERS if provider.name == "openrouter")
|
||||
_OPENROUTER_CHAT_MODEL: Final = "openrouter/openai/gpt-5-mini"
|
||||
|
|
@ -230,6 +262,26 @@ def _upstream_calls(gateway: Gateway, handle: ScenarioHandle) -> list[dict[str,
|
|||
return _calls_to(_observed_requests(gateway), handle)
|
||||
|
||||
|
||||
def _sdk_inference_call(gateway: Gateway, handle: ScenarioHandle, provider: _Provider) -> dict[str, JsonValue]:
|
||||
observed: Final = _upstream_calls(gateway, handle)
|
||||
inference_path: Final = f"/{handle.scenario_id}/v1/systemone"
|
||||
(call,) = tuple(request for request in observed if request["path"] == inference_path)
|
||||
metadata_calls: Final = tuple(request for request in observed if request["path"] != inference_path)
|
||||
assert metadata_calls in (
|
||||
(),
|
||||
(
|
||||
{
|
||||
"path": f"{inference_path}/api/show",
|
||||
"authorization": "",
|
||||
"body": {"name": provider.body_model},
|
||||
"method": "POST",
|
||||
"api_key": "",
|
||||
},
|
||||
),
|
||||
), observed
|
||||
return call
|
||||
|
||||
|
||||
def _spend_row(call_id: str) -> dict[str, JsonValue]:
|
||||
rows: Final = eventually(lambda: read_rows(_SPEND_QUERY, (call_id,)), lambda found: len(found) == 1, seconds=70)
|
||||
return rows[0]
|
||||
|
|
@ -345,6 +397,174 @@ def test_each_provider_gets_its_own_path_key_and_body_and_is_billed(gateway: Gat
|
|||
assert math.isclose(_number(row["spend"]), expected_spend, rel_tol=1e-9), row
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _without_ambient_ollama_settings(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.delenv("OLLAMA_API_KEY", raising=False)
|
||||
monkeypatch.delenv("OLLAMA_API_BASE", raising=False)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", _LOCAL_PROVIDERS)
|
||||
@pytest.mark.parametrize("route", ("/v1/systemone", "/v1/decisions"), ids=("systemone", "decisions"))
|
||||
@pytest.mark.parametrize(
|
||||
("base_suffix", "api_key"),
|
||||
(
|
||||
pytest.param("", None, id="root-no-key"),
|
||||
pytest.param("/", _API_KEY, id="root-slash-key"),
|
||||
pytest.param("/v1", None, id="v1-no-key"),
|
||||
pytest.param("/v1/", _API_KEY, id="v1-slash-key"),
|
||||
),
|
||||
)
|
||||
def test_local_decisions_routes_preserve_wire_auth_usage_and_spend(
|
||||
gateway: Gateway, provider: _Provider, route: str, base_suffix: str, api_key: str | None
|
||||
) -> None:
|
||||
native_reply: Final[Mapping[str, JsonValue]] = {
|
||||
"model": provider.body_model,
|
||||
"answers": dict(_LOCAL_ANSWERS),
|
||||
"usage": _USAGE,
|
||||
}
|
||||
request_body: Final[Mapping[str, JsonValue]] = (
|
||||
{"state": _LOCAL_STATE, "questions": dict(_LOCAL_QUESTIONS)}
|
||||
if route == "/v1/systemone"
|
||||
else {"input": _LOCAL_STATE, "questions": list(_LOCAL_OPENAI_QUESTIONS)}
|
||||
)
|
||||
expected_reply: Final = (
|
||||
native_reply
|
||||
if route == "/v1/systemone"
|
||||
else {"model": provider.body_model, "answers": list(_LOCAL_OPENAI_ANSWERS), "usage": dict(_LOCAL_OPENAI_USAGE)}
|
||||
)
|
||||
with gateway.scenario() as scenario:
|
||||
handle: Final = _register(scenario, dict(native_reply))
|
||||
model: Final = scenario.model(
|
||||
model=provider.model, api_base=f"{handle.api_base()}{base_suffix}", api_key=api_key
|
||||
)
|
||||
response: Final = gateway.request("POST", route, {"model": model, **request_body})
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json() == expected_reply
|
||||
assert response.headers["x-litellm-model-group"] == model
|
||||
assert float(response.headers.get("x-litellm-response-cost", "0")) == 0.0
|
||||
(call,) = _upstream_calls(gateway, handle)
|
||||
assert (call["path"], call["authorization"], call["body"]) == (
|
||||
f"/{handle.scenario_id}/v1/systemone",
|
||||
f"Bearer {api_key}" if api_key else "",
|
||||
{"model": provider.body_model, "state": _LOCAL_STATE, "questions": _LOCAL_QUESTIONS},
|
||||
)
|
||||
row: Final = _spend_row(response.headers["x-litellm-call-id"])
|
||||
assert {key: value for key, value in row.items() if key != "request_tags"} == {
|
||||
"spend": 0.0,
|
||||
"status": "success",
|
||||
"call_type": "asystemone" if route == "/v1/systemone" else "adecisions",
|
||||
"model_group": model,
|
||||
"custom_llm_provider": provider.name,
|
||||
"api_base": f"{handle.api_base()}/v1/systemone",
|
||||
"prompt_tokens": _USAGE["input_tokens"],
|
||||
"completion_tokens": _USAGE["output_tokens"],
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", _LOCAL_PROVIDERS)
|
||||
@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async"))
|
||||
async def test_local_systemone_sdk_preserves_native_wire(
|
||||
gateway: Gateway, provider: _Provider, asynchronous: bool
|
||||
) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
handle: Final = _register(scenario, _answer_body(provider))
|
||||
response: Final = (
|
||||
await litellm.asystemone(
|
||||
model=provider.model, state=_STATE, questions=_SDK_QUESTIONS, api_base=f"{handle.api_base()}/v1/"
|
||||
)
|
||||
if asynchronous
|
||||
else litellm.systemone(
|
||||
model=provider.model, state=_STATE, questions=_SDK_QUESTIONS, api_base=handle.api_base()
|
||||
)
|
||||
)
|
||||
assert response.model_dump(mode="json") == _answer_body(provider)
|
||||
call: Final = _sdk_inference_call(gateway, handle, provider)
|
||||
assert (call["path"], call["authorization"], call["body"]) == (
|
||||
f"/{handle.scenario_id}/v1/systemone",
|
||||
"",
|
||||
{"model": provider.body_model, "state": _STATE, "questions": _QUESTIONS},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", _LOCAL_PROVIDERS)
|
||||
@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async"))
|
||||
async def test_local_decisions_sdk_translates_openai_format(
|
||||
gateway: Gateway, provider: _Provider, asynchronous: bool
|
||||
) -> None:
|
||||
questions: Final = TypeAdapter(list[dict[str, object]]).validate_python(_LOCAL_OPENAI_QUESTIONS)
|
||||
native_reply: Final[Mapping[str, JsonValue]] = {
|
||||
"model": provider.body_model,
|
||||
"answers": dict(_LOCAL_ANSWERS),
|
||||
"usage": _USAGE,
|
||||
}
|
||||
with gateway.scenario() as scenario:
|
||||
handle: Final = _register(scenario, dict(native_reply))
|
||||
response: Final = (
|
||||
await litellm.adecisions(
|
||||
model=provider.model, input=_LOCAL_STATE, questions=questions, api_base=f"{handle.api_base()}/v1/"
|
||||
)
|
||||
if asynchronous
|
||||
else litellm.decisions(
|
||||
model=provider.model, input=_LOCAL_STATE, questions=questions, api_base=handle.api_base()
|
||||
)
|
||||
)
|
||||
assert response.model_dump(mode="json") == {
|
||||
"model": provider.body_model,
|
||||
"answers": list(_LOCAL_OPENAI_ANSWERS),
|
||||
"usage": dict(_LOCAL_OPENAI_USAGE),
|
||||
}
|
||||
call: Final = _sdk_inference_call(gateway, handle, provider)
|
||||
assert (call["path"], call["authorization"], call["body"]) == (
|
||||
f"/{handle.scenario_id}/v1/systemone",
|
||||
"",
|
||||
{"model": provider.body_model, "state": _LOCAL_STATE, "questions": _LOCAL_QUESTIONS},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", _LOCAL_PROVIDERS[:2])
|
||||
@pytest.mark.parametrize(
|
||||
("environment_base", "api_key", "environment_key", "authorization"),
|
||||
(
|
||||
pytest.param(True, None, _ENV_KEY, f"Bearer {_ENV_KEY}", id="environment-base-and-key"),
|
||||
pytest.param(False, _API_KEY, _ENV_KEY, f"Bearer {_API_KEY}", id="explicit-overrides-environment"),
|
||||
pytest.param(False, "fake-api-key", _ENV_KEY, "", id="fake-key-suppresses-environment"),
|
||||
pytest.param(False, None, None, "", id="unrelated-keys-not-forwarded"),
|
||||
),
|
||||
)
|
||||
def test_ollama_decisions_sdk_respects_environment_and_explicit_auth(
|
||||
gateway: Gateway,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
provider: _Provider,
|
||||
environment_base: bool,
|
||||
api_key: str | None,
|
||||
environment_key: str | None,
|
||||
authorization: str,
|
||||
) -> None:
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "synthetic-unrelated-openai-key")
|
||||
monkeypatch.setenv("LITELLM_API_KEY", "synthetic-unrelated-litellm-key")
|
||||
if environment_key is not None:
|
||||
monkeypatch.setenv("OLLAMA_API_KEY", environment_key)
|
||||
with gateway.scenario() as scenario:
|
||||
handle: Final = _register(scenario, _answer_body(provider))
|
||||
unused: Final = _register(scenario, _answer_body(provider))
|
||||
monkeypatch.setenv("OLLAMA_API_BASE", handle.api_base() if environment_base else unused.api_base())
|
||||
response: Final = litellm.systemone(
|
||||
model=provider.model,
|
||||
state=_STATE,
|
||||
questions=_SDK_QUESTIONS,
|
||||
api_base=None if environment_base else f"{handle.api_base()}/v1/",
|
||||
api_key=api_key,
|
||||
)
|
||||
assert response.model_dump(mode="json") == _answer_body(provider)
|
||||
call: Final = _sdk_inference_call(gateway, handle, provider)
|
||||
assert (call["path"], call["authorization"], call["body"]) == (
|
||||
f"/{handle.scenario_id}/v1/systemone",
|
||||
authorization,
|
||||
{"model": provider.body_model, "state": _STATE, "questions": _QUESTIONS},
|
||||
)
|
||||
assert _upstream_calls(gateway, unused) == []
|
||||
|
||||
|
||||
def test_test_connection_evaluation_mode_uses_typesafe_decisions_path(gateway: Gateway) -> None:
|
||||
provider: Final = _PROVIDERS[1]
|
||||
with gateway.scenario() as scenario:
|
||||
|
|
|
|||
|
|
@ -75,9 +75,9 @@ def test_agents_forwards_caller_supplied_key_to_chosen_api_base(gateway: Gateway
|
|||
assert request.headers["x-goog-api-key"] == "synthetic-caller-key"
|
||||
assert _ENV_GEMINI_KEY not in request.headers.values()
|
||||
assert json.loads(request.body) == {
|
||||
"name": _AGENT_NAME,
|
||||
"id": _AGENT_NAME,
|
||||
"base_agent": "waverunner",
|
||||
"instructions": "make slides",
|
||||
"system_instruction": "make slides",
|
||||
}
|
||||
return Reply(body=json.dumps(_AGENT_REPLY).encode())
|
||||
|
||||
|
|
|
|||
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 (
|
||||
|
|
|
|||
|
|
@ -578,8 +578,14 @@ async def test_unknown_provider_is_rejected_before_http(respx_mock: respx.MockRo
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_provider_without_decisions_support_is_rejected_before_http(respx_mock: respx.MockRouter) -> None:
|
||||
with pytest.raises(litellm.BadRequestError, match=r"Unknown Decisions provider 'anthropic'\. Supported providers"):
|
||||
async def test_provider_without_decisions_support_is_rejected_before_http(
|
||||
monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
|
||||
) -> None:
|
||||
monkeypatch.delenv("OLLAMA_API_KEY", raising=False)
|
||||
monkeypatch.delenv("OLLAMA_API_BASE", raising=False)
|
||||
with pytest.raises(
|
||||
litellm.BadRequestError, match=r"Unknown Decisions provider 'anthropic'\. Supported providers"
|
||||
) as error:
|
||||
await litellm.asystemone(
|
||||
model="anthropic/claude-sonnet-4-5",
|
||||
state="review",
|
||||
|
|
@ -587,6 +593,8 @@ async def test_provider_without_decisions_support_is_rejected_before_http(respx_
|
|||
api_key="caller-key",
|
||||
)
|
||||
|
||||
supported: Final = str(error.value).split("Supported providers: ", 1)[1].split(", ")
|
||||
assert {"ollama", "ollama_chat"} <= set(supported)
|
||||
assert len(respx_mock.calls) == 0
|
||||
|
||||
|
||||
|
|
@ -1536,6 +1544,8 @@ _PROVIDERS_WHOSE_CHAT_MODELS_ALSO_DECIDE: Final = frozenset(
|
|||
LlmProviders.OPENROUTER.value,
|
||||
LlmProviders.HOSTED_VLLM.value,
|
||||
LlmProviders.LM_STUDIO.value,
|
||||
LlmProviders.OLLAMA.value,
|
||||
LlmProviders.OLLAMA_CHAT.value,
|
||||
}
|
||||
)
|
||||
_DEDICATED_DECIDER_PROVIDERS: Final[tuple[str, ...]] = tuple(
|
||||
|
|
@ -3911,6 +3921,185 @@ async def test_native_decisions_distinguishes_missing_usage_from_reported_zero(
|
|||
assert (expected == 0) == (usage["input_tokens"] == usage["output_tokens"] == 0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("provider", ("ollama", "ollama_chat"))
|
||||
@pytest.mark.parametrize("use_async", (True, False), ids=("asystemone", "systemone"))
|
||||
@pytest.mark.parametrize("api_base", ("http://ollama.local:11434", "http://ollama.local:11434/v1"))
|
||||
@pytest.mark.parametrize("api_key", (None, "fake-api-key"))
|
||||
async def test_ollama_systemone_preserves_questions_without_unrelated_or_placeholder_auth(
|
||||
provider: str,
|
||||
use_async: bool,
|
||||
api_base: str,
|
||||
api_key: str | None,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
respx_mock: respx.MockRouter,
|
||||
) -> None:
|
||||
monkeypatch.delenv("OLLAMA_API_KEY", raising=False)
|
||||
monkeypatch.delenv("OLLAMA_API_BASE", raising=False)
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "unrelated-openai-key")
|
||||
monkeypatch.setattr(litellm, "api_key", "unrelated-global-key")
|
||||
monkeypatch.setattr(litellm, "openai_key", "unrelated-global-openai-key")
|
||||
route: Final = respx_mock.post("http://ollama.local:11434/v1/systemone").respond(json=_RESPONSE)
|
||||
kwargs: Final = {
|
||||
"model": f"{provider}/d3:4b",
|
||||
"state": "review",
|
||||
"questions": _QUESTIONS,
|
||||
"api_base": api_base,
|
||||
"api_key": api_key,
|
||||
}
|
||||
|
||||
response: Final = await litellm.asystemone(**kwargs) if use_async else litellm.systemone(**kwargs)
|
||||
|
||||
assert json.loads(route.calls.last.request.content) == {
|
||||
"model": "d3:4b",
|
||||
"state": "review",
|
||||
"questions": dict(_QUESTIONS),
|
||||
}
|
||||
assert "authorization" not in route.calls.last.request.headers
|
||||
assert response.answers["is_defect"] == NoulAnswer(type="noul", noul=0.9)
|
||||
assert response.answers["sentiment"].choice == "positive"
|
||||
assert response.answers["severity"].score == 1
|
||||
assert response.usage == DecisionsUsage(input_tokens=_INPUT_TOKENS, output_tokens=_OUTPUT_TOKENS)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("provider", ("ollama", "ollama_chat"))
|
||||
@pytest.mark.parametrize("use_async", (True, False), ids=("adecisions", "decisions"))
|
||||
async def test_ollama_decisions_translates_openai_questions_and_answers(
|
||||
provider: str, use_async: bool, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
|
||||
) -> None:
|
||||
monkeypatch.delenv("OLLAMA_API_KEY", raising=False)
|
||||
monkeypatch.delenv("OLLAMA_API_BASE", raising=False)
|
||||
route: Final = respx_mock.post("http://ollama.local:11434/v1/systemone").respond(json=_RESPONSE)
|
||||
kwargs: Final = {
|
||||
"model": f"{provider}/d3:4b",
|
||||
"input": "review",
|
||||
"questions": [{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}],
|
||||
"api_base": "http://ollama.local:11434/v1",
|
||||
}
|
||||
|
||||
response: Final = await litellm.adecisions(**kwargs) if use_async else litellm.decisions(**kwargs)
|
||||
|
||||
assert json.loads(route.calls.last.request.content) == {
|
||||
"model": "d3:4b",
|
||||
"state": "review",
|
||||
"questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}},
|
||||
}
|
||||
assert "authorization" not in route.calls.last.request.headers
|
||||
assert response.answers[0].type == "predicate"
|
||||
assert response.answers[0].name == "is_defect"
|
||||
assert response.answers[0].probability == 0.9
|
||||
assert response.usage is not None
|
||||
assert response.usage.input_tokens == _INPUT_TOKENS
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("provider", ("ollama", "ollama_chat"))
|
||||
@pytest.mark.parametrize("environment_base", (None, "http://ollama.local:11434/v1"))
|
||||
async def test_ollama_decisions_resolves_environment_base_or_localhost_default(
|
||||
provider: str, environment_base: str | None, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
|
||||
) -> None:
|
||||
monkeypatch.delenv("OLLAMA_API_KEY", raising=False)
|
||||
monkeypatch.delenv("OLLAMA_API_BASE", raising=False)
|
||||
if environment_base is not None:
|
||||
monkeypatch.setenv("OLLAMA_API_BASE", environment_base)
|
||||
url: Final = f"{environment_base or 'http://localhost:11434/v1'}/systemone"
|
||||
route: Final = respx_mock.post(url).respond(json=_RESPONSE)
|
||||
|
||||
response: Final = await litellm.asystemone(model=f"{provider}/d3:4b", state="review", questions=_QUESTIONS)
|
||||
|
||||
assert route.call_count == 1
|
||||
assert "authorization" not in route.calls.last.request.headers
|
||||
assert response.answers["is_defect"] == NoulAnswer(type="noul", noul=0.9)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("provider", ("ollama", "ollama_chat"))
|
||||
@pytest.mark.parametrize("explicit_key", (None, "deployment-key"))
|
||||
async def test_ollama_decisions_resolves_environment_key_and_explicit_connection_overrides(
|
||||
provider: str, explicit_key: str | None, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
|
||||
) -> None:
|
||||
monkeypatch.delenv("OLLAMA_API_KEY", raising=False)
|
||||
monkeypatch.delenv("OLLAMA_API_BASE", raising=False)
|
||||
monkeypatch.setenv("OLLAMA_API_BASE", "http://ignored.local:11434")
|
||||
monkeypatch.setenv("OLLAMA_API_KEY", "environment-key")
|
||||
route: Final = respx_mock.post("http://ollama.local:11434/v1/systemone").respond(json=_RESPONSE)
|
||||
|
||||
response: Final = await litellm.asystemone(
|
||||
model=f"{provider}/d3:4b",
|
||||
state="review",
|
||||
questions=_QUESTIONS,
|
||||
api_base="http://ollama.local:11434",
|
||||
api_key=explicit_key,
|
||||
)
|
||||
|
||||
assert route.calls.last.request.headers["authorization"] == f"Bearer {explicit_key or 'environment-key'}"
|
||||
assert response.answers["is_defect"] == NoulAnswer(type="noul", noul=0.9)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("provider", ("ollama", "ollama_chat"))
|
||||
async def test_ollama_decisions_maps_upstream_bad_requests(
|
||||
provider: str, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
|
||||
) -> None:
|
||||
monkeypatch.delenv("OLLAMA_API_KEY", raising=False)
|
||||
monkeypatch.delenv("OLLAMA_API_BASE", raising=False)
|
||||
route: Final = respx_mock.post("http://ollama.local:11434/v1/systemone").respond(
|
||||
400, json={"error": "decision rejected"}
|
||||
)
|
||||
|
||||
with pytest.raises(litellm.BadRequestError, match="decision rejected") as error:
|
||||
await litellm.asystemone(
|
||||
model=f"{provider}/d3:4b",
|
||||
state="review",
|
||||
questions=_QUESTIONS,
|
||||
api_base="http://ollama.local:11434",
|
||||
num_retries=0,
|
||||
)
|
||||
|
||||
assert error.value.status_code == 400
|
||||
assert route.call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("provider", ("ollama", "ollama_chat"))
|
||||
@pytest.mark.parametrize("drop_params", (False, True))
|
||||
@pytest.mark.parametrize("image", ("data:image/png;base64,AA==", "https://images.example/input.png"))
|
||||
async def test_ollama_decisions_rejects_input_images_before_http_or_drops_them(
|
||||
provider: str, drop_params: bool, image: str, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
|
||||
) -> None:
|
||||
monkeypatch.delenv("OLLAMA_API_KEY", raising=False)
|
||||
monkeypatch.delenv("OLLAMA_API_BASE", raising=False)
|
||||
monkeypatch.setattr(litellm, "drop_params", False)
|
||||
kwargs: Final = {
|
||||
"model": f"{provider}/d3:4b",
|
||||
"input": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "input_text", "text": "review"}, {"type": "input_image", "image_url": image}],
|
||||
}
|
||||
],
|
||||
"questions": [{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}],
|
||||
"api_base": "http://ollama.local:11434",
|
||||
"drop_params": drop_params,
|
||||
}
|
||||
if not drop_params:
|
||||
with pytest.raises(litellm.UnsupportedParamsError, match=r"\['input_image'\].*drop_params"):
|
||||
await litellm.adecisions(**kwargs)
|
||||
assert len(respx_mock.calls) == 0
|
||||
return
|
||||
route: Final = respx_mock.post("http://ollama.local:11434/v1/systemone").respond(json=_RESPONSE)
|
||||
|
||||
response: Final = await litellm.adecisions(**kwargs)
|
||||
|
||||
assert json.loads(route.calls.last.request.content) == {
|
||||
"model": "d3:4b",
|
||||
"state": "review",
|
||||
"questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}},
|
||||
}
|
||||
assert response.answers[0].probability == 0.9
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("use_async", (True, False), ids=("asystemone", "systemone"))
|
||||
@pytest.mark.parametrize("api_base", ("http://lmstudio.local:1234", "http://lmstudio.local:1234/v1"))
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
|||
from litellm.llms.gemini.agents.transformation import GeminiAgentsConfig
|
||||
from litellm.llms.gemini.common_utils import GeminiError
|
||||
from litellm.types.agents import (
|
||||
AgentCreateRequest,
|
||||
AgentCreateResponse,
|
||||
AgentDeleteResult,
|
||||
AgentListResponse,
|
||||
|
|
@ -104,7 +105,7 @@ class TestCreateAgent:
|
|||
|
||||
result = handler.create_agent(
|
||||
agents_api_config=config,
|
||||
name="agent-x",
|
||||
create_request=AgentCreateRequest(id="agent-x"),
|
||||
litellm_params=litellm_params,
|
||||
logging_obj=logging_obj,
|
||||
extra_headers={"X-Test": "1"},
|
||||
|
|
@ -117,8 +118,7 @@ class TestCreateAgent:
|
|||
client.post.assert_called_once()
|
||||
kwargs = client.post.call_args.kwargs
|
||||
assert kwargs["url"].endswith("/v1beta/agents")
|
||||
assert kwargs["json"]["name"] == "agent-x"
|
||||
assert kwargs["json"]["foo"] == "bar"
|
||||
assert kwargs["json"] == {"id": "agent-x", "foo": "bar"}
|
||||
assert kwargs["headers"]["X-Test"] == "1"
|
||||
logging_obj.pre_call.assert_called_once()
|
||||
logging_obj.post_call.assert_called_once()
|
||||
|
|
@ -130,7 +130,7 @@ class TestCreateAgent:
|
|||
|
||||
result = handler.create_agent(
|
||||
agents_api_config=config,
|
||||
name="agent-x",
|
||||
create_request=AgentCreateRequest(id="agent-x"),
|
||||
litellm_params=litellm_params,
|
||||
logging_obj=_make_logging_obj(),
|
||||
client=client,
|
||||
|
|
@ -152,7 +152,7 @@ class TestCreateAgent:
|
|||
with pytest.raises(GeminiError):
|
||||
handler.create_agent(
|
||||
agents_api_config=config,
|
||||
name="agent-x",
|
||||
create_request=AgentCreateRequest(id="agent-x"),
|
||||
litellm_params=litellm_params,
|
||||
logging_obj=_make_logging_obj(),
|
||||
client=client,
|
||||
|
|
@ -169,7 +169,7 @@ class TestCreateAgent:
|
|||
|
||||
result = await handler.async_create_agent(
|
||||
agents_api_config=config,
|
||||
name="agent-y",
|
||||
create_request=AgentCreateRequest(id="agent-y"),
|
||||
litellm_params=litellm_params,
|
||||
logging_obj=_make_logging_obj(),
|
||||
extra_body={"baz": "qux"},
|
||||
|
|
@ -193,7 +193,7 @@ class TestCreateAgent:
|
|||
with pytest.raises(GeminiError):
|
||||
await handler.async_create_agent(
|
||||
agents_api_config=config,
|
||||
name="agent-y",
|
||||
create_request=AgentCreateRequest(id="agent-y"),
|
||||
litellm_params=litellm_params,
|
||||
logging_obj=_make_logging_obj(),
|
||||
client=client,
|
||||
|
|
|
|||
|
|
@ -9,9 +9,14 @@ small helper utilities without touching the network.
|
|||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from collections.abc import Callable, Mapping
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
|
||||
import litellm
|
||||
|
|
@ -137,7 +142,7 @@ class TestSyncEntryPoints:
|
|||
assert response is sentinel
|
||||
handler.create_agent.assert_called_once()
|
||||
kw = handler.create_agent.call_args.kwargs
|
||||
assert kw["name"] == "waverunner"
|
||||
assert kw["create_request"].id == "waverunner"
|
||||
assert kw["_is_async"] is False
|
||||
assert kw["extra_headers"] == {"X-Test": "1"}
|
||||
assert kw["extra_body"] == {"foo": "bar"}
|
||||
|
|
@ -205,7 +210,7 @@ class TestAsyncEntryPoints:
|
|||
|
||||
def fake_create_agent(**kwargs):
|
||||
assert kwargs["_is_async"] is True
|
||||
assert kwargs["name"] == "waverunner"
|
||||
assert kwargs["create_request"].id == "waverunner"
|
||||
return sentinel
|
||||
|
||||
handler = MagicMock()
|
||||
|
|
@ -349,3 +354,250 @@ class TestAsyncErrorWrapping:
|
|||
with patch(_HANDLER_PATH, handler):
|
||||
with pytest.raises(litellm.APIConnectionError):
|
||||
await alist_versions(name="waverunner", api_key="AIza")
|
||||
|
||||
|
||||
_GEMINI_AGENTS_URL: Final = "https://generativelanguage.googleapis.com/v1beta/agents"
|
||||
|
||||
|
||||
def _sent_json(route: respx.Route) -> object:
|
||||
return json.loads(route.calls.last.request.content)
|
||||
|
||||
|
||||
class TestCreateRequestBody:
|
||||
def test_create_sends_every_current_field_verbatim(self, respx_mock: respx.MockRouter) -> None:
|
||||
route: Final = respx_mock.post(_GEMINI_AGENTS_URL).mock(
|
||||
return_value=httpx.Response(200, json={"id": "sentinel", "base_agent": "antigravity-preview-09-2026"})
|
||||
)
|
||||
|
||||
response: Final = create(
|
||||
id="sentinel",
|
||||
base_agent="antigravity-preview-09-2026",
|
||||
agent_config={"type": "antigravity", "model": "gemini-3.8-flash", "max_total_tokens": None},
|
||||
system_instruction="Always answer in exactly three words.",
|
||||
base_environment={"type": "remote", "sources": [{"type": "gcs", "source": "gs://bucket/skill"}]},
|
||||
description="Watches the queue.",
|
||||
tools=[{"type": "code_execution"}],
|
||||
api_key="AIza-test",
|
||||
)
|
||||
|
||||
assert _sent_json(route) == {
|
||||
"id": "sentinel",
|
||||
"base_agent": "antigravity-preview-09-2026",
|
||||
"agent_config": {"type": "antigravity", "model": "gemini-3.8-flash", "max_total_tokens": None},
|
||||
"system_instruction": "Always answer in exactly three words.",
|
||||
"base_environment": {"type": "remote", "sources": [{"type": "gcs", "source": "gs://bucket/skill"}]},
|
||||
"description": "Watches the queue.",
|
||||
"tools": [{"type": "code_execution"}],
|
||||
}
|
||||
assert (response.id, response.name) == ("sentinel", "sentinel")
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("identifier", "instruction", "expected"),
|
||||
[
|
||||
pytest.param(
|
||||
{"name": "slides"},
|
||||
{"instructions": "Make slides."},
|
||||
{"id": "slides", "system_instruction": "Make slides."},
|
||||
id="preview-aliases",
|
||||
),
|
||||
pytest.param(
|
||||
{"id": "slides", "name": "slides"},
|
||||
{"system_instruction": "Make slides.", "instructions": "Make slides."},
|
||||
{"id": "slides", "system_instruction": "Make slides."},
|
||||
id="same-value-under-both-names",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_create_sends_preview_aliases_under_current_names(
|
||||
self,
|
||||
respx_mock: respx.MockRouter,
|
||||
identifier: Mapping[str, str],
|
||||
instruction: Mapping[str, str],
|
||||
expected: Mapping[str, str],
|
||||
) -> None:
|
||||
route: Final = respx_mock.post(_GEMINI_AGENTS_URL).mock(
|
||||
return_value=httpx.Response(200, json={"id": "slides"})
|
||||
)
|
||||
|
||||
create(**identifier, **instruction, base_agent="antigravity-preview-09-2026", api_key="AIza-test")
|
||||
|
||||
assert _sent_json(route) == {**expected, "base_agent": "antigravity-preview-09-2026"}
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"fields",
|
||||
[
|
||||
pytest.param({"base_agent": "antigravity-preview-09-2026"}, id="no-identifier"),
|
||||
pytest.param({"id": "slides", "name": "decks"}, id="different-id-and-name"),
|
||||
pytest.param(
|
||||
{"id": "slides", "system_instruction": "Make slides.", "instructions": "Make decks."},
|
||||
id="different-system-instruction-and-instructions",
|
||||
),
|
||||
pytest.param({"id": "slides", "agent_config": "antigravity"}, id="agent-config-not-an-object"),
|
||||
],
|
||||
)
|
||||
def test_create_rejects_a_malformed_identifier_or_instruction(
|
||||
self, respx_mock: respx.MockRouter, fields: Mapping[str, str]
|
||||
) -> None:
|
||||
with pytest.raises(litellm.BadRequestError) as exc_info:
|
||||
create(**fields, api_key="AIza-test")
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert not respx_mock.calls
|
||||
|
||||
def test_create_with_drop_params_drops_only_the_invalid_optional_fields(
|
||||
self, respx_mock: respx.MockRouter
|
||||
) -> None:
|
||||
route: Final = respx_mock.post(_GEMINI_AGENTS_URL).mock(
|
||||
return_value=httpx.Response(200, json={"id": "slides"})
|
||||
)
|
||||
|
||||
create(
|
||||
id="slides",
|
||||
agent_config="antigravity",
|
||||
tools="code_execution",
|
||||
instructions="Make slides.",
|
||||
drop_params=True,
|
||||
api_key="AIza-test",
|
||||
)
|
||||
|
||||
assert _sent_json(route) == {"id": "slides", "system_instruction": "Make slides."}
|
||||
|
||||
def test_create_with_drop_params_drops_a_clashing_instructions_alias(self, respx_mock: respx.MockRouter) -> None:
|
||||
route: Final = respx_mock.post(_GEMINI_AGENTS_URL).mock(
|
||||
return_value=httpx.Response(200, json={"id": "slides"})
|
||||
)
|
||||
|
||||
create(
|
||||
id="slides",
|
||||
system_instruction="Make slides.",
|
||||
instructions="Make decks.",
|
||||
drop_params=True,
|
||||
api_key="AIza-test",
|
||||
)
|
||||
|
||||
assert _sent_json(route) == {"id": "slides", "system_instruction": "Make slides."}
|
||||
|
||||
def test_create_with_drop_params_still_rejects_a_clashing_identifier(self, respx_mock: respx.MockRouter) -> None:
|
||||
with pytest.raises(litellm.BadRequestError) as exc_info:
|
||||
create(id="slides", name="decks", drop_params=True, api_key="AIza-test")
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert not respx_mock.calls
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("fields", "expected_message"),
|
||||
[
|
||||
pytest.param(
|
||||
{"id": "slides", "instructions": 5, "tools": "code_execution"},
|
||||
"Invalid agent create request: 'instructions' must be a string; 'tools' must be a list of objects."
|
||||
" Set drop_params to drop invalid optional fields.",
|
||||
id="names-the-alias-the-caller-sent",
|
||||
),
|
||||
pytest.param(
|
||||
{"name": 7, "base_agent": "antigravity-preview-09-2026"},
|
||||
"Invalid agent create request: 'name' must be a string.",
|
||||
id="names-the-id-alias-the-caller-sent",
|
||||
),
|
||||
pytest.param(
|
||||
{"id": 7, "agent_config": "antigravity", "drop_params": True},
|
||||
"Invalid agent create request: 'agent_config' must be an object; 'id' must be a string.",
|
||||
id="invalid-id-is-never-dropped",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_create_rejects_invalid_fields_by_the_names_the_caller_sent(
|
||||
self, respx_mock: respx.MockRouter, fields: Mapping[str, object], expected_message: str
|
||||
) -> None:
|
||||
with pytest.raises(litellm.BadRequestError) as exc_info:
|
||||
create(**fields, api_key="AIza-test")
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert exc_info.value.message.endswith(expected_message)
|
||||
assert not respx_mock.calls
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acreate_sends_every_current_field_verbatim(
|
||||
self, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
|
||||
route: Final = respx_mock.post(_GEMINI_AGENTS_URL).mock(
|
||||
return_value=httpx.Response(200, json={"id": "sentinel"})
|
||||
)
|
||||
|
||||
response: Final = await acreate(
|
||||
id="sentinel",
|
||||
agent_config={"type": "antigravity", "model": "gemini-3.8-flash"},
|
||||
system_instruction="Always answer in exactly three words.",
|
||||
api_key="AIza-test",
|
||||
)
|
||||
|
||||
assert _sent_json(route) == {
|
||||
"id": "sentinel",
|
||||
"agent_config": {"type": "antigravity", "model": "gemini-3.8-flash"},
|
||||
"system_instruction": "Always answer in exactly three words.",
|
||||
}
|
||||
assert response.id == "sentinel"
|
||||
|
||||
|
||||
class TestAgentIdentifier:
|
||||
def test_get_by_id_fetches_that_agent(self, respx_mock: respx.MockRouter) -> None:
|
||||
route: Final = respx_mock.get(f"{_GEMINI_AGENTS_URL}/slides").mock(
|
||||
return_value=httpx.Response(200, json={"id": "slides", "base_agent": "antigravity-preview-09-2026"})
|
||||
)
|
||||
|
||||
response: Final = get(id="slides", api_key="AIza-test")
|
||||
|
||||
assert route.call_count == 1
|
||||
assert (response.id, response.base_agent) == ("slides", "antigravity-preview-09-2026")
|
||||
|
||||
def test_delete_by_id_deletes_that_agent(self, respx_mock: respx.MockRouter) -> None:
|
||||
route: Final = respx_mock.delete(f"{_GEMINI_AGENTS_URL}/slides").mock(
|
||||
return_value=httpx.Response(200, json={})
|
||||
)
|
||||
|
||||
result: Final = delete(id="slides", api_key="AIza-test")
|
||||
|
||||
assert route.call_count == 1
|
||||
assert (result.name, result.deleted) == ("slides", True)
|
||||
|
||||
def test_list_versions_by_id_lists_that_agents_versions(self, respx_mock: respx.MockRouter) -> None:
|
||||
version: Final = {"agent": "slides", "name": "agents/slides/versions/v1"}
|
||||
route: Final = respx_mock.get(f"{_GEMINI_AGENTS_URL}/slides/versions").mock(
|
||||
return_value=httpx.Response(200, json={"agentVersions": [version]})
|
||||
)
|
||||
|
||||
response: Final = list_versions(id="slides", api_key="AIza-test")
|
||||
|
||||
assert route.call_count == 1
|
||||
assert response.agent_versions == [version]
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"operation",
|
||||
[pytest.param(get, id="get"), pytest.param(delete, id="delete"), pytest.param(list_versions, id="list_versions")],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"identifier",
|
||||
[
|
||||
pytest.param({}, id="no-identifier"),
|
||||
pytest.param({"id": "slides", "name": "decks"}, id="different-id-and-name"),
|
||||
],
|
||||
)
|
||||
def test_rejects_a_missing_or_conflicting_identifier(
|
||||
self,
|
||||
respx_mock: respx.MockRouter,
|
||||
operation: Callable[..., object],
|
||||
identifier: Mapping[str, str],
|
||||
) -> None:
|
||||
with pytest.raises(litellm.BadRequestError) as exc_info:
|
||||
operation(**identifier, api_key="AIza-test")
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert not respx_mock.calls
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aget_rejects_a_missing_identifier(self, respx_mock: respx.MockRouter) -> None:
|
||||
with pytest.raises(litellm.BadRequestError) as exc_info:
|
||||
await aget(api_key="AIza-test")
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert not respx_mock.calls
|
||||
|
|
|
|||
|
|
@ -236,6 +236,159 @@ class TestGoogleInteractionsCreate:
|
|||
assert json.loads(route.calls.last.request.content) == expected_body
|
||||
assert response.status == "completed"
|
||||
|
||||
@pytest.mark.usefixtures("fake_provider_credentials")
|
||||
def test_create_with_agent_config_sends_it_to_gemini(self, respx_mock: MockRouter) -> None:
|
||||
route: Final = respx_mock.post("https://generativelanguage.googleapis.com/v1beta/interactions").mock(
|
||||
return_value=httpx.Response(
|
||||
200,
|
||||
json={"id": "interaction-agent", "object": "interaction", "status": "completed"},
|
||||
)
|
||||
)
|
||||
|
||||
interactions.create(
|
||||
agent="antigravity-preview-09-2026",
|
||||
agent_config={"type": "antigravity", "model": "gemini-3.8-flash"},
|
||||
input="Say hi in one word.",
|
||||
environment="remote",
|
||||
api_key="gemini-offline",
|
||||
)
|
||||
|
||||
assert json.loads(route.calls.last.request.content) == {
|
||||
"agent": "antigravity-preview-09-2026",
|
||||
"agent_config": {"type": "antigravity", "model": "gemini-3.8-flash"},
|
||||
"input": "Say hi in one word.",
|
||||
"environment": "remote",
|
||||
}
|
||||
|
||||
@pytest.mark.usefixtures("fake_provider_credentials")
|
||||
def test_create_keeps_positional_response_settings_in_place(self, respx_mock: MockRouter) -> None:
|
||||
route: Final = respx_mock.post("https://generativelanguage.googleapis.com/v1beta/interactions").mock(
|
||||
return_value=httpx.Response(200, json={"id": "interaction-model", "status": "completed"})
|
||||
)
|
||||
|
||||
interactions.create(
|
||||
"gemini/gemini-2.5-flash",
|
||||
None,
|
||||
"Describe a cat.",
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
["TEXT"],
|
||||
api_key="gemini-offline",
|
||||
)
|
||||
|
||||
assert json.loads(route.calls.last.request.content) == {
|
||||
"model": "gemini-2.5-flash",
|
||||
"input": "Describe a cat.",
|
||||
"response_modalities": ["TEXT"],
|
||||
}
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("target", "agent_config", "expected_target"),
|
||||
[
|
||||
pytest.param(
|
||||
{"model": "gemini/gemini-2.5-flash"},
|
||||
{"type": "antigravity", "model": "gemini-3.8-flash"},
|
||||
{"model": "gemini-2.5-flash"},
|
||||
id="model-drops-it",
|
||||
),
|
||||
pytest.param(
|
||||
{"agent": "antigravity-preview-09-2026"},
|
||||
{"type": "antigravity", "model": "gemini-3.8-flash"},
|
||||
{
|
||||
"agent": "antigravity-preview-09-2026",
|
||||
"agent_config": {"type": "antigravity", "model": "gemini-3.8-flash"},
|
||||
},
|
||||
id="agent-keeps-an-object",
|
||||
),
|
||||
pytest.param(
|
||||
{"agent": "antigravity-preview-09-2026"},
|
||||
"antigravity",
|
||||
{"agent": "antigravity-preview-09-2026"},
|
||||
id="agent-drops-a-non-object",
|
||||
),
|
||||
pytest.param(
|
||||
{"agent": "antigravity-preview-09-2026"},
|
||||
{"model": "gemini-3.8-flash"},
|
||||
{"agent": "antigravity-preview-09-2026"},
|
||||
id="agent-drops-an-object-without-a-type",
|
||||
),
|
||||
pytest.param(
|
||||
{"agent": "antigravity-preview-09-2026"},
|
||||
{"type": 7, "model": "gemini-3.8-flash"},
|
||||
{"agent": "antigravity-preview-09-2026"},
|
||||
id="agent-drops-a-non-string-type",
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.usefixtures("fake_provider_credentials")
|
||||
def test_create_with_drop_params_drops_agent_config_gemini_cannot_take(
|
||||
self,
|
||||
respx_mock: MockRouter,
|
||||
target: Mapping[str, str],
|
||||
agent_config: object,
|
||||
expected_target: Mapping[str, object],
|
||||
) -> None:
|
||||
route: Final = respx_mock.post("https://generativelanguage.googleapis.com/v1beta/interactions").mock(
|
||||
return_value=httpx.Response(200, json={"id": "interaction-drop", "status": "completed"})
|
||||
)
|
||||
|
||||
interactions.create(
|
||||
**target,
|
||||
agent_config=agent_config,
|
||||
input="Say hi in one word.",
|
||||
drop_params=True,
|
||||
api_key="gemini-offline",
|
||||
)
|
||||
|
||||
assert json.loads(route.calls.last.request.content) == {**expected_target, "input": "Say hi in one word."}
|
||||
|
||||
@pytest.mark.usefixtures("fake_provider_credentials")
|
||||
def test_create_rejects_agent_config_a_bridged_provider_cannot_take(self, respx_mock: MockRouter) -> None:
|
||||
route: Final = respx_mock.post("https://api.openai.com/v1/responses").mock(
|
||||
return_value=httpx.Response(200, json=_openai_response_body("resp-agent-config", "Hi"))
|
||||
)
|
||||
|
||||
with pytest.raises(litellm.UnsupportedParamsError, match="'agent_config' is not supported for openai") as e:
|
||||
interactions.create(
|
||||
model="gpt-5.6",
|
||||
agent_config={"type": "antigravity", "model": "gemini-3.8-flash"},
|
||||
input="Say hi in one word.",
|
||||
api_key="sk-offline",
|
||||
)
|
||||
|
||||
assert e.value.status_code == 400
|
||||
assert not route.called
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"drop_kwargs",
|
||||
[
|
||||
pytest.param({"drop_params": True}, id="drop-params"),
|
||||
pytest.param({"additional_drop_params": ["agent_config"]}, id="additional-drop-params"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.usefixtures("fake_provider_credentials")
|
||||
def test_create_with_drop_params_drops_agent_config_a_bridged_provider_cannot_take(
|
||||
self, respx_mock: MockRouter, drop_kwargs: Mapping[str, object]
|
||||
) -> None:
|
||||
route: Final = respx_mock.post("https://api.openai.com/v1/responses").mock(
|
||||
return_value=httpx.Response(200, json=_openai_response_body("resp-agent-config", "Hi"))
|
||||
)
|
||||
|
||||
interactions.create(
|
||||
model="gpt-5.6",
|
||||
agent_config={"type": "antigravity", "model": "gemini-3.8-flash"},
|
||||
input="Say hi in one word.",
|
||||
api_key="sk-offline",
|
||||
**drop_kwargs,
|
||||
)
|
||||
|
||||
assert json.loads(route.calls.last.request.content) == {"model": "gpt-5.6", "input": "Say hi in one word."}
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "api_key", "url", "expected_body", "sse_body", "expected_chunks"),
|
||||
[
|
||||
|
|
|
|||
|
|
@ -37,6 +37,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
|||
schema_has_unresolvable_ref,
|
||||
schema_property_count,
|
||||
split_concatenated_json_objects,
|
||||
split_media_for_tpm,
|
||||
strip_encrypted_reasoning_from_messages,
|
||||
strip_litellm_response_only_fields,
|
||||
system_messages_first,
|
||||
|
|
@ -2512,3 +2513,118 @@ def test_shared_image_url_agrees_with_responses_normalization(
|
|||
expected: Final = {"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE"}}
|
||||
assert as_openai_image_part(part) == expected
|
||||
assert as_openai_image_part(normalized) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"part,projected,media",
|
||||
[
|
||||
(
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE", "detail": "high"}},
|
||||
{"type": "image_url", "image_url": {"url": "", "detail": "high"}},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE", "detail": "high"}},
|
||||
),
|
||||
(
|
||||
{"type": "input_image", "image_url": "https://example.com/image.png", "detail": "low"},
|
||||
{"type": "input_image", "image_url": "", "detail": "low"},
|
||||
{"type": "image_url", "image_url": {"url": "https://example.com/image.png", "detail": "low"}},
|
||||
),
|
||||
(
|
||||
{"type": "computer_screenshot", "image_url": "data:image/png;base64,IMAGE"},
|
||||
{"type": "computer_screenshot", "image_url": ""},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE", "detail": "auto"}},
|
||||
),
|
||||
(
|
||||
{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "IMAGE"}},
|
||||
{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": ""}},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE", "detail": "auto"}},
|
||||
),
|
||||
(
|
||||
{"type": "input_audio", "input_audio": {"data": "AUDIO", "format": "wav"}},
|
||||
{"type": "input_audio", "input_audio": {"data": "", "format": "wav"}},
|
||||
{"type": "input_audio", "input_audio": {"data": "AUDIO", "format": "wav"}},
|
||||
),
|
||||
(
|
||||
{"type": "file", "file": {"file_data": "IMAGE", "format": "image/png", "filename": "image.png"}},
|
||||
{"type": "file", "file": {"file_data": "", "format": "image/png", "filename": "image.png"}},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE", "detail": "auto"}},
|
||||
),
|
||||
(
|
||||
{"type": "input_file", "file_data": "data:audio/wav;base64,AUDIO", "filename": "audio.wav"},
|
||||
{"type": "input_file", "file_data": "", "filename": "audio.wav"},
|
||||
{"type": "input_audio", "input_audio": {"data": "AUDIO", "format": "wav"}},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_tpm_projection_blanks_only_declared_media_transport(
|
||||
part: Mapping[str, object], projected: Mapping[str, object], media: Mapping[str, object],
|
||||
) -> None:
|
||||
text: Final = {"type": "text", "text": "data:image/png;base64,IMAGE"}
|
||||
request: Final = [{"role": "user", "content": [text, {**part, "cache_control": {"type": "ephemeral"}}]}]
|
||||
original: Final = copy.deepcopy(request)
|
||||
assert split_media_for_tpm(request) == (
|
||||
[{"role": "user", "content": [text, {**projected, "cache_control": {"type": "ephemeral"}}]}],
|
||||
(media,),
|
||||
)
|
||||
assert request == original
|
||||
|
||||
|
||||
def test_tpm_projection_retains_nested_and_native_history_around_media() -> None:
|
||||
image: Final = {"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE", "detail": "high"}}
|
||||
empty_image: Final = {"type": "image_url", "image_url": {"url": "", "detail": "high"}}
|
||||
document: Final = {"type": "document", "title": "keep", "source": {"type": "content", "content": [image]}}
|
||||
empty_document: Final = {**document, "source": {"type": "content", "content": [empty_image]}}
|
||||
tool: Final = {"type": "tool_result", "tool_use_id": "call", "content": [document, {"type": "text", "text": "keep"}]}
|
||||
empty_tool: Final = {**tool, "content": [empty_document, {"type": "text", "text": "keep"}]}
|
||||
request: Final = [
|
||||
{"role": "user", "content": [tool]},
|
||||
{"type": "function_call_output", "call_id": "call", "output": [image]},
|
||||
{"type": "custom_tool_call_output", "call_id": "custom", "output": [image]},
|
||||
{"type": "computer_call_output", "call_id": "computer", "output": {
|
||||
"type": "computer_screenshot", "image_url": "data:image/png;base64,IMAGE",
|
||||
}},
|
||||
{"type": "image_generation_call", "id": "ig", "result": "IMAGE", "status": "completed"},
|
||||
{"type": "code_interpreter_call", "id": "ci", "code": "keep()", "outputs": [
|
||||
{"type": "logs", "logs": "keep"}, {"type": "image", "url": "https://example.com/image.png"},
|
||||
]},
|
||||
]
|
||||
original: Final = copy.deepcopy(request)
|
||||
projected, media = split_media_for_tpm(request)
|
||||
assert projected == [
|
||||
{"role": "user", "content": [empty_tool]},
|
||||
{"type": "function_call_output", "call_id": "call", "output": [empty_image]},
|
||||
{"type": "custom_tool_call_output", "call_id": "custom", "output": [empty_image]},
|
||||
{"type": "computer_call_output", "call_id": "computer", "output": {
|
||||
"type": "computer_screenshot", "image_url": "",
|
||||
}},
|
||||
{"type": "image_generation_call", "id": "ig", "result": "", "status": "completed"},
|
||||
{"type": "code_interpreter_call", "id": "ci", "code": "keep()", "outputs": [
|
||||
{"type": "logs", "logs": "keep"}, {"type": "image", "url": ""},
|
||||
]},
|
||||
]
|
||||
assert media == (
|
||||
image, image, image,
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE", "detail": "auto"}},
|
||||
{"type": "image_url", "image_url": {"url": "data:application/octet-stream;base64,IMAGE", "detail": "auto"}},
|
||||
{"type": "image_url", "image_url": {"url": "https://example.com/image.png", "detail": "auto"}},
|
||||
)
|
||||
assert request == original
|
||||
|
||||
|
||||
def test_tpm_projection_preserves_non_media_identity_and_opaque_values() -> None:
|
||||
image: Final = {"type": "input_image", "image_url": "data:image/png;base64,OPAQUE"}
|
||||
request: Final = [
|
||||
{"role": "user", "content": json.dumps(image)},
|
||||
{"type": "custom_tool_call_output", "call_id": "call", "output": json.dumps(image)},
|
||||
{"type": "function_call", "call_id": "call", "name": "keep", "arguments": {"content": [image]}},
|
||||
{"type": "file_search_call", "results": [{"text": "keep", "content": [image]}]},
|
||||
{"type": "image_generation_call", "id": "ig", "status": "in_progress", "result": None},
|
||||
{"role": "user", "content": [
|
||||
{"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "PDF"}},
|
||||
{"type": "document", "source": {"type": "text", "data": json.dumps(image)}},
|
||||
{"type": "file", "file": {"format": "application/pdf", "file_data": "PDF"}},
|
||||
{"type": "text", "text": "keep", "metadata": {"content": [image]}},
|
||||
]},
|
||||
]
|
||||
projected, media = split_media_for_tpm(request)
|
||||
assert projected is request
|
||||
assert media == ()
|
||||
|
|
|
|||
|
|
@ -2,16 +2,23 @@
|
|||
Test the audio utils functionality in litellm_core_utils/audio_utils/utils.py
|
||||
"""
|
||||
|
||||
import base64
|
||||
import io
|
||||
import os
|
||||
import tempfile
|
||||
from typing import Final
|
||||
from unittest.mock import mock_open, patch
|
||||
|
||||
import pytest
|
||||
import soundfile as sf
|
||||
|
||||
from litellm.litellm_core_utils.audio_utils.utils import (
|
||||
AUDIO_BYTES_PER_TOKEN,
|
||||
AUDIO_TOKENS_PER_SECOND,
|
||||
DEFAULT_AUDIO_TOKEN_ESTIMATE,
|
||||
ProcessedAudioFile,
|
||||
calculate_request_duration,
|
||||
estimate_audio_tokens,
|
||||
get_audio_file_content_hash,
|
||||
get_audio_file_for_health_check,
|
||||
get_audio_file_name,
|
||||
|
|
@ -19,6 +26,63 @@ from litellm.litellm_core_utils.audio_utils.utils import (
|
|||
)
|
||||
|
||||
|
||||
def _encoded_audio(seconds: int, sample_rate: int, channels: int, container: str) -> str:
|
||||
buffer: Final = io.BytesIO()
|
||||
with sf.SoundFile(buffer, "w", samplerate=sample_rate, channels=channels, format=container) as audio:
|
||||
audio.buffer_write(b"\0" * (seconds * sample_rate * channels * 2), dtype="int16")
|
||||
return base64.b64encode(buffer.getvalue()).decode()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("sample_rate", "channels", "container"), [(8000, 1, "WAV"), (48000, 2, "WAV"), (24000, 1, "FLAC")]
|
||||
)
|
||||
def test_audio_estimate_uses_duration_independent_of_sample_rate_channels_and_compression(
|
||||
sample_rate: int, channels: int, container: str
|
||||
) -> None:
|
||||
assert estimate_audio_tokens(_encoded_audio(1, sample_rate, channels, container)) == AUDIO_TOKENS_PER_SECOND
|
||||
assert estimate_audio_tokens(_encoded_audio(31, sample_rate, channels, container)) == 31 * AUDIO_TOKENS_PER_SECOND
|
||||
|
||||
|
||||
# Audio rates verified 2026-10-28: https://ai.google.dev/gemini-api/docs/generate-content/media-resolution and
|
||||
# https://developers.openai.com/api/docs/guides/voice-latency-cost
|
||||
@pytest.mark.parametrize(
|
||||
("models", "fallback", "expected"),
|
||||
[
|
||||
(("openai/gpt-audio-1.5",), "gemini-3.8-flash", 10),
|
||||
(("gpt-audio-1.5", "gemini/gemini-3.8-flash"), None, 25),
|
||||
(("vertex_ai/gemini-3.8-flash",), None, 25),
|
||||
(("openrouter/google/gemini-3.8-flash",), None, 25),
|
||||
(("azure/audio-deployment",), "gpt-audio-1.5", 10),
|
||||
(("opaque-deployment",), "gemini-3.8-flash", 25),
|
||||
((), "gemini-3.8-flash", 25),
|
||||
(("gemini/gemini-2.5-flash", "gemini/gemini-3-pro-preview"), None, 32),
|
||||
(("other/gemini-3.8-flash",), None, 10),
|
||||
],
|
||||
)
|
||||
def test_audio_estimate_prefers_each_backend_rate_over_the_requested_model(
|
||||
models: tuple[str, ...], fallback: str | None, expected: int
|
||||
) -> None:
|
||||
assert estimate_audio_tokens(_encoded_audio(1, 8000, 1, "WAV"), models, fallback) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("data", [None, "", "invalid-base64!", base64.b64encode(b"unreadable audio").decode()])
|
||||
def test_audio_estimate_retains_nonzero_fallback_for_unreadable_audio(data: str | None) -> None:
|
||||
assert estimate_audio_tokens(data) == DEFAULT_AUDIO_TOKEN_ESTIMATE
|
||||
|
||||
|
||||
def test_audio_estimate_retains_size_based_fallback_for_large_unreadable_audio() -> None:
|
||||
data: Final = base64.b64encode(b"\0" * (AUDIO_BYTES_PER_TOKEN * 400)).decode()
|
||||
assert estimate_audio_tokens(data) == 400
|
||||
|
||||
|
||||
def test_audio_estimate_bounds_base64_allocation_before_decoding() -> None:
|
||||
data: Final = base64.b64encode(b"\0" * (16 * 1024 * 1024 + 1)).decode()
|
||||
with patch("base64.b64decode", side_effect=AssertionError("oversized audio must not be decoded")) as decode:
|
||||
estimate: Final = estimate_audio_tokens(data)
|
||||
assert estimate == len(data) * 3 // 4 // AUDIO_BYTES_PER_TOKEN
|
||||
decode.assert_not_called()
|
||||
|
||||
|
||||
class TestProcessAudioFile:
|
||||
"""Test the process_audio_file function with various input types"""
|
||||
|
||||
|
|
|
|||
|
|
@ -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.",
|
||||
|
|
|
|||
|
|
@ -723,6 +723,37 @@ async def test_ahealth_check_probes_hosted_vllm_with_a_choice_question(
|
|||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("provider", ("ollama", "ollama_chat"))
|
||||
async def test_ahealth_check_probes_ollama_decisions_with_one_noul_question_without_auth(
|
||||
provider: str, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
|
||||
) -> None:
|
||||
monkeypatch.delenv("OLLAMA_API_KEY", raising=False)
|
||||
monkeypatch.delenv("OLLAMA_API_BASE", raising=False)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
upstream: Final = respx_mock.post("http://ollama.local:11434/v1/systemone").respond(
|
||||
json={
|
||||
"model": "d3:4b",
|
||||
"answers": {"reachable": {"type": "noul", "noul": 1.0}},
|
||||
"usage": {"input_tokens": 12, "output_tokens": 1},
|
||||
}
|
||||
)
|
||||
|
||||
result: Final = await ahealth_check(
|
||||
{"model": f"{provider}/d3:4b", "api_base": "http://ollama.local:11434"}, mode="decisions"
|
||||
)
|
||||
|
||||
assert "error" not in result, result
|
||||
assert upstream.call_count == 1
|
||||
assert "authorization" not in upstream.calls.last.request.headers
|
||||
assert json.loads(upstream.calls.last.request.content) == {
|
||||
"model": "d3:4b",
|
||||
"state": "health check",
|
||||
"questions": {"reachable": {"type": "noul", "instructions": "Is the service reachable?"}},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("mode", ("evaluation", "decisions"))
|
||||
async def test_ahealth_check_keeps_the_noul_probe_for_other_decisions_providers(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -37,6 +37,7 @@ from litellm.litellm_core_utils.token_counter import (
|
|||
_get_extrapolating_count_function,
|
||||
_get_tiktoken_count_function,
|
||||
calculate_img_tokens,
|
||||
estimate_image_tokens,
|
||||
high_detail_image_token_upper_bound,
|
||||
offload_token_count,
|
||||
)
|
||||
|
|
@ -2152,6 +2153,136 @@ def test_high_detail_image_token_upper_bound_is_reached_by_the_largest_high_res_
|
|||
assert calculate_img_tokens(_png_data_url(1, 1), mode="high") < high_detail_image_token_upper_bound()
|
||||
|
||||
|
||||
# Image sizing verified 2026-10-28: https://developers.openai.com/api/docs/guides/images-vision
|
||||
# and https://platform.claude.com/docs/en/build-with-claude/vision-coordinates
|
||||
@pytest.mark.parametrize(
|
||||
("model", "width", "height", "detail", "expected"),
|
||||
[
|
||||
("gpt-4o-mini", 1024, 1024, "low", 2833),
|
||||
("gpt-4o-mini", 1024, 1024, "high", 25501),
|
||||
("gpt-5.1", 512, 512, "auto", 210),
|
||||
("gpt-6-astra", 2048, 2048, "high", 3000),
|
||||
("gpt-6-astra", 4096, 512, "high", 2458),
|
||||
("gpt-6-astra", 2048, 2048, "auto", 4916),
|
||||
("gpt-5.6-sol", 2048, 2048, "low", 308),
|
||||
("gpt-5.5", 6400, 6400, "original", 12000),
|
||||
("gpt-5.4-mini", 2048, 2048, "low", 4916),
|
||||
("gpt-5.4-nano", 2048, 2048, "auto", 3000),
|
||||
("gpt-5.2", 4096, 512, "auto", 615),
|
||||
("gpt-4.1-mini-2025-04-14", 1024, 1024, "low", 1659),
|
||||
("claude-haiku-5-5", 200, 200, "auto", 64),
|
||||
("anthropic/claude-haiku-5-5", 1024, 1024, "auto", 1369),
|
||||
("claude-haiku-5-5", 1920, 1080, "low", 2691),
|
||||
("claude-haiku-4-5", 1920, 1080, "auto", 1560),
|
||||
("claude-haiku-4-5", 1075, 1520, "auto", 1551),
|
||||
("claude-haiku-5-5", 3840, 2160, "auto", 4784),
|
||||
("claude-haiku-5-5", 8000, 100, "auto", 184),
|
||||
],
|
||||
)
|
||||
def test_image_estimate_uses_model_resolution_and_detail(
|
||||
model: str, width: int, height: int, detail: str, expected: int
|
||||
) -> None:
|
||||
assert estimate_image_tokens(model, _png_data_url(width, height), detail) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("models", "fallback", "expected_model"),
|
||||
[
|
||||
(("azure/deployment-prod",), "gpt-4o", "gpt-4o"),
|
||||
(("openai/gpt-4o",), "gpt-4o-mini", "gpt-4o"),
|
||||
(("azure/deployment-prod", "openai/gpt-4o"), "gpt-4o-mini", "gpt-4o-mini"),
|
||||
(("openai/gpt-4o", "openrouter/openai/gpt-4o-mini"), None, "gpt-4o-mini"),
|
||||
(("azure_ai/openai/gpt-6-astra",), None, "gpt-6-astra"),
|
||||
(("azure/gpt-41",), None, "gpt-4.1"),
|
||||
(("azure/eu/gpt-4o-mini",), None, "gpt-4o-mini"),
|
||||
(("unknown", "gpt-4o-mini"), None, "gpt-4o-mini"),
|
||||
((), "gpt-6-astra", "gpt-6-astra"),
|
||||
(("other/gpt-4o-mini",), None, "unregistered-model"),
|
||||
(("azure/opaque/gpt-4o-mini",), None, "unregistered-model"),
|
||||
(("bedrock/anthropic.claude-haiku-5-5",), None, "claude-haiku-5-5"),
|
||||
(("bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0",), None, "claude-haiku-4-5"),
|
||||
(("vertex_ai/claude-haiku-5-5@default",), None, "claude-haiku-5-5"),
|
||||
(("azure_ai/claude-haiku-5-5",), None, "claude-haiku-5-5"),
|
||||
(("anthropic/claude-haiku-5-5", "openai/gpt-4o"), "gpt-4o-mini", "claude-haiku-5-5"),
|
||||
(("other/claude-haiku-5-5",), None, "unregistered-model"),
|
||||
],
|
||||
)
|
||||
def test_image_estimate_prefers_each_backend_policy_over_the_requested_model(
|
||||
models: tuple[str, ...], fallback: str | None, expected_model: str
|
||||
) -> None:
|
||||
data: Final = _png_data_url(1024, 1024)
|
||||
assert estimate_image_tokens(models, data, "high", fallback_model=fallback) == estimate_image_tokens(
|
||||
expected_model, data, "high"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "detail"),
|
||||
[
|
||||
("unknown", "high"),
|
||||
("claude-unregistered", "auto"),
|
||||
("gpt-4o", "original"),
|
||||
("gpt-5.2", "original"),
|
||||
("gpt-6-astra", "invalid"),
|
||||
],
|
||||
)
|
||||
def test_image_estimate_preserves_default_when_policy_is_unavailable(model: str, detail: str) -> None:
|
||||
assert (
|
||||
estimate_image_tokens(model, _png_data_url(1024, 1024), detail) == litellm.constants.DEFAULT_IMAGE_TOKEN_COUNT
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"data",
|
||||
[
|
||||
"https://image.invalid/photo.png",
|
||||
"data:image/png;base64,invalid!",
|
||||
"data:application/octet-stream;base64,bm90LWFuLWltYWdl",
|
||||
"data:image/jpeg;base64,/9j/4AAQ",
|
||||
"data:image/webp;base64,UklGRgAAAABXRUJQ",
|
||||
_png_data_url(0, 1024),
|
||||
],
|
||||
)
|
||||
# Accepted patch limits verified 2026-10-28: https://developers.openai.com/api/docs/guides/images-vision
|
||||
# and https://platform.claude.com/docs/en/build-with-claude/vision
|
||||
@pytest.mark.parametrize(
|
||||
("model", "detail", "bound"),
|
||||
[
|
||||
("gpt-4o", "high", high_detail_image_token_upper_bound()),
|
||||
("gpt-6-astra", "high", 3000),
|
||||
("gpt-6-astra", "auto", 36000),
|
||||
("claude-haiku-5-5", "auto", 4784),
|
||||
("claude-haiku-4-5", "auto", 1568),
|
||||
],
|
||||
)
|
||||
def test_image_estimate_uses_bounded_fallback_without_fetching_media(
|
||||
data: str, model: str, detail: str, bound: int
|
||||
) -> None:
|
||||
with patch("httpx.Client.send") as send:
|
||||
estimate: Final = estimate_image_tokens(model, data, detail)
|
||||
send.assert_not_called()
|
||||
assert estimate == bound
|
||||
assert estimate > estimate_image_tokens(model, _png_data_url(1024, 1024), detail)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("container", ["PNG", "JPEG", "GIF", "WEBP"])
|
||||
def test_image_estimate_depends_on_dimensions_and_not_compressed_payload_size(container: str) -> None:
|
||||
output: Final = io.BytesIO()
|
||||
Image.new("RGB", (1024, 1024)).save(output, format=container)
|
||||
data: Final = "data:application/octet-stream;base64," + base64.b64encode(output.getvalue()).decode()
|
||||
assert estimate_image_tokens("gpt-6-astra", data) == estimate_image_tokens("gpt-6-astra", _png_data_url(1024, 1024))
|
||||
|
||||
|
||||
def test_image_estimate_bounds_header_decoding_for_large_inline_payloads() -> None:
|
||||
header: Final = base64.b64decode(_png_data_url(1024, 1024).split(",", 1)[1])
|
||||
data: Final = "data:image/png;base64," + base64.b64encode(header + b"\0" * 4_000_000).decode()
|
||||
with patch("base64.b64decode", wraps=base64.b64decode) as decode_header:
|
||||
estimate: Final = estimate_image_tokens("gpt-6-astra", data)
|
||||
assert estimate == estimate_image_tokens("gpt-6-astra", _png_data_url(1024, 1024))
|
||||
assert decode_header.call_count == 1
|
||||
assert len(decode_header.call_args.args[0]) <= 65_536
|
||||
|
||||
|
||||
HUB_SETUP_SCRIPT: Final = """
|
||||
import json
|
||||
import sys
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ import os
|
|||
import time
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock, patch
|
||||
from unittest.mock import ANY, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
|
@ -291,11 +291,9 @@ def test_get_azure_ad_token_from_username_password(mock_get_bearer_token_provide
|
|||
|
||||
mock_credential.assert_called_once_with(client_id=client_id, username=username, password=password)
|
||||
|
||||
mock_get_bearer_token_provider.assert_called_once_with(
|
||||
mock_credential.return_value, "https://cognitiveservices.azure.com/.default"
|
||||
)
|
||||
mock_get_bearer_token_provider.assert_called_once_with(ANY, "https://cognitiveservices.azure.com/.default")
|
||||
|
||||
assert result == mock_token_provider
|
||||
assert result() == "mock-token"
|
||||
|
||||
|
||||
def test_azure_openai_gpt_4o_naming(monkeypatch):
|
||||
|
|
|
|||
|
|
@ -2,9 +2,12 @@ import json
|
|||
import os
|
||||
import threading
|
||||
import traceback
|
||||
import uuid
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from queue import SimpleQueue
|
||||
from typing import Callable, Final, Optional
|
||||
from urllib.parse import parse_qs
|
||||
from unittest.mock import MagicMock, patch
|
||||
from unittest.mock import ANY, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
|
@ -13,6 +16,7 @@ import litellm
|
|||
from litellm.llms.azure.common_utils import (
|
||||
BaseAzureLLM,
|
||||
_cached_azure_ad_token_refresh_provider,
|
||||
_cached_default_azure_credential_provider,
|
||||
_cached_entra_id_token_provider,
|
||||
get_azure_ad_token,
|
||||
get_azure_ad_token_from_entra_id,
|
||||
|
|
@ -41,6 +45,7 @@ def setup_mocks(monkeypatch):
|
|||
monkeypatch.delenv("AZURE_SCOPE", raising=False)
|
||||
monkeypatch.delenv("AZURE_AD_TOKEN", raising=False)
|
||||
_cached_azure_ad_token_refresh_provider.cache_clear()
|
||||
_cached_default_azure_credential_provider.cache_clear()
|
||||
|
||||
with (
|
||||
patch(
|
||||
|
|
@ -86,6 +91,7 @@ def setup_mocks(monkeypatch):
|
|||
"select_url": mock_select_url,
|
||||
}
|
||||
_cached_azure_ad_token_refresh_provider.cache_clear()
|
||||
_cached_default_azure_credential_provider.cache_clear()
|
||||
|
||||
|
||||
def test_initialize_with_api_key(setup_mocks):
|
||||
|
|
@ -1436,6 +1442,33 @@ def test_get_azure_ad_token_with_token_refresh_reuses_provider_per_scope(setup_m
|
|||
]
|
||||
|
||||
|
||||
def test_get_azure_ad_token_reuses_one_default_azure_credential_provider_per_scope(setup_mocks):
|
||||
setup_mocks["litellm"].enable_azure_ad_token_refresh = True
|
||||
default_scope: Final = "https://cognitiveservices.azure.com/.default"
|
||||
other_scope: Final = "api://other-resource/.default"
|
||||
|
||||
def provider_for(azure_scope: str, azure_credential: AzureCredentialType | None = None) -> Callable[[], str]:
|
||||
if azure_credential != AzureCredentialType.DefaultAzureCredential:
|
||||
raise ValueError("No credential provided")
|
||||
return lambda: f"token-for-{azure_scope}"
|
||||
|
||||
setup_mocks["token_provider"].side_effect = provider_for
|
||||
|
||||
tokens: Final = (
|
||||
get_azure_ad_token(GenericLiteLLMParams()),
|
||||
get_azure_ad_token(GenericLiteLLMParams()),
|
||||
get_azure_ad_token(GenericLiteLLMParams(azure_scope=other_scope)),
|
||||
)
|
||||
|
||||
default_credential_builds: Final = tuple(
|
||||
call.kwargs["azure_scope"]
|
||||
for call in setup_mocks["token_provider"].call_args_list
|
||||
if call.kwargs.get("azure_credential") == AzureCredentialType.DefaultAzureCredential
|
||||
)
|
||||
assert tokens == (f"token-for-{default_scope}", f"token-for-{default_scope}", f"token-for-{other_scope}")
|
||||
assert default_credential_builds == (default_scope, other_scope)
|
||||
|
||||
|
||||
def test_get_azure_ad_token_with_token_refresh_error(setup_mocks):
|
||||
"""Test get_azure_ad_token with token refresh enabled but raising an error."""
|
||||
# Reset mocks to ensure clean state
|
||||
|
|
@ -1536,9 +1569,7 @@ def test_get_azure_ad_token_provider_with_default_azure_credential():
|
|||
mock_default_cred.assert_called_once_with()
|
||||
|
||||
# Verify get_bearer_token_provider was called with the right parameters
|
||||
mock_token_provider.assert_called_once_with(
|
||||
mock_credential_instance, "https://cognitiveservices.azure.com/.default"
|
||||
)
|
||||
mock_token_provider.assert_called_once_with(ANY, "https://cognitiveservices.azure.com/.default")
|
||||
|
||||
# Verify the returned token provider works
|
||||
token = token_provider()
|
||||
|
|
@ -2137,6 +2168,42 @@ class TestEntraIdTokenProviderCache:
|
|||
assert first is second
|
||||
assert mock_credential.call_count == 1
|
||||
|
||||
def test_concurrent_first_callers_build_one_credential(self):
|
||||
callers: Final = 8
|
||||
start: Final = threading.Barrier(callers)
|
||||
built: Final = SimpleQueue[None]()
|
||||
everyone_built: Final = threading.Event()
|
||||
release: Final = threading.Event()
|
||||
|
||||
def held_credential(*_: str) -> MagicMock:
|
||||
built.put(None)
|
||||
if built.qsize() >= callers:
|
||||
everyone_built.set()
|
||||
release.wait(5)
|
||||
return MagicMock()
|
||||
|
||||
def first_call() -> Callable[[], str]:
|
||||
start.wait(5)
|
||||
return get_azure_ad_token_from_entra_id(
|
||||
tenant_id="tenant",
|
||||
client_id="client",
|
||||
client_secret="secret",
|
||||
scope="https://cognitiveservices.azure.com/.default",
|
||||
)
|
||||
|
||||
with (
|
||||
patch("azure.identity.ClientSecretCredential", side_effect=held_credential),
|
||||
patch("azure.identity.get_bearer_token_provider", side_effect=lambda credential, scope: lambda: "token"),
|
||||
ThreadPoolExecutor(max_workers=callers) as pool,
|
||||
):
|
||||
calls: Final = tuple(pool.submit(first_call) for _ in range(callers))
|
||||
everyone_built.wait(0.5)
|
||||
release.set()
|
||||
providers: Final = tuple(call.result(timeout=5) for call in calls)
|
||||
|
||||
assert built.qsize() == 1
|
||||
assert all(provider is providers[0] for provider in providers)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"second_call_kwargs",
|
||||
[
|
||||
|
|
@ -2291,6 +2358,71 @@ def test_oidc_token_exchange_closes_its_own_http_client(monkeypatch: pytest.Monk
|
|||
assert token_endpoint.closed
|
||||
|
||||
|
||||
def test_concurrent_oidc_callers_share_one_token_exchange(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
callers: Final = 8
|
||||
start: Final = threading.Barrier(callers)
|
||||
exchanges: Final = SimpleQueue[None]()
|
||||
everyone_exchanging: Final = threading.Event()
|
||||
release: Final = threading.Event()
|
||||
|
||||
class HeldTokenEndpoint(httpx.BaseTransport):
|
||||
def handle_request(self, request: httpx.Request) -> httpx.Response:
|
||||
exchanges.put(None)
|
||||
if exchanges.qsize() >= callers:
|
||||
everyone_exchanging.set()
|
||||
release.wait(5)
|
||||
return httpx.Response(200, json={"access_token": "exchanged-token", "expires_in": 3600})
|
||||
|
||||
monkeypatch.setattr(litellm, "sync_transport", HeldTokenEndpoint(), raising=False)
|
||||
monkeypatch.setenv("AZURE_TEST_OIDC_BURST_ASSERTION", "burst-jwt")
|
||||
tenant: Final = f"tenant-{uuid.uuid4().hex}"
|
||||
|
||||
def exchange() -> str:
|
||||
start.wait(5)
|
||||
return get_azure_ad_token_from_oidc(
|
||||
azure_ad_token="os.environ/AZURE_TEST_OIDC_BURST_ASSERTION",
|
||||
azure_client_id="client",
|
||||
azure_tenant_id=tenant,
|
||||
)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=callers) as pool:
|
||||
calls: Final = tuple(pool.submit(exchange) for _ in range(callers))
|
||||
exchanges.get(timeout=5)
|
||||
everyone_exchanging.wait(0.5)
|
||||
release.set()
|
||||
tokens: Final = tuple(call.result(timeout=5) for call in calls)
|
||||
|
||||
assert tokens == ("exchanged-token",) * callers
|
||||
assert exchanges.qsize() == 0
|
||||
|
||||
|
||||
def test_oidc_tokens_for_different_scopes_are_exchanged_and_cached_separately(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
exchanged_scopes: Final = SimpleQueue[str]()
|
||||
|
||||
class ScopedTokenEndpoint(httpx.BaseTransport):
|
||||
def handle_request(self, request: httpx.Request) -> httpx.Response:
|
||||
scope: Final = parse_qs(request.read().decode())["scope"][-1]
|
||||
exchanged_scopes.put(scope)
|
||||
return httpx.Response(200, json={"access_token": f"token-for:{scope}", "expires_in": 3600})
|
||||
|
||||
monkeypatch.setattr(litellm, "sync_transport", ScopedTokenEndpoint(), raising=False)
|
||||
monkeypatch.setenv("AZURE_TEST_OIDC_SCOPED_ASSERTION", "scoped-jwt")
|
||||
tenant: Final = f"tenant-{uuid.uuid4().hex}"
|
||||
|
||||
def token_for(scope: str) -> str:
|
||||
return get_azure_ad_token_from_oidc(
|
||||
azure_ad_token="os.environ/AZURE_TEST_OIDC_SCOPED_ASSERTION",
|
||||
azure_client_id="client",
|
||||
azure_tenant_id=tenant,
|
||||
scope=scope,
|
||||
)
|
||||
|
||||
tokens: Final = tuple(token_for(scope) for scope in ("api://first/.default", "api://second/.default") * 2)
|
||||
|
||||
assert tokens == ("token-for:api://first/.default", "token-for:api://second/.default") * 2
|
||||
assert exchanged_scopes.qsize() == 2
|
||||
|
||||
|
||||
def test_evicting_an_azure_client_built_on_the_callers_session_leaves_it_open(monkeypatch):
|
||||
"""`initialize_azure_sdk_client` puts `litellm.aclient_session` on the SDK client.
|
||||
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue