mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-09 08:18:12 +00:00
Merge branch 'main' into fix/sandbox-async-for-guards
This commit is contained in:
commit
11a7320381
92 changed files with 5247 additions and 1105 deletions
|
|
@ -45,7 +45,7 @@ jobs:
|
|||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
|
||||
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14", "3.15"]
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
|
|
@ -132,7 +132,7 @@ jobs:
|
|||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
|
||||
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14", "3.15"]
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
|
|
|
|||
8
.github/workflows/test-merge-smoke.yml
vendored
8
.github/workflows/test-merge-smoke.yml
vendored
|
|
@ -33,7 +33,7 @@ jobs:
|
|||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
|
||||
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14", "3.15"]
|
||||
env:
|
||||
LITELLM_LOCAL_MODEL_COST_MAP: "True"
|
||||
steps:
|
||||
|
|
@ -76,6 +76,12 @@ jobs:
|
|||
- name: Import and CLI checks
|
||||
run: $IN_NETNS .venv/bin/python .github/scripts/run_merge_smoke.py cli
|
||||
|
||||
- name: Routing and UUID compatibility
|
||||
run: >-
|
||||
$IN_NETNS .venv/bin/python -m pytest
|
||||
tests/unit/router_strategy/test_complexity_router.py::TestCustomDimensions
|
||||
tests/unit/test_uuid_helper.py -q
|
||||
|
||||
- name: Proxy startup check
|
||||
run: $IN_NETNS .venv/bin/python .github/scripts/run_merge_smoke.py proxy-startup --diagnostics-dir "${RUNNER_TEMP}/smoke-diagnostics"
|
||||
|
||||
|
|
|
|||
|
|
@ -1,10 +1,10 @@
|
|||
# syntax=docker/dockerfile:1.7
|
||||
|
||||
# Base image for building
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
|
||||
|
||||
# Runtime image
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
|
||||
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
|
||||
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
|
||||
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
|
||||
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
|
||||
|
||||
FROM $UV_IMAGE AS uvbin
|
||||
|
|
@ -124,7 +124,12 @@ ENV HOME=/home/nonroot \
|
|||
PYTHONUNBUFFERED=1 \
|
||||
PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries
|
||||
|
||||
COPY --from=builder --chown=nonroot:nonroot /app /app
|
||||
COPY --from=builder --chown=nonroot:nonroot /app/.venv /app/.venv
|
||||
COPY --from=builder --chown=nonroot:nonroot /app/litellm /app/litellm
|
||||
COPY --from=builder --chown=nonroot:nonroot /app/enterprise /app/enterprise
|
||||
COPY --from=builder --chown=nonroot:nonroot /app/backend /app/backend
|
||||
COPY --from=builder --chown=nonroot:nonroot /app/docker/component_entrypoint.sh /app/docker/component_entrypoint.sh
|
||||
COPY --from=builder --chown=nonroot:nonroot /app/schema.prisma /app/schema.prisma
|
||||
COPY --from=builder /opt/prisma /opt/prisma
|
||||
|
||||
RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
|
||||
|
|
|
|||
|
|
@ -1,10 +1,10 @@
|
|||
# syntax=docker/dockerfile:1.7
|
||||
|
||||
# Base image for building
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
|
||||
|
||||
# Runtime image
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
|
||||
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
|
||||
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
|
||||
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
|
||||
|
|
|
|||
|
|
@ -1,8 +1,8 @@
|
|||
# syntax=docker/dockerfile:1.7
|
||||
|
||||
# Base images
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
|
||||
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
|
||||
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
|
||||
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
|
||||
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
|
||||
# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x)
|
||||
ARG PGBOUNCER_VERSION=1.26.0
|
||||
|
|
@ -140,7 +140,12 @@ ENV HOME=/home/nonroot \
|
|||
PYTHONUNBUFFERED=1 \
|
||||
PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries
|
||||
|
||||
COPY --from=builder --chown=nonroot:nonroot /app /app
|
||||
COPY --from=builder --chown=nonroot:nonroot /app/.venv /app/.venv
|
||||
COPY --from=builder --chown=nonroot:nonroot /app/litellm /app/litellm
|
||||
COPY --from=builder --chown=nonroot:nonroot /app/enterprise /app/enterprise
|
||||
COPY --from=builder --chown=nonroot:nonroot /app/gateway /app/gateway
|
||||
COPY --from=builder --chown=nonroot:nonroot /app/docker/component_entrypoint.sh /app/docker/component_entrypoint.sh
|
||||
COPY --from=builder --chown=nonroot:nonroot /app/schema.prisma /app/schema.prisma
|
||||
COPY --from=builder /opt/prisma /opt/prisma
|
||||
COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer
|
||||
|
||||
|
|
|
|||
8
litellm-rust/Cargo.lock
generated
8
litellm-rust/Cargo.lock
generated
|
|
@ -8202,18 +8202,18 @@ dependencies = [
|
|||
|
||||
[[package]]
|
||||
name = "webpki-root-certs"
|
||||
version = "1.0.10"
|
||||
version = "1.0.11"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "58ec80e73fd8baf57bf0d9615554c75348bda990d6d5bbeab964dd068319fecb"
|
||||
checksum = "e49d0e65abc00c142e8f9428da991da9f1a250c39be97459901eb12b9b49f62a"
|
||||
dependencies = [
|
||||
"rustls-pki-types",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "webpki-roots"
|
||||
version = "1.0.10"
|
||||
version = "1.0.11"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1c1c8e20c8a38fc7d65abfe74eb7280e4aa667e8eec8bc84996d2874810748d9"
|
||||
checksum = "f5f280bbb2b287bc8e7de7e3852945a0ff50d1bbcb5ed42b6c7a568c546df171"
|
||||
dependencies = [
|
||||
"rustls-pki-types",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -1,10 +1,14 @@
|
|||
"""
|
||||
Internal unified UUID helper.
|
||||
|
||||
Always uses fastuuid for performance.
|
||||
Uses fastuuid where supported and the standard library on Python 3.15+.
|
||||
"""
|
||||
|
||||
import fastuuid as _uuid
|
||||
import sys
|
||||
import uuid as _uuid
|
||||
|
||||
if sys.version_info < (3, 15):
|
||||
import fastuuid as _uuid
|
||||
|
||||
# Expose a module-like alias so callers can use: uuid.uuid4()
|
||||
uuid = _uuid
|
||||
|
|
|
|||
|
|
@ -42,7 +42,7 @@ class OtelDestination(LiteLLMBaseModel):
|
|||
success_sampling_rate: float | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Share of the request trees forwarded, 0.0..1.0, drawn once per request when its root span ends; "
|
||||
"Share of the request trees forwarded, 0.0..1.0, picked by a hash of the trace id; "
|
||||
"``None`` forwards every one."
|
||||
),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,13 +1,14 @@
|
|||
"""Provider / exporter factory + the Baggage span processor."""
|
||||
|
||||
import queue
|
||||
import random
|
||||
import threading
|
||||
import time
|
||||
import weakref
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
|
||||
from hashlib import sha256
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Literal
|
||||
from typing import TYPE_CHECKING, Final, Literal, TypeAlias
|
||||
|
||||
from opentelemetry import _logs, baggage, metrics, trace
|
||||
from opentelemetry._logs import Logger, LoggerProvider, NoOpLoggerProvider
|
||||
|
|
@ -486,15 +487,15 @@ def is_llm_call_span(span: ReadableSpan) -> bool:
|
|||
return GenAI.OPERATION_NAME in attributes and MCP.METHOD_NAME not in attributes
|
||||
|
||||
|
||||
#: Request trees held back for a sampled destination until the last span of the trace
|
||||
#: Request trees held back for a sampled destination until the last span of the request
|
||||
#: still open ends, the verdicts kept for the spans that start after that (a post-call
|
||||
#: database write), and the spans in flight per trace. All bounded, so a span that never
|
||||
#: ends, or a flood of requests, cannot hold spans for ever: the oldest tree is decided
|
||||
#: on what it has.
|
||||
#: database write), the spans in flight per request, and the request each open span
|
||||
#: belongs to. All bounded, so a span that never ends, or a flood of requests, cannot
|
||||
#: hold spans for ever: the oldest tree is decided on what it has.
|
||||
_MAX_PENDING_TREES: Final = 1024
|
||||
_MAX_PENDING_SPANS_PER_TREE: Final = 512
|
||||
_MAX_REMEMBERED_VERDICTS: Final = 4096
|
||||
_MAX_OPEN_TRACES: Final = 16384
|
||||
_MAX_OPEN_REQUESTS: Final = 16384
|
||||
|
||||
|
||||
class _PendingTree:
|
||||
|
|
@ -512,25 +513,35 @@ def _is_sampled(destination: "OtelDestination") -> bool:
|
|||
return destination.success_sampling_rate is not None or destination.error_sampling_rate is not None
|
||||
|
||||
|
||||
#: A verdict is one request tree's draw for one destination at its rates: the exporter
|
||||
#: A request is its trace id and the id of its root, the span of it whose parent is
|
||||
#: remote or missing. Calls sharing a caller's trace are separate requests, each decided
|
||||
#: on its own outcome.
|
||||
_RequestKey: TypeAlias = tuple[int, int]
|
||||
|
||||
#: A verdict is one request's keep or drop for one destination at its rates: the exporter
|
||||
#: ``cache_key`` leaves the rates out, so two destinations sharing an exporter with
|
||||
#: different rates would otherwise share one draw.
|
||||
_VerdictKey = tuple[int, object, float | None, float | None]
|
||||
#: different rates would otherwise share one verdict.
|
||||
_VerdictKey: TypeAlias = tuple[int, int, object, float | None, float | None]
|
||||
|
||||
|
||||
def _verdict_key(trace_id: int, destination: "OtelDestination") -> _VerdictKey:
|
||||
return (trace_id, destination.cache_key(), destination.success_sampling_rate, destination.error_sampling_rate)
|
||||
def _verdict_key(request: _RequestKey, destination: "OtelDestination") -> _VerdictKey:
|
||||
return (*request, destination.cache_key(), destination.success_sampling_rate, destination.error_sampling_rate)
|
||||
|
||||
|
||||
def _keeps_tree(destination: "OtelDestination", failed: bool, draw: Callable[[], float]) -> bool:
|
||||
"""Whether one draw against the destination's rate keeps a request tree, read the
|
||||
way the legacy Arize callback reads ``arize_success_sampling_rate`` /
|
||||
``arize_error_sampling_rate``: a tree with a failed span answers to the error
|
||||
rate, an unset rate keeps everything and ``0.0`` keeps nothing."""
|
||||
def _keeps_tree(destination: "OtelDestination", failed: bool, trace_id: int) -> bool:
|
||||
"""Whether the destination's rate keeps a request tree of ``trace_id``: a tree with
|
||||
a failed span answers to ``arize_error_sampling_rate``, the rest to
|
||||
``arize_success_sampling_rate``, an unset rate keeps everything and ``0.0`` keeps
|
||||
nothing.
|
||||
|
||||
The first eight bytes of a SHA-256 of the trace id are read against the rate, so
|
||||
every process deciding a request of the trace with the same outcome reaches the same
|
||||
verdict."""
|
||||
rate: Final = destination.error_sampling_rate if failed else destination.success_sampling_rate
|
||||
if rate is None:
|
||||
return True
|
||||
return rate > 0.0 and draw() <= rate
|
||||
digest: Final = sha256(trace_id.to_bytes(16, "big")).digest()
|
||||
return int.from_bytes(digest[:8], "big") < round(rate * 2**64)
|
||||
|
||||
|
||||
def _in_scope(span: ReadableSpan, scope: "OtelSpanScope") -> bool:
|
||||
|
|
@ -667,10 +678,8 @@ class TenantFanOutSpanProcessor(SpanProcessor):
|
|||
default_capture_by_backend: Mapping[str, str] = MappingProxyType({}),
|
||||
pending_drains: int = _MAX_PENDING_DRAINS,
|
||||
drain_pool: _DrainPool | None = None,
|
||||
sampling_draw: Callable[[], float] | None = None,
|
||||
) -> None:
|
||||
self._operator_sinks: Final = operator_sinks
|
||||
self._draw: Final = sampling_draw if sampling_draw is not None else random.random
|
||||
self._excluded_db_systems: Final = excluded_db_systems
|
||||
self._default_capture: Final = default_capture
|
||||
self._default_capture_by_backend: Final = default_capture_by_backend
|
||||
|
|
@ -682,24 +691,32 @@ class TenantFanOutSpanProcessor(SpanProcessor):
|
|||
self._processors: OrderedDict[object, SpanProcessor] = OrderedDict() # mutable-ok: bounded LRU
|
||||
self._retired: OrderedDict[int, SpanProcessor] = OrderedDict() # mutable-ok: drains as exports finish
|
||||
self._exporting: dict[int, int] = {} # mutable-ok: per-processor in-flight export count
|
||||
self._pending: OrderedDict[int, _PendingTree] = OrderedDict() # mutable-ok: bounded, by trace id
|
||||
self._open: OrderedDict[int, int] = OrderedDict() # mutable-ok: bounded, spans in flight by trace id
|
||||
self._pending: OrderedDict[_RequestKey, _PendingTree] = OrderedDict() # mutable-ok: bounded, by request
|
||||
self._open: OrderedDict[_RequestKey, set[int]] = OrderedDict() # mutable-ok: bounded, open span ids by request
|
||||
self._roots: dict[int, _RequestKey] = {} # mutable-ok: holds only the spans ``_open`` holds, request by span id
|
||||
self._span_roots: weakref.WeakKeyDictionary[SDKSpan, int] = weakref.WeakKeyDictionary()
|
||||
self._verdicts: OrderedDict[_VerdictKey, bool] = OrderedDict() # mutable-ok: bounded LRU
|
||||
self._drain: Final = drain_pool if drain_pool is not None else _DrainPool(capacity=pending_drains)
|
||||
|
||||
def on_start(self, span: SDKSpan, parent_context: Context | None = None) -> None:
|
||||
context: Final = span.context
|
||||
trace_id: Final = context.trace_id if context is not None else 0
|
||||
with self._lock:
|
||||
self._open[trace_id] = self._open.get(trace_id, 0) + 1
|
||||
while len(self._open) > _MAX_OPEN_TRACES:
|
||||
self._open.popitem(last=False)
|
||||
request: Final = self._request_of(span, trace.get_current_span(parent_context))
|
||||
if context is not None:
|
||||
self._span_roots[span] = request[1]
|
||||
self._roots[context.span_id] = request
|
||||
self._open.setdefault(request, set()).add(context.span_id)
|
||||
while len(self._open) > _MAX_OPEN_REQUESTS:
|
||||
for forgotten in self._open.popitem(last=False)[1]:
|
||||
self._roots.pop(forgotten, None)
|
||||
|
||||
def on_end(self, span: ReadableSpan) -> None:
|
||||
suppressed: Final = suppressed_backends()
|
||||
attributes: Final = span.attributes or _NO_ATTRIBUTES
|
||||
context: Final = span.context
|
||||
trace_id: Final = context.trace_id if context is not None else 0
|
||||
with self._lock:
|
||||
request: Final = self._request_of(span)
|
||||
if span.context is not None:
|
||||
self._roots.pop(span.context.span_id, None)
|
||||
failed: Final = span.status.status_code is StatusCode.ERROR
|
||||
for destination in request_destinations():
|
||||
if (
|
||||
|
|
@ -716,11 +733,32 @@ class TenantFanOutSpanProcessor(SpanProcessor):
|
|||
if not _is_sampled(destination):
|
||||
self._forward(span, destination)
|
||||
continue
|
||||
for held_span, held_destination in self._route(trace_id, span, destination, failed):
|
||||
for held_span, held_destination in self._route(request, span, destination, failed):
|
||||
self._forward(held_span, held_destination)
|
||||
for held_span, held_destination in self._settle(trace_id, failed):
|
||||
for held_span, held_destination in self._settle(request, span, failed):
|
||||
self._forward(held_span, held_destination)
|
||||
|
||||
def _request_of(self, span: ReadableSpan, parent_span: Span | None = None) -> _RequestKey:
|
||||
"""The request ``span`` belongs to, read under ``_lock``: the request recorded while
|
||||
it is open, else the span itself when its parent is remote or missing, else its
|
||||
parent's request, recorded while the parent is open or on ``parent_span`` for a
|
||||
child starting under a parent that has ended, else the parent itself."""
|
||||
context: Final = span.context
|
||||
if context is None:
|
||||
return (0, 0)
|
||||
recorded: Final = self._roots.get(context.span_id)
|
||||
if recorded is not None:
|
||||
return recorded
|
||||
parent: Final = span.parent
|
||||
if parent is None or parent.is_remote:
|
||||
return (context.trace_id, context.span_id)
|
||||
parent_request: Final = self._roots.get(parent.span_id)
|
||||
if parent_request is not None:
|
||||
return parent_request
|
||||
if isinstance(parent_span, SDKSpan):
|
||||
return (context.trace_id, self._span_roots.get(parent_span, parent.span_id))
|
||||
return (context.trace_id, parent.span_id)
|
||||
|
||||
def _default_capture_for(self, destination: "OtelDestination") -> str:
|
||||
backend: Final = destination.callback_name
|
||||
return (
|
||||
|
|
@ -743,14 +781,14 @@ class TenantFanOutSpanProcessor(SpanProcessor):
|
|||
self._release(processor)
|
||||
|
||||
def _route(
|
||||
self, trace_id: int, span: ReadableSpan, destination: "OtelDestination", failed: bool
|
||||
self, request: _RequestKey, span: ReadableSpan, destination: "OtelDestination", failed: bool
|
||||
) -> "tuple[tuple[ReadableSpan, OtelDestination], ...]":
|
||||
"""The spans to forward now for ``span`` on a sampled ``destination``: the span
|
||||
itself when its tree's verdict is already known and keeps it, nothing while the
|
||||
tree is held for the rest of its trace, or what the bounds made the fan-out decide
|
||||
tree is held for the rest of its request, or what the bounds made the fan-out decide
|
||||
early (a tree that outgrew its cap is decided on what it has, and so is the oldest
|
||||
tree once too many are waiting, so a span that never ends holds nothing back for
|
||||
ever). A failed span marks its tree before any such early draw. Once shutdown
|
||||
ever). A failed span marks its tree before any such early decision. Once shutdown
|
||||
has started flushing, nothing is held: a span ending then is decided at once,
|
||||
since no later flush would reach it.
|
||||
|
||||
|
|
@ -758,58 +796,62 @@ class TenantFanOutSpanProcessor(SpanProcessor):
|
|||
another thread between the two would leave this span in a fresh tree.
|
||||
"""
|
||||
with self._lock:
|
||||
verdict: Final = self._verdicts.get(_verdict_key(trace_id, destination))
|
||||
verdict: Final = self._verdicts.get(_verdict_key(request, destination))
|
||||
if verdict is not None:
|
||||
return ((span, destination),) if verdict else ()
|
||||
tree: Final = self._pending.setdefault(trace_id, _PendingTree())
|
||||
tree: Final = self._pending.setdefault(request, _PendingTree())
|
||||
tree.held.append((span, destination))
|
||||
if failed:
|
||||
tree.failed = True
|
||||
if self._draining or len(tree.held) >= _MAX_PENDING_SPANS_PER_TREE:
|
||||
return self._decide_locked(trace_id)
|
||||
return self._decide_locked(request)
|
||||
if len(self._pending) > _MAX_PENDING_TREES:
|
||||
oldest: Final = next(iter(self._pending))
|
||||
return self._decide_locked(oldest)
|
||||
return ()
|
||||
|
||||
def _decide(self, trace_id: int) -> "tuple[tuple[ReadableSpan, OtelDestination], ...]":
|
||||
def _decide(self, request: _RequestKey) -> "tuple[tuple[ReadableSpan, OtelDestination], ...]":
|
||||
with self._lock:
|
||||
return self._decide_locked(trace_id)
|
||||
return self._decide_locked(request)
|
||||
|
||||
def _settle(self, trace_id: int, failed: bool) -> "tuple[tuple[ReadableSpan, OtelDestination], ...]":
|
||||
"""What the trace's held tree exports once the span that just ended was the last
|
||||
of its trace still open. The model-call span ends in the post-call callback,
|
||||
after the server span, so a tree is decided when its trace goes quiet rather than
|
||||
def _settle(
|
||||
self, request: _RequestKey, span: ReadableSpan, failed: bool
|
||||
) -> "tuple[tuple[ReadableSpan, OtelDestination], ...]":
|
||||
"""What the request's held tree exports once the span that just ended was the last
|
||||
of its request still open. The model-call span ends in the post-call callback,
|
||||
after the server span, so a tree is decided when its request goes quiet rather than
|
||||
when its root ends, and a model call that fails late still answers to the error
|
||||
rate. A trace whose starts this fan-out never saw (attached mid-flight, or forgotten
|
||||
by the bound) is decided as each of its spans ends.
|
||||
rate. A request whose starts this fan-out never saw (attached mid-flight, or
|
||||
forgotten by the bound) is decided as each of its spans ends.
|
||||
"""
|
||||
with self._lock:
|
||||
tree: Final = self._pending.get(trace_id)
|
||||
tree: Final = self._pending.get(request)
|
||||
if tree is not None and failed:
|
||||
tree.failed = True
|
||||
still_open: Final = self._open.pop(trace_id, 1) - 1
|
||||
if still_open > 0:
|
||||
self._open[trace_id] = still_open
|
||||
still_open: Final = self._open.pop(request, set())
|
||||
if span.context is not None:
|
||||
still_open.discard(span.context.span_id)
|
||||
if still_open:
|
||||
self._open[request] = still_open
|
||||
return ()
|
||||
return self._decide_locked(trace_id)
|
||||
return self._decide_locked(request)
|
||||
|
||||
def _decide_locked(self, trace_id: int) -> "tuple[tuple[ReadableSpan, OtelDestination], ...]":
|
||||
"""Draw once per destination for the held tree; the spans the draws keep, and a
|
||||
verdict remembered for the spans of this trace that start later."""
|
||||
tree: Final = self._pending.pop(trace_id, None)
|
||||
def _decide_locked(self, request: _RequestKey) -> "tuple[tuple[ReadableSpan, OtelDestination], ...]":
|
||||
"""Decide once per destination for the held tree; the spans the verdicts keep, and
|
||||
a verdict remembered for the spans of this request that start later."""
|
||||
tree: Final = self._pending.pop(request, None)
|
||||
if tree is None:
|
||||
return ()
|
||||
verdicts: dict[_VerdictKey, bool] = {} # mutable-ok: one draw per destination of this tree
|
||||
verdicts: dict[_VerdictKey, bool] = {} # mutable-ok: one verdict per destination of this tree
|
||||
for _, destination in tree.held:
|
||||
key = _verdict_key(trace_id, destination)
|
||||
key = _verdict_key(request, destination)
|
||||
if key not in verdicts:
|
||||
verdicts[key] = _keeps_tree(destination, tree.failed, self._draw)
|
||||
verdicts[key] = _keeps_tree(destination, tree.failed, request[0])
|
||||
self._verdicts[key] = verdicts[key]
|
||||
while len(self._verdicts) > _MAX_REMEMBERED_VERDICTS:
|
||||
self._verdicts.popitem(last=False)
|
||||
return tuple(
|
||||
(span, destination) for span, destination in tree.held if verdicts[_verdict_key(trace_id, destination)]
|
||||
(span, destination) for span, destination in tree.held if verdicts[_verdict_key(request, destination)]
|
||||
)
|
||||
|
||||
def _operator_already_writes(
|
||||
|
|
@ -850,8 +892,8 @@ class TenantFanOutSpanProcessor(SpanProcessor):
|
|||
with self._lock:
|
||||
self._draining = True
|
||||
undecided: Final = tuple(self._pending)
|
||||
for trace_id in undecided:
|
||||
for held_span, held_destination in self._decide(trace_id):
|
||||
for request in undecided:
|
||||
for held_span, held_destination in self._decide(request):
|
||||
self._forward(held_span, held_destination)
|
||||
deadline: Final = time.monotonic() + self._drain_seconds
|
||||
with self._lock:
|
||||
|
|
|
|||
|
|
@ -17,6 +17,8 @@ from litellm.utils import supports_reasoning, supports_vision
|
|||
|
||||
from ...openai.chat.gpt_transformation import OpenAIGPTConfig
|
||||
|
||||
SUPPORTED_REASONING_EFFORTS: Final = ("minimal", "low", "medium", "high", "xhigh", "max", "ultra")
|
||||
|
||||
|
||||
class DeepSeekChatConfig(OpenAIGPTConfig):
|
||||
def get_supported_openai_params(self, model: str) -> list:
|
||||
|
|
@ -33,7 +35,7 @@ class DeepSeekChatConfig(OpenAIGPTConfig):
|
|||
optional_params: dict,
|
||||
model: str,
|
||||
drop_params: bool,
|
||||
) -> dict:
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
Map OpenAI params to DeepSeek params.
|
||||
|
||||
|
|
@ -60,6 +62,13 @@ class DeepSeekChatConfig(OpenAIGPTConfig):
|
|||
elif reasoning_effort is not None:
|
||||
optional_params["thinking"] = {"type": "disabled" if reasoning_effort == "none" else "enabled"}
|
||||
|
||||
if (
|
||||
isinstance(reasoning_effort, str)
|
||||
and reasoning_effort in SUPPORTED_REASONING_EFFORTS
|
||||
and optional_params["thinking"] == {"type": "enabled"}
|
||||
):
|
||||
return {**optional_params, "reasoning_effort": reasoning_effort}
|
||||
|
||||
return optional_params
|
||||
|
||||
def _fill_reasoning_content(self, messages: list[AllMessageValues]) -> list[AllMessageValues]:
|
||||
|
|
|
|||
|
|
@ -56,9 +56,9 @@ class ElevenLabsTextToSpeechConfig(BaseTextToSpeechConfig):
|
|||
# Response format mappings from OpenAI to ElevenLabs
|
||||
FORMAT_MAPPINGS = {
|
||||
"mp3": "mp3_44100_128",
|
||||
"pcm": "pcm_44100",
|
||||
"pcm": "pcm_24000",
|
||||
"opus": "opus_48000_128",
|
||||
# ElevenLabs does not support WAV, AAC, or FLAC formats.
|
||||
"wav": "wav_24000",
|
||||
}
|
||||
|
||||
ELEVENLABS_QUERY_PARAMS_KEY = "__elevenlabs_query_params__"
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ from pydantic import BaseModel, TypeAdapter, ValidationError
|
|||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import strip_litellm_response_only_fields
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import strip_litellm_internal_message_fields
|
||||
from litellm.llms.openai.chat.gpt_transformation import (
|
||||
OpenAIChatCompletionStreamingHandler,
|
||||
)
|
||||
|
|
@ -154,7 +154,7 @@ class GroqChatConfig(OpenAILikeChatConfig):
|
|||
for k, v in _message.items():
|
||||
if v is not None:
|
||||
new_message[k] = v
|
||||
messages[idx] = strip_litellm_response_only_fields(new_message)
|
||||
messages[idx] = strip_litellm_internal_message_fields(new_message)
|
||||
|
||||
if is_async:
|
||||
return super()._transform_messages(messages=messages, model=model, is_async=True)
|
||||
|
|
|
|||
|
|
@ -90,7 +90,7 @@ class PassThroughEndpointHandler(BaseTranslation):
|
|||
payload_to_check: Final = {
|
||||
k: v
|
||||
for k, v in without_server_streaming_classification(data).items()
|
||||
if not k.startswith("_") and k not in ("metadata", "litellm_logging_obj")
|
||||
if not k.startswith("_") and k not in ("metadata", "litellm_metadata", "litellm_logging_obj")
|
||||
}
|
||||
verbose_proxy_logger.debug("PassThroughEndpointHandler: Using full payload for guardrail")
|
||||
return safe_dumps(payload_to_check)
|
||||
|
|
|
|||
|
|
@ -7,12 +7,14 @@ Handles Authentication and generating request urls for Vertex AI and Google AI S
|
|||
import asyncio
|
||||
import json
|
||||
import threading
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import DEFAULT_MAX_LRU_CACHE_SIZE
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
from litellm.litellm_core_utils.url_utils import join_api_base
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
|
|
@ -70,16 +72,20 @@ class _VertexCredentialsObject(Protocol):
|
|||
def refresh(self, request: object) -> None: ...
|
||||
|
||||
|
||||
_VertexCredentialCacheKey = tuple[str | None, str | None]
|
||||
|
||||
|
||||
class VertexBase:
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.access_token: str | None = None
|
||||
self.refresh_token: str | None = None
|
||||
self._credentials: GoogleCredentialsObject | None = None
|
||||
self._credentials_project_mapping: dict[
|
||||
tuple[VERTEX_CREDENTIALS_TYPES | None, str | None],
|
||||
self._credentials_project_mapping: OrderedDict[
|
||||
_VertexCredentialCacheKey,
|
||||
tuple[_VertexCredentialsObject, str | None],
|
||||
] = {}
|
||||
] = OrderedDict()
|
||||
self._credentials_cache_lock = threading.RLock()
|
||||
self.project_id: str | None = None
|
||||
self.async_handler: AsyncHTTPHandler | None = None
|
||||
# Per-credential-key asyncio.Lock for single-flight async refresh.
|
||||
|
|
@ -91,10 +97,10 @@ class VertexBase:
|
|||
# each lock; the entry is pruned when the count reaches zero, so the
|
||||
# dict stays bounded even in long-running high-cardinality deployments
|
||||
# without depending on any private asyncio internals.
|
||||
self._async_refresh_locks: dict[tuple, asyncio.Lock] = {}
|
||||
self._async_refresh_lock_refcounts: dict[tuple, int] = {}
|
||||
self._async_refresh_locks: dict[_VertexCredentialCacheKey, asyncio.Lock] = {}
|
||||
self._async_refresh_lock_refcounts: dict[_VertexCredentialCacheKey, int] = {}
|
||||
# Tracks in-flight background refresh tasks to avoid duplicate refreshes.
|
||||
self._background_refresh_tasks: dict[tuple, asyncio.Task] = {}
|
||||
self._background_refresh_tasks: dict[_VertexCredentialCacheKey, asyncio.Task] = {}
|
||||
# Protects the sync get_access_token refresh path.
|
||||
# Use RLock so that the reauthentication retry path (which calls
|
||||
# back into get_access_token while still holding the lock) can
|
||||
|
|
@ -378,7 +384,7 @@ class VertexBase:
|
|||
with self._sync_refresh_lock:
|
||||
credentials.refresh(Request())
|
||||
|
||||
def _acquire_async_refresh_lock(self, credential_cache_key: tuple) -> asyncio.Lock:
|
||||
def _acquire_async_refresh_lock(self, credential_cache_key: _VertexCredentialCacheKey) -> asyncio.Lock:
|
||||
"""Increment the refcount and return the lock for ``credential_cache_key``.
|
||||
|
||||
Every call must be paired with ``_release_async_refresh_lock`` once the
|
||||
|
|
@ -391,7 +397,7 @@ class VertexBase:
|
|||
)
|
||||
return lock
|
||||
|
||||
def _release_async_refresh_lock(self, credential_cache_key: tuple, lock: asyncio.Lock) -> None:
|
||||
def _release_async_refresh_lock(self, credential_cache_key: _VertexCredentialCacheKey, lock: asyncio.Lock) -> None:
|
||||
"""Decrement the refcount and drop the lock entry when it reaches zero.
|
||||
|
||||
Must be called only after the caller has released ``lock`` (i.e. once
|
||||
|
|
@ -409,7 +415,7 @@ class VertexBase:
|
|||
|
||||
def _try_get_cached_token(
|
||||
self,
|
||||
credential_cache_key: tuple,
|
||||
credential_cache_key: _VertexCredentialCacheKey,
|
||||
project_id: str | None,
|
||||
) -> tuple[str, str] | None:
|
||||
"""
|
||||
|
|
@ -432,7 +438,7 @@ class VertexBase:
|
|||
|
||||
def _try_get_usable_cached_token(
|
||||
self,
|
||||
credential_cache_key: tuple,
|
||||
credential_cache_key: _VertexCredentialCacheKey,
|
||||
project_id: str | None,
|
||||
) -> tuple[str, str, "TokenState", _VertexCredentialsObject, str | None] | None:
|
||||
"""
|
||||
|
|
@ -458,19 +464,37 @@ class VertexBase:
|
|||
return creds.token, resolved_project, token_state, creds, cached_project_id
|
||||
|
||||
def _unpack_cached_credentials(
|
||||
self, credential_cache_key: tuple
|
||||
self, credential_cache_key: _VertexCredentialCacheKey
|
||||
) -> tuple[_VertexCredentialsObject | None, str | None]:
|
||||
"""
|
||||
Return (credentials, project_id) from the cache, or (None, None) if
|
||||
not cached. Handles both tuple and legacy cache formats.
|
||||
"""
|
||||
if credential_cache_key not in self._credentials_project_mapping:
|
||||
return None, None
|
||||
cached_entry: Final = self._credentials_project_mapping[credential_cache_key]
|
||||
with self._credentials_cache_lock:
|
||||
cached_entry: Final = self._credentials_project_mapping.get(credential_cache_key)
|
||||
if cached_entry is None:
|
||||
return None, None
|
||||
self._credentials_project_mapping.move_to_end(credential_cache_key)
|
||||
if isinstance(cached_entry, tuple):
|
||||
return cached_entry
|
||||
return cached_entry, cached_entry.quota_project_id or getattr(cached_entry, "project_id", None)
|
||||
|
||||
def _cache_credentials(
|
||||
self,
|
||||
credential_cache_key: _VertexCredentialCacheKey,
|
||||
credentials: _VertexCredentialsObject,
|
||||
project_id: str | None,
|
||||
*,
|
||||
replace: bool = True,
|
||||
) -> None:
|
||||
with self._credentials_cache_lock:
|
||||
if not replace and credential_cache_key in self._credentials_project_mapping:
|
||||
return
|
||||
self._credentials_project_mapping[credential_cache_key] = (credentials, project_id)
|
||||
self._credentials_project_mapping.move_to_end(credential_cache_key)
|
||||
while len(self._credentials_project_mapping) > max(0, DEFAULT_MAX_LRU_CACHE_SIZE):
|
||||
self._credentials_project_mapping.popitem(last=False)
|
||||
|
||||
def _get_token_state(self, credentials: _VertexCredentialsObject) -> "TokenState":
|
||||
"""
|
||||
Return the token state using google-auth's TokenState enum.
|
||||
|
|
@ -494,7 +518,7 @@ class VertexBase:
|
|||
self,
|
||||
credentials: VERTEX_CREDENTIALS_TYPES | None,
|
||||
project_id: str | None,
|
||||
credential_cache_key: tuple,
|
||||
credential_cache_key: _VertexCredentialCacheKey,
|
||||
) -> tuple[_VertexCredentialsObject, str | None]:
|
||||
"""Load credentials via load_auth (in thread) and cache the result."""
|
||||
try:
|
||||
|
|
@ -507,17 +531,12 @@ class VertexBase:
|
|||
raise
|
||||
if _credentials is None:
|
||||
raise ValueError("Could not resolve credentials")
|
||||
self._credentials_project_mapping[credential_cache_key] = (
|
||||
_credentials,
|
||||
credential_project_id,
|
||||
)
|
||||
self._cache_credentials(credential_cache_key, _credentials, credential_project_id)
|
||||
return _credentials, credential_project_id
|
||||
|
||||
async def _background_refresh_credentials(
|
||||
self,
|
||||
credentials: _VertexCredentialsObject,
|
||||
credential_cache_key: tuple,
|
||||
credential_project_id: str | None,
|
||||
) -> None:
|
||||
"""
|
||||
Refresh credentials in the background without blocking the calling request.
|
||||
|
|
@ -528,25 +547,13 @@ class VertexBase:
|
|||
try:
|
||||
verbose_logger.debug("Background proactive credential refresh")
|
||||
await asyncify(self.refresh_auth)(credentials)
|
||||
# Only update the cache if it still points at the credentials
|
||||
# object we just refreshed. The per-key async lock is not held
|
||||
# here, so a concurrent INVALID path may have already replaced
|
||||
# this entry (e.g. via _handle_reauthentication_async, which
|
||||
# creates a fresh credentials object). In that case our write
|
||||
# would clobber the newer entry with a stale reference.
|
||||
cached_creds, _ = self._unpack_cached_credentials(credential_cache_key)
|
||||
if cached_creds is credentials:
|
||||
self._credentials_project_mapping[credential_cache_key] = (
|
||||
credentials,
|
||||
credential_project_id,
|
||||
)
|
||||
except Exception:
|
||||
verbose_logger.debug(
|
||||
"Background credential refresh failed, will retry on next request",
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
async def _await_in_flight_background_refresh(self, credential_cache_key: tuple) -> None:
|
||||
async def _await_in_flight_background_refresh(self, credential_cache_key: _VertexCredentialCacheKey) -> None:
|
||||
"""Wait for an in-flight background refresh to finish, if any.
|
||||
|
||||
google-auth's ``Credentials.refresh()`` is not safe to invoke
|
||||
|
|
@ -568,8 +575,7 @@ class VertexBase:
|
|||
def _schedule_background_refresh(
|
||||
self,
|
||||
credentials: _VertexCredentialsObject,
|
||||
credential_cache_key: tuple,
|
||||
credential_project_id: str | None,
|
||||
credential_cache_key: _VertexCredentialCacheKey,
|
||||
) -> None:
|
||||
"""Kick off a single background refresh for ``credential_cache_key``.
|
||||
|
||||
|
|
@ -581,9 +587,7 @@ class VertexBase:
|
|||
if existing is not None and not existing.done():
|
||||
return
|
||||
self._background_refresh_tasks.pop(credential_cache_key, None)
|
||||
task: Final = asyncio.create_task(
|
||||
self._background_refresh_credentials(credentials, credential_cache_key, credential_project_id)
|
||||
)
|
||||
task: Final = asyncio.create_task(self._background_refresh_credentials(credentials))
|
||||
|
||||
def _drop_background_refresh_task(_fut: asyncio.Future[None]) -> None:
|
||||
if self._background_refresh_tasks.get(credential_cache_key) is _fut:
|
||||
|
|
@ -796,7 +800,7 @@ class VertexBase:
|
|||
self,
|
||||
credentials: VERTEX_CREDENTIALS_TYPES | None,
|
||||
project_id: str | None,
|
||||
credential_cache_key: tuple,
|
||||
credential_cache_key: _VertexCredentialCacheKey,
|
||||
error: Exception,
|
||||
) -> tuple[str, str]:
|
||||
"""
|
||||
|
|
@ -822,8 +826,8 @@ class VertexBase:
|
|||
)
|
||||
|
||||
# Clear the cached credentials
|
||||
if credential_cache_key in self._credentials_project_mapping:
|
||||
del self._credentials_project_mapping[credential_cache_key]
|
||||
with self._credentials_cache_lock:
|
||||
self._credentials_project_mapping.pop(credential_cache_key, None)
|
||||
|
||||
# Retry once with _retry_reauth=True to prevent infinite recursion
|
||||
try:
|
||||
|
|
@ -846,7 +850,7 @@ class VertexBase:
|
|||
self,
|
||||
credentials: VERTEX_CREDENTIALS_TYPES | None,
|
||||
project_id: str | None,
|
||||
credential_cache_key: tuple,
|
||||
credential_cache_key: _VertexCredentialCacheKey,
|
||||
error: Exception,
|
||||
) -> tuple[str, str]:
|
||||
"""
|
||||
|
|
@ -856,7 +860,8 @@ class VertexBase:
|
|||
"Handling async reauthentication for project_id: %s. Clearing cache and retrying once.", project_id
|
||||
)
|
||||
|
||||
self._credentials_project_mapping.pop(credential_cache_key, None)
|
||||
with self._credentials_cache_lock:
|
||||
self._credentials_project_mapping.pop(credential_cache_key, None)
|
||||
|
||||
try:
|
||||
(
|
||||
|
|
@ -875,10 +880,7 @@ class VertexBase:
|
|||
# references the OLD credentials object we just replaced, and
|
||||
# leaving it would force the next request to do a redundant
|
||||
# refresh/reauth before realizing the cached creds are stale.
|
||||
self._credentials_project_mapping[resolved_cache_key] = (
|
||||
_credentials,
|
||||
credential_project_id,
|
||||
)
|
||||
self._cache_credentials(resolved_cache_key, _credentials, credential_project_id)
|
||||
|
||||
if _credentials.token is None or not isinstance(_credentials.token, str):
|
||||
raise ValueError(
|
||||
|
|
@ -924,21 +926,11 @@ class VertexBase:
|
|||
# Convert dict credentials to string for caching
|
||||
cache_credentials: Final = json.dumps(credentials) if isinstance(credentials, dict) else credentials
|
||||
credential_cache_key: Final = (cache_credentials, project_id)
|
||||
_credentials: _VertexCredentialsObject | None = None
|
||||
|
||||
verbose_logger.debug("Checking cached credentials for project_id: %s", project_id)
|
||||
|
||||
if credential_cache_key in self._credentials_project_mapping:
|
||||
_credentials, credential_project_id = self._unpack_cached_credentials(credential_cache_key)
|
||||
if _credentials is not None:
|
||||
verbose_logger.debug("Cached credentials found for project_id: %s.", project_id)
|
||||
# Retrieve both credentials and cached project_id
|
||||
cached_entry: Final = self._credentials_project_mapping[credential_cache_key]
|
||||
verbose_logger.debug("cached_entry: %s", cached_entry)
|
||||
if isinstance(cached_entry, tuple):
|
||||
_credentials, credential_project_id = cached_entry
|
||||
else:
|
||||
# Backward compatibility with old cache format
|
||||
_credentials = cached_entry
|
||||
credential_project_id = _credentials.quota_project_id or getattr(_credentials, "project_id", None)
|
||||
verbose_logger.debug(
|
||||
"Using cached credentials for project_id: %s",
|
||||
credential_project_id,
|
||||
|
|
@ -963,10 +955,7 @@ class VertexBase:
|
|||
f"Could not resolve credentials - either dynamically or from environment, for project_id: {project_id}"
|
||||
)
|
||||
# Cache the project_id and credentials from load_auth result (resolved project_id)
|
||||
self._credentials_project_mapping[credential_cache_key] = (
|
||||
_credentials,
|
||||
credential_project_id,
|
||||
)
|
||||
self._cache_credentials(credential_cache_key, _credentials, credential_project_id)
|
||||
|
||||
## VALIDATE CREDENTIALS
|
||||
verbose_logger.debug("Validating credentials")
|
||||
|
|
@ -974,11 +963,7 @@ class VertexBase:
|
|||
project_id = credential_project_id
|
||||
# Update cache with resolved project_id for future lookups
|
||||
resolved_cache_key: Final = (cache_credentials, project_id)
|
||||
if resolved_cache_key not in self._credentials_project_mapping:
|
||||
self._credentials_project_mapping[resolved_cache_key] = (
|
||||
_credentials,
|
||||
credential_project_id,
|
||||
)
|
||||
self._cache_credentials(resolved_cache_key, _credentials, credential_project_id, replace=False)
|
||||
|
||||
# Check if credentials are None before accessing attributes
|
||||
if _credentials is None:
|
||||
|
|
@ -991,10 +976,6 @@ class VertexBase:
|
|||
try:
|
||||
verbose_logger.debug("Credentials expired, refreshing")
|
||||
self.refresh_auth(_credentials)
|
||||
self._credentials_project_mapping[credential_cache_key] = (
|
||||
_credentials,
|
||||
credential_project_id,
|
||||
)
|
||||
except Exception as e:
|
||||
# if refresh fails, it's possible the user has re-authenticated via `gcloud auth application-default login`
|
||||
# in this case, we should try to reload the credentials by clearing the cache and retrying
|
||||
|
|
@ -1043,9 +1024,9 @@ class VertexBase:
|
|||
# callers on the lock just to schedule that refresh.
|
||||
usable: Final = self._try_get_usable_cached_token(credential_cache_key, project_id)
|
||||
if usable is not None:
|
||||
cached_token, resolved_project, token_state, creds, cached_project_id = usable
|
||||
cached_token, resolved_project, token_state, creds, _cached_project_id = usable
|
||||
if token_state == TokenState.STALE:
|
||||
self._schedule_background_refresh(creds, credential_cache_key, cached_project_id)
|
||||
self._schedule_background_refresh(creds, credential_cache_key)
|
||||
return cached_token, resolved_project
|
||||
|
||||
# === SLOW PATH (per-key lock) ===
|
||||
|
|
@ -1075,10 +1056,7 @@ class VertexBase:
|
|||
# reauth that only repopulated the unresolved key), which
|
||||
# would force the next request through an unnecessary
|
||||
# refresh/reauth cycle.
|
||||
self._credentials_project_mapping[resolved_cache_key] = (
|
||||
_credentials,
|
||||
credential_project_id,
|
||||
)
|
||||
self._cache_credentials(resolved_cache_key, _credentials, credential_project_id)
|
||||
|
||||
# Use google-auth's token_state to decide refresh strategy:
|
||||
# - STALE: token is usable but within REFRESH_THRESHOLD (3:45) of
|
||||
|
|
@ -1098,7 +1076,6 @@ class VertexBase:
|
|||
self._schedule_background_refresh(
|
||||
_credentials,
|
||||
credential_cache_key,
|
||||
credential_project_id,
|
||||
)
|
||||
return current_token, project_id
|
||||
|
||||
|
|
@ -1117,10 +1094,6 @@ class VertexBase:
|
|||
try:
|
||||
verbose_logger.debug("Credentials expired, refreshing")
|
||||
await asyncify(self.refresh_auth)(_credentials)
|
||||
self._credentials_project_mapping[credential_cache_key] = (
|
||||
_credentials,
|
||||
credential_project_id,
|
||||
)
|
||||
except Exception as e:
|
||||
if "Reauthentication is needed" in str(e):
|
||||
verbose_logger.debug("Reauthentication needed, clearing cache and retrying")
|
||||
|
|
|
|||
|
|
@ -33698,7 +33698,7 @@
|
|||
"input_cost_per_token": 1e-05,
|
||||
"input_cost_per_token_above_272k_tokens": 2e-05,
|
||||
"litellm_provider": "github_copilot",
|
||||
"max_input_tokens": 272000,
|
||||
"max_input_tokens": 922000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "responses",
|
||||
|
|
@ -33722,7 +33722,7 @@
|
|||
"input_cost_per_token": 1e-07,
|
||||
"input_cost_per_token_above_272k_tokens": 2e-07,
|
||||
"litellm_provider": "github_copilot",
|
||||
"max_input_tokens": 272000,
|
||||
"max_input_tokens": 922000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "responses",
|
||||
|
|
@ -33746,7 +33746,7 @@
|
|||
"input_cost_per_token": 2e-06,
|
||||
"input_cost_per_token_above_272k_tokens": 4e-06,
|
||||
"litellm_provider": "github_copilot",
|
||||
"max_input_tokens": 272000,
|
||||
"max_input_tokens": 922000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "responses",
|
||||
|
|
@ -33770,7 +33770,7 @@
|
|||
"input_cost_per_token": 2e-06,
|
||||
"input_cost_per_token_above_272k_tokens": 4e-06,
|
||||
"litellm_provider": "github_copilot",
|
||||
"max_input_tokens": 272000,
|
||||
"max_input_tokens": 922000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "responses",
|
||||
|
|
@ -36840,10 +36840,14 @@
|
|||
"cache_creation_input_token_cost": 1.25e-05,
|
||||
"cache_creation_input_token_cost_above_12hr": 2.5e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 0.0001,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": 0.0003,
|
||||
"cache_creation_input_token_cost_above_12hr_batches": 1.25e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_flex": 1.25e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_priority": 5e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_ultrafast": 0.00015,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 2.5e-05,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05,
|
||||
|
|
@ -36931,8 +36935,10 @@
|
|||
"cache_creation_input_token_cost": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr": 5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_batches": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_flex": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_priority": 1e-05,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 5e-06,
|
||||
|
|
@ -37014,8 +37020,10 @@
|
|||
"cache_creation_input_token_cost": 1.25e-07,
|
||||
"cache_creation_input_token_cost_above_12hr": 2.5e-07,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-07,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-07,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-07,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 1e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_batches": 1.25e-07,
|
||||
"cache_creation_input_token_cost_above_12hr_flex": 1.25e-07,
|
||||
"cache_creation_input_token_cost_above_12hr_priority": 5e-07,
|
||||
"cache_creation_input_token_cost_decisions": 1.25e-07,
|
||||
|
|
@ -37186,8 +37194,10 @@
|
|||
"cache_creation_input_token_cost": 5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr": 1e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens": 2e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 1e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 1e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 4e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_batches": 5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_flex": 5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_priority": 2e-05,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 1e-05,
|
||||
|
|
@ -37268,8 +37278,10 @@
|
|||
"cache_creation_input_token_cost": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr": 5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_batches": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_flex": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_priority": 1e-05,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 5e-06,
|
||||
|
|
@ -37349,8 +37361,10 @@
|
|||
"cache_creation_input_token_cost": 2.5e-07,
|
||||
"cache_creation_input_token_cost_above_12hr": 5e-07,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-07,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-07,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_batches": 2.5e-07,
|
||||
"cache_creation_input_token_cost_above_12hr_flex": 2.5e-07,
|
||||
"cache_creation_input_token_cost_above_12hr_priority": 1e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 5e-07,
|
||||
|
|
@ -42726,6 +42740,29 @@
|
|||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"mistral/zai-glm-5-3-flash": {
|
||||
"cache_read_input_token_cost": 3e-08,
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"litellm_provider": "mistral",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 5e-07,
|
||||
"reasoning_effort_levels": [
|
||||
"low",
|
||||
"high",
|
||||
"max"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/zai-glm-5-3-flash",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"mistral/zai-glm-5": {
|
||||
"cache_read_input_token_cost": 1.4e-07,
|
||||
"input_cost_per_token": 1.4e-06,
|
||||
|
|
@ -75023,10 +75060,14 @@
|
|||
"cache_creation_input_token_cost": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr": 5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": 6e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_batches": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_flex": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_priority": 1e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_ultrafast": 3e-05,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 5e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_batches": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06,
|
||||
|
|
@ -76416,9 +76457,10 @@
|
|||
"azure_ai/Microsoft-Decision-1": {
|
||||
"input_cost_per_token": 4.2e-08,
|
||||
"litellm_provider": "azure_ai",
|
||||
"max_input_tokens": 32768,
|
||||
"mode": "evaluation",
|
||||
"output_cost_per_token": 0.0,
|
||||
"source": "https://prices.azure.com/api/retail/prices?$filter=productName eq 'Microsoft Decision Models'"
|
||||
"source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure"
|
||||
},
|
||||
"gpt-rosalind-discovery": {
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
|
|
@ -78629,5 +78671,227 @@
|
|||
"prompt_cache_min_tokens": 512,
|
||||
"source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing",
|
||||
"supports_web_search": true
|
||||
},
|
||||
"gpt-6.1-terra": {
|
||||
"cache_creation_input_token_cost": 1.25e-06,
|
||||
"cache_creation_input_token_cost_above_12hr": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 1e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_batches": 1.25e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_flex": 1.25e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_priority": 5e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_batches": 1.25e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 5e-06,
|
||||
"cache_creation_input_token_cost_batches": 6.25e-07,
|
||||
"cache_creation_input_token_cost_flex": 6.25e-07,
|
||||
"cache_creation_input_token_cost_priority": 2.5e-06,
|
||||
"cache_read_input_token_cost": 5e-08,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 1e-07,
|
||||
"cache_read_input_token_cost_above_272k_tokens_batches": 5e-08,
|
||||
"cache_read_input_token_cost_above_272k_tokens_flex": 5e-08,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 2e-07,
|
||||
"cache_read_input_token_cost_batches": 2.5e-08,
|
||||
"cache_read_input_token_cost_flex": 2.5e-08,
|
||||
"cache_read_input_token_cost_priority": 1e-07,
|
||||
"input_cost_per_token": 1e-06,
|
||||
"input_cost_per_token_above_272k_tokens": 2e-06,
|
||||
"input_cost_per_token_above_272k_tokens_batches": 1e-06,
|
||||
"input_cost_per_token_above_272k_tokens_flex": 1e-06,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 4e-06,
|
||||
"input_cost_per_token_batches": 5e-07,
|
||||
"input_cost_per_token_flex": 5e-07,
|
||||
"input_cost_per_token_priority": 2e-06,
|
||||
"litellm_provider": "openai",
|
||||
"max_input_tokens": 922000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 5e-06,
|
||||
"output_cost_per_token_above_272k_tokens": 7.5e-06,
|
||||
"output_cost_per_token_above_272k_tokens_batches": 3.75e-06,
|
||||
"output_cost_per_token_above_272k_tokens_flex": 3.75e-06,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 1.5e-05,
|
||||
"output_cost_per_token_batches": 2.5e-06,
|
||||
"output_cost_per_token_flex": 2.5e-06,
|
||||
"output_cost_per_token_priority": 1e-05,
|
||||
"regional_processing_uplift_multiplier_eu": 1.1,
|
||||
"regional_processing_uplift_multiplier_us": 1.1,
|
||||
"search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01},
|
||||
"source": "https://developers.openai.com/api/docs/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions",
|
||||
"/v1/batch",
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_computer_use": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_native_streaming": true,
|
||||
"supports_none_reasoning_effort": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_cache_breakpoint": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_tool_search": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true,
|
||||
"supports_xhigh_reasoning_effort": true
|
||||
},
|
||||
"chatgpt/gpt-6.1-terra": {
|
||||
"litellm_provider": "chatgpt",
|
||||
"max_input_tokens": 922000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "responses",
|
||||
"source": "https://developers.openai.com/api/docs/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions",
|
||||
"/v1/responses"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_none_reasoning_effort": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_vision": true,
|
||||
"supports_xhigh_reasoning_effort": true
|
||||
},
|
||||
"global.openai.gpt-6.1-terra": {
|
||||
"cache_creation_input_token_cost": 1.25e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 2.5e-06,
|
||||
"cache_read_input_token_cost": 5e-08,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 1e-07,
|
||||
"input_cost_per_token": 1e-06,
|
||||
"input_cost_per_token_above_272k_tokens": 2e-06,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 922000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 5e-06,
|
||||
"output_cost_per_token_above_272k_tokens": 7.5e-06,
|
||||
"source": "https://developers.openai.com/api/docs/models/gpt-6.1-terra",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions",
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_bedrock_converse_tool_result_images": false,
|
||||
"supports_bedrock_runtime_chat_completions_response_format": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_none_reasoning_effort": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_xhigh_reasoning_effort": true
|
||||
},
|
||||
"us.openai.gpt-6.1-terra": {
|
||||
"cache_creation_input_token_cost": 1.375e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 2.75e-06,
|
||||
"cache_read_input_token_cost": 5.5e-08,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 1.1e-07,
|
||||
"input_cost_per_token": 1.1e-06,
|
||||
"input_cost_per_token_above_272k_tokens": 2.2e-06,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 922000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 5.5e-06,
|
||||
"output_cost_per_token_above_272k_tokens": 8.25e-06,
|
||||
"source": "https://developers.openai.com/api/docs/models/gpt-6.1-terra",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions",
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_bedrock_converse_tool_result_images": false,
|
||||
"supports_bedrock_runtime_chat_completions_response_format": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_none_reasoning_effort": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_xhigh_reasoning_effort": true
|
||||
},
|
||||
"baseten/decision-models/baseten-elephant-9b-latest": {
|
||||
"input_cost_per_token": 0.0,
|
||||
"litellm_provider": "baseten",
|
||||
"max_input_tokens": 16384,
|
||||
"mode": "decisions",
|
||||
"output_cost_per_token": 0.0,
|
||||
"source": "https://inference.baseten.co/v1/models",
|
||||
"supports_vision": true
|
||||
},
|
||||
"baseten/decision-models/baseten-mammoth-27b-latest": {
|
||||
"input_cost_per_token": 0.0,
|
||||
"litellm_provider": "baseten",
|
||||
"max_input_tokens": 16384,
|
||||
"mode": "decisions",
|
||||
"output_cost_per_token": 0.0,
|
||||
"source": "https://inference.baseten.co/v1/models",
|
||||
"supports_vision": true
|
||||
},
|
||||
"github_copilot/gpt-6.1-terra": {
|
||||
"cache_creation_input_token_cost": 1.25e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 2.5e-06,
|
||||
"cache_read_input_token_cost": 5e-08,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 1e-07,
|
||||
"input_cost_per_token": 1e-06,
|
||||
"input_cost_per_token_above_272k_tokens": 2e-06,
|
||||
"litellm_provider": "github_copilot",
|
||||
"max_input_tokens": 922000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "responses",
|
||||
"output_cost_per_token": 5e-06,
|
||||
"output_cost_per_token_above_272k_tokens": 7.5e-06,
|
||||
"source": "https://github.blog/changelog/2026-10-26-gpt-6-1-terra-in-github-copilot",
|
||||
"supported_endpoints": [
|
||||
"/v1/responses"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_vision": true
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -3,8 +3,8 @@
|
|||
The aggregate listing deliberately keeps serving the healthy subset when one server fails, but a
|
||||
failed server must contribute a classified outcome instead of silently shrinking the list: an empty
|
||||
contribution with no signal makes a broken upstream indistinguishable from a healthy server with no
|
||||
tools. Outcomes carry only machine fields (category and status code) so nothing from an upstream
|
||||
body crosses the trust boundary; classification is total, so any exception out of a server fetch
|
||||
tools. Outcomes carry machine fields and controlled diagnostics so nothing from an upstream body
|
||||
crosses the trust boundary; classification is total, so any exception out of a server fetch
|
||||
becomes an outcome, never a second failure.
|
||||
"""
|
||||
|
||||
|
|
@ -21,8 +21,10 @@ from urllib.request import parse_http_list, parse_keqv_list
|
|||
import anyio
|
||||
import httpx
|
||||
import httpx2
|
||||
from mcp import MCPError
|
||||
from mcp.types import CONNECTION_CLOSED
|
||||
from mcp.types import Tool as MCPTool
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
from pydantic import ConfigDict, PositiveInt, TypeAdapter
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -65,6 +67,16 @@ class ServerListFault(LiteLLMBaseModel):
|
|||
model_config = ConfigDict(frozen=True)
|
||||
tag: ListFaultCategory
|
||||
status_code: int | None = None
|
||||
sse_event_limit_bytes: PositiveInt | None = None
|
||||
|
||||
@property
|
||||
def public_message(self) -> str | None:
|
||||
if self.sse_event_limit_bytes is None:
|
||||
return None
|
||||
return (
|
||||
f"Upstream MCP SSE event exceeded the {self.sse_event_limit_bytes} byte per-event limit. "
|
||||
"Paginate tools/list responses so each event stays below the limit."
|
||||
)
|
||||
|
||||
|
||||
ServerOutcome: TypeAlias = ServerListOk | ServerListFault
|
||||
|
|
@ -155,6 +167,18 @@ def raise_classified_list_failure(
|
|||
raise MCPServerListError(classify_list_exception(exc), server_name) from exc
|
||||
|
||||
|
||||
def _sse_event_limit(exc: BaseException) -> int | None:
|
||||
"""The SDK replaces decoder exceptions with CONNECTION_CLOSED and a fixed diagnostic."""
|
||||
if not isinstance(exc, (httpx2.SSEError, MCPError)):
|
||||
return None
|
||||
if isinstance(exc, MCPError) and exc.error.code != CONNECTION_CLOSED:
|
||||
return None
|
||||
prefix: Final = "SSE stream failed: " if isinstance(exc, MCPError) else ""
|
||||
message: Final = exc.error.message if isinstance(exc, MCPError) else str(exc)
|
||||
matched: Final = re.fullmatch(prefix + r"Server-sent event exceeded the ([1-9][0-9]{0,18}) byte limit\.", message)
|
||||
return int(matched[1]) if matched is not None else None
|
||||
|
||||
|
||||
def classify_list_exception(exc: BaseException) -> ServerListFault:
|
||||
"""Classify a per-server listing failure into exactly one outcome. Total: an exception this
|
||||
function cannot recognize is the gateway's own fault (``internal``), never a re-raise."""
|
||||
|
|
@ -179,6 +203,11 @@ def classify_list_exception(exc: BaseException) -> ServerListFault:
|
|||
response: Final = _find_upstream_response(exc)
|
||||
if response is not None:
|
||||
return ServerListFault(tag="upstream_error", status_code=response.status_code)
|
||||
sse_event_limit: Final = next(
|
||||
(limit for error in iter_exception_tree(exc) if (limit := _sse_event_limit(error)) is not None), None
|
||||
)
|
||||
if sse_event_limit is not None:
|
||||
return ServerListFault(tag="upstream_error", sse_event_limit_bytes=sse_event_limit)
|
||||
if isinstance(exc, (httpx.TimeoutException, httpx2.TimeoutException)):
|
||||
return ServerListFault(tag="timeout")
|
||||
if isinstance(exc, (httpx.TransportError, httpx2.TransportError)):
|
||||
|
|
@ -295,7 +324,7 @@ def listing_failed_outright(outcomes: Mapping[str, ServerOutcome]) -> bool:
|
|||
|
||||
def outcome_wire_value(outcome: ServerOutcome) -> dict[str, object]:
|
||||
"""The client-visible form of one outcome, for the tools/list result ``_meta`` and the REST
|
||||
response: category plus status code only, never upstream prose or URLs."""
|
||||
response: category, status code and controlled diagnostics, never upstream prose or URLs."""
|
||||
match outcome.tag:
|
||||
case "ok":
|
||||
return {"status": "ok", "tool_count": outcome.tool_count}
|
||||
|
|
@ -303,6 +332,15 @@ def outcome_wire_value(outcome: ServerOutcome) -> dict[str, object]:
|
|||
return {
|
||||
"status": outcome.tag,
|
||||
**({"http_status": outcome.status_code} if outcome.status_code is not None else {}),
|
||||
**(
|
||||
{
|
||||
"reason": "sse_event_too_large",
|
||||
"limit_bytes": outcome.sse_event_limit_bytes,
|
||||
"message": outcome.public_message,
|
||||
}
|
||||
if outcome.sse_event_limit_bytes is not None
|
||||
else {}
|
||||
),
|
||||
}
|
||||
case _:
|
||||
assert_never(outcome.tag)
|
||||
|
|
|
|||
|
|
@ -957,7 +957,7 @@ if MCP_AVAILABLE:
|
|||
status_code=list_fault_http_status(fault),
|
||||
detail={
|
||||
"error": fault.tag,
|
||||
"message": f"Failed to list tools from server {get_server_prefix(server)}",
|
||||
"message": fault.public_message or f"Failed to list tools from server {get_server_prefix(server)}",
|
||||
},
|
||||
) from e
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -1766,8 +1766,10 @@ class BudgetNewRequest(LiteLLMPydanticObjectBase):
|
|||
max_parallel_requests: int | None = Field(
|
||||
default=None, description="Max concurrent requests allowed for this budget id."
|
||||
)
|
||||
tpm_limit: int | None = Field(default=None, description="Max tokens per minute, allowed for this budget id.")
|
||||
rpm_limit: int | None = Field(default=None, description="Max requests per minute, allowed for this budget id.")
|
||||
tpm_limit: int | None = Field(default=None, ge=0, description="Max tokens per minute, allowed for this budget id.")
|
||||
rpm_limit: int | None = Field(
|
||||
default=None, ge=0, description="Max requests per minute, allowed for this budget id."
|
||||
)
|
||||
tpd_limit: int | None = Field(
|
||||
default=None, description="Max tokens per day, charged by batch submissions, allowed for this budget id."
|
||||
)
|
||||
|
|
@ -1878,6 +1880,8 @@ RouterSettingsDict = Annotated[
|
|||
|
||||
|
||||
class NewTeamRequest(TeamBase):
|
||||
tpm_limit: int | None = Field(default=None, ge=0)
|
||||
rpm_limit: int | None = Field(default=None, ge=0)
|
||||
router_settings: RouterSettingsDict | None = None
|
||||
model_aliases: dict | None = None
|
||||
model_max_budget: GenericBudgetConfigType | None = Field(
|
||||
|
|
@ -1957,8 +1961,8 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase):
|
|||
team_alias: str | None = None
|
||||
organization_id: str | None = None
|
||||
metadata: dict | None = None
|
||||
tpm_limit: int | None = None
|
||||
rpm_limit: int | None = None
|
||||
tpm_limit: int | None = Field(default=None, ge=0)
|
||||
rpm_limit: int | None = Field(default=None, ge=0)
|
||||
tpd_limit: int | None = None
|
||||
max_budget: float | None = None
|
||||
soft_budget: float | None = None
|
||||
|
|
|
|||
|
|
@ -119,6 +119,7 @@ from litellm.proxy.common_utils.sse_keepalive import (
|
|||
)
|
||||
from litellm.proxy.dd_span_tagger import DDSpanTagger
|
||||
from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _arm_auto_router_compression
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import get_or_create_request_stash
|
||||
from litellm.proxy.native_compaction import with_proxy_compaction_executor
|
||||
from litellm.proxy.route_llm_request import (
|
||||
route_request,
|
||||
|
|
@ -1416,6 +1417,8 @@ async def open_sse_before_first_byte(
|
|||
if interval is None:
|
||||
return await produce_response
|
||||
|
||||
# Share admission state with stream cleanup before the producer copies this context.
|
||||
get_or_create_request_stash()
|
||||
produce_task: Final = asyncio.ensure_future(produce_response)
|
||||
await asyncio.wait((produce_task,), timeout=interval)
|
||||
if produce_task.done():
|
||||
|
|
|
|||
|
|
@ -31,6 +31,7 @@ _BINARY_CONTENT_TYPES: Final[frozenset[str]] = frozenset({"application/x-protobu
|
|||
|
||||
_ANNOTATION_QUALIFIERS: Final[frozenset[object]] = frozenset({Annotated, NotRequired, ReadOnly, Required})
|
||||
_DECODED_JSON: Final = TypeAdapter(object)
|
||||
_JSON_REQUEST_BODY: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def resolve_inference_model(
|
||||
|
|
@ -272,7 +273,7 @@ async def read_request_body(request: Request | None) -> dict:
|
|||
parsed_body = {}
|
||||
else:
|
||||
try:
|
||||
parsed_body = orjson.loads(body)
|
||||
parsed_body = dict(_JSON_REQUEST_BODY.validate_python(orjson.loads(body)))
|
||||
except orjson.JSONDecodeError as e:
|
||||
# The surrogate-repair fallback below runs two full-body re.sub
|
||||
# passes, which block the event loop on multi-MB malformed bodies.
|
||||
|
|
@ -415,7 +416,9 @@ def rewrite_request_model(
|
|||
if request is None:
|
||||
return
|
||||
cached_body: Final = _safe_get_request_parsed_body(request=request)
|
||||
body: Final = {**cached_body, "model": model} if cached_body is not None else request_data
|
||||
body: Final = (
|
||||
{**_JSON_REQUEST_BODY.validate_python(cached_body), "model": model} if cached_body is not None else request_data
|
||||
)
|
||||
safe_set_request_parsed_body(request=request, parsed_body=body)
|
||||
request._json = body # pyright: ignore[reportPrivateUsage] # Starlette JSON cache
|
||||
request._body = orjson.dumps(body) # pyright: ignore[reportPrivateUsage] # Starlette body cache
|
||||
|
|
|
|||
|
|
@ -65,6 +65,7 @@ from litellm.repositories.table_repositories import (
|
|||
from litellm.repositories.team_repository import TeamRepository
|
||||
from litellm.repositories.unit_of_work import (
|
||||
LinkedSpendResetWrites,
|
||||
SpendResetWrites,
|
||||
budget_cascade_unit_of_work,
|
||||
enduser_spend_reset_unit_of_work,
|
||||
spend_reset_unit_of_work,
|
||||
|
|
@ -235,7 +236,7 @@ def _enduser_invalidation_where(budget_ids: Sequence[str]) -> dict[str, object]:
|
|||
|
||||
|
||||
def _queue_budget_linked_resets(
|
||||
writes: LinkedSpendResetWrites,
|
||||
writes: SpendResetWrites,
|
||||
cascade: "_BudgetCascade",
|
||||
extra: Mapping[str, object] = MappingProxyType({}),
|
||||
) -> None:
|
||||
|
|
@ -908,6 +909,7 @@ class ResetBudgetJob:
|
|||
db_span("reset_budget_cascade", "LiteLLM_BudgetTable"),
|
||||
budget_cascade_unit_of_work(self._new_batch) as uow,
|
||||
):
|
||||
uow.team_memberships.queue_total_spend_repair(cascade.budget_ids)
|
||||
_queue_budget_linked_resets(uow.team_memberships, cascade)
|
||||
_queue_budget_linked_resets(uow.keys, cascade, extra=_LINKED_KEYS_WHERE)
|
||||
_queue_budget_linked_resets(uow.organizations, cascade, extra=_SPENT_ROWS_WHERE)
|
||||
|
|
|
|||
|
|
@ -472,6 +472,8 @@ _TEAM_ADVISORY_LOCK_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext($1)) IS
|
|||
# recreates the member. The rows travel as one JSON document, not as a numeric array: Prisma
|
||||
# types a raw array parameter from the first batch a connection sees, so after an all-$0 batch
|
||||
# (integers) every later fractional batch on that connection failed with "improper binary format".
|
||||
# total_spend is floored at spend before adding: rows written before total_spend existed (migration
|
||||
# 20260421135425, default 0) carry their history only in spend, so a plain increment keeps them behind forever.
|
||||
_TEAM_MEMBER_SPEND_SQL: Final = """
|
||||
INSERT INTO "LiteLLM_TeamMembership" (user_id, team_id, spend, total_spend)
|
||||
SELECT member.user_id, member.team_id, member.cost, member.cost
|
||||
|
|
@ -487,7 +489,7 @@ WHERE EXISTS (
|
|||
)
|
||||
ON CONFLICT (user_id, team_id) DO UPDATE
|
||||
SET spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend,
|
||||
total_spend = "LiteLLM_TeamMembership".total_spend + EXCLUDED.total_spend
|
||||
total_spend = GREATEST("LiteLLM_TeamMembership".total_spend, "LiteLLM_TeamMembership".spend) + EXCLUDED.total_spend
|
||||
"""
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1013,7 +1013,7 @@ async def get_all_team_memberships(
|
|||
|
||||
returned_tm: Final[list[LiteLLM_TeamMembership]] = []
|
||||
for tm in team_memberships:
|
||||
returned_tm.append(LiteLLM_TeamMembership.model_validate(tm.model_dump()))
|
||||
returned_tm.append(with_total_spend_floor(LiteLLM_TeamMembership.model_validate(tm.model_dump())))
|
||||
|
||||
return returned_tm
|
||||
|
||||
|
|
@ -4202,6 +4202,23 @@ def _check_not_resetting_own_spend(user_id: str, user_api_key_dict: UserAPIKeyAu
|
|||
_raise_reset_spend_error(status.HTTP_403_FORBIDDEN, "Cannot reset your own spend. Ask a proxy admin.")
|
||||
|
||||
|
||||
# One statement so the validated spend cannot go stale between the read and the write: the WHERE re-checks
|
||||
# reset_to <= spend, and total_spend is floored at spend before spend drops so the history survives the reset.
|
||||
_RESET_TEAM_MEMBER_SPEND_SQL: Final = """
|
||||
UPDATE "LiteLLM_TeamMembership"
|
||||
SET total_spend = GREATEST(total_spend, spend), spend = $3::float8
|
||||
WHERE user_id = $1 AND team_id = $2 AND spend >= $3::float8
|
||||
"""
|
||||
|
||||
|
||||
def with_total_spend_floor(membership: LiteLLM_TeamMembership) -> LiteLLM_TeamMembership:
|
||||
"""Rows written before total_spend existed (migration 20260421135425) hold their history only in spend."""
|
||||
spend: Final = membership.spend or 0.0
|
||||
if (membership.total_spend or 0.0) >= spend:
|
||||
return membership
|
||||
return membership.model_copy(update={"total_spend": spend})
|
||||
|
||||
|
||||
def _raise_reset_spend_error(status_code: int, message: str) -> NoReturn:
|
||||
detail: Final = {"error": message}
|
||||
raise HTTPException(status_code=status_code, detail=detail)
|
||||
|
|
@ -4286,10 +4303,12 @@ async def reset_team_member_spend_fn(
|
|||
current_spend: Final = membership.spend or 0.0
|
||||
reset_to: Final = _validate_team_member_reset_spend_value(data.reset_to, membership)
|
||||
|
||||
await _team_membership_db(prisma_client).update(
|
||||
where=membership_where,
|
||||
data={"spend": reset_to},
|
||||
)
|
||||
reset_rows: Final = await _execute_raw(prisma_client)(_RESET_TEAM_MEMBER_SPEND_SQL, user_id, team_id, reset_to)
|
||||
if reset_rows == 0:
|
||||
_raise_reset_spend_error(
|
||||
status.HTTP_409_CONFLICT,
|
||||
f"Member spend changed since it was read; reset_to ({reset_to}) must be <= current spend. Retry.",
|
||||
)
|
||||
|
||||
await invalidate_team_member_spend_state(
|
||||
user_id=user_id,
|
||||
|
|
@ -5548,16 +5567,17 @@ async def team_member_me(
|
|||
litellm_budget_table=None,
|
||||
)
|
||||
|
||||
floored: Final = with_total_spend_floor(membership)
|
||||
return TeamMemberInfoResponse(
|
||||
user_id=caller_user_id,
|
||||
team_id=team_id,
|
||||
team_alias=team_table.team_alias,
|
||||
role=member_role,
|
||||
user_email=user_email,
|
||||
spend=membership.spend,
|
||||
total_spend=membership.total_spend,
|
||||
budget_id=membership.budget_id,
|
||||
litellm_budget_table=membership.litellm_budget_table,
|
||||
spend=floored.spend,
|
||||
total_spend=floored.total_spend,
|
||||
budget_id=floored.budget_id,
|
||||
litellm_budget_table=floored.litellm_budget_table,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -171,4 +171,6 @@ class PrismaBatch(Protocol):
|
|||
@property
|
||||
def litellm_projecttable(self) -> BatchTable: ...
|
||||
|
||||
def execute_raw(self, query: str, *args: object) -> None: ...
|
||||
|
||||
async def commit(self) -> None: ...
|
||||
|
|
|
|||
|
|
@ -15,11 +15,11 @@ instead of full-model writes, which trip ``prisma.errors.DataError`` on rows
|
|||
carrying fields the update input type rejects (see #27730).
|
||||
"""
|
||||
|
||||
from collections.abc import AsyncGenerator, Callable, Mapping
|
||||
from collections.abc import AsyncGenerator, Callable, Mapping, Sequence
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import Final
|
||||
from typing import Final, Protocol, runtime_checkable
|
||||
|
||||
from litellm.repositories.prisma_protocols import BatchTable, PrismaBatch
|
||||
|
||||
|
|
@ -62,6 +62,13 @@ class TeamSpendResetWrites:
|
|||
)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class SpendResetWrites(Protocol):
|
||||
def queue_spend_zero(self, where: Mapping[str, object]) -> None: ...
|
||||
|
||||
def queue_spend_decrement(self, where: Mapping[str, object], amount: float) -> None: ...
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class LinkedSpendResetWrites:
|
||||
table: BatchTable
|
||||
|
|
@ -78,6 +85,30 @@ class LinkedSpendResetWrites:
|
|||
)
|
||||
|
||||
|
||||
# Membership rows written before total_spend existed (migration 20260421135425, default 0) hold their
|
||||
# history only in spend. Zeroing spend would erase it, so the floor lands in the same batch, ahead of the reset.
|
||||
_TEAM_MEMBERSHIP_TOTAL_SPEND_REPAIR_SQL: Final = """
|
||||
UPDATE "LiteLLM_TeamMembership"
|
||||
SET total_spend = spend
|
||||
WHERE budget_id = ANY($1::text[]) AND spend > total_spend
|
||||
"""
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TeamMembershipSpendResetWrites:
|
||||
resets: LinkedSpendResetWrites
|
||||
batch: PrismaBatch
|
||||
|
||||
def queue_spend_zero(self, where: Mapping[str, object]) -> None:
|
||||
self.resets.queue_spend_zero(where)
|
||||
|
||||
def queue_spend_decrement(self, where: Mapping[str, object], amount: float) -> None:
|
||||
self.resets.queue_spend_decrement(where, amount)
|
||||
|
||||
def queue_total_spend_repair(self, budget_ids: Sequence[str]) -> None:
|
||||
self.batch.execute_raw(_TEAM_MEMBERSHIP_TOTAL_SPEND_REPAIR_SQL, list(budget_ids))
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BudgetWindowWrites:
|
||||
table: BatchTable
|
||||
|
|
@ -106,7 +137,7 @@ class BudgetCascadeUnitOfWork:
|
|||
``enduser_spend_reset_unit_of_work`` because one tier can gate millions.
|
||||
"""
|
||||
|
||||
team_memberships: LinkedSpendResetWrites
|
||||
team_memberships: TeamMembershipSpendResetWrites
|
||||
keys: LinkedSpendResetWrites
|
||||
organizations: LinkedSpendResetWrites
|
||||
tags: LinkedSpendResetWrites
|
||||
|
|
@ -132,7 +163,9 @@ async def budget_cascade_unit_of_work(
|
|||
) -> AsyncGenerator[BudgetCascadeUnitOfWork, None]:
|
||||
batch = new_batch()
|
||||
yield BudgetCascadeUnitOfWork(
|
||||
team_memberships=LinkedSpendResetWrites(table=batch.litellm_teammembership),
|
||||
team_memberships=TeamMembershipSpendResetWrites(
|
||||
resets=LinkedSpendResetWrites(table=batch.litellm_teammembership), batch=batch
|
||||
),
|
||||
keys=LinkedSpendResetWrites(table=batch.litellm_verificationtoken),
|
||||
organizations=LinkedSpendResetWrites(table=batch.litellm_organizationtable),
|
||||
tags=LinkedSpendResetWrites(table=batch.litellm_tagtable),
|
||||
|
|
|
|||
|
|
@ -7,12 +7,12 @@ All values are configurable via proxy config.yaml.
|
|||
|
||||
import math
|
||||
import re
|
||||
import warnings
|
||||
import sys
|
||||
from collections.abc import Iterable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from types import MappingProxyType
|
||||
from typing import Annotated, Final, Literal, NamedTuple
|
||||
from typing import TYPE_CHECKING, Annotated, Final, Literal, NamedTuple
|
||||
|
||||
from pydantic import (
|
||||
ConfigDict,
|
||||
|
|
@ -27,8 +27,13 @@ from pydantic import (
|
|||
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("ignore", DeprecationWarning)
|
||||
if TYPE_CHECKING:
|
||||
import sre_constants
|
||||
import sre_parse
|
||||
elif sys.version_info >= (3, 11):
|
||||
from re import _constants as sre_constants
|
||||
from re import _parser as sre_parse
|
||||
else:
|
||||
import sre_constants
|
||||
import sre_parse
|
||||
|
||||
|
|
|
|||
|
|
@ -55,6 +55,8 @@ class LiteLLMKeyType(str, enum.Enum):
|
|||
|
||||
|
||||
class GenerateKeyRequest(KeyRequestBase):
|
||||
tpm_limit: int | None = Field(default=None, ge=0)
|
||||
rpm_limit: int | None = Field(default=None, ge=0)
|
||||
soft_budget: float | None = None
|
||||
send_invite_email: bool | None = None
|
||||
key_type: LiteLLMKeyType | None = Field(
|
||||
|
|
@ -82,6 +84,8 @@ def _is_mapping(value: object) -> TypeIs[Mapping[object, object]]: # guard-ok:
|
|||
|
||||
|
||||
class UpdateKeyRequest(KeyRequestBase):
|
||||
tpm_limit: int | None = Field(default=None, ge=0)
|
||||
rpm_limit: int | None = Field(default=None, ge=0)
|
||||
# Note: the defaults of all Params here MUST BE NONE
|
||||
# else they will get overwritten
|
||||
duration: str | None = None
|
||||
|
|
|
|||
|
|
@ -16,6 +16,8 @@ class NewUserRequestTeam(LiteLLMPydanticObjectBase):
|
|||
|
||||
|
||||
class NewUserRequest(GenerateRequestBase):
|
||||
tpm_limit: int | None = Field(default=None, ge=0)
|
||||
rpm_limit: int | None = Field(default=None, ge=0)
|
||||
max_budget: float | None = None
|
||||
user_email: str | None = None
|
||||
user_alias: str | None = None
|
||||
|
|
@ -47,6 +49,8 @@ class NewUserRequest(GenerateRequestBase):
|
|||
|
||||
|
||||
class UpdateUserRequestNoUserIDorEmail(GenerateRequestBase): # shared with BulkUpdateUserRequest
|
||||
tpm_limit: int | None = Field(default=None, ge=0)
|
||||
rpm_limit: int | None = Field(default=None, ge=0)
|
||||
# repr=False keeps the plaintext out of management-endpoint alerts, which str() the request model
|
||||
password: str | None = Field(default=None, repr=False)
|
||||
spend: float | None = None
|
||||
|
|
|
|||
|
|
@ -9571,7 +9571,7 @@ class ProviderConfigManager:
|
|||
|
||||
return get_nvidia_nim_rerank_config(model)
|
||||
elif litellm.LlmProviders.VERTEX_AI == provider:
|
||||
return litellm.VertexAIRerankConfig()
|
||||
return ProviderConfigManager._get_vertex_ai_rerank_config_cached()
|
||||
elif litellm.LlmProviders.FIREWORKS_AI == provider:
|
||||
return litellm.FireworksAIRerankConfig()
|
||||
elif litellm.LlmProviders.VOYAGE == provider:
|
||||
|
|
@ -9590,6 +9590,11 @@ class ProviderConfigManager:
|
|||
return get_dashscope_family_rerank_config(provider.value)
|
||||
return litellm.CohereRerankConfig()
|
||||
|
||||
@staticmethod
|
||||
@lru_cache(maxsize=1)
|
||||
def _get_vertex_ai_rerank_config_cached() -> BaseRerankConfig:
|
||||
return litellm.VertexAIRerankConfig()
|
||||
|
||||
@staticmethod
|
||||
def get_provider_decisions_config(model: str, provider: LlmProviders) -> BaseDecisionsConfig | None:
|
||||
if provider == LlmProviders.PERPLEXITY:
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
|
||||
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
|
||||
|
||||
FROM $UV_IMAGE AS uvbin
|
||||
|
|
@ -137,7 +137,9 @@ ENV HOME=/home/nonroot \
|
|||
PRISMA_CLI_QUERY_ENGINE_TYPE=binary \
|
||||
PRISMA_OFFLINE_MODE=true
|
||||
|
||||
COPY --from=builder --chown=nonroot:nonroot /app /app
|
||||
COPY --from=builder --chown=nonroot:nonroot /app/.venv /app/.venv
|
||||
COPY --from=builder --chown=nonroot:nonroot /app/litellm /app/litellm
|
||||
COPY --from=builder --chown=nonroot:nonroot /app/run.py /app/run.py
|
||||
COPY --from=builder /opt/prisma /opt/prisma
|
||||
|
||||
RUN chmod -R a+rX /opt/prisma && \
|
||||
|
|
|
|||
|
|
@ -33698,7 +33698,7 @@
|
|||
"input_cost_per_token": 1e-05,
|
||||
"input_cost_per_token_above_272k_tokens": 2e-05,
|
||||
"litellm_provider": "github_copilot",
|
||||
"max_input_tokens": 272000,
|
||||
"max_input_tokens": 922000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "responses",
|
||||
|
|
@ -33722,7 +33722,7 @@
|
|||
"input_cost_per_token": 1e-07,
|
||||
"input_cost_per_token_above_272k_tokens": 2e-07,
|
||||
"litellm_provider": "github_copilot",
|
||||
"max_input_tokens": 272000,
|
||||
"max_input_tokens": 922000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "responses",
|
||||
|
|
@ -33746,7 +33746,7 @@
|
|||
"input_cost_per_token": 2e-06,
|
||||
"input_cost_per_token_above_272k_tokens": 4e-06,
|
||||
"litellm_provider": "github_copilot",
|
||||
"max_input_tokens": 272000,
|
||||
"max_input_tokens": 922000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "responses",
|
||||
|
|
@ -33770,7 +33770,7 @@
|
|||
"input_cost_per_token": 2e-06,
|
||||
"input_cost_per_token_above_272k_tokens": 4e-06,
|
||||
"litellm_provider": "github_copilot",
|
||||
"max_input_tokens": 272000,
|
||||
"max_input_tokens": 922000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "responses",
|
||||
|
|
@ -36840,10 +36840,14 @@
|
|||
"cache_creation_input_token_cost": 1.25e-05,
|
||||
"cache_creation_input_token_cost_above_12hr": 2.5e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 0.0001,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": 0.0003,
|
||||
"cache_creation_input_token_cost_above_12hr_batches": 1.25e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_flex": 1.25e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_priority": 5e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_ultrafast": 0.00015,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 2.5e-05,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05,
|
||||
|
|
@ -36931,8 +36935,10 @@
|
|||
"cache_creation_input_token_cost": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr": 5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_batches": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_flex": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_priority": 1e-05,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 5e-06,
|
||||
|
|
@ -37014,8 +37020,10 @@
|
|||
"cache_creation_input_token_cost": 1.25e-07,
|
||||
"cache_creation_input_token_cost_above_12hr": 2.5e-07,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-07,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-07,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-07,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 1e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_batches": 1.25e-07,
|
||||
"cache_creation_input_token_cost_above_12hr_flex": 1.25e-07,
|
||||
"cache_creation_input_token_cost_above_12hr_priority": 5e-07,
|
||||
"cache_creation_input_token_cost_decisions": 1.25e-07,
|
||||
|
|
@ -37186,8 +37194,10 @@
|
|||
"cache_creation_input_token_cost": 5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr": 1e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens": 2e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 1e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 1e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 4e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_batches": 5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_flex": 5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_priority": 2e-05,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 1e-05,
|
||||
|
|
@ -37268,8 +37278,10 @@
|
|||
"cache_creation_input_token_cost": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr": 5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_batches": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_flex": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_priority": 1e-05,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 5e-06,
|
||||
|
|
@ -37349,8 +37361,10 @@
|
|||
"cache_creation_input_token_cost": 2.5e-07,
|
||||
"cache_creation_input_token_cost_above_12hr": 5e-07,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-07,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-07,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_batches": 2.5e-07,
|
||||
"cache_creation_input_token_cost_above_12hr_flex": 2.5e-07,
|
||||
"cache_creation_input_token_cost_above_12hr_priority": 1e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 5e-07,
|
||||
|
|
@ -42726,6 +42740,29 @@
|
|||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"mistral/zai-glm-5-3-flash": {
|
||||
"cache_read_input_token_cost": 3e-08,
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"litellm_provider": "mistral",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 5e-07,
|
||||
"reasoning_effort_levels": [
|
||||
"low",
|
||||
"high",
|
||||
"max"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/zai-glm-5-3-flash",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"mistral/zai-glm-5": {
|
||||
"cache_read_input_token_cost": 1.4e-07,
|
||||
"input_cost_per_token": 1.4e-06,
|
||||
|
|
@ -75023,10 +75060,14 @@
|
|||
"cache_creation_input_token_cost": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr": 5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": 6e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_batches": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_flex": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_priority": 1e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_ultrafast": 3e-05,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 5e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_batches": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06,
|
||||
|
|
@ -76416,9 +76457,10 @@
|
|||
"azure_ai/Microsoft-Decision-1": {
|
||||
"input_cost_per_token": 4.2e-08,
|
||||
"litellm_provider": "azure_ai",
|
||||
"max_input_tokens": 32768,
|
||||
"mode": "evaluation",
|
||||
"output_cost_per_token": 0.0,
|
||||
"source": "https://prices.azure.com/api/retail/prices?$filter=productName eq 'Microsoft Decision Models'"
|
||||
"source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure"
|
||||
},
|
||||
"gpt-rosalind-discovery": {
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
|
|
@ -78629,5 +78671,227 @@
|
|||
"prompt_cache_min_tokens": 512,
|
||||
"source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing",
|
||||
"supports_web_search": true
|
||||
},
|
||||
"gpt-6.1-terra": {
|
||||
"cache_creation_input_token_cost": 1.25e-06,
|
||||
"cache_creation_input_token_cost_above_12hr": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 1e-05,
|
||||
"cache_creation_input_token_cost_above_12hr_batches": 1.25e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_flex": 1.25e-06,
|
||||
"cache_creation_input_token_cost_above_12hr_priority": 5e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_batches": 1.25e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 5e-06,
|
||||
"cache_creation_input_token_cost_batches": 6.25e-07,
|
||||
"cache_creation_input_token_cost_flex": 6.25e-07,
|
||||
"cache_creation_input_token_cost_priority": 2.5e-06,
|
||||
"cache_read_input_token_cost": 5e-08,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 1e-07,
|
||||
"cache_read_input_token_cost_above_272k_tokens_batches": 5e-08,
|
||||
"cache_read_input_token_cost_above_272k_tokens_flex": 5e-08,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 2e-07,
|
||||
"cache_read_input_token_cost_batches": 2.5e-08,
|
||||
"cache_read_input_token_cost_flex": 2.5e-08,
|
||||
"cache_read_input_token_cost_priority": 1e-07,
|
||||
"input_cost_per_token": 1e-06,
|
||||
"input_cost_per_token_above_272k_tokens": 2e-06,
|
||||
"input_cost_per_token_above_272k_tokens_batches": 1e-06,
|
||||
"input_cost_per_token_above_272k_tokens_flex": 1e-06,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 4e-06,
|
||||
"input_cost_per_token_batches": 5e-07,
|
||||
"input_cost_per_token_flex": 5e-07,
|
||||
"input_cost_per_token_priority": 2e-06,
|
||||
"litellm_provider": "openai",
|
||||
"max_input_tokens": 922000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 5e-06,
|
||||
"output_cost_per_token_above_272k_tokens": 7.5e-06,
|
||||
"output_cost_per_token_above_272k_tokens_batches": 3.75e-06,
|
||||
"output_cost_per_token_above_272k_tokens_flex": 3.75e-06,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 1.5e-05,
|
||||
"output_cost_per_token_batches": 2.5e-06,
|
||||
"output_cost_per_token_flex": 2.5e-06,
|
||||
"output_cost_per_token_priority": 1e-05,
|
||||
"regional_processing_uplift_multiplier_eu": 1.1,
|
||||
"regional_processing_uplift_multiplier_us": 1.1,
|
||||
"search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01},
|
||||
"source": "https://developers.openai.com/api/docs/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions",
|
||||
"/v1/batch",
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_computer_use": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_native_streaming": true,
|
||||
"supports_none_reasoning_effort": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_cache_breakpoint": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_tool_search": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true,
|
||||
"supports_xhigh_reasoning_effort": true
|
||||
},
|
||||
"chatgpt/gpt-6.1-terra": {
|
||||
"litellm_provider": "chatgpt",
|
||||
"max_input_tokens": 922000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "responses",
|
||||
"source": "https://developers.openai.com/api/docs/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions",
|
||||
"/v1/responses"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_none_reasoning_effort": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_vision": true,
|
||||
"supports_xhigh_reasoning_effort": true
|
||||
},
|
||||
"global.openai.gpt-6.1-terra": {
|
||||
"cache_creation_input_token_cost": 1.25e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 2.5e-06,
|
||||
"cache_read_input_token_cost": 5e-08,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 1e-07,
|
||||
"input_cost_per_token": 1e-06,
|
||||
"input_cost_per_token_above_272k_tokens": 2e-06,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 922000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 5e-06,
|
||||
"output_cost_per_token_above_272k_tokens": 7.5e-06,
|
||||
"source": "https://developers.openai.com/api/docs/models/gpt-6.1-terra",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions",
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_bedrock_converse_tool_result_images": false,
|
||||
"supports_bedrock_runtime_chat_completions_response_format": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_none_reasoning_effort": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_xhigh_reasoning_effort": true
|
||||
},
|
||||
"us.openai.gpt-6.1-terra": {
|
||||
"cache_creation_input_token_cost": 1.375e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 2.75e-06,
|
||||
"cache_read_input_token_cost": 5.5e-08,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 1.1e-07,
|
||||
"input_cost_per_token": 1.1e-06,
|
||||
"input_cost_per_token_above_272k_tokens": 2.2e-06,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 922000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 5.5e-06,
|
||||
"output_cost_per_token_above_272k_tokens": 8.25e-06,
|
||||
"source": "https://developers.openai.com/api/docs/models/gpt-6.1-terra",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions",
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_bedrock_converse_tool_result_images": false,
|
||||
"supports_bedrock_runtime_chat_completions_response_format": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_none_reasoning_effort": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_xhigh_reasoning_effort": true
|
||||
},
|
||||
"baseten/decision-models/baseten-elephant-9b-latest": {
|
||||
"input_cost_per_token": 0.0,
|
||||
"litellm_provider": "baseten",
|
||||
"max_input_tokens": 16384,
|
||||
"mode": "decisions",
|
||||
"output_cost_per_token": 0.0,
|
||||
"source": "https://inference.baseten.co/v1/models",
|
||||
"supports_vision": true
|
||||
},
|
||||
"baseten/decision-models/baseten-mammoth-27b-latest": {
|
||||
"input_cost_per_token": 0.0,
|
||||
"litellm_provider": "baseten",
|
||||
"max_input_tokens": 16384,
|
||||
"mode": "decisions",
|
||||
"output_cost_per_token": 0.0,
|
||||
"source": "https://inference.baseten.co/v1/models",
|
||||
"supports_vision": true
|
||||
},
|
||||
"github_copilot/gpt-6.1-terra": {
|
||||
"cache_creation_input_token_cost": 1.25e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 2.5e-06,
|
||||
"cache_read_input_token_cost": 5e-08,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 1e-07,
|
||||
"input_cost_per_token": 1e-06,
|
||||
"input_cost_per_token_above_272k_tokens": 2e-06,
|
||||
"litellm_provider": "github_copilot",
|
||||
"max_input_tokens": 922000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "responses",
|
||||
"output_cost_per_token": 5e-06,
|
||||
"output_cost_per_token_above_272k_tokens": 7.5e-06,
|
||||
"source": "https://github.blog/changelog/2026-10-26-gpt-6-1-terra-in-github-copilot",
|
||||
"supported_endpoints": [
|
||||
"/v1/responses"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_vision": true
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -3,14 +3,14 @@ name = "litellm-core"
|
|||
version = "1.105.0"
|
||||
description = "Library to easily interface with LLM API providers"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10, <3.15"
|
||||
requires-python = ">=3.10, <3.16"
|
||||
license = "MIT"
|
||||
license-files = ["LICENSE"]
|
||||
authors = [
|
||||
{ name = "BerriAI" },
|
||||
]
|
||||
dependencies = [
|
||||
"fastuuid>=0.14.0,<1.0",
|
||||
"fastuuid>=0.14.0,<1.0; python_version < '3.15'",
|
||||
"filelock>=3.16.1,<4.0",
|
||||
"httpx[http2]>=0.28.0,<1.0",
|
||||
"openai>=2.20.0,<4.0.0",
|
||||
|
|
@ -20,14 +20,16 @@ dependencies = [
|
|||
"packaging>=24.0",
|
||||
"importlib-metadata>=8.0.0,<9.0",
|
||||
"tiktoken>=0.8.0,<1.0; python_version < '3.14'",
|
||||
"tiktoken>=0.12.0,<1.0; python_version >= '3.14'",
|
||||
"tiktoken>=0.12.0,<1.0; python_version == '3.14'",
|
||||
"tiktoken>=0.14.0,<1.0; python_version >= '3.15'",
|
||||
"click>=8.0.0,<9.0",
|
||||
"jinja2>=3.1.6,<4.0",
|
||||
"aiohttp>=3.14.2,<4.0",
|
||||
"anyio>=4.14.0,<5.0",
|
||||
"async-timeout>=4.0.3,<6.0; python_version < '3.11'",
|
||||
"pydantic>=2.11.0,<3.0.0; python_version < '3.14'",
|
||||
"pydantic>=2.12.0,<3.0.0; python_version >= '3.14'",
|
||||
"pydantic>=2.12.0,<3.0.0; python_version == '3.14'",
|
||||
"pydantic>=2.14.0,<3.0.0; python_version >= '3.15'",
|
||||
"pydantic-settings>=2.14.1,<3.0",
|
||||
"jsonschema>=4.0.0,<5.0",
|
||||
"certifi>=2026.10.19",
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ name = "litellm"
|
|||
version = "1.108.0"
|
||||
description = "Library to easily interface with LLM API providers"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10, <3.15"
|
||||
requires-python = ">=3.10, <3.16"
|
||||
license = "MIT"
|
||||
license-files = ["LICENSE"]
|
||||
authors = [
|
||||
|
|
@ -14,7 +14,7 @@ dependencies = [
|
|||
# deps. Reproducibility for our Docker/CI comes from `uv.lock`.
|
||||
# When changing a floor, verify it installs + imports on every supported
|
||||
# Python with: `uv pip install --resolution=lowest-direct .`
|
||||
"fastuuid>=0.14.0,<1.0",
|
||||
"fastuuid>=0.14.0,<1.0; python_version < '3.15'",
|
||||
"filelock>=3.16.1,<4.0",
|
||||
"httpx[http2]>=0.28.0,<1.0",
|
||||
"openai>=2.20.0,<4.0.0",
|
||||
|
|
@ -23,7 +23,8 @@ dependencies = [
|
|||
"packaging>=24.0",
|
||||
"importlib-metadata>=8.0.0,<9.0",
|
||||
"tiktoken>=0.8.0,<1.0; python_version < '3.14'",
|
||||
"tiktoken>=0.12.0,<1.0; python_version >= '3.14'",
|
||||
"tiktoken>=0.12.0,<1.0; python_version == '3.14'",
|
||||
"tiktoken>=0.14.0,<1.0; python_version >= '3.15'",
|
||||
"tokenizers>=0.21.0,<1.0",
|
||||
"huggingface-hub>=0.34.0,<2.0",
|
||||
"click>=8.0.0,<9.0",
|
||||
|
|
@ -32,7 +33,8 @@ dependencies = [
|
|||
"anyio>=4.14.0,<5.0",
|
||||
"async-timeout>=4.0.3,<6.0; python_version < '3.11'",
|
||||
"pydantic>=2.11.0,<3.0.0; python_version < '3.14'",
|
||||
"pydantic>=2.12.0,<3.0.0; python_version >= '3.14'",
|
||||
"pydantic>=2.12.0,<3.0.0; python_version == '3.14'",
|
||||
"pydantic>=2.14.0,<3.0.0; python_version >= '3.15'",
|
||||
"pydantic-settings>=2.14.1,<3.0",
|
||||
"jsonschema>=4.0.0,<5.0",
|
||||
"boto3>=1.43.117,<2.0",
|
||||
|
|
@ -55,17 +57,20 @@ Documentation = "https://docs.litellm.ai"
|
|||
proxy = [
|
||||
"gunicorn>=23.0.0,<24.0",
|
||||
"uvicorn>=0.33.0,<1.0",
|
||||
"granian>=2.7.4,<3.0",
|
||||
"granian>=2.7.4,<3.0; python_version < '3.15'",
|
||||
"granian>=2.8.0,<3.0; python_version >= '3.15'",
|
||||
"uvloop>=0.22.1,<1.0; sys_platform != 'win32'",
|
||||
"fastapi>=0.136.3,<1.0",
|
||||
"starlette>=1.0.1,<2.0",
|
||||
"backoff>=2.2.1,<3.0",
|
||||
"pyyaml>=6.0.3,<7.0",
|
||||
"rq>=2.7.0,<3.0",
|
||||
"orjson>=3.11.6,<4.0",
|
||||
"orjson>=3.11.6,<4.0; python_version < '3.15'",
|
||||
"orjson>=3.12.0,<4.0; python_version >= '3.15'",
|
||||
# redis-py's C response parser. It arrives with redis (via rq) either way; naming
|
||||
# it here is what makes redis-py select _HiredisParser instead of the Python one.
|
||||
"hiredis>=3.0.0,<4.0",
|
||||
"hiredis>=3.0.0,<4.0; python_version < '3.15'",
|
||||
"hiredis>=3.5.0,<4.0; python_version >= '3.15'",
|
||||
"apscheduler>=3.11.2,<4.0",
|
||||
"fastapi-sso>=0.19.0,<1.0",
|
||||
"PyJWT>=2.13.0,<3.0",
|
||||
|
|
@ -210,8 +215,8 @@ dev = [
|
|||
"tomli==2.4.1; python_version < '3.11'",
|
||||
"pytest-mock==3.15.1",
|
||||
"pytest-asyncio==1.3.0",
|
||||
"psycopg==3.3.3",
|
||||
"psycopg-binary==3.3.3",
|
||||
"psycopg==3.3.6",
|
||||
"psycopg-binary==3.3.6",
|
||||
"pytest-xdist==3.8.0",
|
||||
"requests-mock==1.12.1",
|
||||
"responses==0.26.0",
|
||||
|
|
@ -374,11 +379,6 @@ exclude-newer = "3 days"
|
|||
[tool.uv.exclude-newer-package]
|
||||
# TODO: ryan, remove after 2026-10-23, once certifi 2026.10.19 (drops the compromised eMudhra root) clears the 3-day window
|
||||
certifi = "2026-10-20T03:00:00Z"
|
||||
# TODO(shrey): remove after 2026-10-25, once grpcio 1.84.1 (drops the compromised eMudhra roots from roots.pem) clears the 3-day window
|
||||
grpcio = "2026-10-22T15:00:00Z"
|
||||
# TODO(shrey): remove after 2026-10-24, once boto3/botocore 1.43.117 (drops the compromised eMudhra roots from cacert.pem) clear the 3-day window
|
||||
boto3 = "2026-10-21T19:31:00Z"
|
||||
botocore = "2026-10-21T19:31:00Z"
|
||||
# TODO: remove after 2026-10-25, once brotli 1.2.1 (CVE-2026-111497) clears the 3-day window
|
||||
brotli = "2026-10-22T14:01:00Z"
|
||||
# TODO: ryan, remove after 2026-10-24, once mlflow 3.17.2 (GHSA-wqmh-qccw-pqwf) clears the 3-day window
|
||||
|
|
|
|||
|
|
@ -0,0 +1,723 @@
|
|||
"""A team's Arize sampling rate decides once per trace, so every pod and worker serving a call of it agrees.
|
||||
|
||||
Two proxies share one database, as two pods behind a load balancer do: the first runs one worker, the second
|
||||
two. Calls of one trace carry the caller's ``traceparent`` and land on different pods and workers; the team's
|
||||
Arize space must receive every call of a trace or none of them.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import re
|
||||
import uuid
|
||||
from collections.abc import Callable, Generator, Iterator, Mapping, Sequence
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass, replace
|
||||
from itertools import chain, product
|
||||
from pathlib import Path
|
||||
from typing import Final, Literal, TypeAlias
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import (
|
||||
Scenario,
|
||||
eventually,
|
||||
gateway_from_environment,
|
||||
list_value,
|
||||
object_value,
|
||||
until_asserted,
|
||||
)
|
||||
from integration._support.database import scratch_database
|
||||
from integration._support.otlp_sink import Span, SpanSinks, configure_sink, recorded_requests, recorded_spans
|
||||
from integration._support.process import OwnedProxy, WorkerHealthcheck, owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
MODEL: Final = "team-sampling"
|
||||
SERVER: Final = 2
|
||||
CALL_ID: Final = "litellm.call_id"
|
||||
OPERATION: Final = "gen_ai.operation.name"
|
||||
SPACE: Final = "space-sampled"
|
||||
CREDENTIALS: Final = {"arize_space_key": SPACE, "arize_api_key": "arize-sampled"}
|
||||
FAILING: Final = re.compile(rb"sampling-fail")
|
||||
JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
|
||||
TRACES: Final = 16
|
||||
Endpoint: TypeAlias = Literal["chat", "responses", "messages"]
|
||||
Verdicts: TypeAlias = Mapping[str, tuple[bool, ...]]
|
||||
|
||||
|
||||
def _chat_reply(stream: bool) -> Reply:
|
||||
usage: Final[JsonValue] = {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9}
|
||||
if not stream:
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": "chatcmpl-sampling",
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "gpt-4o-mini",
|
||||
"choices": [
|
||||
{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}
|
||||
],
|
||||
"usage": usage,
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
chunk: Final = {"id": "chatcmpl-sampling", "object": "chat.completion.chunk", "created": 1, "model": "gpt-4o-mini"}
|
||||
deltas: Final[tuple[JsonValue, ...]] = (
|
||||
{**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}}]},
|
||||
{**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]},
|
||||
{**chunk, "choices": [], "usage": usage},
|
||||
)
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=(*(b"data: " + json.dumps(delta).encode() + b"\n\n" for delta in deltas), b"data: [DONE]\n\n"),
|
||||
)
|
||||
|
||||
|
||||
def _responses_reply(stream: bool) -> Reply:
|
||||
response: Final[Mapping[str, JsonValue]] = {
|
||||
"id": "resp_sampling",
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": "completed",
|
||||
"model": "gpt-4o-mini",
|
||||
"output": [
|
||||
{
|
||||
"id": "msg_sampling",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"status": "completed",
|
||||
"content": [{"type": "output_text", "text": "ok", "annotations": []}],
|
||||
}
|
||||
],
|
||||
"usage": {"input_tokens": 7, "output_tokens": 2, "total_tokens": 9},
|
||||
}
|
||||
if not stream:
|
||||
return Reply(body=json.dumps(response).encode())
|
||||
events: Final[tuple[Mapping[str, JsonValue], ...]] = (
|
||||
{
|
||||
"type": "response.created",
|
||||
"sequence_number": 0,
|
||||
"response": {**response, "status": "in_progress", "output": []},
|
||||
},
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"sequence_number": 1,
|
||||
"item_id": "msg_sampling",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": "ok",
|
||||
},
|
||||
{"type": "response.completed", "sequence_number": 2, "response": dict(response)},
|
||||
)
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events),
|
||||
)
|
||||
|
||||
|
||||
def _upstream(request: Request) -> Reply:
|
||||
if FAILING.search(request.body) is not None:
|
||||
return Reply(status=500, body=b'{"error":{"message":"scripted upstream failure","type":"server_error"}}')
|
||||
stream: Final = object_value(JSON.validate_json(request.body)).get("stream") is True
|
||||
if request.target.endswith("/responses"):
|
||||
return _responses_reply(stream)
|
||||
return _chat_reply(stream)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Sent:
|
||||
trace_id: str
|
||||
parent_span_id: str
|
||||
call_id: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Call:
|
||||
path: str
|
||||
body: Mapping[str, JsonValue]
|
||||
headers: Mapping[str, str]
|
||||
sent: Sent
|
||||
expected: int
|
||||
|
||||
|
||||
def _call(
|
||||
key: str, trace_id: str | None, endpoint: Endpoint = "chat", stream: bool = False, content: str = "hello"
|
||||
) -> Call:
|
||||
parent: Final = uuid.uuid4().hex[:16]
|
||||
call_id: Final = str(uuid.uuid4())
|
||||
messages: Final[JsonValue] = [{"role": "user", "content": content}]
|
||||
path, body = {
|
||||
"chat": ("/v1/chat/completions", {"model": MODEL, "messages": messages, "stream": stream}),
|
||||
"responses": ("/v1/responses", {"model": MODEL, "input": content, "stream": stream}),
|
||||
"messages": ("/v1/messages", {"model": MODEL, "max_tokens": 16, "messages": messages, "stream": stream}),
|
||||
}[endpoint]
|
||||
headers: Final = {
|
||||
"Authorization": f"Bearer {key}",
|
||||
"x-litellm-call-id": call_id,
|
||||
"connection": "close",
|
||||
**({} if trace_id is None else {"traceparent": f"00-{trace_id}-{parent}-01"}),
|
||||
}
|
||||
expected: Final = 500 if FAILING.search(content.encode()) else 200
|
||||
return Call(path, body, headers, Sent(trace_id or "", parent if trace_id else "", call_id), expected)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Pods:
|
||||
"""``first`` serves with one worker, ``second`` with two."""
|
||||
|
||||
first: OwnedProxy
|
||||
second: OwnedProxy
|
||||
scenario: Scenario
|
||||
sinks: SpanSinks
|
||||
|
||||
def team_key(self, variables: Mapping[str, str]) -> str:
|
||||
team: Final = self.scenario.team()
|
||||
attached: Final = self.first.gateway.request(
|
||||
"POST", f"/team/{team}/callback", {"callback_name": "arize", "callback_vars": dict(variables)}
|
||||
)
|
||||
assert attached.status_code == 200, attached.text
|
||||
return self.scenario.key(team_id=team)
|
||||
|
||||
def send(self, pod: OwnedProxy, call: Call) -> Sent:
|
||||
with pod.gateway.client.stream("POST", call.path, json=call.body, headers=call.headers) as response:
|
||||
text: Final = response.read().decode()
|
||||
assert response.status_code == call.expected, text
|
||||
return call.sent
|
||||
|
||||
async def send_async(self, client: httpx.AsyncClient, call: Call) -> Sent:
|
||||
async with client.stream("POST", call.path, json=call.body, headers=call.headers) as response:
|
||||
text: Final = (await response.aread()).decode()
|
||||
assert response.status_code == call.expected, text
|
||||
return call.sent
|
||||
|
||||
def on_each_worker(self, pod: OwnedProxy, call: Callable[[], Call]) -> tuple[Sent, ...]:
|
||||
"""One call served by each worker of ``pod``, the others suspended while it runs."""
|
||||
workers: Final = sorted(pod.worker_pids())
|
||||
assert len(workers) == pod.gateway.workers, workers
|
||||
sent: Final[list[Sent]] = [] # mutable-ok: one call per worker, in worker order
|
||||
for serving in workers:
|
||||
with _suspended_all_but(pod, workers, serving):
|
||||
sent.append(self.send(pod, call()))
|
||||
return tuple(sent)
|
||||
|
||||
def fence(self, key: str) -> tuple[Sent, ...]:
|
||||
"""Calls of an unsampled team on the sampled team's exporter, one per worker of both pods.
|
||||
|
||||
Each worker sends spans to that exporter in the order they were handed to it, so all of these
|
||||
arriving means every call any worker decided before them has been sent too.
|
||||
"""
|
||||
return (self.send(self.first, _call(key, uuid.uuid4().hex)),) + self.on_each_worker(
|
||||
self.second, lambda: _call(key, uuid.uuid4().hex)
|
||||
)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _suspended_all_but(pod: OwnedProxy, workers: Sequence[int], serving: int) -> Generator[None]:
|
||||
if len(workers) == 1:
|
||||
yield
|
||||
return
|
||||
others: Final = [pid for pid in workers if pid != serving]
|
||||
with pod.suspended(others[0]):
|
||||
yield
|
||||
|
||||
|
||||
def _space_id(request: Mapping[str, JsonValue]) -> JsonValue:
|
||||
return {name.lower(): value for name, value in object_value(request.get("headers") or {}).items()}.get(
|
||||
"arize-space-id"
|
||||
)
|
||||
|
||||
|
||||
def _space_span_ids(url: str) -> frozenset[str]:
|
||||
"""Span ids the Arize sink received in export requests sent with the team's space id."""
|
||||
ours: Final = (
|
||||
list_value(request.get("span_ids") or []) for request in recorded_requests(url) if _space_id(request) == SPACE
|
||||
)
|
||||
return frozenset(span_id for span_id in chain.from_iterable(ours) if isinstance(span_id, str))
|
||||
|
||||
|
||||
def _team_spans(pods: Pods, since: int) -> tuple[Span, ...]:
|
||||
owned: Final = _space_span_ids(pods.sinks.arize)
|
||||
return tuple(span for span in recorded_spans(pods.sinks.arize, since)[1] if span["span_id"] in owned)
|
||||
|
||||
|
||||
def _delivered(spans: tuple[Span, ...], sent: Sequence[Sent]) -> frozenset[str]:
|
||||
"""Call ids of ``sent`` whose server span is among ``spans``."""
|
||||
by_parent: Final = {(call.trace_id, call.parent_span_id): call.call_id for call in sent}
|
||||
return frozenset(
|
||||
by_parent[(span["trace_id"], span["parent_span_id"])]
|
||||
for span in spans
|
||||
if span["kind"] == SERVER and (span["trace_id"], span["parent_span_id"]) in by_parent
|
||||
)
|
||||
|
||||
|
||||
def _await_logged(pods: Pods, since: int, sent: Sequence[Sent]) -> None:
|
||||
"""Wait until every call of ``sent`` has its model span in the operator's unsampled copy."""
|
||||
expected: Final = {call.call_id for call in sent}
|
||||
eventually(
|
||||
lambda: (
|
||||
{
|
||||
str(span["attributes"][CALL_ID])
|
||||
for span in recorded_spans(pods.sinks.operator, since)[1]
|
||||
if OPERATION in span["attributes"] and CALL_ID in span["attributes"]
|
||||
}
|
||||
& expected
|
||||
),
|
||||
lambda logged: logged == expected,
|
||||
seconds=60,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Cursors:
|
||||
operator: int
|
||||
arize: int
|
||||
|
||||
|
||||
def _cursors(pods: Pods) -> Cursors:
|
||||
return Cursors(recorded_spans(pods.sinks.operator)[0], recorded_spans(pods.sinks.arize)[0])
|
||||
|
||||
|
||||
def _await_flushed(pods: Pods, cursors: Cursors, fence_key: str) -> None:
|
||||
"""Wait until every worker of both pods has sent the team's exporter everything it decided so far."""
|
||||
fences: Final = pods.fence(fence_key)
|
||||
|
||||
def arrived() -> None:
|
||||
delivered: Final = _delivered(_team_spans(pods, cursors.arize), fences)
|
||||
assert {fence.call_id for fence in fences} <= delivered, "the unsampled team's calls have not arrived yet"
|
||||
|
||||
until_asserted(arrived, seconds=30)
|
||||
|
||||
|
||||
def _settled(
|
||||
pods: Pods,
|
||||
cursors: Cursors,
|
||||
by_trace: Mapping[str, Sequence[Sent]],
|
||||
fence_key: str,
|
||||
holds: Callable[[tuple[bool, ...]], bool] = lambda kept: len(set(kept)) == 1,
|
||||
) -> Verdicts:
|
||||
"""Per trace, which of its calls reached the team's space, once every worker has sent what it decided
|
||||
and ``holds`` is true of every trace."""
|
||||
calls: Final = tuple(chain.from_iterable(by_trace.values()))
|
||||
_await_logged(pods, cursors.operator, calls)
|
||||
_await_flushed(pods, cursors, fence_key)
|
||||
|
||||
def check() -> Verdicts:
|
||||
delivered: Final = _delivered(_team_spans(pods, cursors.arize), calls)
|
||||
verdicts: Final = {
|
||||
trace_id: tuple(call.call_id in delivered for call in sent) for trace_id, sent in by_trace.items()
|
||||
}
|
||||
broken: Final = {trace_id: kept for trace_id, kept in verdicts.items() if not holds(kept)}
|
||||
assert broken == {}, f"{len(broken)} of {len(verdicts)} traces were split (per call, kept): {broken}"
|
||||
return verdicts
|
||||
|
||||
return until_asserted(check, seconds=30)
|
||||
|
||||
|
||||
def _trace_ids(cell: str, count: int = TRACES) -> tuple[str, ...]:
|
||||
"""Fixed RFC 4122 shaped trace ids, like the uuid4 hex a caller mints for ``traceparent``."""
|
||||
return tuple(uuid.uuid5(uuid.NAMESPACE_URL, f"team-sampling/{cell}/{index}").hex for index in range(count))
|
||||
|
||||
|
||||
def _kept(verdicts: Verdicts) -> int:
|
||||
return sum(1 for kept in verdicts.values() if all(kept))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("variables", "content"),
|
||||
[({"arize_success_sampling_rate": "0.5"}, "hello"), ({"arize_error_sampling_rate": "0.5"}, "sampling-fail")],
|
||||
ids=["success-rate", "error-rate"],
|
||||
)
|
||||
def test_each_trace_reaches_the_team_space_whole_or_not_at_all(
|
||||
pods: Pods, variables: Mapping[str, str], content: str
|
||||
) -> None:
|
||||
"""The ticket's case: three calls per trace load balanced over the pods."""
|
||||
sampled: Final = pods.team_key({**CREDENTIALS, **variables})
|
||||
cursors: Final = _cursors(pods)
|
||||
by_trace: Final = {
|
||||
trace_id: tuple(
|
||||
pods.send(pod, _call(sampled, trace_id, content=content)) for pod in (pods.first, pods.second, pods.first)
|
||||
)
|
||||
for trace_id in _trace_ids("-".join(sorted(variables)))
|
||||
}
|
||||
|
||||
verdicts: Final = _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS))
|
||||
|
||||
assert 0 < _kept(verdicts) < TRACES, f"a rate of 0.5 kept {_kept(verdicts)} of {TRACES} traces"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("client", ["sync", "async"])
|
||||
def test_every_endpoint_follows_the_verdict_of_its_trace(pods: Pods, client: str) -> None:
|
||||
sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.5"})
|
||||
cursors: Final = _cursors(pods)
|
||||
plan: Final[tuple[tuple[OwnedProxy, Endpoint, bool], ...]] = (
|
||||
(pods.first, "chat", False),
|
||||
(pods.second, "chat", True),
|
||||
(pods.first, "responses", False),
|
||||
(pods.second, "messages", True),
|
||||
)
|
||||
|
||||
async def send_trace(trace_id: str) -> tuple[Sent, ...]:
|
||||
clients: Final = {
|
||||
id(pod): httpx.AsyncClient(base_url=str(pod.gateway.client.base_url), timeout=15, trust_env=False)
|
||||
for pod in (pods.first, pods.second)
|
||||
}
|
||||
try:
|
||||
return tuple(
|
||||
await asyncio.gather(
|
||||
*(
|
||||
pods.send_async(clients[id(pod)], _call(sampled, trace_id, endpoint, stream))
|
||||
for pod, endpoint, stream in plan
|
||||
)
|
||||
)
|
||||
)
|
||||
finally:
|
||||
for opened in clients.values():
|
||||
await opened.aclose()
|
||||
|
||||
by_trace: Final = {
|
||||
trace_id: (
|
||||
tuple(pods.send(pod, _call(sampled, trace_id, endpoint, stream)) for pod, endpoint, stream in plan)
|
||||
if client == "sync"
|
||||
else asyncio.run(send_trace(trace_id))
|
||||
)
|
||||
for trace_id in _trace_ids(f"endpoints-{client}")
|
||||
}
|
||||
|
||||
verdicts: Final = _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS))
|
||||
|
||||
assert 0 < _kept(verdicts) < TRACES
|
||||
|
||||
|
||||
def test_the_workers_of_one_pod_agree(pods: Pods) -> None:
|
||||
sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.5"})
|
||||
cursors: Final = _cursors(pods)
|
||||
by_trace: Final = {
|
||||
trace_id: pods.on_each_worker(pods.second, lambda trace_id=trace_id: _call(sampled, trace_id))
|
||||
for trace_id in _trace_ids("workers")
|
||||
}
|
||||
|
||||
verdicts: Final = _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS))
|
||||
|
||||
assert 0 < _kept(verdicts) < TRACES
|
||||
|
||||
|
||||
def test_a_replaced_worker_reaches_the_verdict_its_predecessor_reached(pods: Pods) -> None:
|
||||
"""A worker that restarts forgets the verdicts it remembered; the calls it serves next for the same
|
||||
traces still answer the way the first ones did."""
|
||||
sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.5"})
|
||||
cursors: Final = _cursors(pods)
|
||||
fence_key: Final = pods.team_key(CREDENTIALS)
|
||||
traces: Final = _trace_ids("replaced")
|
||||
before: Final = {
|
||||
trace_id: pods.on_each_worker(pods.second, lambda t=trace_id: _call(sampled, t)) for trace_id in traces
|
||||
}
|
||||
_await_logged(pods, cursors.operator, tuple(chain.from_iterable(before.values())))
|
||||
_await_flushed(pods, cursors, fence_key)
|
||||
for victim in sorted(pods.second.worker_pids()):
|
||||
pods.second.kill_worker(victim)
|
||||
by_trace: Final = {
|
||||
trace_id: before[trace_id] + pods.on_each_worker(pods.second, lambda t=trace_id: _call(sampled, t))
|
||||
for trace_id in traces
|
||||
}
|
||||
|
||||
verdicts: Final = _settled(pods, cursors, by_trace, fence_key)
|
||||
|
||||
assert 0 < _kept(verdicts) < TRACES
|
||||
|
||||
|
||||
def test_a_failed_call_is_kept_wherever_a_successful_call_of_its_trace_is(pods: Pods) -> None:
|
||||
"""One call of the trace fails on one pod and answers to the error rate, another succeeds on the other
|
||||
pod and answers to the success rate: with the error rate above the success rate, no trace keeps its
|
||||
successful call and drops its failed one."""
|
||||
sampled: Final = pods.team_key(
|
||||
{**CREDENTIALS, "arize_success_sampling_rate": "0.3", "arize_error_sampling_rate": "0.7"}
|
||||
)
|
||||
cursors: Final = _cursors(pods)
|
||||
by_trace: Final = {
|
||||
trace_id: (
|
||||
pods.send(pods.first, _call(sampled, trace_id, content="sampling-fail")),
|
||||
pods.send(pods.second, _call(sampled, trace_id)),
|
||||
)
|
||||
for trace_id in _trace_ids("mixed", 32)
|
||||
}
|
||||
|
||||
verdicts: Final = _settled(
|
||||
pods, cursors, by_trace, pods.team_key(CREDENTIALS), holds=lambda kept: kept[0] or not kept[1]
|
||||
)
|
||||
|
||||
assert any(kept == (True, False) for kept in verdicts.values()), "no trace sat between the two rates"
|
||||
assert any(kept == (True, True) for kept in verdicts.values()), "no trace sat under both rates"
|
||||
|
||||
|
||||
def test_a_failed_call_answers_to_the_error_rate_on_a_pod_that_dropped_its_trace(pods: Pods) -> None:
|
||||
"""A pod that already dropped a successful call of the trace still keeps a later failed call of it at an
|
||||
unset error rate, as the pod that only served the failed call does."""
|
||||
sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.3"})
|
||||
cursors: Final = _cursors(pods)
|
||||
traces: Final = _trace_ids("failed-after-dropped")
|
||||
succeeded: Final = {trace_id: pods.send(pods.first, _call(sampled, trace_id)) for trace_id in traces}
|
||||
_await_logged(pods, cursors.operator, tuple(succeeded.values()))
|
||||
by_trace: Final = {
|
||||
trace_id: (
|
||||
succeeded[trace_id],
|
||||
pods.send(pods.first, _call(sampled, trace_id, content="sampling-fail")),
|
||||
pods.send(pods.second, _call(sampled, trace_id, content="sampling-fail")),
|
||||
)
|
||||
for trace_id in traces
|
||||
}
|
||||
|
||||
verdicts: Final = _settled(
|
||||
pods, cursors, by_trace, pods.team_key(CREDENTIALS), holds=lambda kept: kept[1] and kept[2]
|
||||
)
|
||||
|
||||
assert any(not kept[0] for kept in verdicts.values()), "no successful call was dropped before its trace failed"
|
||||
|
||||
|
||||
def test_a_successful_call_answers_to_the_success_rate_on_a_pod_that_kept_its_failed_call(pods: Pods) -> None:
|
||||
"""A pod that already kept a failed call of the trace drops a later successful call of it at a zero
|
||||
success rate, as the pod that only served the successful call does."""
|
||||
sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.0"})
|
||||
cursors: Final = _cursors(pods)
|
||||
traces: Final = _trace_ids("succeeded-after-kept")
|
||||
failed: Final = {
|
||||
trace_id: pods.send(pods.first, _call(sampled, trace_id, content="sampling-fail")) for trace_id in traces
|
||||
}
|
||||
_await_logged(pods, cursors.operator, tuple(failed.values()))
|
||||
by_trace: Final = {
|
||||
trace_id: (
|
||||
failed[trace_id],
|
||||
pods.send(pods.first, _call(sampled, trace_id)),
|
||||
pods.send(pods.second, _call(sampled, trace_id)),
|
||||
)
|
||||
for trace_id in traces
|
||||
}
|
||||
|
||||
_settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS), holds=lambda kept: kept == (True, False, False))
|
||||
|
||||
|
||||
def test_concurrent_calls_of_one_trace_on_one_worker_each_answer_to_their_own_outcome(pods: Pods) -> None:
|
||||
"""A failed and a successful call of the trace in flight at once on the single worker of the first pod
|
||||
are decided apart: the failed one is kept at an unset error rate, the successful one dropped at a zero
|
||||
success rate."""
|
||||
sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.0"})
|
||||
cursors: Final = _cursors(pods)
|
||||
traces: Final = _trace_ids("concurrent-mixed")
|
||||
plan: Final = tuple(product(traces, ("sampling-fail", "hello")))
|
||||
|
||||
def send(job: tuple[str, str]) -> Sent:
|
||||
return pods.send(pods.first, _call(sampled, job[0], content=job[1]))
|
||||
|
||||
with ThreadPoolExecutor(max_workers=8) as pool:
|
||||
sent: Final = tuple(pool.map(send, plan))
|
||||
by_trace: Final = {trace_id: tuple(call for call in sent if call.trace_id == trace_id) for trace_id in traces}
|
||||
|
||||
_settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS), holds=lambda kept: kept == (True, False))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("variables", "expected"),
|
||||
[({"arize_success_sampling_rate": "1.0"}, TRACES), ({"arize_success_sampling_rate": "0.0"}, 0), ({}, TRACES)],
|
||||
ids=["one", "zero", "unset"],
|
||||
)
|
||||
def test_a_rate_at_either_end_keeps_every_trace_or_none(
|
||||
pods: Pods, variables: Mapping[str, str], expected: int
|
||||
) -> None:
|
||||
sampled: Final = pods.team_key({**CREDENTIALS, **variables})
|
||||
cursors: Final = _cursors(pods)
|
||||
by_trace: Final = {
|
||||
trace_id: tuple(pods.send(pod, _call(sampled, trace_id)) for pod in (pods.first, pods.second))
|
||||
for trace_id in _trace_ids(f"edge-{'-'.join(f'{k}={v}' for k, v in sorted(variables.items()))}")
|
||||
}
|
||||
|
||||
verdicts: Final = _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS))
|
||||
|
||||
assert _kept(verdicts) == expected
|
||||
|
||||
|
||||
def test_a_request_without_a_traceparent_is_sampled_whole_at_the_rate(pods: Pods) -> None:
|
||||
"""With no caller trace the proxy starts one per request: each request's server and model spans travel
|
||||
together and the rate still keeps some requests and drops others."""
|
||||
sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.5"})
|
||||
cursors: Final = _cursors(pods)
|
||||
sent: Final = tuple(
|
||||
pods.send(pod, _call(sampled, None)) for _, pod in product(range(12), (pods.first, pods.second))
|
||||
)
|
||||
_await_logged(pods, cursors.operator, sent)
|
||||
trace_of: Final = {
|
||||
str(span["attributes"][CALL_ID]): span["trace_id"]
|
||||
for span in recorded_spans(pods.sinks.operator, cursors.operator)[1]
|
||||
if CALL_ID in span["attributes"]
|
||||
}
|
||||
fences: Final = pods.fence(pods.team_key(CREDENTIALS))
|
||||
|
||||
def check() -> Mapping[str, tuple[bool, bool]]:
|
||||
spans: Final = _team_spans(pods, cursors.arize)
|
||||
assert {fence.call_id for fence in fences} <= _delivered(spans, fences), (
|
||||
"the unsampled team's calls have not arrived yet"
|
||||
)
|
||||
kinds: Final = {
|
||||
call.call_id: (
|
||||
any(span["kind"] == SERVER for span in spans if span["trace_id"] == trace_of[call.call_id]),
|
||||
any(OPERATION in span["attributes"] for span in spans if span["trace_id"] == trace_of[call.call_id]),
|
||||
)
|
||||
for call in sent
|
||||
}
|
||||
partial: Final = {call_id: kept for call_id, kept in kinds.items() if kept[0] != kept[1]}
|
||||
assert partial == {}, f"requests reached the space without their whole tree (server, model): {partial}"
|
||||
return kinds
|
||||
|
||||
kept: Final = sum(1 for server, _ in until_asserted(check, seconds=30).values() if server)
|
||||
assert 0 < kept < len(sent), f"a rate of 0.5 kept {kept} of {len(sent)} requests"
|
||||
|
||||
|
||||
def test_a_key_level_destination_follows_the_verdict_of_its_trace(pods: Pods) -> None:
|
||||
team: Final = pods.scenario.team()
|
||||
sampled: Final = pods.scenario.key(
|
||||
team_id=team,
|
||||
metadata={
|
||||
"logging": [
|
||||
{
|
||||
"callback_name": "arize",
|
||||
"callback_type": "success",
|
||||
"callback_vars": {**CREDENTIALS, "arize_success_sampling_rate": "0.5"},
|
||||
}
|
||||
]
|
||||
},
|
||||
)
|
||||
cursors: Final = _cursors(pods)
|
||||
by_trace: Final = {
|
||||
trace_id: tuple(pods.send(pod, _call(sampled, trace_id)) for pod in (pods.first, pods.second, pods.first))
|
||||
for trace_id in _trace_ids("key-level")
|
||||
}
|
||||
|
||||
verdicts: Final = _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS))
|
||||
|
||||
assert 0 < _kept(verdicts) < TRACES
|
||||
|
||||
|
||||
def test_a_concurrent_burst_keeps_every_trace_whole(pods: Pods) -> None:
|
||||
"""Four calls per trace, all in flight at once across both pods."""
|
||||
sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.5"})
|
||||
cursors: Final = _cursors(pods)
|
||||
traces: Final = _trace_ids("burst", 24)
|
||||
plan: Final = tuple(product(traces, (pods.first, pods.second, pods.first, pods.second)))
|
||||
|
||||
def send(job: tuple[str, OwnedProxy]) -> Sent:
|
||||
return pods.send(job[1], _call(sampled, job[0]))
|
||||
|
||||
with ThreadPoolExecutor(max_workers=16) as pool:
|
||||
sent: Final = list(pool.map(send, plan))
|
||||
by_trace: Final = {trace_id: tuple(call for call in sent if call.trace_id == trace_id) for trace_id in traces}
|
||||
|
||||
verdicts: Final = _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS))
|
||||
|
||||
assert 0 < _kept(verdicts) < len(traces)
|
||||
|
||||
|
||||
def test_a_stalled_team_space_during_a_burst_keeps_serving_and_each_trace_whole(pods: Pods) -> None:
|
||||
"""The team's Arize endpoint stops answering while a burst of calls runs over both pods: every call is
|
||||
still served, and once the endpoint answers again each trace arrives whole or not at all."""
|
||||
sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.5"})
|
||||
cursors: Final = _cursors(pods)
|
||||
traces: Final = _trace_ids("stalled")
|
||||
plan: Final = tuple(product(traces, (pods.first, pods.second)))
|
||||
|
||||
def send(job: tuple[str, OwnedProxy]) -> Sent:
|
||||
return pods.send(job[1], _call(sampled, job[0]))
|
||||
|
||||
configure_sink(pods.sinks.arize, paused=True)
|
||||
try:
|
||||
with ThreadPoolExecutor(max_workers=8) as pool:
|
||||
sent: Final = tuple(pool.map(send, plan))
|
||||
finally:
|
||||
configure_sink(pods.sinks.arize, paused=False)
|
||||
by_trace: Final = {trace_id: tuple(call for call in sent if call.trace_id == trace_id) for trace_id in traces}
|
||||
|
||||
verdicts: Final = _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS))
|
||||
|
||||
assert 0 < _kept(verdicts) < len(traces)
|
||||
|
||||
|
||||
def _write_config(directory: Path, sinks: SpanSinks, provider: Wire) -> Path:
|
||||
shared: Final = object_value(
|
||||
JSON.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
|
||||
)
|
||||
config: Final = {
|
||||
**shared,
|
||||
"model_list": [
|
||||
*list_value(shared.get("model_list") or []),
|
||||
{
|
||||
"model_name": MODEL,
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"api_key": "integration-provider-key",
|
||||
"api_base": provider.url + "/v1",
|
||||
},
|
||||
},
|
||||
],
|
||||
"litellm_settings": {
|
||||
**object_value(shared.get("litellm_settings") or {}),
|
||||
"callbacks": ["langfuse_otel"],
|
||||
"otel_tenant_destination_mode": "additive",
|
||||
"provider_url_destination_allowed_hosts": [urlparse(sinks.operator).netloc, urlparse(sinks.arize).netloc],
|
||||
},
|
||||
"router_settings": {**object_value(shared.get("router_settings") or {}), "num_retries": 0},
|
||||
"general_settings": {**object_value(shared.get("general_settings") or {}), "user_api_key_cache_ttl": 2},
|
||||
}
|
||||
path: Final = directory / "team-sampling.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
return path
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _two_pods(provider: Wire, sinks: SpanSinks, directory: Path) -> Generator[Pods]:
|
||||
environment: Final = {
|
||||
"LITELLM_OTEL_V2": "1",
|
||||
"OTEL_BSP_SCHEDULE_DELAY": "300",
|
||||
"MAPPER_NAMES": "genai,openinference",
|
||||
"LANGFUSE_HOST": sinks.operator,
|
||||
"LANGFUSE_PUBLIC_KEY": "pk-lf-operator",
|
||||
"LANGFUSE_SECRET_KEY": "sk-lf-operator",
|
||||
"ARIZE_SPACE_KEY": "space-operator",
|
||||
"ARIZE_API_KEY": "arize-operator",
|
||||
"ARIZE_HTTP_ENDPOINT": sinks.arize + "/v1/traces",
|
||||
}
|
||||
removed: Final = ("LITELLM_OTEL_TENANT_DESTINATION_MODE", "OTEL_TRACES_SAMPLER", "OTEL_TRACES_SAMPLER_ARG")
|
||||
config: Final = _write_config(directory, sinks, provider)
|
||||
for pod in ("first", "second"):
|
||||
(directory / pod).mkdir()
|
||||
with gateway_from_environment() as shared, scratch_database("team_sampling") as database_url:
|
||||
overrides: Final = {**environment, "DATABASE_URL": database_url}
|
||||
with (
|
||||
owned_proxy_process(
|
||||
shared, directory / "first", overrides, config=config, remove_environment=removed
|
||||
) as first,
|
||||
owned_proxy_process(
|
||||
shared,
|
||||
directory / "second",
|
||||
overrides,
|
||||
config=config,
|
||||
remove_environment=removed,
|
||||
workers=2,
|
||||
healthcheck=WorkerHealthcheck.LENIENT,
|
||||
) as second,
|
||||
replace(first.gateway, database_url=database_url).scenario() as scenario,
|
||||
):
|
||||
eventually(second.worker_pids, lambda pids: len(pids) == 2, seconds=120)
|
||||
yield Pods(first, second, scenario, sinks)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def provider() -> Iterator[Wire]:
|
||||
with wire_server(_upstream) as wire:
|
||||
yield wire
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def pods(provider: Wire, audit_sinks: SpanSinks, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Pods]:
|
||||
with _two_pods(provider, audit_sinks, tmp_path_factory.mktemp("team-sampling")) as started:
|
||||
yield started
|
||||
|
|
@ -0,0 +1,153 @@
|
|||
import json
|
||||
import uuid
|
||||
from collections.abc import Iterator, Mapping
|
||||
from contextlib import ExitStack
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.anthropic_sse import message_json
|
||||
from integration._support.client import Gateway
|
||||
from integration._support.openai_wire import chat_reply
|
||||
from integration._support.process import OwnedProxy, owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
GENERIC: Final = "passthrough-scan-generic"
|
||||
AIRS: Final = "passthrough-scan-airs"
|
||||
OPENAI_ROUTE: Final = "/openai/chat/completions"
|
||||
ANTHROPIC_ROUTE: Final = "/anthropic/v1/messages"
|
||||
ROUTE_BODIES: Final[Mapping[str, dict[str, JsonValue]]] = MappingProxyType(
|
||||
{
|
||||
OPENAI_ROUTE: {"model": "gpt-5.6", "messages": [{"role": "user", "content": "Say hi in one word"}]},
|
||||
ANTHROPIC_ROUTE: {
|
||||
"model": "claude-sonnet-5-5",
|
||||
"max_tokens": 16,
|
||||
"messages": [{"role": "user", "content": "Say hi in one word"}],
|
||||
},
|
||||
}
|
||||
)
|
||||
_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml"
|
||||
_JSON: Final = TypeAdapter(dict[str, JsonValue])
|
||||
|
||||
|
||||
def _generic_verdict(_request: Request) -> Reply:
|
||||
return Reply(body=json.dumps({"action": "NONE"}).encode())
|
||||
|
||||
|
||||
def _airs_verdict(_request: Request) -> Reply:
|
||||
return Reply(body=json.dumps({"action": "allow", "category": "benign", "scan_id": "s", "report_id": "r"}).encode())
|
||||
|
||||
|
||||
def _openai_reply(_request: Request) -> Reply:
|
||||
return chat_reply("chatcmpl-passthrough-scan", "gpt-5.6", "Hi", stream=False)
|
||||
|
||||
|
||||
def _anthropic_reply(_request: Request) -> Reply:
|
||||
return Reply(body=message_json("msg_passthrough_scan", "claude-sonnet-5-5", "Hi"))
|
||||
|
||||
|
||||
def _parsed_texts(texts: JsonValue, vendor_body: dict[str, JsonValue]) -> list[JsonValue]:
|
||||
assert isinstance(texts, list), vendor_body
|
||||
return [json.loads(text) if isinstance(text, str) else text for text in texts]
|
||||
|
||||
|
||||
def _scanned_payloads(guardrail: str, vendor_body: dict[str, JsonValue]) -> list[JsonValue]:
|
||||
if guardrail == GENERIC:
|
||||
return _parsed_texts(vendor_body["texts"], vendor_body)
|
||||
contents: Final = vendor_body["contents"]
|
||||
assert isinstance(contents, list), vendor_body
|
||||
return _parsed_texts([content.get("prompt") for content in contents if isinstance(content, dict)], vendor_body)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PassThroughRig:
|
||||
owned: OwnedProxy
|
||||
vendors: Mapping[str, Wire]
|
||||
upstreams: Mapping[str, Wire]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[PassThroughRig]:
|
||||
directory: Final = tmp_path_factory.mktemp("passthrough-guardrail-scan")
|
||||
with ExitStack() as stack:
|
||||
vendors: Final = MappingProxyType(
|
||||
{
|
||||
GENERIC: stack.enter_context(wire_server(_generic_verdict)),
|
||||
AIRS: stack.enter_context(wire_server(_airs_verdict)),
|
||||
}
|
||||
)
|
||||
upstreams: Final = MappingProxyType(
|
||||
{
|
||||
OPENAI_ROUTE: stack.enter_context(wire_server(_openai_reply)),
|
||||
ANTHROPIC_ROUTE: stack.enter_context(wire_server(_anthropic_reply)),
|
||||
}
|
||||
)
|
||||
config: Final = _JSON.validate_python(yaml.safe_load(_CONFIG.read_text()))
|
||||
path: Final = directory / "passthrough-guardrail-scan.yaml"
|
||||
path.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
**config,
|
||||
"guardrails": [
|
||||
{
|
||||
"guardrail_name": GENERIC,
|
||||
"litellm_params": {
|
||||
"guardrail": "generic_guardrail_api",
|
||||
"mode": "pre_call",
|
||||
"default_on": False,
|
||||
"api_base": vendors[GENERIC].url,
|
||||
"api_key": "synthetic-guardrail-key",
|
||||
},
|
||||
},
|
||||
{
|
||||
"guardrail_name": AIRS,
|
||||
"litellm_params": {
|
||||
"guardrail": "panw_prisma_airs",
|
||||
"mode": "pre_call",
|
||||
"default_on": False,
|
||||
"api_base": vendors[AIRS].url,
|
||||
"api_key": "synthetic-airs-key",
|
||||
"profile_name": "integration",
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
)
|
||||
)
|
||||
environment: Final = {
|
||||
"OPENAI_API_BASE": upstreams[OPENAI_ROUTE].url,
|
||||
"OPENAI_API_KEY": "synthetic-openai-key",
|
||||
"ANTHROPIC_API_BASE": upstreams[ANTHROPIC_ROUTE].url,
|
||||
"ANTHROPIC_API_KEY": "synthetic-anthropic-key",
|
||||
}
|
||||
owned: Final = stack.enter_context(owned_proxy_process(session_gateway, directory, environment, config=path))
|
||||
yield PassThroughRig(owned, vendors, upstreams)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("guardrail", [GENERIC, AIRS])
|
||||
@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE])
|
||||
def test_pass_through_guardrail_scans_the_client_body_without_the_callers_key_metadata(
|
||||
rig: PassThroughRig, route: str, guardrail: str
|
||||
) -> None:
|
||||
gateway: Final = rig.owned.gateway
|
||||
with gateway.scenario() as scenario:
|
||||
team: Final = scenario.team()
|
||||
user: Final = scenario.user(
|
||||
user_email=f"passthrough-scan-{uuid.uuid4().hex}@example.com", user_role="internal_user", teams=[team]
|
||||
)
|
||||
key: Final = scenario.key(user_id=user, team_id=team)
|
||||
body: Final = {**ROUTE_BODIES[route], "guardrails": [guardrail]}
|
||||
|
||||
response: Final = gateway.request("POST", route, body, key=key)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
scanned: Final = [
|
||||
_scanned_payloads(guardrail, _JSON.validate_json(request.body))
|
||||
for request in rig.vendors[guardrail].drain()
|
||||
]
|
||||
assert scanned == [[body]], scanned
|
||||
assert len(rig.upstreams[route].drain()) == 1
|
||||
|
|
@ -212,3 +212,36 @@ def test_sync_stream_that_raises_mid_stream_releases_the_connection() -> None:
|
|||
assert _open_sessions_settling_to(wire, at_most=0) == 0, (
|
||||
"a stream that raised mid-way left its TCP session open"
|
||||
)
|
||||
|
||||
|
||||
_STUCK: Final = _Shape(
|
||||
"hosted_vllm",
|
||||
"hosted_vllm/stuck-model",
|
||||
_OPENAI_OPENING,
|
||||
(_openai_delta("the model is stuck on this", None),) * (10 * litellm.REPEATED_STREAMING_CHUNK_LIMIT),
|
||||
)
|
||||
|
||||
|
||||
def _still_generating_peer(request: Request) -> Reply:
|
||||
return Reply(content_type="text/event-stream", chunks=(_STUCK.opening, *_STUCK.closing), pause_between_chunks=0.01)
|
||||
|
||||
|
||||
def test_sync_stream_stuck_repeating_one_chunk_releases_the_connection() -> None:
|
||||
with wire_server(_still_generating_peer, keep_alive=True) as wire:
|
||||
stuck: Final = litellm.completion(**_request_kwargs(_STUCK, wire))
|
||||
with pytest.raises(Exception, match="repeating the same chunk"):
|
||||
for _ in stuck:
|
||||
pass
|
||||
assert _open_sessions_settling_to(wire, at_most=0) == 0, "a stream stuck on one chunk left its TCP session open"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_stream_stuck_repeating_one_chunk_releases_the_connection() -> None:
|
||||
with wire_server(_still_generating_peer, keep_alive=True) as wire:
|
||||
stuck: Final = await litellm.acompletion(**_request_kwargs(_STUCK, wire))
|
||||
with pytest.raises(Exception, match="repeating the same chunk"):
|
||||
async for _ in stuck:
|
||||
pass
|
||||
assert await _open_sessions_settling_to_while_the_loop_runs(wire, at_most=0) == 0, (
|
||||
"a stream stuck on one chunk left its TCP session open"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,273 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import uuid
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from tests.integration._support.client import Gateway, Scenario, eventually, list_value, object_value, string_value
|
||||
from tests.integration._support.database import read_rows, write_rows
|
||||
from tests.integration._support.provider import PROVIDER_URL, SharedProvider
|
||||
from tests.integration._support.wire import Reply
|
||||
|
||||
_CALL_COST: Final = 0.02
|
||||
JSON_VALUE: Final = TypeAdapter(JsonValue)
|
||||
_LEGACY_SPEND: Final = 612.4
|
||||
_LEGACY_TOTAL: Final = 148.1
|
||||
_HEALTHY_SPEND: Final = 12.25
|
||||
_HEALTHY_TOTAL: Final = 97.8
|
||||
|
||||
|
||||
def _completion() -> Reply:
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": f"chatcmpl-{uuid.uuid4().hex}",
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "gpt-4o-mini",
|
||||
"choices": [
|
||||
{"index": 0, "message": {"role": "assistant", "content": "scripted"}, "finish_reason": "stop"}
|
||||
],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
def _priced_model(scenario: Scenario) -> str:
|
||||
return scenario.model(api_base=f"{PROVIDER_URL}/v1", input_cost_per_token=0.001, output_cost_per_token=0.002)
|
||||
|
||||
|
||||
def _member_with_budget(gateway: Gateway, scenario: Scenario, team_id: str, *, max_budget: float) -> str:
|
||||
user_id: Final = scenario.user(user_role="internal_user")
|
||||
gateway.post(
|
||||
"/team/member_add",
|
||||
{"team_id": team_id, "member": {"role": "user", "user_id": user_id}, "max_budget_in_team": max_budget},
|
||||
)
|
||||
return user_id
|
||||
|
||||
|
||||
def _seed_membership(user_id: str, team_id: str, *, spend: float, total_spend: float) -> None:
|
||||
write_rows(
|
||||
'UPDATE "LiteLLM_TeamMembership" SET spend = %s::float8, total_spend = %s::float8 '
|
||||
"WHERE user_id = %s AND team_id = %s",
|
||||
(str(spend), str(total_spend), user_id, team_id),
|
||||
)
|
||||
|
||||
|
||||
def _membership(user_id: str, team_id: str) -> tuple[float, float]:
|
||||
rows: Final = read_rows(
|
||||
'SELECT spend, total_spend FROM "LiteLLM_TeamMembership" WHERE user_id = %s AND team_id = %s',
|
||||
(user_id, team_id),
|
||||
)
|
||||
assert len(rows) == 1, rows
|
||||
return float(str(rows[0]["spend"])), float(str(rows[0]["total_spend"]))
|
||||
|
||||
|
||||
def _membership_reaches(user_id: str, team_id: str, *, total_spend: float) -> tuple[float, float]:
|
||||
return eventually(
|
||||
lambda: _membership(user_id, team_id),
|
||||
lambda row: row[1] == pytest.approx(total_spend),
|
||||
seconds=70,
|
||||
)
|
||||
|
||||
|
||||
def _served(gateway: Gateway, provider: SharedProvider, model: str, key: str) -> None:
|
||||
provider.expect(_completion())
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "messages": [{"role": "user", "content": f"reconcile {uuid.uuid4().hex}"}]},
|
||||
key=key,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert len(provider.received()) == 1
|
||||
|
||||
|
||||
def _reset_spend(gateway: Gateway, team_id: str, user_id: str, reset_to: float) -> httpx.Response:
|
||||
return gateway.request("POST", f"/team/{team_id}/member/{user_id}/reset_spend", {"reset_to": reset_to})
|
||||
|
||||
|
||||
def _membership_figures(row: JsonValue) -> tuple[str, tuple[float, float]]:
|
||||
membership: Final = object_value(row)
|
||||
return string_value(membership["user_id"]), (
|
||||
float(str(membership["spend"])),
|
||||
float(str(membership["total_spend"])),
|
||||
)
|
||||
|
||||
|
||||
def _reported_memberships(gateway: Gateway, path: str, params: dict[str, str]) -> dict[str, tuple[float, float]]:
|
||||
"""/team/info answers with one team object, /team/list with a list of them."""
|
||||
response: Final = gateway.request("GET", path, params=params)
|
||||
assert response.status_code == 200, f"GET {path}: {response.status_code} {response.text}"
|
||||
body: Final = JSON_VALUE.validate_json(response.content)
|
||||
teams: Final = [object_value(body)] if isinstance(body, dict) else [object_value(team) for team in list_value(body)]
|
||||
rows: Final = [
|
||||
row for team in teams for row in list_value(team["team_memberships"])
|
||||
] # comprehension-ok: flatten memberships across teams
|
||||
return dict(_membership_figures(row) for row in rows)
|
||||
|
||||
|
||||
def test_a_flush_lifts_a_legacy_total_to_the_cycle_spend_and_leaves_a_healthy_row_on_its_baseline(
|
||||
gateway: Gateway, provider: SharedProvider
|
||||
) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = _priced_model(scenario)
|
||||
team_id: Final = scenario.team(models=[model])
|
||||
legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000)
|
||||
healthy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000)
|
||||
_seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL)
|
||||
_seed_membership(healthy, team_id, spend=_HEALTHY_SPEND, total_spend=_HEALTHY_TOTAL)
|
||||
legacy_key: Final = scenario.key(user_id=legacy, team_id=team_id, models=[model])
|
||||
healthy_key: Final = scenario.key(user_id=healthy, team_id=team_id, models=[model])
|
||||
|
||||
_served(gateway, provider, model, legacy_key)
|
||||
_served(gateway, provider, model, healthy_key)
|
||||
|
||||
legacy_row: Final = _membership_reaches(
|
||||
legacy, team_id, total_spend=max(_LEGACY_SPEND, _LEGACY_TOTAL) + _CALL_COST
|
||||
)
|
||||
assert legacy_row[0] == pytest.approx(_LEGACY_SPEND + _CALL_COST)
|
||||
healthy_row: Final = _membership_reaches(healthy, team_id, total_spend=_HEALTHY_TOTAL + _CALL_COST)
|
||||
assert healthy_row[0] == pytest.approx(_HEALTHY_SPEND + _CALL_COST)
|
||||
|
||||
|
||||
def test_team_info_and_team_list_report_a_legacy_total_no_lower_than_the_cycle_spend(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
team_id: Final = scenario.team()
|
||||
legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000)
|
||||
healthy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000)
|
||||
_seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL)
|
||||
_seed_membership(healthy, team_id, spend=_HEALTHY_SPEND, total_spend=_HEALTHY_TOTAL)
|
||||
|
||||
info: Final = _reported_memberships(gateway, "/team/info", {"team_id": team_id})
|
||||
listed: Final = _reported_memberships(gateway, "/team/list", {"team_id": team_id})
|
||||
|
||||
for reported in (info, listed):
|
||||
assert reported[legacy] == (pytest.approx(_LEGACY_SPEND), pytest.approx(_LEGACY_SPEND))
|
||||
assert reported[healthy] == (pytest.approx(_HEALTHY_SPEND), pytest.approx(_HEALTHY_TOTAL))
|
||||
assert _membership(legacy, team_id) == (pytest.approx(_LEGACY_SPEND), pytest.approx(_LEGACY_TOTAL))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("reset_to", [0.0, 100.0], ids=["to_zero", "to_nonzero"])
|
||||
def test_an_admin_reset_keeps_the_legacy_cycle_spend_as_the_lifetime_floor(gateway: Gateway, reset_to: float) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
team_id: Final = scenario.team()
|
||||
legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000)
|
||||
_seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL)
|
||||
|
||||
response: Final = _reset_spend(gateway, team_id, legacy, reset_to)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert _membership(legacy, team_id) == (pytest.approx(reset_to), pytest.approx(_LEGACY_SPEND))
|
||||
|
||||
|
||||
def test_concurrent_admin_resets_leave_one_winner_and_a_lifetime_floor_at_the_old_cycle_spend(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
"""Two admins reset the same member at once. Each UPDATE re-checks reset_to <= spend, so
|
||||
the higher reset either lands first or is refused with 409 once the lower one has landed.
|
||||
Whatever the order, the row ends at the lower value with the old cycle spend preserved."""
|
||||
with gateway.scenario() as scenario:
|
||||
team_id: Final = scenario.team()
|
||||
legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000)
|
||||
_seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=2) as pool:
|
||||
higher, lower = pool.map(lambda value: _reset_spend(gateway, team_id, legacy, value), (600.0, 500.0))
|
||||
|
||||
assert lower.status_code == 200, lower.text
|
||||
assert higher.status_code in (200, 409), higher.text
|
||||
assert _membership(legacy, team_id) == (pytest.approx(500.0), pytest.approx(_LEGACY_SPEND))
|
||||
|
||||
|
||||
def test_an_admin_reset_racing_a_flush_ends_with_the_lifetime_equal_to_old_cycle_plus_cost(
|
||||
gateway: Gateway, provider: SharedProvider
|
||||
) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = _priced_model(scenario)
|
||||
team_id: Final = scenario.team(models=[model])
|
||||
legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000)
|
||||
_seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL)
|
||||
key: Final = scenario.key(user_id=legacy, team_id=team_id, models=[model])
|
||||
|
||||
_served(gateway, provider, model, key)
|
||||
reset: Final = _reset_spend(gateway, team_id, legacy, 0.0)
|
||||
assert reset.status_code == 200, reset.text
|
||||
|
||||
spend, total = _membership_reaches(legacy, team_id, total_spend=_LEGACY_SPEND + _CALL_COST)
|
||||
assert spend in (pytest.approx(0.0), pytest.approx(_CALL_COST)), (spend, total)
|
||||
assert total >= spend
|
||||
|
||||
|
||||
def _expire_member_budget(user_id: str, team_id: str, *, budget_duration: str) -> None:
|
||||
write_rows(
|
||||
"UPDATE \"LiteLLM_BudgetTable\" SET budget_duration = %s, budget_reset_at = NOW() - INTERVAL '1 hour' "
|
||||
'WHERE budget_id = (SELECT budget_id FROM "LiteLLM_TeamMembership" WHERE user_id = %s AND team_id = %s)',
|
||||
(budget_duration, user_id, team_id),
|
||||
)
|
||||
|
||||
|
||||
async def _run_scheduled_reset() -> None:
|
||||
client: Final = PrismaClient(
|
||||
database_url=os.environ["DATABASE_URL"], proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache())
|
||||
)
|
||||
await client.connect()
|
||||
try:
|
||||
await ResetBudgetJob(
|
||||
ProxyLogging(user_api_key_cache=DualCache()), client
|
||||
).reset_budget_for_litellm_budget_table()
|
||||
finally:
|
||||
await client.db.disconnect()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"spend, max_budget, carried",
|
||||
[(8.0, 10.0, 0.0), (10.0, 10.0, 0.0), (15.0, 10.0, 5.0)],
|
||||
ids=["below_cap", "at_cap", "above_cap_carries_overage"],
|
||||
)
|
||||
def test_a_scheduled_reset_lifts_the_legacy_total_before_zeroing_or_carrying_the_cycle_spend(
|
||||
gateway: Gateway, monkeypatch: pytest.MonkeyPatch, spend: float, max_budget: float, carried: float
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "budget_rollover", True)
|
||||
with gateway.scenario() as scenario:
|
||||
team_id: Final = scenario.team()
|
||||
legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=max_budget)
|
||||
_seed_membership(legacy, team_id, spend=spend, total_spend=3.0)
|
||||
_expire_member_budget(legacy, team_id, budget_duration="1d")
|
||||
|
||||
asyncio.run(_run_scheduled_reset())
|
||||
|
||||
assert _membership(legacy, team_id) == (pytest.approx(carried), pytest.approx(spend))
|
||||
asyncio.run(_run_scheduled_reset())
|
||||
assert _membership(legacy, team_id) == (pytest.approx(carried), pytest.approx(spend))
|
||||
|
||||
|
||||
def test_a_scheduled_reset_racing_a_flush_ends_with_the_lifetime_equal_to_old_cycle_plus_cost(
|
||||
gateway: Gateway, provider: SharedProvider
|
||||
) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = _priced_model(scenario)
|
||||
team_id: Final = scenario.team(models=[model])
|
||||
legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000)
|
||||
_seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL)
|
||||
_expire_member_budget(legacy, team_id, budget_duration="1d")
|
||||
key: Final = scenario.key(user_id=legacy, team_id=team_id, models=[model])
|
||||
|
||||
_served(gateway, provider, model, key)
|
||||
asyncio.run(_run_scheduled_reset())
|
||||
|
||||
spend, total = _membership_reaches(legacy, team_id, total_spend=_LEGACY_SPEND + _CALL_COST)
|
||||
assert spend in (pytest.approx(0.0), pytest.approx(_CALL_COST)), (spend, total)
|
||||
assert total >= spend
|
||||
|
|
@ -1921,6 +1921,62 @@ async def test_list_tools_classifies_upstream_status_for_post_failures(failure_m
|
|||
assert rate_limit is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("paginated", (False, True))
|
||||
async def test_sse_catalog_limit_reports_the_failure_and_allows_smaller_pages(paginated: bool) -> None:
|
||||
from mcp.types import Tool
|
||||
|
||||
limit: Final = 1024
|
||||
tools: Final = tuple(
|
||||
Tool(name=f"tool_{index}", description="private-upstream-secret " * 25, input_schema={"type": "object"})
|
||||
for index in range(2)
|
||||
)
|
||||
catalog: Final = tuple(tool.model_dump(by_alias=True, exclude_none=True) for tool in tools)
|
||||
assert len(json.dumps({"tools": catalog}).encode()) > limit
|
||||
|
||||
def respond(request: httpx2.Request) -> httpx2.Response:
|
||||
if request.method != "POST":
|
||||
return httpx2.Response(200)
|
||||
payload: Final = _JSONRPC_MESSAGE_ADAPTER.validate_json(request.content)
|
||||
assert isinstance(payload, (JSONRPCRequest, JSONRPCNotification))
|
||||
if isinstance(payload, JSONRPCNotification):
|
||||
return httpx2.Response(202)
|
||||
if payload.method == "initialize":
|
||||
return httpx2.Response(200, json={
|
||||
"jsonrpc": "2.0", "id": payload.id, "result": {
|
||||
"protocolVersion": LATEST_HANDSHAKE_VERSION,
|
||||
"capabilities": {"tools": {}}, "serverInfo": {"name": "catalog", "version": "1"},
|
||||
},
|
||||
})
|
||||
assert payload.method == "tools/list"
|
||||
cursor: Final = (payload.params or {}).get("cursor")
|
||||
page: Final = catalog[1:] if cursor == "next-page" else catalog[:1] if paginated else catalog
|
||||
result: Final = {
|
||||
"tools": page,
|
||||
**({"nextCursor": "next-page"} if paginated and cursor is None else {}),
|
||||
}
|
||||
event: Final = f"data: {json.dumps({'jsonrpc': '2.0', 'id': payload.id, 'result': result})}\n\n"
|
||||
assert (len(event.encode()) < limit) == paginated
|
||||
return httpx2.Response(200, headers={"content-type": "text/event-stream"}, content=event)
|
||||
|
||||
def transport_context(url: str, http_client: httpx2.AsyncClient) -> _TransportContext:
|
||||
return streamable_http_client(url, http_client=http_client, max_sse_event_size=limit)
|
||||
|
||||
client: Final = _MockTransportClient(
|
||||
respond, server_url="https://upstream.example/mcp", transport_context=transport_context,
|
||||
)
|
||||
if paginated:
|
||||
assert await client.list_tools(raise_on_error=True) == list(tools)
|
||||
return
|
||||
with pytest.raises(MCPError) as caught:
|
||||
await client.list_tools(raise_on_error=True)
|
||||
fault: Final = classify_list_exception(caught.value)
|
||||
assert fault == ServerListFault(tag="upstream_error", sse_event_limit_bytes=limit)
|
||||
assert list_fault_http_status(fault) == 502
|
||||
assert fault.public_message is not None
|
||||
assert "private-upstream-secret" not in fault.public_message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_status_check_allows_auth_refresh_before_rejecting() -> None:
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import (
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import asyncio
|
|||
import contextvars
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from base64 import b64encode
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import replace
|
||||
|
|
@ -14,10 +15,12 @@ from typing import Final
|
|||
|
||||
import pytest
|
||||
from opentelemetry import trace as trace_api
|
||||
from opentelemetry.context import Context
|
||||
from opentelemetry.sdk.resources import Resource
|
||||
from opentelemetry.sdk.trace import ReadableSpan, TracerProvider
|
||||
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
|
||||
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
|
||||
from opentelemetry.sdk.trace.id_generator import RandomIdGenerator
|
||||
from opentelemetry.trace import Status, StatusCode
|
||||
|
||||
import litellm
|
||||
|
|
@ -1159,12 +1162,36 @@ class TestFanOut:
|
|||
ARIZE_TEAM_PARAMS = {"arize_space_id": "space-team", "arize_api_key": "key-team"}
|
||||
|
||||
|
||||
def caller_context(trace_id: int) -> Context:
|
||||
return trace_api.set_span_in_context(
|
||||
trace_api.NonRecordingSpan(
|
||||
trace_api.SpanContext(trace_id=trace_id, span_id=1, is_remote=True, trace_flags=trace_api.TraceFlags(0x01))
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def arize_destination(**rates: str) -> OtelDestination:
|
||||
destination = destination_for("arize", {**ARIZE_TEAM_PARAMS, **rates})
|
||||
assert destination is not None, "the fixture must resolve for the test to mean anything"
|
||||
return destination
|
||||
|
||||
|
||||
#: Trace ids whose SHA-256 reads 0.40 and 0.80 of the way through its first eight bytes,
|
||||
#: so a rate of 0.5 keeps the first and drops the second.
|
||||
KEPT_AT_HALF: Final = 4
|
||||
DROPPED_AT_HALF: Final = 5
|
||||
|
||||
|
||||
class _FixedTraceId(RandomIdGenerator):
|
||||
"""Every trace the provider starts gets ``trace_id``."""
|
||||
|
||||
def __init__(self, trace_id: int) -> None:
|
||||
self._trace_id = trace_id
|
||||
|
||||
def generate_trace_id(self) -> int:
|
||||
return self._trace_id
|
||||
|
||||
|
||||
class TestTenantSampling:
|
||||
"""A team's ``arize_success_sampling_rate`` and ``arize_error_sampling_rate`` decide how
|
||||
much of its traffic the fan-out delivers to its space, the way they already decide what
|
||||
|
|
@ -1173,14 +1200,17 @@ class TestTenantSampling:
|
|||
NAMES = {"POST /v1/chat/completions", "auth /v1/chat/completions", "chat gpt-4"}
|
||||
|
||||
@staticmethod
|
||||
def _fan_out(dest_exporter, draw=None, global_exporter=None):
|
||||
provider = TracerProvider()
|
||||
def _fan_out(
|
||||
dest_exporter: InMemorySpanExporter,
|
||||
trace_id: int | None = None,
|
||||
global_exporter: InMemorySpanExporter | None = None,
|
||||
) -> TracerProvider:
|
||||
provider: Final = TracerProvider() if trace_id is None else TracerProvider(id_generator=_FixedTraceId(trace_id))
|
||||
if global_exporter is not None:
|
||||
provider.add_span_processor(_OverriddenBackendFilter(SimpleSpanProcessor(global_exporter), "arize"))
|
||||
kwargs = {"processor_factory": lambda _d: SimpleSpanProcessor(dest_exporter)}
|
||||
if draw is not None:
|
||||
kwargs["sampling_draw"] = draw
|
||||
provider.add_span_processor(TenantFanOutSpanProcessor(**kwargs))
|
||||
provider.add_span_processor(
|
||||
TenantFanOutSpanProcessor(processor_factory=lambda _d: SimpleSpanProcessor(dest_exporter))
|
||||
)
|
||||
return provider
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -1224,7 +1254,7 @@ class TestTenantSampling:
|
|||
|
||||
def test_an_unset_rate_exports_everything(self):
|
||||
dest_exporter = InMemorySpanExporter()
|
||||
provider = self._fan_out(dest_exporter, draw=lambda: 0.99)
|
||||
provider: Final = self._fan_out(dest_exporter, trace_id=DROPPED_AT_HALF)
|
||||
|
||||
self._run(provider, (arize_destination(),))
|
||||
|
||||
|
|
@ -1300,10 +1330,10 @@ class TestTenantSampling:
|
|||
|
||||
def test_a_span_that_ends_after_the_root_follows_the_requests_verdict(self):
|
||||
"""A post-call database write that starts after the whole request tree has ended
|
||||
goes where the rest of the tree went, without a second draw."""
|
||||
goes where the rest of the tree went, without a second verdict."""
|
||||
kept_exporter, dropped_exporter = InMemorySpanExporter(), InMemorySpanExporter()
|
||||
kept = self._fan_out(kept_exporter, draw=lambda: 0.1)
|
||||
dropped = self._fan_out(dropped_exporter, draw=lambda: 0.9)
|
||||
kept: Final = self._fan_out(kept_exporter, trace_id=KEPT_AT_HALF)
|
||||
dropped: Final = self._fan_out(dropped_exporter, trace_id=DROPPED_AT_HALF)
|
||||
destinations = (arize_destination(arize_success_sampling_rate="0.5"),)
|
||||
|
||||
def late_tree(provider):
|
||||
|
|
@ -1327,14 +1357,263 @@ class TestTenantSampling:
|
|||
}
|
||||
assert dropped_exporter.get_finished_spans() == ()
|
||||
|
||||
def test_a_span_that_ends_after_a_failed_root_follows_the_failed_requests_verdict(self):
|
||||
"""A post-call database write after a failed request goes where that request went,
|
||||
though on its own it would answer to the success rate."""
|
||||
dest_exporter: Final = InMemorySpanExporter()
|
||||
provider: Final = self._fan_out(dest_exporter)
|
||||
destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),)
|
||||
|
||||
def run():
|
||||
set_request_destinations(destinations)
|
||||
tracer: Final = get_tracer(provider, "litellm")
|
||||
with tracer.start_as_current_span("POST /v1/chat/completions") as root:
|
||||
root.set_status(Status(StatusCode.ERROR, "upstream 500"))
|
||||
root_context: Final = trace_api.set_span_in_context(root)
|
||||
with tracer.start_as_current_span("postgres INSERT LiteLLM_SpendLogs", context=root_context):
|
||||
pass
|
||||
|
||||
in_fresh_context(run)
|
||||
|
||||
assert {s.name for s in dest_exporter.get_finished_spans()} == {
|
||||
"POST /v1/chat/completions",
|
||||
"postgres INSERT LiteLLM_SpendLogs",
|
||||
}
|
||||
|
||||
def test_a_failed_call_after_a_dropped_call_of_its_trace_still_answers_to_the_error_rate(self):
|
||||
"""A process that dropped a successful call of the trace keeps a later failed call of
|
||||
it at an unset error rate, as a process that only served the failed call does."""
|
||||
served_both, served_failed = InMemorySpanExporter(), InMemorySpanExporter()
|
||||
destinations: Final = (arize_destination(arize_success_sampling_rate="0.5"),)
|
||||
both: Final = self._fan_out(served_both, trace_id=DROPPED_AT_HALF)
|
||||
|
||||
self._run(both, destinations)
|
||||
self._run(both, destinations, failed=True)
|
||||
self._run(self._fan_out(served_failed, trace_id=DROPPED_AT_HALF), destinations, failed=True)
|
||||
|
||||
assert [sorted(s.name for s in exporter.get_finished_spans()) for exporter in (served_both, served_failed)] == [
|
||||
sorted(self.NAMES)
|
||||
] * 2
|
||||
|
||||
def test_a_successful_call_after_a_kept_failed_call_of_its_trace_answers_to_the_success_rate(self):
|
||||
"""A process that kept a failed call of the trace drops a later successful call of it
|
||||
at a zero success rate, as a process that only served the successful call does."""
|
||||
served_both, served_success = InMemorySpanExporter(), InMemorySpanExporter()
|
||||
destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),)
|
||||
both: Final = self._fan_out(served_both, trace_id=KEPT_AT_HALF)
|
||||
|
||||
self._run(both, destinations, failed=True)
|
||||
self._run(both, destinations)
|
||||
self._run(self._fan_out(served_success, trace_id=KEPT_AT_HALF), destinations)
|
||||
|
||||
assert len(served_both.get_finished_spans()) == len(self.NAMES), "only the failed call"
|
||||
assert all(s.name in self.NAMES for s in served_both.get_finished_spans())
|
||||
assert served_success.get_finished_spans() == ()
|
||||
|
||||
def test_a_late_child_that_fails_after_its_request_was_dropped_is_dropped_with_it(self):
|
||||
"""A post-call write that fails after its successful request was dropped goes where
|
||||
the request went, rather than reaching the space on its own at the error rate."""
|
||||
dest_exporter: Final = InMemorySpanExporter()
|
||||
provider: Final = self._fan_out(dest_exporter)
|
||||
destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),)
|
||||
|
||||
def run():
|
||||
set_request_destinations(destinations)
|
||||
tracer: Final = get_tracer(provider, "litellm")
|
||||
with tracer.start_as_current_span("POST /v1/chat/completions") as root:
|
||||
root_context: Final = trace_api.set_span_in_context(root)
|
||||
with tracer.start_as_current_span("postgres INSERT LiteLLM_SpendLogs", context=root_context) as write:
|
||||
write.set_status(Status(StatusCode.ERROR, "connection reset"))
|
||||
|
||||
in_fresh_context(run)
|
||||
|
||||
assert dest_exporter.get_finished_spans() == ()
|
||||
|
||||
def test_a_failure_deep_in_a_request_marks_the_whole_request(self):
|
||||
"""A database span that fails under the auth span puts the whole request, root
|
||||
included, on the error rate."""
|
||||
dest_exporter: Final = InMemorySpanExporter()
|
||||
provider: Final = self._fan_out(dest_exporter)
|
||||
destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),)
|
||||
|
||||
def run():
|
||||
set_request_destinations(destinations)
|
||||
tracer: Final = get_tracer(provider, "litellm")
|
||||
with tracer.start_as_current_span("POST /v1/chat/completions"):
|
||||
with tracer.start_as_current_span("auth /v1/chat/completions"):
|
||||
with tracer.start_as_current_span("postgres SELECT LiteLLM_VerificationToken") as lookup:
|
||||
lookup.set_status(Status(StatusCode.ERROR, "connection reset"))
|
||||
with tracer.start_as_current_span("chat gpt-4"):
|
||||
pass
|
||||
|
||||
in_fresh_context(run)
|
||||
|
||||
assert {s.name for s in dest_exporter.get_finished_spans()} == {
|
||||
*self.NAMES,
|
||||
"postgres SELECT LiteLLM_VerificationToken",
|
||||
}
|
||||
|
||||
def test_a_span_two_levels_down_that_outlives_many_other_spans_still_closes_its_request(self):
|
||||
"""A classifier call under the route span that stays open while the worker serves
|
||||
more requests than it tracks at once still counts against its own request, so the
|
||||
request is decided once it ends, on its own outcome."""
|
||||
dest_exporter: Final = InMemorySpanExporter()
|
||||
provider: Final = self._fan_out(dest_exporter)
|
||||
destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),)
|
||||
|
||||
def other_request():
|
||||
set_request_destinations(())
|
||||
with get_tracer(provider, "litellm").start_as_current_span("POST /v1/chat/completions"):
|
||||
pass
|
||||
|
||||
def run():
|
||||
set_request_destinations(destinations)
|
||||
tracer: Final = get_tracer(provider, "litellm")
|
||||
with tracer.start_as_current_span("POST /v1/chat/completions") as root:
|
||||
with tracer.start_as_current_span("route gpt-4"):
|
||||
with tracer.start_as_current_span("chat classifier"):
|
||||
for _ in range(otel_providers._MAX_OPEN_REQUESTS + 1):
|
||||
contextvars.Context().run(other_request)
|
||||
with tracer.start_as_current_span("chat gpt-4"):
|
||||
pass
|
||||
root.set_status(Status(StatusCode.ERROR, "upstream 500"))
|
||||
|
||||
in_fresh_context(run)
|
||||
|
||||
assert {s.name for s in dest_exporter.get_finished_spans()} == {
|
||||
"POST /v1/chat/completions",
|
||||
"route gpt-4",
|
||||
"chat classifier",
|
||||
"chat gpt-4",
|
||||
}
|
||||
|
||||
def test_a_failure_two_levels_down_reaches_its_request_however_many_spans_are_open(self):
|
||||
"""A classifier call that fails while another request holds more spans open than
|
||||
the worker tracks requests marks its own request, which is then kept at the error
|
||||
rate as a whole."""
|
||||
dest_exporter: Final = InMemorySpanExporter()
|
||||
provider: Final = self._fan_out(dest_exporter)
|
||||
destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),)
|
||||
others: Final = contextvars.Context()
|
||||
|
||||
def open_other_request():
|
||||
set_request_destinations(())
|
||||
tracer: Final = get_tracer(provider, "litellm")
|
||||
root: Final = tracer.start_span("POST /v1/chat/completions")
|
||||
parent: Final = trace_api.set_span_in_context(root)
|
||||
return (
|
||||
root,
|
||||
*(tracer.start_span("chat flood", context=parent) for _ in range(otel_providers._MAX_OPEN_REQUESTS)),
|
||||
)
|
||||
|
||||
def run():
|
||||
set_request_destinations(destinations)
|
||||
tracer: Final = get_tracer(provider, "litellm")
|
||||
with tracer.start_as_current_span("POST /v1/chat/completions"):
|
||||
with tracer.start_as_current_span("route gpt-4"):
|
||||
with tracer.start_as_current_span("chat classifier") as classifier:
|
||||
held_open: Final = others.run(open_other_request)
|
||||
classifier.set_status(Status(StatusCode.ERROR, "classifier timed out"))
|
||||
return held_open
|
||||
|
||||
held_open: Final = in_fresh_context(run)
|
||||
|
||||
assert {s.name for s in dest_exporter.get_finished_spans()} == {
|
||||
"POST /v1/chat/completions",
|
||||
"route gpt-4",
|
||||
"chat classifier",
|
||||
}
|
||||
for span in held_open:
|
||||
span.end()
|
||||
|
||||
def test_the_fan_out_forgets_a_span_once_it_ends_and_a_request_once_it_is_pushed_out(self):
|
||||
"""The request lookup holds only spans still open, and only those of the requests
|
||||
the fan-out still counts, so neither ended spans nor a flood of open requests grow
|
||||
it past the open request bound."""
|
||||
fan_out: Final = TenantFanOutSpanProcessor(
|
||||
processor_factory=lambda _d: SimpleSpanProcessor(InMemorySpanExporter())
|
||||
)
|
||||
provider: Final = TracerProvider()
|
||||
provider.add_span_processor(fan_out)
|
||||
|
||||
def run():
|
||||
set_request_destinations((arize_destination(arize_success_sampling_rate="1.0"),))
|
||||
tracer: Final = get_tracer(provider, "litellm")
|
||||
with tracer.start_as_current_span("POST /v1/chat/completions"):
|
||||
with tracer.start_as_current_span("chat gpt-4"):
|
||||
pass
|
||||
assert (fan_out._roots, dict(fan_out._open)) == ({}, {})
|
||||
root: Final = tracer.start_span("POST /v1/chat/completions", context=trace_api.Context())
|
||||
tracer.start_span("route gpt-4", context=trace_api.set_span_in_context(root))
|
||||
return tuple(
|
||||
tracer.start_span(f"POST {index}", context=trace_api.Context())
|
||||
for index in range(otel_providers._MAX_OPEN_REQUESTS)
|
||||
)
|
||||
|
||||
held_open: Final = in_fresh_context(run)
|
||||
|
||||
assert len(fan_out._roots) == otel_providers._MAX_OPEN_REQUESTS
|
||||
for span in held_open:
|
||||
span.end()
|
||||
assert (fan_out._roots, dict(fan_out._open)) == ({}, {})
|
||||
|
||||
def test_a_span_starting_under_an_ended_span_below_the_root_follows_its_request(self):
|
||||
"""A post-call write started under the route span after its failed request ended
|
||||
is kept with the request at the error rate, rather than decided on its own at the
|
||||
success rate."""
|
||||
dest_exporter: Final = InMemorySpanExporter()
|
||||
provider: Final = self._fan_out(dest_exporter)
|
||||
destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),)
|
||||
|
||||
def run():
|
||||
set_request_destinations(destinations)
|
||||
tracer: Final = get_tracer(provider, "litellm")
|
||||
with tracer.start_as_current_span("POST /v1/chat/completions") as root:
|
||||
with tracer.start_as_current_span("route gpt-4") as route:
|
||||
route_context: Final = trace_api.set_span_in_context(route)
|
||||
root.set_status(Status(StatusCode.ERROR, "upstream 500"))
|
||||
with tracer.start_as_current_span("postgres INSERT LiteLLM_SpendLogs", context=route_context):
|
||||
pass
|
||||
|
||||
in_fresh_context(run)
|
||||
|
||||
assert {s.name for s in dest_exporter.get_finished_spans()} == {
|
||||
"POST /v1/chat/completions",
|
||||
"route gpt-4",
|
||||
"postgres INSERT LiteLLM_SpendLogs",
|
||||
}
|
||||
|
||||
def test_concurrent_calls_of_one_trace_each_answer_to_their_own_outcome(self):
|
||||
"""Two calls of the caller's trace in flight at once on one process, one failing,
|
||||
are decided as two processes serving one each would decide them."""
|
||||
dest_exporter: Final = InMemorySpanExporter()
|
||||
provider: Final = self._fan_out(dest_exporter)
|
||||
destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),)
|
||||
remote: Final = trace_api.SpanContext(
|
||||
trace_id=KEPT_AT_HALF, span_id=1, is_remote=True, trace_flags=trace_api.TraceFlags(0x01)
|
||||
)
|
||||
|
||||
def run():
|
||||
set_request_destinations(destinations)
|
||||
tracer: Final = get_tracer(provider, "litellm")
|
||||
caller: Final = trace_api.set_span_in_context(trace_api.NonRecordingSpan(remote))
|
||||
with tracer.start_as_current_span("POST failed", context=caller) as failed:
|
||||
with tracer.start_as_current_span("POST succeeded", context=caller):
|
||||
pass
|
||||
failed.set_status(Status(StatusCode.ERROR, "upstream 500"))
|
||||
|
||||
in_fresh_context(run)
|
||||
|
||||
assert [s.name for s in dest_exporter.get_finished_spans()] == ["POST failed"]
|
||||
|
||||
def test_a_late_span_whose_verdict_was_forgotten_is_decided_on_its_own_not_stranded(self):
|
||||
"""Once enough other requests have been decided to evict a request's verdict, a span
|
||||
of it that starts late (a post-call database write) opens and closes its own count,
|
||||
so it is decided on a draw of its own as soon as it ends instead of waiting in a tree
|
||||
nothing would ever close."""
|
||||
so it is decided on its own as soon as it ends instead of waiting in a tree nothing
|
||||
would ever close."""
|
||||
dest_exporter = InMemorySpanExporter()
|
||||
provider = self._fan_out(dest_exporter, draw=lambda: 0.1)
|
||||
destinations = (arize_destination(arize_success_sampling_rate="0.5"),)
|
||||
provider: Final = self._fan_out(dest_exporter)
|
||||
destinations: Final = (arize_destination(arize_success_sampling_rate="1.0"),)
|
||||
|
||||
def run():
|
||||
set_request_destinations(destinations)
|
||||
|
|
@ -1353,18 +1632,17 @@ class TestTenantSampling:
|
|||
assert names[0] == "POST /v1/chat/completions"
|
||||
assert names[-1] == "postgres INSERT LiteLLM_SpendLogs", "exported as it ended, not held for a root"
|
||||
|
||||
def test_a_callers_traceparent_neither_steers_the_draw_nor_hides_the_root(self):
|
||||
"""The draw is random, not read from the trace id a caller may have chosen, and the
|
||||
server span under a remote parent still closes the tree."""
|
||||
def test_a_callers_traceparent_decides_by_its_trace_id_and_still_closes_the_tree(self):
|
||||
"""A request continuing a caller's trace is decided on that trace id, and the server
|
||||
span under a remote parent still closes the tree."""
|
||||
dest_exporter = InMemorySpanExporter()
|
||||
draws = iter([0.9, 0.1])
|
||||
provider = self._fan_out(dest_exporter, draw=lambda: next(draws))
|
||||
provider: Final = self._fan_out(dest_exporter)
|
||||
destinations = (arize_destination(arize_success_sampling_rate="0.5"),)
|
||||
|
||||
def run():
|
||||
set_request_destinations(destinations)
|
||||
tracer = get_tracer(provider, "litellm")
|
||||
for trace_id in (1, 2):
|
||||
for trace_id in (DROPPED_AT_HALF, KEPT_AT_HALF):
|
||||
remote = trace_api.SpanContext(
|
||||
trace_id=trace_id, span_id=1, is_remote=True, trace_flags=trace_api.TraceFlags(0x01)
|
||||
)
|
||||
|
|
@ -1377,12 +1655,12 @@ class TestTenantSampling:
|
|||
|
||||
kept = dest_exporter.get_finished_spans()
|
||||
assert [s.name for s in kept] == ["chat gpt-4", "POST /v1/chat/completions"]
|
||||
assert all(s.context.trace_id == 2 for s in kept), "the second request, whose draw of 0.1 passed"
|
||||
assert all(s.context.trace_id == KEPT_AT_HALF for s in kept)
|
||||
|
||||
def test_a_root_that_never_ends_holds_the_tree_only_up_to_the_bound(self):
|
||||
dest_exporter = InMemorySpanExporter()
|
||||
provider = self._fan_out(dest_exporter, draw=lambda: 0.0)
|
||||
destinations = (arize_destination(arize_success_sampling_rate="0.5"),)
|
||||
provider: Final = self._fan_out(dest_exporter)
|
||||
destinations: Final = (arize_destination(arize_success_sampling_rate="1.0"),)
|
||||
|
||||
def run():
|
||||
set_request_destinations(destinations)
|
||||
|
|
@ -1397,8 +1675,8 @@ class TestTenantSampling:
|
|||
|
||||
def test_a_flood_of_waiting_trees_decides_the_oldest_on_what_it_has(self):
|
||||
dest_exporter = InMemorySpanExporter()
|
||||
provider = self._fan_out(dest_exporter, draw=lambda: 0.0)
|
||||
destinations = (arize_destination(arize_success_sampling_rate="0.5"),)
|
||||
provider: Final = self._fan_out(dest_exporter)
|
||||
destinations: Final = (arize_destination(arize_success_sampling_rate="1.0"),)
|
||||
|
||||
def run():
|
||||
set_request_destinations(destinations)
|
||||
|
|
@ -1413,21 +1691,21 @@ class TestTenantSampling:
|
|||
|
||||
assert [s.name for s in dest_exporter.get_finished_spans()] == ["chat 0"]
|
||||
|
||||
def test_a_flood_of_open_traces_forgets_the_oldest_and_decides_it_as_its_spans_end(self):
|
||||
"""The fan-out counts open spans for a bounded number of traces, a bound far above
|
||||
what one process keeps in flight, since every trace of the provider is counted, not
|
||||
only the sampled ones. A trace pushed out of that count is decided on what it has as
|
||||
def test_a_flood_of_open_requests_forgets_the_oldest_and_decides_it_as_its_spans_end(self):
|
||||
"""The fan-out counts open spans for a bounded number of requests, a bound far above
|
||||
what one process keeps in flight, since every request of the provider is counted, not
|
||||
only the sampled ones. A request pushed out of that count is decided on what it has as
|
||||
its next span ends, root still open, rather than held until a bound gets to it."""
|
||||
dest_exporter = InMemorySpanExporter()
|
||||
provider = self._fan_out(dest_exporter, draw=lambda: 0.0)
|
||||
destinations = (arize_destination(arize_success_sampling_rate="0.5"),)
|
||||
provider: Final = self._fan_out(dest_exporter)
|
||||
destinations: Final = (arize_destination(arize_success_sampling_rate="1.0"),)
|
||||
|
||||
def run():
|
||||
set_request_destinations(destinations)
|
||||
tracer = get_tracer(provider, "litellm")
|
||||
root = tracer.start_span("POST /v1/chat/completions", context=trace_api.Context())
|
||||
tracer.start_span("auth /v1/chat/completions", context=trace_api.set_span_in_context(root)).end()
|
||||
for index in range(otel_providers._MAX_OPEN_TRACES):
|
||||
for index in range(otel_providers._MAX_OPEN_REQUESTS):
|
||||
tracer.start_span(f"POST {index}", context=trace_api.Context())
|
||||
assert dest_exporter.get_finished_spans() == (), "still held before the count is forgotten"
|
||||
tracer.start_span("chat gpt-4", context=trace_api.set_span_in_context(root)).end()
|
||||
|
|
@ -1443,14 +1721,12 @@ class TestTenantSampling:
|
|||
|
||||
def test_shutdown_decides_what_is_still_held(self):
|
||||
dest_exporter = InMemorySpanExporter()
|
||||
fan_out = TenantFanOutSpanProcessor(
|
||||
processor_factory=lambda _d: SimpleSpanProcessor(dest_exporter), sampling_draw=lambda: 0.0
|
||||
)
|
||||
fan_out: Final = TenantFanOutSpanProcessor(processor_factory=lambda _d: SimpleSpanProcessor(dest_exporter))
|
||||
provider = TracerProvider()
|
||||
provider.add_span_processor(fan_out)
|
||||
|
||||
def run():
|
||||
set_request_destinations((arize_destination(arize_success_sampling_rate="0.5"),))
|
||||
set_request_destinations((arize_destination(arize_success_sampling_rate="1.0"),))
|
||||
tracer = get_tracer(provider, "litellm")
|
||||
with tracer.start_as_current_span("POST /v1/chat/completions"):
|
||||
with tracer.start_as_current_span("chat gpt-4"):
|
||||
|
|
@ -1476,14 +1752,12 @@ class TestTenantSampling:
|
|||
straggler.end()
|
||||
super().on_end(span)
|
||||
|
||||
fan_out = TenantFanOutSpanProcessor(
|
||||
processor_factory=lambda _d: _EndsAStragglerOnExport(dest_exporter), sampling_draw=lambda: 0.0
|
||||
)
|
||||
fan_out: Final = TenantFanOutSpanProcessor(processor_factory=lambda _d: _EndsAStragglerOnExport(dest_exporter))
|
||||
provider = TracerProvider()
|
||||
provider.add_span_processor(fan_out)
|
||||
|
||||
def run():
|
||||
set_request_destinations((arize_destination(arize_success_sampling_rate="0.5"),))
|
||||
set_request_destinations((arize_destination(arize_success_sampling_rate="1.0"),))
|
||||
tracer = get_tracer(provider, "litellm")
|
||||
root = tracer.start_span("POST /v1/chat/completions", context=trace_api.Context())
|
||||
tracer.start_span("chat gpt-4", context=trace_api.set_span_in_context(root)).end()
|
||||
|
|
@ -1495,20 +1769,20 @@ class TestTenantSampling:
|
|||
|
||||
assert {s.name for s in dest_exporter.get_finished_spans()} == {"chat gpt-4", "embed ada"}
|
||||
|
||||
def test_a_zero_rate_drops_even_a_draw_of_zero(self):
|
||||
def test_a_zero_rate_drops_a_trace_a_half_rate_keeps(self):
|
||||
dest_exporter = InMemorySpanExporter()
|
||||
provider = self._fan_out(dest_exporter, draw=lambda: 0.0)
|
||||
provider: Final = self._fan_out(dest_exporter, trace_id=KEPT_AT_HALF)
|
||||
|
||||
self._run(provider, (arize_destination(arize_success_sampling_rate="0.0"),))
|
||||
|
||||
assert dest_exporter.get_finished_spans() == ()
|
||||
|
||||
def test_the_draw_is_compared_to_the_rate_the_way_the_legacy_callback_compares_it(self):
|
||||
def test_the_trace_id_is_read_against_the_rate(self):
|
||||
dropped_exporter, kept_exporter = InMemorySpanExporter(), InMemorySpanExporter()
|
||||
destinations = (arize_destination(arize_success_sampling_rate="0.2"),)
|
||||
destinations: Final = (arize_destination(arize_success_sampling_rate="0.5"),)
|
||||
|
||||
self._run(self._fan_out(dropped_exporter, draw=lambda: 0.3), destinations)
|
||||
self._run(self._fan_out(kept_exporter, draw=lambda: 0.2), destinations)
|
||||
self._run(self._fan_out(dropped_exporter, trace_id=DROPPED_AT_HALF), destinations)
|
||||
self._run(self._fan_out(kept_exporter, trace_id=KEPT_AT_HALF), destinations)
|
||||
|
||||
assert dropped_exporter.get_finished_spans() == ()
|
||||
assert {s.name for s in kept_exporter.get_finished_spans()} == self.NAMES
|
||||
|
|
@ -1527,6 +1801,39 @@ class TestTenantSampling:
|
|||
assert all(names == self.NAMES for names in by_trace.values())
|
||||
assert 0 < len(by_trace) < 64
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("rates", "failed"),
|
||||
[({"arize_success_sampling_rate": "0.5"}, False), ({"arize_error_sampling_rate": "0.5"}, True)],
|
||||
ids=["success-rate", "error-rate"],
|
||||
)
|
||||
def test_two_processes_serving_one_trace_reach_the_same_verdict(self, rates: Mapping[str, str], failed: bool):
|
||||
"""Two pods behind a load balancer each serve one call of a trace the caller
|
||||
continued with its ``traceparent``: the team's space gets both calls or neither,
|
||||
and a rate of 0.5 over UUID-shaped trace ids keeps some traces and drops others."""
|
||||
exporters: Final = (InMemorySpanExporter(), InMemorySpanExporter())
|
||||
pods: Final = tuple(self._fan_out(exporter) for exporter in exporters)
|
||||
destinations: Final = (arize_destination(**rates),)
|
||||
trace_ids: Final = tuple(uuid.uuid5(uuid.NAMESPACE_URL, f"pod/{index}").int for index in range(32))
|
||||
|
||||
def run(provider: TracerProvider):
|
||||
set_request_destinations(destinations)
|
||||
tracer: Final = get_tracer(provider, "litellm")
|
||||
for trace_id in trace_ids:
|
||||
with tracer.start_as_current_span("POST /v1/chat/completions", context=caller_context(trace_id)):
|
||||
with tracer.start_as_current_span("chat gpt-4") as span:
|
||||
if failed:
|
||||
span.set_status(Status(StatusCode.ERROR, "upstream 500"))
|
||||
|
||||
for provider in pods:
|
||||
in_fresh_context(run, provider)
|
||||
|
||||
first, second = (
|
||||
{s.context.trace_id for s in exporter.get_finished_spans() if s.context is not None}
|
||||
for exporter in exporters
|
||||
)
|
||||
assert first == second
|
||||
assert 0 < len(first) < len(trace_ids)
|
||||
|
||||
def test_in_additive_mode_the_operators_own_copy_is_not_sampled(self, monkeypatch):
|
||||
"""The rates are the team's setting for the team's space; the operator's backbone keeps
|
||||
every span."""
|
||||
|
|
@ -1547,7 +1854,7 @@ class TestTenantSampling:
|
|||
|
||||
assert {s.name for s in dest_exporter.get_finished_spans()} == self.NAMES
|
||||
|
||||
def test_two_views_of_one_exporter_each_draw_at_their_own_rate(self):
|
||||
def test_two_views_of_one_exporter_each_answer_to_their_own_rate(self):
|
||||
"""Two destinations to the same space with different rates share an exporter, not a
|
||||
verdict: the 1.0 view gets the tree once and the 0.0 view adds nothing."""
|
||||
dest_exporter = InMemorySpanExporter()
|
||||
|
|
@ -1582,7 +1889,7 @@ class TestTenantSampling:
|
|||
@pytest.mark.parametrize("bad", ["abc", "-0.1", "1.5", "nan"])
|
||||
def test_an_unusable_rate_exports_rather_than_dropping(self, bad):
|
||||
dest_exporter = InMemorySpanExporter()
|
||||
provider = self._fan_out(dest_exporter, draw=lambda: 0.99)
|
||||
provider: Final = self._fan_out(dest_exporter, trace_id=DROPPED_AT_HALF)
|
||||
|
||||
self._run(provider, (arize_destination(arize_success_sampling_rate=bad),))
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import asyncio
|
||||
import datetime
|
||||
import json
|
||||
import queue
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
|
|
@ -2819,6 +2820,123 @@ def test_raise_on_model_repetition_tolerates_empty_choices(
|
|||
wrapper.raise_on_model_repetition()
|
||||
|
||||
|
||||
_STUCK_CHUNK: Final = "the model is stuck on this"
|
||||
_STUCK_PROVIDER_CHUNK_CAP: Final = 10 * litellm.REPEATED_STREAMING_CHUNK_LIMIT
|
||||
|
||||
|
||||
class _StuckProvider:
|
||||
def __init__(self) -> None:
|
||||
self.closed_early = threading.Event()
|
||||
self.chunks_sent = 0
|
||||
|
||||
def _next_chunk(self) -> ModelResponseStream:
|
||||
self.chunks_sent += 1
|
||||
return _make_chunk(_STUCK_CHUNK)
|
||||
|
||||
def sync_stream(self) -> Iterator[ModelResponseStream]:
|
||||
try:
|
||||
yield from (self._next_chunk() for _ in range(_STUCK_PROVIDER_CHUNK_CAP))
|
||||
except GeneratorExit:
|
||||
self.closed_early.set()
|
||||
raise
|
||||
|
||||
async def async_stream(self) -> AsyncIterator[ModelResponseStream]:
|
||||
try:
|
||||
for _ in range(_STUCK_PROVIDER_CHUNK_CAP):
|
||||
yield self._next_chunk()
|
||||
except GeneratorExit:
|
||||
self.closed_early.set()
|
||||
raise
|
||||
|
||||
|
||||
class _StuckStreamFailureRecorder(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.failures: Final = queue.SimpleQueue[str]()
|
||||
|
||||
def log_failure_event(self, kwargs: object, response_obj: object, start_time: object, end_time: object) -> None:
|
||||
self.failures.put(str(cast(Mapping[str, object], kwargs)["exception"]))
|
||||
|
||||
async def async_log_failure_event(
|
||||
self, kwargs: object, response_obj: object, start_time: object, end_time: object
|
||||
) -> None:
|
||||
self.failures.put(str(cast(Mapping[str, object], kwargs)["exception"]))
|
||||
|
||||
|
||||
def _stuck_stream_wrapper(
|
||||
provider_stream: Iterator[ModelResponseStream] | AsyncIterator[ModelResponseStream],
|
||||
call_type: Literal["completion", "acompletion"],
|
||||
recorder: _StuckStreamFailureRecorder,
|
||||
) -> CustomStreamWrapper:
|
||||
logging_obj: Final = Logging(
|
||||
model="test-model",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=True,
|
||||
call_type=call_type,
|
||||
start_time=datetime.datetime(2026, 1, 1),
|
||||
litellm_call_id=f"stuck-stream-{call_type}",
|
||||
function_id=f"stuck-stream-{call_type}",
|
||||
dynamic_failure_callbacks=[recorder],
|
||||
dynamic_async_failure_callbacks=[recorder],
|
||||
)
|
||||
logging_obj.update_environment_variables(
|
||||
model="test-model",
|
||||
optional_params={},
|
||||
litellm_params=get_litellm_params(custom_llm_provider="openai", acompletion=call_type == "acompletion"),
|
||||
)
|
||||
return CustomStreamWrapper(
|
||||
completion_stream=provider_stream,
|
||||
model="test-model",
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_stream_stuck_on_one_chunk_closes_the_provider_stream_and_logs_one_failure():
|
||||
"""https://github.com/BerriAI/litellm/issues/19428 and https://github.com/BerriAI/litellm/issues/13786"""
|
||||
provider: Final = _StuckProvider()
|
||||
recorder: Final = _StuckStreamFailureRecorder()
|
||||
wrapper: Final = _stuck_stream_wrapper(provider.async_stream(), "acompletion", recorder)
|
||||
|
||||
with pytest.raises(Exception, match="repeating the same chunk"):
|
||||
async for _ in wrapper:
|
||||
pass
|
||||
assert provider.closed_early.is_set()
|
||||
assert provider.chunks_sent <= litellm.REPEATED_STREAMING_CHUNK_LIMIT + 1
|
||||
chunks_sent_at_failure: Final = provider.chunks_sent
|
||||
|
||||
with pytest.raises(StopAsyncIteration):
|
||||
await wrapper.__anext__()
|
||||
await wrapper.wait_for_failure_logging(timeout=5)
|
||||
|
||||
assert provider.chunks_sent == chunks_sent_at_failure
|
||||
assert "repeating the same chunk" in recorder.failures.get_nowait()
|
||||
assert recorder.failures.empty()
|
||||
|
||||
|
||||
def test_sync_stream_stuck_on_one_chunk_closes_the_provider_stream_and_logs_one_failure():
|
||||
"""Sync counterpart of the async test above"""
|
||||
provider: Final = _StuckProvider()
|
||||
recorder: Final = _StuckStreamFailureRecorder()
|
||||
wrapper: Final = _stuck_stream_wrapper(provider.sync_stream(), "completion", recorder)
|
||||
|
||||
with pytest.raises(Exception, match="repeating the same chunk"):
|
||||
for _ in wrapper:
|
||||
pass
|
||||
assert provider.closed_early.is_set()
|
||||
assert provider.chunks_sent <= litellm.REPEATED_STREAMING_CHUNK_LIMIT + 1
|
||||
chunks_sent_at_failure: Final = provider.chunks_sent
|
||||
assert "repeating the same chunk" in recorder.failures.get(timeout=5)
|
||||
|
||||
with pytest.raises(StopIteration):
|
||||
next(wrapper)
|
||||
|
||||
assert provider.chunks_sent == chunks_sent_at_failure
|
||||
with pytest.raises(queue.Empty):
|
||||
recorder.failures.get(timeout=0.5)
|
||||
|
||||
|
||||
def test_usage_chunk_after_finish_reason_updates_hidden_params(logging_obj):
|
||||
"""
|
||||
Test that provider-reported usage from a post-finish_reason chunk
|
||||
|
|
|
|||
|
|
@ -511,6 +511,54 @@ class TestDeepSeekVisionMultimodalContent:
|
|||
assert any(block.get("type") == "image_url" for block in content)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_deepseek_httpx_transport")
|
||||
@pytest.mark.parametrize("stream", (False, True))
|
||||
@pytest.mark.parametrize("model", ("deepseek-flash", "deepseek-v4-pro", "deepseek-chat", "deepseek-reasoner"))
|
||||
@pytest.mark.parametrize("effort", ("minimal", "low", "medium", "high", "xhigh", "max", "ultra"))
|
||||
@pytest.mark.parametrize("thinking", (None, "enabled"))
|
||||
@pytest.mark.respx(assert_all_called=True)
|
||||
async def test_deepseek_forwards_native_effort(
|
||||
respx_mock: respx.MockRouter, model: str, effort: str, thinking: str | None, stream: bool
|
||||
) -> None:
|
||||
route: Final = respx_mock.post(DEEPSEEK_CHAT_COMPLETIONS_URL).mock(return_value=_deepseek_response(stream))
|
||||
response: Final = await litellm.acompletion(
|
||||
model=f"deepseek/{model}",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
api_key=DEEPSEEK_API_KEY,
|
||||
reasoning_effort=effort,
|
||||
stream=stream,
|
||||
**({"thinking": {"type": thinking}} if thinking is not None else {}),
|
||||
)
|
||||
if stream:
|
||||
async for _ in response:
|
||||
pass
|
||||
body: Final = json.loads(route.calls[0].request.content)
|
||||
assert body["reasoning_effort"] == effort
|
||||
assert body["thinking"] == {"type": "enabled"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"thinking,effort,expected_thinking",
|
||||
(
|
||||
(None, None, None),
|
||||
(None, "none", "disabled"),
|
||||
("enabled", "none", "enabled"),
|
||||
("disabled", "max", "disabled"),
|
||||
("enabled", None, "enabled"),
|
||||
(None, "default", "enabled"),
|
||||
(None, "bogus", "enabled"),
|
||||
(None, {"effort": "high"}, "enabled"),
|
||||
),
|
||||
)
|
||||
def test_deepseek_preserves_toggle_without_forwarding_unsupported_effort(
|
||||
thinking: str | None, effort: str | dict[str, str] | None, expected_thinking: str | None
|
||||
) -> None:
|
||||
params: Final = {"reasoning_effort": effort, **({"thinking": {"type": thinking}} if thinking is not None else {})}
|
||||
result: Final = DeepSeekChatConfig().map_openai_params(params, {}, "deepseek-flash", False)
|
||||
assert "reasoning_effort" not in result
|
||||
assert result.get("thinking") == ({"type": expected_thinking} if expected_thinking else None)
|
||||
|
||||
|
||||
class TestDeepSeekThinkingParams:
|
||||
"""Test thinking and reasoning_effort parameter handling for DeepSeek."""
|
||||
|
||||
|
|
@ -735,9 +783,7 @@ async def test_deepseek_provider_async_completion(stream):
|
|||
messages = [{"role": "user", "content": "Hello, world!"}]
|
||||
|
||||
# Mock AsyncHTTPHandler.post method for async test
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.llm_http_handler.AsyncHTTPHandler.post"
|
||||
) as mock_post:
|
||||
with patch("litellm.llms.custom_httpx.llm_http_handler.AsyncHTTPHandler.post") as mock_post:
|
||||
mock_response_data = litellm.ModelResponse(
|
||||
choices=[
|
||||
litellm.Choices(
|
||||
|
|
@ -775,9 +821,7 @@ async def test_deepseek_provider_async_completion(stream):
|
|||
# Check request body
|
||||
request_body = json.loads(call_args.kwargs["data"])
|
||||
assert call_args.kwargs["url"] == "https://api.deepseek.com/beta/chat/completions"
|
||||
assert (
|
||||
request_body["model"] == "deepseek-reasoner"
|
||||
) # Model name should be stripped of provider prefix
|
||||
assert request_body["model"] == "deepseek-reasoner" # Model name should be stripped of provider prefix
|
||||
assert request_body["messages"] == messages
|
||||
assert request_body["stream"] == stream
|
||||
|
||||
|
|
|
|||
|
|
@ -314,4 +314,4 @@ class TestElevenLabsTextToSpeechTransformation:
|
|||
)
|
||||
|
||||
assert voice_id in url
|
||||
assert "output_format=pcm_44100" in url
|
||||
assert "output_format=pcm_24000" in url
|
||||
|
|
|
|||
|
|
@ -55,3 +55,31 @@ def test_speech_keeps_an_internal_prefixed_kwarg_out_of_the_elevenlabs_request(r
|
|||
sent: Final = json.loads(respx_mock.calls[0].request.content)
|
||||
assert "_litellm_undeclared_sentinel" not in sent, sent
|
||||
assert sent["text"] == "hi"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("response_format", "expected_output_format"),
|
||||
[("wav", "wav_24000"), ("pcm", "pcm_24000")],
|
||||
)
|
||||
def test_speech_maps_wav_and_pcm_to_elevenlabs_output_formats(
|
||||
respx_mock: respx.MockRouter,
|
||||
response_format: str,
|
||||
expected_output_format: str,
|
||||
) -> None:
|
||||
api_base: Final = "http://localhost:12346"
|
||||
mock_route: Final = respx_mock.post(url__regex=rf"{api_base}/v1/text-to-speech/.*").mock(
|
||||
return_value=httpx.Response(status_code=200, content=b"audio", headers={"content-type": "audio/mpeg"})
|
||||
)
|
||||
|
||||
litellm.speech(
|
||||
model="elevenlabs/eleven_multilingual_v2",
|
||||
input="hi",
|
||||
voice="21m00Tcm4TlvDq8ikWAM",
|
||||
response_format=response_format,
|
||||
api_base=api_base,
|
||||
api_key="fake_elevenlabs_api_key",
|
||||
)
|
||||
|
||||
assert mock_route.called
|
||||
request_url: Final = httpx.URL(str(respx_mock.calls[0].request.url))
|
||||
assert request_url.params["output_format"] == expected_output_format
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import logging
|
||||
from typing import Final
|
||||
from unittest.mock import patch
|
||||
|
||||
import httpx
|
||||
|
|
@ -523,3 +524,68 @@ def test_tool_search_reply_sent_to_groq_drops_litellm_response_fields():
|
|||
}
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("is_async", [False, True])
|
||||
@pytest.mark.parametrize("with_tool_call", [False, True])
|
||||
async def test_streamed_reasoning_reply_can_be_replayed(is_async: bool, with_tool_call: bool) -> None:
|
||||
tool_calls: Final = [
|
||||
{
|
||||
"index": 0,
|
||||
"id": "call_weather",
|
||||
"type": "function",
|
||||
"function": {"name": "get_weather", "arguments": '{"city":"Paris"}'},
|
||||
}
|
||||
]
|
||||
handler: Final = GroqChatCompletionStreamingHandler(streaming_response=iter(()), sync_stream=True)
|
||||
chunk: Final = handler.chunk_parser(
|
||||
{
|
||||
"id": "chatcmpl-reasoning",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1769511767,
|
||||
"model": "openai/gpt-oss-20b",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {
|
||||
"role": "assistant",
|
||||
"reasoning": "I should check the weather.",
|
||||
"content": "" if with_tool_call else "Hello there!",
|
||||
**({"tool_calls": tool_calls} if with_tool_call else {}),
|
||||
},
|
||||
"finish_reason": "tool_calls" if with_tool_call else "stop",
|
||||
}
|
||||
],
|
||||
}
|
||||
)
|
||||
reply: Final = litellm.stream_chunk_builder([chunk]).choices[0].message
|
||||
assert reply.reasoning_content == "I should check the weather."
|
||||
history: Final = [reply.model_dump(exclude_none=True)]
|
||||
config: Final = GroqChatConfig()
|
||||
|
||||
if is_async:
|
||||
replayed = await config._transform_messages(messages=history, model="openai/gpt-oss-20b", is_async=True)
|
||||
else:
|
||||
replayed = config._transform_messages(messages=history, model="openai/gpt-oss-20b", is_async=False)
|
||||
|
||||
assert replayed == [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "" if with_tool_call else "Hello there!",
|
||||
**(
|
||||
{
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_weather",
|
||||
"type": "function",
|
||||
"function": {"name": "get_weather", "arguments": '{"city":"Paris"}'},
|
||||
}
|
||||
]
|
||||
}
|
||||
if with_tool_call
|
||||
else {}
|
||||
),
|
||||
}
|
||||
]
|
||||
assert reply.reasoning_content == "I should check the weather."
|
||||
|
|
|
|||
|
|
@ -19,6 +19,22 @@ def test_full_payload_guardrail_text_excludes_the_server_streaming_marker():
|
|||
assert json.loads(text) == body, text
|
||||
|
||||
|
||||
def test_full_payload_guardrail_text_excludes_the_injected_litellm_metadata():
|
||||
body: Final = {"model": "m", "messages": [{"role": "user", "content": "hi"}], "guardrails": ["g"]}
|
||||
litellm_metadata: Final = {
|
||||
"user_api_key_token": "hashed-key",
|
||||
"user_api_key_user_email": "someone@example.com",
|
||||
"user_api_key_user_role": "proxy_admin",
|
||||
}
|
||||
|
||||
text: Final = PassThroughEndpointHandler()._extract_text_for_guardrail(
|
||||
{**body, "litellm_metadata": litellm_metadata},
|
||||
None,
|
||||
)
|
||||
|
||||
assert json.loads(text) == body, text
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"marker",
|
||||
[
|
||||
|
|
|
|||
|
|
@ -5,13 +5,19 @@ Based on the test patterns from other rerank providers and the current Vertex AI
|
|||
|
||||
import json
|
||||
import os
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import responses
|
||||
import respx
|
||||
|
||||
import litellm
|
||||
from litellm.llms.vertex_ai.rerank.transformation import VertexAIRerankConfig
|
||||
from litellm.types.rerank import RerankResponse
|
||||
from litellm.utils import ProviderConfigManager
|
||||
|
||||
|
||||
class TestVertexAIRerankTransform:
|
||||
|
|
@ -646,3 +652,83 @@ class TestVertexAIRerankTransform:
|
|||
second_call = mock_ensure_access_token.call_args_list[1]
|
||||
assert first_call.kwargs["credentials"] == "path/to/credentials.json"
|
||||
assert second_call.kwargs["credentials"] == "path/to/credentials.json"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def clear_vertex_rerank_cache(request: pytest.FixtureRequest) -> None:
|
||||
ProviderConfigManager._get_vertex_ai_rerank_config_cached.cache_clear()
|
||||
request.addfinalizer(
|
||||
ProviderConfigManager._get_vertex_ai_rerank_config_cached.cache_clear
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("clear_vertex_rerank_cache")
|
||||
@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"])
|
||||
async def test_rerank_reuses_scoped_credentials_and_refreshes_expired_tokens(
|
||||
is_async: bool,
|
||||
) -> None:
|
||||
now: Final = datetime(2025, 1, 1)
|
||||
credentials: Final = tuple(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "authorized_user",
|
||||
"client_id": identity,
|
||||
"client_secret": "synthetic-client-secret",
|
||||
"refresh_token": f"synthetic-refresh-{identity}",
|
||||
}
|
||||
)
|
||||
for identity in ("first-principal", "second-principal")
|
||||
)
|
||||
with (
|
||||
patch("litellm.disable_aiohttp_transport", True),
|
||||
patch("google.auth._helpers.utcnow", return_value=now) as clock,
|
||||
responses.RequestsMock(assert_all_requests_are_fired=False) as token_http,
|
||||
respx.mock as rank_http,
|
||||
):
|
||||
for token in ("initial", "other-credential", "other-project", "renewed"):
|
||||
token_http.post(
|
||||
"https://oauth2.googleapis.com/token",
|
||||
json={"access_token": token, "expires_in": 3600, "token_type": "Bearer"},
|
||||
)
|
||||
rank_http.post(
|
||||
url__regex=r"https://discoveryengine\.googleapis\.com/v1/projects/[^/]+/"
|
||||
r"locations/global/rankingConfigs/default_ranking_config:rank"
|
||||
).respond(200, json={"records": [{"id": "0", "score": 0.9}]})
|
||||
|
||||
async def rank(
|
||||
credential: str, project: str, expected_token: str, refresh_count: int
|
||||
) -> None:
|
||||
params: Final = {
|
||||
"model": "vertex_ai/semantic-ranker-default@latest",
|
||||
"query": f"Query for {expected_token}",
|
||||
"documents": ["A relevant document"],
|
||||
"vertex_credentials": credential,
|
||||
"vertex_project": project,
|
||||
"caching": False,
|
||||
"num_retries": 0,
|
||||
}
|
||||
result: Final = (
|
||||
await litellm.arerank(**params)
|
||||
if is_async
|
||||
else litellm.rerank(**params)
|
||||
)
|
||||
assert isinstance(result, RerankResponse)
|
||||
assert result.results[0]["index"] == 0
|
||||
assert len(token_http.calls) == refresh_count
|
||||
request: Final = rank_http.calls.last.request
|
||||
assert request.headers["Authorization"] == f"Bearer {expected_token}"
|
||||
assert request.headers["X-Goog-User-Project"] == project
|
||||
assert str(request.url) == (
|
||||
f"https://discoveryengine.googleapis.com/v1/projects/{project}/"
|
||||
"locations/global/rankingConfigs/default_ranking_config:rank"
|
||||
)
|
||||
assert json.loads(request.content)["query"] == params["query"]
|
||||
|
||||
await rank(credentials[0], "first-project", "initial", 1)
|
||||
await rank(credentials[0], "first-project", "initial", 1)
|
||||
await rank(credentials[1], "first-project", "other-credential", 2)
|
||||
await rank(credentials[0], "second-project", "other-project", 3)
|
||||
await rank(credentials[0], "first-project", "initial", 3)
|
||||
clock.return_value = now + timedelta(hours=2)
|
||||
await rank(credentials[0], "first-project", "renewed", 4)
|
||||
|
|
|
|||
|
|
@ -6,13 +6,17 @@ import os
|
|||
import pathlib
|
||||
import re
|
||||
import traceback
|
||||
from typing import Final
|
||||
from datetime import datetime, timedelta
|
||||
from threading import Event
|
||||
from typing import Final, Literal
|
||||
from unittest.mock import MagicMock, call, patch
|
||||
|
||||
import google.auth.exceptions
|
||||
import pytest
|
||||
import requests
|
||||
import responses
|
||||
|
||||
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
|
||||
from litellm.constants import DEFAULT_MAX_LRU_CACHE_SIZE, DEFAULT_MAX_RECURSE_DEPTH
|
||||
|
||||
|
||||
import litellm
|
||||
|
|
@ -2589,3 +2593,120 @@ def test_get_access_token_reraises_malformed_service_account_credentials_without
|
|||
assert [record.getMessage().split(". Error: ")[0] for record in caplog.records] == [
|
||||
"Failed to load vertex credentials. Check to see if credentials containing partial/invalid information"
|
||||
]
|
||||
|
||||
|
||||
def _authorized_user_cache_credentials(identity: str) -> str:
|
||||
return json.dumps(
|
||||
{
|
||||
"type": "authorized_user",
|
||||
"client_id": identity,
|
||||
"client_secret": "synthetic-secret",
|
||||
"refresh_token": identity,
|
||||
"quota_project_id": identity,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"])
|
||||
@pytest.mark.parametrize("inferred_project", [False, True], ids=["explicit", "inferred"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_credential_cache_bounds_retention_and_reloads_least_recently_used(
|
||||
is_async: bool, inferred_project: bool
|
||||
) -> None:
|
||||
vertex: Final = VertexBase()
|
||||
limit: Final = max(0, DEFAULT_MAX_LRU_CACHE_SIZE)
|
||||
with (
|
||||
patch("google.auth._helpers.utcnow", return_value=datetime(2025, 1, 1)),
|
||||
responses.RequestsMock() as oauth,
|
||||
):
|
||||
def token_response(request: requests.PreparedRequest) -> tuple[int, dict[str, str], str]:
|
||||
return 200, {}, json.dumps({"access_token": f"token-{len(oauth.calls) + 1}", "expires_in": 3600})
|
||||
|
||||
oauth.add_callback(responses.POST, "https://oauth2.googleapis.com/token", callback=token_response)
|
||||
|
||||
async def access(index: int, infer: bool = False) -> tuple[str, str]:
|
||||
identity: Final = f"project-{index}"
|
||||
credentials: Final = _authorized_user_cache_credentials(identity)
|
||||
project: Final = None if infer else identity
|
||||
result: Final = (
|
||||
await vertex.get_access_token_async(credentials, project)
|
||||
if is_async
|
||||
else vertex.get_access_token(credentials, project)
|
||||
)
|
||||
assert result[1] == identity
|
||||
assert len(vertex._credentials_project_mapping) <= limit
|
||||
return result
|
||||
|
||||
if limit == 0:
|
||||
assert await access(0, inferred_project) == ("token-1", "project-0")
|
||||
assert await access(0, inferred_project) == ("token-2", "project-0")
|
||||
return
|
||||
|
||||
for index in range(limit):
|
||||
assert await access(index) == (f"token-{index + 1}", f"project-{index}")
|
||||
assert await access(0) == ("token-1", "project-0")
|
||||
assert len(oauth.calls) == limit
|
||||
assert await access(limit, inferred_project) == (f"token-{limit + 1}", f"project-{limit}")
|
||||
assert len(oauth.calls) == limit + 1
|
||||
if limit > (2 if inferred_project else 1):
|
||||
assert await access(0) == ("token-1", "project-0")
|
||||
assert len(oauth.calls) == limit + 1
|
||||
evicted: Final = 1 if limit > 1 else 0
|
||||
assert await access(evicted) == (f"token-{limit + 2}", f"project-{evicted}")
|
||||
assert len(oauth.calls) == limit + 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["sync", "async", "background"])
|
||||
@pytest.mark.skipif(DEFAULT_MAX_LRU_CACHE_SIZE < 2, reason="Eviction during refresh requires two retained entries")
|
||||
@pytest.mark.asyncio
|
||||
async def test_evicted_credentials_finish_refresh_without_being_readmitted(
|
||||
mode: Literal["sync", "async", "background"],
|
||||
) -> None:
|
||||
vertex: Final = VertexBase()
|
||||
limit: Final = DEFAULT_MAX_LRU_CACHE_SIZE
|
||||
now: Final = datetime(2025, 1, 1)
|
||||
refresh_started: Final = Event()
|
||||
release_refresh: Final = Event()
|
||||
donor: Final = _authorized_user_cache_credentials("donor")
|
||||
target: Final = _authorized_user_cache_credentials("target")
|
||||
with patch("google.auth._helpers.utcnow", return_value=now) as clock, responses.RequestsMock() as oauth:
|
||||
def token_response(request: requests.PreparedRequest) -> tuple[int, dict[str, str], str]:
|
||||
body: Final = request.body.decode() if isinstance(request.body, bytes) else request.body
|
||||
assert isinstance(body, str)
|
||||
is_target: Final = "refresh_token=target" in body
|
||||
if is_target and len(oauth.calls) == limit:
|
||||
refresh_started.set()
|
||||
assert release_refresh.wait(5), "OAuth refresh was not released"
|
||||
return 200, {}, json.dumps(
|
||||
{"access_token": f"token-{len(oauth.calls) + 1}", "expires_in": 3600 if is_target else 86400}
|
||||
)
|
||||
|
||||
oauth.add_callback(responses.POST, "https://oauth2.googleapis.com/token", callback=token_response)
|
||||
assert vertex.get_access_token(donor, None) == ("token-1", "donor")
|
||||
assert vertex.get_access_token(donor, None) == ("token-1", "donor")
|
||||
for index in range(limit - 2):
|
||||
assert vertex.get_access_token(donor, f"filler-{index}") == (f"token-{index + 2}", f"filler-{index}")
|
||||
assert vertex.get_access_token(target, "target") == (f"token-{limit}", "target")
|
||||
clock.return_value = now + timedelta(seconds=3550 if mode == "background" else 7200)
|
||||
refresh: Final = asyncio.create_task(
|
||||
asyncio.to_thread(vertex.get_access_token, target, "target")
|
||||
if mode == "sync"
|
||||
else vertex.get_access_token_async(target, "target")
|
||||
)
|
||||
try:
|
||||
assert await asyncio.to_thread(refresh_started.wait, 5), "OAuth refresh did not start"
|
||||
for index in range(limit - 2):
|
||||
assert vertex.get_access_token(donor, f"filler-{index}") == (f"token-{index + 2}", f"filler-{index}")
|
||||
assert vertex.get_access_token(donor, None) == ("token-1", "donor")
|
||||
assert len(vertex._credentials_project_mapping) == limit
|
||||
assert (target, "target") not in vertex._credentials_project_mapping
|
||||
finally:
|
||||
release_refresh.set()
|
||||
completed: Final = await refresh
|
||||
await asyncio.gather(*tuple(vertex._background_refresh_tasks.values()))
|
||||
expected_token: Final = f"token-{limit if mode == 'background' else limit + 1}"
|
||||
assert completed == (expected_token, "target")
|
||||
assert len(oauth.calls) == limit + 1
|
||||
assert (target, "target") not in vertex._credentials_project_mapping
|
||||
assert vertex.get_access_token(target, "target") == (f"token-{limit + 2}", "target")
|
||||
assert len(vertex._credentials_project_mapping) == limit
|
||||
|
|
|
|||
|
|
@ -9,9 +9,10 @@ if sys.version_info < (3, 11): # BaseExceptionGroup is a builtin only from 3.11
|
|||
from exceptiongroup import BaseExceptionGroup
|
||||
|
||||
import httpx
|
||||
import httpx2
|
||||
import pytest
|
||||
from mcp import MCPError
|
||||
from mcp.types import ErrorData
|
||||
from mcp.types import CONNECTION_CLOSED, ErrorData
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.exceptions import (
|
||||
MCPServerListError,
|
||||
|
|
@ -102,6 +103,68 @@ def test_unknown_exception_is_internal():
|
|||
assert classify_list_exception(ValueError("who knows")).tag == "internal"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sdk_wrapped", (False, True))
|
||||
@pytest.mark.parametrize("nested", (False, True))
|
||||
def test_sse_event_limit_is_a_safe_upstream_fault(sdk_wrapped: bool, nested: bool) -> None:
|
||||
limit: Final = 1024
|
||||
decoder_message: Final = f"Server-sent event exceeded the {limit} byte limit."
|
||||
failure: Final = (
|
||||
MCPError(code=CONNECTION_CLOSED, message=f"SSE stream failed: {decoder_message}")
|
||||
if sdk_wrapped else httpx2.SSEError(decoder_message)
|
||||
)
|
||||
wrapped: Final = BaseExceptionGroup("private-upstream-secret", [failure]) if nested else failure
|
||||
fault: Final = classify_list_exception(wrapped)
|
||||
message: Final = (
|
||||
f"Upstream MCP SSE event exceeded the {limit} byte per-event limit. "
|
||||
"Paginate tools/list responses so each event stays below the limit."
|
||||
)
|
||||
assert fault == ServerListFault(tag="upstream_error", sse_event_limit_bytes=limit)
|
||||
assert fault.public_message == message
|
||||
assert list_fault_http_status(fault) == 502
|
||||
assert outcome_wire_value(fault) == {
|
||||
"status": "upstream_error", "reason": "sse_event_too_large", "limit_bytes": limit, "message": message,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"failure",
|
||||
(
|
||||
MCPError(code=CONNECTION_CLOSED, message="Connection closed"),
|
||||
MCPError(code=-32603, message="SSE stream failed: Server-sent event exceeded the 1024 byte limit."),
|
||||
MCPError(code=CONNECTION_CLOSED, message="SSE stream failed: Server-sent event exceeded the 0 byte limit."),
|
||||
MCPError(code=CONNECTION_CLOSED, message="SSE stream failed: Server-sent event exceeded the -1 byte limit."),
|
||||
MCPError(code=CONNECTION_CLOSED, message=(
|
||||
"SSE stream failed: Server-sent event exceeded the 10000000000000000000 byte limit."
|
||||
)),
|
||||
MCPError(code=CONNECTION_CLOSED, message=(
|
||||
"SSE stream failed: Server-sent event exceeded the 1024 byte limit. private-secret"
|
||||
)),
|
||||
httpx2.SSEError("private-secret Server-sent event exceeded the 1024 byte limit."),
|
||||
httpx2.SSEError("Expected response with content type 'text/event-stream', got 'private-secret'."),
|
||||
RuntimeError("SSE stream failed: Server-sent event exceeded the 1024 byte limit."),
|
||||
),
|
||||
)
|
||||
def test_unrelated_errors_do_not_claim_an_sse_event_limit(failure: Exception) -> None:
|
||||
fault: Final = classify_list_exception(failure)
|
||||
assert fault.sse_event_limit_bytes is None
|
||||
assert fault.public_message is None
|
||||
assert "private-secret" not in str(outcome_wire_value(fault))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status_code", (401, 403))
|
||||
def test_upstream_auth_takes_precedence_over_sse_event_limit(status_code: int) -> None:
|
||||
request: Final = httpx2.Request("POST", "https://upstream.example/mcp")
|
||||
response: Final = httpx2.Response(status_code, request=request)
|
||||
failure: Final = BaseExceptionGroup("upstream errors", [
|
||||
httpx2.SSEError("Server-sent event exceeded the 1024 byte limit."),
|
||||
httpx2.HTTPStatusError("private-secret", request=request, response=response),
|
||||
])
|
||||
fault: Final = classify_list_exception(failure)
|
||||
assert fault.tag == ("auth_required" if status_code == 401 else "forbidden")
|
||||
assert list_fault_http_status(fault) == status_code
|
||||
assert fault.sse_event_limit_bytes is None
|
||||
|
||||
|
||||
def test_wire_value_carries_no_prose():
|
||||
fault = ServerListFault(tag="upstream_error", status_code=500)
|
||||
assert outcome_wire_value(fault) == {"status": "upstream_error", "http_status": 500}
|
||||
|
|
|
|||
|
|
@ -113,6 +113,9 @@ class MockBatcher:
|
|||
self.litellm_projecttable = _Table("project", self)
|
||||
self.litellm_endusertable = _Table("enduser", self)
|
||||
|
||||
def execute_raw(self, query: str, *args: object) -> None:
|
||||
self.calls.append({"table": "raw", "op": "execute_raw", "query": query, "args": args})
|
||||
|
||||
async def commit(self):
|
||||
self.committed = True
|
||||
return self.calls
|
||||
|
|
@ -2112,6 +2115,7 @@ def test_budget_cascade_writes_land_in_a_single_transaction(reset_budget_job, mo
|
|||
assert {(call["table"], call["op"]) for call in enduser_page.calls} == {("enduser", "update_many")}
|
||||
assert batcher.committed is True
|
||||
assert {(call["table"], call["op"]) for call in batcher.calls} == {
|
||||
("raw", "execute_raw"),
|
||||
("team_membership", "update_many"),
|
||||
("key", "update_many"),
|
||||
("org", "update_many"),
|
||||
|
|
@ -3393,6 +3397,35 @@ def test_budget_cascade_carries_overage_per_tier_when_rollover_enabled(
|
|||
counter_cache.in_memory_cache.delete_cache.assert_any_call(key="spend:team_member:member-1:team-1")
|
||||
|
||||
|
||||
def test_budget_cascade_floors_membership_total_spend_before_zeroing_spend(
|
||||
rollover_enabled: None,
|
||||
reset_budget_job: ResetBudgetJob,
|
||||
mock_prisma_client: MockPrismaClient,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Membership rows written before total_spend existed hold their history only in
|
||||
spend. The cascade lifts total_spend to spend for the expiring tier's rows, in the
|
||||
same batch and ahead of the zero/decrement, so the reset cannot erase it."""
|
||||
_make_counter_invalidation_job(monkeypatch)
|
||||
budget: Final = _budget_row(budget_id="budget-legacy", budget_duration="7d", max_budget=10.0)
|
||||
mock_prisma_client.data["budget"] = [budget]
|
||||
membership: Final = type(
|
||||
"Membership",
|
||||
(),
|
||||
{"user_id": "member-1", "team_id": "team-1", "spend": 15.0, "budget_id": "budget-legacy"},
|
||||
)
|
||||
mock_prisma_client.db.litellm_teammembership.set_find_many_results([membership])
|
||||
|
||||
asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
|
||||
|
||||
writes: Final = mock_prisma_client.db.batch_calls
|
||||
repair_index: Final = next(i for i, call in enumerate(writes) if call["op"] == "execute_raw")
|
||||
first_membership_reset: Final = next(i for i, call in enumerate(writes) if call["table"] == "team_membership")
|
||||
assert repair_index < first_membership_reset, writes
|
||||
assert writes[repair_index]["args"] == (["budget-legacy"],)
|
||||
assert mock_prisma_client.db.batchers[-1].committed
|
||||
|
||||
|
||||
def test_budget_cascade_carries_enduser_overage_when_rollover_enabled(
|
||||
rollover_enabled, reset_budget_job, mock_prisma_client, monkeypatch
|
||||
):
|
||||
|
|
|
|||
|
|
@ -1128,7 +1128,10 @@ async def test_commit_spend_updates_to_db_writes_team_member_spend_in_one_roster
|
|||
assert "members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', member.user_id))" in statement
|
||||
assert "ON CONFLICT (user_id, team_id) DO UPDATE" in statement
|
||||
assert 'spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend' in statement
|
||||
assert 'total_spend = "LiteLLM_TeamMembership".total_spend + EXCLUDED.total_spend' in statement
|
||||
assert (
|
||||
'total_spend = GREATEST("LiteLLM_TeamMembership".total_spend, "LiteLLM_TeamMembership".spend)'
|
||||
" + EXCLUDED.total_spend"
|
||||
) in statement
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@
|
|||
import asyncio
|
||||
import contextlib
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator
|
||||
from types import SimpleNamespace
|
||||
|
|
@ -35,6 +36,7 @@ from litellm.llms.openai.chat.guardrail_translation.handler import (
|
|||
from litellm.llms.openai.responses.guardrail_translation.handler import (
|
||||
OpenAIResponsesHandler,
|
||||
)
|
||||
from litellm.llms.pass_through.guardrail_translation.handler import PassThroughEndpointHandler
|
||||
from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import (
|
||||
MCPGuardrailTranslationHandler,
|
||||
)
|
||||
|
|
@ -3053,3 +3055,34 @@ async def test_pre_call_hook_skips_attachment_scan_when_guardrail_has_none():
|
|||
)
|
||||
|
||||
assert len(guardrail.apply_calls) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("hook", ["async_pre_call_hook", "async_moderation_hook"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_pass_through_guardrail_text_is_the_client_body_without_the_injected_key_metadata(
|
||||
monkeypatch: pytest.MonkeyPatch, hook: str
|
||||
) -> None:
|
||||
_patch_translation_mappings(monkeypatch, {CallTypes.pass_through: PassThroughEndpointHandler})
|
||||
guardrail: Final = RecordingGuardrail()
|
||||
body: Final = {
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "Say hi in one word"}],
|
||||
"guardrails": ["recording-guardrail"],
|
||||
}
|
||||
data: Final = {**body, "guardrail_to_apply": guardrail}
|
||||
user_api_key_dict: Final = UserAPIKeyAuth(
|
||||
token="hashed-key", user_email="someone@example.com", team_id="team-1", user_role="proxy_admin"
|
||||
)
|
||||
|
||||
if hook == "async_pre_call_hook":
|
||||
await UnifiedLLMGuardrails().async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict, cache=DualCache(), data=data, call_type=CallTypes.pass_through.value
|
||||
)
|
||||
else:
|
||||
await UnifiedLLMGuardrails().async_moderation_hook(
|
||||
data=data, user_api_key_dict=user_api_key_dict, call_type=CallTypes.pass_through.value
|
||||
)
|
||||
|
||||
texts: Final = [call["inputs"]["texts"] for call in guardrail.apply_calls]
|
||||
assert [[json.loads(text) for text in call_texts] for call_texts in texts] == [[body]], texts
|
||||
assert data["litellm_metadata"]["user_api_key_user_email"] == "someone@example.com", data["litellm_metadata"]
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ import contextvars
|
|||
import logging
|
||||
import sys
|
||||
import time
|
||||
from collections.abc import Iterator, Sequence
|
||||
from collections.abc import AsyncGenerator, Iterator, Sequence
|
||||
from contextlib import contextmanager
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Dict, Final, List, Optional
|
||||
|
|
@ -4644,6 +4644,140 @@ async def test_success_event_releases_parallel_slot_v3(monkeypatch):
|
|||
)
|
||||
|
||||
|
||||
_SSE_KEEPALIVE_CALLBACK_TIME: Final = datetime(2026, 1, 1)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("after_ping", [False, True], ids=["fast", "pinged"])
|
||||
@pytest.mark.parametrize("fails", [False, True], ids=["success", "failure"])
|
||||
@pytest.mark.parametrize("preadmitted", [False, True], ids=["producer-admission", "existing-admission"])
|
||||
async def test_sse_keepalive_releases_parallel_capacity(
|
||||
after_ping: bool, fails: bool, preadmitted: bool
|
||||
) -> None:
|
||||
from fastapi.responses import StreamingResponse
|
||||
|
||||
from litellm.proxy.common_request_processing import open_sse_before_first_byte
|
||||
from litellm.proxy.common_utils.sse_keepalive import SSE_COMMENT_PING_BYTES
|
||||
|
||||
cache: Final = DualCache()
|
||||
handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
|
||||
api_key: Final = hash_token("sk-keepalive-release")
|
||||
auth: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1)
|
||||
counter: Final = f"{{api_key:{api_key}}}:max_parallel_requests"
|
||||
|
||||
async def request(call_id: str) -> None:
|
||||
data: Final = {"model": "test-model", "litellm_call_id": call_id}
|
||||
upstream_answers: Final = asyncio.Event()
|
||||
if not after_ping:
|
||||
upstream_answers.set()
|
||||
|
||||
async def admit() -> None:
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=auth, cache=cache, data=data, call_type="acompletion"
|
||||
)
|
||||
|
||||
async def failure(exc: Exception) -> None:
|
||||
await handler.async_post_call_failure_hook(
|
||||
request_data=data, original_exception=exc, user_api_key_dict=auth
|
||||
)
|
||||
|
||||
async def stream() -> AsyncGenerator[bytes, None]:
|
||||
yield b"data: [DONE]\n\n"
|
||||
await handler.async_log_success_event(
|
||||
kwargs={
|
||||
"litellm_call_id": call_id,
|
||||
"standard_logging_object": {"metadata": {"user_api_key_hash": api_key}},
|
||||
},
|
||||
response_obj=ModelResponse(usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2)),
|
||||
start_time=_SSE_KEEPALIVE_CALLBACK_TIME,
|
||||
end_time=_SSE_KEEPALIVE_CALLBACK_TIME,
|
||||
)
|
||||
|
||||
async def produce() -> StreamingResponse:
|
||||
if not preadmitted:
|
||||
await admit()
|
||||
await upstream_answers.wait()
|
||||
if fails:
|
||||
raise HTTPException(status_code=502, detail="upstream failed")
|
||||
return StreamingResponse(stream())
|
||||
|
||||
if preadmitted:
|
||||
await admit()
|
||||
response_call: Final = open_sse_before_first_byte(
|
||||
produce(), 0.001 if after_ping else 60.0, on_late_failure=failure
|
||||
)
|
||||
if fails and not after_ping:
|
||||
with pytest.raises(HTTPException) as rejected:
|
||||
await response_call
|
||||
assert rejected.value.status_code == 502
|
||||
await failure(rejected.value)
|
||||
return
|
||||
response: Final = await response_call
|
||||
first: Final = await anext(response.body_iterator)
|
||||
if after_ping:
|
||||
upstream_answers.set()
|
||||
rest: Final = tuple([chunk async for chunk in response.body_iterator])
|
||||
assert first == (SSE_COMMENT_PING_BYTES if after_ping else b"data: [DONE]\n\n")
|
||||
if fails:
|
||||
error_frame: Final = rest[-2]
|
||||
assert isinstance(error_frame, bytes) and b"upstream failed" in error_frame
|
||||
assert rest[-1] == b"data: [DONE]\n\n"
|
||||
elif after_ping:
|
||||
assert rest[-1] == b"data: [DONE]\n\n"
|
||||
else:
|
||||
assert rest == ()
|
||||
|
||||
for index in range(3):
|
||||
await contextvars.Context().run(asyncio.create_task, request(f"keepalive-{index}"))
|
||||
assert handler._gauge_in_flight_from_cache_value(await cache.async_get_cache(counter)) == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sse_keepalive_rejects_overlap_and_reuses_completed_capacity() -> None:
|
||||
from fastapi.responses import StreamingResponse
|
||||
|
||||
from litellm.proxy.common_request_processing import open_sse_before_first_byte
|
||||
|
||||
cache: Final = DualCache()
|
||||
handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
|
||||
api_key: Final = hash_token("sk-keepalive-overlap")
|
||||
auth: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1)
|
||||
admitted: Final = asyncio.Event()
|
||||
finish: Final = asyncio.Event()
|
||||
|
||||
async def request(call_id: str) -> None:
|
||||
async def stream() -> AsyncGenerator[bytes, None]:
|
||||
await finish.wait()
|
||||
yield b"data: [DONE]\n\n"
|
||||
await handler.async_post_call_success_hook(
|
||||
data={"litellm_call_id": call_id}, user_api_key_dict=auth, response=ModelResponse()
|
||||
)
|
||||
|
||||
async def produce() -> StreamingResponse:
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=auth,
|
||||
cache=cache,
|
||||
data={"model": "test-model", "litellm_call_id": call_id},
|
||||
call_type="acompletion",
|
||||
)
|
||||
admitted.set()
|
||||
return StreamingResponse(stream())
|
||||
|
||||
response: Final = await open_sse_before_first_byte(produce(), 60.0)
|
||||
assert tuple([chunk async for chunk in response.body_iterator]) == (b"data: [DONE]\n\n",)
|
||||
|
||||
first: Final = contextvars.Context().run(asyncio.create_task, request("first"))
|
||||
await admitted.wait()
|
||||
try:
|
||||
with pytest.raises(HTTPException) as rejected:
|
||||
await contextvars.Context().run(asyncio.create_task, request("overlapping"))
|
||||
assert rejected.value.status_code == 429
|
||||
finally:
|
||||
finish.set()
|
||||
await first
|
||||
await contextvars.Context().run(asyncio.create_task, request("after-completion"))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_only_gauge_check_counts_without_acquiring_v3():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -12561,6 +12561,36 @@ async def test_team_member_me_returns_caller_membership(mock_db_client):
|
|||
assert call_kwargs["team_id"] == team_id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"spend, stored_total, reported_total",
|
||||
[(612.4, 148.1, 612.4), (12.25, 97.8, 97.8)],
|
||||
ids=["legacy_row_floors_to_cycle_spend", "healthy_row_keeps_its_total"],
|
||||
)
|
||||
async def test_team_member_me_reports_total_spend_no_lower_than_cycle_spend(
|
||||
mock_db_client, spend, stored_total, reported_total
|
||||
):
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy.management_endpoints.team_endpoints import team_member_me
|
||||
|
||||
team_id = "team-me-floor"
|
||||
caller_id = "alice@example.com"
|
||||
team = _build_team_for_me(team_id, [{"user_id": caller_id, "user_email": None, "role": "user"}])
|
||||
membership = _build_membership_for_me(caller_id, team_id, spend=spend).model_copy(
|
||||
update={"total_spend": stored_total}
|
||||
)
|
||||
p_team, p_membership, p_user = _patch_member_me_helpers(team=team, membership=membership, user=None)
|
||||
with p_team, p_membership, p_user:
|
||||
response = await team_member_me(
|
||||
http_request=MagicMock(spec=Request),
|
||||
team_id=team_id,
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id=caller_id),
|
||||
)
|
||||
|
||||
assert (response.spend, response.total_spend) == (spend, reported_total)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_member_me_matches_email_only_member(mock_db_client):
|
||||
"""
|
||||
|
|
@ -16608,7 +16638,7 @@ async def test_reset_team_member_spend_fn_success(monkeypatch):
|
|||
litellm_budget_table=LiteLLM_BudgetTable(budget_id="b1", max_budget=50.0),
|
||||
)
|
||||
mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row)
|
||||
mock_prisma_client.db.litellm_teammembership.update = AsyncMock(return_value=membership_row)
|
||||
mock_prisma_client.db.execute_raw = AsyncMock(return_value=1)
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache)
|
||||
|
|
@ -16631,15 +16661,76 @@ async def test_reset_team_member_spend_fn_success(monkeypatch):
|
|||
assert response["spend"] == 0.0
|
||||
assert response["previous_spend"] == 10.0
|
||||
assert response["max_budget"] == 50.0
|
||||
mock_prisma_client.db.litellm_teammembership.update.assert_awaited_once_with(
|
||||
where={"user_id_team_id": {"user_id": "member-1", "team_id": "team-1"}},
|
||||
data={"spend": 0.0},
|
||||
)
|
||||
statement, reset_user_id, reset_team_id, reset_to = mock_prisma_client.db.execute_raw.await_args.args
|
||||
assert (reset_user_id, reset_team_id, reset_to) == ("member-1", "team-1", 0.0)
|
||||
assert "total_spend = GREATEST(total_spend, spend)" in statement
|
||||
assert "spend >= $3" in statement
|
||||
assert await real_cache.async_get_cache(key="team-1_member-1") is None
|
||||
assert await real_cache.async_get_cache(key="team_membership:member-1:team-1") is None
|
||||
assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 0.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_team_member_spend_fn_stale_spend_returns_409_and_keeps_caches(monkeypatch):
|
||||
"""The reset validates reset_to against a spend it read a moment ago. If the row moved
|
||||
underneath it (a flush landed, another admin reset first) the guarded UPDATE matches no
|
||||
row, and the endpoint must say so instead of reporting a reset that never happened."""
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
mock_prisma_client = MagicMock()
|
||||
real_cache = UserApiKeyCache()
|
||||
await real_cache.async_set_cache(key="team_membership:member-1:team-1", value="cached-membership")
|
||||
real_spend_counter_cache = DualCache()
|
||||
real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:member-1:team-1", value=999.0)
|
||||
|
||||
membership_row = LiteLLM_TeamMembership(user_id="member-1", team_id="team-1", spend=10.0)
|
||||
mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row)
|
||||
mock_prisma_client.db.execute_raw = AsyncMock(return_value=0)
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache)
|
||||
|
||||
with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
|
||||
"litellm.proxy.management_endpoints.team_endpoints.get_team_object",
|
||||
AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")),
|
||||
), pytest.raises(HTTPException) as exc_info:
|
||||
await reset_team_member_spend_fn(
|
||||
team_id="team-1",
|
||||
user_id="member-1",
|
||||
data=ResetSpendRequest(reset_to=5.0),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user"
|
||||
),
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 409
|
||||
assert await real_cache.async_get_cache(key="team_membership:member-1:team-1") == "cached-membership"
|
||||
assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 999.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_all_team_memberships_floors_total_spend_at_current_cycle_spend():
|
||||
"""Rows written before total_spend existed carry their history only in spend. /team/info
|
||||
and /team/list both read through this helper, so a legacy row reports its lifetime spend
|
||||
as at least the current cycle, while a healthy row keeps its own total."""
|
||||
from litellm.proxy.management_endpoints.team_endpoints import get_all_team_memberships
|
||||
|
||||
legacy_row = MagicMock()
|
||||
legacy_row.model_dump = lambda: {"user_id": "legacy", "team_id": "team-1", "spend": 612.4, "total_spend": 148.1}
|
||||
healthy_row = MagicMock()
|
||||
healthy_row.model_dump = lambda: {"user_id": "healthy", "team_id": "team-1", "spend": 12.25, "total_spend": 97.8}
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_teammembership.find_many = AsyncMock(return_value=[legacy_row, healthy_row])
|
||||
|
||||
legacy, healthy = await get_all_team_memberships(mock_prisma_client, ["team-1"])
|
||||
|
||||
assert (legacy.spend, legacy.total_spend) == (612.4, 612.4)
|
||||
assert (healthy.spend, healthy.total_spend) == (12.25, 97.8)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_team_member_spend_fn_membership_not_found(monkeypatch):
|
||||
mock_prisma_client = MagicMock()
|
||||
|
|
@ -16755,7 +16846,7 @@ async def test_reset_team_member_spend_fn_proxy_admin_can_reset_own_spend(monkey
|
|||
|
||||
membership_row = LiteLLM_TeamMembership(user_id="admin-user", team_id="team-1", spend=10.0)
|
||||
mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row)
|
||||
mock_prisma_client.db.litellm_teammembership.update = AsyncMock(return_value=membership_row)
|
||||
mock_prisma_client.db.execute_raw = AsyncMock(return_value=1)
|
||||
|
||||
with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
|
||||
"litellm.proxy.management_endpoints.team_endpoints.get_team_object",
|
||||
|
|
@ -17994,6 +18085,7 @@ async def test_update_team_team_admin_resending_stored_member_defaults_is_not_an
|
|||
disable_audit_logging_for_mocked_team,
|
||||
):
|
||||
import contextlib
|
||||
|
||||
from litellm.models.budget import LiteLLM_BudgetTable
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
|
|
|
|||
|
|
@ -2,23 +2,35 @@ import json
|
|||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
from pydantic import BaseModel, ValidationError
|
||||
|
||||
from litellm.proxy._types import (
|
||||
ROLES_WITHIN_ORG,
|
||||
BudgetNewRequest,
|
||||
ChangePasswordRequest,
|
||||
GenerateKeyRequest,
|
||||
GenerateKeyResponse,
|
||||
KeyRequest,
|
||||
LiteLLM_AuditLogs,
|
||||
LiteLLM_BudgetTable,
|
||||
LiteLLM_EndUserTable,
|
||||
LiteLLM_TeamMembership,
|
||||
LiteLLM_TeamTable,
|
||||
LiteLLM_UserTable,
|
||||
LitellmUserRoles,
|
||||
NewCustomerRequest,
|
||||
NewMCPServerRequest,
|
||||
NewTeamRequest,
|
||||
NewUserRequest,
|
||||
NewUserResponse,
|
||||
OrganizationMemberUpdateRequest,
|
||||
RegenerateKeyRequest,
|
||||
ResetSpendRequest,
|
||||
UpdateKeyRequest,
|
||||
UpdateMCPServerRequest,
|
||||
UpdateTeamRequest,
|
||||
UpdateUserRequest,
|
||||
UpdateUserRequestNoUserIDorEmail,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
|
||||
|
|
@ -213,6 +225,75 @@ def test_a_max_budget_that_was_actually_sent_is_kept(sent):
|
|||
assert GenerateKeyRequest(max_budget=sent).max_budget == sent
|
||||
|
||||
|
||||
RATE_LIMIT_REQUESTS: Final = (
|
||||
(GenerateKeyRequest, {}),
|
||||
(RegenerateKeyRequest, {}),
|
||||
(UpdateKeyRequest, {"key": "sk-test"}),
|
||||
(NewUserRequest, {}),
|
||||
(UpdateUserRequestNoUserIDorEmail, {}),
|
||||
(UpdateUserRequest, {"user_id": "user-test"}),
|
||||
(NewTeamRequest, {}),
|
||||
(UpdateTeamRequest, {"team_id": "team-test"}),
|
||||
(BudgetNewRequest, {}),
|
||||
(NewCustomerRequest, {"user_id": "customer-test"}),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("request_model", "identifiers"), RATE_LIMIT_REQUESTS)
|
||||
@pytest.mark.parametrize("field", ("rpm_limit", "tpm_limit"))
|
||||
@pytest.mark.parametrize("value", (-1, -100))
|
||||
def test_management_requests_reject_negative_rate_limits(
|
||||
request_model: type[BaseModel], identifiers: dict[str, str], field: str, value: int
|
||||
) -> None:
|
||||
with pytest.raises(ValidationError) as error:
|
||||
request_model.model_validate({**identifiers, field: value})
|
||||
assert error.value.errors()[0]["loc"] == (field,)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("request_model", "identifiers"), RATE_LIMIT_REQUESTS)
|
||||
@pytest.mark.parametrize("value", (None, 0, 100))
|
||||
def test_management_requests_preserve_explicit_rate_limits(
|
||||
request_model: type[BaseModel], identifiers: dict[str, str], value: int | None
|
||||
) -> None:
|
||||
limits: Final = {"rpm_limit": value, "tpm_limit": value}
|
||||
request: Final = request_model.model_validate({**identifiers, **limits})
|
||||
assert request.model_dump(include={"rpm_limit", "tpm_limit"}, exclude_unset=True) == limits
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("request_model", "identifiers"), RATE_LIMIT_REQUESTS)
|
||||
def test_management_requests_leave_omitted_rate_limits_unset(
|
||||
request_model: type[BaseModel], identifiers: dict[str, str]
|
||||
) -> None:
|
||||
request: Final = request_model.model_validate({**identifiers})
|
||||
assert request.model_dump(include={"rpm_limit", "tpm_limit"}) == {"rpm_limit": None, "tpm_limit": None}
|
||||
assert request.model_dump(include={"rpm_limit", "tpm_limit"}, exclude_unset=True) == {}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("response_model", "identifiers"), (
|
||||
(GenerateKeyResponse, {"key": "sk-test"}),
|
||||
(NewUserResponse, {"key": "sk-test"}),
|
||||
(LiteLLM_TeamTable, {"team_id": "team-test"}),
|
||||
(LiteLLM_UserTable, {"user_id": "user-test"}),
|
||||
(LiteLLM_BudgetTable, {"budget_id": "budget-test"}),
|
||||
))
|
||||
def test_existing_negative_rate_limits_remain_readable(
|
||||
response_model: type[BaseModel], identifiers: dict[str, str]
|
||||
) -> None:
|
||||
limits: Final = {"rpm_limit": -1, "tpm_limit": -100}
|
||||
response: Final = response_model.model_validate({**identifiers, **limits})
|
||||
assert response.model_dump(include={"rpm_limit", "tpm_limit"}) == limits
|
||||
|
||||
|
||||
def test_existing_customer_negative_rate_limits_remain_readable() -> None:
|
||||
customer: Final = LiteLLM_EndUserTable.model_validate({
|
||||
"user_id": "customer-test", "blocked": False,
|
||||
"litellm_budget_table": {"rpm_limit": -1, "tpm_limit": -100},
|
||||
})
|
||||
assert customer.litellm_budget_table is not None
|
||||
assert customer.litellm_budget_table.rpm_limit == -1
|
||||
assert customer.litellm_budget_table.tpm_limit == -100
|
||||
|
||||
|
||||
USER_IDENTIFIER_CASES = (
|
||||
({"user_id": "user-1"}, True),
|
||||
({"user_email": "user@example.com"}, True),
|
||||
|
|
|
|||
|
|
@ -15207,21 +15207,23 @@ async def test_token_counter_keeps_the_event_loop_free_during_a_huggingface_coun
|
|||
assert_loop_stayed_free(took, lags)
|
||||
|
||||
|
||||
async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeypatch):
|
||||
from litellm.rust_bridge._native import Tokenizer
|
||||
async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from queue import SimpleQueue
|
||||
from threading import Thread, current_thread
|
||||
|
||||
from tokenizers import Tokenizer
|
||||
|
||||
from litellm import Router
|
||||
from tests.unit.litellm_core_utils.event_loop_lag import assert_loop_stayed_free, timed_with_loop_lags
|
||||
|
||||
claude_tokenizer: Final = litellm.utils.select_tokenizer("claude-fable-5")["tokenizer"]
|
||||
claude_tokenizer: Final = Tokenizer.from_str(litellm.utils.claude_json_str)
|
||||
loop_thread: Final = current_thread()
|
||||
loader_threads: Final[SimpleQueue[Thread]] = SimpleQueue()
|
||||
|
||||
class SlowHubTokenizer:
|
||||
@staticmethod
|
||||
def from_pretrained(identifier: str, revision: str = "main", token: str | None = None) -> Tokenizer:
|
||||
time.sleep(0.3)
|
||||
return claude_tokenizer
|
||||
def from_pretrained(identifier: str, revision: str = "main", token: str | None = None) -> Tokenizer:
|
||||
loader_threads.put(current_thread())
|
||||
return claude_tokenizer
|
||||
|
||||
monkeypatch.setattr("litellm.rust_bridge.tokenizer.from_pretrained", SlowHubTokenizer.from_pretrained)
|
||||
monkeypatch.setattr("tokenizers.Tokenizer", types.SimpleNamespace(from_pretrained=from_pretrained))
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.llm_router",
|
||||
Router(
|
||||
|
|
@ -15237,15 +15239,16 @@ async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeyp
|
|||
),
|
||||
)
|
||||
|
||||
response, took, lags = await timed_with_loop_lags(
|
||||
lambda: proxy_server_module.token_counter(
|
||||
litellm.utils._select_custom_tokenizer_helper.cache_clear()
|
||||
try:
|
||||
response: Final = await proxy_server_module.token_counter(
|
||||
TokenCountRequest(model="self-hosted", prompt="count me off the loop")
|
||||
)
|
||||
)
|
||||
|
||||
assert response.tokenizer_type == "huggingface_tokenizer"
|
||||
assert response.total_tokens > 0
|
||||
assert_loop_stayed_free(took, lags)
|
||||
assert loader_threads.get_nowait() is not loop_thread
|
||||
assert response.tokenizer_type == "huggingface_tokenizer"
|
||||
assert response.total_tokens > 0
|
||||
finally:
|
||||
litellm.utils._select_custom_tokenizer_helper.cache_clear()
|
||||
|
||||
|
||||
async def test_token_counter_loads_a_custom_tokenizer_once_per_identifier_revision_and_token(monkeypatch):
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ import pytest
|
|||
|
||||
from litellm.repositories.unit_of_work import (
|
||||
LinkedSpendResetWrites,
|
||||
SpendResetWrites,
|
||||
budget_cascade_unit_of_work,
|
||||
enduser_spend_reset_unit_of_work,
|
||||
spend_reset_unit_of_work,
|
||||
|
|
@ -39,6 +40,9 @@ class FakeBatch:
|
|||
self.litellm_projecttable = FakeBatchTable("litellm_projecttable", self.calls)
|
||||
self.litellm_endusertable = FakeBatchTable("litellm_endusertable", self.calls)
|
||||
|
||||
def execute_raw(self, query: str, *args: object) -> None:
|
||||
self.calls.append(("execute_raw", {"query": query}, {"args": args}))
|
||||
|
||||
async def commit(self) -> None:
|
||||
self.commit_count += 1
|
||||
|
||||
|
|
@ -166,7 +170,7 @@ async def test_every_cascade_dependent_writes_to_its_own_table_on_the_one_batch(
|
|||
async with budget_cascade_unit_of_work(_new_batch) as uow:
|
||||
writes = [getattr(uow, field.name) for field in fields(uow)]
|
||||
for write in writes:
|
||||
if isinstance(write, LinkedSpendResetWrites):
|
||||
if isinstance(write, SpendResetWrites):
|
||||
write.queue_spend_zero(where={"budget_id": "budget-1"})
|
||||
else:
|
||||
write.queue_window_advance(budget_id="budget-1", budget_reset_at=reset_at)
|
||||
|
|
@ -192,3 +196,18 @@ async def test_budget_cascade_raising_inside_block_skips_commit():
|
|||
await _blow_up_mid_transaction()
|
||||
|
||||
assert batch.commit_count == 0
|
||||
|
||||
|
||||
async def test_membership_total_spend_repair_is_queued_on_the_cascade_batch_ahead_of_the_reset():
|
||||
batch = FakeBatch()
|
||||
|
||||
async with budget_cascade_unit_of_work(lambda: batch) as uow:
|
||||
uow.team_memberships.queue_total_spend_repair(["budget-1", "budget-2"])
|
||||
uow.team_memberships.queue_spend_zero(where={"budget_id": "budget-1"})
|
||||
assert batch.commit_count == 0
|
||||
|
||||
assert batch.commit_count == 1
|
||||
repair, reset = batch.calls
|
||||
assert repair[0] == "execute_raw"
|
||||
assert repair[2] == {"args": (["budget-1", "budget-2"],)}
|
||||
assert reset == ("litellm_teammembership.update_many", {"budget_id": "budget-1"}, {"spend": 0})
|
||||
|
|
|
|||
|
|
@ -1065,6 +1065,11 @@ class TestCustomDimensions:
|
|||
"create table widgets (id integer); ALTER TABLE widgets ADD label text;",
|
||||
id="regex",
|
||||
),
|
||||
pytest.param(
|
||||
{"patterns": [r"\b(?P<verb>create|alter)\s{1,4}?table\b"]},
|
||||
"CREATE TABLE widgets (id integer)",
|
||||
id="named-group-alternation-lazy-repeat",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_custom_dimension_changes_only_matching_requests(
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import base64
|
||||
import binascii
|
||||
import json
|
||||
from collections import OrderedDict
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
|
|
@ -120,7 +121,7 @@ def manager(monkeypatch: pytest.MonkeyPatch) -> GoogleSecretManager:
|
|||
monkeypatch.setattr(litellm, "secret_manager_client", None)
|
||||
monkeypatch.setattr(litellm, "_key_management_system", None)
|
||||
monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings())
|
||||
monkeypatch.setattr(litellm.vertex_chat_completion, "_credentials_project_mapping", {})
|
||||
monkeypatch.setattr(litellm.vertex_chat_completion, "_credentials_project_mapping", OrderedDict())
|
||||
monkeypatch.setattr(google.auth, "default", _default_credentials)
|
||||
monkeypatch.setenv("GOOGLE_SECRET_MANAGER_PROJECT_ID", PROJECT_ID)
|
||||
monkeypatch.delenv("GOOGLE_SECRET_MANAGER_ALWAYS_READ_SECRET_MANAGER", raising=False)
|
||||
|
|
|
|||
|
|
@ -455,30 +455,23 @@ def is_active_priced_mistral_chat_row(name: str, entry: Mapping[str, object]) ->
|
|||
)
|
||||
|
||||
|
||||
def cache_read_is_tenth_of_input(entry: Mapping[str, object]) -> bool:
|
||||
def cache_read_is_below_input(entry: Mapping[str, object]) -> bool:
|
||||
cache_read: Final = entry.get("cache_read_input_token_cost")
|
||||
input_cost: Final = entry.get("input_cost_per_token")
|
||||
return (
|
||||
isinstance(cache_read, float)
|
||||
and isinstance(input_cost, (int, float))
|
||||
and 0 < cache_read < input_cost
|
||||
and cache_read == pytest.approx(input_cost / 10)
|
||||
)
|
||||
return isinstance(cache_read, float) and isinstance(input_cost, (int, float)) and 0 < cache_read < input_cost
|
||||
|
||||
|
||||
@pytest.mark.parametrize("path", (PRICES_PATH, BACKUP_PRICES_PATH), ids=("main", "backup"))
|
||||
def test_active_mistral_chat_rows_price_cache_reads_below_input(path: Path):
|
||||
"""A Mistral chat row without a cache-read rate bills cached prompt tokens at zero, so every
|
||||
active priced row must carry one, and it must be cheaper than a fresh input token. Mistral
|
||||
bills cached tokens at 10% of the input price for every model (docs.mistral.ai/studio/
|
||||
conversations/advanced/prompt-caching, read 2026-09-18), so the ratio is checked as well."""
|
||||
active priced row must carry one, and it must be cheaper than a fresh input token."""
|
||||
rows: Mapping[str, object] = json.loads(path.read_text())
|
||||
drifted: Final = [
|
||||
f"{name}: cache_read={entry.get('cache_read_input_token_cost')} input={entry.get('input_cost_per_token')}"
|
||||
for name, entry in rows.items()
|
||||
if isinstance(entry, dict)
|
||||
and is_active_priced_mistral_chat_row(name, entry)
|
||||
and not cache_read_is_tenth_of_input(entry)
|
||||
and not cache_read_is_below_input(entry)
|
||||
]
|
||||
assert drifted == []
|
||||
|
||||
|
|
|
|||
|
|
@ -1086,6 +1086,7 @@ _COST_MAP_ROW_SCHEMA: Final = {
|
|||
"chat",
|
||||
"completion",
|
||||
"container",
|
||||
"decisions",
|
||||
"image_edit",
|
||||
"embedding",
|
||||
"evaluation",
|
||||
|
|
|
|||
|
|
@ -1,14 +1,20 @@
|
|||
import importlib
|
||||
from typing import Final
|
||||
from uuid import NAMESPACE_DNS, UUID, uuid5 as stdlib_uuid5
|
||||
|
||||
from litellm._uuid import uuid, uuid4
|
||||
|
||||
|
||||
def test_uses_fastuuid_and_uuid4_works():
|
||||
mod = importlib.import_module("litellm._uuid")
|
||||
fastuuid_mod = importlib.import_module("fastuuid")
|
||||
assert hasattr(mod, "uuid4")
|
||||
assert hasattr(mod, "uuid")
|
||||
assert mod.uuid is fastuuid_mod
|
||||
def test_uuid4_returns_canonical_random_uuid() -> None:
|
||||
for value in (uuid4(), uuid.uuid4()):
|
||||
encoded: Final = str(value)
|
||||
parsed: Final = UUID(encoded)
|
||||
assert parsed.version == 4
|
||||
assert encoded == str(parsed)
|
||||
assert value.bytes == parsed.bytes
|
||||
|
||||
# Ensure uuid4 returns something that looks like a UUID string
|
||||
val = str(mod.uuid4())
|
||||
assert isinstance(val, str)
|
||||
assert len(val) >= 8
|
||||
|
||||
def test_uuid5_accepts_byte_names_for_stable_ids() -> None:
|
||||
namespace: Final = uuid.UUID(str(NAMESPACE_DNS))
|
||||
name: Final = "guardrail-\u00e9"
|
||||
generated: Final = uuid.uuid5(namespace, name.encode("utf-8"))
|
||||
assert str(generated) == str(stdlib_uuid5(NAMESPACE_DNS, name))
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@
|
|||
# UI container — Next.js static export served by nginx.
|
||||
|
||||
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
|
||||
ARG NGINX_VERSION=1.31.5-alpine3.24@sha256:34f40471dea485273c5e2a04dd5e97a682332ceb4a9adecd67de450dcb2fb390
|
||||
ARG NGINX_VERSION=1.31.8-alpine3.24@sha256:1ef7ca94d7879df82421428ce13bd1eaa412404953859cf7a5b6512c25e67a60
|
||||
|
||||
# ---------- builder ----------
|
||||
FROM ${UI_BUILD_IMAGE} AS builder
|
||||
|
|
@ -27,6 +27,8 @@ RUN npm run build
|
|||
# ---------- runtime ----------
|
||||
FROM nginx:${NGINX_VERSION} AS runtime
|
||||
|
||||
RUN apk upgrade --no-cache brotli-libs
|
||||
|
||||
# Drop the upstream default :80 server; we own the config.
|
||||
RUN rm -f /etc/nginx/conf.d/default.conf
|
||||
|
||||
|
|
|
|||
|
|
@ -2175,11 +2175,6 @@
|
|||
"count": 1
|
||||
}
|
||||
},
|
||||
"src/utils/dataUtils.test.ts": {
|
||||
"max-nested-callbacks": {
|
||||
"count": 1
|
||||
}
|
||||
},
|
||||
"tests/CreateKeyPage.expiredToken.test.tsx": {
|
||||
"@typescript-eslint/no-require-imports": {
|
||||
"count": 3
|
||||
|
|
|
|||
|
|
@ -80,6 +80,42 @@ describe("GeneralSettings General tab", () => {
|
|||
vi.mocked(deleteConfigFieldSetting).mockClear();
|
||||
});
|
||||
|
||||
it.each([
|
||||
["maximum_spend_logs_retention_period", null, "7d"],
|
||||
["maximum_health_check_retention_period", "30d", "14d"],
|
||||
["maximum_daily_tag_spend_retention_period", undefined, "90d"],
|
||||
["maximum_spend_logs_cleanup_run_budget", "5m", "10m"],
|
||||
["maximum_spend_logs_cleanup_batch_timeout", "", "45s"],
|
||||
["mcp_client_id_header", "x-client-id", "X-MCP-Client.ID"],
|
||||
])("displays, updates and resets the String setting %s", async (fieldName, initialValue, updatedValue) => {
|
||||
vi.mocked(getGeneralSettingsCall).mockResolvedValue([
|
||||
...SETTINGS_FIXTURE,
|
||||
{
|
||||
field_name: fieldName,
|
||||
field_type: "String",
|
||||
field_value: initialValue,
|
||||
field_description: "string setting",
|
||||
stored_in_db: initialValue != null,
|
||||
field_default_value: null,
|
||||
},
|
||||
]);
|
||||
const user = userEvent.setup();
|
||||
renderWithProviders(<GeneralSettings accessToken="token" userRole="Admin" userID="user" />);
|
||||
await user.click(screen.getByRole("tab", { name: "General" }));
|
||||
const input = await screen.findByRole("textbox", { name: fieldName });
|
||||
expect(input).toHaveValue(initialValue ?? "");
|
||||
fireEvent.change(input, { target: { value: updatedValue } });
|
||||
const row = screen.getByRole("row", { name: new RegExp(fieldName) });
|
||||
await user.click(within(row).getByRole("button", { name: "Update" }));
|
||||
expect(vi.mocked(updateConfigFieldSetting).mock.calls).toEqual([["token", fieldName, updatedValue]]);
|
||||
expect(within(row).getByText("In DB")).toBeInTheDocument();
|
||||
const resetIcon = row.querySelectorAll("td")[3].querySelector("svg");
|
||||
expect(resetIcon).not.toBeNull();
|
||||
await user.click(resetIcon as unknown as Element);
|
||||
expect(vi.mocked(deleteConfigFieldSetting).mock.calls).toEqual([["token", fieldName]]);
|
||||
expect(input).toHaveValue("");
|
||||
});
|
||||
|
||||
it("updates max_ui_session_budget with its own value, not the value at its filtered index", async () => {
|
||||
const user = userEvent.setup();
|
||||
renderWithProviders(<GeneralSettings accessToken="token" userRole="Admin" userID="user" />);
|
||||
|
|
|
|||
|
|
@ -68,6 +68,15 @@ const SettingValueEditor: React.FC<{
|
|||
/>
|
||||
);
|
||||
}
|
||||
if (setting.field_type === "String") {
|
||||
return (
|
||||
<Input
|
||||
aria-label={setting.field_name}
|
||||
value={setting.field_value ?? ""}
|
||||
onChange={(event) => onChange(setting.field_name, event.target.value)}
|
||||
/>
|
||||
);
|
||||
}
|
||||
if (setting.field_type === "Integer") {
|
||||
return (
|
||||
<Input
|
||||
|
|
|
|||
|
|
@ -1,10 +1,16 @@
|
|||
import { copyToClipboard } from "@/utils/dataUtils";
|
||||
import { describe, it, expect, vi, beforeEach } from "vitest";
|
||||
import userEvent from "@testing-library/user-event";
|
||||
import { render, screen, waitFor, within } from "@testing-library/react";
|
||||
import { render, screen, within } from "@testing-library/react";
|
||||
import moment from "moment";
|
||||
import { AuditLogDrawer } from "./AuditLogDrawer";
|
||||
import { AuditLogEntry } from "../AuditLogsTableColumns";
|
||||
|
||||
vi.mock("@/utils/dataUtils", async (importOriginal) => ({
|
||||
...(await importOriginal<typeof import("@/utils/dataUtils")>()),
|
||||
copyToClipboard: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock("../../../common_components/DefaultProxyAdminTag", () => ({
|
||||
default: ({ userId }: { userId: string }) => <span>{userId}</span>,
|
||||
}));
|
||||
|
|
@ -120,16 +126,17 @@ describe("AuditLogDrawer", () => {
|
|||
expect(screen.queryByText(/"spend"/)).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should copy the json of a block to the clipboard", async () => {
|
||||
it.each([true, false])("confirms copied JSON only when copying succeeds (%s)", async (success) => {
|
||||
const user = userEvent.setup();
|
||||
const writeText = vi.fn().mockResolvedValue(undefined);
|
||||
Object.defineProperty(navigator, "clipboard", { value: { writeText }, configurable: true });
|
||||
Object.defineProperty(window, "isSecureContext", { value: true, configurable: true });
|
||||
|
||||
Object.defineProperty(navigator, "clipboard", { value: undefined, configurable: true });
|
||||
vi.mocked(copyToClipboard).mockResolvedValue(success);
|
||||
render(<AuditLogDrawer {...defaultProps} />);
|
||||
await user.click(within(blockNamed("Before")).getByTitle("Copy JSON"));
|
||||
|
||||
await waitFor(() => expect(writeText).toHaveBeenCalledWith(JSON.stringify({ max_budget: 10 }, null, 2)));
|
||||
const button = within(blockNamed("Before")).getByRole("button", { name: "Copy JSON" });
|
||||
await user.click(button);
|
||||
|
||||
expect(copyToClipboard).toHaveBeenCalledWith(JSON.stringify({ max_budget: 10 }, null, 2));
|
||||
expect(button.querySelector(".lucide-check") !== null).toBe(success);
|
||||
});
|
||||
|
||||
describe("width", () => {
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import { copyToClipboard } from "@/utils/dataUtils";
|
||||
import { Check, Copy } from "lucide-react";
|
||||
import { useState, useCallback } from "react";
|
||||
import moment from "moment";
|
||||
|
|
@ -27,25 +28,9 @@ function CopyableJsonBlock({ label, value }: { label: string; value: Record<stri
|
|||
const [copied, setCopied] = useState(false);
|
||||
|
||||
const handleCopy = useCallback(async () => {
|
||||
try {
|
||||
const text = JSON.stringify(value, null, 2);
|
||||
if (navigator.clipboard && window.isSecureContext) {
|
||||
await navigator.clipboard.writeText(text);
|
||||
} else {
|
||||
const el = document.createElement("textarea");
|
||||
el.value = text;
|
||||
el.style.position = "fixed";
|
||||
el.style.opacity = "0";
|
||||
document.body.appendChild(el);
|
||||
el.focus();
|
||||
el.select();
|
||||
document.execCommand("copy");
|
||||
document.body.removeChild(el);
|
||||
}
|
||||
if (await copyToClipboard(JSON.stringify(value, null, 2))) {
|
||||
setCopied(true);
|
||||
setTimeout(() => setCopied(false), 2000);
|
||||
} catch (e) {
|
||||
console.error("Copy failed:", e);
|
||||
}
|
||||
}, [value]);
|
||||
|
||||
|
|
|
|||
|
|
@ -6,6 +6,12 @@ import { describe, expect, it, vi } from "vitest";
|
|||
import { render } from "../../../../tests/test-utils";
|
||||
import type { LogEntry } from "../types";
|
||||
import { DrawerHeader } from "./DrawerHeader";
|
||||
import { copyToClipboard } from "@/utils/dataUtils";
|
||||
|
||||
vi.mock("@/utils/dataUtils", async (importOriginal) => ({
|
||||
...(await importOriginal<typeof import("@/utils/dataUtils")>()),
|
||||
copyToClipboard: vi.fn(),
|
||||
}));
|
||||
|
||||
const logEntry = (overrides: Partial<LogEntry>): LogEntry =>
|
||||
({
|
||||
|
|
@ -118,3 +124,22 @@ describe("DrawerHeader back link", () => {
|
|||
expect(screen.queryByRole("button", { name: /^Back to/ })).not.toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
||||
describe("DrawerHeader copying", () => {
|
||||
it.each([true, false])("confirms both ids only when copying succeeds (%s)", async (success) => {
|
||||
const user = userEvent.setup();
|
||||
Object.defineProperty(navigator, "clipboard", { value: undefined, configurable: true });
|
||||
vi.mocked(copyToClipboard).mockReset().mockResolvedValue(success);
|
||||
renderHeader(logEntry({ request_id: "request-copy", litellm_call_id: "call-copy" }), false);
|
||||
|
||||
const requestButton = screen.getByRole("button", { name: "Copy Request ID" });
|
||||
const callButton = screen.getByRole("button", { name: "Copy x-litellm-call-id" });
|
||||
await user.click(requestButton);
|
||||
await user.click(callButton);
|
||||
|
||||
expect(copyToClipboard).toHaveBeenNthCalledWith(1, "request-copy");
|
||||
expect(copyToClipboard).toHaveBeenNthCalledWith(2, "call-copy");
|
||||
expect(requestButton).toHaveAccessibleName(success ? "Copied!" : "Copy Request ID");
|
||||
expect(callButton).toHaveAccessibleName(success ? "Copied!" : "Copy x-litellm-call-id");
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import { copyToClipboard } from "@/utils/dataUtils";
|
||||
import { type ReactNode, useState } from "react";
|
||||
import { ArrowLeft, Check, ChevronDown, ChevronUp, Copy, X } from "lucide-react";
|
||||
import moment from "moment";
|
||||
|
|
@ -183,12 +184,9 @@ function CopyableId({
|
|||
const [copied, setCopied] = useState(false);
|
||||
|
||||
const handleCopy = async () => {
|
||||
try {
|
||||
await navigator.clipboard.writeText(value);
|
||||
if (await copyToClipboard(value)) {
|
||||
setCopied(true);
|
||||
setTimeout(() => setCopied(false), 1200);
|
||||
} catch {
|
||||
/* clipboard unavailable in non-secure contexts */
|
||||
}
|
||||
};
|
||||
|
||||
|
|
|
|||
|
|
@ -1,12 +1,15 @@
|
|||
import React from "react";
|
||||
import { render, screen, waitFor, fireEvent, act } from "@testing-library/react";
|
||||
import { render, screen, waitFor } from "@testing-library/react";
|
||||
import userEvent from "@testing-library/user-event";
|
||||
import { describe, it, expect, vi, beforeEach } from "vitest";
|
||||
import { InputCard } from "./InputCard";
|
||||
import { ParsedMessage } from "./prettyMessagesTypes";
|
||||
import { copyToClipboard } from "@/utils/dataUtils";
|
||||
import { toast } from "@/lib/toast";
|
||||
|
||||
vi.mock("@/utils/dataUtils", () => ({ copyToClipboard: vi.fn() }));
|
||||
|
||||
describe("InputCard", () => {
|
||||
const mockWriteText = vi.fn().mockResolvedValue(undefined);
|
||||
const mockMessages: ParsedMessage[] = [
|
||||
{
|
||||
role: "user",
|
||||
|
|
@ -16,13 +19,7 @@ describe("InputCard", () => {
|
|||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
Object.defineProperty(navigator, "clipboard", {
|
||||
value: {
|
||||
writeText: mockWriteText,
|
||||
},
|
||||
writable: true,
|
||||
configurable: true,
|
||||
});
|
||||
vi.mocked(copyToClipboard).mockResolvedValue(true);
|
||||
});
|
||||
|
||||
it("should render the InputCard component", () => {
|
||||
|
|
@ -91,27 +88,16 @@ describe("InputCard", () => {
|
|||
expect(screen.getByText(/Cost: \$0\.001500/)).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should copy last message content when copy button is clicked", async () => {
|
||||
const messages: ParsedMessage[] = [
|
||||
{
|
||||
role: "user",
|
||||
content: "Copy this text",
|
||||
},
|
||||
];
|
||||
it.each([true, false])("delegates copying and feedback to the helper (%s)", async (success) => {
|
||||
const user = userEvent.setup();
|
||||
Object.defineProperty(navigator, "clipboard", { value: undefined, configurable: true });
|
||||
vi.mocked(copyToClipboard).mockResolvedValue(success);
|
||||
render(<InputCard messages={[{ role: "user", content: "Copy this text" }]} />);
|
||||
|
||||
render(<InputCard messages={messages} />);
|
||||
await user.click(screen.getByRole("button", { name: /copy/i }));
|
||||
|
||||
const copyButton = screen.getByRole("button", { name: /copy/i });
|
||||
|
||||
expect(copyButton).toBeInTheDocument();
|
||||
|
||||
await act(async () => {
|
||||
fireEvent.click(copyButton);
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockWriteText).toHaveBeenCalledWith("Copy this text");
|
||||
});
|
||||
expect(copyToClipboard).toHaveBeenCalledWith("Copy this text", "Input copied");
|
||||
expect(toast.success).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("should toggle collapse state when header is clicked", async () => {
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@
|
|||
*/
|
||||
|
||||
import { useState } from "react";
|
||||
import { toast } from "@/lib/toast";
|
||||
import { copyToClipboard } from "@/utils/dataUtils";
|
||||
import { ParsedMessage } from "./prettyMessagesTypes";
|
||||
import { SectionHeader } from "./SectionHeader";
|
||||
import { CollapsibleMessage } from "./CollapsibleMessage";
|
||||
|
|
@ -32,8 +32,7 @@ export function InputCard({ messages, promptTokens, inputCost }: InputCardProps)
|
|||
|
||||
const handleCopy = () => {
|
||||
const content = lastMessage?.content || "";
|
||||
navigator.clipboard.writeText(content);
|
||||
toast.success("Input copied");
|
||||
void copyToClipboard(content, "Input copied");
|
||||
};
|
||||
|
||||
return (
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ import userEvent from "@testing-library/user-event";
|
|||
import { describe, expect, it, vi } from "vitest";
|
||||
import { GuardrailJumpLink, LogDetailContent } from "./LogDetailContent";
|
||||
import type { LogEntry } from "../types";
|
||||
import { toast } from "@/lib/toast";
|
||||
|
||||
vi.mock("./sections/GuardrailViewer/GuardrailViewer", () => ({
|
||||
default: ({ data }: { data: unknown }) => <div data-testid="guardrail-viewer">{JSON.stringify(data)}</div>,
|
||||
|
|
@ -295,6 +296,46 @@ describe("LogDetailContent", () => {
|
|||
expect(screen.queryByRole("tab", { name: "Request" })).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it.each([
|
||||
["Request", true],
|
||||
["Response", true],
|
||||
["Metadata", true],
|
||||
["Request", false],
|
||||
] as const)("copies %s through the fallback and reports its outcome (%s)", async (section, success) => {
|
||||
const user = userEvent.setup();
|
||||
const clipboardDescriptor = Object.getOwnPropertyDescriptor(navigator, "clipboard");
|
||||
const execCommandDescriptor = Object.getOwnPropertyDescriptor(document, "execCommand");
|
||||
const writeText = vi.fn<(value: string) => boolean>().mockReturnValue(success);
|
||||
Object.defineProperty(navigator, "clipboard", { value: undefined, configurable: true });
|
||||
Object.defineProperty(document, "execCommand", {
|
||||
configurable: true,
|
||||
value: (command: string) =>
|
||||
command === "copy" && writeText((document.activeElement as HTMLTextAreaElement).value),
|
||||
});
|
||||
vi.mocked(toast.success).mockClear();
|
||||
const log = createLogEntry();
|
||||
const expected = { Request: log.messages, Response: log.response, Metadata: log.metadata }[section];
|
||||
|
||||
try {
|
||||
render(<LogDetailContent logEntry={log} />);
|
||||
if (section !== "Metadata") {
|
||||
await user.click(screen.getByRole("tab", { name: "JSON", exact: true }));
|
||||
await user.click(screen.getByRole("tab", { name: section, exact: true }));
|
||||
}
|
||||
const label = section === "Metadata" ? "Copy Metadata" : "Copy JSON";
|
||||
const button = screen.getByRole("button", { name: label, exact: true });
|
||||
await user.click(button);
|
||||
|
||||
expect(writeText).toHaveBeenCalledWith(JSON.stringify(expected, null, 2));
|
||||
expect(button).toHaveAccessibleName(success ? "Copied!" : label);
|
||||
expect(toast.success).toHaveBeenCalledTimes(success ? 1 : 0);
|
||||
} finally {
|
||||
if (clipboardDescriptor) Object.defineProperty(navigator, "clipboard", clipboardDescriptor);
|
||||
if (execCommandDescriptor) Object.defineProperty(document, "execCommand", execCommandDescriptor);
|
||||
else Reflect.deleteProperty(document, "execCommand");
|
||||
}
|
||||
});
|
||||
|
||||
it.each(["object", "serialized", "messages only", "null captures"])(
|
||||
"preserves request inspection and copying for classifier logs with %s data",
|
||||
async (shape) => {
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import { copyToClipboard } from "@/utils/dataUtils";
|
||||
import { useState } from "react";
|
||||
import { Check, ChevronDown, ChevronRight, CircleAlert, Copy, Info } from "lucide-react";
|
||||
import moment from "moment";
|
||||
|
|
@ -304,12 +305,9 @@ function CopyButton({
|
|||
const [copied, setCopied] = useState(false);
|
||||
|
||||
const handleCopy = async () => {
|
||||
try {
|
||||
await navigator.clipboard.writeText(getText());
|
||||
if (await copyToClipboard(getText())) {
|
||||
setCopied(true);
|
||||
setTimeout(() => setCopied(false), 1200);
|
||||
} catch {
|
||||
/* clipboard unavailable in non-secure contexts */
|
||||
}
|
||||
};
|
||||
|
||||
|
|
|
|||
|
|
@ -1,3 +1,5 @@
|
|||
import userEvent from "@testing-library/user-event";
|
||||
import { copyToClipboard } from "@/utils/dataUtils";
|
||||
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
|
||||
import { fireEvent, render, screen, waitFor } from "@testing-library/react";
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
|
|
@ -6,6 +8,11 @@ import { sessionSpendLogsCall } from "../../networking";
|
|||
import { LogEntry } from "../types";
|
||||
import { AutoRouterModelGroupsProvider } from "@/components/shared/table_cells";
|
||||
|
||||
vi.mock("@/utils/dataUtils", async (importOriginal) => ({
|
||||
...(await importOriginal<typeof import("@/utils/dataUtils")>()),
|
||||
copyToClipboard: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock("../../networking", () => ({
|
||||
sessionSpendLogsCall: vi.fn(),
|
||||
}));
|
||||
|
|
@ -201,3 +208,18 @@ describe("LogDetailsDrawer session sidebar auto-router icon", () => {
|
|||
expect(directRow.querySelector(".lucide-waypoints")).toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
describe("LogDetailsDrawer copying", () => {
|
||||
it.each([true, false])("confirms a session id only when copying succeeds (%s)", async (success) => {
|
||||
const user = userEvent.setup();
|
||||
Object.defineProperty(navigator, "clipboard", { value: undefined, configurable: true });
|
||||
vi.mocked(copyToClipboard).mockReset().mockResolvedValue(success);
|
||||
renderSessionDrawer();
|
||||
|
||||
const button = await screen.findByRole("button", { name: "Copy trace id" });
|
||||
await user.click(button);
|
||||
|
||||
expect(copyToClipboard).toHaveBeenCalledWith("session-1");
|
||||
expect(button.querySelector(".lucide-check") !== null).toBe(success);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import { copyToClipboard } from "@/utils/dataUtils";
|
||||
import { useEffect, useMemo, useState } from "react";
|
||||
import { Bot, Check, Copy, Sparkles, Wrench } from "lucide-react";
|
||||
import { Sheet, SheetContent, SheetTitle } from "@/components/ui/sheet";
|
||||
|
|
@ -289,12 +290,9 @@ export function LogDetailsDrawer({
|
|||
|
||||
const handleCopyLeftPanelId = async () => {
|
||||
if (!leftPanelId) return;
|
||||
try {
|
||||
await navigator.clipboard.writeText(leftPanelId);
|
||||
if (await copyToClipboard(leftPanelId)) {
|
||||
setCopiedLeftPanelId(true);
|
||||
setTimeout(() => setCopiedLeftPanelId(false), 1200);
|
||||
} catch {
|
||||
/* clipboard unavailable in non-secure contexts */
|
||||
}
|
||||
};
|
||||
|
||||
|
|
|
|||
|
|
@ -1,12 +1,15 @@
|
|||
import React from "react";
|
||||
import { render, screen, waitFor, fireEvent, act } from "@testing-library/react";
|
||||
import { render, screen, waitFor } from "@testing-library/react";
|
||||
import userEvent from "@testing-library/user-event";
|
||||
import { describe, it, expect, vi, beforeEach } from "vitest";
|
||||
import { OutputCard } from "./OutputCard";
|
||||
import { ParsedMessage } from "./prettyMessagesTypes";
|
||||
import { copyToClipboard } from "@/utils/dataUtils";
|
||||
import { toast } from "@/lib/toast";
|
||||
|
||||
vi.mock("@/utils/dataUtils", () => ({ copyToClipboard: vi.fn() }));
|
||||
|
||||
describe("OutputCard", () => {
|
||||
const mockWriteText = vi.fn().mockResolvedValue(undefined);
|
||||
const mockMessage: ParsedMessage = {
|
||||
role: "assistant",
|
||||
content: "This is a test response",
|
||||
|
|
@ -14,13 +17,7 @@ describe("OutputCard", () => {
|
|||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
Object.defineProperty(navigator, "clipboard", {
|
||||
value: {
|
||||
writeText: mockWriteText,
|
||||
},
|
||||
writable: true,
|
||||
configurable: true,
|
||||
});
|
||||
vi.mocked(copyToClipboard).mockResolvedValue(true);
|
||||
});
|
||||
|
||||
it("should render the OutputCard component", () => {
|
||||
|
|
@ -53,20 +50,16 @@ describe("OutputCard", () => {
|
|||
expect(screen.getByText(/Cost: \$0\.002500/)).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should copy message content when copy button is clicked", async () => {
|
||||
it.each([true, false])("delegates copying and feedback to the helper (%s)", async (success) => {
|
||||
const user = userEvent.setup();
|
||||
Object.defineProperty(navigator, "clipboard", { value: undefined, configurable: true });
|
||||
vi.mocked(copyToClipboard).mockResolvedValue(success);
|
||||
render(<OutputCard message={mockMessage} />);
|
||||
|
||||
const copyButton = screen.getByRole("button", { name: /copy/i });
|
||||
await user.click(screen.getByRole("button", { name: /copy/i }));
|
||||
|
||||
expect(copyButton).toBeInTheDocument();
|
||||
|
||||
await act(async () => {
|
||||
fireEvent.click(copyButton);
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockWriteText).toHaveBeenCalledWith("This is a test response");
|
||||
});
|
||||
expect(copyToClipboard).toHaveBeenCalledWith("This is a test response", "Output copied");
|
||||
expect(toast.success).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("should not copy when message is null and copy button is clicked", async () => {
|
||||
|
|
@ -80,7 +73,7 @@ describe("OutputCard", () => {
|
|||
await user.click(copyButton);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockWriteText).not.toHaveBeenCalled();
|
||||
expect(copyToClipboard).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@
|
|||
*/
|
||||
|
||||
import { useState } from "react";
|
||||
import { toast } from "@/lib/toast";
|
||||
import { copyToClipboard } from "@/utils/dataUtils";
|
||||
import { COLOR_BORDER } from "./constants";
|
||||
import { ParsedMessage } from "./prettyMessagesTypes";
|
||||
import { SectionHeader } from "./SectionHeader";
|
||||
|
|
@ -22,8 +22,7 @@ export function OutputCard({ message, completionTokens, outputCost }: OutputCard
|
|||
const handleCopy = () => {
|
||||
if (!message) return;
|
||||
|
||||
navigator.clipboard.writeText(message.content || "");
|
||||
toast.success("Output copied");
|
||||
void copyToClipboard(message.content || "", "Output copied");
|
||||
};
|
||||
|
||||
return (
|
||||
|
|
|
|||
|
|
@ -1,9 +1,12 @@
|
|||
import React from "react";
|
||||
import { render, screen, waitFor } from "@testing-library/react";
|
||||
import userEvent from "@testing-library/user-event";
|
||||
import { copyToClipboard } from "@/utils/dataUtils";
|
||||
import { describe, it, expect, vi, beforeEach } from "vitest";
|
||||
import { RealtimePrettyView, isRealtimeResponse } from "./RealtimePrettyView";
|
||||
|
||||
vi.mock("@/utils/dataUtils", () => ({ copyToClipboard: vi.fn().mockResolvedValue(true) }));
|
||||
|
||||
const sampleRealtimeResponse = {
|
||||
usage: {
|
||||
total_tokens: 587,
|
||||
|
|
@ -142,15 +145,20 @@ describe("isRealtimeResponse", () => {
|
|||
});
|
||||
|
||||
describe("RealtimePrettyView", () => {
|
||||
const mockWriteText = vi.fn().mockResolvedValue(undefined);
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
Object.defineProperty(navigator, "clipboard", {
|
||||
value: { writeText: mockWriteText },
|
||||
writable: true,
|
||||
configurable: true,
|
||||
});
|
||||
});
|
||||
|
||||
it("copies all transcript turns when the Clipboard API is unavailable", async () => {
|
||||
const user = userEvent.setup();
|
||||
Object.defineProperty(navigator, "clipboard", { value: undefined, configurable: true });
|
||||
render(<RealtimePrettyView response={sampleRealtimeResponse} />);
|
||||
|
||||
await user.click(screen.getByRole("button", { name: /copy/i }));
|
||||
|
||||
expect(copyToClipboard).toHaveBeenCalledWith(
|
||||
"assistant: Hello! How's your day going?\nassistant: I'm here to help with information and general questions.",
|
||||
);
|
||||
});
|
||||
|
||||
it("should render the component successfully", () => {
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import { copyToClipboard } from "@/utils/dataUtils";
|
||||
/**
|
||||
* RealtimePrettyView - Structured pretty view for OpenAI Realtime API logs
|
||||
* Displays session config, conversation turns, and token usage
|
||||
|
|
@ -276,7 +277,7 @@ function ConversationCard({
|
|||
(r.output || []).flatMap((o) => (o.content || []).map((c) => `${o.role}: ${c.transcript || c.text || ""}`)),
|
||||
)
|
||||
.join("\n");
|
||||
navigator.clipboard.writeText(transcripts);
|
||||
void copyToClipboard(transcripts);
|
||||
};
|
||||
|
||||
return (
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ import userEvent from "@testing-library/user-event";
|
|||
import { afterEach, describe, expect, it, vi } from "vitest";
|
||||
import type { ComponentProps } from "react";
|
||||
import { copyToClipboard } from "@/utils/dataUtils";
|
||||
import { renderWithProviders, screen, waitFor } from "../../../tests/test-utils";
|
||||
import { renderWithProviders, screen } from "../../../tests/test-utils";
|
||||
import CopyButton from "./CopyButton";
|
||||
|
||||
vi.mock("@/utils/dataUtils", () => ({ copyToClipboard: vi.fn() }));
|
||||
|
|
@ -13,13 +13,9 @@ vi.mock("lucide-react", async (importOriginal) => ({
|
|||
Check: (props: ComponentProps<"svg">) => <svg {...props} role="img" aria-label="Copied icon" />,
|
||||
}));
|
||||
|
||||
const clipboardDescriptor = Object.getOwnPropertyDescriptor(navigator, "clipboard");
|
||||
|
||||
afterEach(() => {
|
||||
vi.useRealTimers();
|
||||
vi.resetAllMocks();
|
||||
if (clipboardDescriptor) Object.defineProperty(navigator, "clipboard", clipboardDescriptor);
|
||||
else Reflect.deleteProperty(navigator, "clipboard");
|
||||
});
|
||||
|
||||
describe("CopyButton", () => {
|
||||
|
|
@ -28,55 +24,37 @@ describe("CopyButton", () => {
|
|||
expect(screen.queryByRole("button")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("shows the confirmation checkmark only after a successful write", async () => {
|
||||
it("shows the icon confirmation only after copying succeeds", async () => {
|
||||
const user = userEvent.setup();
|
||||
const writeText = vi.fn().mockResolvedValue(undefined);
|
||||
Object.defineProperty(navigator, "clipboard", { value: { writeText }, configurable: true });
|
||||
const pendingCopy = Promise.withResolvers<boolean>();
|
||||
vi.mocked(copyToClipboard).mockReturnValue(pendingCopy.promise);
|
||||
renderWithProviders(<CopyButton value="secret-value" label="Copy value" copiedLabel="Value copied" />);
|
||||
|
||||
renderWithProviders(<CopyButton value="secret-value" label="Copy value" />);
|
||||
await user.click(screen.getByRole("button", { name: "Copy value" }));
|
||||
|
||||
const button = screen.getByRole("button", { name: "Copy value" });
|
||||
expect(copyToClipboard).toHaveBeenCalledWith("secret-value", "Value copied");
|
||||
expect(screen.getByRole("img", { name: "Copy icon" })).toBeInTheDocument();
|
||||
expect(screen.queryByRole("img", { name: "Copied icon" })).not.toBeInTheDocument();
|
||||
|
||||
await user.click(button);
|
||||
|
||||
expect(writeText).toHaveBeenCalledWith("secret-value");
|
||||
expect(await screen.findByRole("img", { name: "Copied icon" })).toBeInTheDocument();
|
||||
await act(async () => pendingCopy.resolve(true));
|
||||
expect(screen.getByRole("img", { name: "Copied icon" })).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("does not show the checkmark when the clipboard write is rejected", async () => {
|
||||
it("keeps the icon unchanged when copying fails", async () => {
|
||||
const user = userEvent.setup();
|
||||
const writeText = vi.fn().mockRejectedValue(new Error("permission denied"));
|
||||
Object.defineProperty(navigator, "clipboard", { value: { writeText }, configurable: true });
|
||||
|
||||
vi.mocked(copyToClipboard).mockResolvedValue(false);
|
||||
renderWithProviders(<CopyButton value="secret-value" label="Copy value" />);
|
||||
|
||||
const button = screen.getByRole("button", { name: "Copy value" });
|
||||
await user.click(button);
|
||||
await user.click(screen.getByRole("button", { name: "Copy value" }));
|
||||
|
||||
await waitFor(() => expect(writeText).toHaveBeenCalledWith("secret-value"));
|
||||
expect(copyToClipboard).toHaveBeenCalledWith("secret-value", "Copied");
|
||||
expect(screen.queryByRole("img", { name: "Copied icon" })).not.toBeInTheDocument();
|
||||
expect(screen.getByRole("img", { name: "Copy icon" })).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("does not show the checkmark when the clipboard API is unavailable", async () => {
|
||||
const user = userEvent.setup();
|
||||
Object.defineProperty(navigator, "clipboard", { value: undefined, configurable: true });
|
||||
|
||||
renderWithProviders(<CopyButton value="secret-value" label="Copy value" />);
|
||||
|
||||
const button = screen.getByRole("button", { name: "Copy value" });
|
||||
await user.click(button);
|
||||
|
||||
expect(screen.queryByRole("img", { name: "Copied icon" })).not.toBeInTheDocument();
|
||||
expect(screen.getByRole("img", { name: "Copy icon" })).toBeInTheDocument();
|
||||
expect(copyToClipboard).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("keeps the default icon styling, custom classes, title and 1200ms confirmation", async () => {
|
||||
vi.useFakeTimers();
|
||||
const writeText = vi.fn().mockResolvedValue(undefined);
|
||||
Object.defineProperty(navigator, "clipboard", { value: { writeText }, configurable: true });
|
||||
vi.mocked(copyToClipboard).mockResolvedValue(true);
|
||||
renderWithProviders(
|
||||
<CopyButton value="request-id" label="Copy request ID" className="shrink-0" iconClassName="size-3" />,
|
||||
);
|
||||
|
|
@ -87,8 +65,7 @@ describe("CopyButton", () => {
|
|||
expect(screen.getByRole("img", { name: "Copy icon" })).toHaveClass("size-3");
|
||||
|
||||
await act(async () => fireEvent.click(button));
|
||||
expect(writeText).toHaveBeenCalledWith("request-id");
|
||||
expect(copyToClipboard).not.toHaveBeenCalled();
|
||||
expect(copyToClipboard).toHaveBeenCalledWith("request-id", "Copied");
|
||||
act(() => vi.advanceTimersByTime(1199));
|
||||
expect(screen.getByRole("img", { name: "Copied icon" })).toBeInTheDocument();
|
||||
act(() => vi.advanceTimersByTime(1));
|
||||
|
|
|
|||
|
|
@ -54,17 +54,7 @@ const CopyButton: React.FC<CopyButtonProps> = ({
|
|||
if (value == null || (!value && !isAction)) return null;
|
||||
|
||||
const handleCopy = async () => {
|
||||
if (isAction) {
|
||||
setCopied(await copyToClipboard(value, copiedLabel));
|
||||
return;
|
||||
}
|
||||
if (!navigator.clipboard) return;
|
||||
try {
|
||||
await navigator.clipboard.writeText(value);
|
||||
setCopied(true);
|
||||
} catch {
|
||||
setCopied(false);
|
||||
}
|
||||
setCopied(await copyToClipboard(value, copiedLabel));
|
||||
};
|
||||
|
||||
const button = (
|
||||
|
|
|
|||
|
|
@ -45,5 +45,30 @@ describe("MyUserTab", () => {
|
|||
expect(screen.getByText("member@example.com")).toBeInTheDocument();
|
||||
expect(screen.getByText("model-one")).toBeInTheDocument();
|
||||
expect(screen.getByText("TPM: 1,000")).toBeInTheDocument();
|
||||
expect(screen.getByText("Never resets")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("shows the next reset time when the member budget has a reset window", () => {
|
||||
vi.mocked(useMyTeamMember).mockReturnValue({
|
||||
data: {
|
||||
user_id: "user-1",
|
||||
user_email: "member@example.com",
|
||||
team_id: "team-1",
|
||||
role: "user",
|
||||
spend: 12.5,
|
||||
total_spend: 30,
|
||||
litellm_budget_table: {
|
||||
max_budget: 100,
|
||||
budget_duration: "30d",
|
||||
budget_reset_at: new Date(2026, 6, 15, 12, 0, 0).toISOString(),
|
||||
},
|
||||
},
|
||||
isLoading: false,
|
||||
error: null,
|
||||
} as ReturnType<typeof useMyTeamMember>);
|
||||
|
||||
renderWithProviders(<MyUserTab teamId="team-1" />);
|
||||
|
||||
expect(screen.getByText("Resets Jul 15, 2026 12:00")).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
import { formatBudgetReset } from "@/utils/budgetUtils";
|
||||
import { describeBudgetResetSchedule } from "@/utils/budgetUtils";
|
||||
import { formatNumberWithCommas } from "@/utils/dataUtils";
|
||||
import { SimpleTooltip } from "@/components/ui/tooltip";
|
||||
import { Badge } from "@/components/ui/badge";
|
||||
|
|
@ -67,7 +67,7 @@ export default function MyUserTab({ teamId }: MyUserTabProps) {
|
|||
const totalSpend = data.total_spend ?? 0;
|
||||
const tpmLimit = budgetTable?.tpm_limit ?? null;
|
||||
const rpmLimit = budgetTable?.rpm_limit ?? null;
|
||||
const budgetReset = formatBudgetReset(budgetTable?.budget_reset_at);
|
||||
const budgetReset = describeBudgetResetSchedule(budgetTable?.budget_duration, budgetTable?.budget_reset_at);
|
||||
const allowedModels = budgetTable?.allowed_models ?? null;
|
||||
|
||||
return (
|
||||
|
|
@ -93,17 +93,16 @@ export default function MyUserTab({ teamId }: MyUserTabProps) {
|
|||
<div className="grid grid-cols-1 gap-4 md:grid-cols-2">
|
||||
<Card>
|
||||
<CardContent>
|
||||
{labelWithTooltip(
|
||||
"Current Cycle Spend (USD)",
|
||||
"Spend for the current budget cycle. Resets to $0 when the budget window rolls over.",
|
||||
)}
|
||||
{labelWithTooltip("Current Cycle Spend (USD)", budgetReset.tooltip)}
|
||||
<div className="mt-2">
|
||||
<h3 className="text-2xl font-semibold">${formatNumber(spend, 4)}</h3>
|
||||
<span className="text-muted-foreground">
|
||||
of {maxBudget === null ? "Unlimited" : `$${formatNumber(maxBudget, 4)}`}
|
||||
</span>
|
||||
</div>
|
||||
{budgetReset && <div className="mt-1 text-muted-foreground">Resets {budgetReset}</div>}
|
||||
<div className="mt-1 text-muted-foreground">
|
||||
{budgetReset.resetAt ? `Resets ${budgetReset.label}` : budgetReset.label}
|
||||
</div>
|
||||
</CardContent>
|
||||
</Card>
|
||||
|
||||
|
|
@ -120,7 +119,10 @@ export default function MyUserTab({ teamId }: MyUserTabProps) {
|
|||
|
||||
<Card>
|
||||
<CardContent>
|
||||
{labelWithTooltip("Total Spend (USD)", "Cumulative spend across all budget cycles within this team.")}
|
||||
{labelWithTooltip(
|
||||
"Total Spend (USD)",
|
||||
"Cumulative spend across budget cycles within this team. This is a lower bound: lifetime tracking began 2026-04-21, so cycles completed before then are not included. It is never less than Current Cycle Spend.",
|
||||
)}
|
||||
<h4 className="mt-2 text-xl font-semibold">${formatNumber(totalSpend, 4)}</h4>
|
||||
</CardContent>
|
||||
</Card>
|
||||
|
|
|
|||
|
|
@ -313,7 +313,47 @@ describe("TeamMembersComponent", () => {
|
|||
/>,
|
||||
);
|
||||
|
||||
expect(screen.getByText("Jul 15, 2026")).toBeInTheDocument();
|
||||
expect(screen.getByText("Jul 15, 2026 12:00")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("shows Never resets for a member whose budget has no reset window", () => {
|
||||
const teamData = createMockTeamData();
|
||||
teamData.team_memberships[0].litellm_budget_table!.budget_reset_at = null;
|
||||
renderWithProviders(
|
||||
<TeamMembersComponent
|
||||
teamData={teamData}
|
||||
canEditTeam={false}
|
||||
handleMemberDelete={mockHandleMemberDelete}
|
||||
onMemberSpendReset={mockOnMemberSpendReset}
|
||||
onMemberBudgetReset={mockOnMemberBudgetReset}
|
||||
setSelectedEditMember={mockSetSelectedEditMember}
|
||||
setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible}
|
||||
setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible}
|
||||
/>,
|
||||
);
|
||||
|
||||
expect(screen.queryByText("Jul 15, 2026 12:00")).not.toBeInTheDocument();
|
||||
expect(screen.getAllByText("Never resets").length).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
it("shows Unknown for a member whose budget has a duration but no scheduled reset yet", () => {
|
||||
const teamData = createMockTeamData();
|
||||
teamData.team_memberships[0].litellm_budget_table!.budget_duration = "30d";
|
||||
teamData.team_memberships[0].litellm_budget_table!.budget_reset_at = null;
|
||||
renderWithProviders(
|
||||
<TeamMembersComponent
|
||||
teamData={teamData}
|
||||
canEditTeam={false}
|
||||
handleMemberDelete={mockHandleMemberDelete}
|
||||
onMemberSpendReset={mockOnMemberSpendReset}
|
||||
onMemberBudgetReset={mockOnMemberBudgetReset}
|
||||
setSelectedEditMember={mockSetSelectedEditMember}
|
||||
setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible}
|
||||
setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible}
|
||||
/>,
|
||||
);
|
||||
|
||||
expect(screen.getByText("Unknown")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("shows each member the values from its own membership row", () => {
|
||||
|
|
|
|||
|
|
@ -10,8 +10,9 @@ import MemberTable from "@/components/common_components/MemberTable";
|
|||
import type { StoredModelMaxBudget } from "@/components/key_team_helpers/modelMaxBudgetPayload";
|
||||
import { Member } from "@/components/networking";
|
||||
import { parseErrorMessage } from "@/components/shared/errorUtils";
|
||||
import { DateCell, MoneyCell } from "@/components/shared/table_cells";
|
||||
import { MoneyCell } from "@/components/shared/table_cells";
|
||||
import { toast } from "@/lib/toast";
|
||||
import { describeBudgetResetSchedule } from "@/utils/budgetUtils";
|
||||
import { formatNumberWithCommas } from "@/utils/dataUtils";
|
||||
import { isProxyAdminRole, isUserTeamAdminForSingleTeam } from "@/utils/roles";
|
||||
import { CircleHelp } from "lucide-react";
|
||||
|
|
@ -152,6 +153,16 @@ export default function TeamMemberTab({
|
|||
return findMembership(userId)?.litellm_budget_table?.budget_reset_at ?? null;
|
||||
};
|
||||
|
||||
const renderBudgetReset = (userId: string | null) => {
|
||||
const budget = findMembership(userId)?.litellm_budget_table;
|
||||
const schedule = describeBudgetResetSchedule(budget?.budget_duration ?? null, budget?.budget_reset_at ?? null);
|
||||
return (
|
||||
<SimpleTooltip content={schedule.tooltip}>
|
||||
<span>{schedule.label}</span>
|
||||
</SimpleTooltip>
|
||||
);
|
||||
};
|
||||
|
||||
const extraColumns: NonNullable<ComponentProps<typeof MemberTable>["extraColumns"]> = [
|
||||
{
|
||||
title: (
|
||||
|
|
@ -190,7 +201,7 @@ export default function TeamMemberTab({
|
|||
title: (
|
||||
<span className="flex items-center gap-1">
|
||||
Current Cycle Spend (USD)
|
||||
<SimpleTooltip content="Spend for the current budget cycle. Resets to $0 when the member's budget window rolls over. This is the value checked against the member's budget.">
|
||||
<SimpleTooltip content="Spend in this member's current budget cycle, the value checked against their budget. The Budget Reset column shows when this cycle ends, or Never resets. Spend above a rollover cap at reset carries into the next cycle.">
|
||||
<CircleHelp className="size-4" aria-label="Current cycle spend information" />
|
||||
</SimpleTooltip>
|
||||
</span>
|
||||
|
|
@ -204,7 +215,7 @@ export default function TeamMemberTab({
|
|||
title: (
|
||||
<span className="flex items-center gap-1">
|
||||
Total Spend (USD)
|
||||
<SimpleTooltip content="Cumulative spend by this member within this team, across all budget cycles. Tracking began 2026-04-21; spend from before that date is not included.">
|
||||
<SimpleTooltip content="Cumulative spend by this member within this team, across budget cycles. This is a lower bound: lifetime tracking began 2026-04-21, so cycles completed before then are not included. It is never less than Current Cycle Spend.">
|
||||
<CircleHelp className="size-4" aria-label="Total spend information" />
|
||||
</SimpleTooltip>
|
||||
</span>
|
||||
|
|
@ -255,7 +266,7 @@ export default function TeamMemberTab({
|
|||
title: "Budget Reset",
|
||||
key: "budget_reset",
|
||||
sortValue: (record: Member) => getUserBudgetReset(record.user_id),
|
||||
render: (record: Member) => <DateCell value={getUserBudgetReset(record.user_id)} precision="date" />,
|
||||
render: (record: Member) => renderBudgetReset(record.user_id),
|
||||
},
|
||||
{
|
||||
title: (
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
import { afterEach, describe, expect, it } from "vitest";
|
||||
|
||||
import { formatBudgetReset, formatBudgetResetUtc } from "./budgetUtils";
|
||||
import { describeBudgetResetSchedule, formatBudgetReset, formatBudgetResetUtc } from "./budgetUtils";
|
||||
|
||||
const originalTimezone: string | undefined = process.env.TZ;
|
||||
|
||||
|
|
@ -21,3 +21,26 @@ describe("formatBudgetResetUtc", () => {
|
|||
expect(formatBudgetReset(resetAt)).toBe("Oct 19, 2026");
|
||||
});
|
||||
});
|
||||
|
||||
describe("describeBudgetResetSchedule", () => {
|
||||
it("reports Never resets when the budget has no duration and no scheduled reset", () => {
|
||||
const schedule = describeBudgetResetSchedule(null, null);
|
||||
expect(schedule.label).toBe("Never resets");
|
||||
expect(schedule.resetAt).toBeNull();
|
||||
});
|
||||
|
||||
it("reports Unknown when a duration is set but no reset has been scheduled yet", () => {
|
||||
const schedule = describeBudgetResetSchedule("30d", null);
|
||||
expect(schedule.label).toBe("Unknown");
|
||||
expect(schedule.tooltip).toContain("every 30d");
|
||||
});
|
||||
|
||||
it("shows the next reset with its time zone and the rollover carryover note", () => {
|
||||
const resetAt = new Date(2026, 6, 15, 12, 0, 0);
|
||||
const schedule = describeBudgetResetSchedule("30d", resetAt.toISOString());
|
||||
expect(schedule.label).toBe("Jul 15, 2026 12:00");
|
||||
expect(schedule.resetAt).toBe(resetAt.toISOString());
|
||||
expect(schedule.tooltip).toContain(Intl.DateTimeFormat().resolvedOptions().timeZone);
|
||||
expect(schedule.tooltip).toContain("carries into the next cycle");
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -16,3 +16,41 @@ export function formatBudgetResetUtc(iso: string | null | undefined): string | n
|
|||
if (!resetDate.isValid()) return null;
|
||||
return resetDate.format("MMM D, YYYY");
|
||||
}
|
||||
|
||||
export interface BudgetResetSchedule {
|
||||
label: string;
|
||||
resetAt: string | null;
|
||||
tooltip: string;
|
||||
}
|
||||
|
||||
const ROLLOVER_NOTE = "Spend above a rollover cap at reset carries into the next cycle.";
|
||||
|
||||
export function describeBudgetResetSchedule(
|
||||
budgetDuration: string | null | undefined,
|
||||
budgetResetAt: string | null | undefined,
|
||||
): BudgetResetSchedule {
|
||||
const resetDate = budgetResetAt ? dayjs(budgetResetAt) : null;
|
||||
if (resetDate?.isValid()) {
|
||||
const cadence = budgetDuration ? `every ${budgetDuration}` : "on a schedule";
|
||||
const timeZone = Intl.DateTimeFormat().resolvedOptions().timeZone;
|
||||
const label = resetDate.format("MMM D, YYYY HH:mm");
|
||||
return {
|
||||
label,
|
||||
resetAt: resetDate.toISOString(),
|
||||
tooltip: `Current Cycle Spend resets ${cadence}. Next reset ${label} (${timeZone}). ${ROLLOVER_NOTE}`,
|
||||
};
|
||||
}
|
||||
if (!budgetDuration) {
|
||||
return {
|
||||
label: "Never resets",
|
||||
resetAt: null,
|
||||
tooltip:
|
||||
"This budget has no reset window, so Current Cycle Spend accumulates for as long as the membership exists.",
|
||||
};
|
||||
}
|
||||
return {
|
||||
label: "Unknown",
|
||||
resetAt: null,
|
||||
tooltip: `Resets every ${budgetDuration}, but the next reset time has not been scheduled yet. ${ROLLOVER_NOTE}`,
|
||||
};
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
// @vitest-environment jsdom
|
||||
|
||||
import "@testing-library/jest-dom/vitest";
|
||||
import { describe, it, expect, beforeEach, vi, afterEach } from "vitest";
|
||||
import {
|
||||
copyToClipboard,
|
||||
|
|
@ -16,8 +17,6 @@ const mockToast = vi.mocked(toast);
|
|||
describe("dataUtils", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
// Reset document.execCommand mock
|
||||
delete (document as any).execCommand;
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
|
|
@ -134,194 +133,114 @@ describe("dataUtils", () => {
|
|||
});
|
||||
|
||||
describe("copyToClipboard", () => {
|
||||
describe("when Clipboard API is available (HTTPS scenario)", () => {
|
||||
beforeEach(() => {
|
||||
// Mock modern Clipboard API
|
||||
Object.assign(navigator, {
|
||||
clipboard: {
|
||||
writeText: vi.fn().mockResolvedValue(undefined),
|
||||
},
|
||||
});
|
||||
});
|
||||
|
||||
it("should use navigator.clipboard.writeText when available", async () => {
|
||||
const result = await copyToClipboard("test text");
|
||||
|
||||
expect(navigator.clipboard.writeText).toHaveBeenCalledWith("test text");
|
||||
expect(mockToast.success).toHaveBeenCalledWith("Copied to clipboard");
|
||||
expect(result).toBe(true);
|
||||
});
|
||||
|
||||
it("should use custom message when provided", async () => {
|
||||
await copyToClipboard("test text", "Custom message");
|
||||
|
||||
expect(mockToast.success).toHaveBeenCalledWith("Custom message");
|
||||
});
|
||||
|
||||
it("should return false for null/undefined text", async () => {
|
||||
expect(await copyToClipboard(null)).toBe(false);
|
||||
expect(await copyToClipboard(undefined)).toBe(false);
|
||||
expect(navigator.clipboard.writeText).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("should fall back to execCommand when clipboard API fails", async () => {
|
||||
// Make clipboard API fail
|
||||
navigator.clipboard.writeText = vi.fn().mockRejectedValue(new Error("Permission denied"));
|
||||
|
||||
// Mock successful execCommand
|
||||
document.execCommand = vi.fn().mockReturnValue(true);
|
||||
|
||||
// Mock DOM methods
|
||||
const mockTextArea = {
|
||||
value: "",
|
||||
style: {} as Record<string, string>,
|
||||
setAttribute: vi.fn(),
|
||||
focus: vi.fn(),
|
||||
select: vi.fn(),
|
||||
};
|
||||
document.createElement = vi.fn().mockReturnValue(mockTextArea);
|
||||
document.body.appendChild = vi.fn();
|
||||
document.body.removeChild = vi.fn();
|
||||
|
||||
const result = await copyToClipboard("test text");
|
||||
|
||||
expect(navigator.clipboard.writeText).toHaveBeenCalledWith("test text");
|
||||
expect(document.execCommand).toHaveBeenCalledWith("copy");
|
||||
expect(mockToast.success).toHaveBeenCalledWith("Copied to clipboard");
|
||||
expect(result).toBe(true);
|
||||
const clipboardDescriptor = Object.getOwnPropertyDescriptor(navigator, "clipboard");
|
||||
const execCommandDescriptor = Object.getOwnPropertyDescriptor(document, "execCommand");
|
||||
const writeText = vi.fn<(text: string) => Promise<void>>();
|
||||
const execCommand = vi.fn<(command: string) => boolean>();
|
||||
const selectedText = vi.fn(() => {
|
||||
const element = document.activeElement;
|
||||
return element instanceof HTMLTextAreaElement
|
||||
? element.value.slice(element.selectionStart, element.selectionEnd)
|
||||
: undefined;
|
||||
});
|
||||
const trigger = document.createElement("button");
|
||||
const text = "Request α\nsecond line";
|
||||
|
||||
beforeEach(() => {
|
||||
writeText.mockResolvedValue(undefined);
|
||||
execCommand.mockImplementation(() => {
|
||||
selectedText();
|
||||
return true;
|
||||
});
|
||||
Object.defineProperty(navigator, "clipboard", { configurable: true, value: { writeText } });
|
||||
Object.defineProperty(document, "execCommand", { configurable: true, value: execCommand });
|
||||
vi.spyOn(console, "error").mockImplementation(() => {});
|
||||
document.body.appendChild(trigger);
|
||||
trigger.focus();
|
||||
});
|
||||
|
||||
describe("when Clipboard API is not available (HTTP scenario)", () => {
|
||||
beforeEach(() => {
|
||||
// Mock HTTP scenario - no clipboard API
|
||||
Object.assign(navigator, {
|
||||
clipboard: undefined,
|
||||
});
|
||||
|
||||
// Mock successful execCommand
|
||||
document.execCommand = vi.fn().mockReturnValue(true);
|
||||
|
||||
// Mock DOM methods
|
||||
const mockTextArea = {
|
||||
value: "",
|
||||
style: {} as Record<string, string>,
|
||||
setAttribute: vi.fn(),
|
||||
focus: vi.fn(),
|
||||
select: vi.fn(),
|
||||
};
|
||||
document.createElement = vi.fn().mockReturnValue(mockTextArea);
|
||||
document.body.appendChild = vi.fn();
|
||||
document.body.removeChild = vi.fn();
|
||||
});
|
||||
|
||||
it("should fall back to execCommand when clipboard API is not available", async () => {
|
||||
const result = await copyToClipboard("test text");
|
||||
|
||||
expect(document.createElement).toHaveBeenCalledWith("textarea");
|
||||
expect(document.execCommand).toHaveBeenCalledWith("copy");
|
||||
expect(mockToast.success).toHaveBeenCalledWith("Copied to clipboard");
|
||||
expect(result).toBe(true);
|
||||
});
|
||||
|
||||
it("should set textarea properties correctly", async () => {
|
||||
const mockTextArea = {
|
||||
value: "",
|
||||
style: {} as Record<string, string>,
|
||||
setAttribute: vi.fn(),
|
||||
focus: vi.fn(),
|
||||
select: vi.fn(),
|
||||
};
|
||||
document.createElement = vi.fn().mockReturnValue(mockTextArea);
|
||||
|
||||
await copyToClipboard("test text");
|
||||
|
||||
expect(mockTextArea.value).toBe("test text");
|
||||
/* eslint-disable jest-dom/prefer-to-have-style -- the subject is a plain mock object, not a DOM node, so toHaveStyle cannot read it */
|
||||
expect(mockTextArea.style.position).toBe("fixed");
|
||||
expect(mockTextArea.style.left).toBe("-999999px");
|
||||
expect(mockTextArea.style.top).toBe("-999999px");
|
||||
/* eslint-enable jest-dom/prefer-to-have-style */
|
||||
expect(mockTextArea.setAttribute).toHaveBeenCalledWith("readonly", "");
|
||||
expect(mockTextArea.focus).toHaveBeenCalled();
|
||||
expect(mockTextArea.select).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("should handle execCommand failure", async () => {
|
||||
document.execCommand = vi.fn().mockReturnValue(false);
|
||||
|
||||
const result = await copyToClipboard("test text");
|
||||
|
||||
expect(document.execCommand).toHaveBeenCalledWith("copy");
|
||||
expect(mockToast.fromError).toHaveBeenCalledWith("Failed to copy to clipboard");
|
||||
expect(result).toBe(false);
|
||||
});
|
||||
|
||||
it("should handle DOM manipulation errors", async () => {
|
||||
document.createElement = vi.fn().mockImplementation(() => {
|
||||
throw new Error("DOM error");
|
||||
});
|
||||
|
||||
const result = await copyToClipboard("test text");
|
||||
|
||||
expect(mockToast.fromError).toHaveBeenCalledWith("Failed to copy to clipboard");
|
||||
expect(result).toBe(false);
|
||||
});
|
||||
|
||||
it("should clean up textarea element after successful copy", async () => {
|
||||
const mockTextArea = {
|
||||
value: "",
|
||||
style: {} as Record<string, string>,
|
||||
setAttribute: vi.fn(),
|
||||
focus: vi.fn(),
|
||||
select: vi.fn(),
|
||||
};
|
||||
document.createElement = vi.fn().mockReturnValue(mockTextArea);
|
||||
|
||||
await copyToClipboard("test text");
|
||||
|
||||
expect(document.body.appendChild).toHaveBeenCalledWith(mockTextArea);
|
||||
expect(document.body.removeChild).toHaveBeenCalledWith(mockTextArea);
|
||||
});
|
||||
afterEach(() => {
|
||||
document.body.replaceChildren();
|
||||
if (clipboardDescriptor) Object.defineProperty(navigator, "clipboard", clipboardDescriptor);
|
||||
else Reflect.deleteProperty(navigator, "clipboard");
|
||||
if (execCommandDescriptor) Object.defineProperty(document, "execCommand", execCommandDescriptor);
|
||||
else Reflect.deleteProperty(document, "execCommand");
|
||||
});
|
||||
|
||||
describe("edge cases", () => {
|
||||
beforeEach(() => {
|
||||
// Mock scenario where navigator exists but clipboard is null
|
||||
Object.assign(navigator, {
|
||||
clipboard: null,
|
||||
});
|
||||
it("confirms a modern clipboard write only after it succeeds", async () => {
|
||||
const pendingWrite = Promise.withResolvers<void>();
|
||||
writeText.mockReturnValue(pendingWrite.promise);
|
||||
const result = copyToClipboard(text, "Request copied");
|
||||
|
||||
document.execCommand = vi.fn().mockReturnValue(true);
|
||||
const mockTextArea = {
|
||||
value: "",
|
||||
style: {} as Record<string, string>,
|
||||
setAttribute: vi.fn(),
|
||||
focus: vi.fn(),
|
||||
select: vi.fn(),
|
||||
};
|
||||
document.createElement = vi.fn().mockReturnValue(mockTextArea);
|
||||
document.body.appendChild = vi.fn();
|
||||
document.body.removeChild = vi.fn();
|
||||
expect(writeText).toHaveBeenCalledWith(text);
|
||||
expect(mockToast.success).not.toHaveBeenCalled();
|
||||
pendingWrite.resolve();
|
||||
|
||||
expect(await result).toBe(true);
|
||||
expect(mockToast.success).toHaveBeenCalledWith("Request copied");
|
||||
expect(execCommand).not.toHaveBeenCalled();
|
||||
expect(trigger).toHaveFocus();
|
||||
});
|
||||
|
||||
it.each([null, undefined, ""])("does not copy absent content %s", async (value) => {
|
||||
expect(await copyToClipboard(value)).toBe(false);
|
||||
expect(writeText).not.toHaveBeenCalled();
|
||||
expect(execCommand).not.toHaveBeenCalled();
|
||||
expect(mockToast.success).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it.each([undefined, null, {}])("copies the full selection when the Clipboard API is %s", async (clipboard) => {
|
||||
Object.defineProperty(navigator, "clipboard", { configurable: true, value: clipboard });
|
||||
|
||||
expect(await copyToClipboard(text)).toBe(true);
|
||||
|
||||
expect(execCommand).toHaveBeenCalledWith("copy");
|
||||
expect(selectedText).toHaveReturnedWith(text);
|
||||
expect(mockToast.success).toHaveBeenCalledWith("Copied to clipboard");
|
||||
expect(document.querySelector("textarea")).not.toBeInTheDocument();
|
||||
expect(trigger).toHaveFocus();
|
||||
});
|
||||
|
||||
it("falls back after a rejected modern write", async () => {
|
||||
writeText.mockRejectedValue(new Error("Permission denied"));
|
||||
|
||||
expect(await copyToClipboard(text, "Request copied")).toBe(true);
|
||||
|
||||
expect(writeText).toHaveBeenCalledWith(text);
|
||||
expect(execCommand).toHaveBeenCalledWith("copy");
|
||||
expect(selectedText).toHaveReturnedWith(text);
|
||||
expect(mockToast.success).toHaveBeenCalledExactlyOnceWith("Request copied");
|
||||
expect(document.querySelector("textarea")).not.toBeInTheDocument();
|
||||
expect(trigger).toHaveFocus();
|
||||
});
|
||||
|
||||
it.each(["false", "throw"])("cleans up and reports failure when the fallback returns %s", async (outcome) => {
|
||||
Object.defineProperty(navigator, "clipboard", { configurable: true, value: undefined });
|
||||
execCommand.mockImplementation(() => {
|
||||
selectedText();
|
||||
if (outcome === "throw") throw new Error("Copy denied");
|
||||
return false;
|
||||
});
|
||||
|
||||
it("should handle navigator.clipboard being null", async () => {
|
||||
const result = await copyToClipboard("test text");
|
||||
expect(await copyToClipboard(text)).toBe(false);
|
||||
|
||||
expect(document.execCommand).toHaveBeenCalledWith("copy");
|
||||
expect(result).toBe(true);
|
||||
expect(selectedText).toHaveReturnedWith(text);
|
||||
expect(mockToast.fromError).toHaveBeenCalledWith("Failed to copy to clipboard");
|
||||
expect(mockToast.success).not.toHaveBeenCalled();
|
||||
expect(document.querySelector("textarea")).not.toBeInTheDocument();
|
||||
expect(trigger).toHaveFocus();
|
||||
});
|
||||
|
||||
it("reports failure when the fallback cannot create its selection", async () => {
|
||||
Object.defineProperty(navigator, "clipboard", { configurable: true, value: undefined });
|
||||
vi.spyOn(document, "createElement").mockImplementationOnce(() => {
|
||||
throw new Error("DOM error");
|
||||
});
|
||||
|
||||
it("should handle navigator.clipboard.writeText being undefined", async () => {
|
||||
Object.assign(navigator, {
|
||||
clipboard: {}, // clipboard exists but writeText doesn't
|
||||
});
|
||||
|
||||
const result = await copyToClipboard("test text");
|
||||
|
||||
expect(document.execCommand).toHaveBeenCalledWith("copy");
|
||||
expect(result).toBe(true);
|
||||
});
|
||||
expect(await copyToClipboard(text)).toBe(false);
|
||||
expect(mockToast.fromError).toHaveBeenCalledWith("Failed to copy to clipboard");
|
||||
expect(mockToast.success).not.toHaveBeenCalled();
|
||||
expect(trigger).toHaveFocus();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -103,19 +103,20 @@ const fallbackCopyToClipboard = (text: string, messageText: string): boolean =>
|
|||
textArea.setAttribute("readonly", "");
|
||||
|
||||
document.body.appendChild(textArea);
|
||||
textArea.focus();
|
||||
textArea.select();
|
||||
try {
|
||||
textArea.focus();
|
||||
textArea.select();
|
||||
|
||||
const successful = document.execCommand("copy");
|
||||
document.body.removeChild(textArea);
|
||||
if (previouslyFocused instanceof HTMLElement) previouslyFocused.focus();
|
||||
|
||||
if (successful) {
|
||||
toast.success(messageText);
|
||||
return true;
|
||||
} else {
|
||||
throw new Error("execCommand failed");
|
||||
if (document.execCommand("copy")) {
|
||||
toast.success(messageText);
|
||||
return true;
|
||||
}
|
||||
} finally {
|
||||
textArea.remove();
|
||||
if (previouslyFocused instanceof HTMLElement) previouslyFocused.focus();
|
||||
}
|
||||
toast.fromError("Failed to copy to clipboard");
|
||||
return false;
|
||||
} catch (err) {
|
||||
toast.fromError("Failed to copy to clipboard");
|
||||
console.error("Failed to copy: ", err);
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue