mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-05 08:07:05 +00:00
b7749f67f1
* warn at startup when a proxy-wide budget is set but no DB is connected litellm.max_budget is only enforced via DB-loaded global spend, so a DB-less proxy silently ignores it. Log a one-time startup warning. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): inject max_budget into DB-less budget warning Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): cover DB-less budget warning startup call Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): pin DB-less budget warning call site Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): stabilize budget warning call-site pin Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: tin <tin@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
779 lines
28 KiB
Python
779 lines
28 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
|
|
from typing import List, Optional, Union
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
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,
|
|
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)
|
|
|
|
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,
|
|
}
|
|
assert normalize(observed) == {
|
|
"master_key": None,
|
|
"user_config_file_path": None,
|
|
"user_custom_auth": None,
|
|
"health_check_interval": None,
|
|
"prisma_client": 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_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()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _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],
|
|
"telemetry": 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["telemetry"] 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 = 17
|
|
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)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 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_two_callbacks_when_prisma_set(monkeypatch):
|
|
import litellm
|
|
|
|
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()
|
|
|
|
observed = {
|
|
"added_to_callbacks": len(litellm.callbacks) - before_callbacks,
|
|
"added_to_async_success": len(litellm._async_success_callback) - before_async,
|
|
"prisma_was_set": True,
|
|
}
|
|
assert normalize(observed) == {
|
|
"added_to_callbacks": 1,
|
|
"added_to_async_success": 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
|
|
|
|
|
|
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"
|
|
)
|