mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
chore: merge main with upstream OpenTelemetry test fix
This commit is contained in:
commit
3ce7c8529b
47 changed files with 2713 additions and 153 deletions
|
|
@ -1,14 +1,14 @@
|
|||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
if [ "${GITHUB_ACTIONS:-}" = true ]; then
|
||||
echo "Integration contracts are owned by CircleCI" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
suite="${1:?integration suite required}"
|
||||
mode="${2:-standard}"
|
||||
side="${3:-}"
|
||||
requested_workers="${INTEGRATION_WORKERS:-}"
|
||||
license_env=()
|
||||
if [ -n "${LITELLM_LICENSE:-}" ]; then
|
||||
license_env=("LITELLM_LICENSE=$LITELLM_LICENSE")
|
||||
fi
|
||||
if [ "$mode" = replica ]; then
|
||||
results="test-results/integration-${suite}-replica"
|
||||
elif [ "$mode" = parity ]; then
|
||||
|
|
@ -84,7 +84,7 @@ if [ "$suite" = browser ]; then
|
|||
export LITELLM_UI_PATH="$PWD/ui/litellm-dashboard/out"
|
||||
test -f "$LITELLM_UI_PATH/index.html"
|
||||
fi
|
||||
export INTEGRATION_SEED="$(.venv/bin/python -c 'import hashlib,os; print(int(hashlib.sha256((os.environ.get("CIRCLE_SHA1", "local") + os.environ.get("CIRCLE_WORKFLOW_ID", "local")).encode()).hexdigest()[:8],16))')"
|
||||
export INTEGRATION_SEED="$(.venv/bin/python -c 'import hashlib,os; print(int(hashlib.sha256(((os.environ.get("CIRCLE_SHA1") or os.environ.get("GITHUB_SHA") or "local") + (os.environ.get("CIRCLE_WORKFLOW_ID") or os.environ.get("GITHUB_RUN_ID") or "local")).encode()).hexdigest()[:8],16))')"
|
||||
export INTEGRATION_ORDER_SEED="$INTEGRATION_SEED"
|
||||
|
||||
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma > "$results/prisma-generate.log" 2>&1
|
||||
|
|
@ -107,7 +107,7 @@ sudo iptables -N integration_only
|
|||
guard_created=true
|
||||
sudo iptables -A integration_only -o lo -j ACCEPT
|
||||
sudo iptables -A integration_only -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
|
||||
for service in postgres-db redis-cache; do
|
||||
for service in "${INTEGRATION_POSTGRES_CONTAINER:-postgres-db}" "${INTEGRATION_REDIS_CONTAINER:-redis-cache}"; do
|
||||
address="$(docker inspect --format '{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}' "$service")"
|
||||
sudo iptables -A integration_only -d "$address" -j ACCEPT
|
||||
done
|
||||
|
|
@ -145,6 +145,9 @@ fi
|
|||
if [ "$suite" = mcp ]; then
|
||||
export INTEGRATION_WORKERS=2 INTEGRATION_COVERAGE=1
|
||||
fi
|
||||
if [ -n "$requested_workers" ]; then
|
||||
export INTEGRATION_WORKERS="$requested_workers"
|
||||
fi
|
||||
isolated_workers=false
|
||||
if [ "${INTEGRATION_WORKERS:-1}" -gt 1 ]; then
|
||||
isolated_workers=true
|
||||
|
|
@ -185,7 +188,7 @@ start_proxy() {
|
|||
"${database_env[@]}" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
|
||||
INTEGRATION_UPSTREAM_URL="$INTEGRATION_UPSTREAM_URL" \
|
||||
LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" LITELLM_SALT_KEY="$LITELLM_SALT_KEY" LITELLM_UI_PATH="$LITELLM_UI_PATH" PROXY_BASE_URL="http://127.0.0.1:$port" \
|
||||
LITELLM_LICENSE="${LITELLM_LICENSE:-}" \
|
||||
"${license_env[@]}" \
|
||||
LITELLM_MODE=PRODUCTION STORE_MODEL_IN_DB=True LITELLM_ENABLE_MCP_STDIO=true "${cost_map_env[@]}" \
|
||||
AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 COVERAGE_FILE="$coverage_data" \
|
||||
"${proxy_command[@]}" --config tests/integration/proxy_config.yaml \
|
||||
|
|
@ -257,7 +260,7 @@ env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
|
|||
"${shared_services[@]}" "${worker_coverage[@]}" \
|
||||
INTEGRATION_WORKERS="${INTEGRATION_WORKERS:-1}" \
|
||||
LITELLM_MODE=PRODUCTION \
|
||||
LITELLM_LICENSE="${LITELLM_LICENSE:-}" \
|
||||
"${license_env[@]}" \
|
||||
INTEGRATION_SEED="$INTEGRATION_SEED" \
|
||||
INTEGRATION_ORDER_SEED="$INTEGRATION_ORDER_SEED" \
|
||||
INTEGRATION_CASE_SHARDS="${CIRCLE_NODE_TOTAL:-1}" INTEGRATION_CASE_SHARD_INDEX="${CIRCLE_NODE_INDEX:-0}" \
|
||||
|
|
|
|||
119
.github/scripts/assert_secret_absent.py
vendored
Normal file
119
.github/scripts/assert_secret_absent.py
vendored
Normal file
|
|
@ -0,0 +1,119 @@
|
|||
import base64
|
||||
import io
|
||||
import os
|
||||
import sys
|
||||
import urllib.parse
|
||||
import zipfile
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from functools import reduce
|
||||
from itertools import chain
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final, TextIO
|
||||
|
||||
SECRET_VARIABLE: Final = "SCANNED_SECRET"
|
||||
MIN_ENCODED_LENGTH: Final = 12
|
||||
MAX_ARCHIVE_DEPTH: Final = 3
|
||||
MASK: Final = "***"
|
||||
|
||||
Forms = tuple[tuple[str, str], ...]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Hit:
|
||||
location: str
|
||||
form: str
|
||||
|
||||
|
||||
def _aligned_encodings(raw: bytes, name: str, encode: Callable[[bytes], bytes]) -> Iterator[tuple[str, str]]:
|
||||
for offset in range(3):
|
||||
chunk: Final = raw[offset:]
|
||||
yield f"{name} offset {offset}", encode(chunk[: len(chunk) - len(chunk) % 3]).decode()
|
||||
|
||||
|
||||
def secret_forms(secret: str) -> Forms:
|
||||
encoded: Final = (
|
||||
("url-encoded", urllib.parse.quote(secret, safe="")),
|
||||
*_aligned_encodings(secret.encode(), "base64", base64.b64encode),
|
||||
*_aligned_encodings(secret.encode(), "base64url", base64.urlsafe_b64encode),
|
||||
)
|
||||
by_value: Final = MappingProxyType(
|
||||
{value: name for name, value in encoded if value != secret and len(value) >= MIN_ENCODED_LENGTH}
|
||||
)
|
||||
return (("raw", secret), *((name, value) for value, name in by_value.items()))
|
||||
|
||||
|
||||
def mask(text: str, forms: Forms) -> str:
|
||||
return reduce(lambda masked, form: masked.replace(form[1], MASK), forms, text)
|
||||
|
||||
|
||||
def _name_hits(location: str, name: str, forms: Forms) -> Iterator[Hit]:
|
||||
for form, value in forms:
|
||||
if value in name:
|
||||
yield Hit(mask(location, forms), f"{form} in name")
|
||||
|
||||
|
||||
def _content_hits(location: str, data: bytes, forms: Forms, depth: int) -> Iterator[Hit]:
|
||||
for form, value in forms:
|
||||
if value.encode() in data:
|
||||
yield Hit(mask(location, forms), form)
|
||||
if depth >= MAX_ARCHIVE_DEPTH or not zipfile.is_zipfile(io.BytesIO(data)):
|
||||
return
|
||||
with zipfile.ZipFile(io.BytesIO(data)) as archive:
|
||||
for member in archive.infolist():
|
||||
if member.is_dir():
|
||||
continue
|
||||
yield from _content_hits(f"{location}!{member.filename}", archive.read(member), forms, depth + 1)
|
||||
|
||||
|
||||
def _walk(root: Path) -> Iterator[tuple[Path, bool]]:
|
||||
for directory, subdirectories, files in os.walk(root):
|
||||
yield from ((Path(directory, name), True) for name in subdirectories if Path(directory, name).is_symlink())
|
||||
yield from ((Path(directory, name), False) for name in files)
|
||||
|
||||
|
||||
def _path_hits(root: Path, path: Path, linked_directory: bool, forms: Forms) -> Iterator[Hit]:
|
||||
location: Final = str(path.relative_to(root))
|
||||
if linked_directory:
|
||||
yield Hit(mask(location, forms), "directory symlink, not scanned")
|
||||
return
|
||||
yield from _name_hits(location, location, forms)
|
||||
yield from _content_hits(location, path.read_bytes(), forms, 0)
|
||||
|
||||
|
||||
def scan(root: Path, secret: str) -> tuple[int, tuple[Hit, ...]]:
|
||||
forms: Final = secret_forms(secret)
|
||||
entries: Final = tuple(_walk(root))
|
||||
hits: Final = tuple(chain.from_iterable(_path_hits(root, path, linked, forms) for path, linked in entries))
|
||||
return len(entries), hits
|
||||
|
||||
|
||||
def main(argv: tuple[str, ...], environ: Mapping[str, str], out: TextIO) -> int:
|
||||
if len(argv) != 1:
|
||||
print(f"usage: assert_secret_absent.py <directory>, with the secret in ${SECRET_VARIABLE}", file=out)
|
||||
return 2
|
||||
secret: Final = environ.get(SECRET_VARIABLE, "")
|
||||
root: Final = Path(argv[0])
|
||||
if not secret:
|
||||
print(f"{SECRET_VARIABLE} is empty, nothing to scan for", file=out)
|
||||
return 0
|
||||
if not root.is_dir():
|
||||
print(f"{root} is not a directory, nothing to scan", file=out)
|
||||
return 0
|
||||
try:
|
||||
scanned, hits = scan(root, secret)
|
||||
except Exception as error:
|
||||
print(f"scan failed with {type(error).__name__}, treating the results as unsafe", file=out)
|
||||
return 1
|
||||
for hit in hits:
|
||||
print(f"secret found: {hit.location} ({hit.form})", file=out)
|
||||
if hits:
|
||||
print(f"{len(hits)} copies of the secret under {root}, refusing to upload", file=out)
|
||||
return 1
|
||||
print(f"no copy of the secret in {scanned} files under {root}", file=out)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main(tuple(sys.argv[1:]), os.environ, sys.stdout))
|
||||
199
.github/workflows/test-integration.yml
vendored
Normal file
199
.github/workflows/test-integration.yml
vendored
Normal file
|
|
@ -0,0 +1,199 @@
|
|||
name: integration
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
types: [labeled]
|
||||
schedule:
|
||||
- cron: "5 8-20/3 * * 1-6"
|
||||
timezone: America/Los_Angeles
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
suite:
|
||||
description: Integration group to run
|
||||
type: choice
|
||||
options:
|
||||
- all
|
||||
- providers
|
||||
- extensions
|
||||
- accounting
|
||||
- management
|
||||
- management-replica
|
||||
- database
|
||||
- database-replica
|
||||
- mcp
|
||||
- security
|
||||
- cost
|
||||
- sdk
|
||||
- browser
|
||||
default: all
|
||||
workers:
|
||||
description: Workers per group, blank keeps each group's default
|
||||
type: string
|
||||
default: ""
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}-${{ github.event.label.name }}-${{ inputs.suite }}-${{ inputs.workers }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
integration:
|
||||
name: integration ${{ matrix.suite }}
|
||||
if: github.event_name != 'pull_request' || github.event.label.name == 'run-ci'
|
||||
environment: integration-license
|
||||
runs-on: ${{ fromJSON('{"providers":"ubuntu-latest-64-cores","extensions":"ubuntu-latest-64-cores","accounting":"ubuntu-latest-32-cores","management":"ubuntu-latest-32-cores","management-replica":"ubuntu-latest-32-cores"}')[matrix.suite] || 'ubuntu-latest-16-cores' }}
|
||||
timeout-minutes: 60
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
suite: ${{ fromJSON((inputs.suite || 'all') == 'all' && '["providers","extensions","accounting","management","management-replica","database","database-replica","mcp","security","cost","sdk","browser"]' || format('["{0}"]', inputs.suite)) }}
|
||||
env:
|
||||
UV_PYTHON: "3.12"
|
||||
INTEGRATION_WORKERS: ${{ inputs.workers || fromJSON('{"providers":"32","extensions":"32","accounting":"16","management":"16","management-replica":"16","cost":"16","mcp":"8","security":"8","database":"4","database-replica":"4"}')[matrix.suite] || '' }}
|
||||
services:
|
||||
postgres:
|
||||
image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5
|
||||
command: -c max_connections=1000 -c shared_preload_libraries=pg_stat_statements -c pg_stat_statements.track=all -c pg_stat_statements.max=20000
|
||||
env:
|
||||
POSTGRES_USER: postgres
|
||||
POSTGRES_PASSWORD: postgres
|
||||
POSTGRES_DB: circle_test
|
||||
ports:
|
||||
- 5432:5432
|
||||
options: >-
|
||||
--shm-size 1g
|
||||
--health-cmd "pg_isready -U postgres"
|
||||
--health-interval 5s
|
||||
--health-timeout 5s
|
||||
--health-retries 20
|
||||
redis:
|
||||
image: redis:7-alpine@sha256:7aec734b2bb298a1d769fd8729f13b8514a41bf90fcdd1f38ec52267fbaa8ee6
|
||||
ports:
|
||||
- 6379:6379
|
||||
options: >-
|
||||
--health-cmd "redis-cli ping"
|
||||
--health-interval 5s
|
||||
--health-timeout 5s
|
||||
--health-retries 20
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Python
|
||||
timeout-minutes: 3
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: ${{ env.UV_PYTHON }}
|
||||
|
||||
- name: Set up uv
|
||||
timeout-minutes: 3
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Cache the Rust build
|
||||
timeout-minutes: 5
|
||||
uses: ./.github/actions/cache-cargo-build
|
||||
|
||||
- name: Cache Prisma binaries
|
||||
timeout-minutes: 5
|
||||
uses: ./.github/actions/cache-prisma-binaries
|
||||
|
||||
- name: Install dependencies
|
||||
timeout-minutes: 20
|
||||
run: .github/scripts/uv_sync_with_retries.sh --frozen --all-groups --all-extras
|
||||
|
||||
- name: Install redis-server 7.2.16
|
||||
timeout-minutes: 3
|
||||
run: |
|
||||
mkdir -p "$RUNNER_TEMP/redis-bin"
|
||||
cid="$(docker create redis:7.2.16@sha256:0637954999d01b7c9ce9167db2da50656e2590d3b884f1c600c5f63bb6e6773c)"
|
||||
docker cp "${cid}:/usr/local/bin/redis-server" "$RUNNER_TEMP/redis-bin/redis-server"
|
||||
docker rm "$cid"
|
||||
"$RUNNER_TEMP/redis-bin/redis-server" --version | grep -q 'v=7.2.16'
|
||||
echo "$RUNNER_TEMP/redis-bin" >> "$GITHUB_PATH"
|
||||
|
||||
- name: Install MCP SDK1 peer
|
||||
if: matrix.suite == 'mcp'
|
||||
timeout-minutes: 5
|
||||
run: |
|
||||
uv venv --python 3.12 .venv-mcp-peer
|
||||
uv pip install --no-config --exclude-newer "3 days" --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1'
|
||||
echo "MCP_TEST_PEER_PYTHON=$PWD/.venv-mcp-peer/bin/python" >> "$GITHUB_ENV"
|
||||
|
||||
- name: Set up Node.js
|
||||
if: matrix.suite == 'browser'
|
||||
timeout-minutes: 5
|
||||
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0
|
||||
with:
|
||||
node-version-file: ui/litellm-dashboard/.nvmrc
|
||||
cache: "npm"
|
||||
cache-dependency-path: |
|
||||
ui/litellm-dashboard/package-lock.json
|
||||
tests/e2e/ui/package-lock.json
|
||||
|
||||
- name: Install browser dependencies and build the dashboard
|
||||
if: matrix.suite == 'browser'
|
||||
timeout-minutes: 20
|
||||
run: |
|
||||
(cd ui/litellm-dashboard && npm ci)
|
||||
(cd tests/e2e/ui && npm ci)
|
||||
sudo env PATH="$PATH" DEBIAN_FRONTEND=noninteractive \
|
||||
node tests/e2e/ui/node_modules/@playwright/test/cli.js install-deps chromium
|
||||
node tests/e2e/ui/node_modules/@playwright/test/cli.js install chromium
|
||||
(cd ui/litellm-dashboard && NEXT_TELEMETRY_DISABLED=1 npm run build)
|
||||
|
||||
- name: Run integration contracts
|
||||
timeout-minutes: 45
|
||||
env:
|
||||
SUITE: ${{ matrix.suite }}
|
||||
INTEGRATION_POSTGRES_CONTAINER: ${{ job.services.postgres.id }}
|
||||
INTEGRATION_REDIS_CONTAINER: ${{ job.services.redis.id }}
|
||||
LITELLM_LICENSE: ${{ secrets.LITELLM_LICENSE }}
|
||||
run: |
|
||||
case "$SUITE" in
|
||||
*-replica) bash .circleci/scripts/run_integration.sh "${SUITE%-replica}" replica ;;
|
||||
*) bash .circleci/scripts/run_integration.sh "$SUITE" ;;
|
||||
esac
|
||||
|
||||
- name: Verify every case ran exactly once
|
||||
if: ${{ !cancelled() && matrix.suite != 'browser' }}
|
||||
timeout-minutes: 2
|
||||
run: |
|
||||
record="$(find test-results -mindepth 2 -maxdepth 2 -name execution.json -print -quit)"
|
||||
test -n "$record"
|
||||
mkdir -p "$RUNNER_TEMP/shards/0"
|
||||
cp "$record" "$RUNNER_TEMP/shards/0/execution.json"
|
||||
.venv/bin/python -m tests.integration._support.sharding "$RUNNER_TEMP/shards" --shards 1
|
||||
|
||||
- name: Report MCP coverage
|
||||
if: ${{ !cancelled() && matrix.suite == 'mcp' }}
|
||||
timeout-minutes: 5
|
||||
run: |
|
||||
coverage_dir=test-results/integration-mcp/coverage
|
||||
test -f "$coverage_dir/coverage.txt" && exit 0
|
||||
export COVERAGE_FILE="$PWD/$coverage_dir/data"
|
||||
.venv/bin/python -m coverage combine --rcfile=tests/integration/mcp_coverage.toml
|
||||
.venv/bin/python -m coverage report --rcfile=tests/integration/mcp_coverage.toml > "$coverage_dir/coverage.txt"
|
||||
.venv/bin/python -m coverage html --rcfile=tests/integration/mcp_coverage.toml -d "$coverage_dir/html"
|
||||
tail -n 1 "$coverage_dir/coverage.txt"
|
||||
|
||||
- name: Check results for the license
|
||||
id: secret-scan
|
||||
if: ${{ !cancelled() }}
|
||||
timeout-minutes: 5
|
||||
env:
|
||||
SCANNED_SECRET: ${{ secrets.LITELLM_LICENSE }}
|
||||
run: python .github/scripts/assert_secret_absent.py test-results
|
||||
|
||||
- name: Upload results
|
||||
if: ${{ !cancelled() && steps.secret-scan.outcome == 'success' }}
|
||||
timeout-minutes: 10
|
||||
uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1
|
||||
with:
|
||||
name: integration-${{ matrix.suite }}-${{ github.run_id }}-${{ github.run_attempt }}
|
||||
path: test-results/
|
||||
retention-days: 14
|
||||
|
|
@ -29,6 +29,9 @@ it does not belong to. Queued clients are therefore bucketed by what it takes to
|
|||
close them, and each bucket is ordered by deadline, so a reap walks the entries
|
||||
that are due rather than the whole queue.
|
||||
|
||||
A bucket whose loop has been collected is dropped, and its clients are left to the
|
||||
collector, as they are when the queue is full.
|
||||
|
||||
The queue holds its clients weakly, so waiting out a grace window never keeps
|
||||
alive anything the collector would have reclaimed first.
|
||||
"""
|
||||
|
|
@ -36,6 +39,7 @@ alive anything the collector would have reclaimed first.
|
|||
import asyncio
|
||||
import contextlib
|
||||
import inspect
|
||||
import itertools
|
||||
import threading
|
||||
import time
|
||||
import weakref
|
||||
|
|
@ -55,6 +59,9 @@ _CLOSABLE_ANYWHERE: Final = "closable-anywhere"
|
|||
_CLOSABLE_ON_ANY_LOOP: Final = "closable-on-any-loop"
|
||||
|
||||
_BucketKey = str | int
|
||||
_LOOP_TOKENS: Final[weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, int]] = weakref.WeakKeyDictionary()
|
||||
_LOOP_TOKEN_COUNTER: Final = itertools.count()
|
||||
_LOOP_TOKEN_LOCK: Final = threading.Lock()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -65,12 +72,12 @@ class _PendingClose:
|
|||
collector would otherwise have reclaimed first.
|
||||
|
||||
``needs_loop`` is set for a client whose close is a coroutine; those can only
|
||||
be closed from the event loop they were evicted on, recorded in ``loop_id``.
|
||||
be closed from the event loop they were evicted on, recorded in ``loop_token``.
|
||||
A client that closes synchronously carries neither constraint.
|
||||
"""
|
||||
|
||||
client_ref: "weakref.ref[object]"
|
||||
loop_id: int | None
|
||||
loop_token: int | None
|
||||
needs_loop: bool
|
||||
close_after: float
|
||||
|
||||
|
|
@ -79,16 +86,24 @@ def _bucket_key(pending: _PendingClose) -> _BucketKey:
|
|||
"""Which reaps can close this entry: any at all, any running a loop, or one loop's."""
|
||||
if not pending.needs_loop:
|
||||
return _CLOSABLE_ANYWHERE
|
||||
if pending.loop_id is None:
|
||||
if pending.loop_token is None:
|
||||
return _CLOSABLE_ON_ANY_LOOP
|
||||
return pending.loop_id
|
||||
return pending.loop_token
|
||||
|
||||
|
||||
def _running_loop_id() -> int | None:
|
||||
def running_loop_token() -> int | None:
|
||||
try:
|
||||
return id(asyncio.get_running_loop())
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
return None
|
||||
try:
|
||||
token: Final[int | None] = _LOOP_TOKENS.get(loop)
|
||||
if token is not None:
|
||||
return token
|
||||
with _LOOP_TOKEN_LOCK:
|
||||
return _LOOP_TOKENS.setdefault(loop, next(_LOOP_TOKEN_COUNTER))
|
||||
except TypeError:
|
||||
return id(loop)
|
||||
|
||||
|
||||
def _close_function(client: object) -> Callable[[], object] | None:
|
||||
|
|
@ -186,6 +201,8 @@ class EvictedClientCloser:
|
|||
self._clock = clock
|
||||
self._owned: weakref.WeakSet[object] = weakref.WeakSet()
|
||||
self._buckets: dict[_BucketKey, deque[_PendingClose]] = {} # mutable-ok: deadline-ordered queues
|
||||
self._dead_loop_tokens: deque[int] = deque() # mutable-ok: loop finalizers append without taking the queue lock
|
||||
self._watched_loop_tokens: set[int] = set() # mutable-ok: tracks tokens until their dead buckets are discarded
|
||||
self._pending_count = 0
|
||||
self._queue_lock = threading.Lock() # the cache is reachable from every worker thread's loop
|
||||
self._close_tasks: set[asyncio.Task[None]] = set() # mutable-ok: strong refs to running closes
|
||||
|
|
@ -216,13 +233,19 @@ class EvictedClientCloser:
|
|||
close_fn: Final = _close_function(client)
|
||||
if close_fn is None:
|
||||
return
|
||||
loop_token: Final = running_loop_token()
|
||||
needs_loop: Final = inspect.iscoroutinefunction(close_fn)
|
||||
if self._pending_count >= self._max_pending:
|
||||
return
|
||||
self._discard_dead_loops()
|
||||
if self._pending_count >= self._max_pending:
|
||||
return
|
||||
if needs_loop and loop_token is not None:
|
||||
self._watch_loop(loop_token)
|
||||
self._enqueue(
|
||||
_PendingClose(
|
||||
client_ref=weakref.ref(client),
|
||||
loop_id=_running_loop_id(),
|
||||
needs_loop=inspect.iscoroutinefunction(close_fn),
|
||||
loop_token=loop_token,
|
||||
needs_loop=needs_loop,
|
||||
close_after=self._clock() + self._grace_seconds,
|
||||
)
|
||||
)
|
||||
|
|
@ -236,7 +259,7 @@ class EvictedClientCloser:
|
|||
if not self._pending_count:
|
||||
return
|
||||
now: Final = self._clock()
|
||||
for pending in self._take_due(_running_loop_id(), now):
|
||||
for pending in self._take_due(running_loop_token(), now):
|
||||
client = pending.client_ref()
|
||||
if client is None:
|
||||
continue
|
||||
|
|
@ -249,6 +272,29 @@ class EvictedClientCloser:
|
|||
def pending_count(self) -> int:
|
||||
return self._pending_count
|
||||
|
||||
def _watch_loop(self, loop_token: int) -> None:
|
||||
if loop_token in self._watched_loop_tokens:
|
||||
return
|
||||
try:
|
||||
weakref.finalize(asyncio.get_running_loop(), self._dead_loop_tokens.append, loop_token)
|
||||
except TypeError:
|
||||
return
|
||||
self._watched_loop_tokens.add(loop_token)
|
||||
|
||||
def _discard_dead_loops(self) -> None:
|
||||
with self._queue_lock:
|
||||
self._discard_dead_loops_locked()
|
||||
|
||||
def _discard_dead_loops_locked(self) -> None:
|
||||
while self._dead_loop_tokens:
|
||||
self._discard_dead_loop_token_locked(self._dead_loop_tokens.popleft())
|
||||
|
||||
def _discard_dead_loop_token_locked(self, loop_token: int) -> None:
|
||||
self._watched_loop_tokens.discard(loop_token)
|
||||
bucket: Final = self._buckets.pop(loop_token, None)
|
||||
if bucket is not None:
|
||||
self._pending_count -= len(bucket)
|
||||
|
||||
def _enqueue(self, pending: _PendingClose) -> None:
|
||||
"""Append to the entry's bucket, dropping any dead entries it queues behind.
|
||||
|
||||
|
|
@ -257,6 +303,7 @@ class EvictedClientCloser:
|
|||
the front rather than having to be searched for.
|
||||
"""
|
||||
with self._queue_lock:
|
||||
self._discard_dead_loops_locked()
|
||||
bucket: Final = self._buckets.setdefault(_bucket_key(pending), deque())
|
||||
while bucket and bucket[0].client_ref() is None:
|
||||
bucket.popleft()
|
||||
|
|
@ -264,9 +311,12 @@ class EvictedClientCloser:
|
|||
bucket.append(pending)
|
||||
self._pending_count += 1
|
||||
|
||||
def _take_due(self, loop_id: int | None, now: float) -> tuple[_PendingClose, ...]:
|
||||
buckets = (_CLOSABLE_ANYWHERE,) if loop_id is None else (_CLOSABLE_ANYWHERE, _CLOSABLE_ON_ANY_LOOP, loop_id)
|
||||
def _take_due(self, loop_token: int | None, now: float) -> tuple[_PendingClose, ...]:
|
||||
buckets = (
|
||||
(_CLOSABLE_ANYWHERE,) if loop_token is None else (_CLOSABLE_ANYWHERE, _CLOSABLE_ON_ANY_LOOP, loop_token)
|
||||
)
|
||||
with self._queue_lock:
|
||||
self._discard_dead_loops_locked()
|
||||
return tuple(pending for key in buckets for pending in self._drain_locked(key, now))
|
||||
|
||||
def _drain_locked(self, key: _BucketKey, now: float) -> Iterator[_PendingClose]:
|
||||
|
|
|
|||
|
|
@ -2,10 +2,9 @@
|
|||
Add the event loop to the cache key, to prevent event loop closed errors.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from typing import Final
|
||||
|
||||
from .evicted_client_closer import EvictedClientCloser, default_evicted_client_closer
|
||||
from .evicted_client_closer import EvictedClientCloser, default_evicted_client_closer, running_loop_token
|
||||
from .in_memory_cache import InMemoryCache
|
||||
|
||||
|
||||
|
|
@ -43,24 +42,22 @@ class LLMClientCache(InMemoryCache):
|
|||
self.evicted_client_closer.schedule(evicted)
|
||||
self.evicted_client_closer.reap()
|
||||
|
||||
def update_cache_key_with_event_loop(self, key):
|
||||
def update_cache_key_with_event_loop(self, key: str) -> str:
|
||||
"""
|
||||
Add the event loop to the cache key, to prevent event loop closed errors.
|
||||
If none, use the key as is.
|
||||
"""
|
||||
try:
|
||||
event_loop: Final = asyncio.get_running_loop()
|
||||
stringified_event_loop: Final = str(id(event_loop))
|
||||
return f"{key}-{stringified_event_loop}"
|
||||
except RuntimeError: # handle no current running event loop
|
||||
return key
|
||||
token: Final = running_loop_token()
|
||||
return f"{key}-{token}" if token is not None else key
|
||||
|
||||
def set_cache(self, key: str, value: object, litellm_owned_client: bool = False, **kwargs):
|
||||
def _scoped_key(self, key: str, loop_bound: bool) -> str:
|
||||
return self.update_cache_key_with_event_loop(key) if loop_bound else key
|
||||
|
||||
def set_cache(self, key: str, value: object, litellm_owned_client: bool = False, loop_bound: bool = True, **kwargs):
|
||||
"""``litellm_owned_client`` marks a client litellm built, so it may be closed once evicted."""
|
||||
if litellm_owned_client:
|
||||
self.evicted_client_closer.mark_owned(value)
|
||||
key = self.update_cache_key_with_event_loop(key)
|
||||
return super().set_cache(key, value, **kwargs)
|
||||
return super().set_cache(self._scoped_key(key, loop_bound), value, **kwargs)
|
||||
|
||||
async def async_set_cache(self, key: str, value: object, litellm_owned_client: bool = False, **kwargs):
|
||||
if litellm_owned_client:
|
||||
|
|
@ -68,13 +65,13 @@ class LLMClientCache(InMemoryCache):
|
|||
key = self.update_cache_key_with_event_loop(key)
|
||||
return await super().async_set_cache(key, value, **kwargs)
|
||||
|
||||
def get_cache(self, key, **kwargs):
|
||||
key = self.update_cache_key_with_event_loop(key)
|
||||
def get_cache(self, key: str, loop_bound: bool = True, **kwargs):
|
||||
scoped_key: Final = self._scoped_key(key, loop_bound)
|
||||
self.evicted_client_closer.reap()
|
||||
|
||||
return super().get_cache(key, **kwargs)
|
||||
return super().get_cache(scoped_key, **kwargs)
|
||||
|
||||
async def async_get_cache(self, key, **kwargs):
|
||||
async def async_get_cache(self, key: str, **kwargs):
|
||||
key = self.update_cache_key_with_event_loop(key)
|
||||
|
||||
return await super().async_get_cache(key, **kwargs)
|
||||
|
|
|
|||
|
|
@ -102,18 +102,14 @@ class AgentOps(OpenTelemetry):
|
|||
"Connection": "keep-alive",
|
||||
}
|
||||
|
||||
client: Final = get_httpx_client()
|
||||
try:
|
||||
response: Final = client.post(
|
||||
url=auth_endpoint,
|
||||
headers=headers,
|
||||
json={"api_key": api_key},
|
||||
timeout=10,
|
||||
)
|
||||
response: Final = get_httpx_client().post(
|
||||
url=auth_endpoint,
|
||||
headers=headers,
|
||||
json={"api_key": api_key},
|
||||
timeout=10,
|
||||
)
|
||||
|
||||
if response.status_code != 200:
|
||||
raise Exception(f"Failed to fetch auth token: {response.text}")
|
||||
if response.status_code != 200:
|
||||
raise Exception(f"Failed to fetch auth token: {response.text}")
|
||||
|
||||
return _DECODED_JSON.validate_python(response.json())
|
||||
finally:
|
||||
client.close()
|
||||
return _DECODED_JSON.validate_python(response.json())
|
||||
|
|
|
|||
|
|
@ -35,6 +35,8 @@ from litellm.types.utils import StandardCallbackDynamicParams
|
|||
if TYPE_CHECKING:
|
||||
from opentelemetry.trace import Span as _Span
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
Span = _Span
|
||||
else:
|
||||
Span = Any
|
||||
|
|
@ -126,6 +128,26 @@ def _starts_its_own_trace(span: object, kwargs: Mapping[str, object]) -> bool:
|
|||
)
|
||||
|
||||
|
||||
def _a_destination_records_the_call(kwargs: Mapping[str, object]) -> bool:
|
||||
"""Whether a key's or team's ``langfuse_otel`` destination gets this call from the OTel v2 fan-out, which records
|
||||
every call but one whose logging payload shows it never reached a provider."""
|
||||
from litellm.integrations.otel.model.metadata import LLMCallEvent
|
||||
|
||||
if not _a_destination_carries_the_trace():
|
||||
return False
|
||||
call: Final = LLMCallEvent.from_dict(kwargs)
|
||||
return call.payload is None or call.reached_a_provider
|
||||
|
||||
|
||||
def _a_destination_carries_the_trace() -> bool:
|
||||
"""Whether this request resolved a key's or team's ``langfuse_otel`` destination, whose fan-out sends it the whole
|
||||
trace."""
|
||||
from litellm.integrations.otel.model.config import ExporterOwner
|
||||
from litellm.integrations.otel.plumbing.context import destination_backends
|
||||
|
||||
return ExporterOwner.LANGFUSE_OTEL in destination_backends()
|
||||
|
||||
|
||||
class LangfuseOtelLogger(OpenTelemetry):
|
||||
def __init__(self, config=None, *args, **kwargs):
|
||||
# Prevent LangfuseOtelLogger from modifying global environment variables by constructing config manually
|
||||
|
|
@ -541,6 +563,27 @@ class LangfuseOtelLogger(OpenTelemetry):
|
|||
"""
|
||||
return None
|
||||
|
||||
def _emit_once(
|
||||
self,
|
||||
kwargs: dict[str, object], # mutable-ok: the v1 dedupe marker is written into the request's kwargs
|
||||
*scope: object,
|
||||
) -> bool:
|
||||
"""Whether to emit this span: once per request, and never for a call a Langfuse destination records."""
|
||||
return not _a_destination_records_the_call(kwargs) and super()._emit_once(kwargs, *scope) # pyright: ignore[reportUnknownMemberType] # the v1 hook takes a bare dict
|
||||
|
||||
async def async_post_call_failure_hook(
|
||||
self,
|
||||
request_data: dict[str, object], # mutable-ok: the v1 hook contract
|
||||
original_exception: Exception,
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
traceback_str: str | None = None,
|
||||
) -> None:
|
||||
"""Record the proxy failure of a request no Langfuse destination carries."""
|
||||
if not _a_destination_carries_the_trace():
|
||||
await super().async_post_call_failure_hook( # pyright: ignore[reportUnknownMemberType] # the v1 hook takes a bare dict
|
||||
request_data, original_exception, user_api_key_dict, traceback_str
|
||||
)
|
||||
|
||||
async def async_service_success_hook(self, *args, **kwargs):
|
||||
"""
|
||||
Langfuse should not receive service success logs.
|
||||
|
|
|
|||
|
|
@ -13,9 +13,9 @@ from opentelemetry.trace.status import Status, StatusCode
|
|||
|
||||
from litellm.integrations.otel.mappers import resolve_mappers
|
||||
from litellm.integrations.otel.mappers.base import AttributeMapper, AttrValue, SpanData
|
||||
from litellm.integrations.otel.mappers.langfuse import LANGFUSE_REQUEST_METADATA_ATTRIBUTES
|
||||
from litellm.integrations.otel.mappers.langfuse import LANGFUSE_REQUEST_METADATA_ATTRIBUTES, LangfuseMapper
|
||||
from litellm.integrations.otel.mappers.openinference import fit_indexed_messages
|
||||
from litellm.integrations.otel.model.config import OpenTelemetryV2Config
|
||||
from litellm.integrations.otel.model.config import ExporterOwner, OpenTelemetryV2Config
|
||||
from litellm.integrations.otel.model.payloads import (
|
||||
GuardrailSpanData,
|
||||
LLMCallSpanData,
|
||||
|
|
@ -34,6 +34,7 @@ from litellm.integrations.otel.model.spans import (
|
|||
mcp_tool_call_span_name,
|
||||
service_span_name,
|
||||
)
|
||||
from litellm.integrations.otel.plumbing.context import destination_backends
|
||||
from litellm.integrations.otel.plumbing.events import GenAIEventRecorder
|
||||
from litellm.integrations.otel.plumbing.providers import to_otel_span_kind
|
||||
|
||||
|
|
@ -168,6 +169,9 @@ class SpanEmitter:
|
|||
self._mappers: list[AttributeMapper] = (
|
||||
list(mappers) if mappers is not None else resolve_mappers(config.mapper_names)
|
||||
)
|
||||
self._langfuse: Final = (
|
||||
None if any(isinstance(mapper, LangfuseMapper) for mapper in self._mappers) else LangfuseMapper()
|
||||
)
|
||||
# Bounded LRU (ordered by insertion / most-recent touch). Storing keys
|
||||
# only — the value is unused — so it behaves like a capped set.
|
||||
self._emitted: OrderedDict[tuple[str, SpanRole], None] = OrderedDict()
|
||||
|
|
@ -295,17 +299,27 @@ class SpanEmitter:
|
|||
)
|
||||
else None
|
||||
)
|
||||
mapped: Final = MappingProxyType(
|
||||
{key: value for mapper in self._mappers for key, value in mapper.map(data).items()}
|
||||
configured: Final = {key: value for mapper in self._mappers for key, value in mapper.map(data).items()}
|
||||
for_langfuse: Final = (
|
||||
dict(
|
||||
sorted(
|
||||
self._langfuse.map(data).items(),
|
||||
key=lambda item: item[0] not in LANGFUSE_REQUEST_METADATA_ATTRIBUTES,
|
||||
)
|
||||
)
|
||||
if self._langfuse is not None and ExporterOwner.LANGFUSE_OTEL in destination_backends()
|
||||
else _NO_ATTRIBUTES
|
||||
)
|
||||
mapped: Final = MappingProxyType({**configured, **for_langfuse})
|
||||
stamped_later: Final = error_attributes(error) if error else _NO_ATTRIBUTES
|
||||
reserved: Final = len(stamped_later.keys() - mapped.keys())
|
||||
overwrites: Final = carried_keys(span) & mapped.keys()
|
||||
# The Langfuse request metadata sheds first, so a span at its attribute limit loses nothing else to it.
|
||||
# The Langfuse request metadata, and the Langfuse vocabulary only a destination asked for, shed first,
|
||||
# so a span at its attribute limit loses nothing else to them.
|
||||
fitted: Final = fit_indexed_messages(
|
||||
MappingProxyType({key: value for key, value in mapped.items() if key not in overwrites}),
|
||||
attribute_budget(span, reserved),
|
||||
LANGFUSE_REQUEST_METADATA_ATTRIBUTES,
|
||||
LANGFUSE_REQUEST_METADATA_ATTRIBUTES.union(for_langfuse),
|
||||
)
|
||||
for key, value in mapped.items():
|
||||
if key in overwrites or key in fitted:
|
||||
|
|
|
|||
|
|
@ -198,6 +198,7 @@ class OpenTelemetryV2(CustomLogger):
|
|||
tracer_provider: TracerProvider | None = None,
|
||||
logger_provider: LoggerProvider | None = None,
|
||||
meter_provider: "MeterProvider | None" = None,
|
||||
console_default: bool = False,
|
||||
**kwargs: Unpack[_CustomLoggerOptions],
|
||||
) -> None:
|
||||
super().__init__(**kwargs)
|
||||
|
|
@ -206,7 +207,7 @@ class OpenTelemetryV2(CustomLogger):
|
|||
self._tracer_provider: TracerProvider = (
|
||||
tracer_provider
|
||||
if tracer_provider is not None
|
||||
else build_tracer_provider(self.config, tenant_overrides=True)
|
||||
else build_tracer_provider(self.config, tenant_overrides=True, console_default=console_default)
|
||||
)
|
||||
self.tracer: Tracer = get_tracer(self._tracer_provider, LITELLM_TRACER_NAME)
|
||||
self._metrics_recorder = self._init_metrics(meter_provider)
|
||||
|
|
@ -217,7 +218,9 @@ class OpenTelemetryV2(CustomLogger):
|
|||
mappers=resolve_mappers(self.config.mapper_names),
|
||||
event_recorder=self._init_events(logger_provider),
|
||||
)
|
||||
self._tenant_tracers = TenantTracerCache(self.config, callback_name, LITELLM_TRACER_NAME)
|
||||
self._tenant_tracers = TenantTracerCache(
|
||||
self.config, callback_name, LITELLM_TRACER_NAME, console_default=console_default
|
||||
)
|
||||
self._open_llm_calls: OrderedDict[str, _LLMCallSpan] = OrderedDict()
|
||||
self._init_otel_logger_on_litellm_proxy()
|
||||
|
||||
|
|
@ -541,9 +544,7 @@ class OpenTelemetryV2(CustomLogger):
|
|||
# gate rejection carries ``is_no_upstream_call`` and gets no span. A cache
|
||||
# hit never hands off to a provider, so ``cache_hit`` substitutes for the
|
||||
# upstream stamp, since the payload proves the call completed.
|
||||
if carrier is None and (
|
||||
call.is_no_upstream_call or call.payload is None or (not call.upstream_started and not call.cache_hit)
|
||||
):
|
||||
if carrier is None and not call.reached_a_provider:
|
||||
return None
|
||||
try:
|
||||
return self._finish_carrier(carrier, call, start_time, end_time)
|
||||
|
|
@ -1091,6 +1092,7 @@ def build_otel_v2_logger(
|
|||
logger_provider: LoggerProvider | None = None,
|
||||
meter_provider: "MeterProvider | None" = None,
|
||||
settings: Mapping[str, object] = MappingProxyType({}),
|
||||
console_default: bool = False,
|
||||
) -> OpenTelemetryV2:
|
||||
return _logger_class(config)(
|
||||
config=config,
|
||||
|
|
@ -1098,6 +1100,7 @@ def build_otel_v2_logger(
|
|||
tracer_provider=tracer_provider,
|
||||
logger_provider=logger_provider,
|
||||
meter_provider=meter_provider,
|
||||
console_default=console_default,
|
||||
**settings,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -103,6 +103,18 @@ class LangfuseMapper:
|
|||
),
|
||||
}
|
||||
|
||||
#: Every attribute this mapper can write on a model call.
|
||||
LLM_CALL_ATTRIBUTES: Final[frozenset[str]] = frozenset(
|
||||
{
|
||||
*_LLM_CALL_ATTRS,
|
||||
*_BLOB_ATTRS,
|
||||
LANGFUSE_TRACE_NAME,
|
||||
LANGFUSE_TRACE_USER_ID,
|
||||
LANGFUSE_TRACE_SESSION_ID,
|
||||
LANGFUSE_TRACE_TAGS,
|
||||
}
|
||||
)
|
||||
|
||||
def map(self, data: SpanData) -> AttributeMap:
|
||||
match data:
|
||||
case LLMCallSpanData():
|
||||
|
|
|
|||
|
|
@ -121,6 +121,17 @@ class ExporterSpec(LiteLLMBaseModel):
|
|||
)
|
||||
|
||||
|
||||
def is_unconfigured_placeholder(spec: ExporterSpec) -> bool:
|
||||
"""Whether ``spec`` is the one ``_normalize`` folds in when nothing was configured.
|
||||
|
||||
No field set is what says the operator asked for nothing: an exporter they did
|
||||
configure survives, even ``OTEL_EXPORTER=console`` whose value matches the default,
|
||||
and so does the gated spec ``credential_gated_exporters`` appends, which would
|
||||
otherwise eat itself when one preset layers onto another.
|
||||
"""
|
||||
return not spec.model_fields_set
|
||||
|
||||
|
||||
class _EnvWithoutBareExcludedServices(PydanticBaseSettingsSource):
|
||||
def __init__(self, settings_cls: type[BaseSettings], env_settings: PydanticBaseSettingsSource) -> None:
|
||||
super().__init__(settings_cls)
|
||||
|
|
|
|||
|
|
@ -272,6 +272,11 @@ class LLMCallEvent:
|
|||
trace: TraceControls
|
||||
session_id: str | None
|
||||
|
||||
@property
|
||||
def reached_a_provider(self) -> bool:
|
||||
"""Whether the call was handed to a provider or answered from the response cache."""
|
||||
return not self.is_no_upstream_call and self.payload is not None and (self.upstream_started or self.cache_hit)
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, kwargs: Mapping[str, object]) -> LLMCallEvent:
|
||||
raw_payload: Final = kwargs.get("standard_logging_object")
|
||||
|
|
|
|||
|
|
@ -47,6 +47,7 @@ from litellm.integrations.otel.mappers.langfuse import (
|
|||
LANGFUSE_OBSERVATION_OUTPUT,
|
||||
LANGFUSE_REQUEST_METADATA_ATTRIBUTES,
|
||||
LANGFUSE_TRACE_NAME,
|
||||
LangfuseMapper,
|
||||
)
|
||||
from litellm.integrations.otel.model.config import (
|
||||
CaptureMessageContent,
|
||||
|
|
@ -55,6 +56,7 @@ from litellm.integrations.otel.model.config import (
|
|||
OpenTelemetryV2Config,
|
||||
capture_message_content_from_env,
|
||||
captures_span_content,
|
||||
is_unconfigured_placeholder,
|
||||
)
|
||||
from litellm.integrations.otel.model.semconv import (
|
||||
DB,
|
||||
|
|
@ -490,19 +492,25 @@ def _is_litellm_span(span: ReadableSpan) -> bool:
|
|||
return scope is not None and scope.name == LITELLM_TRACER_NAME
|
||||
|
||||
|
||||
def without_langfuse_request_metadata(span: ReadableSpan) -> ReadableSpan:
|
||||
"""``span`` without the request metadata the Langfuse mapper writes for a Langfuse destination.
|
||||
def without_destination_langfuse_attributes(span: ReadableSpan, langfuse_mapped: bool) -> ReadableSpan:
|
||||
"""``span`` without the Langfuse attributes only a Langfuse destination asked for.
|
||||
|
||||
Only LiteLLM's own spans lose it: any other span keeps every key its own code set.
|
||||
That is the request metadata the Langfuse mapper writes and, on a model call written by
|
||||
a logger whose own mappers leave the Langfuse vocabulary out (``langfuse_mapped`` off),
|
||||
every Langfuse attribute, since that logger wrote them for a key's or team's Langfuse
|
||||
destination alone. Only LiteLLM's own spans lose them: any other span keeps every key
|
||||
its own code set.
|
||||
"""
|
||||
attributes: Final = span.attributes or _NO_ATTRIBUTES
|
||||
if not (_is_litellm_span(span) and any(key in LANGFUSE_REQUEST_METADATA_ATTRIBUTES for key in attributes)):
|
||||
dropped: Final = (
|
||||
LANGFUSE_REQUEST_METADATA_ATTRIBUTES
|
||||
if langfuse_mapped or not is_llm_call_span(span)
|
||||
else LangfuseMapper.LLM_CALL_ATTRIBUTES
|
||||
)
|
||||
if not (_is_litellm_span(span) and any(key in dropped for key in attributes)):
|
||||
return span
|
||||
return _CountedSpanView(
|
||||
span,
|
||||
MappingProxyType(
|
||||
{key: value for key, value in attributes.items() if key not in LANGFUSE_REQUEST_METADATA_ATTRIBUTES}
|
||||
),
|
||||
span, MappingProxyType({key: value for key, value in attributes.items() if key not in dropped})
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -658,6 +666,7 @@ def _for_destination(
|
|||
span: ReadableSpan,
|
||||
destination: "OtelDestination",
|
||||
default_capture: str,
|
||||
langfuse_mapped: bool,
|
||||
) -> ReadableSpan:
|
||||
"""The view of ``span`` a tenant destination receives.
|
||||
|
||||
|
|
@ -672,13 +681,16 @@ def _for_destination(
|
|||
and the query string of the request URL, which can hold the same key. A
|
||||
destination whose effective setting disables span content loses the prompt and
|
||||
response content too, and a destination for any backend but Langfuse loses the
|
||||
request metadata the Langfuse mapper writes on LiteLLM's own spans. The span itself
|
||||
stays, so the tenant still gets the whole trace tree.
|
||||
Langfuse attributes only a Langfuse destination asked for (see
|
||||
:func:`without_destination_langfuse_attributes`). The span itself stays, so the
|
||||
tenant still gets the whole trace tree.
|
||||
"""
|
||||
extra: Final = destination.resource_attributes
|
||||
defaults: Final = destination.resource_defaults
|
||||
source: Final = (
|
||||
span if destination.callback_name == ExporterOwner.LANGFUSE_OTEL else without_langfuse_request_metadata(span)
|
||||
span
|
||||
if destination.callback_name == ExporterOwner.LANGFUSE_OTEL
|
||||
else without_destination_langfuse_attributes(span, langfuse_mapped)
|
||||
)
|
||||
attributes: Final = source.attributes or _NO_ATTRIBUTES
|
||||
database: Final = _is_database_span(attributes)
|
||||
|
|
@ -728,8 +740,10 @@ class TenantFanOutSpanProcessor(SpanProcessor):
|
|||
default_capture_by_backend: Mapping[str, str] = MappingProxyType({}),
|
||||
pending_drains: int = _MAX_PENDING_DRAINS,
|
||||
drain_pool: _DrainPool | None = None,
|
||||
langfuse_mapped: bool = True,
|
||||
) -> None:
|
||||
self._operator_sinks: Final = operator_sinks
|
||||
self._langfuse_mapped: Final = langfuse_mapped
|
||||
self._excluded_db_systems: Final = excluded_db_systems
|
||||
self._default_capture: Final = default_capture
|
||||
self._default_capture_by_backend: Final = default_capture_by_backend
|
||||
|
|
@ -822,7 +836,8 @@ class TenantFanOutSpanProcessor(SpanProcessor):
|
|||
try:
|
||||
processor.on_end(
|
||||
_scoped(
|
||||
_for_destination(span, destination, self._default_capture_for(destination)), destination.span_scope
|
||||
_for_destination(span, destination, self._default_capture_for(destination), self._langfuse_mapped),
|
||||
destination.span_scope,
|
||||
)
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 # one destination's failure must not cost the others their span
|
||||
|
|
@ -1198,17 +1213,19 @@ class _OverriddenBackendFilter(SpanProcessor):
|
|||
return self._inner.force_flush(timeout_millis)
|
||||
|
||||
|
||||
class _LangfuseRequestMetadataFilter(SpanProcessor):
|
||||
"""Strip the Langfuse request metadata from LiteLLM's own spans before ``inner`` exports them."""
|
||||
class _DestinationLangfuseAttributeFilter(SpanProcessor):
|
||||
"""Strip the Langfuse attributes only a Langfuse destination asked for from LiteLLM's own spans before ``inner``
|
||||
exports them."""
|
||||
|
||||
def __init__(self, inner: SpanProcessor) -> None:
|
||||
def __init__(self, inner: SpanProcessor, langfuse_mapped: bool) -> None:
|
||||
self._inner: Final = inner
|
||||
self._langfuse_mapped: Final = langfuse_mapped
|
||||
|
||||
def on_start(self, span: SDKSpan, parent_context: Context | None = None) -> None:
|
||||
self._inner.on_start(span, parent_context)
|
||||
|
||||
def on_end(self, span: ReadableSpan) -> None:
|
||||
self._inner.on_end(without_langfuse_request_metadata(span))
|
||||
self._inner.on_end(without_destination_langfuse_attributes(span, self._langfuse_mapped))
|
||||
|
||||
def shutdown(self) -> None:
|
||||
self._inner.shutdown()
|
||||
|
|
@ -1491,13 +1508,18 @@ def _spec_processor(
|
|||
spec: ExporterSpec,
|
||||
use_simple_processor: bool | None,
|
||||
strip_message_content: bool,
|
||||
strip_langfuse_request_metadata: bool,
|
||||
strip_destination_langfuse_attributes: bool,
|
||||
langfuse_mapped: bool,
|
||||
) -> SpanProcessor:
|
||||
exporting: Final = _processor_for(
|
||||
_exporter_from_spec(spec),
|
||||
(spec.use_simple_processor if spec.use_simple_processor is not None else use_simple_processor),
|
||||
)
|
||||
scoped: Final = _LangfuseRequestMetadataFilter(exporting) if strip_langfuse_request_metadata else exporting
|
||||
scoped: Final = (
|
||||
_DestinationLangfuseAttributeFilter(exporting, langfuse_mapped)
|
||||
if strip_destination_langfuse_attributes
|
||||
else exporting
|
||||
)
|
||||
return _MessageContentFilter(scoped) if strip_message_content else scoped
|
||||
|
||||
|
||||
|
|
@ -1530,6 +1552,7 @@ def build_tracer_provider(
|
|||
use_simple_processor: bool | None = None,
|
||||
tenant_overrides: bool = False,
|
||||
strip_message_content: bool = False,
|
||||
console_default: bool = False,
|
||||
) -> TracerProvider:
|
||||
"""Build the shared :class:`TracerProvider`.
|
||||
|
||||
|
|
@ -1548,11 +1571,16 @@ def build_tracer_provider(
|
|||
|
||||
``config.langfuse_span_scope`` narrows the exporter owned by ``langfuse_otel``
|
||||
alone; a collector or any other backend in the same config keeps the full tree.
|
||||
In a config that writes the Langfuse vocabulary, every exporter but ``langfuse_otel``'s
|
||||
goes without the request metadata the Langfuse mapper writes on LiteLLM's own spans.
|
||||
Every exporter but ``langfuse_otel``'s goes without the Langfuse attributes only a
|
||||
Langfuse destination asked for (see :func:`without_destination_langfuse_attributes`),
|
||||
which any logger writes for a request with a Langfuse destination.
|
||||
``strip_message_content`` removes prompt and response content before every
|
||||
exporter, for a routed provider whose spans can carry content another
|
||||
destination asked for.
|
||||
|
||||
The console spec ``config._normalize`` folds in when nothing was configured
|
||||
prints only under ``console_default``, which the ``otel`` callback passes: a
|
||||
vendor preset or the proxy's background global logger was never asked to print.
|
||||
"""
|
||||
provider: Final = TracerProvider(resource=build_resource(config))
|
||||
provider.sampler = _RequestExclusionSampler(provider.sampler)
|
||||
|
|
@ -1564,19 +1592,20 @@ def build_tracer_provider(
|
|||
provider.add_span_processor(_processor_for(exporter, use_simple_processor))
|
||||
return provider
|
||||
|
||||
langfuse_mapped: Final = _maps_langfuse(config)
|
||||
# ``config._normalize`` guarantees at least one spec (it folds the top-level
|
||||
# ``exporter``/``endpoint``/``headers`` fields in when ``exporters`` is empty).
|
||||
writes_langfuse: Final = "langfuse" in config.mapper_names or any(
|
||||
spec.owner is ExporterOwner.LANGFUSE_OTEL for spec in config.exporters
|
||||
)
|
||||
for spec in config.exporters:
|
||||
if spec.requires_headers and not spec.headers:
|
||||
continue
|
||||
if is_unconfigured_placeholder(spec) and not console_default:
|
||||
continue
|
||||
processor = _spec_processor(
|
||||
spec,
|
||||
use_simple_processor,
|
||||
strip_message_content,
|
||||
strip_langfuse_request_metadata=writes_langfuse and spec.owner is not ExporterOwner.LANGFUSE_OTEL,
|
||||
strip_destination_langfuse_attributes=spec.owner is not ExporterOwner.LANGFUSE_OTEL,
|
||||
langfuse_mapped=langfuse_mapped,
|
||||
)
|
||||
owner = spec.owner.value if tenant_overrides and spec.owner is not None else None
|
||||
scope = _operator_scope(config, spec)
|
||||
|
|
@ -1613,8 +1642,9 @@ def attach_tenant_fan_out(
|
|||
so exactly one fan-out lands. ``configs`` name the operator's own exporters, one
|
||||
config per v2 logger since each keeps its own provider and still writes its
|
||||
account, so an additive destination pointing at any of them is delivered once
|
||||
rather than twice. ``excluded_db_systems`` only filters what the fan-out
|
||||
delivers, never the operator's own exporters.
|
||||
rather than twice. The first is the config of the logger that owns ``provider``,
|
||||
whose model calls the fan-out carries. ``excluded_db_systems`` only filters what
|
||||
the fan-out delivers, never the operator's own exporters.
|
||||
"""
|
||||
with _FAN_OUT_ATTACH_LOCK:
|
||||
if any(isinstance(processor, TenantFanOutSpanProcessor) for processor in _attached_processors(provider)):
|
||||
|
|
@ -1625,10 +1655,16 @@ def attach_tenant_fan_out(
|
|||
excluded_db_systems=excluded_db_systems,
|
||||
default_capture=_operator_capture_mode(configs),
|
||||
default_capture_by_backend=MappingProxyType(dict(_owned_capture_modes(configs))),
|
||||
langfuse_mapped=any(_maps_langfuse(config) for config in configs[:1]),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _maps_langfuse(config: OpenTelemetryV2Config) -> bool:
|
||||
"""Whether a logger built from ``config`` writes the Langfuse vocabulary on every model call."""
|
||||
return "langfuse" in config.mapper_names
|
||||
|
||||
|
||||
def _operator_capture_mode(configs: Sequence[OpenTelemetryV2Config]) -> str:
|
||||
"""The mode of the operator's generic ``otel`` callback, so ``callback_settings.otel`` counts as the global default"""
|
||||
return next(
|
||||
|
|
|
|||
|
|
@ -163,10 +163,12 @@ class TenantTracerCache:
|
|||
config: OpenTelemetryV2Config,
|
||||
callback_name: str | None,
|
||||
tracer_name: str,
|
||||
console_default: bool = False,
|
||||
) -> None:
|
||||
self._config = config
|
||||
self._callback_name = callback_name
|
||||
self._tracer_name = tracer_name
|
||||
self._console_default = console_default
|
||||
# Guards the three mutable structures below: ``pre_call`` can run on
|
||||
# thread-pool workers concurrently with the event loop, so cache
|
||||
# updates, span counts, and retirement must be atomic.
|
||||
|
|
@ -284,6 +286,7 @@ class TenantTracerCache:
|
|||
built: Final = build_tracer_provider(
|
||||
self._routed_config(credential_headers, project_headers, endpoint, service_name),
|
||||
strip_message_content=not self._config.capture_span_content,
|
||||
console_default=self._console_default,
|
||||
)
|
||||
self._providers[cache_key] = built
|
||||
return built
|
||||
|
|
|
|||
|
|
@ -111,7 +111,7 @@ _REQUIRED_HEADERS_BY_CALLBACK: Final[Mapping[str, frozenset[str]]] = MappingProx
|
|||
_NO_ATTRS: Final[Mapping[str, str]] = MappingProxyType({})
|
||||
|
||||
|
||||
def _arize_project() -> Mapping[str, str]:
|
||||
def _arize_project(_params: StandardCallbackDynamicParams) -> Mapping[str, str]:
|
||||
from litellm.integrations.arize.arize import ArizeLogger
|
||||
|
||||
project: Final = ArizeLogger.get_arize_config().project_name
|
||||
|
|
@ -124,10 +124,27 @@ def _arize_default_project() -> Mapping[str, str]:
|
|||
return MappingProxyType({"model_id": ARIZE_DEFAULT_PROJECT})
|
||||
|
||||
|
||||
#: Resource attributes the operator chose for a backend, applied to every export the
|
||||
#: team receives over the span's own resource: ``ARIZE_PROJECT_NAME`` for Arize.
|
||||
_RESOURCE_ATTRIBUTES_BY_CALLBACK: Final[Mapping[str, Callable[[], Mapping[str, str]]]] = MappingProxyType(
|
||||
{"arize": _arize_project}
|
||||
def _stands_in_for_the_legacy_langfuse_logger() -> bool:
|
||||
"""Whether a Langfuse destination replaces the older Langfuse logger's copy, which a proxy without Langfuse keys
|
||||
of its own sends for a key's or team's ``langfuse_otel``."""
|
||||
return not (os.environ.get("LANGFUSE_PUBLIC_KEY") and os.environ.get("LANGFUSE_SECRET_KEY"))
|
||||
|
||||
|
||||
def _langfuse_environment(params: StandardCallbackDynamicParams) -> Mapping[str, str]:
|
||||
"""The key's or team's ``langfuse_environment``, else ``LANGFUSE_TRACING_ENVIRONMENT``, where the older Langfuse
|
||||
logger wrote it."""
|
||||
from litellm.types.integrations.langfuse_otel import LangfuseSpanAttributes
|
||||
|
||||
if not _stands_in_for_the_legacy_langfuse_logger():
|
||||
return _NO_ATTRS
|
||||
environment: Final = params.get("langfuse_environment") or os.environ.get("LANGFUSE_TRACING_ENVIRONMENT")
|
||||
return MappingProxyType({LangfuseSpanAttributes.LANGFUSE_ENVIRONMENT.value: environment} if environment else {})
|
||||
|
||||
|
||||
#: Resource attributes a backend gets on every export the team receives, over the
|
||||
#: span's own resource: ``ARIZE_PROJECT_NAME`` for Arize, the Langfuse environment.
|
||||
_RESOURCE_ATTRIBUTES_BY_CALLBACK: Final[Mapping[str, Callable[[StandardCallbackDynamicParams], Mapping[str, str]]]] = (
|
||||
MappingProxyType({"arize": _arize_project, "langfuse_otel": _langfuse_environment})
|
||||
)
|
||||
|
||||
#: Resource attributes a backend needs on every export to accept it, filled only where
|
||||
|
|
@ -137,12 +154,14 @@ _RESOURCE_DEFAULTS_BY_CALLBACK: Final[Mapping[str, Callable[[], Mapping[str, str
|
|||
)
|
||||
|
||||
|
||||
def _resource_attributes(callback_name: str, service_name: str | None) -> Mapping[str, str]:
|
||||
def _resource_attributes(
|
||||
callback_name: str, service_name: str | None, params: StandardCallbackDynamicParams
|
||||
) -> Mapping[str, str]:
|
||||
backend: Final = _RESOURCE_ATTRIBUTES_BY_CALLBACK.get(callback_name)
|
||||
return MappingProxyType(
|
||||
{
|
||||
**({"service.name": service_name} if service_name else {}),
|
||||
**(backend() if backend is not None else {}),
|
||||
**(backend(params) if backend is not None else {}),
|
||||
}
|
||||
)
|
||||
|
||||
|
|
@ -214,7 +233,7 @@ def destination_for(
|
|||
return OtelDestination(
|
||||
endpoint=endpoint,
|
||||
headers=MappingProxyType(dict(headers)),
|
||||
resource_attributes=_resource_attributes(callback_name, service_name),
|
||||
resource_attributes=_resource_attributes(callback_name, service_name, params),
|
||||
resource_defaults=_resource_defaults(callback_name),
|
||||
callback_name=callback_name,
|
||||
protocol=protocol,
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@
|
|||
from collections.abc import Iterable
|
||||
from typing import Final
|
||||
|
||||
from litellm.integrations.otel.model.config import ExporterOwner, ExporterSpec
|
||||
from litellm.integrations.otel.model.config import ExporterOwner, ExporterSpec, is_unconfigured_placeholder
|
||||
|
||||
|
||||
def ensure_mappers(mapper_names: Iterable[str], *names: str) -> list[str]:
|
||||
|
|
@ -35,14 +35,3 @@ def credential_gated_exporters(
|
|||
*(spec for spec in exporters if not is_unconfigured_placeholder(spec)),
|
||||
ExporterSpec(owner=owner, requires_headers=True),
|
||||
)
|
||||
|
||||
|
||||
def is_unconfigured_placeholder(spec: "ExporterSpec") -> bool:
|
||||
"""Whether ``spec`` is the one ``_normalize`` folds in when nothing was configured.
|
||||
|
||||
No field set is what says the operator asked for nothing: an exporter they did
|
||||
configure survives, even ``OTEL_EXPORTER=console`` whose value matches the default,
|
||||
and so does the gated spec this module appends, which would otherwise eat itself
|
||||
when one preset layers onto another.
|
||||
"""
|
||||
return not spec.model_fields_set
|
||||
|
|
|
|||
|
|
@ -5640,7 +5640,7 @@ def build_custom_logger_compatible_class(
|
|||
return callback
|
||||
otel_settings: Final = _get_custom_logger_settings_from_proxy_server(callback_name=logging_integration)
|
||||
otel_logger_v2: Final = build_otel_v2_logger(
|
||||
config=OpenTelemetryV2Config(**otel_settings), settings=otel_settings
|
||||
config=OpenTelemetryV2Config(**otel_settings), settings=otel_settings, console_default=True
|
||||
)
|
||||
_in_memory_loggers.append(otel_logger_v2)
|
||||
_maybe_auto_initialize_arize_phoenix(_in_memory_loggers)
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ import httpx
|
|||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
|
||||
from litellm.litellm_core_utils.url_utils import join_api_base
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.base_llm.responses.codex_compat import drop_unsupported_tools, normalize_codex_input_items
|
||||
|
|
@ -196,7 +197,9 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
|
|||
params: Final = cast( # cast-ok: the base signature leaves the params dict untyped
|
||||
"ResponsesAPIOptionalRequestParams", response_api_optional_request_params
|
||||
)
|
||||
hoisted: Final = hoist_additional_tools(input, params.get("tools"))
|
||||
hoisted: Final = hoist_additional_tools(
|
||||
input, params.get("tools"), drop_params=drops_unsupported_params(litellm_params.drop_params)
|
||||
)
|
||||
normalized_input, rewritten_types = normalize_codex_input_items(hoisted.input)
|
||||
if rewritten_types:
|
||||
verbose_logger.warning(
|
||||
|
|
|
|||
|
|
@ -1205,6 +1205,17 @@ class AsyncHTTPHandler:
|
|||
# Mirrors LiteLLMAiohttpTransport._background_close_tasks.
|
||||
_finalizer_close_tasks: ClassVar[set["asyncio.Task[None]"]] = set() # mutable-ok: strong refs for pending closes
|
||||
|
||||
@classmethod
|
||||
async def _close_client_at_finalization(cls, client: httpx.AsyncClient) -> None:
|
||||
try:
|
||||
await client.aclose()
|
||||
except (OSError, RuntimeError) as exc:
|
||||
verbose_logger.debug("Error closing client at finalization: %s", exc)
|
||||
finally:
|
||||
task: Final = asyncio.current_task()
|
||||
if task is not None:
|
||||
cls._finalizer_close_tasks.discard(task)
|
||||
|
||||
@classmethod
|
||||
def _on_finalizer_close_done(cls, task: "asyncio.Task[None]") -> None:
|
||||
cls._finalizer_close_tasks.discard(task)
|
||||
|
|
@ -1272,8 +1283,8 @@ class AsyncHTTPHandler:
|
|||
# here is the cross-loop path the transport refuses.
|
||||
self._dispose_wrapped_aiohttp_session()
|
||||
return
|
||||
task: Final = loop.create_task(self._client.aclose())
|
||||
cls: Final = type(self)
|
||||
task: Final = loop.create_task(cls._close_client_at_finalization(self._client))
|
||||
cls._finalizer_close_tasks.add(task)
|
||||
task.add_done_callback(cls._on_finalizer_close_done)
|
||||
except Exception:
|
||||
|
|
@ -1973,7 +1984,7 @@ def get_httpx_client(params: dict | None = None) -> HTTPHandler:
|
|||
cache = LLMClientCache()
|
||||
setattr(litellm, "in_memory_llm_clients_cache", cache)
|
||||
|
||||
_cached_client: Final = cache.get_cache(_cache_key_name)
|
||||
_cached_client: Final = cache.get_cache(_cache_key_name, loop_bound=False)
|
||||
if _cached_client:
|
||||
return _cached_client
|
||||
|
||||
|
|
@ -1989,6 +2000,7 @@ def get_httpx_client(params: dict | None = None) -> HTTPHandler:
|
|||
value=_new_client,
|
||||
ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS,
|
||||
litellm_owned_client=True,
|
||||
loop_bound=False,
|
||||
)
|
||||
return _new_client
|
||||
|
||||
|
|
|
|||
|
|
@ -275,6 +275,10 @@ def build_output_token_limit_response(
|
|||
return e.response.headers, (_aiter_once(chunk) if is_async else _iter_once(chunk))
|
||||
|
||||
|
||||
def _is_loop_bound(client_initialization_params: Mapping[str, object]) -> bool:
|
||||
return client_initialization_params.get("is_async") is not False
|
||||
|
||||
|
||||
class BaseOpenAILLM:
|
||||
"""
|
||||
Base class for OpenAI LLMs for getting their httpx clients and SSL verification settings
|
||||
|
|
@ -289,7 +293,9 @@ class BaseOpenAILLM:
|
|||
client_initialization_params=client_initialization_params,
|
||||
client_type=client_type,
|
||||
)
|
||||
_cached_client: Final = litellm.in_memory_llm_clients_cache.get_cache(_cache_key)
|
||||
_cached_client: Final = litellm.in_memory_llm_clients_cache.get_cache(
|
||||
_cache_key, loop_bound=_is_loop_bound(client_initialization_params)
|
||||
)
|
||||
return _cached_client
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -330,6 +336,7 @@ class BaseOpenAILLM:
|
|||
value=openai_client,
|
||||
ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS,
|
||||
litellm_owned_client=litellm_owned_client,
|
||||
loop_bound=_is_loop_bound(client_initialization_params),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -287,6 +287,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
streaming_sampling_rate: int | None = None,
|
||||
streaming_end_of_stream_only: bool | None = None,
|
||||
streaming_buffer_release_on_scan: bool | None = None,
|
||||
streaming_scan_in_background: bool | None = None,
|
||||
async_handler: AsyncHTTPHandler | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
|
|
@ -301,6 +302,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
"streaming_sampling_rate": streaming_sampling_rate,
|
||||
"streaming_end_of_stream_only": streaming_end_of_stream_only,
|
||||
"streaming_buffer_release_on_scan": streaming_buffer_release_on_scan,
|
||||
"streaming_scan_in_background": streaming_scan_in_background,
|
||||
}
|
||||
)
|
||||
)
|
||||
|
|
@ -366,6 +368,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
self.streaming_sampling_rate = streaming_params.streaming_sampling_rate
|
||||
self.streaming_end_of_stream_only = streaming_params.streaming_end_of_stream_only
|
||||
self.streaming_buffer_release_on_scan = streaming_params.streaming_buffer_release_on_scan
|
||||
self.streaming_scan_in_background = streaming_params.streaming_scan_in_background
|
||||
|
||||
def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None:
|
||||
super().update_in_memory_litellm_params(litellm_params)
|
||||
|
|
|
|||
|
|
@ -137,6 +137,13 @@ def _chunk_choices(item: object) -> Sequence[object]:
|
|||
return choices
|
||||
|
||||
|
||||
def _ends_stream(endpoint_translation: object, item: object) -> bool:
|
||||
has_ended: Final[Callable[[Sequence[object]], bool] | None] = getattr(
|
||||
endpoint_translation, "_check_streaming_has_ended", None
|
||||
)
|
||||
return has_ended is not None and has_ended((item,))
|
||||
|
||||
|
||||
def _held_choices(held_chars_per_choice: Mapping[int, int]) -> frozenset[int]:
|
||||
return frozenset(idx for idx, held in held_chars_per_choice.items() if held > 0)
|
||||
|
||||
|
|
@ -1179,6 +1186,13 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
if buffer_until_moderated and not release_on_scan:
|
||||
end_of_stream_only = True
|
||||
|
||||
# Run each sampled scan as a background task so the stream keeps flowing; a failed scan ends
|
||||
# the stream before the next chunk, and the terminal frames wait for the final verdict.
|
||||
# Buffered and end-only streams already scan before releasing.
|
||||
scan_in_background: Final = bool(_streaming_setting("streaming_scan_in_background", False)) and not (
|
||||
buffer_until_moderated or end_of_stream_only
|
||||
)
|
||||
|
||||
if guardrail_to_apply is None:
|
||||
async for item in response:
|
||||
yield item
|
||||
|
|
@ -1242,6 +1256,10 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
last_scan_key: StreamingScanKey | None = None # rebind-ok: replaced after every scan round
|
||||
tool_calls_in_flight = False # rebind-ok: tracks the latest scan key's unscanned tool calls
|
||||
verdict_settled = False # rebind-ok: set once the end-of-stream scan or a block owns the verdict
|
||||
pending_scan: asyncio.Task[object] | None = None # rebind-ok: the one background scan in flight
|
||||
# Background scans read copies, because the translation rewrites the chunks it scans in place.
|
||||
# Each chunk is copied once, when a scan first covers it.
|
||||
scanned_responses: Final[list[object]] = [] # mutable-ok: grows with the scanned stream
|
||||
|
||||
try:
|
||||
async for item in response:
|
||||
|
|
@ -1291,10 +1309,21 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
withheld_items.append(item)
|
||||
continue
|
||||
|
||||
if scan_in_background and (
|
||||
pending_end_of_stream_items or _ends_stream(mappings[CallTypes(call_type)](), item)
|
||||
):
|
||||
pending_end_of_stream_items.append(item)
|
||||
continue
|
||||
|
||||
if pending_scan is not None and pending_scan.done():
|
||||
if pending_scan.exception() is not None:
|
||||
break # the end-of-stream handling below ends the stream with this scan's verdict
|
||||
pending_scan = None
|
||||
|
||||
# Process chunk based on sampling rate
|
||||
if buffer_until_moderated:
|
||||
withheld_items.append(item)
|
||||
if chunk_counter % sampling_rate == 0:
|
||||
if chunk_counter % sampling_rate == 0 and pending_scan is None:
|
||||
endpoint_translation = mappings[CallTypes(call_type)]()
|
||||
scan_key = endpoint_translation.get_streaming_scan_key(responses_so_far)
|
||||
if scan_key is not None:
|
||||
|
|
@ -1327,18 +1356,30 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
guardrail_to_apply.guardrail_name,
|
||||
)
|
||||
|
||||
if scan_in_background:
|
||||
scanned_responses.extend(copy.deepcopy(responses_so_far[len(scanned_responses) :]))
|
||||
scan = endpoint_translation.process_output_streaming_response(
|
||||
responses_so_far=list(scanned_responses) if scan_in_background else responses_so_far,
|
||||
guardrail_to_apply=guardrail_to_apply,
|
||||
litellm_logging_obj=request_data.get("litellm_logging_obj"),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_data=request_data,
|
||||
)
|
||||
if scan_in_background:
|
||||
pending_scan = asyncio.create_task(scan)
|
||||
if scan_key is not None:
|
||||
last_scan_key = scan_key
|
||||
chunks_yielded = True
|
||||
responses_yielded.append(item)
|
||||
yield item
|
||||
continue
|
||||
|
||||
original_items = (
|
||||
tuple(copy.deepcopy(withheld_items)) if buffer_until_moderated else (copy.deepcopy(item),)
|
||||
)
|
||||
|
||||
try:
|
||||
await endpoint_translation.process_output_streaming_response(
|
||||
responses_so_far=responses_so_far,
|
||||
guardrail_to_apply=guardrail_to_apply,
|
||||
litellm_logging_obj=request_data.get("litellm_logging_obj"),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_data=request_data,
|
||||
)
|
||||
await scan
|
||||
except ModifyResponseException as e:
|
||||
verdict_settled = True
|
||||
if e.original_response is None:
|
||||
|
|
@ -1405,6 +1446,10 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
|
||||
endpoint_translation = mappings[CallTypes(call_type)]()
|
||||
|
||||
if pending_scan is not None:
|
||||
await asyncio.wait((pending_scan,))
|
||||
background_scan_failed: Final = pending_scan is not None and pending_scan.exception() is not None
|
||||
|
||||
buffered_items: Final = (
|
||||
tuple(copy.deepcopy(withheld_items))
|
||||
if buffer_until_moderated and release_on_scan and not end_of_stream_only
|
||||
|
|
@ -1414,8 +1459,10 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
)
|
||||
end_scan_key: Final = endpoint_translation.get_streaming_scan_key(responses_so_far)
|
||||
verdict_settled = True
|
||||
if _is_redundant_scan(end_scan_key, last_scan_key) and not _owes_ended_scan(
|
||||
endpoint_translation, guardrail_to_apply, responses_so_far, last_scan_key
|
||||
if (
|
||||
_is_redundant_scan(end_scan_key, last_scan_key)
|
||||
and not _owes_ended_scan(endpoint_translation, guardrail_to_apply, responses_so_far, last_scan_key)
|
||||
and not background_scan_failed
|
||||
):
|
||||
verbose_proxy_logger.debug(
|
||||
"Skipping end-of-stream scan for guardrail %s: the last sampled round already scanned it all",
|
||||
|
|
@ -1430,6 +1477,8 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
|
||||
try:
|
||||
with anyio.CancelScope(shield=chunks_yielded):
|
||||
if pending_scan is not None:
|
||||
pending_scan.result() # a failed background scan raises into the handlers below
|
||||
await endpoint_translation.process_output_streaming_response(
|
||||
responses_so_far=responses_so_far,
|
||||
guardrail_to_apply=guardrail_to_apply,
|
||||
|
|
@ -1471,6 +1520,11 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
):
|
||||
yield error_item
|
||||
except (GeneratorExit, asyncio.CancelledError):
|
||||
if pending_scan is not None:
|
||||
with anyio.CancelScope(shield=True):
|
||||
await asyncio.wait((pending_scan,))
|
||||
if isinstance(pending_scan.exception(), (ModifyResponseException, HTTPException)):
|
||||
verdict_settled = True # rebind-ok: the background scan already reached a verdict
|
||||
translation_class: Final = None if call_type is None else mappings.get(CallTypes(call_type))
|
||||
if (
|
||||
chunks_yielded
|
||||
|
|
@ -1487,3 +1541,8 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
request_data=typed_request_data,
|
||||
)
|
||||
raise
|
||||
except Exception:
|
||||
if pending_scan is not None:
|
||||
pending_scan.cancel()
|
||||
await asyncio.wait((pending_scan,))
|
||||
raise
|
||||
|
|
|
|||
|
|
@ -46,6 +46,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail):
|
|||
streaming_sampling_rate=streaming_params.streaming_sampling_rate,
|
||||
streaming_end_of_stream_only=streaming_params.streaming_end_of_stream_only,
|
||||
streaming_buffer_release_on_scan=streaming_params.streaming_buffer_release_on_scan,
|
||||
streaming_scan_in_background=streaming_params.streaming_scan_in_background,
|
||||
timeout=litellm_params.timeout,
|
||||
)
|
||||
litellm.logging_callback_manager.add_litellm_callback(_bedrock_callback)
|
||||
|
|
|
|||
|
|
@ -89,6 +89,8 @@ def _unique_tools(tools: Sequence[Mapping[str, object]]) -> tuple[Mapping[str, o
|
|||
def hoist_additional_tools(
|
||||
input: str | ResponseInputParam,
|
||||
tools: Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None,
|
||||
*,
|
||||
drop_params: bool,
|
||||
) -> HoistedAdditionalTools:
|
||||
existing: Final = tuple(tools or ())
|
||||
if isinstance(input, str):
|
||||
|
|
@ -103,6 +105,8 @@ def hoist_additional_tools(
|
|||
len(items),
|
||||
)
|
||||
remaining_input: Final = [item for item in input if not _is_additional_tools_item(item)]
|
||||
if not drop_params:
|
||||
return HoistedAdditionalTools(input=remaining_input, tools=(*existing, *hoisted), hoisted=hoisted)
|
||||
unique: Final = tuple(
|
||||
cast("ALL_RESPONSES_API_TOOL_PARAMS", tool) # cast-ok: a merged namespace is rebuilt from the same tool dicts
|
||||
for tool in _unique_tools((*existing, *hoisted))
|
||||
|
|
|
|||
|
|
@ -114,7 +114,7 @@ def with_provider_file_ids(
|
|||
}
|
||||
|
||||
|
||||
def _with_hoisted_additional_tools(responses_request: Mapping[str, object]) -> Mapping[str, object]:
|
||||
def _with_hoisted_additional_tools(responses_request: Mapping[str, object], drop_params: bool) -> Mapping[str, object]:
|
||||
responses_input: Final = responses_request.get("input")
|
||||
if not isinstance(responses_input, list):
|
||||
return responses_request
|
||||
|
|
@ -122,7 +122,7 @@ def _with_hoisted_additional_tools(responses_request: Mapping[str, object]) -> M
|
|||
top_level_tools: Final = cast( # cast-ok: tools are forwarded as sent, as inference does
|
||||
"tuple[ALL_RESPONSES_API_TOOL_PARAMS, ...]", _ITEMS.validate_python(responses_request.get("tools") or ())
|
||||
)
|
||||
hoisted: Final = hoist_additional_tools(input_items, top_level_tools)
|
||||
hoisted: Final = hoist_additional_tools(input_items, top_level_tools, drop_params=drop_params)
|
||||
if not hoisted.hoisted:
|
||||
return responses_request
|
||||
return {**responses_request, "input": hoisted.input, "tools": list(hoisted.tools)}
|
||||
|
|
@ -196,7 +196,7 @@ async def chat_count_request(
|
|||
chat_body: Final = await _with_previous_turns(
|
||||
responses_batch_body_to_chat_body(
|
||||
_with_compaction_as_summary_text(
|
||||
_with_emulated_file_search(_with_hoisted_additional_tools(responses_request))
|
||||
_with_emulated_file_search(_with_hoisted_additional_tools(responses_request, policy.drop_params))
|
||||
),
|
||||
custom_llm_provider,
|
||||
policy.drop_params,
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ from typing import Final
|
|||
from pydantic import TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
|
||||
from litellm.responses.additional_tools import HoistedAdditionalTools, hoist_additional_tools
|
||||
from litellm.responses.litellm_completion_transformation.client_tool_search import bridge_client_tool_search
|
||||
from litellm.responses.litellm_completion_transformation.compaction_items import holds_compaction_summary
|
||||
|
|
@ -73,7 +74,11 @@ class LiteLLMCompletionTransformationHandler:
|
|||
| BaseResponsesAPIStreamingIterator
|
||||
| Coroutine[object, object, ResponsesAPIResponse | BaseResponsesAPIStreamingIterator]
|
||||
):
|
||||
hoisted: Final = hoist_additional_tools(input, responses_api_request.get("tools"))
|
||||
hoisted: Final = hoist_additional_tools(
|
||||
input,
|
||||
responses_api_request.get("tools"),
|
||||
drop_params=drops_unsupported_params(_REQUEST_KWARGS_ADAPTER.validate_python(kwargs).get("drop_params")),
|
||||
)
|
||||
previous_response_id: Final = responses_api_request.get("previous_response_id")
|
||||
if _is_async and previous_response_id and has_client_tool_search(hoisted.tools):
|
||||
return self._async_response_api_handler_after_earlier_searches(
|
||||
|
|
|
|||
|
|
@ -709,6 +709,12 @@ class BedrockGuardrailStreamingParams(LiteLLMBaseModel):
|
|||
"and release the withheld chunks once the scan passes, instead of holding everything to end of stream. "
|
||||
"Flagged content is never released. Ignored when streaming_end_of_stream_only is true.",
|
||||
)
|
||||
streaming_scan_in_background: bool = Field(
|
||||
default=False,
|
||||
description="When not buffering and not end-of-stream-only, run each sampled ApplyGuardrail scan in the "
|
||||
"background so the sampled chunk and the ones after it stream while the scan runs. One scan runs at a time. "
|
||||
"A flagged scan ends the stream before the next chunk, so content streamed while it ran reaches the client.",
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_extras(cls, extras: Mapping[str, object] | None) -> "BedrockGuardrailStreamingParams":
|
||||
|
|
|
|||
|
|
@ -27,8 +27,11 @@ tests. A product bug the test exposes is `pytest.skip("BUG: <symptom>")` at the
|
|||
fix in the test and not a deletion. Needs no network, proxy, DB or Redis: `tests/unit`
|
||||
|
||||
CircleCI's integration groups run every file here except the OCI and sandbox tests, which no CI runs, and
|
||||
the files in `GITHUB_FILES` in `run.py`. CircleCI runs on main, not on pull requests, so a new file here
|
||||
first runs after its PR merges. Wanting a file to run before merge is no reason to change CI, and a new
|
||||
the files in `GITHUB_FILES` in `run.py`. `test-integration.yml` runs the same groups through the same
|
||||
`run_integration.sh` on GitHub Actions, on main every three hours from 8am to 8pm Pacific Monday to Saturday,
|
||||
and on a pull request when someone adds the `run-ci` label. Neither runs every pull request, so a new file here
|
||||
first runs on main after its PR merges, or earlier if the PR gets `run-ci`. Wanting a file to run before merge
|
||||
is no reason to change CI, and a new
|
||||
dedicated workflow for it is the same mistake as a shared shard. Write the file, leave `GITHUB_FILES` and
|
||||
every workflow alone, and tell the user it first runs on main. `GITHUB_FILES` is only for a file that
|
||||
needs what CircleCI's groups cannot give it, like the images `image-scan.yml` builds or the plain pytest
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
# Integration contracts
|
||||
|
||||
These tests exercise a running gateway, PostgreSQL and Redis with an owned local upstream. CircleCI owns this suite. Tests are grouped by behavior, with no automatic test retries or fallback to paid provider calls
|
||||
These tests exercise a running gateway, PostgreSQL and Redis with an owned local upstream. CircleCI owns this suite, and `.github/workflows/test-integration.yml` runs the same groups through `run_integration.sh` on GitHub Actions for `run-ci` pull requests and scheduled main runs. Tests are grouped by behavior, with no automatic test retries or fallback to paid provider calls
|
||||
|
||||
The ROI database contracts in `database/test_roi_observed.py` run as plain pytest outside `run_integration.sh` in CircleCI's `roi-database` Postgres job. The job uploads coverage with the `roi-postgres` flag. These contracts own temporary databases and script only the external provider transport
|
||||
|
||||
|
|
|
|||
|
|
@ -210,6 +210,8 @@ def langfuse_rig(
|
|||
sink_reply: Callable[[Request], Reply] | None = None,
|
||||
overrides: Mapping[str, str] | None = None,
|
||||
model_list: Callable[[str], Sequence[Mapping[str, object]]] | None = None,
|
||||
remove_environment: tuple[str, ...] = (),
|
||||
operator_keys: bool = True,
|
||||
) -> Iterator[LangfuseRig]:
|
||||
def upstream(request: Request) -> Reply:
|
||||
if request.method == "GET":
|
||||
|
|
@ -253,8 +255,11 @@ def langfuse_rig(
|
|||
gateway,
|
||||
tmp_path,
|
||||
{
|
||||
"LANGFUSE_PUBLIC_KEY": "pk-lf-integration",
|
||||
"LANGFUSE_SECRET_KEY": "sk-lf-integration",
|
||||
**(
|
||||
{"LANGFUSE_PUBLIC_KEY": "pk-lf-integration", "LANGFUSE_SECRET_KEY": "sk-lf-integration"}
|
||||
if operator_keys
|
||||
else {}
|
||||
),
|
||||
"LANGFUSE_HOST": collector.url,
|
||||
"LITELLM_OTEL_V2": "1",
|
||||
**{
|
||||
|
|
@ -264,6 +269,11 @@ def langfuse_rig(
|
|||
},
|
||||
config=path,
|
||||
workers=2,
|
||||
remove_environment=(
|
||||
remove_environment
|
||||
if operator_keys
|
||||
else (*remove_environment, "LANGFUSE_PUBLIC_KEY", "LANGFUSE_SECRET_KEY")
|
||||
),
|
||||
) as owned:
|
||||
yield LangfuseRig(owned.gateway, provider, collector, owned)
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,98 @@
|
|||
import json
|
||||
import os
|
||||
import uuid
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from _langfuse_otel_metadata_support import exported_llm_span, langfuse_rig, scenario_identity, watched_spans
|
||||
from integration._support.client import Gateway, eventually
|
||||
|
||||
CONSOLE_SPAN: Final = '"span_id": "0x'
|
||||
INHERITED_OTEL_ENV: Final = tuple(name for name in os.environ if name.startswith("OTEL_"))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("overrides", "printed"),
|
||||
[({}, False), ({"OTEL_EXPORTER": "console"}, True)],
|
||||
ids=["exporter-unset", "operator-asked-for-console"],
|
||||
)
|
||||
def test_langfuse_otel_prints_spans_to_stdout_only_when_the_operator_asks(
|
||||
gateway: Gateway, tmp_path: Path, overrides: Mapping[str, str], printed: bool
|
||||
) -> None:
|
||||
marker: Final = "lf-console-" + uuid.uuid4().hex
|
||||
with (
|
||||
langfuse_rig(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{},
|
||||
marker,
|
||||
overrides={"OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT": "span_only", **overrides},
|
||||
remove_environment=INHERITED_OTEL_ENV,
|
||||
) as rig,
|
||||
rig.candidate.scenario() as scenario,
|
||||
):
|
||||
key, identity = scenario_identity(scenario, rig.provider, team=False)
|
||||
response: Final = rig.candidate.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": identity["model"], "messages": [{"role": "user", "content": marker}]},
|
||||
key=key,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
|
||||
exported_llm_span(rig.collector, response.json()["id"])
|
||||
|
||||
on_stdout: Final = CONSOLE_SPAN in rig.proxy.log.read_text()
|
||||
assert on_stdout is printed, f"span JSON on proxy stdout: {on_stdout}, see {rig.proxy.log}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("callbacks", [[], ["langfuse_otel"]], ids=["key-only", "callback-listed"])
|
||||
def test_key_langfuse_without_operator_keys_prints_no_spans(
|
||||
gateway: Gateway, tmp_path: Path, callbacks: list[str]
|
||||
) -> None:
|
||||
marker: Final = "lf-console-key-" + uuid.uuid4().hex
|
||||
with (
|
||||
langfuse_rig(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{"callbacks": callbacks},
|
||||
marker,
|
||||
overrides={"OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT": "span_only"},
|
||||
remove_environment=INHERITED_OTEL_ENV,
|
||||
operator_keys=False,
|
||||
) as rig,
|
||||
rig.candidate.scenario() as scenario,
|
||||
):
|
||||
key: Final = scenario.key(
|
||||
key_alias="lf-console-key-" + uuid.uuid4().hex,
|
||||
metadata={
|
||||
"logging": [
|
||||
{
|
||||
"callback_name": "langfuse_otel",
|
||||
"callback_type": "success",
|
||||
"callback_vars": {
|
||||
"langfuse_public_key": "pk-lf-key",
|
||||
"langfuse_secret_key": "sk-lf-key",
|
||||
"langfuse_host": rig.collector.url,
|
||||
},
|
||||
}
|
||||
]
|
||||
},
|
||||
)
|
||||
model: Final = scenario.model(api_base=rig.provider.url + "/v1")
|
||||
response: Final = rig.candidate.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "messages": [{"role": "user", "content": marker}]},
|
||||
key=key,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
|
||||
eventually(
|
||||
watched_spans(rig.collector), lambda spans: any(marker in json.dumps(span) for span in spans), seconds=25
|
||||
)
|
||||
|
||||
on_stdout: Final = CONSOLE_SPAN in rig.proxy.log.read_text()
|
||||
assert not on_stdout, f"span JSON on proxy stdout, see {rig.proxy.log}"
|
||||
|
|
@ -0,0 +1,224 @@
|
|||
"""Without operator Langfuse keys on OTel v2, Langfuse logging the OTel v2 fan-out does not carry still reaches the
|
||||
tenant's own Langfuse project: a key's failure-only entry on a failed call, a config ``default_team_settings`` entry,
|
||||
and a key destination's call the proxy rejected before any provider. A ``success_and_failure`` key entry sends that
|
||||
project exactly one copy of a failed call."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterator
|
||||
from typing import Final, Literal
|
||||
|
||||
import pytest
|
||||
from integration._support.client import Gateway, gateway_from_environment, held
|
||||
from integration._support.otlp_sink import (
|
||||
AuditConfigWriter,
|
||||
Span,
|
||||
SpanSinks,
|
||||
langfuse_callback_variables,
|
||||
recorded_spans,
|
||||
spans_for_trace,
|
||||
)
|
||||
from integration._support.process import owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
OPERATION: Final = "gen_ai.operation.name"
|
||||
TAGS: Final = ["agent:orders", "run:r1"]
|
||||
CONFIG_TEAM: Final = f"lf-config-team-{uuid.uuid4().hex[:8]}"
|
||||
FAIL: Final = "please-fail"
|
||||
|
||||
|
||||
def _chat(request: Request) -> Reply:
|
||||
if FAIL.encode() in request.body:
|
||||
return Reply(
|
||||
status=500, body=json.dumps({"error": {"message": "upstream boom", "type": "server_error"}}).encode()
|
||||
)
|
||||
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": "ok"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def provider() -> Iterator[Wire]:
|
||||
with wire_server(_chat) as wire:
|
||||
yield wire
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def proxy(
|
||||
audit_sinks: SpanSinks,
|
||||
otel_audit_config: AuditConfigWriter,
|
||||
tmp_path_factory: pytest.TempPathFactory,
|
||||
) -> Iterator[Gateway]:
|
||||
directory: Final = tmp_path_factory.mktemp("langfuse-legacy-destinations")
|
||||
settings: Final[dict[str, JsonValue]] = {
|
||||
"callbacks": [],
|
||||
"default_team_settings": [
|
||||
{
|
||||
"team_id": CONFIG_TEAM,
|
||||
"success_callback": ["langfuse_otel"],
|
||||
"callback_vars": dict(langfuse_callback_variables(audit_sinks)),
|
||||
}
|
||||
],
|
||||
}
|
||||
with (
|
||||
gateway_from_environment() as gateway,
|
||||
owned_proxy_process(
|
||||
gateway,
|
||||
directory,
|
||||
{
|
||||
"LITELLM_OTEL_V2": "1",
|
||||
"OTEL_BSP_SCHEDULE_DELAY": "300",
|
||||
"OTEL_EXPORTER": "otlp_http",
|
||||
"OTEL_ENDPOINT": audit_sinks.operator,
|
||||
},
|
||||
config=otel_audit_config(directory, settings),
|
||||
remove_environment=("LANGFUSE_PUBLIC_KEY", "LANGFUSE_SECRET_KEY", "LANGFUSE_HOST"),
|
||||
workers=2,
|
||||
) as owned,
|
||||
):
|
||||
yield owned.gateway
|
||||
|
||||
|
||||
def _tenant_spans(url: str, since: int, trace_id: str) -> tuple[Span, ...]:
|
||||
return spans_for_trace(recorded_spans(url, since)[1], trace_id)
|
||||
|
||||
|
||||
def _tags(value: object) -> object:
|
||||
"""The tags a span carries: the legacy logger writes them as a JSON string, OTel v2 as a list."""
|
||||
return json.loads(value) if isinstance(value, str) else value
|
||||
|
||||
|
||||
def _settled(spans: tuple[Span, ...]) -> bool:
|
||||
"""The model span and the request's root span have both arrived."""
|
||||
return any(OPERATION in span["attributes"] for span in spans) and any(
|
||||
str(span["name"]).startswith("POST ") for span in spans
|
||||
)
|
||||
|
||||
|
||||
def _logging_entry(audit_sinks: SpanSinks, callback_type: str) -> list[JsonValue]:
|
||||
return [
|
||||
{
|
||||
"callback_name": "langfuse_otel",
|
||||
"callback_type": callback_type,
|
||||
"callback_vars": langfuse_callback_variables(audit_sinks),
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def _chat_status(proxy: Gateway, model: str, key: str, trace_id: str, content: str) -> int:
|
||||
body: Final[dict[str, JsonValue]] = {
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": content}],
|
||||
"metadata": {"tags": [*TAGS], "session_id": f"sess-{trace_id}", "trace_user_id": "user-42"},
|
||||
}
|
||||
return proxy.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
body,
|
||||
key=key,
|
||||
headers={"traceparent": f"00-{trace_id}-{uuid.uuid4().hex[:16]}-01"},
|
||||
).status_code
|
||||
|
||||
|
||||
def _trace_fields(spans: tuple[Span, ...]) -> list[tuple[object, object]]:
|
||||
return [
|
||||
(span["attributes"]["user.id"], _tags(span["attributes"].get("langfuse.trace.tags")))
|
||||
for span in spans
|
||||
if "user.id" in span["attributes"]
|
||||
]
|
||||
|
||||
|
||||
def _carries_the_trace_fields_once(spans: tuple[Span, ...]) -> bool:
|
||||
return _trace_fields(spans) == [("user-42", TAGS)]
|
||||
|
||||
|
||||
def _one_failed_call(spans: tuple[Span, ...]) -> bool:
|
||||
return (
|
||||
_settled(spans)
|
||||
and len([span for span in spans if OPERATION in span["attributes"]]) == 1
|
||||
and "litellm_request" not in {span["name"] for span in spans}
|
||||
)
|
||||
|
||||
|
||||
def _send(
|
||||
proxy: Gateway,
|
||||
provider: Wire,
|
||||
audit_sinks: SpanSinks,
|
||||
setup: Literal["failure", "config_team", "both"],
|
||||
expected: Callable[[tuple[Span, ...]], bool],
|
||||
) -> tuple[int, tuple[Span, ...]]:
|
||||
logging_entry: Final = _logging_entry(audit_sinks, "failure" if setup == "failure" else "success_and_failure")
|
||||
with proxy.scenario() as scenario:
|
||||
model: Final = scenario.model(api_base=provider.url + "/v1", num_retries=0)
|
||||
alias: Final = f"lf-legacy-{setup}-{uuid.uuid4().hex}"
|
||||
key: Final = (
|
||||
scenario.key(key_alias=alias, team_id=scenario.team(team_id=CONFIG_TEAM))
|
||||
if setup == "config_team"
|
||||
else scenario.key(key_alias=alias, metadata={"logging": logging_entry})
|
||||
)
|
||||
since: Final = recorded_spans(audit_sinks.tenant)[0]
|
||||
trace_id: Final = uuid.uuid4().hex
|
||||
content: Final = f"lf-legacy-{setup}-{trace_id}" + ("" if setup == "config_team" else f" {FAIL}")
|
||||
status: Final = _chat_status(proxy, model, key, trace_id, content)
|
||||
tenant: Final = held(
|
||||
lambda: _tenant_spans(audit_sinks.tenant, since, trace_id), expected, holding=3, seconds=40
|
||||
)
|
||||
return status, tenant
|
||||
|
||||
|
||||
@pytest.mark.parametrize("setup", ["failure", "config_team"])
|
||||
def test_a_tenant_langfuse_project_without_a_destination_still_gets_the_trace(
|
||||
setup: Literal["failure", "config_team"], proxy: Gateway, provider: Wire, audit_sinks: SpanSinks
|
||||
) -> None:
|
||||
status, tenant = _send(proxy, provider, audit_sinks, setup, _carries_the_trace_fields_once)
|
||||
|
||||
assert status == (500 if setup == "failure" else 200)
|
||||
assert _trace_fields(tenant) == [("user-42", TAGS)], [span["name"] for span in tenant]
|
||||
|
||||
|
||||
def test_a_key_destinations_call_rejected_before_any_provider_still_gets_the_trace_fields(
|
||||
proxy: Gateway, provider: Wire, audit_sinks: SpanSinks
|
||||
) -> None:
|
||||
with proxy.scenario() as scenario:
|
||||
model: Final = scenario.model(api_base=provider.url + "/v1", num_retries=0)
|
||||
key: Final = scenario.key(
|
||||
key_alias=f"lf-legacy-rejected-{uuid.uuid4().hex}",
|
||||
metadata={"logging": _logging_entry(audit_sinks, "success_and_failure")},
|
||||
rpm_limit=1,
|
||||
)
|
||||
since: Final = recorded_spans(audit_sinks.tenant)[0]
|
||||
statuses: Final = {
|
||||
trace_id: _chat_status(proxy, model, key, trace_id, f"lf-legacy-rejected-{trace_id}")
|
||||
for trace_id in [uuid.uuid4().hex for _ in range(4)]
|
||||
}
|
||||
rejected: Final = [trace_id for trace_id, status in statuses.items() if status == 429]
|
||||
assert rejected, statuses
|
||||
tenant: Final = held(
|
||||
lambda: _tenant_spans(audit_sinks.tenant, since, rejected[0]),
|
||||
_carries_the_trace_fields_once,
|
||||
holding=3,
|
||||
seconds=40,
|
||||
)
|
||||
|
||||
assert _trace_fields(tenant) == [("user-42", TAGS)], [span["name"] for span in tenant]
|
||||
|
||||
|
||||
def test_a_failed_call_reaches_a_success_and_failure_destination_once(
|
||||
proxy: Gateway, provider: Wire, audit_sinks: SpanSinks
|
||||
) -> None:
|
||||
status, tenant = _send(proxy, provider, audit_sinks, "both", _one_failed_call)
|
||||
|
||||
assert status == 500
|
||||
assert _one_failed_call(tenant), [span["name"] for span in tenant]
|
||||
|
|
@ -0,0 +1,186 @@
|
|||
"""A key's or team's own Langfuse project gets one copy of the call with the Langfuse trace fields when the operator
|
||||
has no Langfuse keys, whether or not ``langfuse_otel`` is listed in ``litellm_settings.callbacks``."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from typing import Final, Literal, cast
|
||||
|
||||
import pytest
|
||||
from integration._support.client import Gateway, gateway_from_environment, held
|
||||
from integration._support.otlp_sink import (
|
||||
AuditConfigWriter,
|
||||
Span,
|
||||
SpanSinks,
|
||||
langfuse_callback_variables,
|
||||
recorded_spans,
|
||||
spans_for_trace,
|
||||
)
|
||||
from integration._support.process import owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
OPERATION: Final = "gen_ai.operation.name"
|
||||
TAGS: Final = ["agent:orders", "run:r1"]
|
||||
LEGACY_SPANS: Final = frozenset({"litellm_request", "raw_gen_ai_request"})
|
||||
|
||||
|
||||
def _chat(_request: Request) -> 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": "ok"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def provider() -> Iterator[Wire]:
|
||||
with wire_server(_chat) as wire:
|
||||
yield wire
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", params=[["langfuse_otel"], []], ids=["listed", "not-listed"])
|
||||
def proxy(
|
||||
request: pytest.FixtureRequest,
|
||||
audit_sinks: SpanSinks,
|
||||
otel_audit_config: AuditConfigWriter,
|
||||
tmp_path_factory: pytest.TempPathFactory,
|
||||
) -> Iterator[Gateway]:
|
||||
directory: Final = tmp_path_factory.mktemp("langfuse-tenant-only")
|
||||
with (
|
||||
gateway_from_environment() as gateway,
|
||||
owned_proxy_process(
|
||||
gateway,
|
||||
directory,
|
||||
{
|
||||
"LITELLM_OTEL_V2": "1",
|
||||
"OTEL_BSP_SCHEDULE_DELAY": "300",
|
||||
"OTEL_EXPORTER": "otlp_http",
|
||||
"OTEL_ENDPOINT": audit_sinks.operator,
|
||||
},
|
||||
config=otel_audit_config(directory, {"callbacks": cast("list[JsonValue]", request.param)}),
|
||||
remove_environment=("LANGFUSE_PUBLIC_KEY", "LANGFUSE_SECRET_KEY", "LANGFUSE_HOST"),
|
||||
workers=2,
|
||||
) as owned,
|
||||
):
|
||||
yield owned.gateway
|
||||
|
||||
|
||||
def _trace_spans(url: str, since: int, trace_id: str) -> tuple[Span, ...]:
|
||||
return spans_for_trace(recorded_spans(url, since)[1], trace_id)
|
||||
|
||||
|
||||
def _model_spans(spans: tuple[Span, ...]) -> tuple[Span, ...]:
|
||||
return tuple(span for span in spans if OPERATION in span["attributes"])
|
||||
|
||||
|
||||
def _one_model_span(spans: tuple[Span, ...]) -> bool:
|
||||
return len(_model_spans(spans)) == 1 and not LEGACY_SPANS & {str(span["name"]) for span in spans}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("capture", [None, "span_only"])
|
||||
@pytest.mark.parametrize("level", ["key", "team"])
|
||||
def test_a_tenant_langfuse_project_gets_user_session_and_tags_without_operator_keys(
|
||||
level: Literal["key", "team"], capture: str | None, proxy: Gateway, provider: Wire, audit_sinks: SpanSinks
|
||||
) -> None:
|
||||
logging: Final[list[JsonValue]] = [
|
||||
{
|
||||
"callback_name": "langfuse_otel",
|
||||
"callback_type": "success",
|
||||
"callback_vars": {
|
||||
**langfuse_callback_variables(audit_sinks),
|
||||
**({"capture_message_content": capture} if capture else {}),
|
||||
},
|
||||
}
|
||||
]
|
||||
with proxy.scenario() as scenario:
|
||||
model: Final = scenario.model(api_base=provider.url + "/v1")
|
||||
alias: Final = f"lf-tenant-only-{uuid.uuid4().hex}"
|
||||
key: Final = (
|
||||
scenario.key(key_alias=alias, metadata={"logging": logging})
|
||||
if level == "key"
|
||||
else scenario.key(key_alias=alias, team_id=scenario.team(metadata={"logging": logging}))
|
||||
)
|
||||
since: Final = (recorded_spans(audit_sinks.tenant)[0], recorded_spans(audit_sinks.operator)[0])
|
||||
trace_id: Final = uuid.uuid4().hex
|
||||
session: Final = f"sess-{trace_id}"
|
||||
body: Final[dict[str, JsonValue]] = {
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": f"lf-tenant-only-{trace_id}"}],
|
||||
"metadata": {"tags": [*TAGS], "session_id": session, "trace_user_id": "user-42"},
|
||||
}
|
||||
response: Final = proxy.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
body,
|
||||
key=key,
|
||||
headers={"traceparent": f"00-{trace_id}-{uuid.uuid4().hex[:16]}-01"},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
|
||||
tenant: Final = _model_spans(
|
||||
held(lambda: _trace_spans(audit_sinks.tenant, since[0], trace_id), _one_model_span, holding=3, seconds=40)
|
||||
)
|
||||
operator: Final = _model_spans(
|
||||
held(lambda: _trace_spans(audit_sinks.operator, since[1], trace_id), _one_model_span, holding=3, seconds=40)
|
||||
)
|
||||
|
||||
assert [
|
||||
{name: span["attributes"].get(name) for name in ("user.id", "session.id", "langfuse.trace.tags")}
|
||||
for span in tenant
|
||||
] == [{"user.id": "user-42", "session.id": session, "langfuse.trace.tags": TAGS}], [s["name"] for s in tenant]
|
||||
assert "langfuse.observation.metadata" in tenant[0]["attributes"], sorted(tenant[0]["attributes"])
|
||||
assert (f"lf-tenant-only-{trace_id}" in json.dumps(tenant[0]["attributes"])) == (capture == "span_only"), sorted(
|
||||
tenant[0]["attributes"]
|
||||
)
|
||||
assert len(operator) == 1, [span["name"] for span in operator]
|
||||
assert f"lf-tenant-only-{trace_id}" not in json.dumps(operator[0]["attributes"]), sorted(operator[0]["attributes"])
|
||||
assert [
|
||||
name for name in operator[0]["attributes"] if name in ("user.id", "session.id") or name.startswith("langfuse.")
|
||||
] == [], sorted(operator[0]["attributes"])
|
||||
|
||||
|
||||
def test_a_tenant_langfuse_project_gets_the_keys_langfuse_environment(
|
||||
proxy: Gateway, provider: Wire, audit_sinks: SpanSinks
|
||||
) -> None:
|
||||
logging: Final[list[JsonValue]] = [
|
||||
{
|
||||
"callback_name": "langfuse_otel",
|
||||
"callback_type": "success",
|
||||
"callback_vars": {**langfuse_callback_variables(audit_sinks), "langfuse_environment": "prod"},
|
||||
}
|
||||
]
|
||||
with proxy.scenario() as scenario:
|
||||
model: Final = scenario.model(api_base=provider.url + "/v1")
|
||||
key: Final = scenario.key(key_alias=f"lf-tenant-env-{uuid.uuid4().hex}", metadata={"logging": logging})
|
||||
since: Final = (recorded_spans(audit_sinks.tenant)[0], recorded_spans(audit_sinks.operator)[0])
|
||||
trace_id: Final = uuid.uuid4().hex
|
||||
response: Final = proxy.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "messages": [{"role": "user", "content": f"lf-tenant-env-{trace_id}"}]},
|
||||
key=key,
|
||||
headers={"traceparent": f"00-{trace_id}-{uuid.uuid4().hex[:16]}-01"},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
|
||||
tenant: Final = held(
|
||||
lambda: _trace_spans(audit_sinks.tenant, since[0], trace_id), _one_model_span, holding=3, seconds=40
|
||||
)
|
||||
operator: Final = held(
|
||||
lambda: _trace_spans(audit_sinks.operator, since[1], trace_id), _one_model_span, holding=3, seconds=40
|
||||
)
|
||||
|
||||
assert [span["resource"].get("langfuse.environment") for span in tenant] == ["prod"] * len(tenant), [
|
||||
span["name"] for span in tenant
|
||||
]
|
||||
assert not any("langfuse.environment" in span["resource"] for span in operator), [span["name"] for span in operator]
|
||||
|
|
@ -70,7 +70,9 @@ def bedrock_converse_peer(request: Request) -> Reply:
|
|||
)
|
||||
|
||||
|
||||
def send_to_bedrock_converse(gateway: Gateway, history: list[JsonValue]) -> tuple[httpx.Response, tuple[Request, ...]]:
|
||||
def send_to_bedrock_converse(
|
||||
gateway: Gateway, history: list[JsonValue], drop_params: bool = True
|
||||
) -> tuple[httpx.Response, tuple[Request, ...]]:
|
||||
with wire_server(bedrock_converse_peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(
|
||||
model=f"bedrock/converse/{BEDROCK_GPT}",
|
||||
|
|
@ -79,7 +81,15 @@ def send_to_bedrock_converse(gateway: Gateway, history: list[JsonValue]) -> tupl
|
|||
aws_bedrock_runtime_endpoint=wire.url,
|
||||
)
|
||||
response: Final = gateway.request(
|
||||
"POST", "/v1/responses", {"model": model, "input": history, "store": False, "cache": {"no-cache": True}}
|
||||
"POST",
|
||||
"/v1/responses",
|
||||
{
|
||||
"model": model,
|
||||
"input": history,
|
||||
"store": False,
|
||||
"cache": {"no-cache": True},
|
||||
"drop_params": drop_params,
|
||||
},
|
||||
)
|
||||
return response, wire.drain()
|
||||
|
||||
|
|
@ -113,6 +123,20 @@ def test_codex_tool_list_repeated_on_every_turn_reaches_bedrock_converse_once(ga
|
|||
assert sorted(names) == ["collaboration__followup_task", "exec_command"], names
|
||||
|
||||
|
||||
def test_codex_tool_list_repeated_on_every_turn_reaches_bedrock_converse_as_sent_without_drop_params(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
tools: Final = tool_list_item(function_tool("exec_command"))
|
||||
history: Final[list[JsonValue]] = [tools, user_turn(f"first {marker}"), assistant_turn("ok"), tools]
|
||||
|
||||
response, forwarded = send_to_bedrock_converse(gateway, history, drop_params=False)
|
||||
|
||||
assert response.status_code == 400, response.text
|
||||
assert "already defined" in response.text, response.text
|
||||
assert [str(spec["name"]) for spec in converse_tool_specs(forwarded[0])] == ["exec_command", "exec_command"]
|
||||
|
||||
|
||||
def test_newest_copy_of_a_repeated_codex_tool_is_what_bedrock_converse_receives(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
history: Final[list[JsonValue]] = [
|
||||
|
|
@ -195,6 +219,7 @@ def test_codex_tool_list_repeated_on_every_turn_reaches_bedrock_mantle_once(gate
|
|||
user_turn(f"second {marker}"),
|
||||
],
|
||||
"store": False,
|
||||
"drop_params": True,
|
||||
},
|
||||
)
|
||||
forwarded: Final = wire.drain()
|
||||
|
|
|
|||
|
|
@ -0,0 +1,351 @@
|
|||
import json
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass, field
|
||||
from itertools import product
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import JSON_OBJECT, Gateway, gateway_from_environment, list_value, object_value
|
||||
from integration._support.process import graceful_stop_seconds, owned_proxy
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
|
||||
|
||||
_MODEL: Final = "background-scan-chat"
|
||||
_CLAUDE: Final = "background-scan-claude"
|
||||
_RATE: Final = 5
|
||||
_TRIGGER: Final = "FORBIDDEN"
|
||||
_DENIAL: Final = "Output withheld by policy."
|
||||
_CLEAN: Final = tuple(f"w{index} " for index in range(40))
|
||||
_LEAKY: Final = ("w0 ", f"{_TRIGGER} ", *_CLEAN[2:])
|
||||
_OUTAGE: Final = ("w0 ", "OUTAGE ", *_CLEAN[2:])
|
||||
_HOLD_SECONDS: Final = 8.0
|
||||
_ALLOWED: Final = json.dumps({"action": "NONE", "outputs": [], "assessments": []}).encode()
|
||||
_INTERVENED: Final = json.dumps(
|
||||
{
|
||||
"action": "GUARDRAIL_INTERVENED",
|
||||
"outputs": [{"text": _DENIAL}],
|
||||
"assessments": [{"topicPolicy": {"topics": [{"name": "forbidden", "type": "DENY", "action": "BLOCKED"}]}}],
|
||||
}
|
||||
).encode()
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _Scans:
|
||||
"""The scripted ApplyGuardrail peer holds each output scan until `release` is set or the hold runs out."""
|
||||
|
||||
release: threading.Event = field(default_factory=threading.Event)
|
||||
started_at: list[float] = field(default_factory=list)
|
||||
replied_at: list[float] = field(default_factory=list)
|
||||
|
||||
def answer(self, request: Request) -> Reply:
|
||||
self.started_at.append(time.monotonic())
|
||||
body: Final = request.body.decode()
|
||||
if '"OUTPUT"' in body:
|
||||
self.release.wait(_HOLD_SECONDS)
|
||||
self.replied_at.append(time.monotonic())
|
||||
if "OUTAGE" in body:
|
||||
return Reply(status=500, body=b'{"message": "synthetic ApplyGuardrail outage"}')
|
||||
return Reply(body=_INTERVENED if _TRIGGER in body else _ALLOWED)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Rig:
|
||||
gateway: Gateway
|
||||
scans: list[_Scans]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Streamed:
|
||||
pieces: tuple[str, ...]
|
||||
arrived_at: tuple[float, ...]
|
||||
finished_after: tuple[int, ...]
|
||||
errors: tuple[str, ...]
|
||||
|
||||
|
||||
def _chat_sse(pieces: tuple[str, ...]) -> tuple[bytes, ...]:
|
||||
def frame(delta: dict[str, JsonValue], finish: str | None) -> bytes:
|
||||
chunk: Final = {
|
||||
"id": "chatcmpl-background",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": _MODEL,
|
||||
"choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
|
||||
}
|
||||
return b"data: " + json.dumps(chunk).encode() + b"\n\n"
|
||||
|
||||
return (
|
||||
*(frame({"role": "assistant", "content": piece}, None) for piece in pieces),
|
||||
frame({}, "stop"),
|
||||
b"data: [DONE]\n\n",
|
||||
)
|
||||
|
||||
|
||||
def _anthropic_sse(pieces: tuple[str, ...]) -> tuple[bytes, ...]:
|
||||
def event(payload: dict[str, JsonValue]) -> bytes:
|
||||
return f"event: {payload['type']}\ndata: {json.dumps(payload)}\n\n".encode()
|
||||
|
||||
message: Final[dict[str, JsonValue]] = {
|
||||
"id": "msg_background",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-4-6",
|
||||
"content": [],
|
||||
"stop_reason": None,
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 5, "output_tokens": 1},
|
||||
}
|
||||
return (
|
||||
event({"type": "message_start", "message": message}),
|
||||
event({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
|
||||
*(
|
||||
event({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": piece}})
|
||||
for piece in pieces
|
||||
),
|
||||
event({"type": "content_block_stop", "index": 0}),
|
||||
event({"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": len(pieces)}}),
|
||||
event({"type": "message_stop"}),
|
||||
)
|
||||
|
||||
|
||||
def _model_reply(request: Request) -> Reply:
|
||||
prompt: Final = str(object_value(list_value(object_value(json.loads(request.body))["messages"])[-1])["content"])
|
||||
pieces: Final = _LEAKY if "leak" in prompt else _OUTAGE if "outage" in prompt else _CLEAN
|
||||
return Reply(
|
||||
chunks=_anthropic_sse(pieces) if request.target.endswith("/v1/messages") else _chat_sse(pieces),
|
||||
content_type="text/event-stream",
|
||||
pause_between_chunks=0.05,
|
||||
)
|
||||
|
||||
|
||||
def _bedrock_rail(name: str, endpoint: str, background: bool) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"guardrail_name": name,
|
||||
"litellm_params": {
|
||||
"guardrail": "bedrock",
|
||||
"mode": "post_call",
|
||||
"default_on": False,
|
||||
"disable_exception_on_block": True,
|
||||
"streaming_buffer_until_moderated": False,
|
||||
"streaming_end_of_stream_only": False,
|
||||
"streaming_sampling_rate": _RATE,
|
||||
**({"streaming_scan_in_background": True} if background else {}),
|
||||
"guardrailIdentifier": "synthetic" + uuid.uuid4().hex[:8],
|
||||
"guardrailVersion": "DRAFT",
|
||||
"aws_region_name": "us-east-1",
|
||||
"aws_access_key_id": "AKIASYNTHETICGUARDRAIL",
|
||||
"aws_secret_access_key": "synthetic-secret",
|
||||
"aws_bedrock_runtime_endpoint": endpoint,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
|
||||
directory: Final = tmp_path_factory.mktemp("background-stream-scan")
|
||||
scans: Final = [_Scans()]
|
||||
stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
|
||||
config: Final = directory / "background-stream-scan.yaml"
|
||||
with wire_server(lambda request: scans[0].answer(request)) as policy, wire_server(_model_reply) as model:
|
||||
config.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
**stock,
|
||||
"model_list": [
|
||||
*list_value(stock["model_list"]),
|
||||
{
|
||||
"model_name": _MODEL,
|
||||
"litellm_params": {
|
||||
"model": f"openai/{_MODEL}",
|
||||
"api_base": f"{model.url}/v1",
|
||||
"api_key": "synthetic",
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": _CLAUDE,
|
||||
"litellm_params": {
|
||||
"model": "anthropic/claude-sonnet-4-6",
|
||||
"api_base": model.url,
|
||||
"api_key": "synthetic",
|
||||
},
|
||||
},
|
||||
],
|
||||
"guardrails": [
|
||||
_bedrock_rail("background", policy.url, background=True),
|
||||
_bedrock_rail("inline", policy.url, background=False),
|
||||
],
|
||||
}
|
||||
)
|
||||
)
|
||||
with gateway_from_environment() as shared, owned_proxy(shared, directory, {}, config=config) as candidate:
|
||||
yield _Rig(candidate, scans)
|
||||
|
||||
|
||||
def _chat_text(event: dict[str, JsonValue]) -> str | None:
|
||||
texts: Final = [
|
||||
object_value(object_value(choice).get("delta") or {}).get("content")
|
||||
for choice in list_value(event.get("choices") or [])
|
||||
]
|
||||
return "".join(text for text in texts if isinstance(text, str)) or None
|
||||
|
||||
|
||||
def _messages_text(event: dict[str, JsonValue]) -> str | None:
|
||||
delta: Final = object_value(event.get("delta") or {}) if event.get("type") == "content_block_delta" else {}
|
||||
text: Final = delta.get("text")
|
||||
return text if isinstance(text, str) and text else None
|
||||
|
||||
|
||||
def _responses_text(event: dict[str, JsonValue]) -> str | None:
|
||||
delta: Final = event.get("delta") if event.get("type") == "response.output_text.delta" else None
|
||||
return delta if isinstance(delta, str) and delta else None
|
||||
|
||||
|
||||
def _chat_finished(event: dict[str, JsonValue]) -> bool:
|
||||
return any(object_value(choice).get("finish_reason") for choice in list_value(event.get("choices") or []))
|
||||
|
||||
|
||||
def _messages_finished(event: dict[str, JsonValue]) -> bool:
|
||||
return event.get("type") == "message_delta"
|
||||
|
||||
|
||||
def _responses_finished(event: dict[str, JsonValue]) -> bool:
|
||||
return event.get("type") in ("response.completed", "response.incomplete")
|
||||
|
||||
|
||||
def _error_message(event: dict[str, JsonValue]) -> str | None:
|
||||
error: Final = event.get("error")
|
||||
if isinstance(error, dict):
|
||||
return str(error.get("message"))
|
||||
return str(event.get("message")) if event.get("type") == "error" else None
|
||||
|
||||
|
||||
def _chat_body(prompt: str) -> dict[str, JsonValue]:
|
||||
return {"model": _MODEL, "messages": [{"role": "user", "content": prompt}]}
|
||||
|
||||
|
||||
def _messages_body(prompt: str) -> dict[str, JsonValue]:
|
||||
return {**_chat_body(prompt), "model": _CLAUDE, "max_tokens": 256}
|
||||
|
||||
|
||||
def _responses_body(prompt: str) -> dict[str, JsonValue]:
|
||||
return {"model": _CLAUDE, "max_output_tokens": 256, "input": prompt}
|
||||
|
||||
|
||||
_ROUTES: Final = {
|
||||
"/v1/chat/completions": (_chat_body, _chat_text, _chat_finished),
|
||||
"/v1/messages": (_messages_body, _messages_text, _messages_finished),
|
||||
"/v1/responses": (_responses_body, _responses_text, _responses_finished),
|
||||
}
|
||||
|
||||
|
||||
def _stream(
|
||||
rig: _Rig, guardrail: str, prompt: str, release_after: int, route: str = "/v1/chat/completions"
|
||||
) -> tuple[_Streamed, _Scans]:
|
||||
"""Stream `route`, releasing held scans once `release_after` text pieces have arrived."""
|
||||
scans: Final = _Scans()
|
||||
rig.scans[0] = scans
|
||||
return _read(rig, guardrail, prompt, route, scans, release_after), scans
|
||||
|
||||
|
||||
def _read(rig: _Rig, guardrail: str, prompt: str, route: str, scans: _Scans, release_after: int) -> _Streamed:
|
||||
body_of, text_of, finished = _ROUTES[route]
|
||||
pieces: Final[list[str]] = []
|
||||
arrived_at: Final[list[float]] = []
|
||||
finished_after: Final[list[int]] = []
|
||||
errors: Final[list[str]] = []
|
||||
with rig.gateway.client.stream(
|
||||
"POST",
|
||||
route,
|
||||
json={**body_of(f"{prompt} {uuid.uuid4().hex}"), "stream": True, "guardrails": [guardrail]},
|
||||
headers={"Authorization": f"Bearer {rig.gateway.key}"},
|
||||
) as response:
|
||||
assert response.status_code == 200, response.read()
|
||||
for line in response.iter_lines():
|
||||
if not line.startswith("data: ") or line == "data: [DONE]":
|
||||
continue
|
||||
event = object_value(json.loads(line[len("data: ") :]))
|
||||
text = text_of(event)
|
||||
if text is not None:
|
||||
pieces.append(text)
|
||||
arrived_at.append(time.monotonic())
|
||||
if finished(event):
|
||||
finished_after.append(len(pieces))
|
||||
error = _error_message(event)
|
||||
if error is not None:
|
||||
errors.append(error)
|
||||
if len(pieces) >= release_after:
|
||||
scans.release.set()
|
||||
return _Streamed(tuple(pieces), tuple(arrived_at), tuple(finished_after), tuple(errors))
|
||||
|
||||
|
||||
def test_chunks_after_the_nth_keep_streaming_while_its_scan_is_pending(rig: _Rig) -> None:
|
||||
streamed, scans = _stream(rig, "background", "clean", release_after=_RATE + 1)
|
||||
assert streamed.pieces == _CLEAN
|
||||
assert streamed.arrived_at[_RATE] < scans.replied_at[0], "the chunk after the Nth waited for the Nth chunk's scan"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", tuple(_ROUTES))
|
||||
def test_block_verdict_stops_the_stream_at_the_next_chunk(rig: _Rig, route: str) -> None:
|
||||
streamed, scans = _stream(rig, "background", "leak", release_after=_RATE + 1, route=route)
|
||||
*before, last = streamed.pieces
|
||||
assert last == _DENIAL
|
||||
assert _DENIAL not in before
|
||||
assert _RATE < len(before) < len(_LEAKY), before
|
||||
assert tuple(before) == _LEAKY[: len(before)]
|
||||
assert streamed.arrived_at[_RATE] < scans.replied_at[0]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", tuple(_ROUTES))
|
||||
def test_block_verdict_that_lands_after_the_upstream_finished_replaces_the_finish(rig: _Rig, route: str) -> None:
|
||||
streamed, scans = _stream(rig, "background", "leak", release_after=len(_LEAKY) + 1, route=route)
|
||||
assert streamed.pieces == (*_LEAKY, _DENIAL)
|
||||
assert streamed.arrived_at[len(_LEAKY) - 1] < scans.replied_at[0], "the verdict landed before the upstream finished"
|
||||
assert streamed.finished_after == (len(streamed.pieces),), "the stream finished before the block verdict"
|
||||
overlaps: Final = zip(scans.started_at[1:], scans.replied_at[:-1], strict=True)
|
||||
assert all(started >= replied for started, replied in overlaps), "a scan started while another was in flight"
|
||||
|
||||
|
||||
def test_without_the_flag_the_nth_chunk_still_waits_for_its_scan(rig: _Rig) -> None:
|
||||
streamed, scans = _stream(rig, "inline", "clean", release_after=_RATE - 1)
|
||||
assert streamed.pieces == _CLEAN
|
||||
assert scans.replied_at[0] < streamed.arrived_at[_RATE - 1]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", tuple(_ROUTES))
|
||||
def test_a_failed_background_scan_ends_the_stream_with_the_inline_scan_error(rig: _Rig, route: str) -> None:
|
||||
background, _ = _stream(rig, "background", "outage", release_after=_RATE + 1, route=route)
|
||||
inline, _ = _stream(rig, "inline", "outage", release_after=0, route=route)
|
||||
assert len(inline.errors) == 1, inline
|
||||
assert background.errors == inline.errors
|
||||
assert _RATE < len(background.pieces) < len(_OUTAGE), background.pieces
|
||||
assert background.pieces == _OUTAGE[: len(background.pieces)]
|
||||
assert background.finished_after == ()
|
||||
|
||||
|
||||
def test_a_mixed_burst_ends_every_stream_on_its_own_verdict_while_the_scan_peer_fails_some(rig: _Rig) -> None:
|
||||
scans: Final = _Scans()
|
||||
scans.release.set()
|
||||
rig.scans[0] = scans
|
||||
calls: Final = tuple(product(("clean", "leak", "outage"), (*_ROUTES, *_ROUTES)))
|
||||
with ThreadPoolExecutor(len(calls)) as pool:
|
||||
served: Final = tuple(pool.map(lambda call: _read(rig, "background", call[0], call[1], scans, 0), calls))
|
||||
for (kind, route), streamed in zip(calls, served, strict=True):
|
||||
match kind:
|
||||
case "clean":
|
||||
assert (streamed.pieces, streamed.errors) == (_CLEAN, ()), (route, streamed)
|
||||
assert streamed.finished_after == (len(_CLEAN),), (route, streamed)
|
||||
case "leak":
|
||||
assert streamed.pieces[-1] == _DENIAL and streamed.errors == (), (route, streamed)
|
||||
assert streamed.finished_after == (len(streamed.pieces),), (route, streamed)
|
||||
case _:
|
||||
assert len(streamed.errors) == 1 and streamed.finished_after == (), (route, streamed)
|
||||
assert streamed.pieces == _OUTAGE[: len(streamed.pieces)], (route, streamed)
|
||||
after: Final = _read(rig, "background", "clean", "/v1/chat/completions", scans, 0)
|
||||
assert after.pieces == _CLEAN
|
||||
20
tests/unit/caching/_event_loop_reuse.py
Normal file
20
tests/unit/caching/_event_loop_reuse.py
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
from collections.abc import Callable
|
||||
from typing import Final, TypeVar
|
||||
|
||||
import pytest
|
||||
|
||||
LoopResult = TypeVar("LoopResult")
|
||||
|
||||
|
||||
def runs_until_a_loop_address_repeats(
|
||||
run_on_new_loop: Callable[[], tuple[int, LoopResult]],
|
||||
earlier: tuple[tuple[int, LoopResult], ...] = (),
|
||||
limit: int = 50,
|
||||
) -> tuple[tuple[int, LoopResult], ...]:
|
||||
if len(earlier) == limit:
|
||||
pytest.skip(f"none of {limit} collected loops had its address handed to a later loop")
|
||||
latest: Final = run_on_new_loop()
|
||||
runs: Final = (*earlier, latest)
|
||||
if latest[0] in {loop_id for loop_id, _ in earlier}:
|
||||
return runs
|
||||
return runs_until_a_loop_address_repeats(run_on_new_loop, runs, limit)
|
||||
|
|
@ -10,6 +10,7 @@ never closed, because litellm does not own its lifecycle.
|
|||
import asyncio
|
||||
import gc
|
||||
import weakref
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import httpx
|
||||
|
|
@ -18,6 +19,7 @@ from redis.asyncio import ConnectionPool, Redis
|
|||
|
||||
from litellm.caching.evicted_client_closer import EvictedClientCloser
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from tests.unit.caching._event_loop_reuse import runs_until_a_loop_address_repeats
|
||||
|
||||
|
||||
class FakeClock:
|
||||
|
|
@ -41,6 +43,14 @@ class AsyncClient:
|
|||
self.closed = True
|
||||
|
||||
|
||||
class RecordingAsyncClient:
|
||||
def __init__(self) -> None:
|
||||
self.close_calls = 0
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self.close_calls += 1
|
||||
|
||||
|
||||
class SyncClient:
|
||||
def __init__(self) -> None:
|
||||
self.closed = False
|
||||
|
|
@ -266,6 +276,100 @@ def test_sync_client_evicted_outside_an_event_loop_is_still_closed():
|
|||
assert closer.pending_count == 0
|
||||
|
||||
|
||||
def test_async_client_from_collected_loop_is_not_closed_by_fresh_loops() -> None:
|
||||
clock: Final = FakeClock()
|
||||
closer: Final = EvictedClientCloser(grace_seconds=0, clock=clock)
|
||||
|
||||
async def reap_then_schedule(client: RecordingAsyncClient) -> None:
|
||||
closer.reap()
|
||||
await asyncio.sleep(0)
|
||||
closer.schedule(client)
|
||||
|
||||
def run_on_new_loop() -> tuple[int, RecordingAsyncClient]:
|
||||
client: Final = RecordingAsyncClient()
|
||||
closer.mark_owned(client)
|
||||
loop: Final = asyncio.new_event_loop()
|
||||
loop_id: Final = id(loop)
|
||||
try:
|
||||
loop.run_until_complete(reap_then_schedule(client))
|
||||
finally:
|
||||
loop.close()
|
||||
del loop
|
||||
gc.collect()
|
||||
return loop_id, client
|
||||
|
||||
runs: Final = runs_until_a_loop_address_repeats(run_on_new_loop)
|
||||
|
||||
assert all(client.close_calls == 0 for _, client in runs)
|
||||
|
||||
|
||||
def test_async_client_queued_on_a_live_loop_is_closed_only_by_that_loop() -> None:
|
||||
clock: Final = FakeClock()
|
||||
closer: Final = EvictedClientCloser(grace_seconds=0, clock=clock)
|
||||
client: Final = RecordingAsyncClient()
|
||||
closer.mark_owned(client)
|
||||
owner_loop: Final = asyncio.new_event_loop()
|
||||
other_loop: Final = asyncio.new_event_loop()
|
||||
|
||||
async def schedule() -> None:
|
||||
closer.schedule(client)
|
||||
|
||||
async def reap() -> None:
|
||||
closer.reap()
|
||||
await asyncio.sleep(0)
|
||||
|
||||
try:
|
||||
owner_loop.run_until_complete(schedule())
|
||||
other_loop.run_until_complete(reap())
|
||||
closes_after_other_loop_reaped: Final = client.close_calls
|
||||
pending_after_other_loop_reaped: Final = closer.pending_count
|
||||
owner_loop.run_until_complete(reap())
|
||||
finally:
|
||||
other_loop.close()
|
||||
owner_loop.close()
|
||||
|
||||
assert closes_after_other_loop_reaped == 0
|
||||
assert pending_after_other_loop_reaped == 1
|
||||
assert client.close_calls == 1
|
||||
assert closer.pending_count == 0
|
||||
|
||||
|
||||
def test_closes_queued_on_a_collected_loop_free_capacity() -> None:
|
||||
clock: Final = FakeClock()
|
||||
closer: Final = EvictedClientCloser(grace_seconds=0, max_pending=1, clock=clock)
|
||||
first: Final = RecordingAsyncClient()
|
||||
closer.mark_owned(first)
|
||||
loop_a: Final = asyncio.new_event_loop()
|
||||
|
||||
async def schedule_first() -> None:
|
||||
closer.schedule(first)
|
||||
|
||||
try:
|
||||
loop_a.run_until_complete(schedule_first())
|
||||
finally:
|
||||
loop_a.close()
|
||||
del loop_a
|
||||
gc.collect()
|
||||
|
||||
second: Final = RecordingAsyncClient()
|
||||
closer.mark_owned(second)
|
||||
loop_b: Final = asyncio.new_event_loop()
|
||||
|
||||
async def schedule_and_reap_second() -> None:
|
||||
closer.schedule(second)
|
||||
closer.reap()
|
||||
await asyncio.sleep(0)
|
||||
|
||||
try:
|
||||
loop_b.run_until_complete(schedule_and_reap_second())
|
||||
finally:
|
||||
loop_b.close()
|
||||
|
||||
assert first.close_calls == 0
|
||||
assert second.close_calls == 1
|
||||
assert closer.pending_count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_async_client_waits_for_a_loop_rather_than_being_dropped():
|
||||
clock = FakeClock()
|
||||
|
|
@ -289,8 +393,7 @@ async def test_an_async_client_waits_for_a_loop_rather_than_being_dropped():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_client_evicted_on_another_event_loop_is_left_alone():
|
||||
"""Closing a client bound to a different loop would schedule work on that loop."""
|
||||
async def test_a_client_evicted_on_a_collected_loop_is_dropped():
|
||||
clock = FakeClock()
|
||||
closer = make_closer(clock)
|
||||
client = AsyncClient()
|
||||
|
|
@ -310,7 +413,7 @@ async def test_a_client_evicted_on_another_event_loop_is_left_alone():
|
|||
await asyncio.sleep(0.05)
|
||||
|
||||
assert client.closed is False
|
||||
assert closer.pending_count == 1
|
||||
assert closer.pending_count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ See: https://github.com/BerriAI/litellm/pull/22247
|
|||
|
||||
import asyncio, gc, httpx, importlib, litellm, os, threading, time, weakref
|
||||
import warnings
|
||||
from typing import Final
|
||||
|
||||
import aiohttp
|
||||
import pytest
|
||||
|
|
@ -29,6 +30,41 @@ from litellm.llms.custom_httpx.http_handler import(
|
|||
from litellm.types.utils import LlmProviders
|
||||
from litellm.utils import _invalidate_model_cost_lowercase_map
|
||||
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
|
||||
from tests.unit.caching._event_loop_reuse import runs_until_a_loop_address_repeats
|
||||
|
||||
|
||||
async def _get_and_set_cache(cache: LLMClientCache) -> object | None:
|
||||
cached: Final = await cache.async_get_cache("client")
|
||||
await cache.async_set_cache("client", object())
|
||||
return cached
|
||||
|
||||
|
||||
def _cache_result_on_new_loop(cache: LLMClientCache) -> tuple[int, object | None]:
|
||||
loop: Final = asyncio.new_event_loop()
|
||||
loop_id: Final = id(loop)
|
||||
try:
|
||||
result: Final = loop.run_until_complete(_get_and_set_cache(cache))
|
||||
finally:
|
||||
loop.close()
|
||||
del loop
|
||||
gc.collect()
|
||||
return loop_id, result
|
||||
|
||||
|
||||
async def _get_cache_key(cache: LLMClientCache) -> str:
|
||||
return cache.update_cache_key_with_event_loop("client")
|
||||
|
||||
|
||||
def _cache_key_on_new_loop(cache: LLMClientCache) -> tuple[int, str]:
|
||||
loop: Final = asyncio.new_event_loop()
|
||||
loop_id: Final = id(loop)
|
||||
try:
|
||||
key: Final = loop.run_until_complete(_get_cache_key(cache))
|
||||
finally:
|
||||
loop.close()
|
||||
del loop
|
||||
gc.collect()
|
||||
return loop_id, key
|
||||
|
||||
|
||||
class MockAsyncClient:
|
||||
|
|
@ -51,6 +87,44 @@ class MockSyncClient:
|
|||
self.closed = True
|
||||
|
||||
|
||||
def test_async_client_cache_does_not_reuse_clients_between_collected_loops() -> None:
|
||||
cache: Final = LLMClientCache()
|
||||
runs: Final = runs_until_a_loop_address_repeats(lambda: _cache_result_on_new_loop(cache))
|
||||
|
||||
assert all(result is None for _, result in runs)
|
||||
|
||||
|
||||
def test_async_client_cache_keys_differ_for_every_collected_loop() -> None:
|
||||
cache: Final = LLMClientCache()
|
||||
runs: Final = runs_until_a_loop_address_repeats(lambda: _cache_key_on_new_loop(cache))
|
||||
keys: Final = tuple(key for _, key in runs)
|
||||
|
||||
assert len(set(keys)) == len(keys)
|
||||
|
||||
|
||||
class _LoopWithoutWeakReferences:
|
||||
__slots__ = ()
|
||||
|
||||
|
||||
def _cache_key_while_running(cache: LLMClientCache, loop: _LoopWithoutWeakReferences) -> str:
|
||||
asyncio._set_running_loop(loop)
|
||||
try:
|
||||
return cache.update_cache_key_with_event_loop("client")
|
||||
finally:
|
||||
asyncio._set_running_loop(None)
|
||||
|
||||
|
||||
def test_cache_keys_stay_per_loop_on_loops_without_weak_references() -> None:
|
||||
cache: Final = LLMClientCache()
|
||||
loop_a: Final = _LoopWithoutWeakReferences()
|
||||
loop_b: Final = _LoopWithoutWeakReferences()
|
||||
|
||||
first_key_on_a: Final = _cache_key_while_running(cache, loop_a)
|
||||
|
||||
assert _cache_key_while_running(cache, loop_a) == first_key_on_a
|
||||
assert _cache_key_while_running(cache, loop_b) != first_key_on_a
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_remove_key_does_not_close_async_client():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -17,7 +17,7 @@ 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 import ReadableSpan, SpanLimits, 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
|
||||
|
|
@ -27,6 +27,7 @@ import litellm
|
|||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.integrations.otel import GenAIOperation
|
||||
from litellm.integrations.otel import logger as otel_logger
|
||||
from litellm.integrations.otel.emitter import SpanEmitter
|
||||
from litellm.integrations.otel.logger import (
|
||||
OpenTelemetryV2,
|
||||
build_otel_v2_logger,
|
||||
|
|
@ -34,6 +35,7 @@ from litellm.integrations.otel.logger import (
|
|||
publish_global_otel_v2_provider,
|
||||
)
|
||||
from litellm.integrations.otel.mappers import resolve_mappers
|
||||
from litellm.integrations.otel.mappers.langfuse import LANGFUSE_REQUEST_METADATA_ATTRIBUTES, LangfuseMapper
|
||||
from litellm.integrations.otel.model.config import (
|
||||
CaptureMessageContent,
|
||||
ExporterOwner,
|
||||
|
|
@ -52,6 +54,8 @@ from litellm.integrations.otel.model.payloads import (
|
|||
ServerInfo,
|
||||
ToolDefinition,
|
||||
)
|
||||
from litellm.integrations.otel.model.spans import SpanRole
|
||||
from litellm.integrations.otel.model.trace_controls import TraceControls
|
||||
from litellm.integrations.otel.plumbing import providers as otel_providers
|
||||
from litellm.integrations.otel.plumbing.context import (
|
||||
destination_backends,
|
||||
|
|
@ -3348,8 +3352,7 @@ class TestArizeTenantOnly:
|
|||
|
||||
assert capture.received() == {
|
||||
"space_id=space-operator,api_key=key-operator": ("chat gpt-4",),
|
||||
None: ("chat gpt-4",),
|
||||
}, "the operator's space, plus the stdout placeholder every credentialed preset keeps today"
|
||||
}, "the operator's space, with no stdout copy the operator never asked for"
|
||||
|
||||
def test_the_standard_otlp_headers_still_reach_the_operators_space(self, monkeypatch):
|
||||
credential_less_arize(monkeypatch)
|
||||
|
|
@ -3362,8 +3365,7 @@ class TestArizeTenantOnly:
|
|||
|
||||
assert capture.received() == {
|
||||
"space_id=space-operator,api_key=key-operator": ("chat gpt-4",),
|
||||
None: ("chat gpt-4",),
|
||||
}, "the operator's space, plus the stdout placeholder every credentialed preset keeps today"
|
||||
}, "the operator's space, with no stdout copy the operator never asked for"
|
||||
|
||||
def test_a_credential_less_arize_still_delivers_a_team_destination(self, monkeypatch):
|
||||
from litellm.integrations.otel.plumbing.providers import attach_tenant_fan_out
|
||||
|
|
@ -3435,6 +3437,271 @@ class TestArizeTenantOnly:
|
|||
assert loggers == [at_startup]
|
||||
|
||||
|
||||
class TestLangfuseDestinationFields:
|
||||
"""A key's or team's Langfuse destination gets the Langfuse fields from whichever logger the proxy published."""
|
||||
|
||||
def test_a_key_destination_gets_the_langfuse_fields_from_the_generic_startup_logger(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""With ``langfuse_otel`` only in a key's or team's logging, the startup logger is the
|
||||
generic one; the key's project still gets the user, session and tags, and the
|
||||
operator's stdout placeholder still gets no Langfuse request metadata."""
|
||||
credential_less_proxy(monkeypatch)
|
||||
capture: Final = _ExporterCapture()
|
||||
for kind in ("otlp_grpc", "otlp_http", "console"):
|
||||
monkeypatch.setitem(otel_providers._EXPORTER_FACTORIES, kind, capture.build) # pyright: ignore[reportPrivateUsage] # the factory table is the seam for capturing each exporter
|
||||
config: Final = OpenTelemetryV2Config()
|
||||
provider: Final = build_tracer_provider(config, tenant_overrides=True, console_default=True)
|
||||
attach_tenant_fan_out(provider, config)
|
||||
emitter: Final = SpanEmitter(get_tracer(provider, "litellm"), config)
|
||||
call: Final = replace(
|
||||
_MODEL_CALL_WITHOUT_CONTENT,
|
||||
trace=TraceControls(user_id="user-42", session_id="sess-1", tags=("agent:orders",)),
|
||||
request_metadata=MappingProxyType({"user_api_key_alias": "prod-key"}),
|
||||
)
|
||||
|
||||
def with_destination() -> None:
|
||||
set_request_destinations(deliverable_destinations((LANGFUSE_DEST,), provider))
|
||||
emitter.emit(SpanRole.LLM_CALL, call)
|
||||
|
||||
in_fresh_context(with_destination)
|
||||
provider.force_flush()
|
||||
|
||||
received: Final = {
|
||||
headers: exporter.get_finished_spans()[0].attributes or {} for headers, exporter in capture.built
|
||||
}
|
||||
to_langfuse: Final = received[LANGFUSE_DEST.header_string()]
|
||||
assert {key: to_langfuse.get(key) for key in ("user.id", "session.id", "langfuse.trace.tags")} == {
|
||||
"user.id": "user-42",
|
||||
"session.id": "sess-1",
|
||||
"langfuse.trace.tags": ("agent:orders",),
|
||||
}
|
||||
assert "langfuse.observation.metadata" in to_langfuse
|
||||
assert "langfuse.observation.metadata" not in received[None]
|
||||
|
||||
@staticmethod
|
||||
def _received_by_account(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
config: OpenTelemetryV2Config,
|
||||
destinations: tuple[OtelDestination, ...],
|
||||
with_content: bool = False,
|
||||
) -> Mapping[str | None, Mapping[str, object]]:
|
||||
"""The attributes of the model call each account received, keyed by the headers it was sent with."""
|
||||
credential_less_proxy(monkeypatch)
|
||||
capture: Final = _ExporterCapture()
|
||||
for kind in ("otlp_grpc", "otlp_http", "console"):
|
||||
monkeypatch.setitem(otel_providers._EXPORTER_FACTORIES, kind, capture.build) # pyright: ignore[reportPrivateUsage] # the factory table is the seam for capturing each exporter
|
||||
provider: Final = build_tracer_provider(config, tenant_overrides=True, console_default=True)
|
||||
attach_tenant_fan_out(provider, config)
|
||||
emitter: Final = SpanEmitter(get_tracer(provider, "litellm"), config)
|
||||
call: Final = replace(
|
||||
_MODEL_CALL_WITH_CONTENT if with_content else _MODEL_CALL_WITHOUT_CONTENT,
|
||||
trace=TraceControls(user_id="user-42", session_id="sess-1", tags=("agent:orders",)),
|
||||
request_metadata=MappingProxyType({"user_api_key_alias": "prod-key"}),
|
||||
)
|
||||
|
||||
def with_destinations() -> None:
|
||||
set_request_destinations(deliverable_destinations(destinations, provider))
|
||||
emitter.emit(SpanRole.LLM_CALL, call)
|
||||
|
||||
in_fresh_context(with_destinations)
|
||||
provider.force_flush()
|
||||
return MappingProxyType(
|
||||
{headers: dict(exporter.get_finished_spans()[0].attributes or {}) for headers, exporter in capture.built}
|
||||
)
|
||||
|
||||
def test_only_the_langfuse_destination_gets_the_langfuse_fields_written_for_it(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""The operator's own exporter and the key's Arize destination get the model call as they would with no
|
||||
Langfuse destination on the request."""
|
||||
received: Final = self._received_by_account(monkeypatch, OpenTelemetryV2Config(), (LANGFUSE_DEST, ARIZE_DEST))
|
||||
|
||||
assert received[LANGFUSE_DEST.header_string()]["user.id"] == "user-42"
|
||||
assert {
|
||||
headers: sorted(LangfuseMapper.LLM_CALL_ATTRIBUTES & attributes.keys())
|
||||
for headers, attributes in received.items()
|
||||
if headers != LANGFUSE_DEST.header_string()
|
||||
} == {None: [], ARIZE_DEST.header_string(): []}
|
||||
|
||||
def test_an_exporter_whose_config_maps_langfuse_keeps_the_langfuse_fields_but_the_request_metadata(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
received: Final = self._received_by_account(
|
||||
monkeypatch, OpenTelemetryV2Config(mapper_names=["genai", "langfuse"]), (LANGFUSE_DEST,)
|
||||
)
|
||||
|
||||
assert {key: received[None].get(key) for key in ("user.id", "session.id", "langfuse.trace.tags")} == {
|
||||
"user.id": "user-42",
|
||||
"session.id": "sess-1",
|
||||
"langfuse.trace.tags": ("agent:orders",),
|
||||
}
|
||||
assert not LANGFUSE_REQUEST_METADATA_ATTRIBUTES & received[None].keys()
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("key_environment", "tracing_environment", "expected"),
|
||||
[
|
||||
("prod", None, {"langfuse.environment": "prod"}),
|
||||
("prod", "staging", {"langfuse.environment": "prod"}),
|
||||
(None, "staging", {"langfuse.environment": "staging"}),
|
||||
(None, None, {}),
|
||||
],
|
||||
ids=["key", "key-over-env", "env", "neither"],
|
||||
)
|
||||
def test_a_langfuse_destination_names_the_environment_the_legacy_logger_wrote(
|
||||
self,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
key_environment: str | None,
|
||||
tracing_environment: str | None,
|
||||
expected: Mapping[str, str],
|
||||
) -> None:
|
||||
credential_less_proxy(monkeypatch)
|
||||
if tracing_environment is None:
|
||||
monkeypatch.delenv("LANGFUSE_TRACING_ENVIRONMENT", raising=False)
|
||||
else:
|
||||
monkeypatch.setenv("LANGFUSE_TRACING_ENVIRONMENT", tracing_environment)
|
||||
|
||||
destination: Final = destination_for(
|
||||
"langfuse_otel",
|
||||
StandardCallbackDynamicParams(
|
||||
langfuse_public_key="pk", langfuse_secret_key="sk", langfuse_environment=key_environment
|
||||
),
|
||||
)
|
||||
|
||||
assert destination is not None
|
||||
assert dict(destination.resource_attributes) == expected
|
||||
|
||||
def test_a_langfuse_destination_on_a_proxy_with_langfuse_keys_names_no_environment(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""With operator keys the older Langfuse logger never runs, so there is no environment of its to keep."""
|
||||
monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "pk-operator")
|
||||
monkeypatch.setenv("LANGFUSE_SECRET_KEY", "sk-operator")
|
||||
|
||||
destination: Final = destination_for(
|
||||
"langfuse_otel",
|
||||
StandardCallbackDynamicParams(
|
||||
langfuse_public_key="pk", langfuse_secret_key="sk", langfuse_environment="prod"
|
||||
),
|
||||
)
|
||||
|
||||
assert destination is not None
|
||||
assert dict(destination.resource_attributes) == {}
|
||||
|
||||
@pytest.mark.parametrize("operator_keys", [False, True], ids=["no-operator-keys", "operator-keys"])
|
||||
@pytest.mark.parametrize("configured", [None, "no_content", "span_only"])
|
||||
def test_a_langfuse_destination_keeps_its_own_capture_setting_whatever_the_operator_keys(
|
||||
self,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
configured: str | None,
|
||||
operator_keys: bool,
|
||||
) -> None:
|
||||
"""An unset setting stays unset, so the destination follows the global capture mode either way."""
|
||||
credential_less_proxy(monkeypatch)
|
||||
if operator_keys:
|
||||
monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "pk-operator")
|
||||
monkeypatch.setenv("LANGFUSE_SECRET_KEY", "sk-operator")
|
||||
|
||||
destination: Final = destination_for(
|
||||
"langfuse_otel",
|
||||
StandardCallbackDynamicParams(langfuse_public_key="pk", langfuse_secret_key="sk"),
|
||||
capture_message_content=configured,
|
||||
)
|
||||
|
||||
assert destination is not None
|
||||
assert destination.capture_message_content == configured
|
||||
|
||||
@pytest.mark.parametrize(("configured", "captured"), [(None, False), ("span_only", True)])
|
||||
def test_a_langfuse_destination_gets_the_prompt_only_when_its_setting_asks_without_operator_keys(
|
||||
self, monkeypatch: pytest.MonkeyPatch, configured: str | None, captured: bool
|
||||
) -> None:
|
||||
credential_less_proxy(monkeypatch)
|
||||
destination: Final = destination_for(
|
||||
"langfuse_otel",
|
||||
StandardCallbackDynamicParams(langfuse_public_key="pk", langfuse_secret_key="sk"),
|
||||
capture_message_content=configured,
|
||||
)
|
||||
arize: Final = destination_for("arize", StandardCallbackDynamicParams(arize_space_id="s", arize_api_key="k"))
|
||||
assert destination is not None
|
||||
assert arize is not None
|
||||
|
||||
received: Final = self._received_by_account(
|
||||
monkeypatch, OpenTelemetryV2Config(), (destination, arize), with_content=True
|
||||
)
|
||||
|
||||
assert {headers: "langfuse.observation.input" in attributes for headers, attributes in received.items()} == {
|
||||
destination.header_string(): captured,
|
||||
arize.header_string(): False,
|
||||
None: False,
|
||||
}
|
||||
assert not [key for key in received[None] if key.startswith("gen_ai.input")]
|
||||
|
||||
def test_only_the_langfuse_destination_gets_the_keys_environment(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
credential_less_proxy(monkeypatch)
|
||||
capture: Final = _ExporterCapture()
|
||||
for kind in ("otlp_grpc", "otlp_http", "console"):
|
||||
monkeypatch.setitem(otel_providers._EXPORTER_FACTORIES, kind, capture.build) # pyright: ignore[reportPrivateUsage] # the factory table is the seam for capturing each exporter
|
||||
config: Final = OpenTelemetryV2Config()
|
||||
provider: Final = build_tracer_provider(config, tenant_overrides=True, console_default=True)
|
||||
attach_tenant_fan_out(provider, config)
|
||||
emitter: Final = SpanEmitter(get_tracer(provider, "litellm"), config)
|
||||
destination: Final = destination_for(
|
||||
"langfuse_otel", {"langfuse_public_key": "pk", "langfuse_secret_key": "sk", "langfuse_environment": "prod"}
|
||||
)
|
||||
assert destination is not None
|
||||
|
||||
def with_destination() -> None:
|
||||
set_request_destinations(deliverable_destinations((destination,), provider))
|
||||
emitter.emit(SpanRole.LLM_CALL, _MODEL_CALL_WITHOUT_CONTENT)
|
||||
|
||||
in_fresh_context(with_destination)
|
||||
provider.force_flush()
|
||||
|
||||
environments: Final = {
|
||||
headers: exporter.get_finished_spans()[0].resource.attributes.get("langfuse.environment")
|
||||
for headers, exporter in capture.built
|
||||
}
|
||||
assert environments == {destination.header_string(): "prod", None: None}
|
||||
|
||||
@staticmethod
|
||||
def _attributes(
|
||||
destinations: tuple[OtelDestination, ...], limit: int | None, call: LLMCallSpanData
|
||||
) -> dict[str, object]:
|
||||
exporter: Final = InMemorySpanExporter()
|
||||
provider: Final = TracerProvider(span_limits=SpanLimits(max_span_attributes=limit))
|
||||
provider.add_span_processor(SimpleSpanProcessor(exporter))
|
||||
emitter: Final = SpanEmitter(get_tracer(provider, "litellm"), OpenTelemetryV2Config())
|
||||
|
||||
def run() -> None:
|
||||
set_request_destinations(destinations)
|
||||
emitter.emit(SpanRole.LLM_CALL, call)
|
||||
|
||||
in_fresh_context(run)
|
||||
return dict(exporter.get_finished_spans()[0].attributes or {})
|
||||
|
||||
def test_a_span_at_its_attribute_limit_sheds_the_destinations_langfuse_fields_first(self) -> None:
|
||||
generic: Final = self._attributes((), None, _MODEL_CALL_WITH_CONTENT)
|
||||
|
||||
assert "user.id" not in generic
|
||||
assert self._attributes((LANGFUSE_DEST,), len(generic), _MODEL_CALL_WITH_CONTENT) == generic, (
|
||||
"the generic attributes all survive"
|
||||
)
|
||||
|
||||
def test_a_span_at_its_attribute_limit_sheds_the_destinations_request_metadata_before_its_other_fields(
|
||||
self,
|
||||
) -> None:
|
||||
call: Final = replace(
|
||||
_MODEL_CALL_WITH_CONTENT, request_metadata=MappingProxyType({"user_api_key_alias": "prod-key"})
|
||||
)
|
||||
whole: Final = self._attributes((LANGFUSE_DEST,), None, call)
|
||||
|
||||
squeezed: Final = self._attributes((LANGFUSE_DEST,), len(whole) - 1, call)
|
||||
|
||||
assert len(squeezed) == len(whole) - 1
|
||||
assert whole.keys() - squeezed.keys() <= LANGFUSE_REQUEST_METADATA_ATTRIBUTES
|
||||
assert squeezed["langfuse.observation.type"] == "generation"
|
||||
|
||||
|
||||
class TestContextIsolation:
|
||||
def test_destinations_do_not_leak_between_requests(self):
|
||||
def first():
|
||||
|
|
@ -4556,6 +4823,115 @@ class TestCredentialGatedExporters:
|
|||
assert kept[0] is operator_console
|
||||
|
||||
|
||||
class TestConsolePlaceholder:
|
||||
"""The console spec ``_normalize`` folds in for an empty exporter list prints spans
|
||||
only for the ``otel`` callback; no other logger was asked to print."""
|
||||
|
||||
@pytest.fixture
|
||||
def capture(self, monkeypatch: pytest.MonkeyPatch) -> "_ExporterCapture":
|
||||
return TestArizeTenantOnly._capture_exporters(monkeypatch)
|
||||
|
||||
@staticmethod
|
||||
def _received(capture: "_ExporterCapture", provider: TracerProvider) -> Mapping[str | None, tuple[str, ...]]:
|
||||
emit(provider)
|
||||
provider.force_flush()
|
||||
return capture.received()
|
||||
|
||||
@staticmethod
|
||||
def _no_otel_env(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
for name in _OTEL_SHORTHAND_ENV:
|
||||
monkeypatch.delenv(name, raising=False)
|
||||
|
||||
@staticmethod
|
||||
def _operator_langfuse(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
TestConsolePlaceholder._no_otel_env(monkeypatch)
|
||||
monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "pk-lf-operator")
|
||||
monkeypatch.setenv("LANGFUSE_SECRET_KEY", "sk-lf-operator")
|
||||
|
||||
@staticmethod
|
||||
def _isolated_proxy_registration(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
for name in ("service_callback", "input_callback", "_async_success_callback", "_async_failure_callback"):
|
||||
monkeypatch.setattr(litellm, name, [])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.open_telemetry_logger", None)
|
||||
|
||||
def test_a_credentialed_langfuse_exports_to_langfuse_alone(
|
||||
self, monkeypatch: pytest.MonkeyPatch, capture: "_ExporterCapture"
|
||||
) -> None:
|
||||
self._operator_langfuse(monkeypatch)
|
||||
config: Final = langfuse_preset()
|
||||
langfuse_headers: Final = next(
|
||||
spec.headers for spec in config.exporters if spec.owner is ExporterOwner.LANGFUSE_OTEL
|
||||
)
|
||||
|
||||
assert self._received(capture, build_tracer_provider(config)) == {langfuse_headers: ("chat gpt-4",)}
|
||||
|
||||
def test_a_console_exporter_the_operator_named_still_prints(
|
||||
self, monkeypatch: pytest.MonkeyPatch, capture: "_ExporterCapture"
|
||||
) -> None:
|
||||
self._operator_langfuse(monkeypatch)
|
||||
monkeypatch.setenv("OTEL_EXPORTER", "console")
|
||||
|
||||
assert self._received(capture, build_tracer_provider(langfuse_preset()))[None] == ("chat gpt-4",)
|
||||
|
||||
@pytest.fixture
|
||||
def otel_callback(self, monkeypatch: pytest.MonkeyPatch, capture: "_ExporterCapture") -> OpenTelemetryV2:
|
||||
"""The plain ``otel`` callback with nothing else configured, built the way the proxy builds it."""
|
||||
from litellm.litellm_core_utils import litellm_logging as logging_module
|
||||
|
||||
self._no_otel_env(monkeypatch)
|
||||
self._isolated_proxy_registration(monkeypatch)
|
||||
monkeypatch.setenv("LITELLM_OTEL_V2", "true")
|
||||
monkeypatch.setattr(logging_module, "_in_memory_loggers", [])
|
||||
is_otel_v2_enabled.cache_clear()
|
||||
built: Final = logging_module.build_custom_logger_compatible_class(
|
||||
"otel", internal_usage_cache=None, llm_router=None
|
||||
)
|
||||
assert isinstance(built, OpenTelemetryV2)
|
||||
return built
|
||||
|
||||
def test_the_otel_callback_prints_when_nothing_else_is_configured(
|
||||
self, otel_callback: OpenTelemetryV2, capture: "_ExporterCapture"
|
||||
) -> None:
|
||||
assert self._received(capture, otel_callback.tracer_provider) == {None: ("chat gpt-4",)}
|
||||
|
||||
def test_the_otel_callback_prints_a_call_for_a_key_with_its_own_service_name(
|
||||
self, otel_callback: OpenTelemetryV2, capture: "_ExporterCapture"
|
||||
) -> None:
|
||||
"""A key or team ``otel_service_name`` moves the call onto a second provider built for that name."""
|
||||
kwargs: Final = {
|
||||
**_closed_chat_call_kwargs(),
|
||||
"litellm_params": {"metadata": {"user_api_key_auth_metadata": {"otel_service_name": "payments-gateway"}}},
|
||||
}
|
||||
|
||||
otel_callback.log_pre_api_call(model="gpt-4o", messages=[], kwargs=kwargs)
|
||||
asyncio.run(otel_callback.async_log_success_event(kwargs, None, None, None))
|
||||
|
||||
assert capture.received() == {None: ("chat gpt-4o",)}
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("env", "expected"),
|
||||
[({}, {}), ({"OTEL_EXPORTER": "console"}, {None: ("chat gpt-4",)})],
|
||||
ids=["nothing-configured", "operator-asked-for-console"],
|
||||
)
|
||||
def test_the_background_global_logger_prints_only_when_asked(
|
||||
self,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
capture: "_ExporterCapture",
|
||||
env: Mapping[str, str],
|
||||
expected: Mapping[str | None, tuple[str, ...]],
|
||||
) -> None:
|
||||
"""The proxy publishes a generic logger even when no callback built one, which is
|
||||
what key-only and team-only destinations run under."""
|
||||
from litellm.integrations.otel.logger import select_global_otel_v2_logger
|
||||
|
||||
self._no_otel_env(monkeypatch)
|
||||
self._isolated_proxy_registration(monkeypatch)
|
||||
for name, value in env.items():
|
||||
monkeypatch.setenv(name, value)
|
||||
|
||||
assert self._received(capture, select_global_otel_v2_logger([]).tracer_provider) == expected
|
||||
|
||||
|
||||
class TestTenantHostSsrfGuard:
|
||||
"""Anyone who can mint a key can write ``langfuse_host``, so the host it names has
|
||||
to be one the operator approved."""
|
||||
|
|
|
|||
|
|
@ -1,10 +1,14 @@
|
|||
import os
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
|
||||
import litellm
|
||||
from litellm.caching.llm_caching_handler import LLMClientCache
|
||||
from litellm.integrations.agentops.agentops import AgentOps, AgentOpsConfig
|
||||
from litellm.llms.custom_httpx.http_handler import get_httpx_client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -111,3 +115,17 @@ def test_agentops_initialization_no_auth():
|
|||
assert agentops.resource_attributes["deployment.environment"] == "test_env"
|
||||
assert agentops.resource_attributes["telemetry.sdk.name"] == "agentops"
|
||||
assert "project.id" not in agentops.resource_attributes
|
||||
|
||||
|
||||
def test_auth_fetch_leaves_the_shared_sync_client_open_for_other_callers(
|
||||
monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache())
|
||||
auth_endpoint: Final = "https://agentops.example/v3/auth/token"
|
||||
respx_mock.post(auth_endpoint).respond(json={"token": "jwt", "project_id": "project-1"})
|
||||
client_another_caller_is_using: Final = get_httpx_client().client
|
||||
|
||||
agentops: Final = AgentOps(config=AgentOpsConfig(api_key="key", auth_endpoint=auth_endpoint))
|
||||
|
||||
assert agentops.resource_attributes["project.id"] == "project-1"
|
||||
assert not client_another_caller_is_using.is_closed
|
||||
|
|
|
|||
|
|
@ -1,9 +1,12 @@
|
|||
import asyncio
|
||||
import base64
|
||||
import contextvars
|
||||
import json
|
||||
import os
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
|
|
@ -12,12 +15,17 @@ from openai.types.responses import ResponseFunctionToolCall
|
|||
from opentelemetry import trace
|
||||
from opentelemetry.context import Context
|
||||
from opentelemetry.sdk.trace import SpanLimits, TracerProvider
|
||||
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
|
||||
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
|
||||
from opentelemetry.trace import Span
|
||||
|
||||
import litellm
|
||||
from litellm.constants import REDACTED_BY_LITELLM
|
||||
from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL, REDACTED_BY_LITELLM
|
||||
from litellm.integrations.langfuse.langfuse_otel import LangfuseOtelLogger
|
||||
from litellm.integrations.opentelemetry import LITELLM_PROXY_REQUEST_SPAN_NAME, OpenTelemetry, OpenTelemetryConfig
|
||||
from litellm.integrations.otel.model.destination import OtelDestination
|
||||
from litellm.integrations.otel.plumbing.context import set_request_destinations
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
from litellm.types.utils import Choices, ModelResponse, StandardCallbackDynamicParams
|
||||
|
||||
|
|
@ -1507,3 +1515,75 @@ def test_langfuse_otel_does_not_start_proxy_request_span() -> None:
|
|||
finally:
|
||||
langfuse_provider.shutdown()
|
||||
generic_provider.shutdown()
|
||||
|
||||
|
||||
_HOOK_CALL: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
|
||||
def _spans_the_legacy_logger_exports(
|
||||
destinations: tuple[OtelDestination, ...], call: Mapping[str, object] = _HOOK_CALL
|
||||
) -> tuple[str, ...]:
|
||||
exporter: Final = InMemorySpanExporter()
|
||||
provider: Final = TracerProvider()
|
||||
provider.add_span_processor(SimpleSpanProcessor(exporter))
|
||||
with patch.dict(os.environ, LANGFUSE_ENV_ONLY, clear=True):
|
||||
logger: Final = LangfuseOtelLogger(tracer_provider=provider)
|
||||
at: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||||
kwargs: Final[dict[str, object]] = {**CHAT_KWARGS, **call, "litellm_params": {"metadata": {}}}
|
||||
server_span: Final = provider.get_tracer("proxy").start_span("POST /v1/chat/completions")
|
||||
|
||||
async def log() -> None:
|
||||
set_request_destinations(destinations)
|
||||
logger.log_success_event(dict(kwargs), _chat_response("ok"), at, at) # pyright: ignore[reportUnknownMemberType] # the v1 logger hooks are untyped
|
||||
await logger.async_log_failure_event({**kwargs, "exception": ValueError("boom")}, None, at, at) # pyright: ignore[reportUnknownMemberType] # the v1 logger hooks are untyped
|
||||
await logger.async_post_call_failure_hook({}, ValueError("boom"), UserAPIKeyAuth(parent_otel_span=server_span))
|
||||
|
||||
try:
|
||||
contextvars.copy_context().run(asyncio.run, log())
|
||||
return tuple(span.name for span in exporter.get_finished_spans())
|
||||
finally:
|
||||
provider.shutdown()
|
||||
|
||||
|
||||
_LANGFUSE_DESTINATION: Final = OtelDestination(
|
||||
endpoint="http://tenant.local/api/public/otel",
|
||||
headers={"Authorization": "Basic dGVuYW50"},
|
||||
callback_name="langfuse_otel",
|
||||
)
|
||||
|
||||
|
||||
_PROVIDER_CALL: Final[Mapping[str, object]] = MappingProxyType(
|
||||
{"standard_logging_object": {"id": "c1"}, "api_call_start_time": datetime(2026, 1, 1, tzinfo=timezone.utc)}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("call", [_HOOK_CALL, _PROVIDER_CALL], ids=["hook", "provider-call"])
|
||||
def test_a_request_with_a_langfuse_destination_exports_nothing_through_the_legacy_logger(
|
||||
call: Mapping[str, object],
|
||||
) -> None:
|
||||
assert _spans_the_legacy_logger_exports((_LANGFUSE_DESTINATION,), call) == ()
|
||||
|
||||
|
||||
def test_a_call_rejected_before_any_provider_still_exports_through_the_legacy_logger() -> None:
|
||||
rejected: Final = {"standard_logging_object": {"id": "c1"}, LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL: True}
|
||||
|
||||
assert _spans_the_legacy_logger_exports((_LANGFUSE_DESTINATION,), rejected) == (
|
||||
"litellm_request",
|
||||
"raw_gen_ai_request",
|
||||
"litellm_request",
|
||||
)
|
||||
|
||||
|
||||
def test_a_request_without_a_langfuse_destination_still_exports_through_the_legacy_logger() -> None:
|
||||
logged: Final = (
|
||||
"litellm_request",
|
||||
"raw_gen_ai_request",
|
||||
"litellm_request",
|
||||
"Failed Proxy Server Request",
|
||||
"POST /v1/chat/completions",
|
||||
)
|
||||
|
||||
assert _spans_the_legacy_logger_exports(()) == logged
|
||||
assert _spans_the_legacy_logger_exports((_LANGFUSE_DESTINATION.model_copy(update={"callback_name": "arize"}),)) == (
|
||||
logged
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1388,6 +1388,55 @@ async def test_finalizer_with_running_loop_schedules_close_and_holds_task_ref():
|
|||
assert not (AsyncHTTPHandler._finalizer_close_tasks & scheduled)
|
||||
|
||||
|
||||
async def _drop_unused_handler() -> None:
|
||||
AsyncHTTPHandler(timeout=61.0)
|
||||
gc.collect()
|
||||
|
||||
|
||||
def test_finalizer_close_does_not_keep_a_loop_closed_right_after_it_alive():
|
||||
loop: Final = asyncio.new_event_loop()
|
||||
loop.run_until_complete(_drop_unused_handler())
|
||||
loop.close()
|
||||
closed_loop: Final = weakref.ref(loop)
|
||||
del loop
|
||||
gc.collect()
|
||||
|
||||
assert closed_loop() is None
|
||||
|
||||
|
||||
class _TransportFailingToClose(httpx.AsyncBaseTransport):
|
||||
async def aclose(self) -> None:
|
||||
raise RuntimeError("close failed")
|
||||
|
||||
|
||||
async def _drop_handler_whose_close_fails() -> None:
|
||||
AsyncHTTPHandler(timeout=61.0, transport=_TransportFailingToClose())
|
||||
gc.collect()
|
||||
|
||||
|
||||
def test_failed_finalizer_close_in_a_loops_last_step_reports_no_unretrieved_exception():
|
||||
loop: Final = asyncio.new_event_loop()
|
||||
exception_handler: Final = MagicMock()
|
||||
loop.set_exception_handler(exception_handler)
|
||||
loop.run_until_complete(_drop_handler_whose_close_fails())
|
||||
loop.close()
|
||||
del loop
|
||||
gc.collect()
|
||||
|
||||
exception_handler.assert_not_called()
|
||||
|
||||
|
||||
async def _sync_client_with_timeout_62() -> HTTPHandler:
|
||||
return get_httpx_client(params={"timeout": 62.0})
|
||||
|
||||
|
||||
def test_sync_client_is_shared_by_short_lived_loops_and_callers_outside_a_loop():
|
||||
built_in_first_loop: Final = asyncio.run(_sync_client_with_timeout_62())
|
||||
|
||||
assert asyncio.run(_sync_client_with_timeout_62()) is built_in_first_loop
|
||||
assert get_httpx_client(params={"timeout": 62.0}) is built_in_first_loop
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sync_close_helper_respects_session_ownership():
|
||||
"""The loop-less fallback closes only sessions the transport owns; a
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import asyncio
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
|
|
@ -154,6 +155,28 @@ def test_get_openai_client_initialization_param_fields(client_type):
|
|||
assert "self" not in result
|
||||
|
||||
|
||||
def _loop_test_client_params(is_async: bool) -> dict[str, object]:
|
||||
return {"api_key": "sk-loop-scope-test", "is_async": is_async}
|
||||
|
||||
|
||||
async def _cache_openai_client(client: openai.OpenAI | openai.AsyncOpenAI, is_async: bool) -> None:
|
||||
BaseOpenAILLM.set_cached_openai_client(client, "openai", _loop_test_client_params(is_async))
|
||||
|
||||
|
||||
async def _cached_openai_client(is_async: bool) -> object:
|
||||
return BaseOpenAILLM.get_cached_openai_client(_loop_test_client_params(is_async), "openai")
|
||||
|
||||
|
||||
def test_sync_openai_client_cached_in_one_loop_is_reused_by_the_next_while_async_is_not():
|
||||
sync_client: Final = openai.OpenAI(api_key="sk-loop-scope-test")
|
||||
async_client: Final = openai.AsyncOpenAI(api_key="sk-loop-scope-test")
|
||||
asyncio.run(_cache_openai_client(sync_client, is_async=False))
|
||||
asyncio.run(_cache_openai_client(async_client, is_async=True))
|
||||
|
||||
assert asyncio.run(_cached_openai_client(is_async=False)) is sync_client
|
||||
assert asyncio.run(_cached_openai_client(is_async=True)) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("client_type", ["openai", "azure"])
|
||||
def test_get_openai_client_cache_key(client_type):
|
||||
"""Verify get_openai_client_cache_key doesn't raise on tuple + tuple concatenation."""
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import gc
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
|
|
@ -2966,6 +2967,188 @@ class TestStreamingScanDedup:
|
|||
assert [scan["texts"] for scan in guardrail.scans] == [["abc"]]
|
||||
|
||||
|
||||
class _BackgroundScanGuardrail(_MarkerBlockingScanGuardrail):
|
||||
"""Background-scanned guardrail whose scans hold until released, then block BLOCKME and rewrite other text"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__(sampling_rate=2)
|
||||
self.streaming_scan_in_background = True
|
||||
self.released = anyio.Event()
|
||||
self.scanned = anyio.Event()
|
||||
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict[str, object],
|
||||
input_type: Literal["request", "response"],
|
||||
**kwargs: object,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
try:
|
||||
await self.released.wait()
|
||||
recorded = await super().apply_guardrail(inputs, request_data, input_type, **kwargs)
|
||||
finally:
|
||||
self.scanned.set()
|
||||
return {**recorded, "texts": ["REWRITTEN" for _ in recorded.get("texts") or []]}
|
||||
|
||||
|
||||
class TestStreamingBackgroundScan:
|
||||
"""streaming_scan_in_background: the stream keeps flowing while a sampled scan runs, and a
|
||||
failed scan ends it before the next chunk. LIT-9338."""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _use_real_mappings(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_patch_translation_mappings(monkeypatch, load_guardrail_translation_mappings())
|
||||
|
||||
@staticmethod
|
||||
def _hook(guardrail: CustomGuardrail, upstream: AsyncIterator[ModelResponseStream]) -> AsyncGenerator[object, None]:
|
||||
return UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/chat/completions"),
|
||||
response=upstream,
|
||||
request_data={"guardrail_to_apply": guardrail, "model": "gpt-4", "metadata": {}},
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _guarded_stream(cls, guardrail: CustomGuardrail, *texts: str) -> AsyncGenerator[object, None]:
|
||||
async def upstream() -> AsyncIterator[ModelResponseStream]:
|
||||
for text in texts:
|
||||
yield _stream_chunk(text)
|
||||
yield _stream_chunk("", finish_reason="stop")
|
||||
|
||||
return cls._hook(guardrail, upstream())
|
||||
|
||||
@staticmethod
|
||||
async def _receive(stream: AsyncGenerator[object, None], count: int) -> list[object]:
|
||||
with anyio.fail_after(5):
|
||||
return [await stream.__anext__() for _ in range(count)]
|
||||
|
||||
@staticmethod
|
||||
def _block_frames(items: list[object]) -> bytes:
|
||||
assert not any(isinstance(item, ModelResponseStream) for item in items), items
|
||||
return b"".join(item for item in items if isinstance(item, bytes))
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chunks_stream_while_the_scan_runs_and_its_rewrite_never_reaches_them(self) -> None:
|
||||
guardrail: Final = _BackgroundScanGuardrail()
|
||||
stream: Final = self._guarded_stream(guardrail, "a", "b", "c", "d")
|
||||
|
||||
sent: Final = await self._receive(stream, 4)
|
||||
assert guardrail.scans == ()
|
||||
guardrail.released.set()
|
||||
rest: Final = [item async for item in stream]
|
||||
|
||||
assert [_delta_text(item) for item in (*sent, *rest)] == ["a", "b", "c", "d", ""]
|
||||
assert [scan["texts"] for scan in guardrail.scans] == [["ab"], ["abcd"]], guardrail.scans
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_block_ends_the_stream_before_the_next_chunk(self) -> None:
|
||||
guardrail: Final = _BackgroundScanGuardrail()
|
||||
stream: Final = self._guarded_stream(guardrail, "a", "BLOCKME", "c", "d", "e")
|
||||
|
||||
sent: Final = await self._receive(stream, 3)
|
||||
guardrail.released.set()
|
||||
await guardrail.scanned.wait()
|
||||
rest: Final = [item async for item in stream]
|
||||
|
||||
assert [_delta_text(item) for item in sent] == ["a", "BLOCKME", "c"]
|
||||
assert b"blocked" in self._block_frames(rest)
|
||||
assert [scan["texts"] for scan in guardrail.scans] == [["aBLOCKME"]], guardrail.scans
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_block_that_lands_after_the_last_text_replaces_the_finish(self) -> None:
|
||||
guardrail: Final = _BackgroundScanGuardrail()
|
||||
stream: Final = self._guarded_stream(guardrail, "a", "BLOCKME", "c", "d", "e", "f")
|
||||
|
||||
sent: Final = await self._receive(stream, 6)
|
||||
guardrail.released.set()
|
||||
rest: Final = [item async for item in stream]
|
||||
|
||||
assert [_delta_text(item) for item in sent] == ["a", "BLOCKME", "c", "d", "e", "f"]
|
||||
assert b"blocked" in self._block_frames(rest)
|
||||
assert [scan["texts"] for scan in guardrail.scans] == [["aBLOCKME"]], guardrail.scans
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_block_from_a_scan_that_covered_all_the_text_still_ends_the_stream(self) -> None:
|
||||
guardrail: Final = _BackgroundScanGuardrail()
|
||||
guardrail.released.set()
|
||||
|
||||
out: Final = [item async for item in self._guarded_stream(guardrail, "a", "BLOCKME")]
|
||||
|
||||
assert [_delta_text(item) for item in out if isinstance(item, ModelResponseStream)] == ["a", "BLOCKME"]
|
||||
assert b"blocked" in b"".join(item for item in out if isinstance(item, bytes))
|
||||
assert [scan["texts"] for scan in guardrail.scans] == [["aBLOCKME"]], guardrail.scans
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_upstream_error_cancels_the_running_scan(self) -> None:
|
||||
guardrail: Final = _BackgroundScanGuardrail()
|
||||
|
||||
async def upstream() -> AsyncIterator[ModelResponseStream]:
|
||||
yield _stream_chunk("a")
|
||||
yield _stream_chunk("b")
|
||||
await anyio.wait_all_tasks_blocked()
|
||||
raise RuntimeError("upstream dropped")
|
||||
|
||||
with anyio.fail_after(5), pytest.raises(RuntimeError, match="upstream dropped"):
|
||||
_ = [item async for item in self._hook(guardrail, upstream())]
|
||||
assert guardrail.scanned.is_set(), "the scan was still running when the upstream error surfaced"
|
||||
guardrail.released.set()
|
||||
await anyio.wait_all_tasks_blocked()
|
||||
|
||||
assert guardrail.scans == ()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_upstream_error_after_a_failed_scan_reports_only_the_upstream_error(
|
||||
self, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
guardrail: Final = _BackgroundScanGuardrail()
|
||||
guardrail.released.set()
|
||||
|
||||
async def upstream() -> AsyncIterator[ModelResponseStream]:
|
||||
yield _stream_chunk("a")
|
||||
yield _stream_chunk("BLOCKME")
|
||||
await guardrail.scanned.wait()
|
||||
await anyio.wait_all_tasks_blocked()
|
||||
raise RuntimeError("upstream dropped")
|
||||
|
||||
with pytest.raises(RuntimeError, match="upstream dropped"):
|
||||
_ = [item async for item in self._hook(guardrail, upstream())]
|
||||
gc.collect()
|
||||
|
||||
assert "never retrieved" not in caplog.text
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_closing_mid_scan_waits_for_it_then_scans_what_was_sent(self) -> None:
|
||||
guardrail: Final = _BackgroundScanGuardrail()
|
||||
stream: Final = self._guarded_stream(guardrail, "a", "b", "c", "d")
|
||||
|
||||
await self._receive(stream, 3)
|
||||
guardrail.released.set()
|
||||
await stream.aclose()
|
||||
|
||||
assert [scan["texts"] for scan in guardrail.scans] == [["ab"], ["abc"]], guardrail.scans
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_closing_after_a_background_block_does_not_scan_again(self) -> None:
|
||||
guardrail: Final = _BackgroundScanGuardrail()
|
||||
stream: Final = self._guarded_stream(guardrail, "a", "BLOCKME", "c", "d")
|
||||
|
||||
await self._receive(stream, 3)
|
||||
guardrail.released.set()
|
||||
await stream.aclose()
|
||||
|
||||
assert [scan["texts"] for scan in guardrail.scans] == [["aBLOCKME"]], guardrail.scans
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_buffered_release_on_scan_ignores_the_flag(self) -> None:
|
||||
guardrail: Final = _BackgroundScanGuardrail()
|
||||
guardrail.streaming_buffer_until_moderated = True
|
||||
guardrail.streaming_buffer_release_on_scan = True
|
||||
guardrail.released.set()
|
||||
|
||||
out: Final = [item async for item in self._guarded_stream(guardrail, "BLOCKME", "b", "c", "d")]
|
||||
|
||||
assert b"blocked" in self._block_frames(out)
|
||||
|
||||
|
||||
class TestTranslationMappingsAreReadLive:
|
||||
"""The hooks must read the handler map on every call, never memoize it on the module.
|
||||
|
||||
|
|
|
|||
|
|
@ -95,8 +95,8 @@ async def test_claude_count_request_counts_additional_tools_input_items_as_tools
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_claude_count_request_counts_a_tool_list_repeated_on_every_turn_once() -> None:
|
||||
policy: Final = DropPolicy(drop_params=False, additional_drop_params=frozenset())
|
||||
async def test_claude_count_request_counts_a_tool_list_repeated_on_every_turn_once_under_drop_params() -> None:
|
||||
policy: Final = DropPolicy(drop_params=True, additional_drop_params=frozenset())
|
||||
tool_list: Final[dict[str, JsonValue]] = {"type": "additional_tools", "tools": [_WEATHER_TOOL]}
|
||||
|
||||
repeated: Final = await _claude_count(
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ _USER_MESSAGE = {"type": "message", "role": "user", "content": "Run ls"}
|
|||
|
||||
|
||||
def test_string_input_passes_through_with_existing_tools():
|
||||
hoisted = hoist_additional_tools("hello", [_TOP_LEVEL_TOOL])
|
||||
hoisted = hoist_additional_tools("hello", [_TOP_LEVEL_TOOL], drop_params=False)
|
||||
|
||||
assert hoisted.input == "hello"
|
||||
assert hoisted.tools == (_TOP_LEVEL_TOOL,)
|
||||
|
|
@ -25,7 +25,7 @@ def test_string_input_passes_through_with_existing_tools():
|
|||
def test_input_without_additional_tools_items_is_returned_untouched():
|
||||
request_input = [_USER_MESSAGE]
|
||||
|
||||
hoisted = hoist_additional_tools(request_input, None)
|
||||
hoisted = hoist_additional_tools(request_input, None, drop_params=False)
|
||||
|
||||
assert hoisted.input is request_input
|
||||
assert hoisted.tools == ()
|
||||
|
|
@ -39,7 +39,7 @@ def test_additional_tools_items_are_stripped_and_appended_after_top_level_tools_
|
|||
{"type": "additional_tools", "id": "at_2", "role": "developer", "tools": [_WAIT_TOOL]},
|
||||
]
|
||||
|
||||
hoisted = hoist_additional_tools(request_input, [_TOP_LEVEL_TOOL])
|
||||
hoisted = hoist_additional_tools(request_input, [_TOP_LEVEL_TOOL], drop_params=False)
|
||||
|
||||
assert hoisted.input == [_USER_MESSAGE]
|
||||
assert hoisted.tools == (_TOP_LEVEL_TOOL, _EXEC_TOOL, _WAIT_TOOL)
|
||||
|
|
@ -49,7 +49,7 @@ def test_additional_tools_items_are_stripped_and_appended_after_top_level_tools_
|
|||
def test_additional_tools_item_without_a_tools_list_is_stripped_and_contributes_nothing():
|
||||
request_input = [{"type": "additional_tools", "id": "at_1", "role": "developer", "tools": "exec"}, _USER_MESSAGE]
|
||||
|
||||
hoisted = hoist_additional_tools(request_input, None)
|
||||
hoisted = hoist_additional_tools(request_input, None, drop_params=False)
|
||||
|
||||
assert hoisted.input == [_USER_MESSAGE]
|
||||
assert hoisted.tools == ()
|
||||
|
|
@ -77,7 +77,7 @@ def test_tool_list_repeated_on_every_turn_keeps_each_tool_once_with_the_newest_d
|
|||
_USER_MESSAGE,
|
||||
]
|
||||
|
||||
hoisted: Final = hoist_additional_tools(request_input, None)
|
||||
hoisted: Final = hoist_additional_tools(request_input, None, drop_params=True)
|
||||
|
||||
assert hoisted.tools == (collaboration, _function("exec", "new"))
|
||||
|
||||
|
|
@ -89,7 +89,7 @@ def test_namespace_sent_in_several_turns_keeps_every_distinct_member_once() -> N
|
|||
_tool_list(_namespace("mcp__docs__", _function("fetch", "new"), _function("list"))),
|
||||
]
|
||||
|
||||
hoisted: Final = hoist_additional_tools(request_input, None)
|
||||
hoisted: Final = hoist_additional_tools(request_input, None, drop_params=True)
|
||||
|
||||
assert hoisted.tools == (
|
||||
_namespace("mcp__docs__", _function("search"), _function("fetch", "new"), _function("list")),
|
||||
|
|
@ -109,7 +109,7 @@ def test_unnamed_tools_collapse_only_when_identical() -> None:
|
|||
}
|
||||
request_input: Final = [_tool_list(docs_server, wiki_server), _USER_MESSAGE, _tool_list(docs_server)]
|
||||
|
||||
hoisted: Final = hoist_additional_tools(request_input, None)
|
||||
hoisted: Final = hoist_additional_tools(request_input, None, drop_params=True)
|
||||
|
||||
assert hoisted.tools == (docs_server, wiki_server)
|
||||
|
||||
|
|
@ -117,6 +117,20 @@ def test_unnamed_tools_collapse_only_when_identical() -> None:
|
|||
def test_top_level_tool_repeated_in_a_tool_list_item_is_kept_once() -> None:
|
||||
request_input: Final = [_tool_list(_TOP_LEVEL_TOOL, _WAIT_TOOL), _USER_MESSAGE]
|
||||
|
||||
hoisted: Final = hoist_additional_tools(request_input, [_TOP_LEVEL_TOOL])
|
||||
hoisted: Final = hoist_additional_tools(request_input, [_TOP_LEVEL_TOOL], drop_params=True)
|
||||
|
||||
assert hoisted.tools == (_TOP_LEVEL_TOOL, _WAIT_TOOL)
|
||||
|
||||
|
||||
def test_tool_list_repeated_on_every_turn_is_sent_as_is_without_drop_params() -> None:
|
||||
request_input: Final = [
|
||||
_tool_list(_function("exec", "old")),
|
||||
_USER_MESSAGE,
|
||||
_tool_list(_function("exec", "new")),
|
||||
_USER_MESSAGE,
|
||||
]
|
||||
|
||||
hoisted: Final = hoist_additional_tools(request_input, [_TOP_LEVEL_TOOL], drop_params=False)
|
||||
|
||||
assert hoisted.tools == (_TOP_LEVEL_TOOL, _function("exec", "old"), _function("exec", "new"))
|
||||
assert hoisted.input == [_USER_MESSAGE, _USER_MESSAGE]
|
||||
|
|
|
|||
124
tests/unit/test_assert_secret_absent.py
Normal file
124
tests/unit/test_assert_secret_absent.py
Normal file
|
|
@ -0,0 +1,124 @@
|
|||
import base64
|
||||
import importlib.util
|
||||
import io
|
||||
import sys
|
||||
import urllib.parse
|
||||
import zipfile
|
||||
from collections.abc import Callable, Mapping
|
||||
from itertools import chain, product
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
_REPO_ROOT: Final = Path(__file__).resolve().parents[2]
|
||||
_MODULE_PATH: Final = _REPO_ROOT / ".github" / "scripts" / "assert_secret_absent.py"
|
||||
_spec: Final = importlib.util.spec_from_file_location("assert_secret_absent", _MODULE_PATH)
|
||||
scanner: Final = importlib.util.module_from_spec(_spec)
|
||||
sys.modules[_spec.name] = scanner
|
||||
_spec.loader.exec_module(scanner)
|
||||
|
||||
SECRET: Final = "lic/ense+KEY:9f8e~~??>>~?7d6c5b4a3210-zz"
|
||||
PREFIXES: Final = ("", "x", "xy")
|
||||
ENCODERS: Final = (base64.b64encode, base64.urlsafe_b64encode)
|
||||
|
||||
|
||||
def _zip(members: Mapping[str, bytes], compression: int = zipfile.ZIP_DEFLATED) -> bytes:
|
||||
buffer: Final = io.BytesIO()
|
||||
with zipfile.ZipFile(buffer, "w", compression) as archive:
|
||||
for name, data in members.items():
|
||||
archive.writestr(name, data)
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
def _blob(prefix: str) -> bytes:
|
||||
return f"{prefix}{SECRET}tail".encode()
|
||||
|
||||
|
||||
def _planted_forms() -> tuple[str, ...]:
|
||||
encoded: Final = tuple(encode(_blob(prefix)).decode() for encode, prefix in product(ENCODERS, PREFIXES))
|
||||
return (SECRET, urllib.parse.quote(SECRET, safe=""), *encoded)
|
||||
|
||||
|
||||
def _assert_nothing_leaked(output: str) -> None:
|
||||
windows: Final = frozenset(
|
||||
chain.from_iterable((form[i : i + 8] for i in range(len(form) - 7)) for form in _planted_forms())
|
||||
)
|
||||
assert not [window for window in windows if window in output]
|
||||
|
||||
|
||||
def _run(root: Path, secret: str = SECRET) -> tuple[int, str]:
|
||||
out: Final = io.StringIO()
|
||||
code: Final = scanner.main((str(root),), {scanner.SECRET_VARIABLE: secret}, out)
|
||||
return code, out.getvalue()
|
||||
|
||||
|
||||
def _encoded_blob(prefix: str, encode: Callable[[bytes], bytes]) -> bytes:
|
||||
return b"trace " + encode(_blob(prefix)) + b" end"
|
||||
|
||||
|
||||
QUOTED: Final = urllib.parse.quote(SECRET, safe="")
|
||||
PLACEMENTS: Final = MappingProxyType(
|
||||
{
|
||||
"raw in a log": {"proxy.log": f"boot with key {SECRET} ok".encode()},
|
||||
"url-encoded in a url": {"upstream.log": f"GET /verify_license/{QUOTED}".encode()},
|
||||
"base64 aligned": {"a.json": _encoded_blob("", base64.b64encode)},
|
||||
"base64 one byte in": {"a.json": _encoded_blob("x", base64.b64encode)},
|
||||
"base64 two bytes in": {"a.json": _encoded_blob("xy", base64.b64encode)},
|
||||
"base64url two bytes in": {"a.json": _encoded_blob("xy", base64.urlsafe_b64encode)},
|
||||
"inside a zip member": {"trace.zip": _zip({"network.json": f'{{"auth": "{SECRET}"}}'.encode()})},
|
||||
"inside a nested zip": {"outer.zip": _zip({"inner.zip": _zip({"env.txt": SECRET.encode()})})},
|
||||
"named after it in a nested zip": {
|
||||
"outer.zip": _zip({"inner.zip": _zip({f"dump-{QUOTED}.txt": SECRET.encode()})})
|
||||
},
|
||||
"in a file name": {f"run-{QUOTED}.log": b"nothing inside"},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("placement", PLACEMENTS)
|
||||
def test_every_planted_copy_blocks_the_upload_without_printing_the_secret(tmp_path: Path, placement: str) -> None:
|
||||
worker: Final = tmp_path / "gw0"
|
||||
worker.mkdir()
|
||||
for name, data in PLACEMENTS[placement].items():
|
||||
(worker / name).write_bytes(data)
|
||||
code, output = _run(tmp_path)
|
||||
assert code == 1
|
||||
assert "secret found:" in output
|
||||
_assert_nothing_leaked(output)
|
||||
|
||||
|
||||
def test_clean_results_pass_and_count_every_file(tmp_path: Path) -> None:
|
||||
(tmp_path / "junit.xml").write_text("<testsuite/>")
|
||||
(tmp_path / "gw0").mkdir()
|
||||
(tmp_path / "gw0" / "proxy.log").write_text("LITELLM_LICENSE is not set")
|
||||
code, output = _run(tmp_path)
|
||||
assert code == 0
|
||||
assert "no copy of the secret in 2 files" in output
|
||||
|
||||
|
||||
def test_an_unset_secret_has_nothing_to_find(tmp_path: Path) -> None:
|
||||
(tmp_path / "proxy.log").write_text(SECRET)
|
||||
assert _run(tmp_path, secret="")[0] == 0
|
||||
|
||||
|
||||
def test_a_symlinked_directory_fails_closed(tmp_path: Path) -> None:
|
||||
outside: Final = tmp_path / "outside"
|
||||
outside.mkdir()
|
||||
results: Final = tmp_path / "results"
|
||||
results.mkdir()
|
||||
(results / "linked").symlink_to(outside, target_is_directory=True)
|
||||
code, output = _run(results)
|
||||
assert code == 1
|
||||
assert "directory symlink" in output
|
||||
|
||||
|
||||
def test_an_unreadable_archive_fails_without_echoing_its_contents(tmp_path: Path) -> None:
|
||||
archive: Final = _zip({"env.txt": f"LITELLM_LICENSE={SECRET}".encode()}, zipfile.ZIP_STORED)
|
||||
at: Final = archive.index(SECRET.encode())
|
||||
(tmp_path / "trace.zip").write_bytes(archive[:at] + bytes([archive[at] ^ 0xFF]) + archive[at + 1 :])
|
||||
code, output = _run(tmp_path)
|
||||
assert code == 1
|
||||
assert "scan failed with BadZipFile" in output
|
||||
_assert_nothing_leaked(output)
|
||||
Loading…
Add table
Reference in a new issue