mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-10 08:28:25 +00:00
Merge remote-tracking branch 'origin/main' into litellm_openai_ws_frame_model_gate
# Conflicts: # litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
This commit is contained in:
commit
6f2bfb5d87
154 changed files with 9036 additions and 781 deletions
|
|
@ -1,6 +1,8 @@
|
|||
---
|
||||
name: database-access
|
||||
description: Add database queries, extract existing Prisma calls into repositories, or explicitly port queries to Rust and SQLx without expanding migration scope
|
||||
metadata:
|
||||
internal: true
|
||||
---
|
||||
|
||||
# Database access
|
||||
|
|
|
|||
4
.github/workflows/spend-logs-guard.yml
vendored
4
.github/workflows/spend-logs-guard.yml
vendored
|
|
@ -4,10 +4,6 @@ on: # zizmor: ignore[dangerous-triggers] runs the base branch's code only; the P
|
|||
pull_request_target:
|
||||
branches:
|
||||
- main
|
||||
pull_request_review:
|
||||
types:
|
||||
- submitted
|
||||
- dismissed
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
|
|
|||
|
|
@ -2,9 +2,12 @@
|
|||
needs an approving review from the spend-log owner before it can merge.
|
||||
|
||||
CODEOWNERS can only match file paths, and every schema change edits the same schema.prisma copies,
|
||||
so a path rule would put every migration on the owner. This gate reads the diff instead: when no
|
||||
added line in the watched files mentions the table it passes at once, otherwise it passes only once
|
||||
the owner's latest review on the pull request is an approval. The workflow re-runs on every review.
|
||||
so a path rule would put every migration on the owner. This gate reads the diff instead: a schema.prisma
|
||||
hunk counts when it adds or removes a line inside the `model LiteLLM_SpendLogs { ... }` block, a migration
|
||||
hunk counts when an added line names the table. Without such a hunk it passes at once, otherwise it passes
|
||||
only once the owner's latest review on the pull request is an approval that still covers every spend-log line
|
||||
in the head (approvals survive pushes, so a later commit may not add spend-log lines the owner never saw).
|
||||
Reviews do not trigger the workflow, so the check is re-run by hand after approving.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -13,8 +16,10 @@ import argparse
|
|||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
from collections.abc import Sequence
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter
|
||||
|
|
@ -25,7 +30,10 @@ WATCHED: Final = (
|
|||
"litellm-proxy-extras/litellm_proxy_extras/schema.prisma",
|
||||
)
|
||||
MIGRATIONS_PREFIX: Final = "litellm-proxy-extras/litellm_proxy_extras/migrations/"
|
||||
TABLE: Final = re.compile(r"(?<![A-Za-z0-9])LiteLLM_SpendLogs", re.IGNORECASE)
|
||||
TABLE: Final = re.compile(r"(?<![A-Za-z0-9])LiteLLM_SpendLogs(?![A-Za-z0-9])", re.IGNORECASE)
|
||||
DECIDING_STATES: Final = frozenset({"APPROVED", "CHANGES_REQUESTED", "DISMISSED"})
|
||||
MODEL_START: Final = re.compile(r"^model\s+LiteLLM_SpendLogs\s*\{")
|
||||
HUNK: Final = re.compile(r"^@@ -(\d+)(?:,\d+)? \+(\d+)(?:,\d+)? @@[^\n]*\n", re.MULTILINE)
|
||||
|
||||
|
||||
class Reviewer(BaseModel):
|
||||
|
|
@ -37,31 +45,77 @@ class Review(BaseModel):
|
|||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
state: str
|
||||
user: Reviewer
|
||||
commit_id: str | None = None
|
||||
|
||||
|
||||
_REVIEWS: Final = TypeAdapter(list[Review])
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Block:
|
||||
"""Lines of `model LiteLLM_SpendLogs { ... }` in one schema.prisma copy at the merge base and at the head."""
|
||||
|
||||
base: range
|
||||
head: range
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Hunk:
|
||||
removed: tuple[tuple[int, str], ...]
|
||||
added: tuple[tuple[int, str], ...]
|
||||
|
||||
|
||||
def watched(path: str) -> bool:
|
||||
return path in WATCHED or (path.startswith(MIGRATIONS_PREFIX) and path.endswith("/migration.sql"))
|
||||
|
||||
|
||||
def _added_lines(section: str) -> tuple[str, ...]:
|
||||
body: Final = section.split("\n@@", 1)
|
||||
return tuple(line[1:] for line in body[1].splitlines() if line.startswith("+")) if len(body) == 2 else ()
|
||||
def model_block(schema: str) -> range:
|
||||
"""1-based line range of `model LiteLLM_SpendLogs { ... }` including both braces, empty when the schema has none."""
|
||||
lines: Final = schema.splitlines()
|
||||
starts: Final = tuple(number for number, line in enumerate(lines, 1) if MODEL_START.match(line))
|
||||
if not starts:
|
||||
return range(0)
|
||||
ends: Final = tuple(number for number, line in enumerate(lines, 1) if number > starts[0] and line.startswith("}"))
|
||||
return range(starts[0], (ends[0] if ends else len(lines)) + 1)
|
||||
|
||||
|
||||
def spend_log_hits(diff: str) -> tuple[str, ...]:
|
||||
"""`path: added line` for every added line in a watched file that names the table."""
|
||||
def _hunks(section: str) -> tuple[Hunk, ...]:
|
||||
headers: Final = tuple(HUNK.finditer(section))
|
||||
ends: Final = tuple(header.start() for header in headers[1:]) + (len(section),)
|
||||
return tuple(_hunk(header, section[header.end() : end]) for header, end in zip(headers, ends))
|
||||
|
||||
|
||||
def _hunk(header: re.Match[str], body: str) -> Hunk:
|
||||
lines: Final = body.splitlines()
|
||||
return Hunk(
|
||||
removed=tuple(enumerate((line[1:] for line in lines if line.startswith("-")), int(header.group(1)))),
|
||||
added=tuple(enumerate((line[1:] for line in lines if line.startswith("+")), int(header.group(2)))),
|
||||
)
|
||||
|
||||
|
||||
def spend_log_hits(diff: str, blocks: Mapping[str, Block]) -> tuple[str, ...]:
|
||||
"""`path: +line` or `path: -line` for every changed line that touches the table.
|
||||
|
||||
A schema.prisma line counts when its number falls inside that copy's `blocks` entry (head for added lines,
|
||||
merge base for removed ones), a migration line when it is added and names the table."""
|
||||
sections: Final = re.split(r"^diff --git [^\n]*\n", diff, flags=re.MULTILINE)
|
||||
return tuple(f"{path}: {line.strip()}" for path, line in _watched_added_lines(sections) if TABLE.search(line))
|
||||
|
||||
|
||||
def _watched_added_lines(sections: Sequence[str]) -> tuple[tuple[str, str], ...]:
|
||||
pairs: Final = tuple((_path_of(section), section) for section in sections)
|
||||
paths: Final = tuple((_path_of(section), section) for section in sections)
|
||||
return tuple(
|
||||
(path, line) for path, section in pairs if watched(path) for line in _added_lines(section)
|
||||
) # comprehension-ok: path then its lines
|
||||
hit for path, section in paths if watched(path) for hit in _section_hits(path, _hunks(section), blocks)
|
||||
) # comprehension-ok: path then its hits
|
||||
|
||||
|
||||
def _section_hits(path: str, hunks: Sequence[Hunk], blocks: Mapping[str, Block]) -> tuple[str, ...]:
|
||||
if path not in blocks:
|
||||
return tuple(f"{path}: +{line.strip()}" for hunk in hunks for _, line in hunk.added if TABLE.search(line))
|
||||
block: Final = blocks[path]
|
||||
removed: Final = tuple(
|
||||
f"{path}: -{line.strip()}" for hunk in hunks for number, line in hunk.removed if number in block.base
|
||||
) # comprehension-ok: hunk then its lines
|
||||
added: Final = tuple(
|
||||
f"{path}: +{line.strip()}" for hunk in hunks for number, line in hunk.added if number in block.head
|
||||
) # comprehension-ok: hunk then its lines
|
||||
return removed + added
|
||||
|
||||
|
||||
def _path_of(section: str) -> str:
|
||||
|
|
@ -70,11 +124,33 @@ def _path_of(section: str) -> str:
|
|||
|
||||
|
||||
def approved_by(reviews: Sequence[Review], login: str) -> bool:
|
||||
"""True when `login`'s most recent review on the pull request is an approval."""
|
||||
states: Final = tuple(review.state for review in reviews if review.user.login == login)
|
||||
"""True when `login`'s most recent review that changes GitHub's approval state is an approval.
|
||||
|
||||
COMMENTED and PENDING reviews leave an earlier approval in force on GitHub, so they are skipped here too.
|
||||
"""
|
||||
own: Final = tuple(review.state for review in reviews if review.user.login == login)
|
||||
states: Final = tuple(state for state in own if state in DECIDING_STATES)
|
||||
return bool(states) and states[-1] == "APPROVED"
|
||||
|
||||
|
||||
def latest_approval(reviews: Sequence[Review], login: str) -> Review | None:
|
||||
"""`login`'s most recent APPROVED review on the pull request, the one GitHub keeps in force after pushes."""
|
||||
own: Final = tuple(review for review in reviews if review.user.login == login and review.state == "APPROVED")
|
||||
return own[-1] if own else None
|
||||
|
||||
|
||||
def approval_covers(
|
||||
approval: Review, head: str, hits: Sequence[str], hits_at: Callable[[str], Sequence[str] | None]
|
||||
) -> bool:
|
||||
"""True when the approval is on the head, or the head adds no spend-log line the approved commit did not."""
|
||||
if approval.commit_id is None:
|
||||
return False
|
||||
if approval.commit_id == head:
|
||||
return True
|
||||
reviewed: Final = hits_at(approval.commit_id)
|
||||
return reviewed is not None and set(hits) <= set(reviewed)
|
||||
|
||||
|
||||
def added_diff(base: str, head: str) -> str:
|
||||
result: Final = subprocess.run(
|
||||
("git", "diff", "--unified=0", "--no-color", f"{base}..{head}", "--", *WATCHED, MIGRATIONS_PREFIX),
|
||||
|
|
@ -85,6 +161,27 @@ def added_diff(base: str, head: str) -> str:
|
|||
return result.stdout
|
||||
|
||||
|
||||
def file_at(revision: str, path: str) -> str:
|
||||
"""Contents of `path` at `revision`, empty when the file does not exist there."""
|
||||
result: Final = subprocess.run(("git", "show", f"{revision}:{path}"), check=False, capture_output=True, text=True)
|
||||
return result.stdout if result.returncode == 0 else ""
|
||||
|
||||
|
||||
def schema_blocks(base: str, head: str) -> Mapping[str, Block]:
|
||||
return MappingProxyType(
|
||||
{path: Block(base=model_block(file_at(base, path)), head=model_block(file_at(head, path))) for path in WATCHED}
|
||||
)
|
||||
|
||||
|
||||
def fetched(commit: str) -> bool:
|
||||
result: Final = subprocess.run(("git", "fetch", "--no-tags", "--depth=1", "origin", commit), check=False)
|
||||
return result.returncode == 0
|
||||
|
||||
|
||||
def reviewed_hits(base: str, commit: str) -> tuple[str, ...] | None:
|
||||
return spend_log_hits(added_diff(base, commit), schema_blocks(base, commit)) if fetched(commit) else None
|
||||
|
||||
|
||||
class Arguments(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
base: str
|
||||
|
|
@ -100,18 +197,28 @@ def main(argv: Sequence[str]) -> int:
|
|||
parser.add_argument("--approver", required=True, help="GitHub login whose approval unlocks the merge")
|
||||
parser.add_argument("--reviews", required=True, help="JSON file with the pull request's reviews")
|
||||
args: Final = Arguments.model_validate(vars(parser.parse_args(argv)))
|
||||
hits: Final = spend_log_hits(added_diff(args.base, args.head))
|
||||
hits: Final = spend_log_hits(added_diff(args.base, args.head), schema_blocks(args.base, args.head))
|
||||
if not hits:
|
||||
sys.stdout.write("No added line in schema.prisma or a migration names LiteLLM_SpendLogs; nothing to review.\n")
|
||||
sys.stdout.write(
|
||||
"No change inside model LiteLLM_SpendLogs in schema.prisma and no added migration line names it; "
|
||||
"nothing to review.\n"
|
||||
)
|
||||
return 0
|
||||
if approved_by(_REVIEWS.validate_json(args.reviews.read_bytes()), args.approver):
|
||||
reviews: Final = _REVIEWS.validate_json(args.reviews.read_bytes())
|
||||
approval: Final = latest_approval(reviews, args.approver) if approved_by(reviews, args.approver) else None
|
||||
if approval is not None and approval_covers(
|
||||
approval, args.head, hits, lambda commit: reviewed_hits(args.base, commit)
|
||||
):
|
||||
sys.stdout.write(f"LiteLLM_SpendLogs change approved by {args.approver}: {', '.join(hits)}\n")
|
||||
return 0
|
||||
listed: Final = "".join(f" {hit}\n" for hit in hits)
|
||||
sys.stdout.write(
|
||||
reason: Final = (
|
||||
f"This pull request changes LiteLLM_SpendLogs and needs an approving review from @{args.approver}:\n"
|
||||
f"{listed}The check re-runs when a review is submitted.\n"
|
||||
if approval is None
|
||||
else f"LiteLLM_SpendLogs lines changed after @{args.approver} approved {approval.commit_id}, "
|
||||
f"a new approving review is needed:\n"
|
||||
)
|
||||
sys.stdout.write(f"{reason}{listed}After approving, re-run this job.\n")
|
||||
return 1
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ from litellm.integrations.custom_guardrail import (
|
|||
CustomGuardrail,
|
||||
log_guardrail_information,
|
||||
)
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import effective_skip_tool_message_for_guardrail
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.guardrails._content_utils import iter_tool_json_fields, map_tool_json_text, walk_user_text
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
|
@ -753,7 +754,12 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
|
|||
_masked_entity_count.set({})
|
||||
|
||||
# Covers multimodal list content + Responses-API input.
|
||||
inspected = walk_user_text(data, self.redact_text, self.redact_tool_payload) + self._redact_prompt(data)
|
||||
inspected = walk_user_text(
|
||||
data,
|
||||
self.redact_text,
|
||||
self.redact_tool_payload,
|
||||
skip_tool_message=effective_skip_tool_message_for_guardrail(self),
|
||||
) + self._redact_prompt(data)
|
||||
|
||||
if inspected == 0:
|
||||
# Image-only, empty-text, and unsupported payloads inspected
|
||||
|
|
|
|||
|
|
@ -343,11 +343,16 @@ def btree_max_entry_bytes(block_size: int) -> int:
|
|||
|
||||
|
||||
def _value_bytes(column: ColumnType, identifier: "sql.Identifier") -> "sql.Composed":
|
||||
"""The bytes a value takes in the index entry: its length, or its stored size when
|
||||
Postgres compressed it, since `index_form_tuple` keeps the compressed bytes."""
|
||||
from psycopg import sql
|
||||
|
||||
if column.type_name in _OCTET_LENGTH_TYPES:
|
||||
return sql.SQL("octet_length({})").format(identifier)
|
||||
return sql.SQL("octet_length({}::text)").format(identifier)
|
||||
length: Final = (
|
||||
sql.SQL("octet_length({})").format(identifier)
|
||||
if column.type_name in _OCTET_LENGTH_TYPES
|
||||
else sql.SQL("octet_length({}::text)").format(identifier)
|
||||
)
|
||||
return sql.SQL("least({}, pg_column_size({}))").format(length, identifier)
|
||||
|
||||
|
||||
def _aligned(offset: "sql.Composable", align: str) -> "sql.Composed":
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
---
|
||||
name: rust-error
|
||||
description: Define, wrap, convert, classify, report or test Rust error types in litellm-rust with thiserror. Use when adding an error type or variant, mapping one crate's error into another, turning an error into a gateway response or Python exception, deciding retry or fault from an error, or replacing call sites that build errors from strings
|
||||
metadata:
|
||||
internal: true
|
||||
---
|
||||
|
||||
# Rust errors
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
---
|
||||
name: rust-gateway
|
||||
description: Add or change an HTTP route, extractor, response envelope, stream or router mount in the litellm-rust gateway crates (crates/gateway and crates/gateway-*). Use when adding an endpoint, parsing a request body, choosing an error envelope for an API surface, relaying SSE, or writing a router-level test
|
||||
metadata:
|
||||
internal: true
|
||||
---
|
||||
|
||||
# Rust gateway
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
---
|
||||
name: rust-string-enums
|
||||
description: Define or refactor Rust string-valued enums and their Serde adapters in litellm-rust, using Strum and serde_with while preserving parsing, wire values, and schemas
|
||||
metadata:
|
||||
internal: true
|
||||
---
|
||||
|
||||
# Rust string enums
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
---
|
||||
name: rust-tracing
|
||||
description: Add or change Rust diagnostic tracing in litellm-rust, including route spans, subscriber layers, and Python logger delivery
|
||||
metadata:
|
||||
internal: true
|
||||
---
|
||||
|
||||
# Rust tracing
|
||||
|
|
|
|||
|
|
@ -274,7 +274,9 @@ class InMemoryCache(BaseCache):
|
|||
) -> list[float] | None:
|
||||
results: Final = []
|
||||
for increment in increment_list:
|
||||
result = await self.async_increment(increment["key"], increment["increment_value"], **kwargs)
|
||||
result = await self.async_increment(
|
||||
increment["key"], increment["increment_value"], ttl=increment["ttl"], **kwargs
|
||||
)
|
||||
results.append(result)
|
||||
return results
|
||||
|
||||
|
|
|
|||
|
|
@ -1,9 +1,9 @@
|
|||
import asyncio
|
||||
import contextvars
|
||||
import importlib
|
||||
from collections.abc import Coroutine, Mapping
|
||||
from collections.abc import Collection, Coroutine, Mapping
|
||||
from functools import partial
|
||||
from typing import TYPE_CHECKING, Final, Literal, cast, overload
|
||||
from typing import TYPE_CHECKING, Final, Literal, cast, get_type_hints, overload
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
|
|
@ -61,6 +61,7 @@ from litellm.types.utils import (
|
|||
ProviderRoutingHints,
|
||||
)
|
||||
from litellm.utils import (
|
||||
EXTRA_BODY_ROUTING_KEYS,
|
||||
ImageResponse,
|
||||
ProviderConfigManager,
|
||||
get_llm_provider,
|
||||
|
|
@ -466,7 +467,7 @@ def image_generation(
|
|||
litellm.LlmProviders.QWENCLOUD,
|
||||
litellm.LlmProviders.QWEN_AI_PLATFORM,
|
||||
litellm.LlmProviders.EDENAI,
|
||||
):
|
||||
) or (custom_llm_provider == litellm.LlmProviders.TENCENT and image_generation_config is not None):
|
||||
if image_generation_config is None:
|
||||
raise ValueError(f"image generation config is not supported for {custom_llm_provider}")
|
||||
|
||||
|
|
@ -807,6 +808,32 @@ def image_variation(
|
|||
return response
|
||||
|
||||
|
||||
_IMAGE_EDIT_OPENAI_PARAMS: Final = frozenset(get_type_hints(ImageEditOptionalRequestParams))
|
||||
_DROP_PARAM_KWARGS: Final = frozenset({"drop_params", "additional_drop_params"})
|
||||
_ADDITIONAL_DROP_PARAMS: Final = TypeAdapter[tuple[str, ...]](tuple[str, ...])
|
||||
|
||||
|
||||
def _listed_drop_params(additional_drop_params: object) -> tuple[str, ...]:
|
||||
try:
|
||||
return _ADDITIONAL_DROP_PARAMS.validate_python(additional_drop_params)
|
||||
except ValidationError:
|
||||
return ()
|
||||
|
||||
|
||||
def _tokenhub_passthrough_edit_params(
|
||||
non_default_params: Mapping[str, object],
|
||||
extra_body: Mapping[str, object] | None,
|
||||
mapped_params: Collection[str],
|
||||
additional_drop_params: Collection[str] | None,
|
||||
) -> Mapping[str, object]:
|
||||
dropped: Final = frozenset(additional_drop_params or ())
|
||||
unmapped_excluded: Final = _DROP_PARAM_KWARGS | _IMAGE_EDIT_OPENAI_PARAMS | frozenset(mapped_params) | dropped
|
||||
return {
|
||||
**{key: value for key, value in non_default_params.items() if key not in unmapped_excluded},
|
||||
**{key: value for key, value in (extra_body or {}).items() if key not in dropped | EXTRA_BODY_ROUTING_KEYS},
|
||||
}
|
||||
|
||||
|
||||
@client
|
||||
def image_edit(
|
||||
image: FileTypes | list[FileTypes] | None = None,
|
||||
|
|
@ -935,22 +962,24 @@ def image_edit(
|
|||
raise ValueError(f"image edit is not supported for {custom_llm_provider}")
|
||||
|
||||
local_vars.update(kwargs)
|
||||
provider_supported_params: Final[frozenset[str]] = frozenset(
|
||||
image_edit_provider_config.get_supported_openai_params(model)
|
||||
)
|
||||
# Get ImageEditOptionalRequestParams with only valid parameters
|
||||
image_edit_optional_params: Final[ImageEditOptionalRequestParams] = (
|
||||
_get_ImageEditRequestUtils().get_requested_image_edit_optional_param(
|
||||
local_vars,
|
||||
provider_supported_params=frozenset(
|
||||
image_edit_provider_config.get_supported_openai_params(model)
|
||||
).intersection(non_default_params),
|
||||
provider_supported_params=provider_supported_params.intersection(non_default_params),
|
||||
)
|
||||
)
|
||||
# Get optional parameters for the responses API
|
||||
additional_drop_params: Final = _listed_drop_params(kwargs.get("additional_drop_params"))
|
||||
image_edit_request_params: Final[dict] = _get_ImageEditRequestUtils().get_optional_params_image_edit(
|
||||
model=model,
|
||||
image_edit_provider_config=image_edit_provider_config,
|
||||
image_edit_optional_params=image_edit_optional_params,
|
||||
drop_params=kwargs.get("drop_params"),
|
||||
additional_drop_params=kwargs.get("additional_drop_params"),
|
||||
additional_drop_params=additional_drop_params,
|
||||
)
|
||||
|
||||
if image_edit_provider_config.use_multipart_form_data() and (
|
||||
|
|
@ -1039,7 +1068,19 @@ def image_edit(
|
|||
image=images,
|
||||
prompt=prompt,
|
||||
image_edit_provider_config=image_edit_provider_config,
|
||||
image_edit_optional_request_params=image_edit_request_params,
|
||||
image_edit_optional_request_params={
|
||||
**image_edit_request_params,
|
||||
**(
|
||||
_tokenhub_passthrough_edit_params(
|
||||
non_default_params,
|
||||
extra_body if isinstance(extra_body, dict) else None,
|
||||
provider_supported_params,
|
||||
additional_drop_params,
|
||||
)
|
||||
if custom_llm_provider == litellm.LlmProviders.TENCENT.value
|
||||
else {}
|
||||
),
|
||||
},
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
litellm_params=litellm_params,
|
||||
logging_obj=litellm_logging_obj,
|
||||
|
|
|
|||
|
|
@ -1,6 +1,9 @@
|
|||
import base64
|
||||
import os
|
||||
from collections.abc import Collection, Mapping
|
||||
from io import BufferedReader, BytesIO
|
||||
from typing import Any, Final, cast, get_type_hints
|
||||
from pathlib import Path
|
||||
from typing import Any, Final, Protocol, cast, get_type_hints, runtime_checkable
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.token_counter import get_image_type
|
||||
|
|
@ -16,7 +19,7 @@ class ImageEditRequestUtils:
|
|||
image_edit_provider_config: BaseImageEditConfig,
|
||||
image_edit_optional_params: ImageEditOptionalRequestParams,
|
||||
drop_params: bool | None = None,
|
||||
additional_drop_params: list[str] | None = None,
|
||||
additional_drop_params: Collection[str] | None = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Get optional parameters for the image edit API.
|
||||
|
|
@ -145,3 +148,36 @@ class ImageEditRequestUtils:
|
|||
except Exception:
|
||||
# If anything goes wrong, default to PNG
|
||||
return FILE_MIME_TYPES[FileType.PNG]
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class _SeekableBinaryReader(Protocol):
|
||||
def tell(self) -> int: ...
|
||||
|
||||
def seek(self, offset: int) -> int: ...
|
||||
|
||||
def read(self) -> bytes: ...
|
||||
|
||||
|
||||
def _read_image_bytes(image: object) -> bytes:
|
||||
if isinstance(image, bytes):
|
||||
return image
|
||||
if isinstance(image, tuple):
|
||||
return _read_image_bytes(image[1])
|
||||
if isinstance(image, os.PathLike):
|
||||
return Path(image).read_bytes()
|
||||
if isinstance(image, _SeekableBinaryReader):
|
||||
position: Final = image.tell()
|
||||
image.seek(0)
|
||||
data: Final = image.read()
|
||||
image.seek(position)
|
||||
return data
|
||||
raise ValueError(f"Unsupported image type for image edit: {type(image).__name__}")
|
||||
|
||||
|
||||
def image_to_data_url(image: object) -> str:
|
||||
if isinstance(image, str):
|
||||
return image
|
||||
image_bytes: Final = _read_image_bytes(image)
|
||||
mime_type: Final = ImageEditRequestUtils.get_image_content_type(image_bytes)
|
||||
return f"data:{mime_type};base64,{base64.b64encode(image_bytes).decode('utf-8')}"
|
||||
|
|
|
|||
|
|
@ -436,6 +436,7 @@ _POSTGRES_OPERATION_BY_CALL_TYPE: Final[Mapping[str, PostgresOperation]] = Mappi
|
|||
"reset_budget_windows": PostgresOperation("select", None),
|
||||
"write_budget_windows": PostgresOperation("update", None),
|
||||
"roll_window_spend_row": PostgresOperation("update", None),
|
||||
"open_window_spend_row": PostgresOperation("upsert", None),
|
||||
"seed_window_spend": PostgresOperation("select", None),
|
||||
"select_window_spend_rows": PostgresOperation("select", None),
|
||||
"commit_window_spend_updates": PostgresOperation("upsert", None),
|
||||
|
|
|
|||
|
|
@ -97,6 +97,8 @@ WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY: Final = "websearch_native_blocks"
|
|||
|
||||
_RESPONSE_CONTENT_FIELD: Final = "content"
|
||||
|
||||
_CLAUDE_CODE_WEB_SEARCH_PROMPT_PREFIX: Final = "Perform a web search for the query:"
|
||||
|
||||
_ResponseT: Final = TypeVar("_ResponseT")
|
||||
|
||||
|
||||
|
|
@ -363,7 +365,8 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
get_last_user_message,
|
||||
)
|
||||
|
||||
query: Final = get_last_user_message(cast(list[AllMessageValues], messages))
|
||||
last_user_message: Final = get_last_user_message(cast(list[AllMessageValues], messages)) or ""
|
||||
query: Final = last_user_message.removeprefix(_CLAUDE_CODE_WEB_SEARCH_PROMPT_PREFIX).strip()
|
||||
if not query:
|
||||
return None
|
||||
|
||||
|
|
|
|||
|
|
@ -75,6 +75,8 @@ OPTIONAL_KWARGS_KEYS: Final = (
|
|||
"rpm",
|
||||
"itpm",
|
||||
"otpm",
|
||||
"tpd",
|
||||
"rpd",
|
||||
"use_xai_oauth",
|
||||
"chatgpt_auth_file",
|
||||
"github_copilot_auth_type",
|
||||
|
|
|
|||
|
|
@ -2758,7 +2758,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
if entry is None or not is_str_keyed_dict(litellm_params):
|
||||
return
|
||||
metadata: Final[object] = litellm_params.get("metadata")
|
||||
bucket: Final = metadata if is_str_keyed_dict(metadata) else {}
|
||||
bucket: Final = {**metadata} if is_str_keyed_dict(metadata) else {}
|
||||
litellm_params["metadata"] = bucket
|
||||
with_guardrail_entry(bucket, entry)
|
||||
|
||||
|
|
@ -3481,7 +3481,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
print_verbose=print_verbose,
|
||||
level=LogfireLevel.INFO.value,
|
||||
level=LogfireLevel.INFO,
|
||||
)
|
||||
|
||||
if callback == "lunary" and lunaryLogger is not None:
|
||||
|
|
@ -4442,7 +4442,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
response_obj=result,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
level=LogfireLevel.ERROR.value,
|
||||
level=LogfireLevel.ERROR,
|
||||
print_verbose=print_verbose,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ from collections.abc import Mapping, Sequence
|
|||
from dataclasses import dataclass, field, replace
|
||||
from datetime import datetime, timezone, tzinfo
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal, TypedDict, cast
|
||||
from typing import Final, Literal, Protocol, TypedDict, cast
|
||||
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
|
||||
|
||||
from typing_extensions import ReadOnly
|
||||
|
|
@ -2085,6 +2085,41 @@ def calculate_image_response_web_search_cost(
|
|||
)
|
||||
|
||||
|
||||
class _ModelPricedImageCostCalculator(Protocol):
|
||||
def __call__(self, model: str, image_response: ImageResponse, model_info: ModelInfo | None = None) -> float: ...
|
||||
|
||||
|
||||
def _model_priced_image_cost_calculator(
|
||||
custom_llm_provider: str | None, model: str
|
||||
) -> _ModelPricedImageCostCalculator | None:
|
||||
match custom_llm_provider:
|
||||
case litellm.LlmProviders.RECRAFT.value:
|
||||
from litellm.llms.recraft.cost_calculator import cost_calculator as recraft_image_cost_calculator
|
||||
|
||||
return recraft_image_cost_calculator
|
||||
case litellm.LlmProviders.AIML.value:
|
||||
from litellm.llms.aiml.image_generation.cost_calculator import cost_calculator as aiml_image_cost_calculator
|
||||
|
||||
return aiml_image_cost_calculator
|
||||
case litellm.LlmProviders.COMETAPI.value:
|
||||
from litellm.llms.cometapi.image_generation.cost_calculator import (
|
||||
cost_calculator as cometapi_image_cost_calculator,
|
||||
)
|
||||
|
||||
return cometapi_image_cost_calculator
|
||||
case litellm.LlmProviders.RUNWAYML.value:
|
||||
from litellm.llms.runwayml.cost_calculator import cost_calculator as runwayml_image_cost_calculator
|
||||
|
||||
return runwayml_image_cost_calculator
|
||||
case litellm.LlmProviders.TENCENT.value:
|
||||
from litellm.llms.tencent.cost_calculator import image_cost_calculator as tencent_image_cost_calculator
|
||||
from litellm.llms.tencent.image_generation.transformation import is_tokenhub_image_model
|
||||
|
||||
return tencent_image_cost_calculator if is_tokenhub_image_model(model) else None
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
class CostCalculatorUtils:
|
||||
@staticmethod
|
||||
def call_type_has_image_response(call_type: str) -> bool:
|
||||
|
|
@ -2131,9 +2166,6 @@ class CostCalculatorUtils:
|
|||
from litellm.llms.gemini.image_generation.cost_calculator import (
|
||||
cost_calculator as gemini_image_cost_calculator,
|
||||
)
|
||||
from litellm.llms.recraft.cost_calculator import (
|
||||
cost_calculator as recraft_image_cost_calculator,
|
||||
)
|
||||
from litellm.llms.vertex_ai.image_generation.cost_calculator import (
|
||||
cost_calculator as vertex_ai_image_cost_calculator,
|
||||
)
|
||||
|
|
@ -2146,7 +2178,10 @@ class CostCalculatorUtils:
|
|||
)
|
||||
resolved_n: Final = n if n is not None else (len(completion_response.data) if completion_response.data else 0)
|
||||
pricing: Final = deployment_pricing(model_info)
|
||||
model_priced_calculator: Final = _model_priced_image_cost_calculator(custom_llm_provider, model)
|
||||
|
||||
if model_priced_calculator is not None:
|
||||
return model_priced_calculator(model=model, image_response=completion_response, model_info=pricing)
|
||||
if custom_llm_provider == litellm.LlmProviders.VERTEX_AI.value:
|
||||
if isinstance(completion_response, ImageResponse):
|
||||
return vertex_ai_image_cost_calculator(
|
||||
|
|
@ -2164,36 +2199,6 @@ class CostCalculatorUtils:
|
|||
optional_params=optional_params,
|
||||
)
|
||||
raise TypeError("completion_response must be of type ImageResponse for bedrock image cost calculation")
|
||||
elif custom_llm_provider == litellm.LlmProviders.RECRAFT.value:
|
||||
from litellm.llms.recraft.cost_calculator import (
|
||||
cost_calculator as recraft_image_cost_calculator,
|
||||
)
|
||||
|
||||
return recraft_image_cost_calculator(
|
||||
model=model,
|
||||
image_response=completion_response,
|
||||
model_info=pricing,
|
||||
)
|
||||
elif custom_llm_provider == litellm.LlmProviders.AIML.value:
|
||||
from litellm.llms.aiml.image_generation.cost_calculator import (
|
||||
cost_calculator as aiml_image_cost_calculator,
|
||||
)
|
||||
|
||||
return aiml_image_cost_calculator(
|
||||
model=model,
|
||||
image_response=completion_response,
|
||||
model_info=pricing,
|
||||
)
|
||||
elif custom_llm_provider == litellm.LlmProviders.COMETAPI.value:
|
||||
from litellm.llms.cometapi.image_generation.cost_calculator import (
|
||||
cost_calculator as cometapi_image_cost_calculator,
|
||||
)
|
||||
|
||||
return cometapi_image_cost_calculator(
|
||||
model=model,
|
||||
image_response=completion_response,
|
||||
model_info=pricing,
|
||||
)
|
||||
elif custom_llm_provider == litellm.LlmProviders.GEMINI.value:
|
||||
if call_type in (
|
||||
CallTypes.image_edit.value,
|
||||
|
|
@ -2237,16 +2242,6 @@ class CostCalculatorUtils:
|
|||
optional_params=optional_params,
|
||||
model_info=pricing,
|
||||
)
|
||||
elif custom_llm_provider == litellm.LlmProviders.RUNWAYML.value:
|
||||
from litellm.llms.runwayml.cost_calculator import (
|
||||
cost_calculator as runwayml_image_cost_calculator,
|
||||
)
|
||||
|
||||
return runwayml_image_cost_calculator(
|
||||
model=model,
|
||||
image_response=completion_response,
|
||||
model_info=pricing,
|
||||
)
|
||||
elif (
|
||||
custom_llm_provider == litellm.LlmProviders.OPENAI.value
|
||||
or custom_llm_provider == litellm.LlmProviders.AZURE.value
|
||||
|
|
|
|||
|
|
@ -34,7 +34,7 @@ def safe_json_structure(
|
|||
def _transform(key: str | None, value: str) -> str:
|
||||
return value if value_transform is None else value_transform(key, value)
|
||||
|
||||
def _serialize(obj: object, seen: set[int], depth: int, key: str | None = None) -> Any:
|
||||
def _serialize(obj: object, seen: set[int], depth: int, key: str | None = None) -> object:
|
||||
# Check for maximum depth.
|
||||
if depth > max_depth:
|
||||
return "MaxDepthExceeded"
|
||||
|
|
@ -48,32 +48,32 @@ def safe_json_structure(
|
|||
if id(obj) in seen:
|
||||
return "CircularReference Detected"
|
||||
seen.add(id(obj))
|
||||
result: dict | list | tuple | set | str
|
||||
if isinstance(obj, dict):
|
||||
result = {}
|
||||
for k, v in obj.items():
|
||||
if isinstance(k, (str)):
|
||||
clean_k = k.replace("\x00", "") if "\x00" in k else k
|
||||
result[clean_k] = _serialize(v, seen, depth + 1, clean_k)
|
||||
mapping_result: Final = {
|
||||
(clean_k := k.replace("\x00", "") if "\x00" in k else k): _serialize(v, seen, depth + 1, clean_k)
|
||||
for k, v in obj.items()
|
||||
if isinstance(k, str)
|
||||
}
|
||||
seen.remove(id(obj))
|
||||
return result
|
||||
return mapping_result
|
||||
elif isinstance(obj, list):
|
||||
result = [_serialize(item, seen, depth + 1, key) for item in obj]
|
||||
list_result: Final = [_serialize(item, seen, depth + 1, key) for item in obj]
|
||||
seen.remove(id(obj))
|
||||
return result
|
||||
return list_result
|
||||
elif isinstance(obj, tuple):
|
||||
result = tuple(_serialize(item, seen, depth + 1, key) for item in obj)
|
||||
tuple_result: Final = tuple(_serialize(item, seen, depth + 1, key) for item in obj)
|
||||
seen.remove(id(obj))
|
||||
return result
|
||||
return tuple_result
|
||||
elif isinstance(obj, set):
|
||||
result = sorted([_serialize(item, seen, depth + 1, key) for item in obj])
|
||||
sortable_values: Final[list[Any]] = [_serialize(item, seen, depth + 1, key) for item in obj]
|
||||
set_result: Final = sorted(sortable_values)
|
||||
seen.remove(id(obj))
|
||||
return result
|
||||
return set_result
|
||||
elif isinstance(obj, BaseModel):
|
||||
dumped: Final = obj.model_dump()
|
||||
result = _serialize(dumped, seen, depth + 1, key)
|
||||
model_result: Final = _serialize(dumped, seen, depth + 1, key)
|
||||
seen.remove(id(obj))
|
||||
return result
|
||||
return model_result
|
||||
else:
|
||||
# Fall back to string conversion for non-serializable objects.
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -268,6 +268,7 @@ async def caller_within_model_budget(
|
|||
user_id=user_id,
|
||||
user_model_max_budget=user_model_max_budget,
|
||||
model=model,
|
||||
team_id=user_api_key_auth.team_id,
|
||||
)
|
||||
except litellm.BudgetExceededError:
|
||||
return "denied"
|
||||
|
|
@ -309,6 +310,7 @@ async def caller_within_model_budget(
|
|||
end_user_id=end_user_id,
|
||||
end_user_model_max_budget=end_user_model_max_budget,
|
||||
model=model,
|
||||
team_id=user_api_key_auth.team_id,
|
||||
)
|
||||
except litellm.BudgetExceededError:
|
||||
return "denied"
|
||||
|
|
|
|||
|
|
@ -10,6 +10,8 @@ import time
|
|||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, TypeAdapter, with_config
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -34,6 +36,15 @@ else:
|
|||
LiteLLMLoggingObj = Any
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True))
|
||||
class _PegasusResponse(TypedDict):
|
||||
message: ReadOnly[NotRequired[str | None]]
|
||||
finishReason: ReadOnly[NotRequired[str | None]]
|
||||
|
||||
|
||||
_PEGASUS_RESPONSE: Final = TypeAdapter(_PegasusResponse)
|
||||
|
||||
|
||||
class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig):
|
||||
"""
|
||||
Handles transforming OpenAI-style requests into Bedrock InvokeModel requests for
|
||||
|
|
@ -207,7 +218,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig):
|
|||
ModelResponse with choices[0].message.content and finish_reason
|
||||
"""
|
||||
try:
|
||||
completion_response: Final = raw_response.json()
|
||||
completion_response: Final = _PEGASUS_RESPONSE.validate_python(raw_response.json())
|
||||
except Exception as e:
|
||||
raise BedrockError(
|
||||
message=f"Error parsing response: {raw_response.text}, error: {e}",
|
||||
|
|
@ -224,7 +235,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig):
|
|||
message_content: Final = completion_response.get("message", "")
|
||||
|
||||
# Extract finish reason and map to LiteLLM format
|
||||
finish_reason_raw: Final = completion_response.get("finishReason", "stop")
|
||||
finish_reason_raw: Final = completion_response.get("finishReason") or "stop"
|
||||
finish_reason: Final = map_finish_reason(finish_reason_raw)
|
||||
|
||||
# Set the response content
|
||||
|
|
|
|||
|
|
@ -1,11 +1,12 @@
|
|||
import copy
|
||||
import json
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final, cast, get_args
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -48,6 +49,27 @@ from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
|||
|
||||
_GUARDRAIL_CONFIG_VALIDATOR: Final["TypeAdapter[GuardrailConfigBlock]"] = TypeAdapter(GuardrailConfigBlock)
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True))
|
||||
class _CohereGeneration(TypedDict):
|
||||
text: ReadOnly[str | None]
|
||||
finish_reason: ReadOnly[str | None]
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True))
|
||||
class _CohereInvokeResponse(TypedDict):
|
||||
text: NotRequired[ReadOnly[str | None]]
|
||||
generations: NotRequired[ReadOnly[Sequence[_CohereGeneration]]]
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True))
|
||||
class _LlamaInvokeResponse(TypedDict):
|
||||
generation: ReadOnly[str | None]
|
||||
|
||||
|
||||
_COHERE_INVOKE_RESPONSE: Final = TypeAdapter(_CohereInvokeResponse)
|
||||
_LLAMA_INVOKE_RESPONSE: Final = TypeAdapter(_LlamaInvokeResponse)
|
||||
|
||||
_GUARDRAIL_CONFIG_EXPECTED_FORMAT: Final = (
|
||||
"{'guardrailIdentifier': str, 'guardrailVersion': str, 'trace': 'enabled'|'disabled'|'enabled_full'}"
|
||||
)
|
||||
|
|
@ -310,12 +332,13 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
|
|||
outputText: str | None = None
|
||||
try:
|
||||
if provider == "cohere":
|
||||
if "text" in completion_response:
|
||||
outputText = completion_response["text"]
|
||||
elif "generations" in completion_response:
|
||||
outputText = completion_response["generations"][0]["text"]
|
||||
cohere_response: Final = _COHERE_INVOKE_RESPONSE.validate_python(completion_response)
|
||||
if "text" in cohere_response:
|
||||
outputText = cohere_response["text"]
|
||||
elif "generations" in cohere_response:
|
||||
outputText = cohere_response["generations"][0]["text"]
|
||||
model_response.choices[0].finish_reason = map_finish_reason(
|
||||
completion_response["generations"][0]["finish_reason"]
|
||||
cohere_response["generations"][0]["finish_reason"] or "stop"
|
||||
)
|
||||
elif provider == "anthropic":
|
||||
return litellm.AmazonAnthropicClaudeConfig().transform_response(
|
||||
|
|
@ -361,7 +384,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
|
|||
elif provider == "ai21":
|
||||
outputText = completion_response.get("completions")[0].get("data").get("text")
|
||||
elif provider == "meta" or provider == "llama" or provider == "deepseek_r1":
|
||||
outputText = completion_response["generation"]
|
||||
outputText = _LLAMA_INVOKE_RESPONSE.validate_python(completion_response)["generation"]
|
||||
elif provider == "mistral":
|
||||
outputText = litellm.AmazonMistralConfig.get_outputText(completion_response, model_response)
|
||||
else: # amazon titan
|
||||
|
|
|
|||
|
|
@ -53,6 +53,7 @@ from litellm.types.realtime import (
|
|||
RealtimeResponseTypedDict,
|
||||
)
|
||||
|
||||
_JSON_OBJECT_MESSAGE: Final = TypeAdapter(dict[str, JsonValue], config=ConfigDict(strict=True))
|
||||
_JSON_VALUE: Final = TypeAdapter(object)
|
||||
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
|
@ -856,7 +857,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
|
|||
List of Bedrock format messages (JSON strings)
|
||||
"""
|
||||
try:
|
||||
json_message: Final[dict[str, JsonValue]] = json.loads(message)
|
||||
json_message: Final = _JSON_OBJECT_MESSAGE.validate_python(json.loads(message))
|
||||
except json.JSONDecodeError:
|
||||
verbose_logger.warning("Invalid JSON message: %s", message[:200])
|
||||
return []
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ import ssl
|
|||
import sys
|
||||
import typing
|
||||
import urllib.request
|
||||
import weakref
|
||||
from collections.abc import Callable, Generator
|
||||
from typing import ClassVar, Final
|
||||
|
||||
|
|
@ -109,6 +110,24 @@ def map_aiohttp_exceptions() -> Generator[None, None, None]:
|
|||
raise mapped_exc(message) from exc
|
||||
|
||||
|
||||
def _release_transport_anchor(_transport: object) -> None:
|
||||
"""Release a transport anchored to a response stream. See ``_anchor_transport_to``."""
|
||||
|
||||
|
||||
def _anchor_transport_to(stream: "AiohttpResponseStream", transport: object) -> None:
|
||||
"""Keep the transport alive for as long as the stream can still read.
|
||||
|
||||
The transport finalizer closes the session it owns, and nothing on the path from a
|
||||
live response back to that session references the transport, so a client dropped
|
||||
mid-stream would close the session under the body. ``weakref.finalize`` holds the
|
||||
transport in its own registry rather than on the stream, which keeps the transport
|
||||
out of the response's reference cycle (an exception carrying the response, for one)
|
||||
so it is finalized by refcount once the anchor drops and can still schedule an async
|
||||
close, instead of being collected in the same cyclic pass as the session it closes.
|
||||
"""
|
||||
weakref.finalize(stream, _release_transport_anchor, transport)
|
||||
|
||||
|
||||
class AiohttpResponseStream(httpx.AsyncByteStream):
|
||||
CHUNK_SIZE = 1024 * 16
|
||||
|
||||
|
|
@ -463,10 +482,12 @@ class LiteLLMAiohttpTransport(AiohttpTransport):
|
|||
# Re-raise if it's a different RuntimeError
|
||||
raise
|
||||
|
||||
stream: Final = AiohttpResponseStream(response)
|
||||
_anchor_transport_to(stream, self)
|
||||
return httpx.Response(
|
||||
status_code=response.status,
|
||||
headers=response.headers,
|
||||
stream=AiohttpResponseStream(response),
|
||||
stream=stream,
|
||||
request=request,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -14,6 +14,8 @@ Docs - https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api
|
|||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, TypeAdapter, with_config
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.url_utils import ensure_api_route
|
||||
|
|
@ -28,6 +30,23 @@ from ..common_utils import DashScopeError
|
|||
DEFAULT_API_BASE: Final = "https://dashscope.aliyuncs.com/compatible-mode/v1"
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True))
|
||||
class _DashScopeUsage(TypedDict, total=False):
|
||||
prompt_tokens: NotRequired[ReadOnly[int]]
|
||||
total_tokens: NotRequired[ReadOnly[int]]
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True))
|
||||
class _DashScopeEmbeddingResponse(TypedDict, total=False):
|
||||
data: NotRequired[ReadOnly[list[object]]]
|
||||
model: NotRequired[ReadOnly[str | None]]
|
||||
usage: NotRequired[ReadOnly[_DashScopeUsage | None]]
|
||||
id: NotRequired[ReadOnly[object]]
|
||||
|
||||
|
||||
_DASHSCOPE_EMBEDDING_RESPONSE: Final = TypeAdapter(_DashScopeEmbeddingResponse)
|
||||
|
||||
|
||||
class DashScopeEmbeddingConfig(BaseEmbeddingConfig):
|
||||
"""
|
||||
Reference: https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api
|
||||
|
|
@ -152,11 +171,12 @@ class DashScopeEmbeddingConfig(BaseEmbeddingConfig):
|
|||
message=message,
|
||||
)
|
||||
|
||||
parsed_response: Final = _DASHSCOPE_EMBEDDING_RESPONSE.validate_python(response_json)
|
||||
model_response.object = "list"
|
||||
model_response.data = response_json.get("data", [])
|
||||
model_response.model = response_json.get("model", model)
|
||||
model_response.data = parsed_response.get("data", [])
|
||||
model_response.model = parsed_response.get("model", model)
|
||||
|
||||
usage: Final = response_json.get("usage") or {}
|
||||
usage: Final = parsed_response.get("usage") or {}
|
||||
prompt_tokens: Final = usage.get("prompt_tokens", 0)
|
||||
total_tokens: Final = usage.get("total_tokens", prompt_tokens)
|
||||
setattr(
|
||||
|
|
@ -169,8 +189,8 @@ class DashScopeEmbeddingConfig(BaseEmbeddingConfig):
|
|||
),
|
||||
)
|
||||
|
||||
if "id" in response_json:
|
||||
setattr(model_response, "id", response_json["id"])
|
||||
if "id" in parsed_response:
|
||||
setattr(model_response, "id", parsed_response["id"])
|
||||
|
||||
return model_response
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
import json
|
||||
from typing import Final
|
||||
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm import verbose_logger
|
||||
from litellm.types.llms.openai import (
|
||||
|
|
@ -23,6 +25,9 @@ def _usage_block(usage: Usage) -> ChatCompletionUsageBlock:
|
|||
)
|
||||
|
||||
|
||||
_JSON_CHUNK: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True))
|
||||
|
||||
|
||||
class ModelResponseIterator:
|
||||
def __init__(self, streaming_response, sync_stream: bool):
|
||||
self.streaming_response = streaming_response
|
||||
|
|
@ -99,7 +104,7 @@ class ModelResponseIterator:
|
|||
self.response_iterator = self.streaming_response
|
||||
chunk: Final = self._next_sse_data()
|
||||
try:
|
||||
json_chunk: Final = json.loads(chunk)
|
||||
json_chunk: Final = _JSON_CHUNK.validate_python(json.loads(chunk))
|
||||
return self.chunk_parser(chunk=json_chunk)
|
||||
except ValueError as e:
|
||||
verbose_logger.debug(
|
||||
|
|
@ -126,7 +131,7 @@ class ModelResponseIterator:
|
|||
if chunk == "[DONE]":
|
||||
raise StopAsyncIteration
|
||||
try:
|
||||
json_chunk: Final = json.loads(chunk)
|
||||
json_chunk: Final = _JSON_CHUNK.validate_python(json.loads(chunk))
|
||||
return self.chunk_parser(chunk=json_chunk)
|
||||
except ValueError as e:
|
||||
verbose_logger.debug(
|
||||
|
|
|
|||
|
|
@ -4,13 +4,14 @@ from typing import Final
|
|||
|
||||
from httpx._types import RequestFiles
|
||||
|
||||
from litellm.images.utils import image_to_data_url
|
||||
from litellm.litellm_core_utils.url_utils import join_api_base
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.images.main import ImageEditOptionalRequestParams
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import FileTypes
|
||||
|
||||
from .transformation import DEFAULT_BASE_URL, FalAIImageEditConfig, to_data_url
|
||||
from .transformation import DEFAULT_BASE_URL, FalAIImageEditConfig
|
||||
|
||||
FLUX_LORA_DEPTH_ENDPOINT: Final[str] = "fal-ai/flux-lora-depth"
|
||||
SUPPORTED_OPENAI_PARAMS: Final[tuple[str, ...]] = ("n", "size")
|
||||
|
|
@ -68,7 +69,7 @@ class FalAIFluxLoraDepthEditConfig(FalAIImageEditConfig):
|
|||
)
|
||||
request_body: Final[dict[str, object]] = { # mutable-ok: base class contract returns a dict
|
||||
"prompt": prompt,
|
||||
"image_url": to_data_url(next(iter(images))),
|
||||
"image_url": image_to_data_url(next(iter(images))),
|
||||
**provider_params,
|
||||
}
|
||||
return request_body, ()
|
||||
|
|
|
|||
|
|
@ -1,15 +1,12 @@
|
|||
import base64
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Protocol, runtime_checkable
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
import httpx
|
||||
from httpx._types import RequestFiles
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
|
||||
from litellm.images.utils import ImageEditRequestUtils
|
||||
from litellm.images.utils import image_to_data_url
|
||||
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
|
||||
from litellm.llms.fal_ai.image_generation.gpt_image_2_transformation import (
|
||||
map_gpt_image_quality,
|
||||
|
|
@ -39,39 +36,6 @@ PARAM_TRANSLATION: Final[Mapping[str, str]] = MappingProxyType(
|
|||
)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class _SeekableBinaryReader(Protocol):
|
||||
def tell(self) -> int: ...
|
||||
|
||||
def seek(self, offset: int) -> int: ...
|
||||
|
||||
def read(self) -> bytes: ...
|
||||
|
||||
|
||||
def _read_image_bytes(image: object) -> bytes:
|
||||
if isinstance(image, bytes):
|
||||
return image
|
||||
if isinstance(image, tuple):
|
||||
return _read_image_bytes(image[1])
|
||||
if isinstance(image, os.PathLike):
|
||||
return Path(image).read_bytes()
|
||||
if isinstance(image, _SeekableBinaryReader):
|
||||
position: Final = image.tell()
|
||||
image.seek(0)
|
||||
data: Final = image.read()
|
||||
image.seek(position)
|
||||
return data
|
||||
raise ValueError(f"Unsupported image type for Fal AI image edit: {type(image).__name__}")
|
||||
|
||||
|
||||
def to_data_url(image: object) -> str:
|
||||
if isinstance(image, str):
|
||||
return image
|
||||
image_bytes: Final = _read_image_bytes(image)
|
||||
mime_type: Final = ImageEditRequestUtils.get_image_content_type(image_bytes)
|
||||
return f"data:{mime_type};base64,{base64.b64encode(image_bytes).decode('utf-8')}"
|
||||
|
||||
|
||||
def _first(value: object) -> object:
|
||||
return value[0] if isinstance(value, list) and value else value
|
||||
|
||||
|
|
@ -146,14 +110,14 @@ class FalAIImageEditConfig(BaseImageEditConfig):
|
|||
raise ValueError("Fal AI image edit requires at least one input image")
|
||||
mask: Final = _first(image_edit_optional_request_params.get("mask"))
|
||||
mask_field: Final[Mapping[str, str]] = (
|
||||
MappingProxyType({"mask_url": to_data_url(mask)}) if mask is not None else MappingProxyType({})
|
||||
MappingProxyType({"mask_url": image_to_data_url(mask)}) if mask is not None else MappingProxyType({})
|
||||
)
|
||||
provider_params: Final[Mapping[str, object]] = MappingProxyType(
|
||||
{key: value for key, value in image_edit_optional_request_params.items() if key != "mask"}
|
||||
)
|
||||
request_body: Final[dict[str, object]] = { # mutable-ok: base class contract returns a dict
|
||||
"prompt": prompt,
|
||||
"image_urls": tuple(to_data_url(img) for img in images),
|
||||
"image_urls": tuple(image_to_data_url(img) for img in images),
|
||||
**mask_field,
|
||||
**provider_params,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -12,6 +12,8 @@ from types import MappingProxyType
|
|||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, TypeAdapter, with_config
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.caching.caching import InMemoryCache
|
||||
|
|
@ -37,6 +39,17 @@ _token_cache: Final = InMemoryCache()
|
|||
_NO_LITELLM_PARAMS: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True))
|
||||
class _GigaChatTokenResponse(TypedDict):
|
||||
tok: NotRequired[ReadOnly[str | None]]
|
||||
access_token: NotRequired[ReadOnly[str | None]]
|
||||
exp: NotRequired[ReadOnly[int | float | str | None]]
|
||||
expires_at: NotRequired[ReadOnly[int | float | str | None]]
|
||||
|
||||
|
||||
_TOKEN_RESPONSE: Final = TypeAdapter(_GigaChatTokenResponse)
|
||||
|
||||
|
||||
class GigaChatAuthError(BaseLLMException):
|
||||
"""GigaChat authentication error."""
|
||||
|
||||
|
|
@ -230,7 +243,7 @@ async def _request_token_async(
|
|||
|
||||
def _parse_token_response(response: httpx.Response) -> tuple[str, int]:
|
||||
"""Parse OAuth token response."""
|
||||
data: Final = response.json()
|
||||
data: Final = _TOKEN_RESPONSE.validate_python(response.json())
|
||||
|
||||
# GigaChat returns either 'tok'/'exp' or 'access_token'/'expires_at'
|
||||
access_token: Final = data.get("tok") or data.get("access_token")
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ This requires websockets, and is currently only supported on LiteLLM Proxy.
|
|||
"""
|
||||
|
||||
import ssl
|
||||
from typing import Any, Final, cast
|
||||
from typing import TYPE_CHECKING, Final, cast
|
||||
|
||||
from litellm._logging import redact_string, verbose_logger
|
||||
from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES
|
||||
|
|
@ -23,6 +23,9 @@ from ....llms.custom_httpx.http_handler import get_shared_realtime_ssl_context
|
|||
from ..common_utils import is_openai_backed_api_base
|
||||
from ..openai import OpenAIChatCompletion
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastapi import WebSocket
|
||||
|
||||
|
||||
class OpenAIRealtime(OpenAIChatCompletion):
|
||||
"""
|
||||
|
|
@ -131,7 +134,7 @@ class OpenAIRealtime(OpenAIChatCompletion):
|
|||
async def async_realtime(
|
||||
self,
|
||||
model: str,
|
||||
websocket: Any,
|
||||
websocket: "WebSocket",
|
||||
logging_obj: LiteLLMLogging,
|
||||
api_base: str | None = None,
|
||||
api_key: str | None = None,
|
||||
|
|
|
|||
|
|
@ -12,12 +12,14 @@ This module decodes them into float arrays for OpenAI-compatible responses.
|
|||
"""
|
||||
|
||||
import base64
|
||||
import builtins
|
||||
import struct
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Final, Literal
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
from pydantic import ConfigDict, TypeAdapter, with_config
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.url_utils import join_api_base
|
||||
|
|
@ -31,6 +33,17 @@ _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_i
|
|||
_INT: Final = TypeAdapter(int)
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True))
|
||||
class _PerplexityEmbeddingResponse(TypedDict, total=False):
|
||||
model: ReadOnly[str | None]
|
||||
object: ReadOnly[Literal["list"]]
|
||||
data: ReadOnly[Sequence[builtins.object] | Mapping[str, builtins.object] | Literal[""]]
|
||||
usage: ReadOnly[builtins.object]
|
||||
|
||||
|
||||
_EMBEDDING_RESPONSE_ADAPTER: Final = TypeAdapter(_PerplexityEmbeddingResponse)
|
||||
|
||||
|
||||
class PerplexityEmbeddingError(BaseLLMException):
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -153,14 +166,18 @@ class PerplexityEmbeddingConfig(BaseEmbeddingConfig):
|
|||
litellm_params: dict = {},
|
||||
) -> EmbeddingResponse:
|
||||
try:
|
||||
raw_response_json: Final = raw_response.json()
|
||||
raw_response_json: Final[object] = raw_response.json()
|
||||
except Exception:
|
||||
raise PerplexityEmbeddingError(message=raw_response.text, status_code=raw_response.status_code)
|
||||
|
||||
model_response.model = raw_response_json.get("model", model)
|
||||
model_response.object = raw_response_json.get("object", "list")
|
||||
response_data: Final[_PerplexityEmbeddingResponse] = _EMBEDDING_RESPONSE_ADAPTER.validate_python(
|
||||
raw_response_json
|
||||
)
|
||||
|
||||
raw_data: Final = raw_response_json.get("data", [])
|
||||
model_response.model = response_data.get("model", model)
|
||||
model_response.object = response_data.get("object", "list")
|
||||
|
||||
raw_data: Final = response_data.get("data", [])
|
||||
decoded_data: Final[list[dict[str, object]]] = []
|
||||
for item in raw_data:
|
||||
decoded_item = dict(_JSON_OBJECT.validate_python(item))
|
||||
|
|
@ -168,7 +185,7 @@ class PerplexityEmbeddingConfig(BaseEmbeddingConfig):
|
|||
decoded_data.append(decoded_item)
|
||||
model_response.data = decoded_data
|
||||
|
||||
usage_data: Final = _JSON_OBJECT.validate_python(raw_response_json.get("usage", {}))
|
||||
usage_data: Final = _JSON_OBJECT.validate_python(response_data.get("usage", {}))
|
||||
usage: Final = Usage(
|
||||
prompt_tokens=_INT.validate_python(
|
||||
usage_data.get("prompt_tokens", 0) or usage_data.get("total_tokens", 0) or 0
|
||||
|
|
|
|||
|
|
@ -1,6 +1,49 @@
|
|||
from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
|
||||
from litellm.types.utils import Usage
|
||||
from typing import Final
|
||||
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import (
|
||||
calculate_image_response_cost_from_usage,
|
||||
flat_image_cost,
|
||||
generic_cost_per_token,
|
||||
resolve_image_model_info,
|
||||
)
|
||||
from litellm.types.utils import ImageResponse, ModelInfo, Usage
|
||||
|
||||
|
||||
def cost_per_token(model: str, usage: Usage) -> tuple[float, float]:
|
||||
return generic_cost_per_token(model=model, usage=usage, custom_llm_provider="tencent")
|
||||
|
||||
|
||||
def _deployment_image_token_pricing(model_info: ModelInfo | None) -> ModelInfo | None:
|
||||
if model_info is None or model_info.get("output_cost_per_image_token") is not None:
|
||||
return model_info
|
||||
output_cost_per_token: Final = model_info.get("output_cost_per_token")
|
||||
if output_cost_per_token is None:
|
||||
return model_info
|
||||
priced: Final[ModelInfo] = {**model_info, "output_cost_per_image_token": output_cost_per_token}
|
||||
return priced
|
||||
|
||||
|
||||
def _deployment_prices_per_image_only(model_info: ModelInfo | None) -> bool:
|
||||
return (
|
||||
model_info is not None
|
||||
and model_info.get("output_cost_per_image") is not None
|
||||
and model_info.get("output_cost_per_image_token") is None
|
||||
and model_info.get("output_cost_per_token") is None
|
||||
)
|
||||
|
||||
|
||||
def image_cost_calculator(model: str, image_response: ImageResponse, model_info: ModelInfo | None = None) -> float:
|
||||
resolved_model_info: Final = resolve_image_model_info(
|
||||
model=model, custom_llm_provider="tencent", model_info=_deployment_image_token_pricing(model_info)
|
||||
)
|
||||
if _deployment_prices_per_image_only(model_info):
|
||||
return flat_image_cost(resolved_model_info, image_response)
|
||||
token_based_cost: Final = calculate_image_response_cost_from_usage(
|
||||
model=model,
|
||||
image_response=image_response,
|
||||
custom_llm_provider="tencent",
|
||||
model_info=resolved_model_info,
|
||||
)
|
||||
if token_based_cost is not None:
|
||||
return token_based_cost
|
||||
return flat_image_cost(resolved_model_info, image_response)
|
||||
|
|
|
|||
0
litellm/llms/tencent/image_edit/__init__.py
Normal file
0
litellm/llms/tencent/image_edit/__init__.py
Normal file
147
litellm/llms/tencent/image_edit/transformation.py
Normal file
147
litellm/llms/tencent/image_edit/transformation.py
Normal file
|
|
@ -0,0 +1,147 @@
|
|||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, TypeAlias
|
||||
|
||||
import httpx
|
||||
from httpx._types import RequestFiles
|
||||
from pydantic import TypeAdapter
|
||||
from typing_extensions import override
|
||||
|
||||
from litellm.images.utils import image_to_data_url
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
|
||||
from litellm.llms.tencent.image_generation.transformation import (
|
||||
TOKENHUB_IMAGE_PARAMS,
|
||||
tokenhub_image_headers,
|
||||
tokenhub_image_params,
|
||||
tokenhub_image_response,
|
||||
tokenhub_image_url,
|
||||
)
|
||||
from litellm.types.images.main import ImageEditOptionalRequestParams
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import ImageResponse
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
_IMAGES: Final = TypeAdapter(tuple[object, ...])
|
||||
_FORM_BOOLEANS: Final = MappingProxyType({"true": True, "false": False})
|
||||
_OPENAI_PARAMS: Final = frozenset(TOKENHUB_IMAGE_PARAMS)
|
||||
|
||||
_FormEntries: TypeAlias = tuple[tuple[tuple[str, ...], object], ...]
|
||||
|
||||
|
||||
def _form_scalar(value: object) -> object:
|
||||
match value:
|
||||
case str() if value.lower() in _FORM_BOOLEANS:
|
||||
return _FORM_BOOLEANS[value.lower()]
|
||||
case str() if value.removeprefix("-").isdigit():
|
||||
return int(value)
|
||||
case _:
|
||||
return value
|
||||
|
||||
|
||||
def _form_path(key: str) -> tuple[str, ...]:
|
||||
root, bracket, rest = key.partition("[")
|
||||
if not bracket or not rest.endswith("]"):
|
||||
return (key,)
|
||||
return (root, *rest.removesuffix("]").split("]["))
|
||||
|
||||
|
||||
def _nested_form_value(entries: _FormEntries) -> object:
|
||||
match entries:
|
||||
case (((), value),):
|
||||
return _form_scalar(value)
|
||||
case _:
|
||||
return _nested_form_fields(tuple(entry for entry in entries if entry[0]))
|
||||
|
||||
|
||||
def _nested_form_fields(entries: _FormEntries) -> Mapping[str, object]:
|
||||
roots: Final = tuple(dict.fromkeys(path[0] for path, _ in entries))
|
||||
return {
|
||||
root: _nested_form_value(tuple((path[1:], value) for path, value in entries if path[0] == root))
|
||||
for root in roots
|
||||
}
|
||||
|
||||
|
||||
def _tokenhub_form_params(params: Mapping[str, object]) -> Mapping[str, object]:
|
||||
tokenhub_entries: Final = tuple(
|
||||
(_form_path(key), value) for key, value in params.items() if key not in _OPENAI_PARAMS
|
||||
)
|
||||
return {
|
||||
**{key: value for key, value in params.items() if key in _OPENAI_PARAMS},
|
||||
**_nested_form_fields(tokenhub_entries),
|
||||
}
|
||||
|
||||
|
||||
class TencentImageEditConfig(BaseImageEditConfig):
|
||||
@override
|
||||
def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: mirrors override contract
|
||||
return list(TOKENHUB_IMAGE_PARAMS)
|
||||
|
||||
@override
|
||||
def map_openai_params(
|
||||
self,
|
||||
image_edit_optional_params: ImageEditOptionalRequestParams,
|
||||
model: str,
|
||||
drop_params: bool,
|
||||
) -> dict[str, object]: # mutable-ok: mirrors override contract
|
||||
return dict(tokenhub_image_params(image_edit_optional_params, model, drop_params))
|
||||
|
||||
@override
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: Mapping[str, str],
|
||||
model: str,
|
||||
api_key: str | None = None,
|
||||
litellm_params: Mapping[str, object] | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> dict[str, str]: # mutable-ok: mirrors override contract
|
||||
return dict(tokenhub_image_headers(headers, api_key, model))
|
||||
|
||||
@override
|
||||
def get_complete_url(
|
||||
self,
|
||||
model: str,
|
||||
api_base: str | None,
|
||||
litellm_params: Mapping[str, object],
|
||||
) -> str:
|
||||
return tokenhub_image_url(api_base)
|
||||
|
||||
@override
|
||||
def use_multipart_form_data(self) -> bool:
|
||||
return False
|
||||
|
||||
@override
|
||||
def transform_image_edit_request(
|
||||
self,
|
||||
model: str,
|
||||
prompt: str | None,
|
||||
image: object,
|
||||
image_edit_optional_request_params: Mapping[str, object],
|
||||
litellm_params: GenericLiteLLMParams,
|
||||
headers: Mapping[str, str],
|
||||
) -> tuple[dict[str, object], RequestFiles]: # mutable-ok: mirrors override contract
|
||||
images: Final = _IMAGES.validate_python(image) if isinstance(image, list) else (image,)
|
||||
return {
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"images": tuple(image_to_data_url(img) for img in images if img is not None),
|
||||
**_tokenhub_form_params(image_edit_optional_request_params),
|
||||
}, ()
|
||||
|
||||
@override
|
||||
def transform_image_edit_response(
|
||||
self,
|
||||
model: str,
|
||||
raw_response: httpx.Response,
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
) -> ImageResponse:
|
||||
image_response: Final = tokenhub_image_response(raw_response)
|
||||
if image_response is None:
|
||||
raise BaseLLMException(
|
||||
status_code=raw_response.status_code,
|
||||
message=raw_response.text,
|
||||
headers=raw_response.headers,
|
||||
)
|
||||
return image_response
|
||||
0
litellm/llms/tencent/image_generation/__init__.py
Normal file
0
litellm/llms/tencent/image_generation/__init__.py
Normal file
200
litellm/llms/tencent/image_generation/transformation.py
Normal file
200
litellm/llms/tencent/image_generation/transformation.py
Normal file
|
|
@ -0,0 +1,200 @@
|
|||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
|
||||
from typing_extensions import override
|
||||
|
||||
from litellm.exceptions import AuthenticationError, UnsupportedParamsError
|
||||
from litellm.litellm_core_utils.url_utils import ensure_api_route, join_api_base
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.base_llm.image_generation.transformation import (
|
||||
BaseImageGenerationConfig,
|
||||
)
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
OpenAIImageGenerationOptionalParams,
|
||||
)
|
||||
from litellm.types.utils import (
|
||||
ImageObject,
|
||||
ImageResponse,
|
||||
ImageUsage,
|
||||
ImageUsageInputTokensDetails,
|
||||
LlmProviders,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
TOKENHUB_API_BASE: Final = "https://tokenhub-intl.tencentcloudmaas.com/v1"
|
||||
TOKENHUB_IMAGE_PATH: Final = "/wand/si-image/generation"
|
||||
TOKENHUB_IMAGE_MODEL_PREFIX: Final = "seedream-"
|
||||
TOKENHUB_IMAGE_PARAMS: Final[tuple[OpenAIImageGenerationOptionalParams, ...]] = (
|
||||
"n",
|
||||
"size",
|
||||
"response_format",
|
||||
"output_format",
|
||||
"background",
|
||||
"user",
|
||||
)
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
|
||||
|
||||
|
||||
class _TokenHubImage(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
url: str | None = None
|
||||
b64_json: str | None = None
|
||||
size: str | None = None
|
||||
output_format: str | None = None
|
||||
|
||||
|
||||
class _TokenHubUsage(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
total_tokens: int
|
||||
|
||||
|
||||
class _TokenHubImageResponse(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, hide_input_in_errors=True)
|
||||
|
||||
created: int
|
||||
data: tuple[_TokenHubImage, ...]
|
||||
tokenhub_usage: _TokenHubUsage | None = None
|
||||
|
||||
|
||||
def is_tokenhub_image_model(model: str) -> bool:
|
||||
return model.removeprefix(f"{LlmProviders.TENCENT.value}/").startswith(TOKENHUB_IMAGE_MODEL_PREFIX)
|
||||
|
||||
|
||||
def tokenhub_image_url(api_base: str | None) -> str:
|
||||
base: Final = api_base or get_secret_str("TENCENT_API_BASE") or TOKENHUB_API_BASE
|
||||
return join_api_base(ensure_api_route(base, "v1"), TOKENHUB_IMAGE_PATH)
|
||||
|
||||
|
||||
def tokenhub_image_headers(headers: Mapping[str, str], api_key: str | None, model: str) -> Mapping[str, str]:
|
||||
resolved_api_key: Final = api_key or get_secret_str("TENCENT_API_KEY")
|
||||
if not resolved_api_key:
|
||||
raise AuthenticationError(
|
||||
message="TENCENT_API_KEY is not set. Set it in the environment or pass api_key.",
|
||||
llm_provider="tencent",
|
||||
model=model,
|
||||
)
|
||||
return {**headers, "Authorization": f"Bearer {resolved_api_key}", "Content-Type": "application/json"}
|
||||
|
||||
|
||||
def tokenhub_image_params(params: Mapping[str, object], model: str, drop_params: bool) -> Mapping[str, object]:
|
||||
n: Final = params.get("n")
|
||||
if n is not None and n != 1 and not drop_params:
|
||||
raise UnsupportedParamsError(
|
||||
status_code=400,
|
||||
message=f"Tencent TokenHub returns one image per request, so n={n} is not supported for {model}. "
|
||||
+ "To send n=1 instead, set `litellm.drop_params = True`.",
|
||||
)
|
||||
return {
|
||||
key: value
|
||||
for key, value in params.items()
|
||||
if key != "n" and value is not None and not (key == "background" and value == "auto")
|
||||
}
|
||||
|
||||
|
||||
def tokenhub_image_response(raw_response: httpx.Response) -> ImageResponse | None:
|
||||
try:
|
||||
parsed: Final = _TokenHubImageResponse.model_validate_json(raw_response.content)
|
||||
except ValidationError:
|
||||
return None
|
||||
total_tokens: Final = parsed.tokenhub_usage.total_tokens if parsed.tokenhub_usage is not None else 0
|
||||
first_image: Final = parsed.data[0] if parsed.data else _TokenHubImage()
|
||||
return ImageResponse(
|
||||
created=parsed.created,
|
||||
data=[ImageObject(url=image.url, b64_json=image.b64_json) for image in parsed.data],
|
||||
size=first_image.size,
|
||||
output_format=first_image.output_format,
|
||||
usage=ImageUsage(
|
||||
input_tokens=0,
|
||||
input_tokens_details=ImageUsageInputTokensDetails(image_tokens=0, text_tokens=0),
|
||||
output_tokens=total_tokens,
|
||||
total_tokens=total_tokens,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class TencentImageGenerationConfig(BaseImageGenerationConfig):
|
||||
@override
|
||||
def get_supported_openai_params(
|
||||
self, model: str
|
||||
) -> list[OpenAIImageGenerationOptionalParams]: # mutable-ok: mirrors override contract
|
||||
return list(TOKENHUB_IMAGE_PARAMS)
|
||||
|
||||
@override
|
||||
def map_openai_params(
|
||||
self,
|
||||
non_default_params: Mapping[str, object],
|
||||
optional_params: Mapping[str, object],
|
||||
model: str,
|
||||
drop_params: bool,
|
||||
) -> dict[str, object]: # mutable-ok: mirrors override contract
|
||||
return {**optional_params, **tokenhub_image_params(non_default_params, model, drop_params)}
|
||||
|
||||
@override
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: str | None,
|
||||
api_key: str | None,
|
||||
model: str,
|
||||
optional_params: Mapping[str, object],
|
||||
litellm_params: Mapping[str, object],
|
||||
stream: bool | None = None,
|
||||
) -> str:
|
||||
return tokenhub_image_url(api_base)
|
||||
|
||||
@override
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: Mapping[str, str],
|
||||
model: str,
|
||||
messages: Sequence[AllMessageValues],
|
||||
optional_params: Mapping[str, object],
|
||||
litellm_params: Mapping[str, object],
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> dict[str, str]: # mutable-ok: mirrors override contract
|
||||
return dict(tokenhub_image_headers(headers, api_key, model))
|
||||
|
||||
@override
|
||||
def transform_image_generation_request(
|
||||
self,
|
||||
model: str,
|
||||
prompt: str,
|
||||
optional_params: Mapping[str, object],
|
||||
litellm_params: Mapping[str, object],
|
||||
headers: Mapping[str, str],
|
||||
) -> dict[str, object]: # mutable-ok: mirrors override contract
|
||||
extra_body: Final = _JSON_OBJECT.validate_python(optional_params.get("extra_body") or {})
|
||||
params: Final = {key: value for key, value in optional_params.items() if key != "extra_body"}
|
||||
return {"model": model, "prompt": prompt, **params, **extra_body}
|
||||
|
||||
@override
|
||||
def transform_image_generation_response(
|
||||
self,
|
||||
model: str,
|
||||
raw_response: httpx.Response,
|
||||
model_response: ImageResponse,
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
request_data: Mapping[str, object],
|
||||
optional_params: Mapping[str, object],
|
||||
litellm_params: Mapping[str, object],
|
||||
encoding: object,
|
||||
api_key: str | None = None,
|
||||
json_mode: bool | None = None,
|
||||
) -> ImageResponse:
|
||||
image_response: Final = tokenhub_image_response(raw_response)
|
||||
if image_response is None:
|
||||
raise BaseLLMException(
|
||||
status_code=raw_response.status_code,
|
||||
message=raw_response.text,
|
||||
headers=raw_response.headers,
|
||||
)
|
||||
return image_response
|
||||
|
|
@ -2743,10 +2743,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
chunk_indices = support.get("groundingChunkIndices", [])
|
||||
|
||||
if start_index is not None and end_index is not None and chunk_indices:
|
||||
# Use the first chunk's URL for the annotation
|
||||
first_chunk_idx = chunk_indices[0]
|
||||
if first_chunk_idx in chunk_to_uri_map:
|
||||
uri_info = chunk_to_uri_map[first_chunk_idx]
|
||||
for chunk_index in chunk_indices:
|
||||
if chunk_index not in chunk_to_uri_map:
|
||||
continue
|
||||
uri_info = chunk_to_uri_map[chunk_index]
|
||||
|
||||
url_citation: ChatCompletionAnnotationURLCitation = {
|
||||
"start_index": start_index,
|
||||
|
|
|
|||
|
|
@ -9,10 +9,11 @@ The actual message transformation reuses OpenAIGPTConfig since Gemma uses OpenAI
|
|||
"""
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
from typing import TYPE_CHECKING, Final, cast
|
||||
|
||||
import httpx
|
||||
from pydantic import ValidationError
|
||||
from pydantic import ConfigDict, TypeAdapter, ValidationError
|
||||
from typing_extensions import NotRequired, ReadOnly, TypeAliasType, TypedDict
|
||||
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
|
|
@ -32,6 +33,17 @@ if TYPE_CHECKING:
|
|||
from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer
|
||||
|
||||
|
||||
class _PredictionsEnvelope(TypedDict):
|
||||
predictions: NotRequired[ReadOnly[dict[str, object]]]
|
||||
|
||||
|
||||
_PredictionsEnvelopeType = TypeAliasType("_PredictionsEnvelopeType", _PredictionsEnvelope)
|
||||
_PREDICTIONS_ENVELOPE: Final[TypeAdapter[_PredictionsEnvelope]] = TypeAdapter(
|
||||
_PredictionsEnvelopeType,
|
||||
config=ConfigDict(strict=True, extra="allow"),
|
||||
)
|
||||
|
||||
|
||||
def parse_vertex_gemma_container_error(predictions: object) -> VertexGemmaContainerError | None:
|
||||
try:
|
||||
return VertexGemmaContainerError.model_validate(predictions)
|
||||
|
|
@ -142,7 +154,7 @@ class VertexGemmaConfig(OpenAIGPTConfig):
|
|||
|
||||
def _unwrap_predictions_response(
|
||||
self,
|
||||
response_json: dict[str, Any],
|
||||
response_json: _PredictionsEnvelope,
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
Unwrap the Vertex Gemma predictions format to OpenAI format.
|
||||
|
|
@ -354,7 +366,7 @@ class VertexGemmaConfig(OpenAIGPTConfig):
|
|||
message=f"Request failed: {response.text}",
|
||||
)
|
||||
|
||||
response_json: Final = response.json()
|
||||
response_json: Final = _PREDICTIONS_ENVELOPE.validate_python(response.json())
|
||||
|
||||
# Unwrap predictions to get OpenAI-compatible response
|
||||
openai_response: Final = self._unwrap_predictions_response(response_json)
|
||||
|
|
@ -449,7 +461,7 @@ class VertexGemmaConfig(OpenAIGPTConfig):
|
|||
message=f"Request failed: {response.text}",
|
||||
)
|
||||
|
||||
response_json: Final = response.json()
|
||||
response_json: Final = _PREDICTIONS_ENVELOPE.validate_python(response.json())
|
||||
|
||||
# Unwrap predictions to get OpenAI-compatible response
|
||||
openai_response: Final = self._unwrap_predictions_response(response_json)
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ from types import UnionType
|
|||
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, Union, get_args, get_origin
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
from pydantic import fields as pyd_fields
|
||||
|
||||
import litellm
|
||||
|
|
@ -37,6 +38,8 @@ if TYPE_CHECKING:
|
|||
else:
|
||||
LiteLLMLoggingObj = Any
|
||||
|
||||
_RESPONSE_BODY: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True))
|
||||
|
||||
|
||||
class _EventModelClass(Protocol):
|
||||
@property
|
||||
|
|
@ -420,7 +423,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig):
|
|||
|
||||
@staticmethod
|
||||
def _parsed_response_body(raw_response: httpx.Response) -> dict[str, object]:
|
||||
return raw_response.json()
|
||||
return _RESPONSE_BODY.validate_python(raw_response.json())
|
||||
|
||||
@staticmethod
|
||||
def _annotation_origin(annotation: object) -> object:
|
||||
|
|
|
|||
|
|
@ -21589,7 +21589,13 @@
|
|||
],
|
||||
"supports_vision": true,
|
||||
"supports_audio_input": true,
|
||||
"supports_video_input": true
|
||||
"supports_video_input": true,
|
||||
"input_cost_per_token_priority": 5.6252e-07,
|
||||
"output_cost_per_token_priority": 3.37498e-06,
|
||||
"cache_read_input_token_cost_priority": 5.628e-08,
|
||||
"cache_creation_input_token_cost_priority": 5.6252e-07,
|
||||
"input_cost_per_audio_token_priority": 1.12497e-06,
|
||||
"cache_read_input_audio_token_cost_priority": 1.1249e-07
|
||||
},
|
||||
"databricks/databricks-gemini-3-1-flash-image": {
|
||||
"cache_creation_input_token_cost": 6.2503e-07,
|
||||
|
|
@ -21676,7 +21682,14 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_video_input": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 4.50002e-06,
|
||||
"output_cost_per_token_priority": 2.699998e-05,
|
||||
"cache_read_input_token_cost_priority": 4.5003e-07,
|
||||
"cache_creation_input_token_cost_priority": 4.50002e-06,
|
||||
"input_cost_per_token_above_200k_tokens_priority": 8.99997e-06,
|
||||
"output_cost_per_token_above_200k_tokens_priority": 4.049997e-05,
|
||||
"cache_read_input_token_cost_above_200k_tokens_priority": 8.9999e-07
|
||||
},
|
||||
"databricks/databricks-gemini-3-flash": {
|
||||
"cache_creation_input_token_cost": 6.2503e-07,
|
||||
|
|
@ -21710,7 +21723,13 @@
|
|||
],
|
||||
"supports_vision": true,
|
||||
"supports_audio_input": true,
|
||||
"supports_video_input": true
|
||||
"supports_video_input": true,
|
||||
"input_cost_per_token_priority": 1.12497e-06,
|
||||
"output_cost_per_token_priority": 6.75003e-06,
|
||||
"cache_read_input_token_cost_priority": 1.1249e-07,
|
||||
"cache_creation_input_token_cost_priority": 1.12497e-06,
|
||||
"input_cost_per_audio_token_priority": 2.25001e-06,
|
||||
"cache_read_input_audio_token_cost_priority": 2.2498e-07
|
||||
},
|
||||
"databricks/databricks-gemini-3-pro": {
|
||||
"cache_creation_input_token_cost": 2.49998e-06,
|
||||
|
|
@ -21757,7 +21776,11 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.35002e-06,
|
||||
"output_cost_per_token_priority": 6.75003e-06,
|
||||
"cache_read_input_token_cost_priority": 1.3503e-07,
|
||||
"cache_creation_input_token_cost_priority": 1.35002e-06
|
||||
},
|
||||
"databricks/databricks-gemini-3-7-flash": {
|
||||
"cache_creation_input_token_cost": 7.4998e-07,
|
||||
|
|
@ -21785,10 +21808,14 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.35002e-06,
|
||||
"output_cost_per_token_priority": 6.75003e-06,
|
||||
"cache_read_input_token_cost_priority": 1.3503e-07,
|
||||
"cache_creation_input_token_cost_priority": 1.35002e-06
|
||||
},
|
||||
"databricks/databricks-gemini-3-6-flash": {
|
||||
"cache_creation_input_token_cost": 1.87502e-06,
|
||||
"cache_creation_input_token_cost": 7.4998e-07,
|
||||
"cache_read_input_token_cost": 7.497e-08,
|
||||
"input_cost_per_token": 7.4998e-07,
|
||||
"input_dbu_cost_per_token": 1.0714e-05,
|
||||
|
|
@ -21816,7 +21843,11 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.35002e-06,
|
||||
"output_cost_per_token_priority": 6.75003e-06,
|
||||
"cache_read_input_token_cost_priority": 1.3503e-07,
|
||||
"cache_creation_input_token_cost_priority": 1.35002e-06
|
||||
},
|
||||
"databricks/databricks-gemini-3-5-flash": {
|
||||
"cache_creation_input_token_cost": 1.87502e-06,
|
||||
|
|
@ -21846,7 +21877,11 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 3.37498e-06,
|
||||
"output_cost_per_token_priority": 2.025002e-05,
|
||||
"cache_read_input_token_cost_priority": 3.3747e-07,
|
||||
"cache_creation_input_token_cost_priority": 3.37498e-06
|
||||
},
|
||||
"databricks/databricks-gemini-3-5-flash-lite": {
|
||||
"cache_creation_input_token_cost": 3.7499e-07,
|
||||
|
|
@ -21878,7 +21913,11 @@
|
|||
"supports_response_schema": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 6.7501e-07,
|
||||
"output_cost_per_token_priority": 5.62499e-06,
|
||||
"cache_read_input_token_cost_priority": 6.748e-08,
|
||||
"cache_creation_input_token_cost_priority": 6.7501e-07
|
||||
},
|
||||
"databricks/databricks-gemma-3-12b": {
|
||||
"cache_creation_input_token_cost": 1.5001e-07,
|
||||
|
|
@ -21921,7 +21960,11 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false
|
||||
"supports_vision": false,
|
||||
"input_cost_per_token_priority": 2.45e-06,
|
||||
"output_cost_per_token_priority": 7.7e-06,
|
||||
"cache_read_input_token_cost_priority": 4.55e-07,
|
||||
"cache_creation_input_token_cost_priority": 2.45e-06
|
||||
},
|
||||
"databricks/databricks-glm-5-3": {
|
||||
"cache_creation_input_token_cost": 1.4e-06,
|
||||
|
|
@ -21951,7 +21994,11 @@
|
|||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false,
|
||||
"thinking_always_on": true
|
||||
"thinking_always_on": true,
|
||||
"input_cost_per_token_priority": 2.45e-06,
|
||||
"output_cost_per_token_priority": 7.7e-06,
|
||||
"cache_read_input_token_cost_priority": 4.55e-07,
|
||||
"cache_creation_input_token_cost_priority": 2.45e-06
|
||||
},
|
||||
"databricks/databricks-glm-5-3-flash": {
|
||||
"cache_creation_input_token_cost": 1.5001e-07,
|
||||
|
|
@ -21980,7 +22027,11 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"thinking_always_on": true
|
||||
"thinking_always_on": true,
|
||||
"input_cost_per_token_priority": 2.625e-07,
|
||||
"output_cost_per_token_priority": 8.75e-07,
|
||||
"cache_read_input_token_cost_priority": 5.25e-08,
|
||||
"cache_creation_input_token_cost_priority": 2.625e-07
|
||||
},
|
||||
"databricks/databricks-gpt-5": {
|
||||
"cache_creation_input_token_cost": 1.24999e-06,
|
||||
|
|
@ -21999,7 +22050,11 @@
|
|||
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_tool_choice": true
|
||||
"supports_tool_choice": true,
|
||||
"input_cost_per_token_priority": 2.49998e-06,
|
||||
"output_cost_per_token_priority": 1.999998e-05,
|
||||
"cache_read_input_token_cost_priority": 2.4997e-07,
|
||||
"cache_creation_input_token_cost_priority": 2.49998e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-1": {
|
||||
"cache_creation_input_token_cost": 1.24999e-06,
|
||||
|
|
@ -22018,7 +22073,11 @@
|
|||
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_tool_choice": true
|
||||
"supports_tool_choice": true,
|
||||
"input_cost_per_token_priority": 2.49998e-06,
|
||||
"output_cost_per_token_priority": 1.999998e-05,
|
||||
"cache_read_input_token_cost_priority": 2.4997e-07,
|
||||
"cache_creation_input_token_cost_priority": 2.49998e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-2": {
|
||||
"cache_creation_input_token_cost": 1.75e-06,
|
||||
|
|
@ -22037,7 +22096,11 @@
|
|||
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_tool_choice": true
|
||||
"supports_tool_choice": true,
|
||||
"input_cost_per_token_priority": 3.5e-06,
|
||||
"output_cost_per_token_priority": 2.8e-05,
|
||||
"cache_read_input_token_cost_priority": 3.5e-07,
|
||||
"cache_creation_input_token_cost_priority": 3.5e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-3-codex": {
|
||||
"cache_creation_input_token_cost": 1.75e-06,
|
||||
|
|
@ -22056,7 +22119,11 @@
|
|||
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_tool_choice": true
|
||||
"supports_tool_choice": true,
|
||||
"input_cost_per_token_priority": 3.5e-06,
|
||||
"output_cost_per_token_priority": 2.8e-05,
|
||||
"cache_read_input_token_cost_priority": 3.5e-07,
|
||||
"cache_creation_input_token_cost_priority": 3.5e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-4": {
|
||||
"cache_creation_input_token_cost": 2.49998e-06,
|
||||
|
|
@ -22087,7 +22154,14 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 5.00003e-06,
|
||||
"output_cost_per_token_priority": 2.999997e-05,
|
||||
"cache_read_input_token_cost_priority": 5.0001e-07,
|
||||
"cache_creation_input_token_cost_priority": 5.00003e-06,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 5.00003e-06,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 2.999997e-05,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 5.0001e-07
|
||||
},
|
||||
"databricks/databricks-gpt-5-4-mini": {
|
||||
"cache_creation_input_token_cost": 7.4998e-07,
|
||||
|
|
@ -22115,7 +22189,11 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.50003e-06,
|
||||
"output_cost_per_token_priority": 8.99997e-06,
|
||||
"cache_read_input_token_cost_priority": 1.5001e-07,
|
||||
"cache_creation_input_token_cost_priority": 1.50003e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-4-nano": {
|
||||
"cache_creation_input_token_cost": 1.9999e-07,
|
||||
|
|
@ -22175,7 +22253,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 8.00002e-06,
|
||||
"output_cost_per_token_priority": 4.000003e-05,
|
||||
"cache_read_input_token_cost_priority": 8.0003e-07,
|
||||
"cache_creation_input_token_cost_priority": 9.99999e-06,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 1.599997e-05,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 6.000001e-05,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 1.59999e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 1.999998e-05
|
||||
},
|
||||
"databricks/databricks-gpt-5-6-terra": {
|
||||
"cache_creation_input_token_cost": 2.49998e-06,
|
||||
|
|
@ -22207,7 +22293,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 4.00001e-06,
|
||||
"output_cost_per_token_priority": 2.399999e-05,
|
||||
"cache_read_input_token_cost_priority": 3.9998e-07,
|
||||
"cache_creation_input_token_cost_priority": 5.00003e-06,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 8.00002e-06,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 3.600002e-05,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 8.0003e-07,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-6-luna": {
|
||||
"cache_creation_input_token_cost": 2.4997e-07,
|
||||
|
|
@ -22239,7 +22333,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 3.9998e-07,
|
||||
"output_cost_per_token_priority": 2.40002e-06,
|
||||
"cache_read_input_token_cost_priority": 3.997e-08,
|
||||
"cache_creation_input_token_cost_priority": 5.0001e-07,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 8.0003e-07,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 3.60003e-06,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 8.001e-08,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 1.00002e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-5": {
|
||||
"cache_creation_input_token_cost": 5.00003e-06,
|
||||
|
|
@ -22270,7 +22372,14 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.249997e-05,
|
||||
"output_cost_per_token_priority": 7.50001e-05,
|
||||
"cache_read_input_token_cost_priority": 1.24999e-06,
|
||||
"cache_creation_input_token_cost_priority": 1.249997e-05,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 1.249997e-05,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 7.50001e-05,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 1.24999e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-5-pro": {
|
||||
"cache_creation_input_token_cost": 2.999997e-05,
|
||||
|
|
@ -22319,7 +22428,11 @@
|
|||
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_tool_choice": true
|
||||
"supports_tool_choice": true,
|
||||
"input_cost_per_token_priority": 4.5003e-07,
|
||||
"output_cost_per_token_priority": 3.60003e-06,
|
||||
"cache_read_input_token_cost_priority": 4.501e-08,
|
||||
"cache_creation_input_token_cost_priority": 4.5003e-07
|
||||
},
|
||||
"databricks/databricks-gpt-5-nano": {
|
||||
"cache_creation_input_token_cost": 4.998e-08,
|
||||
|
|
@ -22469,7 +22582,11 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 5.25e-06,
|
||||
"output_cost_per_token_priority": 2.625e-05,
|
||||
"cache_read_input_token_cost_priority": 5.25e-07,
|
||||
"cache_creation_input_token_cost_priority": 5.25e-06
|
||||
},
|
||||
"databricks/system.ai.kimi-k3": {
|
||||
"litellm_provider": "databricks",
|
||||
|
|
@ -22551,7 +22668,11 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false,
|
||||
"thinking_always_on": true
|
||||
"thinking_always_on": true,
|
||||
"input_cost_per_token_priority": 3.85e-07,
|
||||
"output_cost_per_token_priority": 3.85e-06,
|
||||
"cache_read_input_token_cost_priority": 3.85e-07,
|
||||
"cache_creation_input_token_cost_priority": 3.85e-07
|
||||
},
|
||||
"databricks/databricks-qwen3-next-80b-a3b-instruct": {
|
||||
"cache_creation_input_token_cost": 1.5001e-07,
|
||||
|
|
@ -61865,6 +61986,39 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"tencent/seedream-image-v5.0-flash": {
|
||||
"litellm_provider": "tencent",
|
||||
"mode": "image_generation",
|
||||
"output_cost_per_image": 0.018,
|
||||
"output_cost_per_image_token": 1.6e-06,
|
||||
"source": "https://cloud.tencent.com/document/product/1823/130055",
|
||||
"supported_endpoints": [
|
||||
"/v1/images/generations",
|
||||
"/v1/images/edits"
|
||||
]
|
||||
},
|
||||
"tencent/seedream-image-v5.0-lite": {
|
||||
"litellm_provider": "tencent",
|
||||
"mode": "image_generation",
|
||||
"output_cost_per_image": 0.035,
|
||||
"output_cost_per_image_token": 1.6e-06,
|
||||
"source": "https://www.tencentcloud.com/document/product/1300/78937",
|
||||
"supported_endpoints": [
|
||||
"/v1/images/generations",
|
||||
"/v1/images/edits"
|
||||
]
|
||||
},
|
||||
"tencent/seedream-image-v5.0-pro": {
|
||||
"litellm_provider": "tencent",
|
||||
"mode": "image_generation",
|
||||
"output_cost_per_image": 0.09,
|
||||
"output_cost_per_image_token": 1.6e-06,
|
||||
"source": "https://www.tencentcloud.com/document/product/1300/78937",
|
||||
"supported_endpoints": [
|
||||
"/v1/images/generations",
|
||||
"/v1/images/edits"
|
||||
]
|
||||
},
|
||||
"cognition/swe-1.6": {
|
||||
"input_cost_per_token": 5e-07,
|
||||
"output_cost_per_token": 2.5e-06,
|
||||
|
|
@ -77774,7 +77928,11 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 5.25e-07,
|
||||
"output_cost_per_token_priority": 2.1e-06,
|
||||
"cache_read_input_token_cost_priority": 5.25e-08,
|
||||
"cache_creation_input_token_cost_priority": 5.25e-07
|
||||
},
|
||||
"databricks/databricks-deepseek-v4-1-flash-1017": {
|
||||
"cache_creation_input_token_cost": 3.0002e-07,
|
||||
|
|
@ -77897,7 +78055,11 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_video_input": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 4.50002e-06,
|
||||
"output_cost_per_token_priority": 2.250003e-05,
|
||||
"cache_read_input_token_cost_priority": 2.2498e-07,
|
||||
"cache_creation_input_token_cost_priority": 4.50002e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-4-pro": {
|
||||
"cache_creation_input_token_cost": 2.999997e-05,
|
||||
|
|
@ -77959,7 +78121,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 4.00001e-06,
|
||||
"output_cost_per_token_priority": 1.99997e-05,
|
||||
"cache_read_input_token_cost_priority": 1.9999e-07,
|
||||
"cache_creation_input_token_cost_priority": 5.00003e-06,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 8.00002e-06,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 2.99999e-05,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 3.9998e-07,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06
|
||||
},
|
||||
"databricks/databricks-gpt-6-1-terra": {
|
||||
"cache_creation_input_token_cost": 1.24999e-06,
|
||||
|
|
@ -77991,7 +78161,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.99997e-06,
|
||||
"output_cost_per_token_priority": 9.99999e-06,
|
||||
"cache_read_input_token_cost_priority": 1.0003e-07,
|
||||
"cache_creation_input_token_cost_priority": 2.49998e-06,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 4.00001e-06,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 1.500002e-05,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 1.9999e-07,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 5.00003e-06
|
||||
},
|
||||
"databricks/databricks-gpt-6-astra": {
|
||||
"cache_creation_input_token_cost": 1.249997e-05,
|
||||
|
|
@ -78023,7 +78201,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.999998e-05,
|
||||
"output_cost_per_token_priority": 9.99999e-05,
|
||||
"cache_read_input_token_cost_priority": 1.99997e-06,
|
||||
"cache_creation_input_token_cost_priority": 2.500001e-05,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 4.000003e-05,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 0.0001500002,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 4.00001e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 5.000002e-05
|
||||
},
|
||||
"databricks/databricks-gpt-6-luna": {
|
||||
"cache_creation_input_token_cost": 1.2502e-07,
|
||||
|
|
@ -78055,7 +78241,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.9999e-07,
|
||||
"output_cost_per_token_priority": 1.0003e-06,
|
||||
"cache_read_input_token_cost_priority": 2.002e-08,
|
||||
"cache_creation_input_token_cost_priority": 2.4997e-07,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 3.9998e-07,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 1.5001e-06,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 3.997e-08,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 5.0001e-07
|
||||
},
|
||||
"databricks/databricks-gpt-6-sol": {
|
||||
"cache_creation_input_token_cost": 2.49998e-06,
|
||||
|
|
@ -78087,7 +78281,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 4.00001e-06,
|
||||
"output_cost_per_token_priority": 1.99997e-05,
|
||||
"cache_read_input_token_cost_priority": 3.9998e-07,
|
||||
"cache_creation_input_token_cost_priority": 5.00003e-06,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 8.00002e-06,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 2.99999e-05,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 8.0003e-07,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06
|
||||
},
|
||||
"databricks/databricks-gpt-image-1": {
|
||||
"cache_creation_input_token_cost": 5.00003e-06,
|
||||
|
|
|
|||
|
|
@ -3793,8 +3793,8 @@
|
|||
"responses": true,
|
||||
"text_completion": true,
|
||||
"embeddings": false,
|
||||
"image_generations": false,
|
||||
"image_edits": false,
|
||||
"image_generations": true,
|
||||
"image_edits": true,
|
||||
"audio_transcriptions": false,
|
||||
"audio_speech": false,
|
||||
"rerank": false,
|
||||
|
|
|
|||
|
|
@ -3487,6 +3487,13 @@
|
|||
"title": "Entities",
|
||||
"type": "object"
|
||||
},
|
||||
"guardrails": {
|
||||
"additionalProperties": {
|
||||
"$ref": "#/components/schemas/GuardrailMetric"
|
||||
},
|
||||
"title": "Guardrails",
|
||||
"type": "object"
|
||||
},
|
||||
"mcp_servers": {
|
||||
"additionalProperties": {
|
||||
"$ref": "#/components/schemas/MetricWithMetadata"
|
||||
|
|
@ -3802,6 +3809,74 @@
|
|||
"title": "ExportType",
|
||||
"type": "string"
|
||||
},
|
||||
"GuardrailKeyMetric": {
|
||||
"properties": {
|
||||
"metadata": {
|
||||
"$ref": "#/components/schemas/KeyMetadata"
|
||||
},
|
||||
"metrics": {
|
||||
"$ref": "#/components/schemas/GuardrailSpendMetrics"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"metrics",
|
||||
"metadata"
|
||||
],
|
||||
"title": "GuardrailKeyMetric",
|
||||
"type": "object"
|
||||
},
|
||||
"GuardrailMetadata": {
|
||||
"properties": {
|
||||
"guardrail_name": {
|
||||
"title": "Guardrail Name",
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"guardrail_name"
|
||||
],
|
||||
"title": "GuardrailMetadata",
|
||||
"type": "object"
|
||||
},
|
||||
"GuardrailMetric": {
|
||||
"properties": {
|
||||
"api_key_breakdown": {
|
||||
"additionalProperties": {
|
||||
"$ref": "#/components/schemas/GuardrailKeyMetric"
|
||||
},
|
||||
"title": "Api Key Breakdown",
|
||||
"type": "object"
|
||||
},
|
||||
"metadata": {
|
||||
"$ref": "#/components/schemas/GuardrailMetadata"
|
||||
},
|
||||
"metrics": {
|
||||
"$ref": "#/components/schemas/GuardrailSpendMetrics"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"metrics",
|
||||
"metadata"
|
||||
],
|
||||
"title": "GuardrailMetric",
|
||||
"type": "object"
|
||||
},
|
||||
"GuardrailSpendMetrics": {
|
||||
"properties": {
|
||||
"spend": {
|
||||
"default": 0.0,
|
||||
"title": "Spend",
|
||||
"type": "number"
|
||||
},
|
||||
"untracked_units": {
|
||||
"default": 0,
|
||||
"title": "Untracked Units",
|
||||
"type": "integer"
|
||||
}
|
||||
},
|
||||
"title": "GuardrailSpendMetrics",
|
||||
"type": "object"
|
||||
},
|
||||
"HTTPAuthSecurityScheme": {
|
||||
"description": "Defines a security scheme using HTTP authentication.",
|
||||
"properties": {
|
||||
|
|
|
|||
|
|
@ -370,7 +370,7 @@ class _KeyModelBudgetLimiter(Protocol):
|
|||
|
||||
class _UserModelBudgetLimiter(Protocol):
|
||||
async def is_user_within_model_budget(
|
||||
self, user_id: str, user_model_max_budget: Mapping[str, object], model: str
|
||||
self, user_id: str, user_model_max_budget: Mapping[str, object], model: str, team_id: str | None = None
|
||||
) -> bool: ...
|
||||
|
||||
|
||||
|
|
@ -388,11 +388,15 @@ class _EveryModelBudgetLimiter(Protocol):
|
|||
async def is_key_within_model_budget(self, user_api_key_dict: UserAPIKeyAuth, model: str) -> bool: ...
|
||||
|
||||
async def is_user_within_model_budget(
|
||||
self, user_id: str, user_model_max_budget: Mapping[str, object], model: str
|
||||
self, user_id: str, user_model_max_budget: Mapping[str, object], model: str, team_id: str | None = None
|
||||
) -> bool: ...
|
||||
|
||||
async def is_end_user_within_model_budget(
|
||||
self, end_user_id: str, end_user_model_max_budget: Mapping[str, object], model: str
|
||||
self,
|
||||
end_user_id: str,
|
||||
end_user_model_max_budget: Mapping[str, object],
|
||||
model: str,
|
||||
team_id: str | None = None,
|
||||
) -> bool: ...
|
||||
|
||||
async def is_team_within_model_budget(
|
||||
|
|
@ -503,6 +507,7 @@ async def _check_user_model_budget(
|
|||
user_id=valid_token.user_id,
|
||||
user_model_max_budget=user_model_max_budget,
|
||||
model=model_name,
|
||||
team_id=valid_token.team_id,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -547,6 +552,7 @@ async def _is_within_every_model_budget(
|
|||
user_id=valid_token.user_id,
|
||||
user_model_max_budget=user_model_max_budget,
|
||||
model=model,
|
||||
team_id=valid_token.team_id,
|
||||
)
|
||||
if (
|
||||
valid_token.end_user_id is not None
|
||||
|
|
@ -557,6 +563,7 @@ async def _is_within_every_model_budget(
|
|||
end_user_id=valid_token.end_user_id,
|
||||
end_user_model_max_budget=end_user_model_max_budget,
|
||||
model=model,
|
||||
team_id=valid_token.team_id,
|
||||
)
|
||||
if valid_token.team_id is not None and isinstance(team_model_max_budget, Mapping) and team_model_max_budget:
|
||||
await model_max_budget_limiter.is_team_within_model_budget(
|
||||
|
|
@ -3154,6 +3161,7 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e
|
|||
end_user_id=valid_token.end_user_id,
|
||||
end_user_model_max_budget=end_user_mmb,
|
||||
model=model_name,
|
||||
team_id=valid_token.team_id,
|
||||
)
|
||||
|
||||
# Check 6: Additional Common Checks across jwt + key auth
|
||||
|
|
@ -4733,6 +4741,7 @@ async def _run_post_custom_auth_checks(
|
|||
end_user_id=valid_token.end_user_id,
|
||||
end_user_model_max_budget=end_user_mmb,
|
||||
model=model_name,
|
||||
team_id=valid_token.team_id,
|
||||
)
|
||||
|
||||
# team / user / end_user / project context objects are fetched by
|
||||
|
|
|
|||
|
|
@ -3,9 +3,12 @@ from collections.abc import Mapping
|
|||
from typing import Any, Final
|
||||
|
||||
import requests
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
|
||||
from .exceptions import NotFoundError, UnauthorizedError
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True))
|
||||
|
||||
|
||||
class ModelsManagementClient:
|
||||
def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30):
|
||||
|
|
@ -109,7 +112,7 @@ class ModelsManagementClient:
|
|||
try:
|
||||
response: Final = session.send(request.prepare(), timeout=self._timeout)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
return _JSON_OBJECT.validate_python(response.json())
|
||||
except requests.exceptions.HTTPError as e:
|
||||
if e.response.status_code == 401:
|
||||
raise UnauthorizedError(e)
|
||||
|
|
@ -145,7 +148,7 @@ class ModelsManagementClient:
|
|||
try:
|
||||
response: Final = session.send(request.prepare(), timeout=self._timeout)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
return _JSON_OBJECT.validate_python(response.json())
|
||||
except requests.exceptions.HTTPError as e:
|
||||
if e.response.status_code == 401:
|
||||
raise UnauthorizedError(e)
|
||||
|
|
@ -287,7 +290,7 @@ class ModelsManagementClient:
|
|||
try:
|
||||
response: Final = session.send(request.prepare(), timeout=self._timeout)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
return _JSON_OBJECT.validate_python(response.json())
|
||||
except requests.exceptions.HTTPError as e:
|
||||
if e.response.status_code == 401:
|
||||
raise UnauthorizedError(e)
|
||||
|
|
|
|||
|
|
@ -266,6 +266,8 @@ from litellm.proxy.anthropic_endpoints.streaming_model_restamp import (
|
|||
AnthropicStreamModelRestamper,
|
||||
)
|
||||
from litellm.proxy.litellm_pre_call_utils import (
|
||||
ROUTER_SCHEDULED_ROUTE_TYPES,
|
||||
SCHEDULER_PRIORITY_HEADER,
|
||||
LiteLLMProxyRequestSetup,
|
||||
add_litellm_data_to_request,
|
||||
get_metadata_variable_name,
|
||||
|
|
@ -2478,14 +2480,20 @@ class ProxyBaseLLMRequestProcessing:
|
|||
self.data["model"] = alias_target
|
||||
|
||||
aliased_model: Final = self.data.get("model")
|
||||
if llm_router is not None and route_type in ROUTER_SCHEDULED_ROUTE_TYPES and isinstance(aliased_model, str):
|
||||
self.data.update(
|
||||
llm_router.priority_header_update(
|
||||
model=aliased_model, kwargs=self.data, header=request.headers.get(SCHEDULER_PRIORITY_HEADER)
|
||||
)
|
||||
)
|
||||
|
||||
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
|
||||
self.data["litellm_call_id"] = litellm_call_id
|
||||
bind_telemetry_call_id(litellm_call_id)
|
||||
DDSpanTagger.tag_call_id(self.data.get("litellm_call_id"))
|
||||
DDSpanTagger.tag_call_id(litellm_call_id)
|
||||
DDSpanTagger.tag_request(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
requested_model=self.data.get("model"),
|
||||
requested_model=aliased_model,
|
||||
)
|
||||
|
||||
### AUTO STREAM USAGE TRACKING ###
|
||||
|
|
|
|||
269
litellm/proxy/db/accounting_ledger.py
Normal file
269
litellm/proxy/db/accounting_ledger.py
Normal file
|
|
@ -0,0 +1,269 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
from collections.abc import AsyncGenerator, Callable
|
||||
from contextlib import AbstractAsyncContextManager, asynccontextmanager
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import TYPE_CHECKING, Final, Literal, TypeAlias
|
||||
|
||||
from pydantic import AwareDatetime, ConfigDict, JsonValue, TypeAdapter
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm.accounting import CostEvent, DeliveryTarget, EventKind, InboundUsageEvent, OutboundUsageEvent
|
||||
from litellm.proxy.db.db_span import db_span
|
||||
from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import Prisma
|
||||
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
AccountingEvent: TypeAlias = InboundUsageEvent | OutboundUsageEvent | CostEvent
|
||||
_JSON: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_INSTANT: Final[TypeAdapter[datetime]] = TypeAdapter(AwareDatetime)
|
||||
_STRINGS: Final = TypeAdapter(tuple[dict[str, str], ...])
|
||||
_TIMEOUT: Final = timedelta(seconds=10)
|
||||
|
||||
|
||||
class AccountingConflictError(ValueError):
|
||||
pass
|
||||
|
||||
|
||||
_Transaction: TypeAlias = AbstractAsyncContextManager["Prisma"]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _AcceptedEvent:
|
||||
event_id: str
|
||||
kind: EventKind
|
||||
source_event_id: str | None
|
||||
occurred_at: datetime
|
||||
payload: str
|
||||
payload_hash: str
|
||||
targets: tuple[str, ...]
|
||||
schema_version: int
|
||||
version: int
|
||||
|
||||
|
||||
def _json(value: object) -> str:
|
||||
return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False)
|
||||
|
||||
|
||||
def _prepare(
|
||||
event: AccountingEvent,
|
||||
targets: tuple[DeliveryTarget, ...],
|
||||
envelope: str,
|
||||
schema_version: int,
|
||||
version: int,
|
||||
) -> _AcceptedEvent:
|
||||
if type(schema_version) is not int or type(version) is not int or min(schema_version, version) < 1:
|
||||
raise ValueError("Accounting versions must be positive integers")
|
||||
validated: Final = type(event).model_validate_json(event.model_dump_json())
|
||||
payload: Final = _JSON.validate_json(envelope)
|
||||
if payload.get("event") != validated.model_dump(mode="json"):
|
||||
raise ValueError("Accounting envelope must contain the matching public event")
|
||||
destinations: Final = tuple(sorted(frozenset(DeliveryTarget(target).value for target in targets)))
|
||||
serialized: Final = _json(payload)
|
||||
identity: Final = _json((schema_version, version, destinations, payload))
|
||||
kind: Final = (
|
||||
EventKind.COST
|
||||
if isinstance(validated, CostEvent)
|
||||
else EventKind.INBOUND_USAGE
|
||||
if isinstance(validated, InboundUsageEvent)
|
||||
else EventKind.OUTBOUND_USAGE
|
||||
)
|
||||
return _AcceptedEvent(
|
||||
event_id=validated.event_id,
|
||||
kind=kind,
|
||||
source_event_id=validated.usage_event_id if isinstance(validated, CostEvent) else None,
|
||||
occurred_at=validated.occurred_at if isinstance(validated, CostEvent) else validated.ended_at,
|
||||
payload=serialized,
|
||||
payload_hash=hashlib.sha256(identity.encode()).hexdigest(),
|
||||
targets=destinations,
|
||||
schema_version=schema_version,
|
||||
version=version,
|
||||
)
|
||||
|
||||
|
||||
_INSERT_EVENT: Final = """
|
||||
INSERT INTO "LiteLLM_AccountingEvent"
|
||||
(event_id, kind, source_event_id, schema_version, payload_hash, occurred_at, payload)
|
||||
VALUES ($1, $2::"AccountingEventKind", $3, $4::int, $5, $6::timestamp, $7::jsonb)
|
||||
ON CONFLICT (event_id) DO NOTHING RETURNING event_id
|
||||
"""
|
||||
_INSERT_DELIVERIES: Final = """
|
||||
INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version)
|
||||
SELECT $1, value::"AccountingDeliveryTarget", $3::int FROM jsonb_array_elements_text($2::jsonb)
|
||||
"""
|
||||
_LOCK_SOURCE: Final = """
|
||||
SELECT event_id FROM "LiteLLM_AccountingEvent"
|
||||
WHERE event_id = $1 AND kind = 'outbound_usage' FOR UPDATE
|
||||
"""
|
||||
_EXISTING_COST: Final = """
|
||||
SELECT event_id FROM "LiteLLM_AccountingEvent" WHERE kind = 'cost' AND source_event_id = $1
|
||||
"""
|
||||
_LOCK_EVENT: Final = """
|
||||
SELECT payload_hash FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 FOR UPDATE
|
||||
"""
|
||||
_COMPLETED_FAMILY: Final = """
|
||||
root.source_event_id IS NULL AND root.kind IN ('inbound_usage', 'outbound_usage')
|
||||
AND root.recorded_at < $1::timestamp
|
||||
AND (root.kind = 'inbound_usage' OR EXISTS (
|
||||
SELECT 1 FROM "LiteLLM_AccountingEvent" AS cost WHERE cost.source_event_id = root.event_id AND cost.kind = 'cost'
|
||||
))
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM "LiteLLM_AccountingEvent" AS member
|
||||
WHERE (member.event_id = root.event_id OR (member.kind = 'cost' AND member.source_event_id = root.event_id))
|
||||
AND (member.recorded_at >= $1::timestamp
|
||||
OR EXISTS (SELECT 1 FROM "LiteLLM_AccountingDelivery" AS delivery
|
||||
WHERE delivery.event_id = member.event_id AND delivery.state != 'applied'))
|
||||
)
|
||||
"""
|
||||
_RETENTION_CANDIDATES: Final = (
|
||||
'SELECT root.event_id FROM "LiteLLM_AccountingEvent" AS root WHERE '
|
||||
+ _COMPLETED_FAMILY
|
||||
+ " ORDER BY root.recorded_at, root.event_id LIMIT $2::int FOR UPDATE OF root SKIP LOCKED"
|
||||
)
|
||||
_ELIGIBLE_GROUP: Final = (
|
||||
'SELECT root.event_id FROM "LiteLLM_AccountingEvent" AS root WHERE root.event_id = $2 AND ' + _COMPLETED_FAMILY
|
||||
)
|
||||
_DELETE_DELIVERIES: Final = """
|
||||
DELETE FROM "LiteLLM_AccountingDelivery" WHERE event_id IN (
|
||||
SELECT event_id FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 OR (kind = 'cost' AND source_event_id = $1)
|
||||
)
|
||||
"""
|
||||
_DELETE_EVENTS: Final = """
|
||||
DELETE FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 OR (kind = 'cost' AND source_event_id = $1)
|
||||
"""
|
||||
|
||||
|
||||
_Conflict: TypeAlias = Literal["missing_source", "competing_cost", "identity_conflict"]
|
||||
|
||||
|
||||
async def _accept(db: Prisma, record: _AcceptedEvent) -> Literal["inserted", "duplicate"] | _Conflict:
|
||||
if record.source_event_id is not None:
|
||||
source: Final = _STRINGS.validate_python(await db.query_raw(_LOCK_SOURCE, record.source_event_id))
|
||||
if not source:
|
||||
return "missing_source"
|
||||
costs: Final = _STRINGS.validate_python(tuple(await db.query_raw(_EXISTING_COST, record.source_event_id)))
|
||||
if any(cost["event_id"] != record.event_id for cost in costs):
|
||||
return "competing_cost"
|
||||
inserted: Final = _STRINGS.validate_python(
|
||||
await db.query_raw(
|
||||
_INSERT_EVENT,
|
||||
record.event_id,
|
||||
record.kind.value,
|
||||
record.source_event_id,
|
||||
record.schema_version,
|
||||
record.payload_hash,
|
||||
record.occurred_at.astimezone(timezone.utc).isoformat(),
|
||||
record.payload,
|
||||
)
|
||||
)
|
||||
if not inserted:
|
||||
existing: Final = _STRINGS.validate_python(tuple(await db.query_raw(_LOCK_EVENT, record.event_id)))
|
||||
if not existing or existing[0]["payload_hash"] != record.payload_hash:
|
||||
return "identity_conflict"
|
||||
return "duplicate"
|
||||
await db.execute_raw(_INSERT_DELIVERIES, record.event_id, _json(record.targets), record.version)
|
||||
return "inserted"
|
||||
|
||||
|
||||
async def _delete_completed(db: Prisma, event_id: str, cutoff: str) -> int:
|
||||
if not _STRINGS.validate_python(await db.query_raw(_ELIGIBLE_GROUP, cutoff, event_id)):
|
||||
return 0
|
||||
await db.execute_raw(_DELETE_DELIVERIES, event_id)
|
||||
return await db.execute_raw(_DELETE_EVENTS, event_id)
|
||||
|
||||
|
||||
class AccountingLedger:
|
||||
def __init__(self, transaction: Callable[[], _Transaction]) -> None:
|
||||
self._transaction: Final = transaction
|
||||
|
||||
@classmethod
|
||||
def for_client(cls, client: PrismaClient) -> AccountingLedger:
|
||||
from prisma.client import TransactionManager
|
||||
|
||||
managers: Final = TypeAdapter(TransactionManager, config=ConfigDict(arbitrary_types_allowed=True))
|
||||
|
||||
@asynccontextmanager
|
||||
async def transaction() -> AsyncGenerator[Prisma]:
|
||||
primary: Final = writer_wrapper(client.db)
|
||||
manager: Final = managers.validate_python(primary.tx(timeout=_TIMEOUT))
|
||||
async with db_span("accounting_ledger", "LiteLLM_AccountingEvent"), manager as db:
|
||||
await db.execute_raw("SET LOCAL statement_timeout = 5000")
|
||||
await db.execute_raw("SET LOCAL lock_timeout = 1000")
|
||||
yield db
|
||||
|
||||
return cls(transaction)
|
||||
|
||||
async def insert_event(
|
||||
self,
|
||||
event: AccountingEvent,
|
||||
targets: tuple[DeliveryTarget, ...],
|
||||
*,
|
||||
envelope: str,
|
||||
schema_version: int = 1,
|
||||
version: int = 1,
|
||||
) -> Literal["inserted", "duplicate"]:
|
||||
"""Accept T5/T6's serialized envelope, containing its public event under `event`.
|
||||
|
||||
The immutable acceptance hash includes the target set and delivery version.
|
||||
Changing either for an accepted event conflicts. T7 must provide a separate
|
||||
add-target or new-version workflow that preserves receipts and prevents
|
||||
duplicate contributions; redelivery cannot silently create or reset work.
|
||||
|
||||
The Prisma adapter opens one interactive transaction per event, including
|
||||
two SET LOCAL statements and a source-row lock for costs. I/O is async,
|
||||
but awaiting this method inline adds commit latency per event, potentially
|
||||
several times per request. T5 owns bounded asynchronous capture and an
|
||||
explicit batching strategy; looping over this method does not batch commits.
|
||||
Success acknowledges commit, independently of HTTP response acknowledgement.
|
||||
"""
|
||||
record: Final = _prepare(event, targets, envelope, schema_version, version)
|
||||
async with self._transaction() as db:
|
||||
result: Final = await _accept(db, record)
|
||||
match result:
|
||||
case "inserted" | "duplicate":
|
||||
return result
|
||||
case "missing_source":
|
||||
raise AccountingConflictError("Cost source must be a durably accepted outbound usage event")
|
||||
case "competing_cost":
|
||||
raise AccountingConflictError("Usage event already has an authoritative cost")
|
||||
case "identity_conflict":
|
||||
raise AccountingConflictError("Event identity conflicts with its accepted content or delivery set")
|
||||
case _:
|
||||
assert_never(result)
|
||||
|
||||
async def purge_completed(
|
||||
self,
|
||||
*,
|
||||
redelivery_before: datetime | None = None,
|
||||
recovery_before: datetime | None = None,
|
||||
limit: int = 100,
|
||||
) -> int:
|
||||
"""Delete at most `limit` completed families once both retention owners permit it.
|
||||
|
||||
T8 must supply a coherent recovery watermark including outstanding holds.
|
||||
Omitting either watermark disables deletion. Each usage row and its cost
|
||||
remain together so a retained charge never loses its source or vice versa.
|
||||
Applied receipts are deleted with the family, ending its retained dedup
|
||||
protection. The caller must guarantee events recorded before
|
||||
redelivery_before cannot be redelivered through supported paths.
|
||||
T8's recovery_before must also protect checkpoints,
|
||||
outstanding holds and unfinished settlements; age alone is insufficient.
|
||||
"""
|
||||
if type(limit) is not int or not 1 <= limit <= 1000:
|
||||
raise ValueError("Accounting retention batch size must be between 1 and 1000")
|
||||
if redelivery_before is None or recovery_before is None:
|
||||
return 0
|
||||
cutoff: Final = min(
|
||||
_INSTANT.validate_python(redelivery_before).astimezone(timezone.utc),
|
||||
_INSTANT.validate_python(recovery_before).astimezone(timezone.utc),
|
||||
).isoformat()
|
||||
async with self._transaction() as db:
|
||||
roots: Final = _STRINGS.validate_python(tuple(await db.query_raw(_RETENTION_CANDIDATES, cutoff, limit)))
|
||||
deleted: Final = tuple([await _delete_completed(db, root["event_id"], cutoff) for root in roots])
|
||||
return sum(deleted)
|
||||
|
|
@ -64,6 +64,16 @@ _ROLL_WINDOW_SPEND_SQL: Final = (
|
|||
"AND window_start < ($4::timestamptz AT TIME ZONE 'UTC')"
|
||||
)
|
||||
|
||||
_OPEN_WINDOW_SPEND_SQL: Final = (
|
||||
'INSERT INTO "LiteLLM_BudgetWindowSpend" '
|
||||
"(entity_type, entity_id, window_duration, window_start, spend, created_at, updated_at) "
|
||||
"VALUES ($1, $2, $3, ($4::timestamptz AT TIME ZONE 'UTC'), 0, "
|
||||
"($4::timestamptz AT TIME ZONE 'UTC'), ($4::timestamptz AT TIME ZONE 'UTC')) "
|
||||
"ON CONFLICT (entity_type, entity_id, window_duration) DO UPDATE SET "
|
||||
"window_start = EXCLUDED.window_start, spend = 0, updated_at = EXCLUDED.updated_at "
|
||||
"WHERE \"LiteLLM_BudgetWindowSpend\".window_start <= ($5::timestamptz AT TIME ZONE 'UTC')"
|
||||
)
|
||||
|
||||
_SEED_FROM_SPEND_LOGS_KEY_SQL: Final = (
|
||||
"SELECT COALESCE(SUM(spend), 0.0) AS total, "
|
||||
"COALESCE(SUM(spend) FILTER (WHERE \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')), 0.0) AS before_batch "
|
||||
|
|
@ -389,3 +399,29 @@ async def roll_window_spend_row(
|
|||
to_naive_utc(new_window_start),
|
||||
to_naive_utc(datetime.now(timezone.utc)),
|
||||
)
|
||||
|
||||
|
||||
async def open_window_spend_row(
|
||||
prisma_client: "PrismaClient",
|
||||
entity_type: str,
|
||||
entity_id: str,
|
||||
window_duration: str,
|
||||
window_start: datetime,
|
||||
replaces_window_start: datetime | None,
|
||||
) -> None:
|
||||
"""Start a window at `window_start` unless another writer already started one.
|
||||
|
||||
The row is created when it does not exist, and moved onto the new window only
|
||||
while it still holds `replaces_window_start`, the ended window the caller saw.
|
||||
A writer that lost the race leaves the winner's window in place, so reading
|
||||
the row back gives every writer the same window start.
|
||||
"""
|
||||
async with db_span("open_window_spend_row", "LiteLLM_BudgetWindowSpend"):
|
||||
await prisma_client.db.execute_raw(
|
||||
_OPEN_WINDOW_SPEND_SQL,
|
||||
entity_type,
|
||||
entity_id,
|
||||
window_duration,
|
||||
to_naive_utc(window_start),
|
||||
None if replaces_window_start is None else to_naive_utc(replaces_window_start),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -305,6 +305,11 @@ def iter_message_text(data: Mapping[str, object]) -> Iterator[str]:
|
|||
yield from _iter_text_parts_in_content(message.get("content"))
|
||||
|
||||
|
||||
def _is_tool_message(message: object) -> bool:
|
||||
fields: Final = _text_object(message)
|
||||
return fields is not None and str(fields.get("role") or "").lower() == "tool"
|
||||
|
||||
|
||||
def _call_item(item: object) -> tuple[str, object] | None:
|
||||
"""Return ``(field, value)`` for a Responses tool-call input item, else None."""
|
||||
fields: Final = _text_object(item)
|
||||
|
|
@ -316,11 +321,14 @@ def walk_user_text(
|
|||
data: dict[str, object],
|
||||
visit: Callable[[str], str],
|
||||
visit_tool: Callable[[JsonValue], JsonValue] | None = None,
|
||||
*,
|
||||
skip_tool_message: bool = False,
|
||||
) -> int:
|
||||
"""Rewrite every text fragment in place via ``visit``.
|
||||
|
||||
Mutates ``data["messages"]`` and ``data["input"]`` (including tool payloads), ``data["system"]`` and
|
||||
``data["instructions"]``. Returns the number
|
||||
``data["instructions"]``. ``skip_tool_message`` leaves ``role: tool`` messages and ``tool_result``
|
||||
blocks as sent. Returns the number
|
||||
of fragments visited so callers can short-circuit when nothing was
|
||||
inspected.
|
||||
"""
|
||||
|
|
@ -353,6 +361,8 @@ def walk_user_text(
|
|||
|
||||
def _rewrite_part(part: Mapping[str, object], native: bool) -> Mapping[str, object]:
|
||||
nonlocal visited
|
||||
if skip_tool_message and part.get("type") == "tool_result":
|
||||
return part
|
||||
text: Final = _part_text(part)
|
||||
rewritten: Final = {**part, **({"text": visit(text)} if text is not None else {})}
|
||||
if text is not None:
|
||||
|
|
@ -384,18 +394,17 @@ def walk_user_text(
|
|||
return content
|
||||
|
||||
messages: Final = data.get("messages")
|
||||
if isinstance(messages, list):
|
||||
for message in messages:
|
||||
if not isinstance(message, dict):
|
||||
continue
|
||||
if "content" in message:
|
||||
message["content"] = _rewrite_content(message["content"])
|
||||
calls = message.get("tool_calls")
|
||||
if isinstance(calls, list):
|
||||
for call in calls:
|
||||
function = call.get("function") if isinstance(call, dict) else None
|
||||
if isinstance(function, dict) and "arguments" in function:
|
||||
function["arguments"] = _rewrite_arguments(function["arguments"])
|
||||
for message in _object_list(messages) or ():
|
||||
if (skip_tool_message and _is_tool_message(message)) or not isinstance(message, dict):
|
||||
continue
|
||||
if "content" in message:
|
||||
message["content"] = _rewrite_content(message["content"])
|
||||
calls = message.get("tool_calls")
|
||||
if isinstance(calls, list):
|
||||
for call in calls:
|
||||
function = call.get("function") if isinstance(call, dict) else None
|
||||
if isinstance(function, dict) and "arguments" in function:
|
||||
function["arguments"] = _rewrite_arguments(function["arguments"])
|
||||
|
||||
if "system" in data:
|
||||
data["system"] = _rewrite_content(data["system"]) # rebind-ok: preserve the native hook's in-place contract
|
||||
|
|
|
|||
|
|
@ -169,7 +169,7 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail):
|
|||
|
||||
from litellm.types.llms.openai import OpenAIModerationResponse
|
||||
|
||||
return OpenAIModerationResponse(**response.json())
|
||||
return OpenAIModerationResponse.model_validate(response.json())
|
||||
|
||||
def _check_moderation_result(self, moderation_response: "OpenAIModerationResponse") -> None:
|
||||
"""
|
||||
|
|
|
|||
176
litellm/proxy/hooks/model_budget_shared_windows.py
Normal file
176
litellm/proxy/hooks/model_budget_shared_windows.py
Normal file
|
|
@ -0,0 +1,176 @@
|
|||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING, Final, Protocol, TypeAlias
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.proxy.db.budget_window_spend_writer import open_window_spend_row
|
||||
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
|
||||
WindowSpendTransaction,
|
||||
build_window_spend_transaction,
|
||||
)
|
||||
from litellm.repositories.table_repositories import BudgetWindowSpendRepository
|
||||
from litellm.router_strategy.budget_limiter import budget_window_seconds
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
MODEL_BUDGET_WINDOW_ENTITY_TYPE: Final = "model_max_budget"
|
||||
_READ_TTL_SECONDS: Final = 5
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SharedWindow:
|
||||
started_at: float
|
||||
spend: float
|
||||
|
||||
def is_open(self, budget_duration: str, now: float) -> bool:
|
||||
return now < self.started_at + budget_window_seconds(self.started_at, budget_duration)
|
||||
|
||||
|
||||
class LocalTTLCache(Protocol):
|
||||
def get_cache(self, key: str) -> object: ...
|
||||
|
||||
def set_cache(self, key: str, value: object, *, ttl: float) -> None: ...
|
||||
|
||||
def delete_cache(self, key: str) -> None: ...
|
||||
|
||||
|
||||
class SharedWindows(Protocol):
|
||||
async def read(self, spend_key: str, budget_duration: str) -> SharedWindow | None: ...
|
||||
|
||||
async def open(self, spend_key: str, budget_duration: str, now: float) -> float | None: ...
|
||||
|
||||
async def add(self, spend_key: str, budget_duration: str, started_at: float, spend: float) -> None: ...
|
||||
|
||||
|
||||
PrismaClientSource: TypeAlias = Callable[[], "PrismaClient | None"]
|
||||
WindowSpendSink: TypeAlias = Callable[[WindowSpendTransaction], Awaitable[None]]
|
||||
|
||||
|
||||
def _proxy_prisma_client() -> "PrismaClient | None":
|
||||
try:
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
except ImportError:
|
||||
return None
|
||||
return prisma_client
|
||||
|
||||
|
||||
async def _queue_on_proxy_spend_writer(transaction: WindowSpendTransaction) -> None:
|
||||
from litellm.proxy.proxy_server import proxy_logging_obj
|
||||
|
||||
await proxy_logging_obj.db_spend_update_writer.window_spend_update_queue.add_update(transaction)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Reading:
|
||||
window: SharedWindow | None
|
||||
|
||||
|
||||
class DatabaseSharedWindows:
|
||||
def __init__(
|
||||
self,
|
||||
prisma_client_source: PrismaClientSource = _proxy_prisma_client,
|
||||
sink: WindowSpendSink = _queue_on_proxy_spend_writer,
|
||||
read_ttl_seconds: int = _READ_TTL_SECONDS,
|
||||
) -> None:
|
||||
self._prisma_client_source = prisma_client_source
|
||||
self._sink = sink
|
||||
self._read_ttl_seconds = read_ttl_seconds
|
||||
self._readings: LocalTTLCache = InMemoryCache(max_size_in_memory=10_000, default_ttl=read_ttl_seconds)
|
||||
|
||||
async def read(self, spend_key: str, budget_duration: str) -> SharedWindow | None:
|
||||
cached: Final = self._readings.get_cache(key=spend_key)
|
||||
if isinstance(cached, _Reading):
|
||||
return cached.window
|
||||
window: Final = await self._read_row(spend_key=spend_key, budget_duration=budget_duration)
|
||||
self._readings.set_cache(key=spend_key, value=_Reading(window=window), ttl=self._read_ttl_seconds)
|
||||
return window
|
||||
|
||||
async def open(self, spend_key: str, budget_duration: str, now: float) -> float | None:
|
||||
"""When the window every replica shares started, opening one at `now` if none is open.
|
||||
|
||||
Without a database every window is this replica's own, so it starts at `now`. None means the database did
|
||||
not answer, so this replica cannot know which window the others are in yet.
|
||||
"""
|
||||
prisma_client: Final = self._prisma_client_source()
|
||||
if prisma_client is None:
|
||||
return now
|
||||
try:
|
||||
seen: Final = await _fetch_row(
|
||||
prisma_client=prisma_client, spend_key=spend_key, budget_duration=budget_duration, use_writer=True
|
||||
)
|
||||
if seen is not None and seen.is_open(budget_duration=budget_duration, now=now):
|
||||
return seen.started_at
|
||||
await open_window_spend_row(
|
||||
prisma_client=prisma_client,
|
||||
entity_type=MODEL_BUDGET_WINDOW_ENTITY_TYPE,
|
||||
entity_id=spend_key,
|
||||
window_duration=budget_duration,
|
||||
window_start=datetime.fromtimestamp(now, tz=timezone.utc),
|
||||
replaces_window_start=None
|
||||
if seen is None
|
||||
else datetime.fromtimestamp(seen.started_at, tz=timezone.utc),
|
||||
)
|
||||
opened: Final = await _fetch_row(
|
||||
prisma_client=prisma_client, spend_key=spend_key, budget_duration=budget_duration, use_writer=True
|
||||
)
|
||||
except Exception: # noqa: BLE001 # an unreachable database leaves this replica on its own window
|
||||
verbose_proxy_logger.exception("Unable to open the shared model budget window")
|
||||
return None
|
||||
if opened is None or not opened.is_open(budget_duration=budget_duration, now=now):
|
||||
return None
|
||||
return opened.started_at
|
||||
|
||||
async def add(self, spend_key: str, budget_duration: str, started_at: float, spend: float) -> None:
|
||||
if self._prisma_client_source() is None:
|
||||
return
|
||||
try:
|
||||
await self._sink(
|
||||
build_window_spend_transaction(
|
||||
entity_type=MODEL_BUDGET_WINDOW_ENTITY_TYPE,
|
||||
entity_id=spend_key,
|
||||
window_duration=budget_duration,
|
||||
window_start=datetime.fromtimestamp(started_at, tz=timezone.utc),
|
||||
spend=spend,
|
||||
)
|
||||
)
|
||||
except Exception: # noqa: BLE001 # spend tracking must never fail the success callback
|
||||
verbose_proxy_logger.debug("Unable to queue shared model budget spend", exc_info=True)
|
||||
|
||||
async def _read_row(self, spend_key: str, budget_duration: str) -> SharedWindow | None:
|
||||
prisma_client: Final = self._prisma_client_source()
|
||||
if prisma_client is None:
|
||||
return None
|
||||
try:
|
||||
return await _fetch_row(
|
||||
prisma_client=prisma_client, spend_key=spend_key, budget_duration=budget_duration, use_writer=False
|
||||
)
|
||||
except Exception: # noqa: BLE001 # an unreadable row leaves enforcement on this replica's own counter
|
||||
verbose_proxy_logger.exception("Unable to read shared model budget spend")
|
||||
return None
|
||||
|
||||
|
||||
async def _fetch_row(
|
||||
prisma_client: "PrismaClient", spend_key: str, budget_duration: str, use_writer: bool
|
||||
) -> SharedWindow | None:
|
||||
row: Final = await BudgetWindowSpendRepository(prisma_client, use_writer=use_writer).table.find_unique(
|
||||
where={
|
||||
"entity_type_entity_id_window_duration": {
|
||||
"entity_type": MODEL_BUDGET_WINDOW_ENTITY_TYPE,
|
||||
"entity_id": spend_key,
|
||||
"window_duration": budget_duration,
|
||||
}
|
||||
}
|
||||
)
|
||||
if row is None:
|
||||
return None
|
||||
window_start: Final = row.window_start
|
||||
started_at: Final = (
|
||||
window_start if window_start.tzinfo is not None else window_start.replace(tzinfo=timezone.utc)
|
||||
).timestamp()
|
||||
return SharedWindow(started_at=started_at, spend=float(row.spend or 0.0))
|
||||
|
||||
|
||||
PROXY_SHARED_WINDOWS: Final = DatabaseSharedWindows()
|
||||
File diff suppressed because it is too large
Load diff
|
|
@ -118,6 +118,9 @@ _SESSION_ID_VALUE_RE: Final = re.compile(r"^[a-zA-Z0-9_\-]{8,}$")
|
|||
|
||||
_SHA256_HEX_RE: Final = re.compile(r"^[0-9a-f]{64}$")
|
||||
|
||||
SCHEDULER_PRIORITY_HEADER: Final = "x-litellm-priority"
|
||||
ROUTER_SCHEDULED_ROUTE_TYPES: Final = frozenset({"acompletion", "atext_completion"})
|
||||
|
||||
# W3C Trace Context traceparent header: https://www.w3.org/TR/trace-context/
|
||||
# e.g. "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
|
||||
_TRACEPARENT_RE: Final = re.compile(r"^[0-9a-f]{2}-([0-9a-f]{32})-[0-9a-f]{16}-[0-9a-f]{2}$", re.IGNORECASE)
|
||||
|
|
@ -688,7 +691,7 @@ def _metadata_string_tags(metadata_objects: Iterable[Mapping[str, object]]) -> t
|
|||
|
||||
|
||||
def _strip_untrusted_request_header_controls(
|
||||
headers: Any,
|
||||
headers: object,
|
||||
*,
|
||||
allow_client_message_redaction_opt_out: bool = False,
|
||||
) -> None:
|
||||
|
|
@ -1692,7 +1695,7 @@ class LiteLLMProxyRequestSetup:
|
|||
|
||||
@staticmethod
|
||||
def add_internal_user_from_user_mapping(
|
||||
general_settings: dict | None,
|
||||
general_settings: Mapping[str, object] | None,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
headers: dict,
|
||||
) -> UserAPIKeyAuth:
|
||||
|
|
@ -1711,7 +1714,7 @@ class LiteLLMProxyRequestSetup:
|
|||
return user_api_key_dict
|
||||
|
||||
@staticmethod
|
||||
def get_user_from_headers(headers: dict, general_settings: dict | None = None) -> str | None:
|
||||
def get_user_from_headers(headers: dict, general_settings: Mapping[str, object] | None = None) -> str | None:
|
||||
"""
|
||||
Get the user from the specified header if `general_settings.user_header_name` is set.
|
||||
"""
|
||||
|
|
@ -1732,7 +1735,9 @@ class LiteLLMProxyRequestSetup:
|
|||
return user
|
||||
|
||||
@staticmethod
|
||||
def get_openai_org_id_from_headers(headers: dict, general_settings: dict | None = None) -> str | None:
|
||||
def get_openai_org_id_from_headers(
|
||||
headers: dict, general_settings: Mapping[str, object] | None = None
|
||||
) -> str | None:
|
||||
"""
|
||||
Get the OpenAI Org ID from the headers.
|
||||
"""
|
||||
|
|
@ -1823,7 +1828,7 @@ class LiteLLMProxyRequestSetup:
|
|||
headers: dict,
|
||||
request_data: Mapping[str, object],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
general_settings: dict[str, object] | None = None,
|
||||
general_settings: Mapping[str, object] | None = None,
|
||||
) -> LitellmDataForBackendLLMCall:
|
||||
"""
|
||||
- Adds user from headers
|
||||
|
|
|
|||
|
|
@ -231,9 +231,12 @@ def _shadow_eval_attempts(prisma_client: "PrismaClient") -> _ShadowEvalAttemptTa
|
|||
return _ShadowEvalAttemptRepository(prisma_client).table
|
||||
|
||||
|
||||
_RAW_ROWS: Final = TypeAdapter(Sequence[dict[str, object]], config=ConfigDict(strict=True))
|
||||
|
||||
|
||||
async def _query_raw(prisma_client: "PrismaClient", query: str, *args: object) -> Sequence[Mapping[str, object]]:
|
||||
async with db_span("auto_router_report_query", sql_relation(query)):
|
||||
return await prisma_client.db.query_raw(query, *args)
|
||||
return _RAW_ROWS.validate_python(await prisma_client.db.query_raw(query, *args))
|
||||
|
||||
|
||||
async def _execute_raw(prisma_client: "PrismaClient", query: str, *args: object) -> int:
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ from collections.abc import Awaitable, Callable, Mapping, Sequence
|
|||
from collections.abc import Set as AbstractSet
|
||||
from dataclasses import dataclass, replace
|
||||
from datetime import date, datetime, timedelta
|
||||
from itertools import groupby
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal, NoReturn, Protocol
|
||||
|
||||
|
|
@ -28,6 +29,10 @@ from litellm.types.proxy.management_endpoints.common_daily_activity import (
|
|||
DailySpendData,
|
||||
DailySpendMetadata,
|
||||
GroupedData,
|
||||
GuardrailKeyMetric,
|
||||
GuardrailMetadata,
|
||||
GuardrailMetric,
|
||||
GuardrailSpendMetrics,
|
||||
KeyMetadata,
|
||||
KeyMetricWithMetadata,
|
||||
MetricWithMetadata,
|
||||
|
|
@ -39,6 +44,7 @@ from litellm.types.repositories.daily_activity import (
|
|||
DailyActivityTable,
|
||||
EntityRollupRow,
|
||||
GroupingSetsRow,
|
||||
GuardrailSpendRow,
|
||||
KeyMetadataRow,
|
||||
RollupMetricsRow,
|
||||
SpendLogsWindow,
|
||||
|
|
@ -1013,17 +1019,68 @@ def _fold_entity_rollups_sync(
|
|||
)
|
||||
|
||||
|
||||
def _guardrail_metric(day: DailySpendData, rows: Sequence[GuardrailSpendRow]) -> GuardrailMetric:
|
||||
return GuardrailMetric(
|
||||
metrics=GuardrailSpendMetrics(
|
||||
spend=sum(row.spend for row in rows),
|
||||
untracked_units=sum(row.untracked_units for row in rows),
|
||||
),
|
||||
metadata=GuardrailMetadata(guardrail_name=rows[0].guardrail_name),
|
||||
api_key_breakdown={
|
||||
row.api_key: GuardrailKeyMetric(
|
||||
metrics=GuardrailSpendMetrics(spend=row.spend, untracked_units=row.untracked_units),
|
||||
metadata=day.breakdown.api_keys[row.api_key].metadata,
|
||||
)
|
||||
for row in rows
|
||||
if row.api_key in day.breakdown.api_keys
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _day_with_guardrails(day: DailySpendData, rows: Sequence[GuardrailSpendRow]) -> DailySpendData:
|
||||
by_guardrail: Final = groupby(sorted(rows, key=lambda row: row.guardrail_id), key=lambda row: row.guardrail_id)
|
||||
guardrails: Final = {guardrail_id: _guardrail_metric(day, tuple(group)) for guardrail_id, group in by_guardrail}
|
||||
return day.model_copy(update={"breakdown": day.breakdown.model_copy(update={"guardrails": guardrails})})
|
||||
|
||||
|
||||
def _with_guardrail_breakdown(
|
||||
results: Sequence[DailySpendData], rows: Sequence[GuardrailSpendRow]
|
||||
) -> tuple[DailySpendData, ...]:
|
||||
rows_by_date: Final = MappingProxyType(
|
||||
{day: tuple(group) for day, group in groupby(sorted(rows, key=lambda row: row.date), key=lambda row: row.date)}
|
||||
)
|
||||
return tuple(_day_with_guardrails(day, rows_by_date.get(day.date.strftime("%Y-%m-%d"), ())) for day in results)
|
||||
|
||||
|
||||
async def _guardrail_spend_or_empty(
|
||||
repository: DailyActivityRepository, scope: DailyActivityScope
|
||||
) -> tuple[GuardrailSpendRow, ...]:
|
||||
from prisma.errors import PrismaError
|
||||
|
||||
try:
|
||||
return await repository.guardrail_spend(scope)
|
||||
except PrismaError as exc:
|
||||
verbose_proxy_logger.warning("Returning daily activity without its guardrail breakdown: %s", exc)
|
||||
return ()
|
||||
|
||||
|
||||
async def get_daily_activity_aggregated(
|
||||
repository: DailyActivityRepository,
|
||||
scope: DailyActivityScope,
|
||||
*,
|
||||
entity_metadata_field: Mapping[str, dict[str, object]] | None = None,
|
||||
include_entity_breakdown: bool = False,
|
||||
include_guardrail_breakdown: bool = False,
|
||||
api_key_limit: int = constants.USAGE_TOP_API_KEYS_DEFAULT,
|
||||
) -> SpendAnalyticsPaginatedResponse:
|
||||
try:
|
||||
aggregated_rows: Final = await repository.aggregated(
|
||||
scope, include_entity_breakdown=include_entity_breakdown, api_key_limit=api_key_limit
|
||||
aggregated_rows, guardrail_rows = await asyncio.gather(
|
||||
repository.aggregated(
|
||||
scope, include_entity_breakdown=include_entity_breakdown, api_key_limit=api_key_limit
|
||||
),
|
||||
_guardrail_spend_or_empty(repository, scope)
|
||||
if include_guardrail_breakdown and not scope.model
|
||||
else asyncio.sleep(0, result=()),
|
||||
)
|
||||
records: Final = aggregated_rows.grouping_rows
|
||||
aggregated: Final = await _aggregate_grouping_sets_records(
|
||||
|
|
@ -1058,8 +1115,13 @@ async def get_daily_activity_aggregated(
|
|||
api_key_metadata=entity_key_metadata,
|
||||
entity_metadata_field=entity_metadata_field,
|
||||
)
|
||||
results: Final = (
|
||||
list(await asyncio.to_thread(_with_guardrail_breakdown, aggregated["results"], guardrail_rows))
|
||||
if guardrail_rows
|
||||
else aggregated["results"]
|
||||
)
|
||||
return SpendAnalyticsPaginatedResponse(
|
||||
results=aggregated["results"],
|
||||
results=results,
|
||||
metadata=DailySpendMetadata(
|
||||
total_spend=aggregated["totals"].spend,
|
||||
total_flat_cost=aggregated["totals"].flat_cost,
|
||||
|
|
|
|||
|
|
@ -300,6 +300,7 @@ def _register_aggregated_route(router: APIRouter, resolver: EntityScopeResolver,
|
|||
resolved.scope,
|
||||
entity_metadata_field=resolved.entity_metadata,
|
||||
include_entity_breakdown=resolver.include_entity_breakdown,
|
||||
include_guardrail_breakdown=resolver.include_guardrail_breakdown,
|
||||
api_key_limit=api_key_limit,
|
||||
)
|
||||
if resolver.entity != "team":
|
||||
|
|
|
|||
|
|
@ -56,6 +56,7 @@ class EntityScopeResolver:
|
|||
resolve: EntityScopeResolve
|
||||
include_entity_breakdown: bool
|
||||
operation_names: OperationNames
|
||||
include_guardrail_breakdown: bool = False
|
||||
|
||||
|
||||
def _query_ids(value: str | None) -> tuple[str, ...] | None:
|
||||
|
|
@ -364,6 +365,7 @@ USER_RESOLVER = EntityScopeResolver(
|
|||
query=_user_query,
|
||||
resolve=_resolve_user,
|
||||
include_entity_breakdown=False,
|
||||
include_guardrail_breakdown=True,
|
||||
operation_names=MappingProxyType(
|
||||
{
|
||||
"aggregated": "get_user_daily_activity_aggregated",
|
||||
|
|
@ -385,6 +387,7 @@ TEAM_RESOLVER = EntityScopeResolver(
|
|||
query=_team_query,
|
||||
resolve=_resolve_team,
|
||||
include_entity_breakdown=True,
|
||||
include_guardrail_breakdown=True,
|
||||
operation_names=MappingProxyType(
|
||||
{
|
||||
"aggregated": "get_team_daily_activity_aggregated",
|
||||
|
|
@ -423,6 +426,7 @@ ORGANIZATION_RESOLVER = EntityScopeResolver(
|
|||
query=_organization_query,
|
||||
resolve=_resolve_organization,
|
||||
include_entity_breakdown=True,
|
||||
include_guardrail_breakdown=True,
|
||||
operation_names=MappingProxyType(
|
||||
{
|
||||
"aggregated": "get_organization_daily_activity_aggregated",
|
||||
|
|
|
|||
|
|
@ -499,7 +499,7 @@ async def new_user(
|
|||
- permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking.
|
||||
- metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }
|
||||
- max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x.
|
||||
- model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys)
|
||||
- model_max_budget: Optional[dict] - Model-specific max budget for user, keyed by model name, `*` pattern, or model access group. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys)
|
||||
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
- model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)
|
||||
- mcp_rpm_limit: Optional[dict] - Per-MCP-server rpm limit, keyed by MCP server name {"github": 100, "slack": 200}. Enforced for keys and teams only; values set on a user are stored but not enforced per user.
|
||||
|
|
@ -1702,7 +1702,7 @@ async def user_update(
|
|||
- permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking.
|
||||
- metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }
|
||||
- max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x.
|
||||
- model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys)
|
||||
- model_max_budget: Optional[dict] - Model-specific max budget for user, keyed by model name, `*` pattern, or model access group. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys)
|
||||
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
- model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)
|
||||
- mcp_rpm_limit: Optional[dict] - Per-MCP-server rpm limit, keyed by MCP server name {"github": 100, "slack": 200}. Enforced for keys and teams only; values set on a user are stored but not enforced per user.
|
||||
|
|
|
|||
|
|
@ -2328,7 +2328,7 @@ async def generate_key_fn(
|
|||
- throttle_on_budget_exceeded: Optional[bool] - When the key exceeds its max_budget, throttle its tpm/rpm to the global budget_exceeded_throttle_percentage instead of blocking the key entirely.
|
||||
- enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Supported Claude models on Anthropic, Bedrock, Vertex AI, and Azure AI only.
|
||||
- permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false}
|
||||
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
|
||||
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets keyed by model name, `*` pattern like "claude-opus*", or model access group, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
|
||||
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
- model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit.
|
||||
- model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit.
|
||||
|
|
@ -2538,7 +2538,7 @@ async def generate_service_account_key_fn(
|
|||
- metadata: Optional[dict] - Metadata for key, store information for key. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }
|
||||
- guardrails: Optional[List[str]] - List of active guardrails for the key
|
||||
- permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false}
|
||||
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
|
||||
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets keyed by model name, `*` pattern like "claude-opus*", or model access group, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
|
||||
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
- model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit.
|
||||
- model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit.
|
||||
|
|
@ -3855,7 +3855,7 @@ async def update_key_fn(
|
|||
- enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests)
|
||||
- spend: Optional[float] - Amount spent by key
|
||||
- max_budget: Optional[float] - Max budget for key
|
||||
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
|
||||
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets keyed by model name, `*` pattern like "claude-opus*", or model access group, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
|
||||
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
- budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.)
|
||||
- soft_budget: Optional[float] - Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached. Set to null to remove the soft budget.
|
||||
|
|
@ -6303,7 +6303,7 @@ async def regenerate_key_fn(
|
|||
- tags: Optional[List[str]] - Tags for organizing keys (Enterprise only)
|
||||
- spend: Optional[float] - Amount spent by key
|
||||
- max_budget: Optional[float] - Max budget for key
|
||||
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
|
||||
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets keyed by model name, `*` pattern like "claude-opus*", or model access group, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
|
||||
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
- budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.)
|
||||
- soft_budget: Optional[float] - Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached.
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ GET /router/fields - Get router settings field definitions without values (for U
|
|||
import inspect
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, get_args
|
||||
from typing import Final, get_args
|
||||
|
||||
from fastapi import APIRouter, Depends
|
||||
from pydantic import Field
|
||||
|
|
@ -32,7 +32,7 @@ router: Final = APIRouter()
|
|||
|
||||
class RouterSettingsResponse(LiteLLMBaseModel):
|
||||
fields: list[RouterSettingsField] = Field(description="List of all configurable router settings with metadata")
|
||||
current_values: dict[str, Any] = Field(description="Current values of router settings")
|
||||
current_values: dict[str, object] = Field(description="Current values of router settings")
|
||||
routing_strategy_descriptions: dict[str, str] = Field(description="Descriptions for each routing strategy option")
|
||||
routing_group_strategies: tuple[str, ...] = Field(
|
||||
description="Strategies supported when constructing a routing group"
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ All /team management endpoints
|
|||
|
||||
import asyncio
|
||||
import copy
|
||||
import itertools
|
||||
import json
|
||||
import math
|
||||
import traceback
|
||||
|
|
@ -139,7 +140,12 @@ from litellm.proxy.common_utils.user_api_key_cache import (
|
|||
from litellm.proxy.db.db_span import db_span
|
||||
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
|
||||
from litellm.proxy.hooks.model_max_budget_limiter import (
|
||||
RoutedModels,
|
||||
build_model_max_budget_usage,
|
||||
entries_outranked_by,
|
||||
entry_names_model,
|
||||
is_enforced_budget_entry,
|
||||
proxy_routed_models,
|
||||
resolve_model_budget,
|
||||
)
|
||||
from litellm.proxy.management.teams.authz import TEAM_OR_ORG_ADMIN, TeamRole, is_team_admin, team_access_denied
|
||||
|
|
@ -1467,8 +1473,10 @@ def _check_team_model_budget_update_authority(
|
|||
data: UpdateTeamRequest,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
existing_model_max_budget: Mapping[str, object] | None,
|
||||
routed_models_source: Callable[[str | None], RoutedModels] = proxy_routed_models,
|
||||
) -> None:
|
||||
"""Like `_check_team_budget_update_authority`: only a proxy admin may raise, re-window or drop a per-model cap."""
|
||||
"""Like `_check_team_budget_update_authority`: only a proxy admin may raise, re-window or drop a per-model cap, or
|
||||
add an entry that takes precedence over one."""
|
||||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
|
||||
return
|
||||
if "model_max_budget" not in data.model_fields_set or not existing_model_max_budget:
|
||||
|
|
@ -1487,29 +1495,97 @@ def _check_team_model_budget_update_authority(
|
|||
)
|
||||
},
|
||||
)
|
||||
for model_name, proposed in requested.items():
|
||||
governing = resolve_model_budget(model=model_name, model_max_budget=existing_model_max_budget)
|
||||
if governing is None:
|
||||
continue
|
||||
cap = governing.budget_config
|
||||
if cap.max_budget is None:
|
||||
continue
|
||||
if (
|
||||
proposed.max_budget is None
|
||||
or proposed.max_budget > cap.max_budget
|
||||
or proposed.budget_duration != cap.budget_duration
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": (
|
||||
f"Only a proxy admin can raise a team's model_max_budget for {model_name!r} or change its "
|
||||
f"budget_duration. Current max_budget={cap.max_budget} per {cap.budget_duration} "
|
||||
f"(entry {governing.budget_model!r}), requested={proposed.max_budget} per "
|
||||
f"{proposed.budget_duration}."
|
||||
)
|
||||
},
|
||||
routed_models: Final = routed_models_source(data.team_id)
|
||||
for entry, proposed in requested.items():
|
||||
existing = _existing_model_cap(existing_model_max_budget.get(entry))
|
||||
if existing is not None:
|
||||
_check_existing_model_cap_kept(entry=entry, cap_entry=entry, proposed=proposed, existing=existing)
|
||||
if is_enforced_budget_entry(existing) or not is_enforced_budget_entry(proposed):
|
||||
continue
|
||||
governing = resolve_model_budget(model=entry, model_max_budget=existing_model_max_budget, team_id=data.team_id)
|
||||
if governing is not None and entry_names_model(entry=governing.budget_model, model=entry):
|
||||
_check_existing_model_cap_kept(
|
||||
entry=entry, cap_entry=governing.budget_model, proposed=proposed, existing=governing.budget_config
|
||||
)
|
||||
_check_new_entry_outranks_no_cap(
|
||||
entry=entry,
|
||||
existing_model_max_budget=existing_model_max_budget,
|
||||
routed_models=routed_models,
|
||||
team_id=data.team_id,
|
||||
)
|
||||
|
||||
|
||||
def _check_existing_model_cap_kept(entry: str, cap_entry: str, proposed: BudgetConfig, existing: BudgetConfig) -> None:
|
||||
if existing.max_budget is None or not _looser_model_cap(proposed=proposed, cap=existing):
|
||||
return
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": (
|
||||
f"Only a proxy admin can raise a team's model_max_budget for {entry!r} or change its "
|
||||
f"budget_duration. Current max_budget={existing.max_budget} per {existing.budget_duration} "
|
||||
f"(entry {cap_entry!r}), requested={proposed.max_budget} per {proposed.budget_duration}."
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _check_new_entry_outranks_no_cap(
|
||||
entry: str,
|
||||
existing_model_max_budget: Mapping[str, object],
|
||||
routed_models: RoutedModels,
|
||||
team_id: str | None,
|
||||
) -> None:
|
||||
"""A new entry gets its own counter, so the models it takes from a capped entry stop drawing from that cap."""
|
||||
outranked: Final = next(
|
||||
_capped_entries_outranked_by(
|
||||
entry=entry,
|
||||
existing_model_max_budget=existing_model_max_budget,
|
||||
routed_models=routed_models,
|
||||
team_id=team_id,
|
||||
),
|
||||
None,
|
||||
)
|
||||
if outranked is None:
|
||||
return
|
||||
outranked_name, cap = outranked
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": (
|
||||
f"Only a proxy admin can add a team model_max_budget entry {entry!r} that takes precedence over "
|
||||
f"{outranked_name!r}, since the models they share would stop drawing from its max_budget="
|
||||
f"{cap.max_budget} per {cap.budget_duration}."
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _capped_entries_outranked_by(
|
||||
entry: str,
|
||||
existing_model_max_budget: Mapping[str, object],
|
||||
routed_models: RoutedModels,
|
||||
team_id: str | None,
|
||||
) -> Iterator[tuple[str, BudgetConfig]]:
|
||||
governing: Final = resolve_model_budget(model=entry, model_max_budget=existing_model_max_budget, team_id=team_id)
|
||||
outranked_names: Final = itertools.chain(
|
||||
()
|
||||
if governing is None or entry_names_model(entry=governing.budget_model, model=entry)
|
||||
else (governing.budget_model,),
|
||||
entries_outranked_by(entry=entry, model_max_budget=existing_model_max_budget, routed_models=routed_models),
|
||||
)
|
||||
outranked_caps: Final = ((name, _existing_model_cap(existing_model_max_budget[name])) for name in outranked_names)
|
||||
yield from (
|
||||
(name, cap) for name, cap in outranked_caps if name != entry and cap is not None and cap.max_budget is not None
|
||||
)
|
||||
|
||||
|
||||
def _looser_model_cap(proposed: BudgetConfig, cap: BudgetConfig) -> bool:
|
||||
return (
|
||||
proposed.max_budget is None
|
||||
or (cap.max_budget is not None and proposed.max_budget > cap.max_budget)
|
||||
or (is_enforced_budget_entry(cap) and proposed.budget_duration != cap.budget_duration)
|
||||
)
|
||||
|
||||
|
||||
def _should_auto_add_team_creator(
|
||||
|
|
@ -1572,7 +1648,7 @@ async def new_team(
|
|||
- prompts: Optional[List[str]] - List of prompts that the team is allowed to use.
|
||||
- organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`.
|
||||
- model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias)
|
||||
- model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}}
|
||||
- model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model, keyed by model name, `*` pattern, or model access group. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}}
|
||||
- guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails)
|
||||
- policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies)
|
||||
- disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the team. Proxy admin only.
|
||||
|
|
@ -2350,7 +2426,7 @@ async def update_team(
|
|||
- tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing).
|
||||
- organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`.
|
||||
- model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias)
|
||||
- model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}}
|
||||
- model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model, keyed by model name, `*` pattern, or model access group. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}}
|
||||
- guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails)
|
||||
- policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies)
|
||||
- disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the team. Proxy admin only.
|
||||
|
|
|
|||
|
|
@ -40,6 +40,7 @@ from litellm.proxy.management_helpers.access_group_team_sync import (
|
|||
)
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.repositories.chunked_in import delete_many_in, find_many_in, update_many_in
|
||||
from litellm.repositories.team_repository import TeamRepository
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import Prisma
|
||||
|
|
@ -144,7 +145,7 @@ def _changes_nothing(plan: RosterPlan) -> bool:
|
|||
|
||||
|
||||
async def _unchanged_roster(prisma_client: PrismaClient, team_id: str) -> RosterSync | TeamGone:
|
||||
team_row: Final = await prisma_client.writer_db.litellm_teamtable.find_unique(where={"team_id": team_id})
|
||||
team_row: Final = await TeamRepository(prisma_client, use_writer=True).table.find_unique(where={"team_id": team_id})
|
||||
if team_row is None:
|
||||
return TeamGone(team_id=team_id)
|
||||
team: Final = LiteLLM_TeamTable.model_validate(team_row.model_dump())
|
||||
|
|
|
|||
|
|
@ -300,12 +300,13 @@ def decode_state_hash(encrypted_state: str) -> dict:
|
|||
if decrypted_json is None:
|
||||
raise ValueError("Failed to decrypt state parameter")
|
||||
|
||||
state_data: Final = json.loads(decrypted_json)
|
||||
state_data: Final = _DECODED_STATE.validate_python(json.loads(decrypted_json))
|
||||
return state_data
|
||||
|
||||
|
||||
_MANAGED_AUTH_CODE_PREFIX: Final = "llm_scope_"
|
||||
_DECODED_JSON: Final = TypeAdapter(object)
|
||||
_DECODED_STATE: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True))
|
||||
|
||||
|
||||
class _ManagedAuthorizationCode(LiteLLMBaseModel):
|
||||
|
|
|
|||
|
|
@ -33,6 +33,7 @@ from pydantic import TypeAdapter
|
|||
from starlette.datastructures import UploadFile as StarletteUploadFile
|
||||
from starlette.routing import BaseRoute, Route
|
||||
from starlette.websockets import WebSocketState
|
||||
from typing_extensions import ReadOnly
|
||||
from websockets.asyncio.client import ClientConnection, connect
|
||||
from websockets.exceptions import (
|
||||
ConnectionClosedError,
|
||||
|
|
@ -179,8 +180,17 @@ def _is_mergeable_client_metadata(key: str, value: object) -> bool:
|
|||
return not key.startswith("user_api_key_") and (key != "model_info" or isinstance(value, Mapping))
|
||||
|
||||
|
||||
class _RegisteredPassThroughRoute(TypedDict):
|
||||
endpoint_id: ReadOnly[str]
|
||||
path: ReadOnly[str | None]
|
||||
type: ReadOnly[str]
|
||||
methods: ReadOnly[str | bool | list[str] | tuple[str, ...] | Mapping[str, object]]
|
||||
auth: ReadOnly[bool]
|
||||
passthrough_params: ReadOnly[Mapping[str, object]]
|
||||
|
||||
|
||||
# Global registry to track registered pass-through routes and prevent memory leaks
|
||||
_registered_pass_through_routes: Final[dict[str, dict[str, str | bool | list[str] | Mapping[str, object]]]] = {}
|
||||
_registered_pass_through_routes: Final[dict[str, _RegisteredPassThroughRoute]] = {}
|
||||
|
||||
|
||||
def get_response_body(response: httpx.Response) -> dict | None:
|
||||
|
|
@ -2245,7 +2255,7 @@ def create_pass_through_route(
|
|||
status_code=status.HTTP_405_METHOD_NOT_ALLOWED,
|
||||
detail=f"Method {request.method} is not allowed for pass-through endpoint {path}.",
|
||||
)
|
||||
target_params: Final = {
|
||||
target_params: Final[dict[str, object]] = {
|
||||
"target": target,
|
||||
"custom_headers": custom_headers,
|
||||
"forward_headers": _forward_headers,
|
||||
|
|
@ -2277,7 +2287,9 @@ def create_pass_through_route(
|
|||
|
||||
# Ensure custom_headers is a dict. Botocore returns a HeadersDict
|
||||
# for SigV4-prepared requests, which is a Mapping but not a dict.
|
||||
headers_dict: Final = dict(param_custom_headers) if isinstance(param_custom_headers, Mapping) else {}
|
||||
headers_dict: Final[dict[str, object]] = (
|
||||
{**param_custom_headers} if isinstance(param_custom_headers, Mapping) else {}
|
||||
)
|
||||
|
||||
# Ensure query_params and custom_body are dicts or None
|
||||
final_query_params: Final = query_params_data if isinstance(query_params_data, dict) else {}
|
||||
|
|
@ -3350,9 +3362,7 @@ class SafeRouteAdder:
|
|||
return f"{path}/{{subpath:path}}"
|
||||
return None
|
||||
|
||||
def registered_route_methods(
|
||||
route_info: Mapping[str, str | bool | Sequence[str] | Mapping[str, object]],
|
||||
) -> tuple[tuple[str, str], ...]:
|
||||
def registered_route_methods(route_info: _RegisteredPassThroughRoute) -> tuple[tuple[str, str], ...]:
|
||||
starlette_path: Final = registered_route_path(route_info)
|
||||
methods: Final = route_info.get("methods")
|
||||
if starlette_path is None or not isinstance(methods, (list, tuple)):
|
||||
|
|
@ -3708,7 +3718,7 @@ class InitPassThroughEndpointHelpers:
|
|||
return (comparison_route, *forwarded)
|
||||
|
||||
@staticmethod
|
||||
def get_registered_pass_through_route(route: str, method: str | None = None) -> dict[str, Any] | None:
|
||||
def get_registered_pass_through_route(route: str, method: str | None = None) -> _RegisteredPassThroughRoute | None:
|
||||
"""Get passthrough params for a given route and optionally filter by HTTP method"""
|
||||
comparison_route: Final = InitPassThroughEndpointHelpers._route_for_registry_lookup(route)
|
||||
for key in _registered_pass_through_routes:
|
||||
|
|
|
|||
|
|
@ -666,7 +666,11 @@ from litellm.proxy.list_api.common import (
|
|||
problem_response,
|
||||
request_validation_problem,
|
||||
)
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, add_litellm_data_to_request
|
||||
from litellm.proxy.litellm_pre_call_utils import (
|
||||
SCHEDULER_PRIORITY_HEADER,
|
||||
LiteLLMProxyRequestSetup,
|
||||
add_litellm_data_to_request,
|
||||
)
|
||||
from litellm.proxy.logging_endpoints.callback_logs_endpoints import (
|
||||
rust_control_plane_router,
|
||||
)
|
||||
|
|
@ -18481,6 +18485,11 @@ async def async_queue_request(
|
|||
raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value})
|
||||
router: Final = llm_router
|
||||
|
||||
data.update(
|
||||
router.priority_header_update(
|
||||
model=data["model"], kwargs=data, header=request.headers.get(SCHEDULER_PRIORITY_HEADER)
|
||||
)
|
||||
)
|
||||
queued_priority, _ = router.split_request_priority(model=data["model"], kwargs=data)
|
||||
if "stream" in data and data["stream"] is True: # use generate_responses to stream responses
|
||||
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ from datetime import datetime, timedelta
|
|||
from types import MappingProxyType
|
||||
from typing import Final, Literal, TypeVar
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -155,6 +155,7 @@ class _DailyUserSpendOwnerRow(LiteLLMBaseModel):
|
|||
last_owner: str | None = None
|
||||
|
||||
|
||||
_RAW_ROWS: Final = TypeAdapter(Sequence[dict[str, object]], config=ConfigDict(strict=True))
|
||||
_TOKEN_DIGEST_ROWS: Final = TypeAdapter(tuple[_TokenDigestRow, ...])
|
||||
_SPEND_LOG_DIGEST_ROWS: Final = TypeAdapter(tuple[_SpendLogDigestRow, ...])
|
||||
_DAILY_USER_SPEND_OWNER_ROWS: Final = TypeAdapter(tuple[_DailyUserSpendOwnerRow, ...])
|
||||
|
|
@ -194,7 +195,7 @@ async def _rows_within_the_statement_timeout(
|
|||
await transaction.execute_raw(_SPEND_LOG_STATEMENT_TIMEOUT_SQL)
|
||||
for setting in planner_settings:
|
||||
await transaction.execute_raw(setting)
|
||||
return await transaction.query_raw(sql, *params)
|
||||
return _RAW_ROWS.validate_python(await transaction.query_raw(sql, *params))
|
||||
|
||||
|
||||
async def _reverse_hash_key_metadata(
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ from litellm.repositories.daily_activity_sql import (
|
|||
build_cache_leakage_keys_sql,
|
||||
build_entity_rollup_sql,
|
||||
build_export_sql,
|
||||
build_guardrail_spend_sql,
|
||||
build_key_page_sql,
|
||||
build_key_search_sql,
|
||||
build_model_top_keys_sql,
|
||||
|
|
@ -36,6 +37,7 @@ from litellm.types.repositories.daily_activity import (
|
|||
ExportRow,
|
||||
ExportType,
|
||||
GroupingSetsRow,
|
||||
GuardrailSpendRow,
|
||||
KeyMetadataRow,
|
||||
KeyPage,
|
||||
KeySpendRow,
|
||||
|
|
@ -100,6 +102,7 @@ class DailyActivityDatabase(Protocol):
|
|||
|
||||
_GROUPING_ADAPTER: Final = TypeAdapter(tuple[GroupingSetsRow, ...])
|
||||
_ENTITY_ADAPTER: Final = TypeAdapter(tuple[EntityRollupRow, ...])
|
||||
_GUARDRAIL_SPEND_ADAPTER: Final = TypeAdapter(tuple[GuardrailSpendRow, ...])
|
||||
_KEY_SPEND_ADAPTER: Final = TypeAdapter(tuple[KeySpendRow, ...])
|
||||
_KEY_PAGE_TOTAL_ADAPTER: Final[TypeAdapter[int]] = TypeAdapter(int)
|
||||
_USER_SPEND_ADAPTER: Final = TypeAdapter(tuple[UserSpendRow, ...])
|
||||
|
|
@ -191,6 +194,9 @@ class DailyActivityRepository:
|
|||
distinct_api_keys=distinct_api_keys,
|
||||
)
|
||||
|
||||
async def guardrail_spend(self, scope: DailyActivityScope) -> tuple[GuardrailSpendRow, ...]:
|
||||
return _GUARDRAIL_SPEND_ADAPTER.validate_python(await self._query(build_guardrail_spend_sql(scope)))
|
||||
|
||||
async def search_keys(self, scope: DailyActivityScope, *, search: str, limit: int) -> tuple[str, ...]:
|
||||
if not 1 <= limit <= constants.USAGE_KEY_SEARCH_MAX:
|
||||
raise ValueError(f"limit must be between 1 and {constants.USAGE_KEY_SEARCH_MAX}")
|
||||
|
|
|
|||
|
|
@ -261,6 +261,25 @@ def build_entity_rollup_sql(scope: DailyActivityScope, *, api_key_limit: int) ->
|
|||
return SqlQuery(sql=sql, params=(*where_params, PTU_SENTINEL_API_KEY, api_key_limit))
|
||||
|
||||
|
||||
def build_guardrail_spend_sql(scope: DailyActivityScope) -> SqlQuery:
|
||||
pg_table: Final = PRISMA_TO_PG_TABLE[scope.table]
|
||||
where_clause, where_params = build_where_clause(scope)
|
||||
sql: Final = f"""
|
||||
SELECT g.date, g.guardrail_id, COALESCE(t.guardrail_name, g.guardrail_id) AS guardrail_name, g.api_key,
|
||||
COALESCE(SUM(g.cost), 0)::float AS spend,
|
||||
COALESCE(SUM(CASE WHEN g.cost IS NULL THEN g.units ELSE g.untracked_units END), 0)::bigint AS untracked_units
|
||||
FROM "LiteLLM_DailyGuardrailUsageUnits" g
|
||||
LEFT JOIN "LiteLLM_GuardrailsTable" t ON t.guardrail_id = g.guardrail_id
|
||||
WHERE g.date >= $1 AND g.date <= $2
|
||||
AND EXISTS (
|
||||
SELECT 1 FROM "{pg_table}"
|
||||
WHERE {where_clause} AND "{pg_table}".date = g.date AND "{pg_table}".api_key = g.api_key
|
||||
)
|
||||
GROUP BY g.date, g.guardrail_id, t.guardrail_name, g.api_key
|
||||
"""
|
||||
return SqlQuery(sql=sql, params=where_params)
|
||||
|
||||
|
||||
def _key_spend_select() -> str:
|
||||
return """
|
||||
COALESCE(SUM(spend), 0)::float AS spend,
|
||||
|
|
|
|||
|
|
@ -359,6 +359,7 @@ from litellm.router_utils.request_priority import (
|
|||
InvalidPriority,
|
||||
parse_default_priority,
|
||||
priority_error,
|
||||
priority_header_fields,
|
||||
request_drops_params,
|
||||
resolve_request_priority,
|
||||
unqueueable_priority_message,
|
||||
|
|
@ -1369,7 +1370,7 @@ class Router:
|
|||
},
|
||||
]
|
||||
|
||||
router = Router(model_list=model_list, fallbacks=[{"azure-gpt-3.5-turbo": "openai-gpt-3.5-turbo"}])
|
||||
router = Router(model_list=model_list, fallbacks=[{"azure-gpt-3.5-turbo": ["openai-gpt-3.5-turbo"]}])
|
||||
```
|
||||
"""
|
||||
|
||||
|
|
@ -1595,6 +1596,10 @@ class Router:
|
|||
|
||||
if default_fallbacks is not None or litellm.default_fallbacks is not None:
|
||||
_fallbacks = default_fallbacks or litellm.default_fallbacks
|
||||
if isinstance(_fallbacks, str):
|
||||
raise TypeError(
|
||||
f"default_fallbacks must be a list of model names, got '{_fallbacks}'. Use ['{_fallbacks}']"
|
||||
)
|
||||
if self.fallbacks is not None:
|
||||
self.fallbacks.append({"*": _fallbacks})
|
||||
else:
|
||||
|
|
@ -2838,7 +2843,7 @@ class Router:
|
|||
self._initialize_core_endpoints()
|
||||
self._initialize_specialized_endpoints()
|
||||
|
||||
def validate_fallbacks(self, fallback_param: list | None):
|
||||
def validate_fallbacks(self, fallback_param: Sequence[Mapping[str, object]] | None):
|
||||
"""
|
||||
Validate the fallbacks parameter.
|
||||
"""
|
||||
|
|
@ -2851,6 +2856,13 @@ class Router:
|
|||
raise ValueError(
|
||||
f"Dictionary '{fallback_dict}' must have exactly one key, but has {len(fallback_dict)} keys."
|
||||
)
|
||||
((model_group, targets),) = fallback_dict.items()
|
||||
if isinstance(targets, list) or check_non_standard_fallback_format(fallbacks=[fallback_dict]):
|
||||
continue
|
||||
raise TypeError(
|
||||
f"Fallbacks for '{model_group}' must be a list of model names, got '{targets}'. "
|
||||
f"Use {{'{model_group}': ['{targets}']}}"
|
||||
)
|
||||
|
||||
def _add_encrypted_content_affinity_check(self, enable_global_affinity: bool) -> None:
|
||||
def _move_before_deployment_affinity(
|
||||
|
|
@ -5411,14 +5423,22 @@ class Router:
|
|||
request_priority: Final = resolve_request_priority(
|
||||
requested=kwargs.get("priority"),
|
||||
default_priority=self.default_priority,
|
||||
drops_params=lambda: request_drops_params(
|
||||
kwargs, self.default_litellm_params, self._request_deployment_params(model=model, kwargs=kwargs)
|
||||
),
|
||||
drops_params=lambda: self._drops_params(model=model, kwargs=kwargs),
|
||||
)
|
||||
if isinstance(request_priority, InvalidPriority):
|
||||
raise priority_error(model=model, message=request_priority.message)
|
||||
return request_priority, {key: value for key, value in kwargs.items() if key != "priority"}
|
||||
|
||||
def priority_header_update(
|
||||
self, model: str, kwargs: Mapping[str, object], header: str | None
|
||||
) -> Mapping[str, int | str]:
|
||||
return priority_header_fields(header, drops_params=lambda: self._drops_params(model=model, kwargs=kwargs))
|
||||
|
||||
def _drops_params(self, model: str, kwargs: Mapping[str, object]) -> bool:
|
||||
return request_drops_params(
|
||||
kwargs, self.default_litellm_params, self._request_deployment_params(model=model, kwargs=kwargs)
|
||||
)
|
||||
|
||||
async def _schedule_factory(
|
||||
self,
|
||||
model: str,
|
||||
|
|
@ -11850,6 +11870,11 @@ class Router:
|
|||
indices: Final = self.model_name_to_deployment_indices.get(model_name)
|
||||
return None if indices is None else tuple(self.model_list[idx] for idx in indices)
|
||||
|
||||
def team_public_model_names_of(self, team_id: str) -> frozenset[str]:
|
||||
"""The public names `team_id` serves its own deployments under."""
|
||||
with self._deployment_index_lock:
|
||||
return frozenset(name for owner, name in self.team_model_to_deployment_indices if owner == team_id)
|
||||
|
||||
def _team_deployments_named(self, team_id: str, public_model_name: str) -> tuple[DeploymentTypedDict, ...] | None:
|
||||
"""The deployments `team_id` serves under `public_model_name`, or None when that pair is not indexed."""
|
||||
with self._deployment_index_lock:
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ from datetime import datetime
|
|||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm._internal_context import service_target
|
||||
|
|
@ -417,17 +417,35 @@ def _should_run_cooldown_logic(
|
|||
return True
|
||||
|
||||
|
||||
_FALLBACK_ENTRIES_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, list[str | Mapping[str, object]]] | str, ...])
|
||||
_FALLBACK_ENTRIES_ADAPTER: Final = TypeAdapter(tuple[object, ...])
|
||||
_FALLBACK_ENTRY_ADAPTER: Final = TypeAdapter(Mapping[str, list[str | Mapping[str, object]] | str])
|
||||
|
||||
|
||||
def _fallback_target_name(target: str | Mapping[str, object]) -> str:
|
||||
return target if isinstance(target, str) else str(target.get("model", ""))
|
||||
|
||||
|
||||
def _fallback_target_names(
|
||||
targets: Sequence[str | Mapping[str, object]] | str,
|
||||
) -> list[str]: # mutable-ok: mirrors get_fallback_model_group's fallbacks shape
|
||||
return [targets] if isinstance(targets, str) else [_fallback_target_name(target) for target in targets]
|
||||
|
||||
|
||||
def _fallback_chain_names(
|
||||
entry: Mapping[str, Sequence[str | Mapping[str, object]]],
|
||||
) -> Mapping[str, list[str]]: # mutable-ok: mirrors get_fallback_model_group's fallbacks shape
|
||||
return {key: [_fallback_target_name(target) for target in targets] for key, targets in entry.items()}
|
||||
entry: object,
|
||||
) -> Mapping[str, list[str]] | None: # mutable-ok: mirrors get_fallback_model_group's fallbacks shape
|
||||
try:
|
||||
validated: Final = _FALLBACK_ENTRY_ADAPTER.validate_python(entry)
|
||||
except ValidationError:
|
||||
return None
|
||||
return {key: _fallback_target_names(targets) for key, targets in validated.items()}
|
||||
|
||||
|
||||
def _fallback_entries(configured: object) -> tuple[object, ...]:
|
||||
try:
|
||||
return _FALLBACK_ENTRIES_ADAPTER.validate_python(configured)
|
||||
except ValidationError:
|
||||
return ()
|
||||
|
||||
|
||||
def _has_fallback_group(litellm_router_instance: LitellmRouter, requested_model_group: str | None) -> bool:
|
||||
|
|
@ -437,9 +455,7 @@ def _has_fallback_group(litellm_router_instance: LitellmRouter, requested_model_
|
|||
if requested_model_group is None or not configured:
|
||||
return False
|
||||
chains: Final = tuple(
|
||||
_fallback_chain_names(entry)
|
||||
for entry in _FALLBACK_ENTRIES_ADAPTER.validate_python(configured)
|
||||
if isinstance(entry, Mapping)
|
||||
chain for chain in map(_fallback_chain_names, _fallback_entries(configured)) if chain is not None
|
||||
)
|
||||
fallback_group, _ = get_fallback_model_group(fallbacks=chains, model_group=requested_model_group)
|
||||
return bool(fallback_group)
|
||||
|
|
|
|||
|
|
@ -5,15 +5,19 @@ When enabled via router_settings.optional_pre_call_checks: ["enforce_model_rate_
|
|||
|
||||
- tpm/rpm: combined TPM + optional RPM (legacy)
|
||||
- itpm/otpm: separate input/output tokens per minute
|
||||
- tpd/rpd: tokens and requests per UTC calendar day
|
||||
|
||||
When a deployment sets both itpm/otpm and tpm/rpm, both are enforced. A warning
|
||||
is logged the first time such a deployment is seen.
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
from collections.abc import Callable
|
||||
from datetime import datetime, timedelta
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
import litellm
|
||||
from litellm._internal_context import with_service_target
|
||||
|
|
@ -33,6 +37,7 @@ from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import (
|
|||
io_token_refund_failure,
|
||||
)
|
||||
from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET
|
||||
from litellm.types.caching import RedisPipelineIncrementOperation
|
||||
from litellm.types.router import RouterErrors
|
||||
from litellm.types.utils import StandardLoggingPayload
|
||||
from litellm.utils import get_utc_datetime
|
||||
|
|
@ -47,6 +52,62 @@ else:
|
|||
|
||||
class RoutingArgs:
|
||||
ttl: int = 60 # 1min (RPM/TPM expire key)
|
||||
daily_ttl: int = 86400
|
||||
|
||||
|
||||
DAY_FORMAT: Final = "%Y-%m-%d"
|
||||
|
||||
|
||||
def seconds_until_utc_midnight(now: datetime) -> int:
|
||||
next_midnight: Final = (now + timedelta(days=1)).replace(hour=0, minute=0, second=0, microsecond=0)
|
||||
return max(1, int((next_midnight - now).total_seconds()))
|
||||
|
||||
|
||||
def _first_set(*values: int | None) -> int | None:
|
||||
return next((value for value in values if value is not None), None)
|
||||
|
||||
|
||||
class _DailyLimitFields(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore", from_attributes=True)
|
||||
tpd: int | None = None
|
||||
rpd: int | None = None
|
||||
|
||||
|
||||
class _DailyLimitParams(_DailyLimitFields):
|
||||
model: str | None = None
|
||||
|
||||
|
||||
class _DailyLimitModelInfo(_DailyLimitFields):
|
||||
id: str | None = None
|
||||
|
||||
|
||||
class _DailyLimitDeployment(_DailyLimitFields):
|
||||
model_name: str = ""
|
||||
litellm_params: _DailyLimitParams = _DailyLimitParams()
|
||||
model_info: _DailyLimitModelInfo = _DailyLimitModelInfo()
|
||||
|
||||
def limits(self) -> tuple[int | None, int | None]:
|
||||
return (
|
||||
_first_set(self.tpd, self.litellm_params.tpd, self.model_info.tpd),
|
||||
_first_set(self.rpd, self.litellm_params.rpd, self.model_info.rpd),
|
||||
)
|
||||
|
||||
def cache_keys(self, current_day: str) -> tuple[str, str]:
|
||||
prefix: Final = f"{self.model_info.id}:{self.litellm_params.model}"
|
||||
return f"{prefix}:tpd:{current_day}", f"{prefix}:rpd:{current_day}"
|
||||
|
||||
def limit_error(self, kind: Literal["tpd", "rpd"], limit: int, usage: int, now: datetime) -> litellm.RateLimitError:
|
||||
return litellm.RateLimitError(
|
||||
message=f"Model rate limit exceeded. {kind.upper()} limit={limit}, current usage={usage}",
|
||||
llm_provider="",
|
||||
model=self.litellm_params.model,
|
||||
response=httpx.Response(
|
||||
status_code=429,
|
||||
content=f"{RouterErrors.user_defined_ratelimit_error.value} {kind} limit={limit}. current usage={usage}. id={self.model_info.id}, model_group={self.model_name}",
|
||||
headers={"retry-after": str(seconds_until_utc_midnight(now))},
|
||||
request=httpx.Request(method="model_rate_limit_check", url="https://github.com/BerriAI/litellm"),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class ModelRateLimitingCheck(CustomLogger):
|
||||
|
|
@ -60,12 +121,60 @@ class ModelRateLimitingCheck(CustomLogger):
|
|||
this check actively enforces those limits across ALL routing strategies.
|
||||
"""
|
||||
|
||||
def __init__(self, dual_cache: DualCache):
|
||||
def __init__(self, dual_cache: DualCache, clock: Callable[[], datetime] = get_utc_datetime):
|
||||
self.dual_cache = dual_cache
|
||||
self.clock = clock
|
||||
# model_ids already warned about itpm/otpm + tpm/rpm on the same deployment,
|
||||
# so the warning is logged once per deployment rather than per request.
|
||||
self._io_token_conflict_warned_ids: set[str] = set()
|
||||
|
||||
def _increment(self, key: str, value: int, ttl: int) -> int:
|
||||
return self.dual_cache.increment_cache(key=key, value=value, ttl=ttl)
|
||||
|
||||
async def _async_increment(
|
||||
self, key: str, value: int, ttl: int, parent_otel_span: Span | None = None
|
||||
) -> int | None:
|
||||
count: Final = await self.dual_cache.async_increment_cache(
|
||||
key=key, value=value, ttl=ttl, parent_otel_span=parent_otel_span
|
||||
)
|
||||
return None if count is None else int(count)
|
||||
|
||||
@staticmethod
|
||||
def _request_count_operations(
|
||||
rpm_key: str | None, rpd_key: str | None, increment_value: int
|
||||
) -> tuple[RedisPipelineIncrementOperation, ...]:
|
||||
return tuple(
|
||||
RedisPipelineIncrementOperation(key=key, increment_value=increment_value, ttl=ttl)
|
||||
for key, ttl in ((rpm_key, RoutingArgs.ttl), (rpd_key, RoutingArgs.daily_ttl))
|
||||
if key is not None
|
||||
)
|
||||
|
||||
async def _async_refund_request_counts(
|
||||
self, rpm_key: str | None, rpd_key: str | None, parent_otel_span: Span | None
|
||||
) -> None:
|
||||
operations: Final = self._request_count_operations(rpm_key, rpd_key, -1)
|
||||
if operations:
|
||||
await self.dual_cache.async_increment_cache_pipeline(
|
||||
increment_list=list(operations), parent_otel_span=parent_otel_span
|
||||
)
|
||||
|
||||
async def _async_increment_request_counts(
|
||||
self, rpm_key: str | None, rpd_key: str | None, parent_otel_span: Span | None
|
||||
) -> tuple[int | None, int | None]:
|
||||
operations: Final = self._request_count_operations(rpm_key, rpd_key, 1)
|
||||
if not operations:
|
||||
return None, None
|
||||
counts: Final = await self.dual_cache.async_increment_cache_pipeline(
|
||||
increment_list=list(operations), parent_otel_span=parent_otel_span
|
||||
)
|
||||
if counts is None or len(counts) != len(operations):
|
||||
return None, None
|
||||
count_by_key: Final = {op["key"]: int(count) for op, count in zip(operations, counts)}
|
||||
return (
|
||||
None if rpm_key is None else count_by_key.get(rpm_key),
|
||||
None if rpd_key is None else count_by_key.get(rpd_key),
|
||||
)
|
||||
|
||||
def _warn_io_token_and_tpm_rpm_coexist_once(self, deployment: dict) -> None:
|
||||
tpm_limit, rpm_limit = self._get_deployment_limits(deployment)
|
||||
if tpm_limit is None and rpm_limit is None:
|
||||
|
|
@ -129,6 +238,21 @@ class ModelRateLimitingCheck(CustomLogger):
|
|||
|
||||
return tpm, rpm
|
||||
|
||||
def _check_daily_limits(self, deployment: _DailyLimitDeployment, now: datetime) -> None:
|
||||
tpd_limit, rpd_limit = deployment.limits()
|
||||
if tpd_limit is None and rpd_limit is None:
|
||||
return
|
||||
tpd_key, rpd_key = deployment.cache_keys(now.strftime(DAY_FORMAT))
|
||||
if tpd_limit is not None:
|
||||
current_tpd: Final = self._get_current_tpm(tpd_key, tpd_limit)
|
||||
if current_tpd is not None and current_tpd >= tpd_limit:
|
||||
raise deployment.limit_error("tpd", tpd_limit, current_tpd, now)
|
||||
if rpd_limit is not None:
|
||||
current_rpd: Final = self._increment(rpd_key, 1, RoutingArgs.daily_ttl)
|
||||
if current_rpd > rpd_limit:
|
||||
self._increment(rpd_key, -1, RoutingArgs.daily_ttl)
|
||||
raise deployment.limit_error("rpd", rpd_limit, current_rpd, now)
|
||||
|
||||
def _get_cache_keys(self, deployment: dict, current_minute: str) -> tuple[str, str]:
|
||||
"""Get the cache keys for TPM and RPM tracking."""
|
||||
model_id: Final = deployment.get("model_info", {}).get("id")
|
||||
|
|
@ -178,13 +302,10 @@ class ModelRateLimitingCheck(CustomLogger):
|
|||
)
|
||||
io_reservation_made = True
|
||||
|
||||
dt: Final = self.clock()
|
||||
daily: Final = _DailyLimitDeployment.model_validate(deployment)
|
||||
tpm_limit, rpm_limit = self._get_deployment_limits(deployment)
|
||||
|
||||
# If no limits are set, allow the request
|
||||
if tpm_limit is None and rpm_limit is None:
|
||||
return deployment
|
||||
|
||||
dt: Final = get_utc_datetime()
|
||||
current_minute: Final = dt.strftime("%H-%M")
|
||||
tpm_key, rpm_key = self._get_cache_keys(deployment, current_minute)
|
||||
|
||||
|
|
@ -213,8 +334,9 @@ class ModelRateLimitingCheck(CustomLogger):
|
|||
|
||||
# Check RPM limit (atomic increment-first to avoid race conditions)
|
||||
if rpm_limit is not None:
|
||||
current_rpm: Final = self.dual_cache.increment_cache(key=rpm_key, value=1, ttl=RoutingArgs.ttl)
|
||||
current_rpm: Final = self._increment(rpm_key, 1, RoutingArgs.ttl)
|
||||
if current_rpm is not None and current_rpm > rpm_limit:
|
||||
self._increment(rpm_key, -1, RoutingArgs.ttl)
|
||||
raise litellm.RateLimitError(
|
||||
message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}",
|
||||
llm_provider="",
|
||||
|
|
@ -230,6 +352,12 @@ class ModelRateLimitingCheck(CustomLogger):
|
|||
),
|
||||
)
|
||||
|
||||
try:
|
||||
self._check_daily_limits(daily, dt)
|
||||
except litellm.RateLimitError:
|
||||
if rpm_limit is not None:
|
||||
self._increment(rpm_key, -1, RoutingArgs.ttl)
|
||||
raise
|
||||
return deployment
|
||||
|
||||
except litellm.RateLimitError:
|
||||
|
|
@ -259,15 +387,14 @@ class ModelRateLimitingCheck(CustomLogger):
|
|||
)
|
||||
io_reservation_made = True
|
||||
|
||||
dt: Final = self.clock()
|
||||
daily: Final = _DailyLimitDeployment.model_validate(deployment)
|
||||
tpd_limit, rpd_limit = daily.limits()
|
||||
tpm_limit, rpm_limit = self._get_deployment_limits(deployment)
|
||||
|
||||
# If no limits are set, allow the request
|
||||
if tpm_limit is None and rpm_limit is None:
|
||||
return deployment
|
||||
|
||||
dt: Final = get_utc_datetime()
|
||||
current_minute: Final = dt.strftime("%H-%M")
|
||||
tpm_key, rpm_key = self._get_cache_keys(deployment, current_minute)
|
||||
tpd_key, rpd_key = daily.cache_keys(dt.strftime(DAY_FORMAT))
|
||||
|
||||
model_id: Final = deployment.get("model_info", {}).get("id")
|
||||
model_name: Final = deployment.get("litellm_params", {}).get("model")
|
||||
|
|
@ -293,30 +420,37 @@ class ModelRateLimitingCheck(CustomLogger):
|
|||
num_retries=0, # Don't retry - return 429 immediately
|
||||
)
|
||||
|
||||
# Check RPM limit (atomic increment-first to avoid race conditions)
|
||||
if rpm_limit is not None:
|
||||
current_rpm: Final = await self.dual_cache.async_increment_cache(
|
||||
key=rpm_key,
|
||||
value=1,
|
||||
ttl=RoutingArgs.ttl,
|
||||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
if current_rpm is not None and current_rpm > rpm_limit:
|
||||
raise litellm.RateLimitError(
|
||||
message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}",
|
||||
llm_provider="",
|
||||
model=model_name,
|
||||
response=httpx.Response(
|
||||
status_code=429,
|
||||
content=f"{RouterErrors.user_defined_ratelimit_error.value} rpm limit={rpm_limit}. current usage={current_rpm}. id={model_id}, model_group={model_group}",
|
||||
headers={"retry-after": str(60)},
|
||||
request=httpx.Request(
|
||||
method="model_rate_limit_check",
|
||||
url="https://github.com/BerriAI/litellm",
|
||||
),
|
||||
if tpd_limit is not None:
|
||||
current_tpd: Final = await self._async_get_current_tpm(tpd_key, tpd_limit, parent_otel_span)
|
||||
if current_tpd is not None and current_tpd >= tpd_limit:
|
||||
raise daily.limit_error("tpd", tpd_limit, current_tpd, dt)
|
||||
|
||||
charged_rpm_key: Final = rpm_key if rpm_limit is not None else None
|
||||
charged_rpd_key: Final = rpd_key if rpd_limit is not None else None
|
||||
current_rpm, current_rpd = await self._async_increment_request_counts(
|
||||
charged_rpm_key, charged_rpd_key, parent_otel_span
|
||||
)
|
||||
if rpm_limit is not None and current_rpm is not None and current_rpm > rpm_limit:
|
||||
await self._async_refund_request_counts(charged_rpm_key, charged_rpd_key, parent_otel_span)
|
||||
raise litellm.RateLimitError(
|
||||
message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}",
|
||||
llm_provider="",
|
||||
model=model_name,
|
||||
response=httpx.Response(
|
||||
status_code=429,
|
||||
content=f"{RouterErrors.user_defined_ratelimit_error.value} rpm limit={rpm_limit}. current usage={current_rpm}. id={model_id}, model_group={model_group}",
|
||||
headers={"retry-after": str(60)},
|
||||
request=httpx.Request(
|
||||
method="model_rate_limit_check",
|
||||
url="https://github.com/BerriAI/litellm",
|
||||
),
|
||||
num_retries=0, # Don't retry - return 429 immediately
|
||||
)
|
||||
),
|
||||
num_retries=0, # Don't retry - return 429 immediately
|
||||
)
|
||||
|
||||
if rpd_limit is not None and current_rpd is not None and current_rpd > rpd_limit:
|
||||
await self._async_refund_request_counts(charged_rpm_key, charged_rpd_key, parent_otel_span)
|
||||
raise daily.limit_error("rpd", rpd_limit, current_rpd, dt)
|
||||
|
||||
return deployment
|
||||
|
||||
|
|
@ -373,17 +507,16 @@ class ModelRateLimitingCheck(CustomLogger):
|
|||
if not model or not total_tokens:
|
||||
return
|
||||
|
||||
dt: Final = get_utc_datetime()
|
||||
dt: Final = self.clock()
|
||||
current_minute: Final = dt.strftime("%H-%M")
|
||||
tpm_key: Final = f"{model_id}:{model}:tpm:{current_minute}"
|
||||
tpd_key: Final = f"{model_id}:{model}:tpd:{dt.strftime(DAY_FORMAT)}"
|
||||
|
||||
verbose_router_logger.debug("[TPM TRACKING] Incrementing %s by %s", tpm_key, total_tokens)
|
||||
|
||||
await self.dual_cache.async_increment_cache(
|
||||
key=tpm_key,
|
||||
value=total_tokens,
|
||||
ttl=RoutingArgs.ttl,
|
||||
)
|
||||
tokens: Final = int(total_tokens)
|
||||
await self._async_increment(tpm_key, tokens, RoutingArgs.ttl)
|
||||
await self._async_increment(tpd_key, tokens, RoutingArgs.daily_ttl)
|
||||
|
||||
except Exception as e:
|
||||
verbose_router_logger.debug("Error in ModelRateLimitingCheck.async_log_success_event: %s", e)
|
||||
|
|
@ -434,15 +567,14 @@ class ModelRateLimitingCheck(CustomLogger):
|
|||
if not model or not total_tokens:
|
||||
return
|
||||
|
||||
dt: Final = get_utc_datetime()
|
||||
dt: Final = self.clock()
|
||||
current_minute: Final = dt.strftime("%H-%M")
|
||||
tpm_key: Final = f"{model_id}:{model}:tpm:{current_minute}"
|
||||
tpd_key: Final = f"{model_id}:{model}:tpd:{dt.strftime(DAY_FORMAT)}"
|
||||
|
||||
self.dual_cache.increment_cache(
|
||||
key=tpm_key,
|
||||
value=total_tokens,
|
||||
ttl=RoutingArgs.ttl,
|
||||
)
|
||||
tokens: Final = int(total_tokens)
|
||||
self._increment(tpm_key, tokens, RoutingArgs.ttl)
|
||||
self._increment(tpd_key, tokens, RoutingArgs.daily_ttl)
|
||||
|
||||
except Exception as e:
|
||||
verbose_router_logger.debug("Error in ModelRateLimitingCheck.log_success_event: %s", e)
|
||||
|
|
|
|||
|
|
@ -47,6 +47,15 @@ def resolve_request_priority(
|
|||
return priority
|
||||
|
||||
|
||||
def priority_header_fields(header: str | None, drops_params: Callable[[], bool]) -> Mapping[str, int | str]:
|
||||
if not header:
|
||||
return {}
|
||||
priority: Final = _as_priority(header)
|
||||
if isinstance(priority, InvalidPriority):
|
||||
return {} if drops_params() else {"priority": header}
|
||||
return {"priority": priority}
|
||||
|
||||
|
||||
def parse_default_priority(value: object) -> int | None | InvalidPriority:
|
||||
if value is None:
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -21,6 +21,8 @@ from types import MappingProxyType
|
|||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, TypeAdapter, with_config
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -41,6 +43,15 @@ if TYPE_CHECKING:
|
|||
from botocore.awsrequest import HTTPHeaders
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True))
|
||||
class _SecretStringResponse(TypedDict):
|
||||
SecretString: ReadOnly[str]
|
||||
|
||||
|
||||
_SECRET_STRING_RESPONSE: Final = TypeAdapter(_SecretStringResponse)
|
||||
_AWS_SECRET_MANAGER_RESPONSE: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True))
|
||||
|
||||
|
||||
class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -154,7 +165,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
|
|||
try:
|
||||
response: Final = await async_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8"))
|
||||
response.raise_for_status()
|
||||
return response.json()["SecretString"]
|
||||
return _SECRET_STRING_RESPONSE.validate_python(response.json())["SecretString"]
|
||||
except httpx.TimeoutException:
|
||||
raise ValueError("Timeout error occurred")
|
||||
except Exception as e:
|
||||
|
|
@ -204,7 +215,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
|
|||
|
||||
try:
|
||||
response: Final = sync_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8"))
|
||||
return response.json()["SecretString"]
|
||||
return _SECRET_STRING_RESPONSE.validate_python(response.json())["SecretString"]
|
||||
except httpx.TimeoutException:
|
||||
raise ValueError("Timeout error occurred")
|
||||
except httpx.HTTPStatusError as e:
|
||||
|
|
@ -424,7 +435,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
|
|||
try:
|
||||
response: Final = await async_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8"))
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
return _AWS_SECRET_MANAGER_RESPONSE.validate_python(response.json())
|
||||
except httpx.HTTPStatusError as err:
|
||||
raise ValueError(f"HTTP error occurred: {err.response.text}")
|
||||
except httpx.TimeoutException:
|
||||
|
|
@ -481,7 +492,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
|
|||
try:
|
||||
response: Final = await async_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8"))
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
return _AWS_SECRET_MANAGER_RESPONSE.validate_python(response.json())
|
||||
except httpx.HTTPStatusError as err:
|
||||
raise ValueError(f"HTTP error occurred: {err.response.text}")
|
||||
except httpx.TimeoutException:
|
||||
|
|
@ -533,7 +544,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
|
|||
try:
|
||||
response: Final = await async_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8"))
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
return _AWS_SECRET_MANAGER_RESPONSE.validate_python(response.json())
|
||||
except httpx.HTTPStatusError as err:
|
||||
raise ValueError(f"HTTP error occurred: {err.response.text}")
|
||||
except httpx.TimeoutException:
|
||||
|
|
@ -613,7 +624,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
|
|||
try:
|
||||
response: Final = await async_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8"))
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
return _AWS_SECRET_MANAGER_RESPONSE.validate_python(response.json())
|
||||
except httpx.HTTPStatusError as err:
|
||||
raise ValueError(f"HTTP error occurred: {err.response.text}")
|
||||
except httpx.TimeoutException:
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from typing import Any, Literal
|
||||
from typing import Literal
|
||||
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
|
|
@ -40,7 +40,7 @@ class BedrockKBRetrievalResult(TypedDict, total=False):
|
|||
content: BedrockKBContent | None
|
||||
location: BedrockKBLocation | None
|
||||
score: float | None
|
||||
metadata: dict[str, Any] | None
|
||||
metadata: dict[str, object] | None
|
||||
|
||||
|
||||
class BedrockKBResponse(TypedDict, total=False):
|
||||
|
|
@ -87,21 +87,21 @@ class BedrockKBMetadataConfiguration(TypedDict, total=False):
|
|||
class BedrockKBModelConfiguration(TypedDict, total=False):
|
||||
"""Model configuration for reranking."""
|
||||
|
||||
additionalModelRequestFields: dict[str, Any] | None
|
||||
additionalModelRequestFields: dict[str, object] | None
|
||||
modelArn: str | None
|
||||
|
||||
|
||||
class BedrockKBRerankingConfiguration(TypedDict, total=False):
|
||||
"""Configuration for reranking in vector search."""
|
||||
|
||||
bedrockRerankingConfiguration: dict[str, Any] | None # This could be further typed if needed
|
||||
bedrockRerankingConfiguration: dict[str, object] | None # This could be further typed if needed
|
||||
type: str | None
|
||||
|
||||
|
||||
class BedrockKBVectorSearchConfiguration(TypedDict, total=False):
|
||||
"""Configuration for vector search."""
|
||||
|
||||
filter: dict[str, Any] | None
|
||||
filter: dict[str, object] | None
|
||||
implicitFilterConfiguration: BedrockKBImplicitFilterConfiguration | None
|
||||
numberOfResults: int | None
|
||||
overrideSearchType: str | None
|
||||
|
|
|
|||
|
|
@ -189,6 +189,8 @@ class DeploymentOptions:
|
|||
tpm: int | None = None
|
||||
itpm: int | None = None
|
||||
otpm: int | None = None
|
||||
rpd: int | None = None
|
||||
tpd: int | None = None
|
||||
kubernetes_pod_discovery: bool | None = None
|
||||
default_api_key_rpm_limit: int | None = None
|
||||
default_api_key_tpm_limit: int | None = None
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ Type definitions for AWS Bedrock Invoke Agent API responses.
|
|||
https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_InvokeAgent.html
|
||||
"""
|
||||
|
||||
from typing import Any, Final
|
||||
from typing import Final
|
||||
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
|
|
@ -39,7 +39,7 @@ class InvokeAgentModelInvocationInput(TypedDict, total=False):
|
|||
"""Model invocation input details."""
|
||||
|
||||
foundationModel: str | None
|
||||
inferenceConfiguration: dict[str, Any] | None
|
||||
inferenceConfiguration: dict[str, object] | None
|
||||
text: str | None
|
||||
traceId: str | None
|
||||
type: str | None
|
||||
|
|
@ -49,9 +49,9 @@ class InvokeAgentModelInvocationOutput(TypedDict, total=False):
|
|||
"""Model invocation output details."""
|
||||
|
||||
metadata: InvokeAgentMetadata | None
|
||||
parsedResponse: dict[str, Any] | None
|
||||
rawResponse: dict[str, Any] | None
|
||||
reasoningContent: dict[str, Any] | None
|
||||
parsedResponse: dict[str, object] | None
|
||||
rawResponse: dict[str, object] | None
|
||||
reasoningContent: dict[str, object] | None
|
||||
traceId: str | None
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
from collections.abc import Mapping
|
||||
from enum import Enum
|
||||
from typing import Any, Literal
|
||||
from typing import Literal
|
||||
|
||||
from typing_extensions import Required, TypedDict
|
||||
|
||||
|
|
@ -235,7 +236,7 @@ class GeminiVideoGenerationInstance(TypedDict, total=False):
|
|||
"""Instance data for Gemini video generation request"""
|
||||
|
||||
prompt: Required[str]
|
||||
image: dict[str, Any]
|
||||
image: Mapping[str, object]
|
||||
|
||||
|
||||
class GeminiVideoGenerationParameters(LiteLLMBaseModel):
|
||||
|
|
@ -263,7 +264,7 @@ class GeminiVideoGenerationParameters(LiteLLMBaseModel):
|
|||
negativePrompt: str | None = None
|
||||
"""Text describing what not to include in the video."""
|
||||
|
||||
lastFrame: Any | None = None
|
||||
lastFrame: object | None = None
|
||||
"""
|
||||
The final image for interpolation video to transition.
|
||||
Should be used with the 'image' parameter.
|
||||
|
|
@ -275,7 +276,7 @@ class GeminiVideoGenerationParameters(LiteLLMBaseModel):
|
|||
Only supported in Veo 3.1 (list of VideoGenerationReferenceImage objects).
|
||||
"""
|
||||
|
||||
video: Any | None = None
|
||||
video: object | None = None
|
||||
"""
|
||||
Video to be used for video extension (Video object).
|
||||
Only supported in Veo 3.1 & Veo 3 Fast.
|
||||
|
|
@ -354,5 +355,5 @@ class GeminiLongRunningOperationResponse(LiteLLMBaseModel):
|
|||
response: GeminiOperationResponse | None = None
|
||||
"""Response object when operation is complete"""
|
||||
|
||||
error: dict[str, Any] | None = None
|
||||
error: dict[str, object] | None = None
|
||||
"""Error details if operation failed"""
|
||||
|
|
|
|||
|
|
@ -4,7 +4,8 @@ Type definitions for LangGraph API.
|
|||
LangGraph provides a streaming and non-streaming API for running agents.
|
||||
"""
|
||||
|
||||
from typing import Any, Literal
|
||||
from collections.abc import Mapping
|
||||
from typing import Literal
|
||||
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
|
|
@ -29,8 +30,8 @@ class LangGraphRequest(TypedDict, total=False):
|
|||
assistant_id: str
|
||||
input: LangGraphInput
|
||||
stream_mode: str | None
|
||||
config: dict[str, Any] | None
|
||||
metadata: dict[str, Any] | None
|
||||
config: Mapping[str, object] | None
|
||||
metadata: Mapping[str, object] | None
|
||||
|
||||
|
||||
# Response Types - Streaming
|
||||
|
|
@ -38,7 +39,7 @@ class LangGraphStreamEvent(TypedDict, total=False):
|
|||
"""Single event in a LangGraph stream response."""
|
||||
|
||||
event: str
|
||||
data: Any
|
||||
data: object
|
||||
|
||||
|
||||
# Response Types - Non-streaming
|
||||
|
|
@ -55,7 +56,7 @@ class LangGraphResponse(TypedDict, total=False):
|
|||
"""Non-streaming response structure from LangGraph."""
|
||||
|
||||
messages: list[LangGraphResponseMessage]
|
||||
values: dict[str, Any]
|
||||
values: dict[str, object]
|
||||
|
||||
|
||||
# Parsed response for internal use
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
from collections.abc import Callable
|
||||
from typing import Any, ClassVar
|
||||
from typing import ClassVar
|
||||
|
||||
from pydantic import ConfigDict, Field
|
||||
|
||||
|
|
@ -10,7 +10,7 @@ class MCPTool(LiteLLMBaseModel):
|
|||
model_config: ClassVar[ConfigDict] = ConfigDict(arbitrary_types_allowed=True)
|
||||
name: str
|
||||
description: str
|
||||
input_schema: dict[str, Any]
|
||||
input_schema: dict[str, object]
|
||||
handler: Callable
|
||||
server_id: str | None = Field(default=None, frozen=True)
|
||||
|
||||
|
|
@ -18,18 +18,18 @@ class MCPTool(LiteLLMBaseModel):
|
|||
class ToolSchema(LiteLLMBaseModel):
|
||||
name: str
|
||||
description: str
|
||||
inputSchema: dict[str, Any]
|
||||
inputSchema: dict[str, object]
|
||||
|
||||
|
||||
class ListToolsResponse(LiteLLMBaseModel):
|
||||
tools: list[ToolSchema]
|
||||
nextCursor: str | None = None
|
||||
_meta: dict[str, Any] | None = None
|
||||
_meta: dict[str, object] | None = None
|
||||
|
||||
|
||||
class CallToolRequest(LiteLLMBaseModel):
|
||||
method: str = "tools/call"
|
||||
params: dict[str, Any]
|
||||
params: dict[str, object]
|
||||
|
||||
|
||||
class ContentItem(LiteLLMBaseModel):
|
||||
|
|
|
|||
|
|
@ -1,5 +1,3 @@
|
|||
from typing import Any
|
||||
|
||||
from pydantic import Field
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
|
|
@ -20,7 +18,7 @@ class EnkryptAIPolicyViolationDetail(TypedDict, total=False):
|
|||
class EnkryptAIPIIDetail(TypedDict, total=False):
|
||||
"""Details for PII detection."""
|
||||
|
||||
pii: dict[str, Any]
|
||||
pii: dict[str, object]
|
||||
|
||||
|
||||
class EnkryptAIToxicityDetail(TypedDict, total=False):
|
||||
|
|
@ -85,8 +83,8 @@ class EnkryptAIResponseDetails(TypedDict, total=False):
|
|||
toxicity: EnkryptAIToxicityDetail | None
|
||||
keyword_detected: EnkryptAIKeywordDetail | None
|
||||
bias: EnkryptAIBiasDetail | None
|
||||
prompt_injection: dict[str, Any] | None
|
||||
jailbreak: dict[str, Any] | None
|
||||
prompt_injection: dict[str, object] | None
|
||||
jailbreak: dict[str, object] | None
|
||||
|
||||
|
||||
class EnkryptAIResponse(TypedDict, total=False):
|
||||
|
|
@ -107,7 +105,7 @@ class EnkryptAIProcessedResult(TypedDict):
|
|||
| EnkryptAIToxicityDetail
|
||||
| EnkryptAIKeywordDetail
|
||||
| EnkryptAIBiasDetail
|
||||
| dict[str, Any],
|
||||
| dict[str, object],
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -65,6 +65,34 @@ class MetricWithMetadata(MetricBase):
|
|||
api_key_breakdown: dict[str, KeyMetricWithMetadata] = Field(default_factory=dict) # api_key -> {metrics, metadata}
|
||||
|
||||
|
||||
class GuardrailSpendMetrics(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
spend: float = 0.0
|
||||
untracked_units: int = 0
|
||||
|
||||
|
||||
class GuardrailKeyMetric(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
metrics: GuardrailSpendMetrics
|
||||
metadata: KeyMetadata
|
||||
|
||||
|
||||
class GuardrailMetadata(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
guardrail_name: str
|
||||
|
||||
|
||||
class GuardrailMetric(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
metrics: GuardrailSpendMetrics
|
||||
metadata: GuardrailMetadata
|
||||
api_key_breakdown: Mapping[str, GuardrailKeyMetric] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class BreakdownMetrics(LiteLLMBaseModel):
|
||||
"""Breakdown of spend by different dimensions"""
|
||||
|
||||
|
|
@ -75,6 +103,7 @@ class BreakdownMetrics(LiteLLMBaseModel):
|
|||
endpoints: dict[str, MetricWithMetadata] = Field(default_factory=dict) # endpoint -> {metrics, metadata}
|
||||
api_keys: dict[str, KeyMetricWithMetadata] = Field(default_factory=dict) # api_key -> {metrics, metadata}
|
||||
entities: dict[str, MetricWithMetadata] = Field(default_factory=dict) # entity -> {metrics, metadata}
|
||||
guardrails: Mapping[str, GuardrailMetric] = Field(default_factory=dict) # guardrail_id -> spend
|
||||
|
||||
|
||||
class DailySpendData(LiteLLMBaseModel):
|
||||
|
|
|
|||
|
|
@ -168,6 +168,16 @@ class EntityRollupRow(RollupMetricsRow):
|
|||
distinct_api_keys: int | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class GuardrailSpendRow:
|
||||
date: str
|
||||
guardrail_id: str
|
||||
guardrail_name: str
|
||||
api_key: str
|
||||
spend: float
|
||||
untracked_units: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AggregatedRows:
|
||||
grouping_rows: tuple[GroupingSetsRow, ...]
|
||||
|
|
|
|||
|
|
@ -533,6 +533,8 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
|
|||
rpm: int | None = None
|
||||
itpm: int | None = None
|
||||
otpm: int | None = None
|
||||
tpd: int | None = None
|
||||
rpd: int | None = None
|
||||
timeout: float | str | httpx.Timeout | SDKTimeout | None = None # if str, pass in as os.environ/
|
||||
# read timeout of a live stream, a faked stream takes the longer of this and timeout. If str, pass in as os.environ/
|
||||
stream_timeout: float | str | None = None
|
||||
|
|
@ -738,6 +740,8 @@ class LiteLLMParamsTypedDict(TypedDict, total=False):
|
|||
rpm: int | None
|
||||
itpm: int | None
|
||||
otpm: int | None
|
||||
tpd: ReadOnly[int | None]
|
||||
rpd: ReadOnly[int | None]
|
||||
order: int | None
|
||||
weight: int | None
|
||||
max_parallel_requests: int | None
|
||||
|
|
|
|||
|
|
@ -10418,6 +10418,13 @@ class ProviderConfigManager:
|
|||
return get_modelscope_image_generation_config(model)
|
||||
elif LlmProviders.EDENAI == provider:
|
||||
return litellm.EdenAIImageGenerationConfig()
|
||||
elif LlmProviders.TENCENT == provider:
|
||||
from litellm.llms.tencent.image_generation.transformation import (
|
||||
TencentImageGenerationConfig,
|
||||
is_tokenhub_image_model,
|
||||
)
|
||||
|
||||
return TencentImageGenerationConfig() if is_tokenhub_image_model(model) else None
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -10606,6 +10613,15 @@ class ProviderConfigManager:
|
|||
)
|
||||
|
||||
return get_openrouter_image_edit_config(model)
|
||||
elif LlmProviders.TENCENT == provider:
|
||||
from litellm.llms.tencent.image_edit.transformation import (
|
||||
TencentImageEditConfig,
|
||||
)
|
||||
from litellm.llms.tencent.image_generation.transformation import (
|
||||
is_tokenhub_image_model,
|
||||
)
|
||||
|
||||
return TencentImageEditConfig() if is_tokenhub_image_model(model) else None
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -21589,7 +21589,13 @@
|
|||
],
|
||||
"supports_vision": true,
|
||||
"supports_audio_input": true,
|
||||
"supports_video_input": true
|
||||
"supports_video_input": true,
|
||||
"input_cost_per_token_priority": 5.6252e-07,
|
||||
"output_cost_per_token_priority": 3.37498e-06,
|
||||
"cache_read_input_token_cost_priority": 5.628e-08,
|
||||
"cache_creation_input_token_cost_priority": 5.6252e-07,
|
||||
"input_cost_per_audio_token_priority": 1.12497e-06,
|
||||
"cache_read_input_audio_token_cost_priority": 1.1249e-07
|
||||
},
|
||||
"databricks/databricks-gemini-3-1-flash-image": {
|
||||
"cache_creation_input_token_cost": 6.2503e-07,
|
||||
|
|
@ -21676,7 +21682,14 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_video_input": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 4.50002e-06,
|
||||
"output_cost_per_token_priority": 2.699998e-05,
|
||||
"cache_read_input_token_cost_priority": 4.5003e-07,
|
||||
"cache_creation_input_token_cost_priority": 4.50002e-06,
|
||||
"input_cost_per_token_above_200k_tokens_priority": 8.99997e-06,
|
||||
"output_cost_per_token_above_200k_tokens_priority": 4.049997e-05,
|
||||
"cache_read_input_token_cost_above_200k_tokens_priority": 8.9999e-07
|
||||
},
|
||||
"databricks/databricks-gemini-3-flash": {
|
||||
"cache_creation_input_token_cost": 6.2503e-07,
|
||||
|
|
@ -21710,7 +21723,13 @@
|
|||
],
|
||||
"supports_vision": true,
|
||||
"supports_audio_input": true,
|
||||
"supports_video_input": true
|
||||
"supports_video_input": true,
|
||||
"input_cost_per_token_priority": 1.12497e-06,
|
||||
"output_cost_per_token_priority": 6.75003e-06,
|
||||
"cache_read_input_token_cost_priority": 1.1249e-07,
|
||||
"cache_creation_input_token_cost_priority": 1.12497e-06,
|
||||
"input_cost_per_audio_token_priority": 2.25001e-06,
|
||||
"cache_read_input_audio_token_cost_priority": 2.2498e-07
|
||||
},
|
||||
"databricks/databricks-gemini-3-pro": {
|
||||
"cache_creation_input_token_cost": 2.49998e-06,
|
||||
|
|
@ -21757,7 +21776,11 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.35002e-06,
|
||||
"output_cost_per_token_priority": 6.75003e-06,
|
||||
"cache_read_input_token_cost_priority": 1.3503e-07,
|
||||
"cache_creation_input_token_cost_priority": 1.35002e-06
|
||||
},
|
||||
"databricks/databricks-gemini-3-7-flash": {
|
||||
"cache_creation_input_token_cost": 7.4998e-07,
|
||||
|
|
@ -21785,10 +21808,14 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.35002e-06,
|
||||
"output_cost_per_token_priority": 6.75003e-06,
|
||||
"cache_read_input_token_cost_priority": 1.3503e-07,
|
||||
"cache_creation_input_token_cost_priority": 1.35002e-06
|
||||
},
|
||||
"databricks/databricks-gemini-3-6-flash": {
|
||||
"cache_creation_input_token_cost": 1.87502e-06,
|
||||
"cache_creation_input_token_cost": 7.4998e-07,
|
||||
"cache_read_input_token_cost": 7.497e-08,
|
||||
"input_cost_per_token": 7.4998e-07,
|
||||
"input_dbu_cost_per_token": 1.0714e-05,
|
||||
|
|
@ -21816,7 +21843,11 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.35002e-06,
|
||||
"output_cost_per_token_priority": 6.75003e-06,
|
||||
"cache_read_input_token_cost_priority": 1.3503e-07,
|
||||
"cache_creation_input_token_cost_priority": 1.35002e-06
|
||||
},
|
||||
"databricks/databricks-gemini-3-5-flash": {
|
||||
"cache_creation_input_token_cost": 1.87502e-06,
|
||||
|
|
@ -21846,7 +21877,11 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 3.37498e-06,
|
||||
"output_cost_per_token_priority": 2.025002e-05,
|
||||
"cache_read_input_token_cost_priority": 3.3747e-07,
|
||||
"cache_creation_input_token_cost_priority": 3.37498e-06
|
||||
},
|
||||
"databricks/databricks-gemini-3-5-flash-lite": {
|
||||
"cache_creation_input_token_cost": 3.7499e-07,
|
||||
|
|
@ -21878,7 +21913,11 @@
|
|||
"supports_response_schema": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 6.7501e-07,
|
||||
"output_cost_per_token_priority": 5.62499e-06,
|
||||
"cache_read_input_token_cost_priority": 6.748e-08,
|
||||
"cache_creation_input_token_cost_priority": 6.7501e-07
|
||||
},
|
||||
"databricks/databricks-gemma-3-12b": {
|
||||
"cache_creation_input_token_cost": 1.5001e-07,
|
||||
|
|
@ -21921,7 +21960,11 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false
|
||||
"supports_vision": false,
|
||||
"input_cost_per_token_priority": 2.45e-06,
|
||||
"output_cost_per_token_priority": 7.7e-06,
|
||||
"cache_read_input_token_cost_priority": 4.55e-07,
|
||||
"cache_creation_input_token_cost_priority": 2.45e-06
|
||||
},
|
||||
"databricks/databricks-glm-5-3": {
|
||||
"cache_creation_input_token_cost": 1.4e-06,
|
||||
|
|
@ -21951,7 +21994,11 @@
|
|||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false,
|
||||
"thinking_always_on": true
|
||||
"thinking_always_on": true,
|
||||
"input_cost_per_token_priority": 2.45e-06,
|
||||
"output_cost_per_token_priority": 7.7e-06,
|
||||
"cache_read_input_token_cost_priority": 4.55e-07,
|
||||
"cache_creation_input_token_cost_priority": 2.45e-06
|
||||
},
|
||||
"databricks/databricks-glm-5-3-flash": {
|
||||
"cache_creation_input_token_cost": 1.5001e-07,
|
||||
|
|
@ -21980,7 +22027,11 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"thinking_always_on": true
|
||||
"thinking_always_on": true,
|
||||
"input_cost_per_token_priority": 2.625e-07,
|
||||
"output_cost_per_token_priority": 8.75e-07,
|
||||
"cache_read_input_token_cost_priority": 5.25e-08,
|
||||
"cache_creation_input_token_cost_priority": 2.625e-07
|
||||
},
|
||||
"databricks/databricks-gpt-5": {
|
||||
"cache_creation_input_token_cost": 1.24999e-06,
|
||||
|
|
@ -21999,7 +22050,11 @@
|
|||
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_tool_choice": true
|
||||
"supports_tool_choice": true,
|
||||
"input_cost_per_token_priority": 2.49998e-06,
|
||||
"output_cost_per_token_priority": 1.999998e-05,
|
||||
"cache_read_input_token_cost_priority": 2.4997e-07,
|
||||
"cache_creation_input_token_cost_priority": 2.49998e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-1": {
|
||||
"cache_creation_input_token_cost": 1.24999e-06,
|
||||
|
|
@ -22018,7 +22073,11 @@
|
|||
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_tool_choice": true
|
||||
"supports_tool_choice": true,
|
||||
"input_cost_per_token_priority": 2.49998e-06,
|
||||
"output_cost_per_token_priority": 1.999998e-05,
|
||||
"cache_read_input_token_cost_priority": 2.4997e-07,
|
||||
"cache_creation_input_token_cost_priority": 2.49998e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-2": {
|
||||
"cache_creation_input_token_cost": 1.75e-06,
|
||||
|
|
@ -22037,7 +22096,11 @@
|
|||
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_tool_choice": true
|
||||
"supports_tool_choice": true,
|
||||
"input_cost_per_token_priority": 3.5e-06,
|
||||
"output_cost_per_token_priority": 2.8e-05,
|
||||
"cache_read_input_token_cost_priority": 3.5e-07,
|
||||
"cache_creation_input_token_cost_priority": 3.5e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-3-codex": {
|
||||
"cache_creation_input_token_cost": 1.75e-06,
|
||||
|
|
@ -22056,7 +22119,11 @@
|
|||
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_tool_choice": true
|
||||
"supports_tool_choice": true,
|
||||
"input_cost_per_token_priority": 3.5e-06,
|
||||
"output_cost_per_token_priority": 2.8e-05,
|
||||
"cache_read_input_token_cost_priority": 3.5e-07,
|
||||
"cache_creation_input_token_cost_priority": 3.5e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-4": {
|
||||
"cache_creation_input_token_cost": 2.49998e-06,
|
||||
|
|
@ -22087,7 +22154,14 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 5.00003e-06,
|
||||
"output_cost_per_token_priority": 2.999997e-05,
|
||||
"cache_read_input_token_cost_priority": 5.0001e-07,
|
||||
"cache_creation_input_token_cost_priority": 5.00003e-06,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 5.00003e-06,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 2.999997e-05,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 5.0001e-07
|
||||
},
|
||||
"databricks/databricks-gpt-5-4-mini": {
|
||||
"cache_creation_input_token_cost": 7.4998e-07,
|
||||
|
|
@ -22115,7 +22189,11 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.50003e-06,
|
||||
"output_cost_per_token_priority": 8.99997e-06,
|
||||
"cache_read_input_token_cost_priority": 1.5001e-07,
|
||||
"cache_creation_input_token_cost_priority": 1.50003e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-4-nano": {
|
||||
"cache_creation_input_token_cost": 1.9999e-07,
|
||||
|
|
@ -22175,7 +22253,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 8.00002e-06,
|
||||
"output_cost_per_token_priority": 4.000003e-05,
|
||||
"cache_read_input_token_cost_priority": 8.0003e-07,
|
||||
"cache_creation_input_token_cost_priority": 9.99999e-06,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 1.599997e-05,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 6.000001e-05,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 1.59999e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 1.999998e-05
|
||||
},
|
||||
"databricks/databricks-gpt-5-6-terra": {
|
||||
"cache_creation_input_token_cost": 2.49998e-06,
|
||||
|
|
@ -22207,7 +22293,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 4.00001e-06,
|
||||
"output_cost_per_token_priority": 2.399999e-05,
|
||||
"cache_read_input_token_cost_priority": 3.9998e-07,
|
||||
"cache_creation_input_token_cost_priority": 5.00003e-06,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 8.00002e-06,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 3.600002e-05,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 8.0003e-07,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-6-luna": {
|
||||
"cache_creation_input_token_cost": 2.4997e-07,
|
||||
|
|
@ -22239,7 +22333,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 3.9998e-07,
|
||||
"output_cost_per_token_priority": 2.40002e-06,
|
||||
"cache_read_input_token_cost_priority": 3.997e-08,
|
||||
"cache_creation_input_token_cost_priority": 5.0001e-07,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 8.0003e-07,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 3.60003e-06,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 8.001e-08,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 1.00002e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-5": {
|
||||
"cache_creation_input_token_cost": 5.00003e-06,
|
||||
|
|
@ -22270,7 +22372,14 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.249997e-05,
|
||||
"output_cost_per_token_priority": 7.50001e-05,
|
||||
"cache_read_input_token_cost_priority": 1.24999e-06,
|
||||
"cache_creation_input_token_cost_priority": 1.249997e-05,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 1.249997e-05,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 7.50001e-05,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 1.24999e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-5-pro": {
|
||||
"cache_creation_input_token_cost": 2.999997e-05,
|
||||
|
|
@ -22319,7 +22428,11 @@
|
|||
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_tool_choice": true
|
||||
"supports_tool_choice": true,
|
||||
"input_cost_per_token_priority": 4.5003e-07,
|
||||
"output_cost_per_token_priority": 3.60003e-06,
|
||||
"cache_read_input_token_cost_priority": 4.501e-08,
|
||||
"cache_creation_input_token_cost_priority": 4.5003e-07
|
||||
},
|
||||
"databricks/databricks-gpt-5-nano": {
|
||||
"cache_creation_input_token_cost": 4.998e-08,
|
||||
|
|
@ -22469,7 +22582,11 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 5.25e-06,
|
||||
"output_cost_per_token_priority": 2.625e-05,
|
||||
"cache_read_input_token_cost_priority": 5.25e-07,
|
||||
"cache_creation_input_token_cost_priority": 5.25e-06
|
||||
},
|
||||
"databricks/system.ai.kimi-k3": {
|
||||
"litellm_provider": "databricks",
|
||||
|
|
@ -22551,7 +22668,11 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false,
|
||||
"thinking_always_on": true
|
||||
"thinking_always_on": true,
|
||||
"input_cost_per_token_priority": 3.85e-07,
|
||||
"output_cost_per_token_priority": 3.85e-06,
|
||||
"cache_read_input_token_cost_priority": 3.85e-07,
|
||||
"cache_creation_input_token_cost_priority": 3.85e-07
|
||||
},
|
||||
"databricks/databricks-qwen3-next-80b-a3b-instruct": {
|
||||
"cache_creation_input_token_cost": 1.5001e-07,
|
||||
|
|
@ -61865,6 +61986,39 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"tencent/seedream-image-v5.0-flash": {
|
||||
"litellm_provider": "tencent",
|
||||
"mode": "image_generation",
|
||||
"output_cost_per_image": 0.018,
|
||||
"output_cost_per_image_token": 1.6e-06,
|
||||
"source": "https://cloud.tencent.com/document/product/1823/130055",
|
||||
"supported_endpoints": [
|
||||
"/v1/images/generations",
|
||||
"/v1/images/edits"
|
||||
]
|
||||
},
|
||||
"tencent/seedream-image-v5.0-lite": {
|
||||
"litellm_provider": "tencent",
|
||||
"mode": "image_generation",
|
||||
"output_cost_per_image": 0.035,
|
||||
"output_cost_per_image_token": 1.6e-06,
|
||||
"source": "https://www.tencentcloud.com/document/product/1300/78937",
|
||||
"supported_endpoints": [
|
||||
"/v1/images/generations",
|
||||
"/v1/images/edits"
|
||||
]
|
||||
},
|
||||
"tencent/seedream-image-v5.0-pro": {
|
||||
"litellm_provider": "tencent",
|
||||
"mode": "image_generation",
|
||||
"output_cost_per_image": 0.09,
|
||||
"output_cost_per_image_token": 1.6e-06,
|
||||
"source": "https://www.tencentcloud.com/document/product/1300/78937",
|
||||
"supported_endpoints": [
|
||||
"/v1/images/generations",
|
||||
"/v1/images/edits"
|
||||
]
|
||||
},
|
||||
"cognition/swe-1.6": {
|
||||
"input_cost_per_token": 5e-07,
|
||||
"output_cost_per_token": 2.5e-06,
|
||||
|
|
@ -77774,7 +77928,11 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 5.25e-07,
|
||||
"output_cost_per_token_priority": 2.1e-06,
|
||||
"cache_read_input_token_cost_priority": 5.25e-08,
|
||||
"cache_creation_input_token_cost_priority": 5.25e-07
|
||||
},
|
||||
"databricks/databricks-deepseek-v4-1-flash-1017": {
|
||||
"cache_creation_input_token_cost": 3.0002e-07,
|
||||
|
|
@ -77897,7 +78055,11 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_video_input": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 4.50002e-06,
|
||||
"output_cost_per_token_priority": 2.250003e-05,
|
||||
"cache_read_input_token_cost_priority": 2.2498e-07,
|
||||
"cache_creation_input_token_cost_priority": 4.50002e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-4-pro": {
|
||||
"cache_creation_input_token_cost": 2.999997e-05,
|
||||
|
|
@ -77959,7 +78121,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 4.00001e-06,
|
||||
"output_cost_per_token_priority": 1.99997e-05,
|
||||
"cache_read_input_token_cost_priority": 1.9999e-07,
|
||||
"cache_creation_input_token_cost_priority": 5.00003e-06,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 8.00002e-06,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 2.99999e-05,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 3.9998e-07,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06
|
||||
},
|
||||
"databricks/databricks-gpt-6-1-terra": {
|
||||
"cache_creation_input_token_cost": 1.24999e-06,
|
||||
|
|
@ -77991,7 +78161,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.99997e-06,
|
||||
"output_cost_per_token_priority": 9.99999e-06,
|
||||
"cache_read_input_token_cost_priority": 1.0003e-07,
|
||||
"cache_creation_input_token_cost_priority": 2.49998e-06,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 4.00001e-06,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 1.500002e-05,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 1.9999e-07,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 5.00003e-06
|
||||
},
|
||||
"databricks/databricks-gpt-6-astra": {
|
||||
"cache_creation_input_token_cost": 1.249997e-05,
|
||||
|
|
@ -78023,7 +78201,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.999998e-05,
|
||||
"output_cost_per_token_priority": 9.99999e-05,
|
||||
"cache_read_input_token_cost_priority": 1.99997e-06,
|
||||
"cache_creation_input_token_cost_priority": 2.500001e-05,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 4.000003e-05,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 0.0001500002,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 4.00001e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 5.000002e-05
|
||||
},
|
||||
"databricks/databricks-gpt-6-luna": {
|
||||
"cache_creation_input_token_cost": 1.2502e-07,
|
||||
|
|
@ -78055,7 +78241,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.9999e-07,
|
||||
"output_cost_per_token_priority": 1.0003e-06,
|
||||
"cache_read_input_token_cost_priority": 2.002e-08,
|
||||
"cache_creation_input_token_cost_priority": 2.4997e-07,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 3.9998e-07,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 1.5001e-06,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 3.997e-08,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 5.0001e-07
|
||||
},
|
||||
"databricks/databricks-gpt-6-sol": {
|
||||
"cache_creation_input_token_cost": 2.49998e-06,
|
||||
|
|
@ -78087,7 +78281,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 4.00001e-06,
|
||||
"output_cost_per_token_priority": 1.99997e-05,
|
||||
"cache_read_input_token_cost_priority": 3.9998e-07,
|
||||
"cache_creation_input_token_cost_priority": 5.00003e-06,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 8.00002e-06,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 2.99999e-05,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 8.0003e-07,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06
|
||||
},
|
||||
"databricks/databricks-gpt-image-1": {
|
||||
"cache_creation_input_token_cost": 5.00003e-06,
|
||||
|
|
|
|||
|
|
@ -3793,8 +3793,8 @@
|
|||
"responses": true,
|
||||
"text_completion": true,
|
||||
"embeddings": false,
|
||||
"image_generations": false,
|
||||
"image_edits": false,
|
||||
"image_generations": true,
|
||||
"image_edits": true,
|
||||
"audio_transcriptions": false,
|
||||
"audio_speech": false,
|
||||
"rerank": false,
|
||||
|
|
|
|||
205
skills/litellm-setup/SKILL.md
Normal file
205
skills/litellm-setup/SKILL.md
Normal file
|
|
@ -0,0 +1,205 @@
|
|||
---
|
||||
name: litellm-setup
|
||||
description: Set up a self-hosted LiteLLM proxy, the OpenAI-compatible LLM gateway, on this machine, from install through provider keys, config.yaml, a master key, optional Postgres and Admin UI, and a verified chat completion. Use when the user asks to set up LiteLLM, install the LiteLLM proxy, self-host an LLM gateway, or put one endpoint in front of OpenAI, Anthropic, Gemini, Azure, Bedrock, Ollama, or other providers. Not for using the LiteLLM Python SDK inside an app
|
||||
---
|
||||
|
||||
# Set up a LiteLLM proxy
|
||||
|
||||
Take the user from nothing to a running LiteLLM proxy that has answered one chat completion. Run every command yourself unless a step says the user does it
|
||||
|
||||
Never print a secret. Provider keys, the master key, the salt key, the database URL, and virtual keys live only in `.env`, and `config.yaml` reads them through `os.environ/NAME`. Check that a value is there with `grep -c '^NAME=.' .env`, never with `cat .env` or `echo "$NAME"`. Load `.env` into a shell with this line, which exports each line as it is, so a `&` or `$` in a value stays literal, and first drops any database settings the shell already had and sets `LITELLM_MODE=PRODUCTION`, which stops `litellm` from also loading a `.env` it finds above its install folder, such as `~/.env`, so only this `.env` picks the database. The steps that need it start with it:
|
||||
|
||||
```bash
|
||||
unset DATABASE_URL DIRECT_URL DATABASE_HOST; export LITELLM_MODE=PRODUCTION; while IFS= read -r line || [ -n "$line" ]; do case "$line" in [A-Za-z_]*=*) export "$line" ;; esac; done < .env
|
||||
```
|
||||
|
||||
## 1. Ask once
|
||||
|
||||
Ask everything in one message, with defaults the user can accept by saying "go":
|
||||
|
||||
- Providers and models. Default is OpenAI with its newest flagship model, using the `OPENAI_API_KEY` in your environment when `[ -n "$OPENAI_API_KEY" ] && echo set` says it is there. Run that check before asking, and say in the question that the default uses that variable, naming it but never its value, since a key in your environment can belong to another account
|
||||
- Whether they want virtual keys, teams, spend tracking, or the Admin UI. Those need Postgres, plain routing does not. Default is no
|
||||
- uv (default) or Docker. With Postgres, Docker runs the repo's quickstart, which brings its own database, and uv needs a Postgres server they already have, with a new, empty database for LiteLLM, so someone with their own Postgres server picks uv
|
||||
- Folder and port. Defaults are `./litellm-proxy` (`./litellm-gateway` for the quickstart) and 4000
|
||||
|
||||
Then follow the steps for their answers. Plain routing skips steps 3 and 6. The quickstart is steps 2, 3, 8, and 9. uv with their own Postgres skips step 3
|
||||
|
||||
In the same message, say that someone who picks uv and would rather answer prompts can run `litellm --setup` in a separate terminal window once LiteLLM is installed, not through Claude Code. Never run it yourself, since it needs a real terminal and without one prints "Setup cancelled." and writes nothing. Give them the command as `LITELLM_MODE=PRODUCTION <bin>/litellm --setup`, with `<bin>` replaced by the folder `uv tool dir --bin` prints, to run from the folder they want the config in. It writes `litellm_config.yaml` there with the provider keys and master key in plain text and offers to start the proxy, which then listens on every network interface and stops with Ctrl+C in that window. When the user is done, run this in that folder. It makes the file readable only by them, keeps it out of git (the wizard does neither), and prints only the model names. Never print that file or run `grep`, `sed`, or `cat` on it:
|
||||
|
||||
```bash
|
||||
chmod 600 litellm_config.yaml; grep -qx 'litellm_config.yaml' .gitignore 2>/dev/null || echo 'litellm_config.yaml' >> .gitignore
|
||||
"$(uv tool dir)/litellm/bin/python" -I -c 'import yaml; print(*(m["model_name"] for m in yaml.safe_load(open("litellm_config.yaml"))["model_list"]), sep="\n")'
|
||||
```
|
||||
|
||||
Then skip to step 8 and run it against their proxy, starting that shell command with this line instead of the loader line, which reads the master key without printing it:
|
||||
|
||||
```bash
|
||||
export LITELLM_MASTER_KEY="$("$(uv tool dir)/litellm/bin/python" -I -c 'import yaml; print(yaml.safe_load(open("litellm_config.yaml"))["general_settings"]["master_key"])')"
|
||||
```
|
||||
|
||||
## 2. Check prerequisites
|
||||
|
||||
```bash
|
||||
uv --version; docker info --format '{{.ServerVersion}}'; docker compose version; curl --version | head -1; openssl version; command -v python3
|
||||
```
|
||||
|
||||
On native Windows, ask the user to run Claude Code inside WSL, since every command here is POSIX shell. Every path needs curl, openssl, and python3, so ask the user to install any that is missing. The uv path needs uv, which downloads a supported Python (LiteLLM supports 3.10 to 3.14) when the `--python` range in step 4 rules out the one on the machine. If uv is missing, offer to install it, and only after the user says yes run `curl -LsSf https://astral.sh/uv/install.sh | env UV_NO_MODIFY_PATH=1 sh`, which leaves their shell profile alone (https://docs.astral.sh/uv/getting-started/installation/). It puts uv in `~/.local/bin`, so start every later command that runs `uv` or `litellm` with `export PATH="$HOME/.local/bin:$PATH";`. The Docker path needs a running Docker daemon, and the quickstart also needs Docker Compose v2. If `docker info` says `permission denied`, the user's account cannot use Docker without sudo. Offer the uv path, or ask them to add themselves to the `docker` group the way https://docs.docker.com/engine/install/linux-postinstall/ shows and start a new login session before you check again
|
||||
|
||||
Check the port is free with `curl -s -o /dev/null --max-time 2 http://127.0.0.1:4000; echo $?`, which prints `7` (could not connect) only when nothing listens there. Don't use `lsof` for this, since it misses ports that root or Docker hold. If it is taken, pick another port and use it everywhere below, because a proxy asked for a busy 4000 silently moves to a random port
|
||||
|
||||
## 3. Postgres with Docker: run the quickstart
|
||||
|
||||
The repo's quickstart starts the gateway, Postgres, and the Admin UI in Docker, listening on this machine only, and puts a `litellm-gateway` command in that folder that stops, starts, updates, or removes them. It asks nothing when it has no terminal. It writes the master key, salt key, and database password to its own `.env` and never prints them. Download it, then run it with the folder and port from step 1. It returns once the gateway is ready:
|
||||
|
||||
```bash
|
||||
curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/quickstart.sh -o quickstart.sh
|
||||
LITELLM_DIR="$PWD/litellm-gateway" LITELLM_PORT=4000 sh quickstart.sh --yes
|
||||
cd litellm-gateway
|
||||
```
|
||||
|
||||
Models live in the database here, so add each one through the API instead of a `config.yaml`, naming it the way step 5 says. Pipe the request body in, so the provider key never shows up in output or on a command line. Unless the user agreed to the key in your environment, as step 5 says, ask them to add the line `OPENAI_API_KEY=<their key>` to this `.env` themselves. Run the loader line first:
|
||||
|
||||
```bash
|
||||
python3 -c 'import json, os, sys; json.dump({"model_name": "gpt-6.1-sol", "litellm_params": {"model": "openai/gpt-6.1-sol", "api_key": os.environ["OPENAI_API_KEY"]}}, sys.stdout)' \
|
||||
| curl -sS -o /dev/null -w '%{http_code}\n' http://127.0.0.1:4000/model/new \
|
||||
-H "Authorization: Bearer $LITELLM_MASTER_KEY" -H "Content-Type: application/json" -d @-
|
||||
```
|
||||
|
||||
Run that in one shell command that starts with the loader line, since the master key is in `.env`. `200` means the model was added. Then go to step 8
|
||||
|
||||
## 4. Install
|
||||
|
||||
uv:
|
||||
|
||||
```bash
|
||||
uv tool install --python '>=3.10,<3.15' 'litellm[proxy]'
|
||||
export PATH="$(uv tool dir --bin):$PATH"; litellm --version
|
||||
```
|
||||
|
||||
With their own Postgres, install `'litellm[proxy,extra_proxy]'` instead, with the same `--python` range, which adds the Prisma client the proxy uses for the database. Running `uv tool install` again with the new extras replaces the old install. Start every later command that runs `litellm` with `export PATH="$(uv tool dir --bin):$PATH";`, so this install runs even when the shell has no `litellm` or an older one from a venv or `pip install litellm`. `uv tool update-shell` puts it on the PATH of new shells by editing the user's shell profile, so run it only after the user says yes
|
||||
|
||||
Docker:
|
||||
|
||||
```bash
|
||||
docker pull docker.litellm.ai/berriai/litellm:main-stable
|
||||
```
|
||||
|
||||
## 5. Write `.env` and `config.yaml`
|
||||
|
||||
Run every command from here on in the folder from step 1. Create `.env` readable only by the user, keep it out of git, and add a master key unless it already has one. The proxy refuses to start without one, and `sk-` is the format of every key it issues:
|
||||
|
||||
```bash
|
||||
mkdir -p litellm-proxy && cd litellm-proxy
|
||||
touch .env && chmod 600 .env; grep -qx '.env' .gitignore 2>/dev/null || echo '.env' >> .gitignore
|
||||
grep -q '^LITELLM_MASTER_KEY=.' .env || printf 'LITELLM_MASTER_KEY=sk-%s\n' "$(openssl rand -hex 32)" >> .env
|
||||
```
|
||||
|
||||
Copy a provider key that is already in your environment only after the user agreed to that variable, in their step 1 answer or when you ask now, and copy it without printing it, as in `printf 'OPENAI_API_KEY=%s\n' "$OPENAI_API_KEY" >> .env`. Otherwise ask the user to add the line `OPENAI_API_KEY=<their key>` to `.env` themselves and tell you when it is there. If they paste a key into the chat anyway, write it to `.env` and never repeat it. Keep one `NAME=value` per line with no quotes, so `docker --env-file` and the loader line read the same value
|
||||
|
||||
Confirm each model's current LiteLLM name in the cost map instead of trusting memory, since model names go stale fast. This prints the chat models of the provider named last (`openai`, `anthropic`, `gemini`, `azure`, `bedrock`, `vertex_ai`, and so on) the way `config.yaml` takes them, with the provider in front. It matches by prefix, so `bedrock` also lists the newer Bedrock models and `vertex_ai` lists every Vertex family. For Ollama, use a model from `ollama list` instead:
|
||||
|
||||
```bash
|
||||
curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json | python3 -c '
|
||||
import json, sys
|
||||
for name, row in json.load(sys.stdin).items():
|
||||
provider = str(row.get("litellm_provider", "")) if isinstance(row, dict) else ""
|
||||
if provider.startswith(sys.argv[1]) and row.get("mode") == "chat":
|
||||
prefix = "bedrock" if provider == "bedrock_converse" else "vertex_ai" if provider.startswith("vertex_ai") else provider
|
||||
print(name if name.startswith(prefix + "/") else prefix + "/" + name)
|
||||
' openai
|
||||
```
|
||||
|
||||
In `config.yaml`, `model` is a name from that list, as in `openai/gpt-6.1-sol` or `bedrock/us.anthropic.<name>`, and `model_name` is what clients will send. If the folder already has a `config.yaml`, show the user its `model_name` lines and ask before replacing it. This example uses the newest GPT model in the cost map when this skill was written, so replace it with what the user picked:
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: gpt-6.1-sol
|
||||
litellm_params:
|
||||
model: openai/gpt-6.1-sol
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
|
||||
general_settings:
|
||||
master_key: os.environ/LITELLM_MASTER_KEY
|
||||
```
|
||||
|
||||
Add one entry per model. On Azure, ask the user for each deployment's name and the model it runs. Use `azure/<deployment name>` as `model`, since Azure routes by deployment and a name from the lookup works only when the deployment has that same name, and add `model_info:` with `base_model: azure/<model it runs>` to the entry, since LiteLLM prices each call by that name. Providers that need more than a key, like Azure or Bedrock, list their `litellm_params` on their page under https://docs.litellm.ai/docs/providers. Point every secret at `.env` through `os.environ/NAME` the same way
|
||||
|
||||
## 6. Their own Postgres, on the uv path
|
||||
|
||||
Ask the user to create a new, empty database for LiteLLM on their server (`CREATE DATABASE litellm;`) and add `DATABASE_URL=postgresql://...` for it to `.env` themselves. The proxy creates its tables there on first start and stops on a database that already holds other tables. If that URL goes through a transaction pooler such as PgBouncer, they also add `DIRECT_URL` for the same database without the pooler, which the migrations use. Then add a salt key. It encrypts the provider credentials the proxy stores in the database and falls back to the master key when unset, so set it once and never change it:
|
||||
|
||||
```bash
|
||||
grep -q '^LITELLM_SALT_KEY=.' .env || printf 'LITELLM_SALT_KEY=sk-%s\n' "$(openssl rand -hex 32)" >> .env
|
||||
```
|
||||
|
||||
Generate the Prisma client once inside the tool's environment. A pip or uv install doesn't ship it generated (Docker images do):
|
||||
|
||||
```bash
|
||||
tool="$(uv tool dir)/litellm"
|
||||
schema=$("$tool/bin/python" -I -c 'import importlib.util, os; print(os.path.join(importlib.util.find_spec("litellm").submodule_search_locations[0], "proxy", "schema.prisma"))')
|
||||
PATH="$tool/bin:$PATH" prisma generate --schema "$schema"
|
||||
```
|
||||
|
||||
## 7. Start the proxy
|
||||
|
||||
Bind to `127.0.0.1` unless the user wants other machines to reach it
|
||||
|
||||
uv, in one shell command that starts with the loader line, since the installed `litellm` does not read `.env` by itself:
|
||||
|
||||
```bash
|
||||
export PATH="$HOME/.local/bin:$PATH"; export PATH="$(uv tool dir --bin):$PATH"
|
||||
nohup litellm --config config.yaml --host 127.0.0.1 --port 4000 > litellm.log 2>&1 &
|
||||
echo $! > litellm.pid
|
||||
```
|
||||
|
||||
Docker:
|
||||
|
||||
```bash
|
||||
docker run -d --name litellm --env-file .env -v "$PWD/config.yaml:/app/config.yaml:ro" -p 127.0.0.1:4000:4000 \
|
||||
docker.litellm.ai/berriai/litellm:main-stable --config /app/config.yaml --port 4000
|
||||
```
|
||||
|
||||
Wait on the health route, stopping early if the proxy exits. The first start with Postgres applies the database migrations and takes longer. Define `alive` for the path you took, `alive() { kill -0 "$(cat litellm.pid)" 2>/dev/null; }` for uv or `alive() { [ "$(docker inspect -f '{{.State.Running}}' litellm 2>/dev/null)" = true ]; }` for Docker, in the same shell command as the loop:
|
||||
|
||||
```bash
|
||||
for i in $(seq 1 180); do curl -fsS http://127.0.0.1:4000/health/liveliness >/dev/null 2>&1 && break; alive || break; sleep 1; done
|
||||
curl -sS http://127.0.0.1:4000/health/readiness
|
||||
```
|
||||
|
||||
Readiness reports `"db":"connected"` with Postgres and `"db":"Not connected"` without it. If the proxy never comes up, read the end of `litellm.log` (or `docker logs litellm`) and check the troubleshooting list
|
||||
|
||||
## 8. Verify with a chat completion
|
||||
|
||||
Use a `model_name` from `config.yaml` (or step 3), in one shell command that starts with the loader line:
|
||||
|
||||
```bash
|
||||
curl -sS http://127.0.0.1:4000/v1/chat/completions \
|
||||
-H "Authorization: Bearer $LITELLM_MASTER_KEY" -H "Content-Type: application/json" \
|
||||
-d '{"model": "gpt-6.1-sol", "messages": [{"role": "user", "content": "Say hello in five words."}]}'
|
||||
```
|
||||
|
||||
Show the user the reply in `choices[0].message.content`. An error here means the setup is not done, so fix it before handing off
|
||||
|
||||
## 9. Hand off
|
||||
|
||||
Tell the user, without printing any secret:
|
||||
|
||||
- The base URL is `http://127.0.0.1:4000`. Any OpenAI-compatible client works with it as the base URL and the master key or a virtual key as the API key
|
||||
- The master key is `LITELLM_MASTER_KEY` in the folder's `.env`, or `general_settings.master_key` in `litellm_config.yaml` when they used `litellm --setup`
|
||||
- To stop it, `kill "$(cat litellm.pid)"` (uv), `docker stop litellm` (Docker), `<folder>/litellm-gateway stop` (quickstart), or Ctrl+C in its window (`litellm --setup`). To start it again, `docker start litellm`, `<folder>/litellm-gateway start`, on uv the whole step 7 command as you ran it, loader line and PATH export included, wrapped in `( ... )` so their shell keeps its own settings, shown as a code block and never paraphrased, since a new shell has neither, or for `litellm --setup` `LITELLM_MODE=PRODUCTION <bin>/litellm --config litellm_config.yaml --host 127.0.0.1 --port <port>` in a terminal window in that folder, with `<bin>` and `<port>` filled in, since running the wizard again writes a new config with a new master key. `update` and `uninstall` update or remove a quickstart install. Give the quickstart commands with the folder's full path, since a bare `litellm-gateway` is on PATH only when `~/.local/bin` is, and can belong to an earlier install
|
||||
- With Postgres, the Admin UI is at `http://127.0.0.1:4000/ui` (user `admin`, password the master key), and `POST /key/generate` with the master key creates virtual keys with their own budgets and model lists. Without Postgres, both need step 3 or 6 first
|
||||
- The Admin UI warns that the master-key login is a shared secret, which is fine while only this machine reaches the proxy. Before anyone else can reach it, invite an admin with their own password on the Internal Users page, then set `general_settings.disable_env_credential_login: true` in the proxy config
|
||||
- To manage the running proxy from a coding agent (keys, teams, users, models, MCP servers), install `npx skills add BerriAI/litellm-skills -g`
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
- `refused to start: no master key is set` means the loader line did not run in the shell command that started `litellm`, or `--env-file .env` is missing from `docker run`
|
||||
- `a database URL is set but the prisma CLI is neither on PATH nor importable` means the install lacks `extra_proxy`, so rerun step 4 with it, or, on plain routing, that the command that started `litellm` did not start with the loader line
|
||||
- `The Client hasn't been generated yet` or `Unable to find Prisma binaries` means the Prisma generate command in step 6 has not run
|
||||
- `Incorrect API key provided` or another provider `AuthenticationError` means a provider key is wrong. Fix it in `.env` and restart, since the proxy reads the environment only at startup. On uv, stop it and rerun step 7. On Docker, `docker rm -f litellm` and rerun step 7, since `docker restart` keeps the environment the container was created with. On the quickstart the key is saved with the model, so delete the model and add it again: `GET /model/info` lists each model's `model_info.id`, and `POST /model/delete` with `{"id": "<id>"}` deletes it
|
||||
- `Cannot automatically baseline this database` means `DATABASE_URL` points at a database that already holds tables LiteLLM did not create. Point it at a new, empty database
|
||||
- `LLM Provider NOT provided` means a `model` in `config.yaml` (or step 3) lacks its provider in front. Use the name as the step 5 lookup prints it
|
||||
- `Invalid model name passed in model=...` means the request's `model` is not a `model_name` in `config.yaml` or step 3. `GET /v1/models` lists the ones the key can use
|
||||
- `No connected db.` on a proxy without Postgres means the bearer token is not the master key
|
||||
- `Conflict. The container name "/litellm" is already in use` means an earlier run left a container. `docker start litellm` reuses it, or with the user's OK `docker rm -f litellm` and rerun step 7
|
||||
- Inside Docker, `localhost` in `config.yaml` (an Ollama `api_base`, say) is the container itself. Use `host.docker.internal`, and on Linux also pass `--add-host=host.docker.internal:host-gateway`
|
||||
|
|
@ -175,6 +175,8 @@ ignored_function_names = [
|
|||
"_responses_continuation", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py
|
||||
"_responses_continuation_ceilings", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py
|
||||
"_team_deployments_named", # Tested through the team lookups in the concurrent delete tests in test_router.py
|
||||
"priority_header_update", # Tested through the priority header tests in test_common_request_processing.py
|
||||
"_drops_params", # Read by split_request_priority and priority_header_update, tested through both
|
||||
"_responses_event_announces_background_run", # Tested through the background Responses stream tests in test_router.py
|
||||
"_aresponses_background_stream_can_release", # Tested through the background Responses stream tests in test_router.py
|
||||
"_raise_released_stream_error", # Tested through the background Responses stream content-policy tests in test_router.py
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import hashlib
|
||||
import os
|
||||
import queue
|
||||
import secrets
|
||||
|
|
@ -973,6 +974,7 @@ def test_the_migration_job_refuses_before_the_first_migration_when_spend_log_key
|
|||
_insert_spend_log_with_key(conn, "legacy-1", _legacy_bearer_token())
|
||||
_insert_spend_log_with_key(conn, "legacy-2", _legacy_bearer_token())
|
||||
_insert_spend_log_with_key(conn, "hashed-1", "a" * 64)
|
||||
_insert_spend_log_with_key(conn, "repetitive-1", "z" * 3000)
|
||||
ledger_before: Final = _ledger(scratch_database)
|
||||
|
||||
with pytest.raises(RuntimeError) as refused:
|
||||
|
|
@ -1109,6 +1111,68 @@ def test_the_entry_size_check_draws_the_line_exactly_where_postgres_does(
|
|||
assert finding.detail == f"with an index entry over 2704 bytes, largest 2712 bytes ({described_as})"
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_a_value_postgres_compresses_under_the_limit_is_not_counted_while_random_text_of_that_length_is(
|
||||
scratch_database: str,
|
||||
) -> None:
|
||||
"""Postgres compresses a long value before it checks the entry size, so 3000 repeated
|
||||
bytes index fine while 3000 random hex bytes do not. The check must draw the same line."""
|
||||
create_index: Final = sql.SQL('CREATE INDEX "compress_probe_idx" ON "compress_probe" ("k")')
|
||||
with psycopg.connect(scratch_database, autocommit=True) as conn:
|
||||
conn.execute('CREATE TABLE "compress_probe" ("k" TEXT)')
|
||||
conn.execute('INSERT INTO "compress_probe" ("k") VALUES (%s), (%s)', ("a" * 3000, secrets.token_hex(1500)))
|
||||
findings: Final = migration_data_checks.find_blockers(
|
||||
conn, "public", migration_data_checks.checks_for_statement(create_index.as_string(conn), "probe")
|
||||
)
|
||||
with pytest.raises(psycopg.errors.ProgramLimitExceeded, match="index row size 3016 exceeds"):
|
||||
conn.execute(create_index)
|
||||
conn.execute('DELETE FROM "compress_probe" WHERE "k" <> %s', ("a" * 3000,))
|
||||
conn.execute(create_index)
|
||||
stored: Final = conn.execute('SELECT octet_length("k"), pg_column_size("k") FROM "compress_probe"').fetchone()
|
||||
|
||||
assert stored is not None and stored[0] == 3000 and stored[1] < 2704, stored
|
||||
(finding,) = findings
|
||||
assert isinstance(finding, migration_data_checks.BlockedChange)
|
||||
assert finding.rows == 1, "only the random key is too long for the index once Postgres compresses the other"
|
||||
assert finding.detail == "with an index entry over 2704 bytes, largest 3016 bytes (k up to 3000 bytes)"
|
||||
|
||||
|
||||
def _postgres_rejects_an_index_on(conn: "psycopg.Connection[tuple[object, ...]]", key: str) -> bool:
|
||||
conn.execute('DROP TABLE IF EXISTS "one_key"')
|
||||
conn.execute('CREATE TABLE "one_key" ("k" TEXT)')
|
||||
conn.execute('INSERT INTO "one_key" ("k") VALUES (%s)', (key,))
|
||||
try:
|
||||
conn.execute('CREATE INDEX "one_key_idx" ON "one_key" ("k")')
|
||||
except psycopg.errors.ProgramLimitExceeded:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_the_check_flags_exactly_the_compressed_keys_postgres_rejects_at_the_entry_limit(scratch_database: str) -> None:
|
||||
"""Keys that compress to a stored size just under and just over the entry limit go through the
|
||||
check and through Postgres one by one. The check must flag a key exactly when Postgres rejects it."""
|
||||
noise: Final = "".join(hashlib.sha256(str(i).encode()).hexdigest() for i in range(40))
|
||||
keys: Final = tuple(noise[:n] + "a" * 1500 for n in range(2400, 2500))
|
||||
create_index: Final = sql.SQL('CREATE INDEX "stored_probe_idx" ON "stored_probe" ("k")')
|
||||
with psycopg.connect(scratch_database, autocommit=True) as conn:
|
||||
conn.execute('CREATE TABLE "stored_probe" ("k" TEXT)')
|
||||
conn.execute('INSERT INTO "stored_probe" ("k") SELECT unnest(%s::text[])', (list(keys),))
|
||||
compressed: Final = conn.execute(
|
||||
'SELECT count(*) FROM "stored_probe" WHERE pg_column_size("k") < octet_length("k") AND octet_length("k") > 2704'
|
||||
).fetchone()
|
||||
findings: Final = migration_data_checks.find_blockers(
|
||||
conn, "public", migration_data_checks.checks_for_statement(create_index.as_string(conn), "probe")
|
||||
)
|
||||
rejected: Final = tuple(key for key in keys if _postgres_rejects_an_index_on(conn, key))
|
||||
|
||||
assert compressed == (len(keys),), "every probe key must be longer than the limit yet stored compressed"
|
||||
assert 0 < len(rejected) < len(keys), f"the probe must straddle the limit, Postgres rejected {len(rejected)}"
|
||||
(finding,) = findings
|
||||
assert isinstance(finding, migration_data_checks.BlockedChange)
|
||||
assert finding.rows == len(rejected), f"the check must flag the {len(rejected)} keys Postgres rejects and no other"
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_a_scan_that_cannot_finish_in_time_is_reported_as_unverified_rather_than_passed(scratch_database: str) -> None:
|
||||
with psycopg.connect(scratch_database, autocommit=True) as conn:
|
||||
|
|
|
|||
|
|
@ -1,16 +1,48 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from collections.abc import Iterator
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Iterator, Mapping, Sequence
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import datetime, timezone
|
||||
from decimal import Decimal
|
||||
from itertools import chain
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from typing import TYPE_CHECKING, Final, TypeAlias
|
||||
|
||||
import psycopg
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from psycopg import sql
|
||||
from pydantic import JsonValue
|
||||
|
||||
from litellm.accounting import DeliveryState, DeliveryTarget, EventKind
|
||||
from litellm.accounting import (
|
||||
CostEvent,
|
||||
DeliveryState,
|
||||
DeliveryTarget,
|
||||
EventKind,
|
||||
InboundUsageEvent,
|
||||
Metric,
|
||||
MetricName,
|
||||
OutboundUsageEvent,
|
||||
Outcome,
|
||||
PricingSnapshot,
|
||||
)
|
||||
from litellm.proxy.db import accounting_ledger
|
||||
from litellm.proxy.db.accounting_ledger import AccountingConflictError, AccountingEvent, AccountingLedger
|
||||
from litellm.types.utils import CallTypes
|
||||
from tests.integration._support.database import scratch_database
|
||||
from tests.integration._support.prisma_clients import connected_proxy_prisma
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import Prisma
|
||||
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
_Writer: TypeAlias = tuple["PrismaClient", AccountingLedger]
|
||||
|
||||
_SCHEMA: Final = Path(__file__).resolve().parents[3] / "litellm-proxy-extras/litellm_proxy_extras/schema.prisma"
|
||||
_EVENT_INSERT: Final = (
|
||||
|
|
@ -126,3 +158,377 @@ def test_new_delivery_defaults_to_pending_with_no_attempts(ledger: psycopg.Conne
|
|||
(DeliveryTarget.KEY_SPEND.value,),
|
||||
).fetchone()
|
||||
assert row == ("pending", 0)
|
||||
|
||||
|
||||
_AT: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||||
_CUTOFF: Final = datetime(2026, 2, 1, tzinfo=timezone.utc)
|
||||
_TARGETS: Final = (DeliveryTarget.USER_SPEND, DeliveryTarget.TEAM_SPEND)
|
||||
|
||||
|
||||
def _usage(event_id: str = "usage", *, inbound: bool = False) -> InboundUsageEvent | OutboundUsageEvent:
|
||||
fields: Final = dict(
|
||||
event_id=event_id,
|
||||
request_id="request",
|
||||
call_type=CallTypes.completion,
|
||||
started_at=_AT,
|
||||
ended_at=_AT,
|
||||
outcome=Outcome.SUCCEEDED,
|
||||
dimensions=(),
|
||||
)
|
||||
if inbound:
|
||||
return InboundUsageEvent(**fields, metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)),))
|
||||
return OutboundUsageEvent(**fields, attempt_id=event_id, retry_of_attempt_id=None, metrics=None)
|
||||
|
||||
|
||||
def _cost(source: str = "usage", event_id: str = "cost", amount: str = "0.125") -> CostEvent:
|
||||
return CostEvent(
|
||||
event_id=event_id,
|
||||
usage_event_id=source,
|
||||
occurred_at=_AT,
|
||||
dimensions=(),
|
||||
charges=(Metric(name=MetricName.SPEND, value=Decimal(amount), unit="USD"),),
|
||||
snapshot=PricingSnapshot(
|
||||
calculator_version="fixture",
|
||||
model="fixture",
|
||||
provider=None,
|
||||
pricing_source="cost_map",
|
||||
priced_model_key="fixture",
|
||||
total_cost=Decimal(amount),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _envelope(event: AccountingEvent, *, context: Mapping[str, JsonValue] | None = None) -> str:
|
||||
return json.dumps({"event": event.model_dump(mode="json"), "context": context or {}})
|
||||
|
||||
|
||||
async def _insert(
|
||||
store: AccountingLedger,
|
||||
event: AccountingEvent,
|
||||
targets: tuple[DeliveryTarget, ...] = _TARGETS,
|
||||
*,
|
||||
context: Mapping[str, JsonValue] | None = None,
|
||||
) -> str:
|
||||
return await store.insert_event(event, targets, envelope=_envelope(event, context=context))
|
||||
|
||||
|
||||
def _rows(database_url: str, statement: str, parameters: tuple[object, ...] = ()) -> Sequence[tuple[object, ...]]:
|
||||
with psycopg.connect(database_url) as connection:
|
||||
cursor: Final = connection.execute(statement, parameters)
|
||||
return cursor.fetchall() if cursor.description is not None else []
|
||||
|
||||
|
||||
def _complete_and_age(database_url: str) -> None:
|
||||
with psycopg.connect(database_url) as connection:
|
||||
connection.execute('UPDATE "LiteLLM_AccountingEvent" SET recorded_at = %s', (_AT,))
|
||||
connection.execute("UPDATE \"LiteLLM_AccountingDelivery\" SET state = 'applied', applied_at = %s", (_AT,))
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(loop_scope="function")
|
||||
async def writer(ledger_database: str, monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[_Writer]:
|
||||
monkeypatch.setenv("DATABASE_URL", ledger_database)
|
||||
monkeypatch.delenv("DATABASE_URL_READ_REPLICA", raising=False)
|
||||
_rows(ledger_database, 'TRUNCATE "LiteLLM_AccountingDelivery", "LiteLLM_AccountingEvent"')
|
||||
async with connected_proxy_prisma(ledger_database) as client:
|
||||
yield client, AccountingLedger.for_client(client)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_event_and_complete_delivery_set_commit_together(writer: _Writer, ledger_database: str) -> None:
|
||||
_, store = writer
|
||||
event: Final = _usage()
|
||||
assert await _insert(store, event, context={"requested_region": "eu", "quantity": "1.25"}) == "inserted"
|
||||
rows: Final = _rows(ledger_database, 'SELECT kind::text, source_event_id, payload FROM "LiteLLM_AccountingEvent"')
|
||||
assert rows == [
|
||||
("outbound_usage", None, json.loads(_envelope(event, context={"requested_region": "eu", "quantity": "1.25"})))
|
||||
]
|
||||
deliveries: Final = _rows(
|
||||
ledger_database, 'SELECT target::text, state::text, attempts FROM "LiteLLM_AccountingDelivery"'
|
||||
)
|
||||
assert set(deliveries) == {(target.value, "pending", 0) for target in _TARGETS}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_duplicates_preserve_completed_deliveries(writer: _Writer, ledger_database: str) -> None:
|
||||
_, store = writer
|
||||
event: Final = _usage(inbound=True)
|
||||
results: Final = await asyncio.gather(*(_insert(store, event) for _ in range(8)))
|
||||
assert results.count("inserted") == 1 and results.count("duplicate") == 7
|
||||
_complete_and_age(ledger_database)
|
||||
before: Final = _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY target')
|
||||
reordered: Final = json.dumps(json.loads(_envelope(event)), sort_keys=True, indent=2)
|
||||
assert await store.insert_event(event, tuple(reversed(_TARGETS)) + _TARGETS, envelope=reordered) == "duplicate"
|
||||
assert _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY target') == before
|
||||
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(1,)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("change", ["payload", "context", "targets", "schema_version", "version"])
|
||||
async def test_conflicting_acceptance_cannot_overwrite_or_add_work(
|
||||
writer: _Writer, ledger_database: str, change: str
|
||||
) -> None:
|
||||
_, store = writer
|
||||
event: Final = _usage()
|
||||
await _insert(store, event)
|
||||
altered: Final = event.model_copy(update={"outcome": Outcome.FAILED}) if change == "payload" else event
|
||||
targets: Final = (DeliveryTarget.KEY_SPEND,) if change == "targets" else _TARGETS
|
||||
context: Final = {"override": "0.5"} if change == "context" else {}
|
||||
options: Final = {change: 2} if change in ("schema_version", "version") else {}
|
||||
with pytest.raises(AccountingConflictError, match="conflicts"):
|
||||
await store.insert_event(altered, targets, envelope=_envelope(altered, context=context), **options)
|
||||
assert _rows(ledger_database, 'SELECT payload FROM "LiteLLM_AccountingEvent"') == [(json.loads(_envelope(event)),)]
|
||||
assert set(_rows(ledger_database, 'SELECT target::text FROM "LiteLLM_AccountingDelivery"')) == {
|
||||
(t.value,) for t in _TARGETS
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_competing_cost_identities_accept_only_one_authoritative_charge(
|
||||
writer: _Writer, ledger_database: str
|
||||
) -> None:
|
||||
_, store = writer
|
||||
await _insert(store, _usage())
|
||||
results: Final = await asyncio.gather(
|
||||
_insert(store, _cost()), _insert(store, _cost(event_id="other-cost")), return_exceptions=True
|
||||
)
|
||||
assert sum(result == "inserted" for result in results) == 1
|
||||
assert sum(isinstance(result, AccountingConflictError) for result in results) == 1
|
||||
accepted: Final = _rows(
|
||||
ledger_database, "SELECT event_id, payload FROM \"LiteLLM_AccountingEvent\" WHERE kind = 'cost'"
|
||||
)
|
||||
assert len(accepted) == 1
|
||||
event: Final = CostEvent.model_validate_json(json.dumps(accepted[0][1]["event"]))
|
||||
assert await _insert(store, event) == "duplicate"
|
||||
assert event.snapshot.total_cost == Decimal("0.125")
|
||||
with pytest.raises(AccountingConflictError):
|
||||
await _insert(store, _cost(event_id=event.event_id, amount="0.5"))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("source", ["missing", "inbound", "cost"])
|
||||
async def test_cost_requires_an_existing_outbound_source(writer: _Writer, ledger_database: str, source: str) -> None:
|
||||
_, store = writer
|
||||
if source == "inbound":
|
||||
await _insert(store, _usage("source", inbound=True))
|
||||
if source == "cost":
|
||||
await _insert(store, _usage())
|
||||
await _insert(store, _cost(event_id="source"))
|
||||
with pytest.raises(AccountingConflictError, match="outbound usage"):
|
||||
await _insert(store, _cost(source="source", event_id="invalid"))
|
||||
assert _rows(ledger_database, "SELECT count(*) FROM \"LiteLLM_AccountingEvent\" WHERE event_id = 'invalid'") == [
|
||||
(0,)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("failure", ["before_commit", "after_commit", "cancellation"])
|
||||
async def test_interruption_rolls_back_or_retry_recognizes_commit(
|
||||
writer: _Writer, ledger_database: str, failure: str
|
||||
) -> None:
|
||||
client, store = writer
|
||||
|
||||
@asynccontextmanager
|
||||
async def interrupted() -> AsyncGenerator[Prisma]:
|
||||
async with client.writer_db.tx() as tx:
|
||||
yield tx
|
||||
if failure == "before_commit":
|
||||
raise RuntimeError("injected pre-commit failure")
|
||||
if failure == "cancellation":
|
||||
raise asyncio.CancelledError
|
||||
raise RuntimeError("injected lost acknowledgement")
|
||||
|
||||
error: Final = asyncio.CancelledError if failure == "cancellation" else RuntimeError
|
||||
with pytest.raises(error):
|
||||
await _insert(AccountingLedger(interrupted), _usage())
|
||||
expected: Final = 1 if failure == "after_commit" else 0
|
||||
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(expected,)]
|
||||
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(expected * len(_TARGETS),)]
|
||||
assert await _insert(store, _usage()) == ("duplicate" if expected else "inserted")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delivery_insert_failure_rolls_back_event(writer: _Writer, ledger_database: str) -> None:
|
||||
from prisma.errors import RawQueryError
|
||||
|
||||
_, store = writer
|
||||
with psycopg.connect(ledger_database) as connection:
|
||||
connection.execute(
|
||||
'ALTER TABLE "LiteLLM_AccountingDelivery" ADD CONSTRAINT reject_delivery CHECK (version < 1)'
|
||||
)
|
||||
try:
|
||||
with pytest.raises(RawQueryError, match="reject_delivery"):
|
||||
await _insert(store, _usage())
|
||||
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(0,)]
|
||||
finally:
|
||||
with psycopg.connect(ledger_database) as connection:
|
||||
connection.execute('ALTER TABLE "LiteLLM_AccountingDelivery" DROP CONSTRAINT reject_delivery')
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("blocked", ["pending", "parked", "unpriced", "recent_cost", "recovery"])
|
||||
async def test_retention_preserves_incomplete_or_recovery_needed_families(
|
||||
writer: _Writer, ledger_database: str, blocked: str
|
||||
) -> None:
|
||||
_, store = writer
|
||||
await _insert(store, _usage())
|
||||
if blocked != "unpriced":
|
||||
await _insert(store, _cost())
|
||||
_complete_and_age(ledger_database)
|
||||
with psycopg.connect(ledger_database) as connection:
|
||||
if blocked in ("pending", "parked"):
|
||||
connection.execute(
|
||||
"UPDATE \"LiteLLM_AccountingDelivery\" SET state = %s WHERE event_id = 'cost'", (blocked,)
|
||||
)
|
||||
if blocked == "recent_cost":
|
||||
connection.execute(
|
||||
"UPDATE \"LiteLLM_AccountingEvent\" SET recorded_at = %s WHERE event_id = 'cost'", (_CUTOFF,)
|
||||
)
|
||||
before: Final = _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id')
|
||||
assert (
|
||||
await store.purge_completed(
|
||||
redelivery_before=_CUTOFF, recovery_before=_AT if blocked == "recovery" else _CUTOFF
|
||||
)
|
||||
== 0
|
||||
)
|
||||
assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == before
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retention_removes_completed_family_and_receipts_together(writer: _Writer, ledger_database: str) -> None:
|
||||
_, store = writer
|
||||
await _insert(store, _usage())
|
||||
await _insert(store, _cost())
|
||||
_complete_and_age(ledger_database)
|
||||
assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2
|
||||
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(0,)]
|
||||
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(0,)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retention_respects_batch_bound_and_skips_locked_families(writer: _Writer, ledger_database: str) -> None:
|
||||
client, store = writer
|
||||
for name in ("a", "b", "c"):
|
||||
await _insert(store, _usage(name, inbound=True))
|
||||
_complete_and_age(ledger_database)
|
||||
async with client.writer_db.tx() as tx:
|
||||
await tx.query_raw("SELECT event_id FROM \"LiteLLM_AccountingEvent\" WHERE event_id = 'a' FOR UPDATE")
|
||||
assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF, limit=1) == 1
|
||||
assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == [
|
||||
("a",),
|
||||
("c",),
|
||||
]
|
||||
assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retention_makes_progress_past_pending_and_unpriced_events(writer: _Writer, ledger_database: str) -> None:
|
||||
_, store = writer
|
||||
await _insert(store, _usage("a", inbound=True))
|
||||
await _insert(store, _usage("b"))
|
||||
await _insert(store, _usage("c", inbound=True))
|
||||
_complete_and_age(ledger_database)
|
||||
_rows(ledger_database, "UPDATE \"LiteLLM_AccountingDelivery\" SET state = 'pending' WHERE event_id = 'a'")
|
||||
assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF, limit=1) == 1
|
||||
assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == [
|
||||
("a",),
|
||||
("b",),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_zero_cost_and_empty_delivery_sets_are_durable(writer: _Writer, ledger_database: str) -> None:
|
||||
_, store = writer
|
||||
await _insert(store, _usage(), targets=())
|
||||
event: Final = _cost(amount="0")
|
||||
assert await _insert(store, event, targets=()) == "inserted"
|
||||
assert await _insert(store, event, targets=()) == "duplicate"
|
||||
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(0,)]
|
||||
_complete_and_age(ledger_database)
|
||||
assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retention_cannot_delete_source_while_cost_acceptance_is_uncommitted(
|
||||
writer: _Writer, ledger_database: str
|
||||
) -> None:
|
||||
client, store = writer
|
||||
await _insert(store, _usage())
|
||||
_complete_and_age(ledger_database)
|
||||
inserted: Final = asyncio.Event()
|
||||
release: Final = asyncio.Event()
|
||||
|
||||
@asynccontextmanager
|
||||
async def paused_commit() -> AsyncGenerator[Prisma]:
|
||||
async with client.writer_db.tx() as tx:
|
||||
yield tx
|
||||
inserted.set()
|
||||
await release.wait()
|
||||
|
||||
task: Final = asyncio.create_task(_insert(AccountingLedger(paused_commit), _cost()))
|
||||
try:
|
||||
await asyncio.wait_for(inserted.wait(), timeout=5)
|
||||
assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 0
|
||||
finally:
|
||||
release.set()
|
||||
await task
|
||||
assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == [
|
||||
("cost",),
|
||||
("usage",),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retention_failure_rolls_back_event_and_receipt_deletion(writer: _Writer, ledger_database: str) -> None:
|
||||
client, store = writer
|
||||
await _insert(store, _usage(inbound=True))
|
||||
_complete_and_age(ledger_database)
|
||||
|
||||
@asynccontextmanager
|
||||
async def interrupted() -> AsyncGenerator[Prisma]:
|
||||
async with client.writer_db.tx() as tx:
|
||||
yield tx
|
||||
raise RuntimeError("injected cleanup failure")
|
||||
|
||||
with pytest.raises(RuntimeError, match="cleanup failure"):
|
||||
await AccountingLedger(interrupted).purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF)
|
||||
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(1,)]
|
||||
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(len(_TARGETS),)]
|
||||
assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ledger_reads_and_writes_use_primary_when_replica_is_configured(
|
||||
writer: _Writer, ledger_database: str, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
with scratch_database("replica") as replica_url:
|
||||
monkeypatch.setenv("DATABASE_URL_READ_REPLICA", replica_url)
|
||||
async with connected_proxy_prisma(ledger_database) as client:
|
||||
store: Final = AccountingLedger.for_client(client)
|
||||
await _insert(store, _usage())
|
||||
assert await _insert(store, _usage()) == "duplicate"
|
||||
assert await _insert(store, _cost()) == "inserted"
|
||||
_complete_and_age(ledger_database)
|
||||
assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2
|
||||
assert _rows(replica_url, "SELECT to_regclass('public.\"LiteLLM_AccountingEvent\"')") == [(None,)]
|
||||
|
||||
|
||||
def _index_conditions(plan: object) -> tuple[tuple[object, object], ...]:
|
||||
if isinstance(plan, dict):
|
||||
own: Final = ((plan.get("Index Name"), plan.get("Index Cond")),) if "Index Name" in plan else ()
|
||||
children: Final = (plan.get("Plan"), *plan.get("Plans", ()))
|
||||
return own + tuple(chain.from_iterable(_index_conditions(child) for child in children if child is not None))
|
||||
return tuple(chain.from_iterable(_index_conditions(item) for item in plan)) if isinstance(plan, list) else ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("statement", [accounting_ledger._DELETE_EVENTS, accounting_ledger._DELETE_DELIVERIES])
|
||||
def test_family_cleanup_lookups_use_the_kind_source_index(ledger_database: str, statement: str) -> None:
|
||||
explained: Final = "EXPLAIN (FORMAT JSON) " + statement.replace("$1", "%s")
|
||||
with psycopg.connect(ledger_database) as connection:
|
||||
connection.execute("SET enable_seqscan = off")
|
||||
plan: Final = connection.execute(explained, ("usage",) * statement.count("$1")).fetchone()
|
||||
assert plan is not None
|
||||
conditions: Final = _index_conditions(plan[0])
|
||||
assert any(
|
||||
name == "LiteLLM_AccountingEvent_kind_source_event_id_idx" and "source_event_id" in str(condition)
|
||||
for name, condition in conditions
|
||||
), conditions
|
||||
|
|
|
|||
|
|
@ -0,0 +1,65 @@
|
|||
import json
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.client import Gateway
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
|
||||
|
||||
def _grounded_reply(request: Request) -> Reply:
|
||||
body: Final = {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {"role": "model", "parts": [{"text": "The passage"}]},
|
||||
"finishReason": "STOP",
|
||||
"groundingMetadata": {
|
||||
"groundingChunks": [
|
||||
{"retrievedContext": {"uri": "gs://documents/source"}},
|
||||
{"web": {"uri": "https://example.com/first", "title": "First"}},
|
||||
{"web": {"uri": "https://example.com/second", "title": "Second"}},
|
||||
],
|
||||
"groundingSupports": [
|
||||
{"segment": {"startIndex": 4, "endIndex": 11}, "groundingChunkIndices": [0, 99, 2, 1]}
|
||||
],
|
||||
},
|
||||
}
|
||||
],
|
||||
"usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 2, "totalTokenCount": 7},
|
||||
}
|
||||
assert json.loads(request.body)["contents"] == [{"role": "user", "parts": [{"text": "Find sources"}]}]
|
||||
if ":streamGenerateContent" in request.target:
|
||||
return Reply(content_type="text/event-stream", chunks=(f"data: {json.dumps(body)}\n\n".encode(),))
|
||||
return Reply(body=json.dumps(body).encode())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"])
|
||||
def test_chat_grounding_annotations_include_every_web_source(gateway: Gateway, stream: bool) -> None:
|
||||
with wire_server(_grounded_reply) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model="gemini/gemini-3.8-flash", api_base=wire.url, api_key="synthetic-key")
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "messages": [{"role": "user", "content": "Find sources"}], "stream": stream},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
payloads: Final = (
|
||||
[json.loads(line[6:]) for line in response.text.splitlines() if line.startswith("data: {")]
|
||||
if stream
|
||||
else [response.json()]
|
||||
)
|
||||
messages: Final = tuple(payload["choices"][0]["delta" if stream else "message"] for payload in payloads)
|
||||
annotations: Final = tuple(message["annotations"] for message in messages if message.get("annotations"))
|
||||
assert annotations == (
|
||||
[
|
||||
{
|
||||
"type": "url_citation",
|
||||
"url_citation": {
|
||||
"start_index": 4,
|
||||
"end_index": 11,
|
||||
"url": f"https://example.com/{name.lower()}",
|
||||
"title": name,
|
||||
},
|
||||
}
|
||||
for name in ("Second", "First")
|
||||
],
|
||||
), response.text
|
||||
|
|
@ -32,6 +32,7 @@ GEMINI: Final = "gemini-3-pro-preview"
|
|||
ZAI_GLM: Final = "glm-5"
|
||||
ECHOED_KEY_MODEL: Final = f"{KIMI}-echoed-key"
|
||||
BEDROCK_CLAUDE: Final = f"us.anthropic.{CLAUDE}"
|
||||
DATABRICKS_CLAUDE: Final = f"databricks-{CLAUDE}"
|
||||
AZURE_API_VERSION: Final = "2025-04-01-preview"
|
||||
SYNTHETIC_KEY: Final = "synthetic-provider-key"
|
||||
_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})")
|
||||
|
|
@ -215,6 +216,10 @@ def backend(kind: str, wire_url: str) -> Backend:
|
|||
f"bedrock/invoke/{BEDROCK_CLAUDE}",
|
||||
MappingProxyType({"api_base": wire_url, "api_key": SYNTHETIC_KEY, "aws_region_name": "us-east-1"}),
|
||||
)
|
||||
case "databricks_claude":
|
||||
return Backend(
|
||||
f"databricks/{DATABRICKS_CLAUDE}", MappingProxyType({"api_base": wire_url, "api_key": SYNTHETIC_KEY})
|
||||
)
|
||||
case _:
|
||||
raise AssertionError(f"unknown backend kind {kind}")
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,96 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from integration._support.client import gateway_from_environment
|
||||
from integration._support.openai_wire import chat_reply
|
||||
from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from integration.providers._cache_control_marks_support import owned_config
|
||||
from integration.routing.test_payment_required_mapping import OPENAI_MODEL, is_model_info_probe, model_list_reply
|
||||
from pydantic import JsonValue
|
||||
|
||||
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
|
||||
|
||||
GROUP: Final = "daily-limited-group"
|
||||
LIMITED_ID: Final = "free-tier-two-a-day"
|
||||
UNLIMITED_ID: Final = "paid-tier"
|
||||
RPD: Final = 2
|
||||
EXTRA_REQUESTS: Final = 2
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Rig:
|
||||
owned: OwnedProxy
|
||||
limited: Wire
|
||||
unlimited: Wire
|
||||
|
||||
|
||||
def _reply(request: Request) -> Reply:
|
||||
if is_model_info_probe(request):
|
||||
return model_list_reply()
|
||||
return chat_reply("chatcmpl-daily", OPENAI_MODEL, "served", stream=False)
|
||||
|
||||
|
||||
def _deployment(model_id: str, wire: Wire, **params: JsonValue) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"model_name": GROUP,
|
||||
"litellm_params": {
|
||||
"model": OPENAI_MODEL,
|
||||
"api_base": wire.url + "/v1",
|
||||
"api_key": "synthetic-openai-key",
|
||||
**params,
|
||||
},
|
||||
"model_info": {"id": model_id},
|
||||
}
|
||||
|
||||
|
||||
def _config(directory: Path, limited: Wire, unlimited: Wire) -> Path:
|
||||
return owned_config(
|
||||
directory,
|
||||
[_deployment(LIMITED_ID, limited, rpd=RPD, order=1), _deployment(UNLIMITED_ID, unlimited, order=2)],
|
||||
litellm_settings={"cache": False},
|
||||
router_settings={"optional_pre_call_checks": ["enforce_model_rate_limits"], "num_retries": 1},
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
directory: Final = tmp_path_factory.mktemp("daily-limits")
|
||||
with (
|
||||
gateway_from_environment() as gateway,
|
||||
wire_server(_reply) as limited,
|
||||
wire_server(_reply) as unlimited,
|
||||
owned_proxy_process(gateway, directory, {}, config=_config(directory, limited, unlimited)) as owned,
|
||||
):
|
||||
yield Rig(owned, limited, unlimited)
|
||||
|
||||
|
||||
def _chat(rig: Rig) -> httpx.Response:
|
||||
return rig.owned.gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": GROUP, "messages": [{"role": "user", "content": f"daily {uuid.uuid4().hex}"}]},
|
||||
)
|
||||
|
||||
|
||||
def _forwarded(wire: Wire) -> int:
|
||||
return sum(1 for request in wire.drain() if not is_model_info_probe(request))
|
||||
|
||||
|
||||
def test_requests_past_the_deployment_rpd_are_served_by_the_other_deployment_in_the_group(rig: Rig) -> None:
|
||||
responses: Final = tuple(_chat(rig) for _ in range(RPD + EXTRA_REQUESTS))
|
||||
|
||||
assert [response.status_code for response in responses] == [200] * (RPD + EXTRA_REQUESTS), [
|
||||
response.text for response in responses
|
||||
]
|
||||
assert [response.headers["x-litellm-model-id"] for response in responses] == [LIMITED_ID] * RPD + [
|
||||
UNLIMITED_ID
|
||||
] * EXTRA_REQUESTS
|
||||
assert (_forwarded(rig.limited), _forwarded(rig.unlimited)) == (RPD, EXTRA_REQUESTS)
|
||||
|
|
@ -1,4 +1,5 @@
|
|||
import asyncio
|
||||
import json
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
|
|
@ -8,9 +9,11 @@ import anthropic
|
|||
import httpx
|
||||
import openai
|
||||
import pytest
|
||||
from integration._support.anthropic_sse import message_json
|
||||
from integration._support.client import Gateway, eventually, list_value, object_value, string_value
|
||||
from integration._support.openai_wire import openai_error
|
||||
from integration._support.wire import Request, wire_server
|
||||
from integration._support.vertex import service_account_json
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from integration.routing._reasoning_effort_levels_support import (
|
||||
AZURE_API_VERSION,
|
||||
CLAUDE,
|
||||
|
|
@ -34,6 +37,7 @@ from integration.routing._reasoning_effort_levels_support import (
|
|||
marker_of,
|
||||
messages_request,
|
||||
new_marker,
|
||||
path_of,
|
||||
prompt,
|
||||
provider_reply,
|
||||
responses_request,
|
||||
|
|
@ -41,6 +45,8 @@ from integration.routing._reasoning_effort_levels_support import (
|
|||
)
|
||||
from pydantic import JsonValue
|
||||
|
||||
from litellm.constants import DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS
|
||||
|
||||
CHAT: Final = "/v1/chat/completions"
|
||||
MESSAGES: Final = "/v1/messages"
|
||||
RESPONSES: Final = "/v1/responses"
|
||||
|
|
@ -49,6 +55,7 @@ _ADAPTIVE: Final[dict[str, JsonValue]] = {"type": "adaptive"}
|
|||
_STRING_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": "2048"}
|
||||
_XHIGH_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": 9000}
|
||||
_BOOL_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": True}
|
||||
_FALSE_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": False}
|
||||
_NULL_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": None}
|
||||
_LIST_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": [2048]}
|
||||
_EMPTY_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": ""}
|
||||
|
|
@ -67,6 +74,8 @@ _KIMI_INFO: Final[Mapping[str, JsonValue]] = MappingProxyType({"reasoning_effort
|
|||
_CLAUDE_INFO: Final[Mapping[str, JsonValue]] = MappingProxyType({"reasoning_effort_levels": list(CLAUDE_LEVELS)})
|
||||
_GPT_INFO: Final[Mapping[str, JsonValue]] = MappingProxyType({"reasoning_effort_levels": list(GPT_LEVELS)})
|
||||
_UNDECLARED: Final[Mapping[str, JsonValue]] = MappingProxyType({})
|
||||
_EMPTY_DECLARATION: Final[Mapping[str, JsonValue]] = MappingProxyType({"reasoning_effort_levels": []})
|
||||
_NONE_ONLY_DECLARATION: Final[Mapping[str, JsonValue]] = MappingProxyType({"reasoning_effort_levels": ["none"]})
|
||||
_ECHOED_MAP_KEY: Final[Mapping[str, JsonValue]] = MappingProxyType(
|
||||
{"key": "azure/chat-latest", "litellm_provider": "azure"}
|
||||
)
|
||||
|
|
@ -134,7 +143,7 @@ def _invoke_chat(effort: str, thinking: Mapping[str, JsonValue], betas: tuple[st
|
|||
"messages": [{"role": "user", "content": [{"type": "text", "text": prompt(marker)}]}],
|
||||
"thinking": dict(thinking),
|
||||
"output_config": {"effort": effort},
|
||||
"max_tokens": 4096,
|
||||
"max_tokens": DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS,
|
||||
"anthropic_version": "bedrock-2023-05-31",
|
||||
"anthropic_beta": list(betas),
|
||||
}
|
||||
|
|
@ -144,6 +153,14 @@ def _with_output_config(expected: Expected, output_config: JsonValue) -> Expecte
|
|||
return lambda marker: {**expected(marker), "output_config": output_config}
|
||||
|
||||
|
||||
def _with_fields(expected: Expected, fields: Mapping[str, JsonValue]) -> Expected:
|
||||
return lambda marker: {**expected(marker), **fields}
|
||||
|
||||
|
||||
def _without(expected: Expected, field: str) -> Expected:
|
||||
return lambda marker: {key: value for key, value in expected(marker).items() if key != field}
|
||||
|
||||
|
||||
def _converse_chat(effort: str, thinking: Mapping[str, JsonValue]) -> Expected:
|
||||
return lambda marker: {
|
||||
"messages": [{"role": "user", "content": [{"text": prompt(marker)}]}],
|
||||
|
|
@ -320,6 +337,7 @@ _OPENAI_RESPONSES: Final = "/v1/responses"
|
|||
_AZURE_RESPONSES: Final = f"/openai/responses?api-version={AZURE_API_VERSION}"
|
||||
_UNSUPPORTED: Final = "does not support parameters"
|
||||
_GEMINI: Final = f"/models/{GEMINI}:generateContent"
|
||||
_DATABRICKS: Final = "/chat/completions"
|
||||
|
||||
CELLS: Final[Mapping[str, Cell]] = MappingProxyType(
|
||||
{
|
||||
|
|
@ -1750,6 +1768,139 @@ CELLS: Final[Mapping[str, Cell]] = MappingProxyType(
|
|||
_anthropic_chat_without_effort(_ADAPTIVE),
|
||||
_ANTHROPIC,
|
||||
),
|
||||
"chat-anthropic-false-budget-400": Cell(
|
||||
CHAT, "anthropic", _CLAUDE_INFO, {"thinking": _FALSE_BUDGET}, None, _ANTHROPIC, 400, _NON_NUMERIC_BUDGET
|
||||
),
|
||||
"chat-anthropic-undeclared-null-budget-400": Cell(
|
||||
CHAT, "anthropic", _UNDECLARED, {"thinking": _NULL_BUDGET}, None, _ANTHROPIC, 400, _NON_NUMERIC_BUDGET
|
||||
),
|
||||
"chat-anthropic-undeclared-false-budget-drop-params-sends-adaptive-without-effort": Cell(
|
||||
CHAT,
|
||||
"anthropic",
|
||||
_UNDECLARED,
|
||||
{"thinking": _FALSE_BUDGET, "drop_params": True},
|
||||
_anthropic_chat_without_effort(_ADAPTIVE),
|
||||
_ANTHROPIC,
|
||||
),
|
||||
"chat-anthropic-undeclared-missing-budget-sends-low": Cell(
|
||||
CHAT,
|
||||
"anthropic",
|
||||
_UNDECLARED,
|
||||
{"thinking": _MISSING_BUDGET},
|
||||
_anthropic_chat("low", _ADAPTIVE),
|
||||
_ANTHROPIC,
|
||||
),
|
||||
"messages-anthropic-undeclared-false-budget-400": Cell(
|
||||
MESSAGES, "anthropic", _UNDECLARED, {"thinking": _FALSE_BUDGET}, None, _ANTHROPIC, 400, _NON_NUMERIC_BUDGET
|
||||
),
|
||||
"messages-anthropic-empty-string-budget-400": Cell(
|
||||
MESSAGES, "anthropic", _CLAUDE_INFO, {"thinking": _EMPTY_BUDGET}, None, _ANTHROPIC, 400, _NON_NUMERIC_BUDGET
|
||||
),
|
||||
"messages-anthropic-undeclared-null-budget-drop-params-sends-adaptive-without-effort": Cell(
|
||||
MESSAGES,
|
||||
"anthropic",
|
||||
_UNDECLARED,
|
||||
{"thinking": _NULL_BUDGET, "drop_params": True},
|
||||
_native_messages(None, _ADAPTIVE),
|
||||
_ANTHROPIC,
|
||||
),
|
||||
"responses-anthropic-false-budget-400": Cell(
|
||||
RESPONSES,
|
||||
"anthropic",
|
||||
_CLAUDE_INFO,
|
||||
{"thinking": _FALSE_BUDGET},
|
||||
None,
|
||||
_ANTHROPIC,
|
||||
400,
|
||||
_NON_NUMERIC_BUDGET,
|
||||
),
|
||||
"chat-converse-undeclared-false-budget-400": Cell(
|
||||
CHAT,
|
||||
"bedrock_converse",
|
||||
_UNDECLARED,
|
||||
{"thinking": _FALSE_BUDGET},
|
||||
None,
|
||||
_CONVERSE,
|
||||
400,
|
||||
_NON_NUMERIC_BUDGET,
|
||||
),
|
||||
"messages-invoke-undeclared-null-budget-400": Cell(
|
||||
MESSAGES, "bedrock_invoke", _UNDECLARED, {"thinking": _NULL_BUDGET}, None, _INVOKE, 400, _NON_NUMERIC_BUDGET
|
||||
),
|
||||
"chat-databricks-claude-undeclared-false-budget-400": Cell(
|
||||
CHAT,
|
||||
"databricks_claude",
|
||||
_UNDECLARED,
|
||||
{"thinking": _FALSE_BUDGET},
|
||||
None,
|
||||
_DATABRICKS,
|
||||
400,
|
||||
_NON_NUMERIC_BUDGET,
|
||||
),
|
||||
"chat-anthropic-empty-declaration-budget-drop-params-sends-adaptive-without-effort": Cell(
|
||||
CHAT,
|
||||
"anthropic",
|
||||
_EMPTY_DECLARATION,
|
||||
{"thinking": _BUDGET, "drop_params": True},
|
||||
_anthropic_chat_without_effort(_ADAPTIVE),
|
||||
_ANTHROPIC,
|
||||
),
|
||||
"chat-anthropic-none-only-declaration-budget-drop-params-sends-adaptive-without-effort": Cell(
|
||||
CHAT,
|
||||
"anthropic",
|
||||
_NONE_ONLY_DECLARATION,
|
||||
{"thinking": _BUDGET, "drop_params": True},
|
||||
_anthropic_chat_without_effort(_ADAPTIVE),
|
||||
_ANTHROPIC,
|
||||
),
|
||||
"chat-anthropic-empty-declaration-budget-keeps-its-effort": Cell(
|
||||
CHAT,
|
||||
"anthropic",
|
||||
_EMPTY_DECLARATION,
|
||||
{"thinking": _BUDGET},
|
||||
_anthropic_chat("medium", _ADAPTIVE),
|
||||
_ANTHROPIC,
|
||||
),
|
||||
"chat-converse-empty-declaration-budget-drop-params-sends-adaptive-without-effort": Cell(
|
||||
CHAT,
|
||||
"bedrock_converse",
|
||||
_EMPTY_DECLARATION,
|
||||
{"thinking": _BUDGET, "drop_params": True},
|
||||
_converse_chat_without_effort(_ADAPTIVE),
|
||||
_CONVERSE,
|
||||
),
|
||||
"messages-anthropic-empty-declaration-budget-drop-params-sends-adaptive-without-effort": Cell(
|
||||
MESSAGES,
|
||||
"anthropic",
|
||||
_EMPTY_DECLARATION,
|
||||
{"thinking": _BUDGET, "drop_params": True},
|
||||
_native_messages(None, _ADAPTIVE),
|
||||
_ANTHROPIC,
|
||||
),
|
||||
"messages-anthropic-none-only-declaration-budget-drop-params-sends-adaptive-without-effort": Cell(
|
||||
MESSAGES,
|
||||
"anthropic",
|
||||
_NONE_ONLY_DECLARATION,
|
||||
{"thinking": _BUDGET, "drop_params": True},
|
||||
_native_messages(None, _ADAPTIVE),
|
||||
_ANTHROPIC,
|
||||
),
|
||||
"messages-anthropic-empty-declaration-budget-keeps-its-effort": Cell(
|
||||
MESSAGES,
|
||||
"anthropic",
|
||||
_EMPTY_DECLARATION,
|
||||
{"thinking": _BUDGET},
|
||||
_native_messages({"effort": "medium"}, _ADAPTIVE),
|
||||
_ANTHROPIC,
|
||||
),
|
||||
"messages-converse-empty-declaration-budget-drop-params-sends-adaptive-without-effort": Cell(
|
||||
MESSAGES,
|
||||
"bedrock_converse",
|
||||
_EMPTY_DECLARATION,
|
||||
{"thinking": _BUDGET, "drop_params": True},
|
||||
_converse_messages(None),
|
||||
_CONVERSE,
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
|
@ -1815,20 +1966,65 @@ def _bedrock_native_schema(body: Mapping[str, JsonValue], *, invoke: bool) -> Js
|
|||
return JSON_OBJECT.validate_json(string_value(json_schema["schema"]))
|
||||
|
||||
|
||||
_CONVERSE_NATIVE_OUTPUT: Final[dict[str, JsonValue]] = {
|
||||
"textFormat": {
|
||||
"type": "json_schema",
|
||||
"structure": {"jsonSchema": {"name": "answer", "schema": json.dumps(_NATIVE_OUTPUT_SCHEMA)}},
|
||||
}
|
||||
}
|
||||
_INVOKE_NATIVE_OUTPUT: Final[dict[str, JsonValue]] = {
|
||||
"effort": "high",
|
||||
"format": {"type": "json_schema", "schema": _NATIVE_OUTPUT_SCHEMA},
|
||||
}
|
||||
_JSON_TOOL: Final = "json_tool_call"
|
||||
_CONVERSE_TOOL_SCHEMA: Final[dict[str, JsonValue]] = {
|
||||
key: value for key, value in _NATIVE_OUTPUT_SCHEMA.items() if key != "additionalProperties"
|
||||
}
|
||||
_CONVERSE_WITH_NATIVE_OUTPUT: Final = _with_fields(
|
||||
_converse_chat("high", _ADAPTIVE), {"outputConfig": _CONVERSE_NATIVE_OUTPUT}
|
||||
)
|
||||
_INVOKE_WITH_NATIVE_OUTPUT: Final = _with_output_config(
|
||||
_invoke_chat("high", _ADAPTIVE, ("structured-outputs-2025-11-13",)), _INVOKE_NATIVE_OUTPUT
|
||||
)
|
||||
_CONVERSE_WITH_JSON_TOOL: Final = _with_fields(
|
||||
_converse_chat("high", _ADAPTIVE),
|
||||
{
|
||||
"toolConfig": {
|
||||
"tools": [
|
||||
{
|
||||
"toolSpec": {
|
||||
"name": _JSON_TOOL,
|
||||
"description": _JSON_TOOL,
|
||||
"inputSchema": {"json": _CONVERSE_TOOL_SCHEMA},
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
)
|
||||
_INVOKE_WITH_JSON_TOOL: Final = _with_fields(
|
||||
_without(_invoke_chat("high", _ADAPTIVE, ()), "anthropic_beta"),
|
||||
{"tools": [{"name": _JSON_TOOL, "input_schema": _NATIVE_OUTPUT_SCHEMA}]},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("kind", "target", "fields"),
|
||||
("kind", "target", "fields", "expected"),
|
||||
[
|
||||
pytest.param("bedrock_converse", _CONVERSE, {"thinking": _BUDGET}, id="converse-budget"),
|
||||
pytest.param(
|
||||
"bedrock_converse", _CONVERSE, {"thinking": _BUDGET}, _CONVERSE_WITH_NATIVE_OUTPUT, id="converse-budget"
|
||||
),
|
||||
pytest.param(
|
||||
"bedrock_invoke",
|
||||
_INVOKE,
|
||||
{"thinking": _BUDGET, "output_config": {"effort": "medium"}},
|
||||
_INVOKE_WITH_NATIVE_OUTPUT,
|
||||
id="invoke-output-config-beside-budget",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_a_deployment_opted_into_bedrock_native_output_keeps_it_beside_its_declared_level(
|
||||
gateway: Gateway, kind: str, target: str, fields: Mapping[str, JsonValue]
|
||||
gateway: Gateway, kind: str, target: str, fields: Mapping[str, JsonValue], expected: Expected
|
||||
) -> None:
|
||||
marker: Final = new_marker()
|
||||
response_format: Final[dict[str, JsonValue]] = {
|
||||
|
|
@ -1844,10 +2040,116 @@ def test_a_deployment_opted_into_bedrock_native_output_keeps_it_beside_its_decla
|
|||
assert response.status_code == 200, response.text[:2000]
|
||||
assert answer(marker) in response.text, response.text[:2000]
|
||||
assert [(request.method, request.target) for request in received] == [("POST", target)]
|
||||
body: Final = body_of(received[0])
|
||||
assert _bedrock_native_schema(body, invoke=target == _INVOKE) == _NATIVE_OUTPUT_SCHEMA, body
|
||||
assert {"toolConfig", "tools"}.isdisjoint(body), body
|
||||
assert sent_effort(received[0]) == "high", body
|
||||
assert body_of(received[0]) == expected(marker)
|
||||
|
||||
|
||||
_NATIVE_OUTPUT_REJECTION: Final = (
|
||||
"The model returned the following errors: output_config.format: Extra inputs are not permitted"
|
||||
)
|
||||
|
||||
|
||||
def _rejecting_native_output(request: Request) -> Reply:
|
||||
native: Final = _bedrock_native_schema(body_of(request), invoke=path_of(request).endswith("/invoke"))
|
||||
if native is None:
|
||||
return provider_reply(request)
|
||||
return Reply(status=400, body=json.dumps({"message": _NATIVE_OUTPUT_REJECTION}).encode())
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("kind", "target", "rejected", "retried"),
|
||||
[
|
||||
pytest.param(
|
||||
"bedrock_converse", _CONVERSE, _CONVERSE_WITH_NATIVE_OUTPUT, _CONVERSE_WITH_JSON_TOOL, id="converse"
|
||||
),
|
||||
pytest.param("bedrock_invoke", _INVOKE, _INVOKE_WITH_NATIVE_OUTPUT, _INVOKE_WITH_JSON_TOOL, id="invoke"),
|
||||
],
|
||||
)
|
||||
def test_the_retry_after_bedrock_rejects_native_output_keeps_the_declared_level(
|
||||
gateway: Gateway, kind: str, target: str, rejected: Expected, retried: Expected
|
||||
) -> None:
|
||||
marker: Final = new_marker()
|
||||
response_format: Final[dict[str, JsonValue]] = {
|
||||
"type": "json_schema",
|
||||
"json_schema": {"name": "answer", "schema": _NATIVE_OUTPUT_SCHEMA},
|
||||
}
|
||||
with wire_server(_rejecting_native_output) as wire, gateway.scenario() as scenario:
|
||||
model: Final = deployment(scenario, kind, wire, _CLAUDE_INFO_WITH_NATIVE_OUTPUT)
|
||||
response: Final = gateway.request(
|
||||
"POST", CHAT, chat_request(model, marker, response_format=response_format, thinking=_BUDGET)
|
||||
)
|
||||
received: Final = wire.drain()
|
||||
assert response.status_code == 200, response.text[:2000]
|
||||
assert [(request.method, request.target) for request in received] == [("POST", target), ("POST", target)]
|
||||
assert [body_of(request) for request in received] == [rejected(marker), retried(marker)]
|
||||
|
||||
|
||||
_VERTEX_PROJECT: Final = "scripted-project"
|
||||
_VERTEX_LOCATION: Final = "us-east5"
|
||||
_VERTEX_TOKEN: Final = "/_oauth/token"
|
||||
_VERTEX_RAW_PREDICT: Final = (
|
||||
f"/v1/projects/{_VERTEX_PROJECT}/locations/{_VERTEX_LOCATION}/publishers/anthropic/models/{CLAUDE}:rawPredict"
|
||||
)
|
||||
_VERTEX_ACCESS_TOKEN: Final = json.dumps({"access_token": "scripted", "token_type": "Bearer", "expires_in": 3600})
|
||||
|
||||
|
||||
def _vertex(expected: Expected) -> Expected:
|
||||
return _with_fields(_without(expected, "model"), {"anthropic_version": "vertex-2023-10-16"})
|
||||
|
||||
|
||||
def _vertex_reply(request: Request) -> Reply:
|
||||
if request.target == _VERTEX_TOKEN:
|
||||
return Reply(body=_VERTEX_ACCESS_TOKEN.encode())
|
||||
marker: Final = marker_of(request)
|
||||
return Reply(body=message_json(f"msg_{marker}", CLAUDE, answer(marker)))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("route", "fields", "expected", "status"),
|
||||
[
|
||||
pytest.param(
|
||||
CHAT,
|
||||
{"reasoning_effort": "medium"},
|
||||
_vertex(_anthropic_chat("high", {"type": "adaptive", "display": "summarized"})),
|
||||
200,
|
||||
id="chat-medium-to-high",
|
||||
),
|
||||
pytest.param(
|
||||
MESSAGES,
|
||||
{"thinking": _BUDGET},
|
||||
_vertex(_native_messages({"effort": "high"}, _ADAPTIVE)),
|
||||
200,
|
||||
id="messages-budget-to-high",
|
||||
),
|
||||
pytest.param(CHAT, {"thinking": _FALSE_BUDGET}, None, 400, id="chat-false-budget-400"),
|
||||
],
|
||||
)
|
||||
def test_a_vertex_ai_claude_deployment_gets_one_of_its_declared_levels(
|
||||
gateway: Gateway, route: str, fields: dict[str, JsonValue], expected: Expected | None, status: int
|
||||
) -> None:
|
||||
marker: Final = new_marker()
|
||||
with wire_server(_vertex_reply) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(
|
||||
model=f"vertex_ai/{CLAUDE}",
|
||||
api_base=wire.url,
|
||||
api_key=None,
|
||||
vertex_project=_VERTEX_PROJECT,
|
||||
vertex_location=_VERTEX_LOCATION,
|
||||
vertex_credentials=service_account_json(_VERTEX_PROJECT, wire.url),
|
||||
model_info=dict(_CLAUDE_INFO),
|
||||
)
|
||||
body: Final = (
|
||||
chat_request(model, marker, **fields) if route == CHAT else messages_request(model, marker, **fields)
|
||||
)
|
||||
response: Final = gateway.request("POST", route, body)
|
||||
received: Final = tuple(request for request in wire.drain() if request.target != _VERTEX_TOKEN)
|
||||
assert response.status_code == status, response.text[:2000]
|
||||
if expected is None:
|
||||
assert _NON_NUMERIC_BUDGET in response.text, response.text[:2000]
|
||||
assert received == (), [request.body[:300] for request in received]
|
||||
return
|
||||
assert answer(marker) in response.text, response.text[:2000]
|
||||
assert [(request.method, request.target) for request in received] == [("POST", _VERTEX_RAW_PREDICT)]
|
||||
assert body_of(received[0]) == expected(marker)
|
||||
|
||||
|
||||
def _single(received: tuple[Request, ...], marker: str) -> Request:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,67 @@
|
|||
"""An SDK caller that reuses one ``metadata=`` dict across ``litellm.completion`` calls to a Bedrock deployment whose
|
||||
``guardrailConfig`` has ``trace: enabled`` is billed, on each call, for that call's guardrail trace only, and the
|
||||
caller's dict is left as the caller wrote it (LIT-10239)"""
|
||||
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.client import eventually
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
import litellm
|
||||
|
||||
_MODEL: Final = "bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0"
|
||||
_PROMPT_TOKENS: Final = 5
|
||||
_COMPLETION_TOKENS: Final = 1
|
||||
_USAGE: Final[Mapping[str, int]] = {"contentPolicyUnits": 1, "topicPolicyUnits": 1}
|
||||
_GUARDRAIL_IDS: Final = ("guardrail-one", "guardrail-two")
|
||||
|
||||
|
||||
def _bedrock(request: Request) -> Reply:
|
||||
"""Bedrock's Converse runtime: a one-token answer and the trace of the guardrail the request named."""
|
||||
guardrail_id: Final = json.loads(request.body)["guardrailConfig"]["guardrailIdentifier"]
|
||||
body: Final[dict[str, JsonValue]] = {
|
||||
"output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}},
|
||||
"stopReason": "end_turn",
|
||||
"usage": {"inputTokens": _PROMPT_TOKENS, "outputTokens": _COMPLETION_TOKENS, "totalTokens": 6},
|
||||
"metrics": {"latencyMs": 3},
|
||||
"trace": {"guardrail": {"inputAssessment": {guardrail_id: {"invocationMetrics": {"usage": dict(_USAGE)}}}}},
|
||||
}
|
||||
return Reply(body=json.dumps(body).encode())
|
||||
|
||||
|
||||
def _guardrail_cost() -> float:
|
||||
rates: Final = litellm.model_cost["bedrock/guardrails"]["guardrail_cost_per_unit"]
|
||||
return sum(units * rates[counter] for counter, units in _USAGE.items())
|
||||
|
||||
|
||||
def _llm_cost() -> float:
|
||||
prices: Final = litellm.model_cost[_MODEL.removeprefix("bedrock/")]
|
||||
return _PROMPT_TOKENS * prices["input_cost_per_token"] + _COMPLETION_TOKENS * prices["output_cost_per_token"]
|
||||
|
||||
|
||||
def test_each_call_sharing_one_metadata_dict_is_billed_its_own_trace_only() -> None:
|
||||
billed: list[float] = [] # mutable-ok: the success callback appends what each call was billed
|
||||
litellm.success_callback.append(
|
||||
lambda kwargs, response, start, end: billed.append(kwargs["standard_logging_object"]["response_cost"])
|
||||
)
|
||||
shared_metadata: Final = {"tag": "x"}
|
||||
with wire_server(_bedrock) as wire:
|
||||
for guardrail_id in _GUARDRAIL_IDS:
|
||||
litellm.completion(
|
||||
model=_MODEL,
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
api_base=wire.url,
|
||||
aws_access_key_id="AKIASCRIPTED",
|
||||
aws_secret_access_key="scripted",
|
||||
aws_region_name="us-east-1",
|
||||
guardrailConfig={"guardrailIdentifier": guardrail_id, "guardrailVersion": "1", "trace": "enabled"},
|
||||
metadata=shared_metadata,
|
||||
)
|
||||
costs: Final = eventually(lambda: tuple(billed), lambda seen: len(seen) == len(_GUARDRAIL_IDS), seconds=30)
|
||||
|
||||
assert costs == pytest.approx((_llm_cost() + _guardrail_cost(),) * len(_GUARDRAIL_IDS))
|
||||
assert shared_metadata == {"tag": "x"}
|
||||
|
|
@ -12,6 +12,7 @@ from pydantic import JsonValue, TypeAdapter
|
|||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.router_utils.cooldown_handlers import async_get_cooldown_deployments
|
||||
from litellm.types.router import DeploymentTypedDict
|
||||
|
||||
_MODEL: Final = "gpt-5.4"
|
||||
|
|
@ -118,3 +119,30 @@ async def test_cooldown_skips_only_the_failed_duplicate_model_deployment() -> No
|
|||
"/primary/chat/completions",
|
||||
"/backup/chat/completions",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"router_kwargs",
|
||||
[{"fallbacks": [{"primary": "backup"}]}, {"default_fallbacks": "backup"}],
|
||||
ids=["string_fallback_target", "string_default_fallbacks"],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_string_valued_fallbacks_still_cool_down_the_failing_primary(router_kwargs: dict[str, object]) -> None:
|
||||
with wire_server(_peer) as wire:
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
_deployment("primary", f"{wire.url}/primary", "primary-deployment"),
|
||||
_deployment("backup", f"{wire.url}/backup", "backup-deployment"),
|
||||
],
|
||||
allowed_fails=0,
|
||||
num_retries=0,
|
||||
**router_kwargs,
|
||||
)
|
||||
with pytest.raises(litellm.AuthenticationError):
|
||||
await router.acompletion(
|
||||
model="primary",
|
||||
messages=[{"role": "user", "content": f"cooldown {uuid.uuid4().hex}"}],
|
||||
)
|
||||
cooled: Final = await async_get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None)
|
||||
|
||||
assert tuple(cooled) == ("primary-deployment",)
|
||||
|
|
|
|||
|
|
@ -413,3 +413,71 @@ def test_resent_tool_history_is_scanned_once(single_worker: Rig, shape: str) ->
|
|||
second: Final = timed(lambda: single_worker.proxy.request("POST", path, following))
|
||||
assert first.value.status_code == 200 and second.value.status_code == 200, (first.value.text, second.value.text)
|
||||
assert_within_factor_of_control(second.seconds, control=first.seconds, factor=0.5)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("placement", ("guardrail", "global"))
|
||||
def test_skip_tool_message_in_guardrail_forwards_tool_results(tmp_path: Path, placement: str) -> None:
|
||||
def configure(configuration: dict[str, JsonValue]) -> None:
|
||||
guardrails: Final = configuration["guardrails"]
|
||||
settings: Final = configuration["litellm_settings"]
|
||||
assert isinstance(guardrails, list) and isinstance(guardrails[0], dict) and isinstance(settings, dict)
|
||||
params: Final = guardrails[0]["litellm_params"]
|
||||
assert isinstance(params, dict)
|
||||
(params if placement == "guardrail" else settings)["skip_tool_message_in_guardrail"] = True
|
||||
|
||||
with tool_rig(tmp_path, configure=configure) as rig:
|
||||
rig.wire.drain()
|
||||
replies: Final = (
|
||||
rig.proxy.request(
|
||||
"POST",
|
||||
"/v1/messages",
|
||||
{
|
||||
"model": "claude",
|
||||
"max_tokens": 32,
|
||||
"cache": {"no-cache": True},
|
||||
"messages": [
|
||||
{"role": "user", "content": TEXT},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "tool_use", "id": "tool_1", "name": "read_file", "input": {"env": TEXT}}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "tool_result", "tool_use_id": "tool_1", "content": TEXT}],
|
||||
},
|
||||
],
|
||||
},
|
||||
),
|
||||
rig.proxy.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": "gpt",
|
||||
"cache": {"no-cache": True},
|
||||
"messages": [
|
||||
{"role": "user", "content": TEXT},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {"name": "read_file", "arguments": json.dumps({"env": TEXT})},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": TEXT},
|
||||
],
|
||||
},
|
||||
),
|
||||
)
|
||||
assert all(reply.status_code == 200 for reply in replies), [reply.text for reply in replies]
|
||||
sent: Final = rig.wire.drain()
|
||||
assert len(sent) == 2
|
||||
for request in sent:
|
||||
messages = json.loads(request.body)["messages"]
|
||||
assert SECRET not in json.dumps(messages[:2])
|
||||
assert SECRET in json.dumps(messages[2])
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
|
|
@ -26,7 +26,7 @@
|
|||
},
|
||||
"results": [
|
||||
{
|
||||
"breakdown": {
|
||||
"breakdown": {"guardrails": {},
|
||||
"api_keys": {
|
||||
"key-a": {
|
||||
"metadata": {
|
||||
|
|
|
|||
|
|
@ -26,7 +26,7 @@
|
|||
},
|
||||
"results": [
|
||||
{
|
||||
"breakdown": {
|
||||
"breakdown": {"guardrails": {},
|
||||
"api_keys": {
|
||||
"key-a": {
|
||||
"metadata": {
|
||||
|
|
|
|||
|
|
@ -26,7 +26,7 @@
|
|||
},
|
||||
"results": [
|
||||
{
|
||||
"breakdown": {
|
||||
"breakdown": {"guardrails": {},
|
||||
"api_keys": {
|
||||
"key-a": {
|
||||
"metadata": {
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue