mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
* feat(langfuse): migrate the sdk callback to langfuse v4
Replace the v2 trace()/generation()/span() calls with SDK v4 observations exported over OpenTelemetry, with one isolated tracer provider per Langfuse credential set, a discarding exporter for mock mode, and v4 trace and observation id normalization. Keeps the session-header trace provenance logic from main so each call under a session alias still gets its own trace
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): drop the always-true prompt client check now that v4 get_prompt is non-optional
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(langfuse): isolate the e2e sync test from cached clients and log the real sdk major
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* refactor(langfuse): type the slack trace-url lookup and drop dead v2 test shims
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(slack): cover the langfuse trace url built from the logger host
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* build(docker): pin langfuse to the locked 4.15.2 in the pip image
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): hash all-zero trace and observation ids instead of passing them through
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* feat(langfuse): honour caller generation ids and assert v4 OTLP exports in legacy tests
v2 accepted generation(id=...). v4 derives the observation id from the OTel
span id, so the isolated tracer provider now carries an id generator that
hands out the id start_generation asked for through a context variable, and
the callback passes the resolved generation_id metadata into it.
The legacy e2e suite patched httpx.Client.post and compared v2 ingestion
batches; it now patches requests.Session.post, decodes the OTLP protobuf
and compares the exported generation against regenerated fixtures. The
local readback test replaces the removed get_generations() with
api.observations.get_many() and polls Langfuse Cloud instead of sleeping.
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* refactor(langfuse): read the sdk version header from package metadata
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): propagate trace_metadata as trace-level attributes in v4
v2 wrote trace(metadata=...) onto the trace object. In v4 the trace only
carries what the observations propagate, so a continuation request with
update_trace_keys=["trace_metadata"] updated the generation's metadata
while the trace kept its stale values. Coerce each entry to the SDK's
string limit and hand it to propagate_attributes(metadata=...).
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): propagate interrupts raised during deferred client teardown
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): honor ssl_verify=False and SSL_VERIFY on the v4 OTLP exporter
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): fall back to the default CA when the configured bundle path is missing
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): renew the client when eviction lands before the callback lease
The cache can evict a logger between handing it to the callback and the callback taking its
lease. Such a lease now hands back a fresh client acquired through the same parameters, so that
callback exports through a live tracer provider instead of one teardown already shut down.
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): emit litellm_call_id and response_id as generation metadata
v2 put the provider response id inside the generation id. v4 observation ids are 16 hex chars derived from that string, so the ids move to generation metadata to keep generations searchable by response id
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* refactor(langfuse): read the response id through a typed protocol
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): do not claim trace root when continuing an existing trace
Langfuse derives a trace's name and I/O from any observation flagged
langfuse.internal.as_root, so a request carrying existing_trace_id
renamed the trace to the generation name and replaced the trace input
and output on every continuation. v2 only updated the keys listed in
update_trace_keys. Continuations now export as plain children of the
remote parent and keep the explicit langfuse.trace.* attributes for the
fields they do want changed.
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): iterate lease renewal instead of recursing, monkeypatch update_trace_keys flag in tests
The recursive lease fallback tripped tests/code_coverage_tests/recursive_detector.py; the renewal
candidates are now walked with itertools.chain. The six update_trace_keys tests set the litellm
global through pytest monkeypatch so the TQ008 budget stays within its ceiling
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): retry raised OTLP exports and honor LANGFUSE_TIMEOUT
The OTLP http exporter only retries 429 and 5xx; a connect or read timeout
propagates and BatchSpanProcessor drops the batch. Wrap the exporter in
RetryingSpanExporter (three backoff retries, as the v2 consumer did) and
build it on every path so the default and private-CA deployments share the
same channel, timeout and retry behaviour
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): sample on a hash of the full trace id and tolerate bad LANGFUSE_SAMPLE_RATE
TraceIdRatioBased reads the low 64 bits of the trace id. litellm trace ids are
UUIDs, whose variant bits sit at the top of that word, so every fractional rate
up to 0.5 dropped all traces. A SHA-256 of the full id gives an unbiased,
deterministic decision. Values outside [0, 1] or non numeric now warn and export
everything instead of raising during callback construction, which surfaced as a
500 on the first request of each worker
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): put the Langfuse trace link back into Slack alerts
The proxy registers LangfusePromptManagement for callbacks: ["langfuse"], so the alert helper never saw the literal "langfuse" string and returned before looking up the trace id, and the prompt management logger never stored the trace id it got back from log_event_on_langfuse. Recognize LangFuseLogger instances in the callback list, record the returned trace id in the shared service trace id cache, and skip the link when no trace id arrives
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* chore(deps): relock langfuse 4.15.2 and opentelemetry 1.33.1 on current main
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* chore(langfuse): mark the deliberate blind except in client teardown for the strict ruff gate
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* refactor(langfuse): pass the resource attributes mapping straight to Resource.create
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): warn about ignored UPSTREAM_LANGFUSE_* on the shared client init path too
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): normalise the OTLP export path so a trailing host slash never yields a double slash
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): nest guardrail and grounding spans under the generation
Langfuse v4 derives the trace name and I/O from every observation marked as_root, and the one with the latest start time wins. Guardrail and grounding spans used to claim root next to the generation, so a post_call guardrail could replace the model's request and response on the trace with its own. Only the generation claims root now; the sibling spans become its children
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): rebuild the cached bundle when mock mode or sample rate changes
The SDK keys resource bundles on the public key alone, so a bundle built with the discarding exporter for LANGFUSE_MOCK, or with an earlier LANGFUSE_SAMPLE_RATE, was handed back to a client that asked for a live exporter or a different rate. Compare both when deciding whether the cached bundle is still valid
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): keep trace_public true when a guardrail span is exported
Langfuse folds langfuse.trace.public across every observation in the trace and reads a missing attribute as false, so a guardrail child span without the flag turned a trace_public: true request private on Langfuse Cloud. Child spans now repeat the generation's value
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* refactor(langfuse): emit observations as plain OTel spans, keep the SDK for prompts and auth
The callback now owns an isolated TracerProvider and OTLP exporter and builds generation and child spans with public OpenTelemetry APIs plus the LangfuseOtelSpanAttributes constants. Caller trace ids, generation ids, parent observation ids and historical start and end times are honoured through the OTel id generator, remote SpanContext and explicit span timestamps, so no private Langfuse SDK tracing handle is used any more. The Langfuse client stays only for get_prompt and auth_check
This also resolves the gauntlet findings on the previous draft: fresh traces start from an empty context so caller application spans are never stamped, the Slack trace link is read from the request logging state instead of constructing a logger per alert, a truthy non-mapping trace_metadata is serialized instead of raising, trace_input and trace_output land on the root generation, discarding a cached client is done under the lock, and the prompt cache no longer leaks a task manager because the client cache no longer tears down shared providers
Fixtures under tests/logging_callback_tests lose the SDK-private langfuse.internal.as_root marker; every other exported attribute is unchanged
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): hand the SDK client a validated sample rate so an unusable LANGFUSE_SAMPLE_RATE no longer breaks the callback
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): gate the SDK version before importing the OTel module in prompt management
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): flush every export channel on proxy shutdown and use the callback's host in Slack trace links
The shutdown hook imported litellm.utils.langFuseLogger, a global the callback registry never assigns, so a graceful restart dropped the spans still queued in the batch processors. Shutdown now calls flush_langfuse_tracing, which force-flushes every acquired channel. The Slack alert link falls back to the registered LangFuseLogger's langfuse_host when the request carries no dynamic host, and the export endpoint tests pin that scheme-relative or absolute LANGFUSE_OTEL_TRACES_EXPORT_PATH values stay on the configured host
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): store resolved credentials on LangfusePromptManagement
The Slack alert trace link reads langfuse_host from every registered LangFuseLogger. Prompt management subclasses it without calling the parent constructor, so it never set the attribute and the alerting handler crashed before posting
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): flush every export channel concurrently under one shutdown deadline
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): flush export channels on daemon threads so a stuck channel cannot hold up interpreter exit
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): own the tracer config and drop the SDK client for prompts and auth
The callback's TracerProvider now sets its sampler, span limits and id generator explicitly so unrelated OTEL_* variables no longer change what Langfuse receives, and trace metadata is written once on the trace instead of folded into the generation, which kept input and output under the attribute cap. Spans are emitted under the langfuse-sdk scope so Langfuse renders them natively, the batch processor queues 100k spans and honors LANGFUSE_FLUSH_AT, and the proxy shutdown flush runs off the event loop with a 10s deadline and logs a miss.
Prompts, auth_check and the project id now go through LangfuseAPI directly with a litellm-owned TTL cache, so no Langfuse() client is built and a host application's client on the same public key is left alone. Dead attributes, the unreachable exporter branch and the export list are cleaned up, and the client-budget eviction behavior is documented.
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* refactor(langfuse): export OTLP spans and fetch prompts through litellm's HTTPHandler instead of a private requests session
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): gate the SDK version before importing the tracing module and retire unheld export channels
An installed v2 SDK used to fail inside the langfuse_sdk import and surface as "Langfuse not installed"; the version check now runs first so v2 users get the upgrade message, and only PackageNotFoundError means the package is missing
Export channels are now leased per credential set: acquire adds a holder, LangFuseLogger.stop (called by DynamicLoggingCache on expiry) releases one, and a channel with no holders is flushed and shut down after a 60 s grace, so rotating key or team credentials no longer grows one batch thread per credential set for the life of the process
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): end the generation when a child span fails, take the client slot last, keep prompt cache keys structured
Generation spans now end in a finally block so a bad guardrail or provider entry cannot strand the trace. The logger acquires its export channel and REST client before counting a client slot and releases the channel synchronously if the REST client fails to build, so retries after a bad config do not exhaust the budget. LANGFUSE_TIMEOUT accepts decimals for the REST client like it already did for OTLP export. The prompt cache keys on (name, version, label) so a missing label and the literal label None stay apart
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): claim the cache entry before releasing its slot and channel hold on eviction
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): coerce generation names, keep v2 release, timeout and retry defaults, refresh stale prompts off the loop
A non-string metadata generation_name reached the OTLP encoder and took the whole batch down; it is now exported as its text and the exporter drops only the span the encoder rejects. LANGFUSE_RELEASE falls back to the deploy platform's commit variable again, the export deadline is back to the v2 default of 20 s and LANGFUSE_MAX_RETRIES sizes the retry ladder. An expired prompt is served at once while one background thread refreshes it, a re-acquired export channel cancels the pending retire timer, and flush reports delivery rather than a drained queue
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(proxy): assert the current Langfuse shutdown flush warning
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): keep host OTel resource out, carry big metadata ints, tolerate bad flush and TTL env, stamp trace I/O under a parent
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): name a malformed prompt cache TTL before the SDK import, keep metadata ints JSON safe, retry every 5xx export
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): name the auth check failure, split a 413 export, wire LANGFUSE_DEBUG, stamp error output under a parent, send the ingestion version header
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): honor LANGFUSE_DEBUG on the callbacks path, cap retry backoff, name the auth failure status and body
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): cap LANGFUSE_MAX_RETRIES at 1000 so an absurd value cannot stall callback init
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* refactor(langfuse): fold 413 halving into bounded rounds instead of recursion
The code-quality recursive-function gate flagged LangfuseSpanExporter.export. A batch of n spans settles within n.bit_length() halving rounds, so the split is a reduce over a frozen round state with the same posts, logs and results. The TTL gate test now asserts the gate returns without raising
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): truncate a single oversized span like v2 instead of dropping it, no retries on REST auth and project lookups
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): write the metadata truncation marker under a flattened key so Langfuse keeps it
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(langfuse): patch the HTTPHandler export path and sync the metadata fixture and lease registry with the v4 callback
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* refactor(langfuse): give the 413 split helpers a single explicit return path
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): url-encode prompt names and fetch cold prompts without client retries
A cold get_prompt runs inline on the event loop; the generated v4 client's default two retries slept through
Retry-After (up to 60 s per attempt) and held the loop. The wrapper also passed the raw name into
api/public/v2/prompts/{name}, so 'what?' fetched prompt 'what' and folder names left the route. Quote the
name with safe='' like the v4 SDK's own get_prompt and pass max_retries=0 like the projects.get calls
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(langfuse): retry a cold prompt miss once and drop upstream headers from prompt errors
A cold prompt fetch makes one immediate second attempt after a 5xx or a
transport failure, as the v2 client did, still with the generated client's
sleeping retries and Retry-After handling off so the event loop never stalls.
A failed fetch raises LangfusePromptError carrying only the status and body,
so the proxy no longer forwards Langfuse's response headers to its client
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(langfuse): stub the logger in the health auth_check test instead of dialing a closed port
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(langfuse): integration test for OTLP v4 delivery and prompt fetch through a real proxy
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* build(docker): keep the pip image's langfuse and otel pins on the v2 line its litellm 1.83.0 wheel expects
The image validates the published PyPI artifact, whose langfuse callback still
reads langfuse.version, so the 4.15.2 pin broke that callback. The pins move
together with the next LITELLM_VERSION bump. Also rewords the trace_version
precedence test docstring: v2 carried two version fields, v4 has one per span
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---------
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
1273 lines
49 KiB
Python
1273 lines
49 KiB
Python
"""Behavior pins for proxy_server lifecycle, helpers, and small utilities.
|
|
|
|
Pins covered:
|
|
- ``proxy_startup_event``
|
|
- ``proxy_shutdown_event``
|
|
- ``_initialize_shared_aiohttp_session``
|
|
- ``cleanup_router_config_variables``
|
|
- ``save_worker_config``
|
|
- ``initialize``
|
|
- ``load_from_azure_key_vault``
|
|
- ``cost_tracking``
|
|
- ``_resolve_typed_dict_type``
|
|
- ``_resolve_pydantic_type``
|
|
- ``get_litellm_model_info``
|
|
- ``run_ollama_serve``
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import inspect
|
|
import json
|
|
import logging
|
|
import os
|
|
import subprocess
|
|
from collections.abc import Awaitable, Callable
|
|
from typing import List, Optional, Union
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
|
from fastapi import FastAPI
|
|
from pydantic import BaseModel
|
|
from typing_extensions import TypedDict
|
|
|
|
import litellm.proxy.proxy_server as ps
|
|
from litellm.proxy.proxy_server import (
|
|
ProxyStartupEvent,
|
|
_initialize_shared_aiohttp_session,
|
|
_resolve_pydantic_type,
|
|
_resolve_typed_dict_type,
|
|
cleanup_router_config_variables,
|
|
cost_tracking,
|
|
get_litellm_model_info,
|
|
initialize,
|
|
initialize_from_worker_config,
|
|
load_from_azure_key_vault,
|
|
proxy_shutdown_event,
|
|
proxy_startup_event,
|
|
run_ollama_serve,
|
|
save_worker_config,
|
|
)
|
|
|
|
from .conftest import normalize
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# cleanup_router_config_variables
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_cleanup_router_config_variables_resets_globals(monkeypatch):
|
|
monkeypatch.setattr(ps, "master_key", "sk-sentinel", raising=False)
|
|
monkeypatch.setattr(ps, "user_config_file_path", "/tmp/config.yaml", raising=False)
|
|
monkeypatch.setattr(ps, "user_custom_auth", lambda x: x, raising=False)
|
|
monkeypatch.setattr(ps, "health_check_interval", 42, raising=False)
|
|
monkeypatch.setattr(ps, "prisma_client", MagicMock(), raising=False)
|
|
monkeypatch.setattr(ps, "heuristic_v1_tuning_baselines", {"router": "baseline"}, raising=False)
|
|
|
|
cleanup_router_config_variables()
|
|
|
|
observed = {
|
|
"master_key": ps.master_key,
|
|
"user_config_file_path": ps.user_config_file_path,
|
|
"user_custom_auth": ps.user_custom_auth,
|
|
"health_check_interval": ps.health_check_interval,
|
|
"prisma_client": ps.prisma_client,
|
|
"heuristic_v1_tuning_baselines": ps.heuristic_v1_tuning_baselines,
|
|
}
|
|
assert normalize(observed) == {
|
|
"master_key": None,
|
|
"user_config_file_path": None,
|
|
"user_custom_auth": None,
|
|
"health_check_interval": None,
|
|
"prisma_client": None,
|
|
"heuristic_v1_tuning_baselines": None,
|
|
}
|
|
|
|
|
|
def test_cleanup_router_config_variables_fails_on_unknown_attr_raises():
|
|
"""The function only writes documented globals — accessing a non-existent
|
|
one after cleanup should still raise AttributeError."""
|
|
cleanup_router_config_variables()
|
|
with pytest.raises(AttributeError):
|
|
_ = ps.this_attribute_should_not_exist_xyz
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# proxy_shutdown_event
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_shutdown_event_disconnects_prisma_and_resets(monkeypatch):
|
|
fake_prisma = MagicMock()
|
|
fake_prisma.disconnect = AsyncMock()
|
|
monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
|
|
monkeypatch.setattr(ps, "master_key", "sk-x", raising=False)
|
|
|
|
fake_jwt = MagicMock()
|
|
fake_jwt.close = AsyncMock()
|
|
monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False)
|
|
monkeypatch.setattr(ps, "db_writer_client", None, raising=False)
|
|
|
|
import litellm
|
|
|
|
monkeypatch.setattr(litellm, "cache", None, raising=False)
|
|
monkeypatch.setattr(litellm, "success_callback", [], raising=False)
|
|
|
|
await proxy_shutdown_event()
|
|
|
|
observed = {
|
|
"disconnect_called": fake_prisma.disconnect.await_count == 1,
|
|
"jwt_closed": fake_jwt.close.await_count == 1,
|
|
"master_key_reset": ps.master_key,
|
|
"prisma_reset": ps.prisma_client,
|
|
}
|
|
assert normalize(observed) == {
|
|
"disconnect_called": True,
|
|
"jwt_closed": True,
|
|
"master_key_reset": None,
|
|
"prisma_reset": None,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_shutdown_flushes_every_langfuse_export_channel(monkeypatch):
|
|
"""A generation finished just before a graceful restart is still queued in its batch
|
|
processor, so shutdown must flush every acquired export channel."""
|
|
from litellm.integrations.langfuse import langfuse_sdk
|
|
|
|
flushed = MagicMock(return_value=True)
|
|
monkeypatch.setattr(langfuse_sdk, "flush_langfuse_tracing", flushed)
|
|
monkeypatch.setattr(ps, "prisma_client", None, raising=False)
|
|
monkeypatch.setattr(ps, "jwt_handler", MagicMock(close=AsyncMock()), raising=False)
|
|
monkeypatch.setattr(ps, "db_writer_client", None, raising=False)
|
|
|
|
import litellm
|
|
|
|
monkeypatch.setattr(litellm, "cache", None, raising=False)
|
|
monkeypatch.setattr(litellm, "success_callback", [], raising=False)
|
|
|
|
await proxy_shutdown_event()
|
|
|
|
assert flushed.call_count == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_shutdown_flushes_langfuse_off_the_event_loop_and_logs_a_timeout(monkeypatch, caplog):
|
|
"""The flush blocks on OTLP exports for up to its deadline, so it must run on a worker thread
|
|
with the shutdown deadline, and a channel that misses it is reported instead of ignored."""
|
|
import threading
|
|
|
|
from litellm.constants import LANGFUSE_SHUTDOWN_FLUSH_TIMEOUT_MILLIS
|
|
from litellm.integrations.langfuse import langfuse_sdk
|
|
|
|
ran_on = MagicMock()
|
|
|
|
def flushed(timeout_millis: int) -> bool:
|
|
ran_on(threading.current_thread(), timeout_millis)
|
|
return False
|
|
|
|
monkeypatch.setattr(langfuse_sdk, "flush_langfuse_tracing", flushed)
|
|
monkeypatch.setattr(ps, "prisma_client", None, raising=False)
|
|
monkeypatch.setattr(ps, "jwt_handler", MagicMock(close=AsyncMock()), raising=False)
|
|
monkeypatch.setattr(ps, "db_writer_client", None, raising=False)
|
|
|
|
import litellm
|
|
|
|
monkeypatch.setattr(litellm, "cache", None, raising=False)
|
|
monkeypatch.setattr(litellm, "success_callback", [], raising=False)
|
|
|
|
with caplog.at_level("WARNING", logger="LiteLLM Proxy"):
|
|
await proxy_shutdown_event()
|
|
|
|
(flush_thread, timeout_millis), _ = ran_on.call_args
|
|
assert flush_thread is not threading.main_thread()
|
|
assert timeout_millis == LANGFUSE_SHUTDOWN_FLUSH_TIMEOUT_MILLIS
|
|
assert any("Langfuse shutdown flush incomplete" in record.getMessage() for record in caplog.records)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_shutdown_drains_gateway_requests_before_disconnecting(monkeypatch):
|
|
"""
|
|
The gateway request fold lives in memory, so shutdown drains it to the database.
|
|
|
|
That drain has to happen while prisma is still connected: a write attempted
|
|
after ``disconnect()`` raises ClientNotConnectedError, the flush swallows it
|
|
and merges the counts back onto an accumulator the process is about to
|
|
discard, and the final interval is lost silently on every restart. Ordering is
|
|
the whole behavior here, so assert the order rather than that both ran.
|
|
"""
|
|
calls: list = [] # mutable-ok: records call order, which is the assertion
|
|
|
|
fake_prisma = MagicMock()
|
|
fake_prisma.disconnect = AsyncMock(side_effect=lambda: calls.append("disconnect"))
|
|
monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
|
|
|
|
async def _record_flush(client, accumulator):
|
|
calls.append("flush")
|
|
assert client is fake_prisma
|
|
|
|
monkeypatch.setattr(ps, "flush_gateway_requests", _record_flush, raising=False)
|
|
|
|
fake_jwt = MagicMock()
|
|
fake_jwt.close = AsyncMock()
|
|
monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False)
|
|
monkeypatch.setattr(ps, "db_writer_client", None, raising=False)
|
|
|
|
import litellm
|
|
|
|
monkeypatch.setattr(litellm, "cache", None, raising=False)
|
|
monkeypatch.setattr(litellm, "success_callback", [], raising=False)
|
|
|
|
await proxy_shutdown_event()
|
|
|
|
assert calls == ["flush", "disconnect"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_shutdown_skips_gateway_flush_without_a_database(monkeypatch):
|
|
"""No prisma client means nothing to drain to, and no attempt is made."""
|
|
flush = AsyncMock()
|
|
monkeypatch.setattr(ps, "flush_gateway_requests", flush, raising=False)
|
|
monkeypatch.setattr(ps, "prisma_client", None, raising=False)
|
|
|
|
fake_jwt = MagicMock()
|
|
fake_jwt.close = AsyncMock()
|
|
monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False)
|
|
monkeypatch.setattr(ps, "db_writer_client", None, raising=False)
|
|
|
|
import litellm
|
|
|
|
monkeypatch.setattr(litellm, "cache", None, raising=False)
|
|
monkeypatch.setattr(litellm, "success_callback", [], raising=False)
|
|
|
|
await proxy_shutdown_event()
|
|
|
|
assert flush.await_count == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_shutdown_event_prisma_disconnect_raises_error(monkeypatch):
|
|
fake_prisma = MagicMock()
|
|
fake_prisma.disconnect = AsyncMock(side_effect=RuntimeError("db gone"))
|
|
monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
|
|
|
|
fake_jwt = MagicMock()
|
|
fake_jwt.close = AsyncMock()
|
|
monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False)
|
|
|
|
import litellm
|
|
|
|
monkeypatch.setattr(litellm, "cache", None, raising=False)
|
|
monkeypatch.setattr(litellm, "success_callback", [], raising=False)
|
|
|
|
with pytest.raises(RuntimeError, match="db gone"):
|
|
await proxy_shutdown_event()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _flush_spend_logs_queue_on_shutdown
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_flush_spend_logs_queue_on_shutdown_drains_before_disconnect(monkeypatch):
|
|
fake_prisma = MagicMock()
|
|
monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
|
|
monkeypatch.setattr(ps, "db_writer_client", None, raising=False)
|
|
|
|
drain = AsyncMock()
|
|
import litellm.proxy.utils as utils_mod
|
|
|
|
monkeypatch.setattr(utils_mod, "drain_spend_logs_queue", drain)
|
|
|
|
await ps._flush_spend_logs_queue_on_shutdown()
|
|
|
|
observed = {
|
|
"drain_calls": drain.await_count,
|
|
"drain_prisma": drain.await_args.kwargs["prisma_client"] is fake_prisma,
|
|
}
|
|
assert observed == {
|
|
"drain_calls": 1,
|
|
"drain_prisma": True,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_flush_spend_logs_queue_on_shutdown_swallows_drain_errors(monkeypatch):
|
|
monkeypatch.setattr(ps, "prisma_client", MagicMock(), raising=False)
|
|
monkeypatch.setattr(ps, "db_writer_client", None, raising=False)
|
|
|
|
import litellm.proxy.utils as utils_mod
|
|
|
|
monkeypatch.setattr(
|
|
utils_mod,
|
|
"drain_spend_logs_queue",
|
|
AsyncMock(side_effect=RuntimeError("db gone")),
|
|
)
|
|
|
|
await ps._flush_spend_logs_queue_on_shutdown()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_flush_spend_counters_on_shutdown_commits_buffered_spend(monkeypatch):
|
|
fake_prisma = MagicMock()
|
|
monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
|
|
commit = AsyncMock()
|
|
monkeypatch.setattr(ps.proxy_logging_obj.db_spend_update_writer, "db_update_spend_transaction_handler", commit)
|
|
|
|
await ps.flush_spend_counters_on_shutdown()
|
|
|
|
observed = {
|
|
"commit_calls": commit.await_count,
|
|
"commit_prisma": commit.await_args.kwargs["prisma_client"] is fake_prisma,
|
|
"commit_proxy_logging": commit.await_args.kwargs["proxy_logging_obj"] is ps.proxy_logging_obj,
|
|
}
|
|
assert observed == {"commit_calls": 1, "commit_prisma": True, "commit_proxy_logging": True}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_flush_spend_counters_on_shutdown_logs_and_swallows_commit_errors(monkeypatch, caplog):
|
|
monkeypatch.setattr(ps, "prisma_client", MagicMock(), raising=False)
|
|
monkeypatch.setattr(
|
|
ps.proxy_logging_obj.db_spend_update_writer,
|
|
"db_update_spend_transaction_handler",
|
|
AsyncMock(side_effect=RuntimeError("db gone")),
|
|
)
|
|
|
|
with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"):
|
|
await ps.flush_spend_counters_on_shutdown()
|
|
|
|
assert "Error flushing spend counters on shutdown: db gone" in caplog.text
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _initialize_shared_aiohttp_session
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_initialize_shared_aiohttp_session_returns_client_session():
|
|
from aiohttp import ClientSession
|
|
|
|
session = await _initialize_shared_aiohttp_session()
|
|
try:
|
|
observed = {
|
|
"is_client_session": isinstance(session, ClientSession),
|
|
"is_closed": session.closed,
|
|
"has_connector": session.connector is not None,
|
|
}
|
|
assert normalize(observed) == {
|
|
"is_client_session": True,
|
|
"is_closed": False,
|
|
"has_connector": True,
|
|
}
|
|
finally:
|
|
if session is not None:
|
|
await session.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_initialize_shared_aiohttp_session_aiohttp_missing_returns_none_on_failure(
|
|
monkeypatch,
|
|
):
|
|
"""If aiohttp import fails, the function catches and returns None — no raise."""
|
|
import builtins
|
|
|
|
real_import = builtins.__import__
|
|
|
|
def _raise_for_aiohttp(name, *args, **kwargs):
|
|
if name == "aiohttp":
|
|
raise ImportError("simulated missing aiohttp")
|
|
return real_import(name, *args, **kwargs)
|
|
|
|
monkeypatch.setattr(builtins, "__import__", _raise_for_aiohttp)
|
|
result = await _initialize_shared_aiohttp_session()
|
|
assert result is None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# save_worker_config
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_save_worker_config_writes_json_to_environ(monkeypatch):
|
|
monkeypatch.delenv("WORKER_CONFIG", raising=False)
|
|
|
|
save_worker_config(model="gpt-4", config="/tmp/c.yaml", debug=True)
|
|
|
|
payload = json.loads(os.environ["WORKER_CONFIG"])
|
|
assert normalize(payload) == {
|
|
"model": "gpt-4",
|
|
"config": "/tmp/c.yaml",
|
|
"debug": True,
|
|
}
|
|
|
|
|
|
def test_save_worker_config_invalid_no_kwargs_yields_empty(monkeypatch):
|
|
monkeypatch.delenv("WORKER_CONFIG", raising=False)
|
|
|
|
save_worker_config()
|
|
assert os.environ["WORKER_CONFIG"] == "{}"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _redact_worker_config_for_logging (LIT-4152)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
_LIT4152_SECRETS = (
|
|
"sk-lit4152-regression-master-key-abcdef1234567890",
|
|
"leak_password_9090",
|
|
"sk-lit4152-provider-api-key-abcdef",
|
|
"postgresql://leak_user:leak_password_9090@leak-host.internal:5432/leak_db",
|
|
)
|
|
|
|
|
|
def _lit4152_worker_config_dict():
|
|
return {
|
|
"model": "openai/gpt-4o-mini",
|
|
"config": "/tmp/c.yaml",
|
|
"master_key": _LIT4152_SECRETS[0],
|
|
"database_url": _LIT4152_SECRETS[3],
|
|
"api_key": _LIT4152_SECRETS[2],
|
|
"drop_params": True,
|
|
}
|
|
|
|
|
|
def test__redact_worker_config_for_logging_dict_masks_all_secret_shapes():
|
|
"""LIT-4152 regression: dict-form worker_config must not embed any raw
|
|
secret. Covers the segment-matched fields (`master_key`, `api_key`) and the
|
|
URL-with-credentials field (`database_url`), which the segment masker
|
|
misses because neither segment matches its sensitive-pattern set.
|
|
"""
|
|
from litellm.proxy.proxy_server import _redact_worker_config_for_logging
|
|
|
|
redacted = _redact_worker_config_for_logging(_lit4152_worker_config_dict())
|
|
rendered = repr(redacted)
|
|
for secret in _LIT4152_SECRETS:
|
|
assert secret not in rendered, f"leak: {secret} in {rendered!r}"
|
|
assert isinstance(redacted, dict)
|
|
assert redacted["model"] == "openai/gpt-4o-mini"
|
|
assert redacted["drop_params"] is True
|
|
|
|
|
|
def test__redact_worker_config_for_logging_json_string_round_trips_masked():
|
|
"""Docker/K8s deployments hand the proxy a JSON string via ``WORKER_CONFIG``.
|
|
Confirm the string path also masks and that the returned value re-parses
|
|
into a dict with the sensitive fields masked.
|
|
"""
|
|
from litellm.proxy.proxy_server import _redact_worker_config_for_logging
|
|
|
|
payload = json.dumps(_lit4152_worker_config_dict())
|
|
redacted = _redact_worker_config_for_logging(payload)
|
|
assert isinstance(redacted, str)
|
|
for secret in _LIT4152_SECRETS:
|
|
assert secret not in redacted, f"leak: {secret} in {redacted!r}"
|
|
parsed = json.loads(redacted)
|
|
assert parsed["model"] == "openai/gpt-4o-mini"
|
|
|
|
|
|
def test__redact_worker_config_for_logging_passthrough_for_none_and_non_json_string():
|
|
"""Non-dict, non-JSON-parseable string is passed through verbatim (nothing
|
|
to mask) and ``None`` returns ``None``.
|
|
"""
|
|
from litellm.proxy.proxy_server import _redact_worker_config_for_logging
|
|
|
|
assert _redact_worker_config_for_logging(None) is None
|
|
assert _redact_worker_config_for_logging("/tmp/some_config.yaml") == "/tmp/some_config.yaml"
|
|
|
|
|
|
def test__redact_worker_config_for_logging_masks_non_string_url_webhook_values():
|
|
"""The URL/webhook fields the segment masker cannot catch by key name
|
|
(``alert_to_webhook_url``, ``pass_through_endpoints``,
|
|
``database_extra_connection_params``) can hold non-string shapes:
|
|
``alert_to_webhook_url`` is typed as ``Optional[Dict]`` and can nest
|
|
secret query params under keys the segment masker also misses. Confirm
|
|
the whole value is replaced regardless of shape so a nested webhook or
|
|
Bearer token under a non-segment-matched key does not slip through.
|
|
"""
|
|
from litellm.proxy.proxy_server import _redact_worker_config_for_logging
|
|
|
|
nested_webhook_secret = "https://hooks.slack.com/services/T0/B0/nested-webhook-secret-xyz"
|
|
data = {
|
|
"master_key": "sk-should-be-masked",
|
|
"alert_to_webhook_url": {"budget_alerts": nested_webhook_secret},
|
|
"pass_through_endpoints": [
|
|
{
|
|
"path": "/upstream",
|
|
"target": "https://api.provider.com",
|
|
"headers": {"Authorization": "Bearer nested-token-should-be-gone"},
|
|
}
|
|
],
|
|
"database_extra_connection_params": {"password": "extra-db-password-abc"},
|
|
}
|
|
redacted = _redact_worker_config_for_logging(data)
|
|
rendered = repr(redacted)
|
|
for secret in (
|
|
"sk-should-be-masked",
|
|
nested_webhook_secret,
|
|
"nested-token-should-be-gone",
|
|
"extra-db-password-abc",
|
|
):
|
|
assert secret not in rendered, f"leak: {secret} in {rendered!r}"
|
|
|
|
|
|
def test__redact_worker_config_for_logging_masks_nested_secret_fields():
|
|
"""LIT-4152 nested regression: the URL/webhook credential fields the segment
|
|
masker cannot catch by name (``database_url``,
|
|
``database_extra_connection_params``, ``pass_through_endpoints``,
|
|
``alert_to_webhook_url``) must be redacted at any depth, not just the top
|
|
level. A worker_config that nests ``general_settings`` under a parent key
|
|
must not leak a nested ``database_url`` or webhook secret; the earlier
|
|
top-level-only redaction would have passed these through raw.
|
|
"""
|
|
from litellm.proxy.proxy_server import _redact_worker_config_for_logging
|
|
|
|
nested_db_url = "postgresql://nested_user:nested_pw_4152@nested-host:5432/db"
|
|
nested_webhook = "https://hooks.slack.com/services/T0/B0/nested-4152-webhook"
|
|
nested_extra_pw = "nested-extra-conn-pw-4152"
|
|
nested_bearer = "Bearer nested-passthrough-token-4152"
|
|
data = {
|
|
"config": {
|
|
"general_settings": {
|
|
"database_url": nested_db_url,
|
|
"database_extra_connection_params": {"password": nested_extra_pw},
|
|
"alert_to_webhook_url": {"budget_alerts": nested_webhook},
|
|
"pass_through_endpoints": [{"path": "/up", "headers": {"Authorization": nested_bearer}}],
|
|
}
|
|
}
|
|
}
|
|
redacted = _redact_worker_config_for_logging(data)
|
|
rendered = repr(redacted)
|
|
for secret in (nested_db_url, nested_webhook, nested_extra_pw, nested_bearer):
|
|
assert secret not in rendered, f"nested leak: {secret} in {rendered!r}"
|
|
|
|
inner = redacted["config"]["general_settings"]
|
|
assert inner["database_url"] == "REDACTED"
|
|
assert inner["pass_through_endpoints"] == "REDACTED"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# initialize
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_initialize_signature_is_async_with_expected_params():
|
|
sig = inspect.signature(initialize)
|
|
# Hard-coded so a signature change (param added/removed) trips the gate.
|
|
expected_param_count = 16
|
|
observed = {
|
|
"is_async": inspect.iscoroutinefunction(initialize),
|
|
"param_count": len(sig.parameters),
|
|
"has_model": "model" in sig.parameters,
|
|
"has_config": "config" in sig.parameters,
|
|
}
|
|
assert normalize(observed) == {
|
|
"is_async": True,
|
|
"param_count": expected_param_count,
|
|
"has_model": True,
|
|
"has_config": True,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_initialize_invalid_unexpected_kwarg_raises_type_error():
|
|
with pytest.raises(TypeError):
|
|
await initialize(this_is_not_a_real_kwarg=True)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_initialize_from_worker_config_drops_legacy_telemetry_key():
|
|
with pytest.raises(TypeError):
|
|
await initialize(telemetry=True)
|
|
await initialize_from_worker_config({"telemetry": True, "request_timeout": 77})
|
|
assert ps.user_request_timeout == 77
|
|
with pytest.raises(TypeError):
|
|
await initialize_from_worker_config({"this_is_not_a_real_kwarg": True})
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# load_from_azure_key_vault
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_load_from_azure_key_vault_disabled_no_side_effect(monkeypatch):
|
|
import litellm
|
|
|
|
sentinel_secret_mgr = object()
|
|
monkeypatch.setattr(litellm, "secret_manager_client", sentinel_secret_mgr, raising=False)
|
|
|
|
result = load_from_azure_key_vault(use_azure_key_vault=False)
|
|
|
|
observed = {
|
|
"return_value": result,
|
|
"secret_manager_unchanged": litellm.secret_manager_client is sentinel_secret_mgr,
|
|
"called_with": False,
|
|
}
|
|
assert normalize(observed) == {
|
|
"return_value": None,
|
|
"secret_manager_unchanged": True,
|
|
"called_with": False,
|
|
}
|
|
|
|
|
|
def test_load_from_azure_key_vault_missing_uri_failure_is_swallowed(monkeypatch):
|
|
"""Enabled but AZURE_KEY_VAULT_URI unset / azure libs likely unavailable —
|
|
function catches Exception and does not raise."""
|
|
monkeypatch.delenv("AZURE_KEY_VAULT_URI", raising=False)
|
|
|
|
result = load_from_azure_key_vault(use_azure_key_vault=True)
|
|
assert result is None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# cost_tracking
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_cost_tracking_adds_db_and_shadow_eval_callbacks_when_prisma_set(monkeypatch):
|
|
import litellm
|
|
from litellm.integrations.shadow_eval_logger import ShadowEvalLogger
|
|
|
|
fake_prisma = MagicMock()
|
|
monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
|
|
monkeypatch.setattr(litellm, "callbacks", [], raising=False)
|
|
monkeypatch.setattr(litellm, "_async_success_callback", [], raising=False)
|
|
|
|
before_callbacks = len(litellm.callbacks)
|
|
before_async = len(litellm._async_success_callback)
|
|
|
|
cost_tracking()
|
|
cost_tracking()
|
|
|
|
observed = {
|
|
"added_to_callbacks": len(litellm.callbacks) - before_callbacks,
|
|
"added_to_async_success": len(litellm._async_success_callback) - before_async,
|
|
"shadow_eval_loggers": sum(isinstance(cb, ShadowEvalLogger) for cb in litellm.callbacks),
|
|
"prisma_was_set": True,
|
|
}
|
|
assert normalize(observed) == {
|
|
"added_to_callbacks": 2,
|
|
"added_to_async_success": 1,
|
|
"shadow_eval_loggers": 1,
|
|
"prisma_was_set": True,
|
|
}
|
|
|
|
|
|
def test_cost_tracking_no_op_when_prisma_missing(monkeypatch):
|
|
"""Without a prisma_client cost_tracking is a no-op — not an error."""
|
|
import litellm
|
|
|
|
monkeypatch.setattr(ps, "prisma_client", None, raising=False)
|
|
monkeypatch.setattr(litellm, "callbacks", [], raising=False)
|
|
monkeypatch.setattr(litellm, "_async_success_callback", [], raising=False)
|
|
|
|
cost_tracking()
|
|
|
|
assert litellm.callbacks == []
|
|
assert litellm._async_success_callback == []
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _resolve_typed_dict_type
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _SampleTD(TypedDict):
|
|
a: int
|
|
b: str
|
|
|
|
|
|
def test_resolve_typed_dict_type_finds_class_in_optional():
|
|
typ = Optional[_SampleTD]
|
|
result = _resolve_typed_dict_type(typ)
|
|
|
|
observed = {
|
|
"input_repr": "Optional[_SampleTD]",
|
|
"result_is_sample_td": result is _SampleTD,
|
|
"result_is_class": isinstance(result, type),
|
|
}
|
|
assert normalize(observed) == {
|
|
"input_repr": "Optional[_SampleTD]",
|
|
"result_is_sample_td": True,
|
|
"result_is_class": True,
|
|
}
|
|
|
|
|
|
def test_resolve_typed_dict_type_invalid_plain_type_returns_none():
|
|
"""A non-TypedDict, non-Union input returns None — not an error."""
|
|
assert _resolve_typed_dict_type(int) is None
|
|
assert _resolve_typed_dict_type(str) is None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _resolve_pydantic_type
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _SampleModelA(BaseModel):
|
|
x: int
|
|
|
|
|
|
class _SampleModelB(BaseModel):
|
|
y: str
|
|
|
|
|
|
def test_resolve_pydantic_type_extracts_non_none_args_from_union():
|
|
typ = Union[_SampleModelA, _SampleModelB, None]
|
|
result = _resolve_pydantic_type(typ)
|
|
|
|
observed = {
|
|
"result_type": type(result).__name__,
|
|
"result_len": len(result),
|
|
"contains_a": _SampleModelA in result,
|
|
"contains_b": _SampleModelB in result,
|
|
}
|
|
assert normalize(observed) == {
|
|
"result_type": "list",
|
|
"result_len": 2,
|
|
"contains_a": True,
|
|
"contains_b": True,
|
|
}
|
|
|
|
|
|
def test_resolve_pydantic_type_invalid_non_union_non_model_returns_empty():
|
|
"""When given a non-Union and non-BaseModel input the function returns [].
|
|
|
|
This is the silent-empty fallback path — error-ish by behavior."""
|
|
result = _resolve_pydantic_type(int)
|
|
assert result == []
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# get_litellm_model_info
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_get_litellm_model_info_uses_base_model_for_lookup(monkeypatch):
|
|
import litellm
|
|
|
|
expected_info = {"max_tokens": 8192, "input_cost_per_token": 0.00003}
|
|
fake_get = MagicMock(return_value=expected_info)
|
|
monkeypatch.setattr(litellm, "get_model_info", fake_get, raising=False)
|
|
|
|
model = {
|
|
"model_info": {"base_model": "gpt-4"},
|
|
"litellm_params": {"model": "azure/my-deployment"},
|
|
}
|
|
result = get_litellm_model_info(model=model)
|
|
|
|
observed = {
|
|
"called_arg": (
|
|
fake_get.call_args.args[0] if fake_get.call_args.args else fake_get.call_args.kwargs.get("model")
|
|
),
|
|
"returned_max_tokens": result.get("max_tokens"),
|
|
"returned_cost": result.get("input_cost_per_token"),
|
|
}
|
|
assert normalize(observed) == {
|
|
"called_arg": "gpt-4",
|
|
"returned_max_tokens": 8192,
|
|
"returned_cost": 0.00003,
|
|
}
|
|
|
|
|
|
def test_get_litellm_model_info_invalid_empty_dict_returns_empty():
|
|
"""Empty input means model_to_lookup is None — internal exception is caught
|
|
and the function returns {}."""
|
|
result = get_litellm_model_info(model={})
|
|
assert result == {}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# run_ollama_serve
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_run_ollama_serve_invokes_subprocess_popen(monkeypatch):
|
|
fake_popen = MagicMock()
|
|
monkeypatch.setattr(ps.subprocess, "Popen", fake_popen)
|
|
|
|
run_ollama_serve()
|
|
|
|
args, kwargs = fake_popen.call_args
|
|
observed = {
|
|
"popen_called": fake_popen.call_count == 1,
|
|
"command": args[0] if args else kwargs.get("args"),
|
|
"has_stdout_kw": "stdout" in kwargs,
|
|
"has_stderr_kw": "stderr" in kwargs,
|
|
}
|
|
assert normalize(observed) == {
|
|
"popen_called": True,
|
|
"command": ["ollama", "serve"],
|
|
"has_stdout_kw": True,
|
|
"has_stderr_kw": True,
|
|
}
|
|
|
|
|
|
def test_run_ollama_serve_popen_failure_is_swallowed(monkeypatch):
|
|
"""Popen raising OSError must NOT propagate — function logs and returns."""
|
|
monkeypatch.setattr(ps.subprocess, "Popen", MagicMock(side_effect=OSError("no ollama binary")))
|
|
|
|
result = run_ollama_serve()
|
|
assert result is None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# proxy_startup_event
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_startup_event_is_async_context_manager_with_expected_signature():
|
|
"""proxy_startup_event is the FastAPI lifespan. Verify its surface without
|
|
actually running the heavy init path (DB, Router, OTEL, etc.)."""
|
|
sig = inspect.signature(proxy_startup_event)
|
|
wrapped = getattr(proxy_startup_event, "__wrapped__", None)
|
|
observed = {
|
|
"param_count": len(sig.parameters),
|
|
"has_app_param": "app" in sig.parameters,
|
|
"wrapped_is_async": inspect.iscoroutinefunction(wrapped) or inspect.isasyncgenfunction(wrapped),
|
|
"has_asynccontextmanager_wrapper": wrapped is not None,
|
|
}
|
|
assert normalize(observed) == {
|
|
"param_count": 1,
|
|
"has_app_param": True,
|
|
"wrapped_is_async": True,
|
|
"has_asynccontextmanager_wrapper": True,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_startup_event_invalid_missing_app_arg_raises():
|
|
"""Calling the lifespan with no FastAPI app argument must fail."""
|
|
with pytest.raises(TypeError):
|
|
# Intentionally invoke the underlying async generator function with
|
|
# no arguments — the decorator preserves the missing-arg TypeError.
|
|
async with proxy_startup_event(): # type: ignore[call-arg]
|
|
pass
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_startup_event_prunes_dead_workers_live_gauges(tmp_path):
|
|
"""With PROMETHEUS_MULTIPROC_DIR set, a booting worker drops the live-gauge files of pids that no longer
|
|
exist, so a crashed worker's in-flight samples leave the aggregate as soon as its replacement starts."""
|
|
exited = subprocess.Popen(["true"])
|
|
assert exited.wait(timeout=30) == 0
|
|
stale = tmp_path / f"gauge_livesum_{exited.pid}.db"
|
|
stale.touch()
|
|
counter = tmp_path / f"counter_{exited.pid}.db"
|
|
counter.touch()
|
|
|
|
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
|
|
clean_env["PROMETHEUS_MULTIPROC_DIR"] = str(tmp_path)
|
|
with patch.dict(os.environ, clean_env, clear=True):
|
|
try:
|
|
async with proxy_startup_event(app=None):
|
|
pass
|
|
except Exception:
|
|
pass
|
|
|
|
assert not stale.exists()
|
|
assert counter.exists()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(("disable_model_info_refresh", "job_scheduled"), [(True, False), (False, True)])
|
|
async def test_proxy_startup_event_honors_disable_model_info_refresh(
|
|
disable_model_info_refresh: bool, job_scheduled: bool
|
|
) -> None:
|
|
"""``general_settings.disable_model_info_refresh: true`` keeps the proxy from polling every
|
|
OpenAI-compatible deployment's ``/v1/models`` in the background, so a proxy fronting a replay
|
|
fixture (or a metered upstream) makes only the calls its clients asked for."""
|
|
scheduler = AsyncIOScheduler()
|
|
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} | {
|
|
"LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY": "true"
|
|
}
|
|
with (
|
|
patch.dict(os.environ, clean_env, clear=True),
|
|
patch.object(ps, "scheduler", scheduler),
|
|
patch.dict(ps.general_settings, {"disable_model_info_refresh": disable_model_info_refresh}),
|
|
):
|
|
try:
|
|
async with proxy_startup_event(app=None):
|
|
job = scheduler.get_job("refresh_model_info")
|
|
finally:
|
|
if scheduler.running:
|
|
scheduler.shutdown(wait=False)
|
|
|
|
assert (job is not None) is job_scheduled, (
|
|
f"disable_model_info_refresh={disable_model_info_refresh} but refresh_model_info job is {job}"
|
|
)
|
|
|
|
|
|
def test_otel_global_provider_published_after_callback_init():
|
|
"""The OTel V2 global-provider publish must run after callback
|
|
initialization in ``proxy_startup_event``.
|
|
|
|
Regression for the orphan span: a preset (arize, langfuse, …) builds its
|
|
single folded logger during ``_initialize_startup_logging``. Publishing the
|
|
global ``TracerProvider`` before that ran found no logger and built a second
|
|
generic one whose provider became the global, so the FastAPI server span and
|
|
the preset's gen-ai spans exported through different providers and the LLM
|
|
span was orphaned. The publish (``publish_global_otel_v2_provider``) must
|
|
therefore appear after ``_initialize_startup_logging`` in the lifespan source.
|
|
"""
|
|
wrapped = getattr(proxy_startup_event, "__wrapped__", proxy_startup_event)
|
|
source = inspect.getsource(wrapped)
|
|
init_pos = source.find("_initialize_startup_logging(")
|
|
publish_pos = source.find("publish_global_otel_v2_provider(")
|
|
assert init_pos != -1, "callback init call not found in proxy_startup_event"
|
|
assert publish_pos != -1, "OTEL global publish not found in proxy_startup_event"
|
|
assert init_pos < publish_pos, (
|
|
"OTEL global provider is published before callbacks are initialized; a "
|
|
"preset logger will not exist yet and a second generic logger will own "
|
|
"the global provider, orphaning gen-ai spans"
|
|
)
|
|
|
|
|
|
def test_startup_warns_for_global_budget_without_database(caplog):
|
|
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
|
ProxyStartupEvent._warn_budget_without_db(max_budget=100.0, prisma_client=None)
|
|
|
|
assert "litellm.max_budget=100.0" in caplog.text
|
|
assert "will NOT be enforced" in caplog.text
|
|
assert "requests will never be blocked" in caplog.text
|
|
|
|
|
|
def test_startup_does_not_warn_for_global_budget_with_database(caplog):
|
|
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
|
ProxyStartupEvent._warn_budget_without_db(max_budget=100.0, prisma_client=MagicMock())
|
|
|
|
assert "litellm.max_budget" not in caplog.text
|
|
|
|
|
|
@pytest.mark.parametrize("max_budget", [0, None])
|
|
def test_startup_does_not_warn_without_global_budget(caplog, max_budget):
|
|
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
|
ProxyStartupEvent._warn_budget_without_db(max_budget=max_budget, prisma_client=None)
|
|
|
|
assert "litellm.max_budget" not in caplog.text
|
|
|
|
|
|
def test_proxy_startup_event_warns_for_global_budget_without_database():
|
|
"""Pin the lifespan call that prevents silent DB-less budgets.
|
|
|
|
The call must follow Prisma setup so DB-backed deployments do not false-positive.
|
|
Direct ``_warn_budget_without_db`` tests cover the warning behavior itself.
|
|
"""
|
|
wrapped = getattr(proxy_startup_event, "__wrapped__", proxy_startup_event)
|
|
source = inspect.getsource(wrapped)
|
|
budget_check_pos = source.find("if prisma_client is not None and litellm.max_budget > 0:")
|
|
warn_pos = source.find("_warn_budget_without_db(")
|
|
next_startup_section_pos = source.find(
|
|
"await ProxyStartupEvent.initialize_scheduled_background_jobs(",
|
|
budget_check_pos,
|
|
)
|
|
|
|
assert budget_check_pos != -1, "global budget startup block not found"
|
|
assert warn_pos != -1, "DB-less budget warning call not found"
|
|
assert next_startup_section_pos != -1, "startup section after budget block not found"
|
|
assert budget_check_pos < warn_pos < next_startup_section_pos, (
|
|
"DB-less budget warning must run after Prisma setup and the DB-backed budget block"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tuning_baseline_v3_is_created_alongside_the_legacy_row():
|
|
from litellm.router_utils.auto_router_tuning_baseline import DEFAULT_TUNING_FINGERPRINT
|
|
|
|
prisma_client = MagicMock()
|
|
prisma_client.db.litellm_config.find_unique = AsyncMock(return_value=None)
|
|
prisma_client.db.litellm_config.create = AsyncMock()
|
|
deployment = {
|
|
"model_name": "a",
|
|
"litellm_params": {"model": "auto_router/complexity_router", "complexity_router_config": {}},
|
|
}
|
|
|
|
result = await ProxyStartupEvent._load_heuristic_v1_tuning_baselines(prisma_client, [deployment])
|
|
|
|
assert result == {'yaml:["a",[]]': DEFAULT_TUNING_FINGERPRINT}
|
|
assert prisma_client.db.litellm_config.create.await_args.kwargs["data"] == {
|
|
"param_name": "auto_router_tuning_baseline_v3",
|
|
"param_value": json.dumps(dict(result)),
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_scorer_baseline_upgrade_preserves_existing_routers_and_is_not_refreshed_on_restart():
|
|
from litellm.router_utils.auto_router_tuning_baseline import mutable_tuned_identities, snapshot_tuning_baselines
|
|
|
|
deployments = [
|
|
{
|
|
"model_name": name,
|
|
"litellm_params": {
|
|
"model": "auto_router/complexity_router",
|
|
"complexity_router_config": {"tiers": {"SIMPLE": name}, "code_keywords": [name]},
|
|
},
|
|
}
|
|
for name in ("a", "b")
|
|
]
|
|
prisma_client = MagicMock()
|
|
prisma_client.db.litellm_config.find_unique = AsyncMock(
|
|
side_effect=lambda where: (
|
|
MagicMock(param_value='{"legacy-router":"old-combined-hash"}')
|
|
if where["param_name"] == "auto_router_tuning_baseline_v2"
|
|
else None
|
|
)
|
|
)
|
|
prisma_client.db.litellm_config.create = AsyncMock()
|
|
|
|
baseline = await ProxyStartupEvent._load_heuristic_v1_tuning_baselines(prisma_client, deployments)
|
|
|
|
assert baseline == snapshot_tuning_baselines(deployments)
|
|
assert mutable_tuned_identities(deployments, baseline) == frozenset()
|
|
prisma_client.db.litellm_config.create.assert_awaited_once_with(
|
|
data={"param_name": "auto_router_tuning_baseline_v3", "param_value": json.dumps(dict(baseline))}
|
|
)
|
|
prisma_client.db.litellm_config.find_unique.side_effect = None
|
|
prisma_client.db.litellm_config.find_unique.return_value = MagicMock(param_value=json.dumps(dict(baseline)))
|
|
changed = [
|
|
{
|
|
"model_name": "a",
|
|
"litellm_params": {
|
|
"model": "auto_router/complexity_router",
|
|
"complexity_router_config": {"tiers": {"SIMPLE": "different-model"}, "code_keywords": ["new-rule"]},
|
|
},
|
|
}
|
|
]
|
|
|
|
reloaded = await ProxyStartupEvent._load_heuristic_v1_tuning_baselines(prisma_client, changed)
|
|
|
|
assert reloaded == baseline
|
|
assert mutable_tuned_identities(changed, reloaded) == frozenset({'yaml:["a",[]]'})
|
|
prisma_client.db.litellm_config.create.assert_awaited_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tuning_baseline_waits_for_a_complete_db_model_census(monkeypatch):
|
|
prisma_client = MagicMock()
|
|
monkeypatch.setattr(ps.proxy_config, "_get_models_from_db", AsyncMock(return_value=None))
|
|
|
|
result = await ProxyStartupEvent.enforce_heuristic_v1_tuning_baseline(
|
|
prisma_client=prisma_client,
|
|
llm_router=None,
|
|
limit=1,
|
|
)
|
|
|
|
assert result is None
|
|
prisma_client.db.litellm_config.find_unique.assert_not_called()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _initialize_slack_alerting_jobs — spend-report pod locking (issue #14809)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
SlackAlertingJobs = dict[str, Callable[[], Awaitable[None]]]
|
|
|
|
|
|
def _make_slack_alerting_proxy_logging(acquire_lock_result: bool | None) -> MagicMock:
|
|
proxy_logging_obj = MagicMock()
|
|
proxy_logging_obj.slack_alerting_instance.alerting = ["slack"]
|
|
proxy_logging_obj.slack_alerting_instance.send_weekly_spend_report = AsyncMock()
|
|
proxy_logging_obj.slack_alerting_instance.send_monthly_spend_report = AsyncMock()
|
|
proxy_logging_obj.slack_alerting_instance.send_fallback_stats_from_prometheus = AsyncMock()
|
|
pod_lock_manager = proxy_logging_obj.db_spend_update_writer.pod_lock_manager
|
|
pod_lock_manager.acquire_lock = AsyncMock(return_value=acquire_lock_result)
|
|
pod_lock_manager.release_lock = AsyncMock()
|
|
return proxy_logging_obj
|
|
|
|
|
|
async def _init_slack_alerting_jobs(
|
|
acquire_lock_result: bool | None,
|
|
spend_report_frequency: str = "7d",
|
|
) -> tuple[SlackAlertingJobs, MagicMock]:
|
|
scheduler = MagicMock()
|
|
proxy_logging_obj = _make_slack_alerting_proxy_logging(acquire_lock_result)
|
|
|
|
await ProxyStartupEvent._initialize_slack_alerting_jobs(
|
|
scheduler=scheduler,
|
|
general_settings={"spend_report_frequency": spend_report_frequency},
|
|
proxy_logging_obj=proxy_logging_obj,
|
|
prisma_client=MagicMock(),
|
|
)
|
|
|
|
jobs = {call.kwargs["id"]: call.args[0] for call in scheduler.add_job.call_args_list}
|
|
return jobs, proxy_logging_obj
|
|
|
|
|
|
@pytest.mark.parametrize("spend_report_frequency", ["0d", "-1d", "7h"])
|
|
@pytest.mark.asyncio
|
|
async def test_initialize_slack_alerting_jobs_invalid_frequency_raises(spend_report_frequency: str):
|
|
"""A non-positive window used to become an every-second APScheduler interval, and now also
|
|
computes a negative lock TTL that expires instantly and suppresses the report for good.
|
|
match= is load-bearing: drop the guard and "-1d" still raises, but from duration_in_seconds."""
|
|
with pytest.raises(ValueError, match="positive number of days"):
|
|
await _init_slack_alerting_jobs(
|
|
acquire_lock_result=True,
|
|
spend_report_frequency=spend_report_frequency,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_weekly_spend_report_skipped_when_another_pod_holds_the_lock():
|
|
"""regression: issue #14809 - every pod ran its own weekly spend report job."""
|
|
jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=False)
|
|
|
|
await jobs["weekly_spend_report_job"]()
|
|
|
|
proxy_logging_obj.slack_alerting_instance.send_weekly_spend_report.assert_not_awaited()
|
|
proxy_logging_obj.db_spend_update_writer.pod_lock_manager.acquire_lock.assert_awaited_once_with(
|
|
cronjob_id="weekly_spend_report_job",
|
|
ttl=7 * 86400 - 3600,
|
|
allow_reentrant=False,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("acquire_lock_result", [True, None])
|
|
@pytest.mark.asyncio
|
|
async def test_weekly_spend_report_sent_when_the_lock_is_free_or_absent(acquire_lock_result):
|
|
"""None means redis isn't configured; a single-pod deploy must still report."""
|
|
jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=acquire_lock_result)
|
|
|
|
await jobs["weekly_spend_report_job"]()
|
|
|
|
proxy_logging_obj.slack_alerting_instance.send_weekly_spend_report.assert_awaited_once_with("7d")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_weekly_spend_report_lock_ttl_tracks_the_configured_window():
|
|
"""TTL is the window less an hour: long enough that no second pod re-sends inside the
|
|
window, short enough that the lock is gone before the next one opens. A fixed TTL would
|
|
break one end or the other as soon as spend_report_frequency changes."""
|
|
jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=True, spend_report_frequency="1d")
|
|
|
|
await jobs["weekly_spend_report_job"]()
|
|
|
|
proxy_logging_obj.db_spend_update_writer.pod_lock_manager.acquire_lock.assert_awaited_once_with(
|
|
cronjob_id="weekly_spend_report_job",
|
|
ttl=86400 - 3600,
|
|
allow_reentrant=False,
|
|
)
|
|
proxy_logging_obj.slack_alerting_instance.send_weekly_spend_report.assert_awaited_once_with("1d")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_monthly_spend_report_skipped_when_another_pod_holds_the_lock():
|
|
jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=False)
|
|
|
|
await jobs["monthly_spend_report_job"]()
|
|
|
|
proxy_logging_obj.slack_alerting_instance.send_monthly_spend_report.assert_not_awaited()
|
|
proxy_logging_obj.db_spend_update_writer.pod_lock_manager.acquire_lock.assert_awaited_once_with(
|
|
cronjob_id="monthly_spend_report_job",
|
|
ttl=3600,
|
|
allow_reentrant=False,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("acquire_lock_result", [True, None])
|
|
@pytest.mark.asyncio
|
|
async def test_monthly_spend_report_sent_when_the_lock_is_free_or_absent(acquire_lock_result):
|
|
jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=acquire_lock_result)
|
|
|
|
await jobs["monthly_spend_report_job"]()
|
|
|
|
proxy_logging_obj.slack_alerting_instance.send_monthly_spend_report.assert_awaited_once_with()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_spend_report_locks_are_never_released():
|
|
"""The lock is a per-window marker, not a mutex: releasing it lets the next pod re-send."""
|
|
jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=True)
|
|
|
|
await jobs["weekly_spend_report_job"]()
|
|
await jobs["monthly_spend_report_job"]()
|
|
|
|
proxy_logging_obj.db_spend_update_writer.pod_lock_manager.release_lock.assert_not_awaited()
|
|
|
|
|
|
def _init_daily_global_spend_reconcile_job() -> tuple[AsyncIOScheduler, MagicMock, MagicMock]:
|
|
scheduler = AsyncIOScheduler()
|
|
proxy_logging_obj = MagicMock()
|
|
proxy_logging_obj.alerting_handler = AsyncMock()
|
|
prisma_client = MagicMock()
|
|
ProxyStartupEvent._initialize_daily_global_spend_reconcile_job(
|
|
scheduler=scheduler,
|
|
proxy_logging_obj=proxy_logging_obj,
|
|
prisma_client=prisma_client,
|
|
)
|
|
return scheduler, proxy_logging_obj, prisma_client
|
|
|
|
|
|
def test_daily_global_spend_reconcile_job_is_scheduled_nightly_with_an_immediate_catch_up_run():
|
|
"""Startup schedules the LiteLLM_DailyGlobalSpend backfill a couple of minutes out, so a
|
|
fresh deploy switches usage reads to the global table without waiting for the nightly
|
|
run, and after that it fires once a day at 00:30 UTC, when the previous UTC day is closed."""
|
|
from datetime import datetime, timedelta, timezone
|
|
|
|
from litellm.constants import DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID
|
|
|
|
scheduler, _, _ = _init_daily_global_spend_reconcile_job()
|
|
job = scheduler.get_job(DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID)
|
|
assert job is not None
|
|
|
|
assert timedelta(0) < job.next_run_time - datetime.now(timezone.utc) <= timedelta(minutes=2)
|
|
after_catch_up = datetime(2026, 9, 16, 12, 0, tzinfo=timezone.utc)
|
|
assert job.trigger.get_next_fire_time(None, after_catch_up) == datetime(2026, 9, 17, 0, 30, tzinfo=timezone.utc)
|
|
just_after_a_run = datetime(2026, 9, 17, 0, 30, 1, tzinfo=timezone.utc)
|
|
assert job.trigger.get_next_fire_time(None, just_after_a_run) == datetime(2026, 9, 18, 0, 30, tzinfo=timezone.utc)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_daily_global_spend_reconcile_job_runs_under_the_pod_lock_and_alerts_through_the_proxy(monkeypatch):
|
|
from litellm.constants import DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID
|
|
|
|
scheduler, proxy_logging_obj, prisma_client = _init_daily_global_spend_reconcile_job()
|
|
run = AsyncMock()
|
|
monkeypatch.setattr(ps, "run_scheduled_daily_global_spend_reconcile", run)
|
|
|
|
await scheduler.get_job(DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID).func()
|
|
|
|
run.assert_awaited_once()
|
|
assert run.await_args.args == (prisma_client,)
|
|
assert run.await_args.kwargs["pod_lock_manager"] is proxy_logging_obj.db_spend_update_writer.pod_lock_manager
|
|
await run.await_args.kwargs["alert"]("day 2026-09-01 failed")
|
|
proxy_logging_obj.alerting_handler.assert_awaited_once()
|
|
assert proxy_logging_obj.alerting_handler.await_args.kwargs["message"] == "day 2026-09-01 failed"
|
|
assert proxy_logging_obj.alerting_handler.await_args.kwargs["level"] == "High"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_prometheus_fallback_stats_job_skipped_when_another_pod_holds_the_lock(monkeypatch):
|
|
"""The boot-time send goes through the same gate, so a losing pod sends nothing at all:
|
|
startup and the scheduled job both stay at zero."""
|
|
monkeypatch.setenv("PROMETHEUS_URL", "http://prometheus.invalid")
|
|
jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=False)
|
|
send_fallback_stats = proxy_logging_obj.slack_alerting_instance.send_fallback_stats_from_prometheus
|
|
assert send_fallback_stats.await_count == 0
|
|
|
|
await jobs["prometheus_fallback_stats_job"]()
|
|
|
|
assert send_fallback_stats.await_count == 0
|
|
proxy_logging_obj.db_spend_update_writer.pod_lock_manager.acquire_lock.assert_awaited_with(
|
|
cronjob_id="prometheus_fallback_stats_job",
|
|
ttl=3600,
|
|
allow_reentrant=False,
|
|
)
|
|
assert proxy_logging_obj.db_spend_update_writer.pod_lock_manager.acquire_lock.await_count == 2
|
|
|
|
|
|
@pytest.mark.parametrize("acquire_lock_result", [True, None])
|
|
@pytest.mark.asyncio
|
|
async def test_prometheus_fallback_stats_job_runs_when_the_lock_is_free_or_absent(monkeypatch, acquire_lock_result):
|
|
monkeypatch.setenv("PROMETHEUS_URL", "http://prometheus.invalid")
|
|
jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=acquire_lock_result)
|
|
send_fallback_stats = proxy_logging_obj.slack_alerting_instance.send_fallback_stats_from_prometheus
|
|
assert send_fallback_stats.await_count == 1
|
|
|
|
await jobs["prometheus_fallback_stats_job"]()
|
|
|
|
assert send_fallback_stats.await_count == 2
|