mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-25 05:54:53 +00:00
c80f1c2ec8
* fix(responses): write one spend-log row per native WebSocket turn Each response.create turn on a native Responses WebSocket connection now logs its own call as it finishes, with its own usage and the resp_ id the client saw as request_id. A failed turn is booked alone and no longer flips the other turns. * fix(responses): bill the usage a failed native WebSocket turn reports * fix(proxy): pin the limiter slot owner to the proxy call id and refund unsettled reservations on disconnect The v3 limiter stashed whichever litellm_call_id the request body carried as the slot owner, so a client-supplied id on a pass-through body left the proxy's own success and failure events unable to release the slot. The owner is now the proxy's logging call id. A WebSocket connection whose turns log under their own ids never settles the TPM, ITPM, or OTPM estimate the connection reserved, so the disconnect release now refunds any reservation no callback settled, once. * fix(proxy): bill failed WebSocket turns and refund only the reserved window on disconnect A failed turn's reported usage now counts toward the key's TPM and the project ITPM/OTPM scopes the same way a success does, and closing the socket hands back only the estimate still unused in the window that holds it, so a window that already rolled is left alone. An id-less error event is booked against the newest pending turn, the one the backend rejected at parse time, and a turn that nests its input under response logs its user message. A connection's one budget reservation goes to the first turn that settles it, every later turn charging its cost directly. * test(proxy): expect the settling call id on a pass-through budget reservation * fix(proxy): charge successful Responses WebSocket turns to TPM and ITPM/OTPM A native Responses WebSocket turn reaches the rate limiter's success callback with its raw event list as the response and no combined_usage_object, so the turn's tokens never reached the key's TPM window or the project's input and output token limits. Read the usage from the turn's events when the call type is a Responses WebSocket turn. * fix(proxy): release a rejected pass-through's slot and count each websocket turn once toward project quotas The limiter's failure hook looked up the parallel slot by the request body's litellm_call_id, while the slot belongs to the proxy's logging call id, so a pass-through request a pre-call hook rejected held its slot for an hour when its body carried an id. Both now resolve the id the same way. Each response.create frame already charges its estimate to project ITPM/OTPM, so the turn's own usage no longer adds a second count. The cost callback's failure hook read the body's litellm_call_id through a strict model and raised on a non-string value, which dropped the failure row. It now reads a string id or none. * fix(proxy): keep per-turn WebSocket limits, request bodies and collector reservations right Free a rejected request's parallel slot through the request it runs in, and skip native Responses WebSocket connections, whose disconnect frees the slot. Count WebSocket turns toward the connection's tag TPM. In collector mode, claim the shared budget reservation before the spend event is built, so only the first turn settles it. Store each turn's own frame as its request body, and send an id-less error to the oldest waiting turn unless an earlier turn is streaming. * fix(proxy): store a failed WebSocket turn's own request and count tag TPM only for calls that carry the tag * fix(responses): log the client's turn frame, keep the payload id and reservation shape for non-WebSocket calls * fix(proxy): keep the legacy limiter's WebSocket slot until the socket closes * fix(router): count native Responses WebSocket turns toward deployment limits and budgets * fix(router): re-check deployment budget and RPM on each Responses WebSocket turn * test(router): list the WebSocket turn admission method as covered through its partial * fix(responses): give a refused WebSocket turn's error frame a 429 status so Codex ends the turn * fix(router): leave ITPM/OTPM to socket open and record a refused WebSocket turn as the socket's failure * fix(proxy): log the routed model group on a WebSocket turn's failure row * fix(router): count WebSocket turn tokens toward TPM and re-run routing strategy checks on each turn * test(integration): cover per-turn admission on Responses WebSockets Adds integration cells for budget, RPM, TPM, a mid-socket TPM edit, usage-based-routing-v2, the project quota frame, and the open-time failure frame on the native and managed relays, plus HTTP controls for TPM counting and the failure row model_group * fix(responses): count WebSocket turn tokens once and end a socket whose turn check fails A native Responses WebSocket turn is now logged as an async call, so a proxy with a sync success callback such as the response cache no longer runs each CustomLogger's sync success hook on top of the async one. That second run counted every turn's tokens twice toward deployment TPM. A turn whose admission check raises an unexpected error now gets a 500 error frame, the socket closes with 1011, and the error becomes the socket's failure, on both the native and the managed relay, instead of the socket dropping with no frame and no failure row. A refused turn's failure row now carries the request's tags from the router's metadata, the same way it already carries the model group. * test(integration): cover a crashed WebSocket turn check and failure-row tags A turn whose admission check raises gets a 500 server_error frame, the socket closes with 1011 and its one failure row carries the error, on both relays. The refused turn's failure row keeps the caller's tags, and HTTP failure rows on chat, responses and messages keep them too. * fix(responses): keep a WebSocket open after a crashed turn check and keep only string tags on failure rows A turn whose admission check raises an unexpected error still gets a 500 error frame, but the socket now stays open, like a refused turn. Closing it cut any turn still streaming on the same socket, which a client that sends its next response.create early would lose. A refused turn's failure row now takes the router's tags only when every tag is a string, so the fallback never copies a caller's unvalidated tags onto the row. * test(integration): cover a turn still streaming when the next WebSocket turn's check crashes A native turn still streaming when the next turn's admission check raises now has its own cell. It must complete and be billed at the turn's price. The crashed-check cell also keeps the fault set for one more turn, so a relay that stopped checking turns after a crash goes red. * fix(responses): run per-turn guardrails before the deployment's turn check on WebSockets A turn after the first now runs the project quota check, then the proxy's pre_call guardrails, then the deployment's turn admission, on both the native and the managed relay, the order HTTP calls use. A turn a guardrail blocks gets the guardrail's error even past the budget and no longer counts toward the deployment's rpm or runs its pre-call checks. * fix(responses): bill a WebSocket turn's guardrail cost when the deployment refuses it A turn the per-turn guardrail passes and the deployment then refuses now writes its own failed spend row with that turn's guardrail information, on both relays, the way a guardrail-blocked frame and an HTTP refusal already do * fix(responses): bill a refused WebSocket turn on one row only A refused turn whose guardrails recorded a cost is already billed on its own row, so it no longer also becomes the connection's close-time failure row * fix(responses): bill a refused WebSocket turn's guardrail cost on the socket's failure row The turns a deployment refuses after their guardrails passed now add their guardrail records to the socket's close-time failure row instead of a row of their own, so that row's failure hooks, model group and call type cover them * fix(responses): bill a WebSocket's refused turns when the session limit closes it * fix(responses): bill refused WebSocket turns when the session limit cancels the relay * fix(responses): free the WebSocket slot at the session limit without waiting on the failure row * fix(responses): keep the session limit's second cancel so a relay that swallows the first frees the slot * fix(responses): free the WebSocket slot before the close-time failure row and drain pending rows at shutdown * fix(proxy): commit spend before waiting on pending logging tasks at shutdown and again after * fix(responses): start the close-time failure row before the slot release so a late cancel cannot drop it * fix(responses): build the close-time failure row inside its task so a bad guardrail record cannot skip the crash close * fix(proxy): always commit spend again after the shutdown wait for pending logging tasks --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
3815 lines
133 KiB
Python
3815 lines
133 KiB
Python
import os
|
|
import sys
|
|
import traceback
|
|
from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping
|
|
from contextlib import asynccontextmanager
|
|
from typing import Final, Literal, NoReturn, cast
|
|
from unittest import mock
|
|
|
|
from dotenv import load_dotenv
|
|
|
|
from tests._master_key import MASTER_KEY
|
|
|
|
import litellm.proxy
|
|
import litellm.proxy.proxy_server
|
|
|
|
load_dotenv()
|
|
import io
|
|
import json
|
|
|
|
# this file is to test litellm/proxy
|
|
|
|
import asyncio
|
|
import logging
|
|
|
|
import pytest
|
|
|
|
import litellm
|
|
from litellm import RateLimitError, Timeout, completion, completion_cost, embedding
|
|
|
|
# Configure logging
|
|
logging.basicConfig(
|
|
level=logging.DEBUG, # Set the desired logging level
|
|
format="%(asctime)s - %(levelname)s - %(message)s",
|
|
)
|
|
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
from fastapi import FastAPI, HTTPException, WebSocket
|
|
|
|
# test /chat/completion request to the proxy
|
|
from fastapi.testclient import TestClient
|
|
|
|
from litellm.integrations.custom_logger import CustomLogger
|
|
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
|
from litellm.proxy.proxy_server import ( # Replace with the actual module where your FastAPI router is defined
|
|
app,
|
|
initialize,
|
|
save_worker_config,
|
|
)
|
|
from litellm.proxy.utils import ProxyLogging
|
|
|
|
@pytest.fixture
|
|
def admin_mcp_proxy(monkeypatch: pytest.MonkeyPatch) -> FastAPI:
|
|
from litellm.proxy import proxy_server
|
|
from litellm.proxy.auth.litellm_license import LicenseCheck
|
|
|
|
monkeypatch.setenv("LITELLM_ENABLE_ADMIN_MCP", "true")
|
|
monkeypatch.setenv("LITELLM_MASTER_KEY", "sk-" + "1234567890abcdef" * 4)
|
|
for name in ("WORKER_CONFIG", "CONFIG_FILE_PATH", "DATABASE_URL", "LITELLM_LICENSE"):
|
|
monkeypatch.delenv(name, raising=False)
|
|
monkeypatch.setattr(proxy_server, "_license_check", LicenseCheck())
|
|
monkeypatch.setattr(proxy_server, "premium_user", True)
|
|
monkeypatch.setattr(proxy_server, "prisma_client", None)
|
|
monkeypatch.setattr(proxy_server, "general_settings", {"disable_model_info_refresh": True})
|
|
monkeypatch.setattr(proxy_server, "scheduler", None)
|
|
return FastAPI(lifespan=proxy_server.proxy_startup_event)
|
|
|
|
|
|
@pytest.mark.skipif(sys.version_info < (3, 12), reason="Admin MCP requires Python 3.12+")
|
|
@pytest.mark.parametrize("failure_phase", ["startup", "serving", "shutdown", "cancelled", "license"])
|
|
async def test_admin_mcp_failure_still_closes_proxy_resources(
|
|
admin_mcp_proxy: FastAPI,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
failure_phase: Literal["startup", "serving", "shutdown", "cancelled", "license"],
|
|
) -> None:
|
|
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
|
from apscheduler.schedulers.base import STATE_PAUSED
|
|
from litellm_admin_mcp import server as connector_server
|
|
from litellm_admin_mcp.gateway import Gateway
|
|
|
|
from litellm.proxy import proxy_server
|
|
from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager
|
|
from litellm.proxy.shutdown.scheduled_jobs import AwaitableAsyncIOExecutor
|
|
|
|
monkeypatch.setattr(proxy_server, "premium_user", failure_phase != "license")
|
|
executor: Final = AwaitableAsyncIOExecutor()
|
|
scheduler: Final = AsyncIOScheduler(executors={"default": executor})
|
|
monkeypatch.setattr(proxy_server, "scheduler", scheduler)
|
|
monkeypatch.setattr(proxy_server, "scheduler_executor", executor)
|
|
scheduler.start()
|
|
|
|
@asynccontextmanager
|
|
async def failing_connector(_app: FastAPI) -> AsyncGenerator[None, None]:
|
|
if failure_phase == "startup":
|
|
raise RuntimeError("startup failed")
|
|
yield
|
|
assert scheduler.state == STATE_PAUSED
|
|
assert GracefulShutdownManager.is_shutting_down()
|
|
assert proxy_server.shared_aiohttp_session is not None
|
|
assert not proxy_server.shared_aiohttp_session.closed
|
|
if failure_phase == "shutdown":
|
|
raise RuntimeError("shutdown failed")
|
|
|
|
def connector_app(_gateway: Gateway) -> FastAPI:
|
|
return FastAPI(lifespan=failing_connector)
|
|
|
|
monkeypatch.setattr(connector_server, "create_http_app", connector_app)
|
|
expected_error: Final = (
|
|
HTTPException if failure_phase == "license"
|
|
else asyncio.CancelledError if failure_phase == "cancelled"
|
|
else RuntimeError
|
|
)
|
|
message: Final = "LITELLM_LICENSE" if failure_phase == "license" else f"{failure_phase} failed"
|
|
async def run_lifespan() -> None:
|
|
async with admin_mcp_proxy.router.lifespan_context(admin_mcp_proxy) as state:
|
|
assert state == {"tracing_receiver": None}
|
|
if failure_phase == "serving":
|
|
raise RuntimeError("serving failed")
|
|
if failure_phase == "cancelled":
|
|
raise asyncio.CancelledError("cancelled failed")
|
|
|
|
with pytest.raises(expected_error, match=message):
|
|
await run_lifespan()
|
|
|
|
assert proxy_server.shared_aiohttp_session is not None
|
|
assert proxy_server.shared_aiohttp_session.closed
|
|
assert proxy_server.master_key is None
|
|
assert all(route.name != "admin_mcp" for route in admin_mcp_proxy.routes)
|
|
|
|
|
|
@pytest.mark.skipif(sys.version_info < (3, 12), reason="Admin MCP requires Python 3.12+")
|
|
async def test_proxy_shutdown_drains_active_admin_tool_before_closing_connector(
|
|
admin_mcp_proxy: FastAPI, monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
import httpx2
|
|
|
|
from litellm.proxy.middleware.in_flight_requests_middleware import InFlightRequestsMiddleware
|
|
from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager
|
|
|
|
monkeypatch.setenv("LITELLM_ADMIN_TOOLS", "list_teams")
|
|
admin_mcp_proxy.add_middleware(InFlightRequestsMiddleware)
|
|
started: Final = asyncio.Event()
|
|
release: Final = asyncio.Event()
|
|
|
|
@admin_mcp_proxy.get("/user/info")
|
|
async def user_info() -> dict[str, object]:
|
|
return {"user_id": "admin", "user_info": {"user_id": "admin", "user_role": "proxy_admin"}}
|
|
|
|
@admin_mcp_proxy.get("/team/list", operation_id="list_team_team_list_get")
|
|
async def list_teams() -> dict[str, object]:
|
|
started.set()
|
|
await release.wait()
|
|
return {"teams": ["completed-before-shutdown"]}
|
|
|
|
async def complete_during_drain() -> None:
|
|
async with asyncio.timeout(5):
|
|
while not GracefulShutdownManager.is_shutting_down():
|
|
await asyncio.sleep(0)
|
|
release.set()
|
|
|
|
async with httpx2.AsyncClient(
|
|
transport=httpx2.ASGITransport(app=admin_mcp_proxy), base_url="http://localhost:4000"
|
|
) as client:
|
|
async with admin_mcp_proxy.router.lifespan_context(admin_mcp_proxy) as state:
|
|
assert state == {"tracing_receiver": None}
|
|
request: Final = asyncio.create_task(client.post(
|
|
"/admin/mcp",
|
|
headers={"Authorization": "Bearer admin", "Accept": "application/json, text/event-stream"},
|
|
json={"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "list_teams"}},
|
|
))
|
|
await asyncio.wait_for(started.wait(), timeout=5)
|
|
completion: Final = asyncio.create_task(complete_during_drain())
|
|
await asyncio.wait_for(completion, timeout=5)
|
|
response: Final = await asyncio.wait_for(request, timeout=5)
|
|
|
|
assert response.status_code == 200, response.text
|
|
assert response.json()["result"]["isError"] is False
|
|
assert json.loads(response.json()["result"]["content"][0]["text"]) == {
|
|
"teams": ["completed-before-shutdown"]
|
|
}
|
|
assert InFlightRequestsMiddleware.get_count() == 0
|
|
|
|
|
|
@pytest.mark.skipif(sys.version_info < (3, 12), reason="Admin MCP requires Python 3.12+")
|
|
async def test_proxy_shutdown_closes_admin_connector_when_drain_is_cancelled(
|
|
admin_mcp_proxy: FastAPI,
|
|
) -> None:
|
|
import httpx2
|
|
|
|
from litellm.proxy import proxy_server
|
|
from litellm.proxy.middleware.in_flight_requests_middleware import InFlightRequestsMiddleware
|
|
from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager
|
|
|
|
admin_mcp_proxy.add_middleware(InFlightRequestsMiddleware)
|
|
ready: Final = asyncio.Event()
|
|
shutdown: Final = asyncio.Event()
|
|
started: Final = asyncio.Event()
|
|
release: Final = asyncio.Event()
|
|
|
|
@admin_mcp_proxy.get("/hold")
|
|
async def hold_request() -> dict[str, bool]:
|
|
started.set()
|
|
await release.wait()
|
|
return {"complete": True}
|
|
|
|
async def serve() -> None:
|
|
async with admin_mcp_proxy.router.lifespan_context(admin_mcp_proxy):
|
|
ready.set()
|
|
await shutdown.wait()
|
|
|
|
async with httpx2.AsyncClient(
|
|
transport=httpx2.ASGITransport(app=admin_mcp_proxy), base_url="http://localhost:4000"
|
|
) as client:
|
|
serving: Final = asyncio.create_task(serve())
|
|
await asyncio.wait_for(ready.wait(), timeout=5)
|
|
assert any(route.name == "admin_mcp" for route in admin_mcp_proxy.routes)
|
|
request: Final = asyncio.create_task(client.get("/hold"))
|
|
try:
|
|
await asyncio.wait_for(started.wait(), timeout=5)
|
|
shutdown.set()
|
|
async with asyncio.timeout(5):
|
|
while not GracefulShutdownManager.is_shutting_down():
|
|
await asyncio.sleep(0)
|
|
serving.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await serving
|
|
assert all(route.name != "admin_mcp" for route in admin_mcp_proxy.routes)
|
|
assert proxy_server.shared_aiohttp_session is not None
|
|
assert proxy_server.shared_aiohttp_session.closed
|
|
finally:
|
|
release.set()
|
|
await asyncio.wait_for(request, timeout=5)
|
|
|
|
assert InFlightRequestsMiddleware.get_count() == 0
|
|
|
|
|
|
# Your bearer token
|
|
@pytest.mark.asyncio
|
|
async def test_shutdown_commits_spend_a_logging_task_queued_while_the_first_commit_was_writing():
|
|
from litellm.responses.streaming_iterator import spawn_logging_task
|
|
|
|
queued: list[str] = []
|
|
committed: list[str] = []
|
|
|
|
async def commit_spend() -> None:
|
|
committed.extend(queued)
|
|
queued.clear()
|
|
await asyncio.sleep(0.05)
|
|
|
|
async def failure_row() -> None:
|
|
await asyncio.sleep(0.01)
|
|
queued.append("failure row spend")
|
|
|
|
spawn_logging_task(asyncio.get_running_loop(), failure_row(), task_name="failure row")
|
|
|
|
await litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown(flush_spend=commit_spend)
|
|
|
|
assert committed == ["failure row spend"]
|
|
|
|
|
|
token = MASTER_KEY
|
|
|
|
headers = {"Authorization": f"Bearer {token}"}
|
|
|
|
example_completion_result = {
|
|
"choices": [
|
|
{
|
|
"message": {
|
|
"content": "Whispers of the wind carry dreams to me.",
|
|
"role": "assistant",
|
|
}
|
|
}
|
|
],
|
|
}
|
|
example_embedding_result = {
|
|
"object": "list",
|
|
"data": [
|
|
{
|
|
"object": "embedding",
|
|
"index": 0,
|
|
"embedding": [
|
|
-0.006929283495992422,
|
|
-0.005336422007530928,
|
|
-4.547132266452536e-05,
|
|
-0.024047505110502243,
|
|
-0.006929283495992422,
|
|
-0.005336422007530928,
|
|
-4.547132266452536e-05,
|
|
-0.024047505110502243,
|
|
-0.006929283495992422,
|
|
-0.005336422007530928,
|
|
-4.547132266452536e-05,
|
|
-0.024047505110502243,
|
|
],
|
|
}
|
|
],
|
|
"model": "text-embedding-3-small",
|
|
"usage": {"prompt_tokens": 5, "total_tokens": 5},
|
|
}
|
|
example_image_generation_result = {
|
|
"created": 1589478378,
|
|
"data": [{"url": "https://..."}, {"url": "https://..."}],
|
|
}
|
|
|
|
|
|
def mock_patch_acompletion():
|
|
return mock.patch(
|
|
"litellm.proxy.proxy_server.llm_router.acompletion",
|
|
return_value=example_completion_result,
|
|
)
|
|
|
|
|
|
def mock_patch_aembedding():
|
|
return mock.patch(
|
|
"litellm.proxy.proxy_server.llm_router.aembedding",
|
|
return_value=example_embedding_result,
|
|
)
|
|
|
|
|
|
def mock_patch_aimage_generation():
|
|
return mock.patch(
|
|
"litellm.proxy.proxy_server.llm_router.aimage_generation",
|
|
return_value=example_image_generation_result,
|
|
)
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def fake_env_vars(monkeypatch):
|
|
# Set some fake environment variables
|
|
monkeypatch.setenv("OPENAI_API_KEY", "fake_openai_api_key")
|
|
monkeypatch.setenv("OPENAI_API_BASE", "http://fake-openai-api-base")
|
|
monkeypatch.setenv("AZURE_AI_API_BASE", "http://fake-azure-api-base")
|
|
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "fake_azure_openai_api_key")
|
|
monkeypatch.setenv("AZURE_SWEDEN_API_BASE", "http://fake-azure-sweden-api-base")
|
|
monkeypatch.setenv("REDIS_HOST", "localhost")
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def client_no_auth(fake_env_vars):
|
|
# Assuming litellm.proxy.proxy_server is an object
|
|
from litellm.proxy.proxy_server import cleanup_router_config_variables
|
|
|
|
cleanup_router_config_variables()
|
|
filepath = os.path.dirname(os.path.abspath(__file__))
|
|
config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml"
|
|
# initialize can get run in parallel, it sets specific variables for the fast api app, sinc eit gets run in parallel different tests use the wrong variables
|
|
asyncio.run(initialize(config=config_fp, debug=True))
|
|
return TestClient(app)
|
|
|
|
|
|
@mock_patch_acompletion()
|
|
def test_chat_completion(mock_acompletion, client_no_auth):
|
|
global headers
|
|
try:
|
|
# Your test data
|
|
test_data = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [
|
|
{"role": "user", "content": "hi"},
|
|
],
|
|
"max_tokens": 10,
|
|
}
|
|
|
|
print("testing proxy server with chat completions")
|
|
response = client_no_auth.post("/v1/chat/completions", json=test_data)
|
|
mock_acompletion.assert_called_once_with(
|
|
model="gpt-3.5-turbo",
|
|
messages=[
|
|
{"role": "user", "content": "hi"},
|
|
],
|
|
max_tokens=10,
|
|
litellm_call_id=mock.ANY,
|
|
litellm_logging_obj=mock.ANY,
|
|
request_timeout=mock.ANY,
|
|
specific_deployment=True,
|
|
metadata=mock.ANY,
|
|
proxy_server_request=mock.ANY,
|
|
secret_fields=mock.ANY,
|
|
)
|
|
print(f"response - {response.text}")
|
|
assert response.status_code == 200
|
|
result = response.json()
|
|
print(f"Received response: {result}")
|
|
except Exception as e:
|
|
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
|
|
|
|
|
|
def test_chat_completion_malformed_messages_returns_400(client_no_auth):
|
|
"""
|
|
Test that malformed messages (strings instead of dicts) return 400 instead of 500.
|
|
|
|
This test verifies that when a client sends messages as raw strings instead of
|
|
{role, content} objects, LiteLLM returns a 400 invalid_request_error instead
|
|
of a 500 Internal Server Error.
|
|
"""
|
|
global headers
|
|
try:
|
|
# Test data with malformed messages (string instead of dict)
|
|
test_data = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [
|
|
"hi how are you"
|
|
], # Invalid: should be [{"role": "user", "content": "hi how are you"}]
|
|
}
|
|
|
|
print("testing proxy server with malformed messages")
|
|
response = client_no_auth.post(
|
|
"/v1/chat/completions", json=test_data, headers=headers
|
|
)
|
|
|
|
print(f"response status: {response.status_code}")
|
|
print(f"response text: {response.text}")
|
|
|
|
# Should return 400, not 500
|
|
assert (
|
|
response.status_code == 400
|
|
), f"Expected 400, got {response.status_code}. Response: {response.text}"
|
|
|
|
# Verify error format
|
|
result = response.json()
|
|
assert "error" in result, "Response should contain 'error' key"
|
|
error = result["error"]
|
|
|
|
# Verify error type and message
|
|
assert (
|
|
error.get("type") == "invalid_request_error" or error.get("type") is None
|
|
), f"Expected invalid_request_error or None, got {error.get('type')}"
|
|
assert (
|
|
error.get("code") == "400" or error.get("code") == 400
|
|
), f"Expected code 400, got {error.get('code')}"
|
|
|
|
# Error message should indicate invalid request format
|
|
error_message = error.get("message", "")
|
|
assert len(error_message) > 0, "Error message should not be empty"
|
|
|
|
except Exception as e:
|
|
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
|
|
|
|
|
|
def test_get_settings_request_timeout(client_no_auth):
|
|
"""
|
|
When no timeout is set, it should use the litellm.request_timeout value
|
|
"""
|
|
# Set a known value for litellm.request_timeout
|
|
import litellm
|
|
|
|
# Make a GET request to /settings
|
|
response = client_no_auth.get("/settings")
|
|
|
|
# Check if the request was successful
|
|
assert response.status_code == 200
|
|
|
|
# Parse the JSON response
|
|
settings = response.json()
|
|
print("settings", settings)
|
|
|
|
assert settings["litellm.request_timeout"] == litellm.request_timeout
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"litellm_key_header_name",
|
|
["x-litellm-key", None],
|
|
)
|
|
def test_add_headers_to_request(litellm_key_header_name):
|
|
from fastapi import Request
|
|
from starlette.datastructures import URL
|
|
import json
|
|
from litellm.proxy.litellm_pre_call_utils import (
|
|
clean_headers,
|
|
LiteLLMProxyRequestSetup,
|
|
)
|
|
|
|
headers = {
|
|
"Authorization": "Bearer 1234",
|
|
"X-Custom-Header": "Custom-Value",
|
|
"X-Stainless-Header": "Stainless-Value",
|
|
"anthropic-beta": "beta-value",
|
|
}
|
|
request = Request(scope={"type": "http", "method": "POST", "path": "/chat/completions", "headers": []})
|
|
request._url = URL(url="/chat/completions")
|
|
request._body = json.dumps({"model": "gpt-3.5-turbo"}).encode("utf-8")
|
|
request_headers = clean_headers(headers, litellm_key_header_name)
|
|
forwarded_headers = LiteLLMProxyRequestSetup._get_forwardable_headers(
|
|
request_headers
|
|
)
|
|
assert forwarded_headers == {
|
|
"X-Custom-Header": "Custom-Value",
|
|
"anthropic-beta": "beta-value",
|
|
}
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"litellm_key_header_name",
|
|
["x-litellm-key", None],
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"forward_headers",
|
|
[True, False],
|
|
)
|
|
@mock_patch_acompletion()
|
|
def test_chat_completion_forward_headers(
|
|
mock_acompletion, client_no_auth, litellm_key_header_name, forward_headers
|
|
):
|
|
global headers
|
|
try:
|
|
if forward_headers:
|
|
gs = getattr(litellm.proxy.proxy_server, "general_settings")
|
|
gs["forward_client_headers_to_llm_api"] = True
|
|
setattr(litellm.proxy.proxy_server, "general_settings", gs)
|
|
if litellm_key_header_name is not None:
|
|
gs = getattr(litellm.proxy.proxy_server, "general_settings")
|
|
gs["litellm_key_header_name"] = litellm_key_header_name
|
|
setattr(litellm.proxy.proxy_server, "general_settings", gs)
|
|
# Your test data
|
|
test_data = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [
|
|
{"role": "user", "content": "hi"},
|
|
],
|
|
"max_tokens": 10,
|
|
}
|
|
|
|
headers_to_forward = {
|
|
"X-Custom-Header": "Custom-Value",
|
|
"X-Another-Header": "Another-Value",
|
|
}
|
|
|
|
if litellm_key_header_name is not None:
|
|
headers_to_not_forward = {litellm_key_header_name: "Bearer 1234"}
|
|
else:
|
|
headers_to_not_forward = {"Authorization": "Bearer 1234"}
|
|
|
|
received_headers = {**headers_to_forward, **headers_to_not_forward}
|
|
|
|
print("testing proxy server with chat completions")
|
|
response = client_no_auth.post(
|
|
"/v1/chat/completions", json=test_data, headers=received_headers
|
|
)
|
|
if not forward_headers:
|
|
assert "headers" not in mock_acompletion.call_args.kwargs
|
|
else:
|
|
assert mock_acompletion.call_args.kwargs["headers"] == {
|
|
"x-custom-header": "Custom-Value",
|
|
"x-another-header": "Another-Value",
|
|
}
|
|
|
|
print(f"response - {response.text}")
|
|
assert response.status_code == 200
|
|
result = response.json()
|
|
print(f"Received response: {result}")
|
|
except Exception as e:
|
|
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
|
|
|
|
|
|
@pytest.mark.parametrize("forward_llm_auth_headers", [True, False])
|
|
@mock_patch_acompletion()
|
|
def test_chat_completion_forward_llm_provider_auth_headers(
|
|
mock_acompletion, client_no_auth, forward_llm_auth_headers
|
|
):
|
|
"""
|
|
Test that LLM provider auth headers (x-api-key, x-goog-api-key) are forwarded
|
|
when forward_llm_provider_auth_headers=True.
|
|
|
|
This allows clients to send their own LLM provider API keys through the proxy.
|
|
"""
|
|
try:
|
|
# Configure general settings
|
|
gs = getattr(litellm.proxy.proxy_server, "general_settings")
|
|
gs["forward_client_headers_to_llm_api"] = True
|
|
gs["forward_llm_provider_auth_headers"] = forward_llm_auth_headers
|
|
setattr(litellm.proxy.proxy_server, "general_settings", gs)
|
|
|
|
# Test data
|
|
test_data = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [
|
|
{"role": "user", "content": "hello"},
|
|
],
|
|
"max_tokens": 10,
|
|
}
|
|
|
|
# Headers including LLM provider auth
|
|
request_headers = {
|
|
"Authorization": "Bearer sk-proxy-auth-123", # Proxy auth (should be stripped)
|
|
"x-api-key": "sk-ant-api03-test-anthropic-key", # Anthropic API key
|
|
"x-goog-api-key": "google-api-key-123", # Google API key
|
|
"X-Custom-Header": "custom-value", # Custom header (should be forwarded)
|
|
}
|
|
|
|
# Make request
|
|
response = client_no_auth.post(
|
|
"/v1/chat/completions", json=test_data, headers=request_headers
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
|
|
# Check forwarded headers
|
|
forwarded_headers = mock_acompletion.call_args.kwargs.get("headers", {})
|
|
|
|
if forward_llm_auth_headers:
|
|
# LLM provider auth headers should be forwarded
|
|
assert "x-api-key" in forwarded_headers
|
|
assert forwarded_headers["x-api-key"] == "sk-ant-api03-test-anthropic-key"
|
|
assert "x-goog-api-key" in forwarded_headers
|
|
assert forwarded_headers["x-goog-api-key"] == "google-api-key-123"
|
|
else:
|
|
# LLM provider auth headers should be stripped
|
|
assert "x-api-key" not in forwarded_headers
|
|
assert "x-goog-api-key" not in forwarded_headers
|
|
|
|
# Custom headers should always be forwarded (when forward_client_headers_to_llm_api=True)
|
|
assert "x-custom-header" in forwarded_headers
|
|
assert forwarded_headers["x-custom-header"] == "custom-value"
|
|
|
|
# Proxy Authorization should never be forwarded
|
|
assert "authorization" not in forwarded_headers
|
|
|
|
print(
|
|
f"✓ Test passed with forward_llm_provider_auth_headers={forward_llm_auth_headers}"
|
|
)
|
|
print(f" Forwarded headers: {list(forwarded_headers.keys())}")
|
|
|
|
except Exception as e:
|
|
pytest.fail(
|
|
f"Test failed with forward_llm_auth_headers={forward_llm_auth_headers}: {str(e)}"
|
|
)
|
|
finally:
|
|
# Clean up
|
|
gs = getattr(litellm.proxy.proxy_server, "general_settings")
|
|
gs.pop("forward_llm_provider_auth_headers", None)
|
|
setattr(litellm.proxy.proxy_server, "general_settings", gs)
|
|
|
|
|
|
@mock_patch_acompletion()
|
|
@pytest.mark.asyncio
|
|
async def test_team_disable_guardrails(mock_acompletion, client_no_auth, monkeypatch):
|
|
"""
|
|
If team not allowed to turn on/off guardrails
|
|
|
|
Raise 403 forbidden error, if request is made by team on `/key/generate` or `/chat/completions`.
|
|
"""
|
|
import asyncio
|
|
import json
|
|
import time
|
|
|
|
from fastapi import HTTPException, Request
|
|
from starlette.datastructures import URL
|
|
|
|
from litellm.proxy._types import (
|
|
LiteLLM_TeamTable,
|
|
LiteLLM_TeamTableCachedObj,
|
|
ProxyException,
|
|
UserAPIKeyAuth,
|
|
)
|
|
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
|
from litellm.proxy.proxy_server import hash_token
|
|
|
|
user_api_key_cache: Final = UserApiKeyCache()
|
|
_team_id = "1234"
|
|
user_key = "sk-98765678"
|
|
|
|
valid_token = UserAPIKeyAuth(
|
|
team_id=_team_id,
|
|
team_blocked=True,
|
|
token=hash_token(user_key),
|
|
last_refreshed_at=time.time(),
|
|
)
|
|
await asyncio.sleep(1)
|
|
team_obj = LiteLLM_TeamTableCachedObj(
|
|
team_id=_team_id,
|
|
blocked=False,
|
|
last_refreshed_at=time.time(),
|
|
metadata={"guardrails": {"modify_guardrails": False}},
|
|
)
|
|
user_api_key_cache.set_cache(key=hash_token(user_key), value=valid_token)
|
|
user_api_key_cache.set_cache(key="team_id:{}".format(_team_id), value=team_obj)
|
|
|
|
monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
|
|
setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
|
|
setattr(litellm.proxy.proxy_server, "prisma_client", "hello-world")
|
|
|
|
request = Request(scope={"type": "http", "method": "POST", "path": "/chat/completions", "headers": []})
|
|
request._url = URL(url="/chat/completions")
|
|
|
|
body = {"metadata": {"guardrails": {"hide_secrets": False}}}
|
|
json_bytes = json.dumps(body).encode("utf-8")
|
|
|
|
request._body = json_bytes
|
|
|
|
with pytest.raises(ProxyException) as exc_info:
|
|
await user_api_key_auth(request=request, api_key="Bearer " + user_key)
|
|
e = exc_info.value
|
|
assert e.code == str(403)
|
|
|
|
|
|
from tests.unit.proxy.test_custom_callback_input import CompletionCustomHandler
|
|
|
|
|
|
@mock_patch_acompletion()
|
|
def test_custom_logger_failure_handler(mock_acompletion, client_no_auth, monkeypatch):
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from litellm.proxy.proxy_server import hash_token
|
|
|
|
user_api_key_cache: Final = UserApiKeyCache()
|
|
rpm_limit = 0
|
|
|
|
mock_api_key = "sk-my-test-key"
|
|
cache_value = UserAPIKeyAuth(token=hash_token(mock_api_key), rpm_limit=rpm_limit)
|
|
|
|
user_api_key_cache.set_cache(key=hash_token(mock_api_key), value=cache_value)
|
|
|
|
mock_logger = CustomLogger()
|
|
mock_logger_unit_tests = CompletionCustomHandler()
|
|
proxy_logging_obj: ProxyLogging = getattr(
|
|
litellm.proxy.proxy_server, "proxy_logging_obj"
|
|
)
|
|
|
|
litellm.callbacks = [mock_logger, mock_logger_unit_tests]
|
|
proxy_logging_obj._init_litellm_callbacks(llm_router=None)
|
|
|
|
monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
|
|
setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
|
|
setattr(litellm.proxy.proxy_server, "prisma_client", "FAKE-VAR")
|
|
setattr(litellm.proxy.proxy_server, "proxy_logging_obj", proxy_logging_obj)
|
|
|
|
with patch.object(
|
|
mock_logger, "async_log_failure_event", new=AsyncMock()
|
|
) as mock_failed_alert:
|
|
# Your test data
|
|
test_data = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [
|
|
{"role": "user", "content": "hi"},
|
|
],
|
|
"max_tokens": 10,
|
|
}
|
|
|
|
print("testing proxy server with chat completions")
|
|
response = client_no_auth.post(
|
|
"/v1/chat/completions",
|
|
json=test_data,
|
|
headers={"Authorization": "Bearer {}".format(mock_api_key)},
|
|
)
|
|
assert response.status_code == 429
|
|
|
|
# confirm async_log_failure_event is called
|
|
mock_failed_alert.assert_called()
|
|
|
|
assert len(mock_logger_unit_tests.errors) == 0
|
|
|
|
|
|
@mock_patch_acompletion()
|
|
def test_engines_model_chat_completions(mock_acompletion, client_no_auth):
|
|
global headers
|
|
try:
|
|
# Your test data
|
|
test_data = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [
|
|
{"role": "user", "content": "hi"},
|
|
],
|
|
"max_tokens": 10,
|
|
}
|
|
|
|
print("testing proxy server with chat completions")
|
|
response = client_no_auth.post(
|
|
"/engines/gpt-3.5-turbo/chat/completions", json=test_data
|
|
)
|
|
mock_acompletion.assert_called_once_with(
|
|
model="gpt-3.5-turbo",
|
|
messages=[
|
|
{"role": "user", "content": "hi"},
|
|
],
|
|
max_tokens=10,
|
|
litellm_call_id=mock.ANY,
|
|
litellm_logging_obj=mock.ANY,
|
|
request_timeout=mock.ANY,
|
|
specific_deployment=True,
|
|
metadata=mock.ANY,
|
|
proxy_server_request=mock.ANY,
|
|
secret_fields=mock.ANY,
|
|
)
|
|
print(f"response - {response.text}")
|
|
assert response.status_code == 200
|
|
result = response.json()
|
|
print(f"Received response: {result}")
|
|
except Exception as e:
|
|
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
|
|
|
|
|
|
@mock_patch_acompletion()
|
|
def test_chat_completion_azure(mock_acompletion, client_no_auth):
|
|
global headers
|
|
try:
|
|
# Your test data
|
|
test_data = {
|
|
"model": "azure/gpt-4.1-mini",
|
|
"messages": [
|
|
{"role": "user", "content": "write 1 sentence poem"},
|
|
],
|
|
"max_tokens": 10,
|
|
}
|
|
|
|
print("testing proxy server with Azure Request /chat/completions")
|
|
response = client_no_auth.post("/v1/chat/completions", json=test_data)
|
|
|
|
mock_acompletion.assert_called_once_with(
|
|
model="azure/gpt-4.1-mini",
|
|
messages=[
|
|
{"role": "user", "content": "write 1 sentence poem"},
|
|
],
|
|
max_tokens=10,
|
|
litellm_call_id=mock.ANY,
|
|
litellm_logging_obj=mock.ANY,
|
|
request_timeout=mock.ANY,
|
|
specific_deployment=True,
|
|
metadata=mock.ANY,
|
|
proxy_server_request=mock.ANY,
|
|
secret_fields=mock.ANY,
|
|
)
|
|
assert response.status_code == 200
|
|
result = response.json()
|
|
print(f"Received response: {result}")
|
|
assert len(result["choices"][0]["message"]["content"]) > 0
|
|
except Exception as e:
|
|
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
|
|
|
|
|
|
# Run the test
|
|
# test_chat_completion_azure()
|
|
|
|
|
|
@mock_patch_acompletion()
|
|
def test_openai_deployments_model_chat_completions_azure(
|
|
mock_acompletion, client_no_auth
|
|
):
|
|
global headers
|
|
try:
|
|
# Your test data
|
|
test_data = {
|
|
"model": "azure/gpt-4.1-mini",
|
|
"messages": [
|
|
{"role": "user", "content": "write 1 sentence poem"},
|
|
],
|
|
"max_tokens": 10,
|
|
}
|
|
|
|
url = "/openai/deployments/azure/gpt-4.1-mini/chat/completions"
|
|
print(f"testing proxy server with Azure Request {url}")
|
|
response = client_no_auth.post(url, json=test_data)
|
|
|
|
mock_acompletion.assert_called_once_with(
|
|
model="azure/gpt-4.1-mini",
|
|
messages=[
|
|
{"role": "user", "content": "write 1 sentence poem"},
|
|
],
|
|
max_tokens=10,
|
|
litellm_call_id=mock.ANY,
|
|
litellm_logging_obj=mock.ANY,
|
|
request_timeout=mock.ANY,
|
|
specific_deployment=True,
|
|
metadata=mock.ANY,
|
|
proxy_server_request=mock.ANY,
|
|
secret_fields=mock.ANY,
|
|
)
|
|
assert response.status_code == 200
|
|
result = response.json()
|
|
print(f"Received response: {result}")
|
|
assert len(result["choices"][0]["message"]["content"]) > 0
|
|
except Exception as e:
|
|
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
|
|
|
|
|
|
# Run the test
|
|
# test_openai_deployments_model_chat_completions_azure()
|
|
|
|
|
|
### EMBEDDING
|
|
@mock_patch_aembedding()
|
|
def test_embedding(mock_aembedding, client_no_auth):
|
|
global headers
|
|
from litellm.proxy.proxy_server import user_custom_auth
|
|
|
|
try:
|
|
test_data = {
|
|
"model": "azure/text-embedding-ada-002",
|
|
"input": ["good morning from litellm"],
|
|
}
|
|
|
|
async def _pre_call_hook_side_effect(**kwargs):
|
|
data = kwargs["data"]
|
|
metadata = {**(data.get("metadata") or {}), "source": "unit-test"}
|
|
data["metadata"] = metadata
|
|
proxy_request = {**(data.get("proxy_server_request") or {})}
|
|
proxy_request["path"] = "/v1/embeddings"
|
|
data["proxy_server_request"] = proxy_request
|
|
return data
|
|
|
|
async def _post_call_success_side_effect(**kwargs):
|
|
return kwargs["response"]
|
|
|
|
with (
|
|
patch.object(
|
|
litellm.proxy.proxy_server.proxy_logging_obj,
|
|
"pre_call_hook",
|
|
new=AsyncMock(side_effect=_pre_call_hook_side_effect),
|
|
) as mock_pre_call_hook,
|
|
patch.object(
|
|
litellm.proxy.proxy_server.proxy_logging_obj,
|
|
"during_call_hook",
|
|
new=AsyncMock(return_value=None),
|
|
) as mock_during_hook,
|
|
patch.object(
|
|
litellm.proxy.proxy_server.proxy_logging_obj,
|
|
"post_call_success_hook",
|
|
new=AsyncMock(side_effect=_post_call_success_side_effect),
|
|
),
|
|
):
|
|
response = client_no_auth.post("/v1/embeddings", json=test_data)
|
|
|
|
mock_aembedding.assert_called_once_with(
|
|
model="azure/text-embedding-ada-002",
|
|
input=["good morning from litellm"],
|
|
specific_deployment=True,
|
|
litellm_call_id=mock.ANY,
|
|
litellm_logging_obj=mock.ANY,
|
|
request_timeout=mock.ANY,
|
|
metadata=mock.ANY,
|
|
proxy_server_request=mock.ANY,
|
|
secret_fields=mock.ANY,
|
|
)
|
|
assert response.status_code == 200
|
|
result = response.json()
|
|
print(len(result["data"][0]["embedding"]))
|
|
assert len(result["data"][0]["embedding"]) > 10 # this usually has len==1536 so
|
|
|
|
call_metadata = mock_aembedding.call_args.kwargs["metadata"]
|
|
assert call_metadata.get("source") == "unit-test"
|
|
|
|
pre_call_kwargs = mock_pre_call_hook.await_args_list[0].kwargs
|
|
assert (
|
|
pre_call_kwargs.get("call_type") == "aembedding"
|
|
), f"expected pre_call_hook to receive call_type='aembedding', got {pre_call_kwargs.get('call_type')}"
|
|
|
|
except Exception as e:
|
|
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
|
|
|
|
|
|
@mock_patch_aembedding()
|
|
def test_bedrock_embedding(mock_aembedding, client_no_auth):
|
|
global headers
|
|
from litellm.proxy.proxy_server import user_custom_auth
|
|
|
|
try:
|
|
test_data = {
|
|
"model": "amazon-embeddings",
|
|
"input": ["good morning from litellm"],
|
|
}
|
|
|
|
response = client_no_auth.post("/v1/embeddings", json=test_data)
|
|
|
|
mock_aembedding.assert_called_once_with(
|
|
model="amazon-embeddings",
|
|
input=["good morning from litellm"],
|
|
litellm_call_id=mock.ANY,
|
|
litellm_logging_obj=mock.ANY,
|
|
request_timeout=mock.ANY,
|
|
metadata=mock.ANY,
|
|
proxy_server_request=mock.ANY,
|
|
secret_fields=mock.ANY,
|
|
)
|
|
assert response.status_code == 200
|
|
print(response.status_code, response.text)
|
|
result = response.json()
|
|
print(len(result["data"][0]["embedding"]))
|
|
assert len(result["data"][0]["embedding"]) > 10 # this usually has len==1536 so
|
|
except Exception as e:
|
|
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
|
|
|
|
|
|
@pytest.mark.skip(reason="AWS Suspended Account")
|
|
def test_sagemaker_embedding(client_no_auth):
|
|
global headers
|
|
from litellm.proxy.proxy_server import user_custom_auth
|
|
|
|
try:
|
|
test_data = {
|
|
"model": "GPT-J 6B - Sagemaker Text Embedding (Internal)",
|
|
"input": ["good morning from litellm"],
|
|
}
|
|
|
|
response = client_no_auth.post("/v1/embeddings", json=test_data)
|
|
|
|
assert response.status_code == 200
|
|
result = response.json()
|
|
print(len(result["data"][0]["embedding"]))
|
|
assert len(result["data"][0]["embedding"]) > 10 # this usually has len==1536 so
|
|
except Exception as e:
|
|
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
|
|
|
|
|
|
# Run the test
|
|
# test_embedding()
|
|
#### IMAGE GENERATION
|
|
|
|
|
|
@mock_patch_aimage_generation()
|
|
def test_img_gen(mock_aimage_generation, client_no_auth):
|
|
global headers
|
|
from litellm.proxy.proxy_server import user_custom_auth
|
|
|
|
try:
|
|
test_data = {
|
|
"model": "dall-e-3",
|
|
"prompt": "A cute baby sea otter",
|
|
"n": 1,
|
|
"size": "1024x1024",
|
|
"imageConfig": {"aspectRatio": "9:16", "imageSize": "1K"},
|
|
}
|
|
|
|
response = client_no_auth.post("/v1/images/generations", json=test_data)
|
|
|
|
mock_aimage_generation.assert_called_once_with(
|
|
model="dall-e-3",
|
|
prompt="A cute baby sea otter",
|
|
n=1,
|
|
size="1024x1024",
|
|
imageConfig={"aspectRatio": "9:16", "imageSize": "1K"},
|
|
litellm_call_id=mock.ANY,
|
|
metadata=mock.ANY,
|
|
proxy_server_request=mock.ANY,
|
|
secret_fields=mock.ANY,
|
|
)
|
|
assert response.status_code == 200
|
|
result = response.json()
|
|
print(len(result["data"][0]["url"]))
|
|
assert len(result["data"][0]["url"]) > 10
|
|
except Exception as e:
|
|
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
|
|
|
|
|
|
#### ADDITIONAL
|
|
@pytest.mark.skip(reason="test via docker tests. Requires prisma client.")
|
|
def test_add_new_model(client_no_auth):
|
|
global headers
|
|
try:
|
|
test_data = {
|
|
"model_name": "test_openai_models",
|
|
"litellm_params": {
|
|
"model": "gpt-3.5-turbo",
|
|
},
|
|
"model_info": {"description": "this is a test openai model"},
|
|
}
|
|
client_no_auth.post("/model/new", json=test_data, headers=headers)
|
|
response = client_no_auth.get("/model/info", headers=headers)
|
|
assert response.status_code == 200
|
|
result = response.json()
|
|
print(f"response: {result}")
|
|
model_info = None
|
|
for m in result["data"]:
|
|
if m["model_name"] == "test_openai_models":
|
|
model_info = m["model_info"]
|
|
assert model_info["description"] == "this is a test openai model"
|
|
except Exception as e:
|
|
pytest.fail(f"LiteLLM Proxy test failed. Exception {str(e)}")
|
|
|
|
|
|
@pytest.mark.xdist_group("proxy_heavy")
|
|
def test_health(client_no_auth):
|
|
global headers
|
|
import logging
|
|
import time
|
|
|
|
from litellm._logging import verbose_logger, verbose_proxy_logger
|
|
|
|
verbose_proxy_logger.setLevel(logging.DEBUG)
|
|
|
|
try:
|
|
response = client_no_auth.get("/health")
|
|
assert response.status_code == 200
|
|
except Exception as e:
|
|
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
|
|
|
|
|
|
# test_add_new_model()
|
|
|
|
|
|
|
|
class MyCustomHandler(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
print(f"Pre-API Call")
|
|
|
|
def log_success_event(self, kwargs, response_obj, start_time, end_time):
|
|
print(f"On Success")
|
|
assert kwargs["user"] == "proxy-user"
|
|
assert kwargs["model"] == "gpt-3.5-turbo"
|
|
assert kwargs["max_tokens"] == 10
|
|
|
|
|
|
customHandler = MyCustomHandler()
|
|
|
|
|
|
@mock_patch_acompletion()
|
|
def test_chat_completion_optional_params(mock_acompletion, client_no_auth):
|
|
# [PROXY: PROD TEST] - DO NOT DELETE
|
|
# This tests if all the /chat/completion params are passed to litellm
|
|
try:
|
|
# Your test data
|
|
litellm.set_verbose = True
|
|
test_data = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [
|
|
{"role": "user", "content": "hi"},
|
|
],
|
|
"max_tokens": 10,
|
|
"user": "proxy-user",
|
|
}
|
|
|
|
litellm.callbacks = [customHandler]
|
|
print("testing proxy server: optional params")
|
|
response = client_no_auth.post("/v1/chat/completions", json=test_data)
|
|
mock_acompletion.assert_called_once_with(
|
|
model="gpt-3.5-turbo",
|
|
messages=[
|
|
{"role": "user", "content": "hi"},
|
|
],
|
|
max_tokens=10,
|
|
user="proxy-user",
|
|
litellm_call_id=mock.ANY,
|
|
litellm_logging_obj=mock.ANY,
|
|
request_timeout=mock.ANY,
|
|
specific_deployment=True,
|
|
metadata=mock.ANY,
|
|
proxy_server_request=mock.ANY,
|
|
secret_fields=mock.ANY,
|
|
)
|
|
assert response.status_code == 200
|
|
result = response.json()
|
|
print(f"Received response: {result}")
|
|
except Exception as e:
|
|
pytest.fail("LiteLLM Proxy test failed. Exception", e)
|
|
|
|
|
|
# Run the test
|
|
# test_chat_completion_optional_params()
|
|
|
|
|
|
# Test Reading config.yaml file
|
|
from litellm.proxy.proxy_server import ProxyConfig
|
|
|
|
|
|
@pytest.mark.skip(reason="local variable conflicts. needs to be refactored.")
|
|
@mock.patch("litellm.proxy.proxy_server.litellm.Cache")
|
|
def test_load_router_config(mock_cache, fake_env_vars):
|
|
mock_cache.return_value.cache.__dict__ = {"redis_client": None}
|
|
mock_cache.return_value.supported_call_types = [
|
|
"completion",
|
|
"acompletion",
|
|
"embedding",
|
|
"aembedding",
|
|
"atranscription",
|
|
"transcription",
|
|
]
|
|
|
|
try:
|
|
import asyncio
|
|
|
|
print("testing reading config")
|
|
# this is a basic config.yaml with only a model
|
|
filepath = os.path.dirname(os.path.abspath(__file__))
|
|
proxy_config = ProxyConfig()
|
|
result = asyncio.run(
|
|
proxy_config.load_config(
|
|
router=None,
|
|
config_file_path=f"{filepath}/example_config_yaml/simple_config.yaml",
|
|
)
|
|
)
|
|
print(result)
|
|
assert len(result[1]) == 1
|
|
|
|
# this is a load balancing config yaml
|
|
result = asyncio.run(
|
|
proxy_config.load_config(
|
|
router=None,
|
|
config_file_path=f"{filepath}/example_config_yaml/azure_config.yaml",
|
|
)
|
|
)
|
|
print(result)
|
|
assert len(result[1]) == 2
|
|
|
|
# config with general settings - custom callbacks
|
|
result = asyncio.run(
|
|
proxy_config.load_config(
|
|
router=None,
|
|
config_file_path=f"{filepath}/example_config_yaml/azure_config.yaml",
|
|
)
|
|
)
|
|
print(result)
|
|
assert len(result[1]) == 2
|
|
|
|
# tests for litellm.cache set from config
|
|
print("testing reading proxy config for cache")
|
|
litellm.cache = None
|
|
asyncio.run(
|
|
proxy_config.load_config(
|
|
router=None,
|
|
config_file_path=f"{filepath}/example_config_yaml/cache_no_params.yaml",
|
|
)
|
|
)
|
|
assert litellm.cache is not None
|
|
assert "redis_client" in vars(
|
|
litellm.cache.cache
|
|
) # it should default to redis on proxy
|
|
assert litellm.cache.supported_call_types == [
|
|
"completion",
|
|
"acompletion",
|
|
"embedding",
|
|
"aembedding",
|
|
"atranscription",
|
|
"transcription",
|
|
] # init with all call types
|
|
|
|
litellm.disable_cache()
|
|
|
|
print("testing reading proxy config for cache with params")
|
|
mock_cache.return_value.supported_call_types = [
|
|
"embedding",
|
|
"aembedding",
|
|
]
|
|
asyncio.run(
|
|
proxy_config.load_config(
|
|
router=None,
|
|
config_file_path=f"{filepath}/example_config_yaml/cache_with_params.yaml",
|
|
)
|
|
)
|
|
assert litellm.cache is not None
|
|
print(litellm.cache)
|
|
print(litellm.cache.supported_call_types)
|
|
print(vars(litellm.cache.cache))
|
|
assert "redis_client" in vars(
|
|
litellm.cache.cache
|
|
) # it should default to redis on proxy
|
|
assert litellm.cache.supported_call_types == [
|
|
"embedding",
|
|
"aembedding",
|
|
] # init with all call types
|
|
|
|
except Exception as e:
|
|
pytest.fail(
|
|
f"Proxy: Got exception reading config: {str(e)}\n{traceback.format_exc()}"
|
|
)
|
|
|
|
|
|
# test_load_router_config()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_team_update_redis():
|
|
"""
|
|
Tests if team update, updates the redis cache if set
|
|
"""
|
|
from litellm.caching.caching import DualCache, RedisCache
|
|
from litellm.proxy._types import LiteLLM_TeamTableCachedObj
|
|
from litellm.proxy.auth.auth_checks import cache_team_object
|
|
|
|
proxy_logging_obj: ProxyLogging = getattr(
|
|
litellm.proxy.proxy_server, "proxy_logging_obj"
|
|
)
|
|
|
|
redis_cache = RedisCache(host="localhost")
|
|
|
|
with patch.object(
|
|
redis_cache,
|
|
"async_set_cache",
|
|
new=AsyncMock(),
|
|
) as mock_client:
|
|
await cache_team_object(
|
|
team_id="1234",
|
|
team_table=LiteLLM_TeamTableCachedObj(team_id="1234"),
|
|
user_api_key_cache=DualCache(redis_cache=redis_cache),
|
|
proxy_logging_obj=proxy_logging_obj,
|
|
)
|
|
|
|
mock_client.assert_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_team_redis(client_no_auth):
|
|
"""
|
|
Tests if get_team_object gets value from redis cache, if set
|
|
"""
|
|
from litellm.caching.caching import DualCache, RedisCache
|
|
from litellm.proxy.auth.auth_checks import get_team_object
|
|
|
|
proxy_logging_obj: ProxyLogging = getattr(
|
|
litellm.proxy.proxy_server, "proxy_logging_obj"
|
|
)
|
|
|
|
redis_cache = RedisCache()
|
|
|
|
from fastapi import HTTPException
|
|
|
|
with patch.object(
|
|
redis_cache,
|
|
"async_get_cache",
|
|
new=AsyncMock(),
|
|
) as mock_client:
|
|
try:
|
|
await get_team_object(
|
|
team_id="1234",
|
|
user_api_key_cache=DualCache(redis_cache=redis_cache),
|
|
parent_otel_span=None,
|
|
proxy_logging_obj=proxy_logging_obj,
|
|
prisma_client=AsyncMock(),
|
|
)
|
|
except HTTPException:
|
|
pass
|
|
|
|
mock_client.assert_called_once()
|
|
|
|
|
|
import random
|
|
from litellm._uuid import uuid
|
|
from unittest.mock import PropertyMock
|
|
|
|
from litellm.proxy._types import (
|
|
LitellmUserRoles,
|
|
NewUserRequest,
|
|
TeamMemberAddRequest,
|
|
UserAPIKeyAuth,
|
|
)
|
|
from litellm.proxy.management_endpoints.internal_user_endpoints import new_user
|
|
from litellm.proxy.management_endpoints.team_endpoints import team_member_add
|
|
from tests.unit.proxy.management_endpoints.test_key_generate_prisma import prisma_client
|
|
@pytest.fixture
|
|
def mock_prisma_client():
|
|
client = MagicMock()
|
|
client.connect = AsyncMock()
|
|
client.disconnect = AsyncMock()
|
|
return client
|
|
|
|
|
|
@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).")
|
|
@pytest.mark.parametrize(
|
|
"user_role",
|
|
[LitellmUserRoles.INTERNAL_USER.value, LitellmUserRoles.PROXY_ADMIN.value],
|
|
)
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).")
|
|
async def test_create_user_default_budget(prisma_client, user_role): # noqa: F811 # pytest fixture, not a redefinition
|
|
|
|
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
|
|
setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
|
|
setattr(litellm, "max_internal_user_budget", 10)
|
|
setattr(litellm, "internal_user_budget_duration", "5m")
|
|
await litellm.proxy.proxy_server.prisma_client.connect()
|
|
user = f"ishaan {uuid.uuid4().hex}"
|
|
request = NewUserRequest(
|
|
user_id=user, user_role=user_role
|
|
) # create a key with no budget
|
|
with patch.object(
|
|
litellm.proxy.proxy_server.prisma_client, "insert_data", new=AsyncMock()
|
|
) as mock_client:
|
|
await new_user(
|
|
request,
|
|
)
|
|
|
|
mock_client.assert_called()
|
|
|
|
print(f"mock_client.call_args: {mock_client.call_args}")
|
|
print("mock_client.call_args.kwargs: {}".format(mock_client.call_args.kwargs))
|
|
|
|
if user_role == LitellmUserRoles.INTERNAL_USER.value:
|
|
assert (
|
|
mock_client.call_args.kwargs["data"]["max_budget"]
|
|
== litellm.max_internal_user_budget
|
|
)
|
|
assert (
|
|
mock_client.call_args.kwargs["data"]["budget_duration"]
|
|
== litellm.internal_user_budget_duration
|
|
)
|
|
|
|
else:
|
|
assert mock_client.call_args.kwargs["data"]["max_budget"] is None
|
|
assert mock_client.call_args.kwargs["data"]["budget_duration"] is None
|
|
|
|
|
|
def _member_add_tx_cm(team_table):
|
|
"""Transaction whose member writes land on whatever tables are mocked on `prisma_client.db`"""
|
|
|
|
class _Tx:
|
|
query_raw = AsyncMock(return_value=[{"members_with_roles": []}])
|
|
litellm_teamtable = team_table
|
|
|
|
def __getattr__(self, table_name):
|
|
return getattr(litellm.proxy.proxy_server.prisma_client.db, table_name)
|
|
|
|
tx_cm = MagicMock()
|
|
tx_cm.__aenter__ = AsyncMock(return_value=_Tx())
|
|
tx_cm.__aexit__ = AsyncMock(return_value=None)
|
|
return tx_cm
|
|
|
|
|
|
@pytest.mark.parametrize("new_member_method", ["user_id", "user_email"])
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).")
|
|
async def test_create_team_member_add(prisma_client, new_member_method): # noqa: F811 # pytest fixture, not a redefinition
|
|
import time
|
|
|
|
from fastapi import Request
|
|
|
|
from litellm.proxy._types import LiteLLM_TeamTableCachedObj, LiteLLM_UserTable
|
|
from litellm.proxy.proxy_server import hash_token, user_api_key_cache
|
|
|
|
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
|
|
setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
|
|
setattr(litellm, "max_internal_user_budget", 10)
|
|
setattr(litellm, "internal_user_budget_duration", "5m")
|
|
await litellm.proxy.proxy_server.prisma_client.connect()
|
|
user = f"ishaan {uuid.uuid4().hex}"
|
|
_team_id = "litellm-test-client-id-new"
|
|
team_obj = LiteLLM_TeamTableCachedObj(
|
|
team_id=_team_id,
|
|
blocked=False,
|
|
last_refreshed_at=time.time(),
|
|
metadata={"guardrails": {"modify_guardrails": False}},
|
|
)
|
|
# user_api_key_cache.set_cache(key=hash_token(user_key), value=valid_token)
|
|
user_api_key_cache.set_cache(key="team_id:{}".format(_team_id), value=team_obj)
|
|
|
|
setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
|
|
if new_member_method == "user_id":
|
|
data = {
|
|
"team_id": _team_id,
|
|
"member": [{"role": "user", "user_id": user}],
|
|
}
|
|
elif new_member_method == "user_email":
|
|
data = {
|
|
"team_id": _team_id,
|
|
"member": [{"role": "user", "user_email": user}],
|
|
}
|
|
team_member_add_request = TeamMemberAddRequest(**data)
|
|
|
|
with (
|
|
patch(
|
|
"litellm.proxy.proxy_server.prisma_client.db.litellm_usertable",
|
|
new_callable=AsyncMock,
|
|
) as mock_litellm_usertable,
|
|
patch(
|
|
"litellm.proxy.auth.auth_checks._get_team_object_from_user_api_key_cache",
|
|
new=AsyncMock(return_value=team_obj),
|
|
) as mock_team_obj,
|
|
patch(
|
|
"litellm.proxy.proxy_server.prisma_client.get_data",
|
|
new=AsyncMock(return_value=[]),
|
|
) as mock_get_data,
|
|
):
|
|
|
|
mock_client = AsyncMock(
|
|
return_value=LiteLLM_UserTable(
|
|
user_id="1234", max_budget=100, user_email="1234"
|
|
)
|
|
)
|
|
mock_litellm_usertable.upsert = mock_client
|
|
mock_litellm_usertable.find_many = AsyncMock(return_value=[])
|
|
# Mock find_first for user_email validation (returns None for new users)
|
|
mock_litellm_usertable.find_first = AsyncMock(return_value=None)
|
|
# Mock find_unique for user_id validation (returns None for new users)
|
|
mock_litellm_usertable.find_unique = AsyncMock(return_value=None)
|
|
team_mock_client = AsyncMock()
|
|
original_val = getattr(
|
|
litellm.proxy.proxy_server.prisma_client.db, "litellm_teamtable"
|
|
)
|
|
litellm.proxy.proxy_server.prisma_client.db.litellm_teamtable = team_mock_client
|
|
|
|
team_mock_client.update = AsyncMock(
|
|
return_value=LiteLLM_TeamTableCachedObj(team_id="1234")
|
|
)
|
|
|
|
tx_cm = _member_add_tx_cm(team_mock_client)
|
|
original_tx = litellm.proxy.proxy_server.prisma_client.tx
|
|
litellm.proxy.proxy_server.prisma_client.tx = MagicMock(
|
|
return_value=tx_cm
|
|
)
|
|
|
|
print(f"team_member_add_request={team_member_add_request}")
|
|
await team_member_add(
|
|
data=team_member_add_request,
|
|
user_api_key_dict=UserAPIKeyAuth(user_role="proxy_admin"),
|
|
)
|
|
|
|
mock_client.assert_called()
|
|
|
|
print(f"mock_client.call_args: {mock_client.call_args}")
|
|
print("mock_client.call_args.kwargs: {}".format(mock_client.call_args.kwargs))
|
|
|
|
assert (
|
|
mock_client.call_args.kwargs["data"]["create"]["max_budget"]
|
|
== litellm.max_internal_user_budget
|
|
)
|
|
assert (
|
|
mock_client.call_args.kwargs["data"]["create"]["budget_duration"]
|
|
== litellm.internal_user_budget_duration
|
|
)
|
|
|
|
litellm.proxy.proxy_server.prisma_client.db.litellm_teamtable = original_val
|
|
litellm.proxy.proxy_server.prisma_client.tx = original_tx
|
|
|
|
|
|
@pytest.mark.parametrize("team_member_role", ["admin", "user"])
|
|
@pytest.mark.parametrize("team_route", ["/team/member_add", "/team/member_delete"])
|
|
@pytest.mark.asyncio
|
|
async def test_create_team_member_add_team_admin_user_api_key_auth(
|
|
prisma_client, team_member_role, team_route, monkeypatch # noqa: F811 # pytest fixture, not a redefinition
|
|
):
|
|
import time
|
|
|
|
from fastapi import Request
|
|
|
|
from litellm.proxy._types import LiteLLM_TeamTableCachedObj, Member
|
|
from litellm.proxy.proxy_server import (
|
|
ProxyException,
|
|
hash_token,
|
|
user_api_key_auth,
|
|
)
|
|
|
|
user_api_key_cache: Final = UserApiKeyCache()
|
|
|
|
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
|
|
setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
|
|
setattr(litellm, "max_internal_user_budget", 10)
|
|
setattr(litellm, "internal_user_budget_duration", "5m")
|
|
user = f"ishaan {uuid.uuid4().hex}"
|
|
_team_id = "litellm-test-client-id-new"
|
|
user_key = "sk-98765678"
|
|
|
|
valid_token = UserAPIKeyAuth(
|
|
team_id=_team_id,
|
|
token=hash_token(user_key),
|
|
team_member=Member(role=team_member_role, user_id=user),
|
|
last_refreshed_at=time.time(),
|
|
)
|
|
user_api_key_cache.set_cache(key=hash_token(user_key), value=valid_token)
|
|
|
|
team_obj = LiteLLM_TeamTableCachedObj(
|
|
team_id=_team_id,
|
|
blocked=False,
|
|
last_refreshed_at=time.time(),
|
|
metadata={"guardrails": {"modify_guardrails": False}},
|
|
)
|
|
|
|
user_api_key_cache.set_cache(key="team_id:{}".format(_team_id), value=team_obj)
|
|
|
|
monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
|
|
|
|
## TEST IF TEAM ADMIN ALLOWED TO CALL /MEMBER_ADD ENDPOINT
|
|
import json
|
|
|
|
from starlette.datastructures import URL
|
|
|
|
request = Request(scope={"type": "http", "method": "POST", "path": team_route, "headers": []})
|
|
request._url = URL(url=team_route)
|
|
|
|
body = {}
|
|
json_bytes = json.dumps(body).encode("utf-8")
|
|
|
|
request._body = json_bytes
|
|
|
|
## ALLOWED BY USER_API_KEY_AUTH
|
|
await user_api_key_auth(request=request, api_key="Bearer " + user_key)
|
|
|
|
|
|
@pytest.mark.parametrize("new_member_method", ["user_id", "user_email"])
|
|
@pytest.mark.parametrize("user_role", ["admin", "user"])
|
|
@pytest.mark.asyncio
|
|
async def test_create_team_member_add_team_admin(
|
|
prisma_client, new_member_method, user_role # noqa: F811 # pytest fixture, not a redefinition
|
|
):
|
|
"""
|
|
Relevant issue - https://github.com/BerriAI/litellm/issues/5300
|
|
|
|
Allow team admins to:
|
|
- Add and remove team members
|
|
- raise error if team member not an existing 'internal_user'
|
|
"""
|
|
import time
|
|
|
|
from fastapi import Request
|
|
|
|
from litellm.proxy._types import (
|
|
LiteLLM_TeamMembership,
|
|
LiteLLM_TeamTableCachedObj,
|
|
LiteLLM_UserTable,
|
|
Member,
|
|
)
|
|
from litellm.proxy.proxy_server import (
|
|
HTTPException,
|
|
ProxyException,
|
|
hash_token,
|
|
user_api_key_auth,
|
|
user_api_key_cache,
|
|
)
|
|
|
|
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
|
|
setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
|
|
setattr(litellm, "max_internal_user_budget", 10)
|
|
setattr(litellm, "internal_user_budget_duration", "5m")
|
|
user = f"ishaan {uuid.uuid4().hex}"
|
|
_team_id = "litellm-test-client-id-new"
|
|
user_key = "sk-98765678"
|
|
team_admin = f"krrish {uuid.uuid4().hex}"
|
|
|
|
valid_token = UserAPIKeyAuth(
|
|
team_id=_team_id,
|
|
user_id=team_admin,
|
|
token=hash_token(user_key),
|
|
last_refreshed_at=time.time(),
|
|
)
|
|
user_api_key_cache.set_cache(key=hash_token(user_key), value=valid_token)
|
|
|
|
team_obj = LiteLLM_TeamTableCachedObj(
|
|
team_id=_team_id,
|
|
blocked=False,
|
|
last_refreshed_at=time.time(),
|
|
members_with_roles=[Member(role=user_role, user_id=team_admin)],
|
|
metadata={"guardrails": {"modify_guardrails": False}},
|
|
)
|
|
|
|
user_api_key_cache.set_cache(key="team_id:{}".format(_team_id), value=team_obj)
|
|
|
|
setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
|
|
if new_member_method == "user_id":
|
|
data = {
|
|
"team_id": _team_id,
|
|
"member": [{"role": "user", "user_id": user}],
|
|
}
|
|
elif new_member_method == "user_email":
|
|
data = {
|
|
"team_id": _team_id,
|
|
"member": [{"role": "user", "user_email": user}],
|
|
}
|
|
team_member_add_request = TeamMemberAddRequest(**data)
|
|
|
|
with (
|
|
patch(
|
|
"litellm.proxy.proxy_server.prisma_client.db.litellm_usertable",
|
|
new_callable=AsyncMock,
|
|
) as mock_litellm_usertable,
|
|
patch(
|
|
"litellm.proxy.auth.auth_checks._get_team_object_from_user_api_key_cache",
|
|
new=AsyncMock(return_value=team_obj),
|
|
) as mock_team_obj,
|
|
patch(
|
|
"litellm.proxy.proxy_server.prisma_client.get_data",
|
|
new=AsyncMock(return_value=[]),
|
|
) as mock_get_data,
|
|
):
|
|
mock_client = AsyncMock(
|
|
return_value=LiteLLM_UserTable(
|
|
user_id="1234", max_budget=100, user_email="1234"
|
|
)
|
|
)
|
|
mock_litellm_usertable.upsert = mock_client
|
|
mock_litellm_usertable.find_many = AsyncMock(return_value=[])
|
|
# Mock find_first for user_email validation (returns None for new users)
|
|
mock_litellm_usertable.find_first = AsyncMock(return_value=None)
|
|
# Mock find_unique for user_id validation (returns None for new users)
|
|
mock_litellm_usertable.find_unique = AsyncMock(return_value=None)
|
|
|
|
team_mock_client = AsyncMock()
|
|
team_mock_client.update = AsyncMock(
|
|
return_value=LiteLLM_TeamTableCachedObj(team_id="1234")
|
|
)
|
|
membership_mock_client = AsyncMock()
|
|
membership_mock_client.upsert = AsyncMock(
|
|
return_value=LiteLLM_TeamMembership(user_id="1234", team_id=_team_id)
|
|
)
|
|
|
|
tx_cm = _member_add_tx_cm(team_mock_client)
|
|
|
|
with (
|
|
patch.object(
|
|
litellm.proxy.proxy_server.prisma_client.db,
|
|
"litellm_teamtable",
|
|
team_mock_client,
|
|
),
|
|
patch.object( # test-quality-ok: legacy test swaps the prisma table on the module-level client
|
|
litellm.proxy.proxy_server.prisma_client.db,
|
|
"litellm_teammembership",
|
|
membership_mock_client,
|
|
),
|
|
patch.object(
|
|
litellm.proxy.proxy_server.prisma_client,
|
|
"tx",
|
|
MagicMock(return_value=tx_cm),
|
|
),
|
|
):
|
|
error = None
|
|
try:
|
|
await team_member_add(
|
|
data=team_member_add_request,
|
|
user_api_key_dict=valid_token,
|
|
)
|
|
except HTTPException as e:
|
|
error = e
|
|
|
|
if error is not None:
|
|
assert user_role == "user" or new_member_method == "user_id"
|
|
assert error.status_code == 403
|
|
return
|
|
|
|
mock_client.assert_called()
|
|
|
|
assert (
|
|
mock_client.call_args.kwargs["data"]["create"]["max_budget"]
|
|
== litellm.max_internal_user_budget
|
|
)
|
|
assert (
|
|
mock_client.call_args.kwargs["data"]["create"]["budget_duration"]
|
|
== litellm.internal_user_budget_duration
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).")
|
|
async def test_user_info_team_list(prisma_client): # noqa: F811 # pytest fixture, not a redefinition
|
|
"""Assert user_info for admin calls team_list function"""
|
|
from litellm.proxy._types import LiteLLM_UserTable
|
|
|
|
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
|
|
setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
|
|
await litellm.proxy.proxy_server.prisma_client.connect()
|
|
|
|
from litellm.proxy.management_endpoints.internal_user_endpoints import user_info
|
|
|
|
with patch(
|
|
"litellm.proxy.management_endpoints.team_endpoints.list_team",
|
|
new_callable=AsyncMock,
|
|
) as mock_client:
|
|
|
|
prisma_client.get_data = AsyncMock(
|
|
return_value=LiteLLM_UserTable(
|
|
user_role="proxy_admin",
|
|
user_id="default_user_id",
|
|
max_budget=None,
|
|
user_email="",
|
|
)
|
|
)
|
|
|
|
try:
|
|
await user_info(
|
|
request=MagicMock(),
|
|
user_id=None,
|
|
user_api_key_dict=UserAPIKeyAuth(
|
|
api_key=MASTER_KEY, user_id="default_user_id"
|
|
),
|
|
)
|
|
except Exception:
|
|
pass
|
|
|
|
mock_client.assert_called()
|
|
|
|
|
|
@pytest.mark.skip(reason="Local test")
|
|
@pytest.mark.asyncio
|
|
async def test_add_callback_via_key(prisma_client): # noqa: F811 # pytest fixture, not a redefinition
|
|
"""
|
|
Test if callback specified in key, is used.
|
|
"""
|
|
global headers
|
|
import json
|
|
|
|
from fastapi import HTTPException, Request, Response
|
|
from starlette.datastructures import URL
|
|
|
|
from litellm.proxy.proxy_server import chat_completion
|
|
|
|
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
|
|
setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
|
|
await litellm.proxy.proxy_server.prisma_client.connect()
|
|
|
|
litellm.set_verbose = True
|
|
|
|
try:
|
|
# Your test data
|
|
test_data = {
|
|
"model": "azure/gpt-4.1-mini",
|
|
"messages": [
|
|
{"role": "user", "content": "write 1 sentence poem"},
|
|
],
|
|
"max_tokens": 10,
|
|
"mock_response": "Hello world",
|
|
"api_key": "my-fake-key",
|
|
}
|
|
|
|
request = Request(scope={"type": "http", "method": "POST", "headers": {}})
|
|
request._url = URL(url="/chat/completions")
|
|
|
|
json_bytes = json.dumps(test_data).encode("utf-8")
|
|
|
|
request._body = json_bytes
|
|
|
|
with patch.object(
|
|
litellm.litellm_core_utils.litellm_logging,
|
|
"LangFuseLogger",
|
|
new=MagicMock(),
|
|
) as mock_client:
|
|
resp = await chat_completion(
|
|
request=request,
|
|
fastapi_response=Response(),
|
|
user_api_key_dict=UserAPIKeyAuth(
|
|
metadata={
|
|
"allow_client_mock_response": True,
|
|
"logging": [
|
|
{
|
|
"callback_name": "langfuse", # 'otel', 'langfuse', 'lunary'
|
|
"callback_type": "success", # set, if required by integration - future improvement, have logging tools work for success + failure by default
|
|
"callback_vars": {
|
|
"langfuse_public_key": "os.environ/LANGFUSE_PUBLIC_KEY",
|
|
"langfuse_secret_key": "os.environ/LANGFUSE_SECRET_KEY",
|
|
"langfuse_host": "https://us.cloud.langfuse.com",
|
|
},
|
|
}
|
|
],
|
|
}
|
|
),
|
|
)
|
|
print(resp)
|
|
mock_client.assert_called()
|
|
mock_client.return_value.log_event.assert_called()
|
|
args, kwargs = mock_client.return_value.log_event.call_args
|
|
kwargs = kwargs["kwargs"]
|
|
assert "user_api_key_metadata" in kwargs["litellm_params"]["metadata"]
|
|
assert (
|
|
"logging"
|
|
in kwargs["litellm_params"]["metadata"]["user_api_key_metadata"]
|
|
)
|
|
checked_keys = False
|
|
for item in kwargs["litellm_params"]["metadata"]["user_api_key_metadata"][
|
|
"logging"
|
|
]:
|
|
for k, v in item["callback_vars"].items():
|
|
print("k={}, v={}".format(k, v))
|
|
if "key" in k:
|
|
assert "os.environ" in v
|
|
checked_keys = True
|
|
|
|
assert checked_keys
|
|
except Exception as e:
|
|
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"callback_type, expected_success_callbacks, expected_failure_callbacks",
|
|
[
|
|
("success", ["langfuse"], []),
|
|
("failure", [], ["langfuse"]),
|
|
("success_and_failure", ["langfuse"], ["langfuse"]),
|
|
],
|
|
)
|
|
async def test_add_callback_via_key_litellm_pre_call_utils(
|
|
mock_prisma_client,
|
|
callback_type,
|
|
expected_success_callbacks,
|
|
expected_failure_callbacks,
|
|
):
|
|
import json
|
|
|
|
from fastapi import HTTPException, Request, Response
|
|
from starlette.datastructures import URL
|
|
|
|
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
|
|
|
|
setattr(litellm.proxy.proxy_server, "prisma_client", mock_prisma_client)
|
|
setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
|
|
|
|
proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config")
|
|
|
|
request = Request(scope={"type": "http", "method": "POST", "headers": {}})
|
|
request._url = URL(url="/chat/completions")
|
|
|
|
test_data = {
|
|
"model": "azure/gpt-4.1-mini",
|
|
"messages": [
|
|
{"role": "user", "content": "write 1 sentence poem"},
|
|
],
|
|
"max_tokens": 10,
|
|
"mock_response": "Hello world",
|
|
"api_key": "my-fake-key",
|
|
}
|
|
|
|
json_bytes = json.dumps(test_data).encode("utf-8")
|
|
|
|
request._body = json_bytes
|
|
|
|
data = {
|
|
"data": {
|
|
"model": "azure/gpt-4.1-mini",
|
|
"messages": [{"role": "user", "content": "write 1 sentence poem"}],
|
|
"max_tokens": 10,
|
|
"mock_response": "Hello world",
|
|
"api_key": "my-fake-key",
|
|
},
|
|
"request": request,
|
|
"user_api_key_dict": UserAPIKeyAuth(
|
|
token=None,
|
|
key_name=None,
|
|
key_alias=None,
|
|
spend=0.0,
|
|
max_budget=None,
|
|
expires=None,
|
|
models=[],
|
|
aliases={},
|
|
config={},
|
|
user_id=None,
|
|
team_id=None,
|
|
max_parallel_requests=None,
|
|
metadata={
|
|
"allow_client_mock_response": True,
|
|
"logging": [
|
|
{
|
|
"callback_name": "langfuse",
|
|
"callback_type": callback_type,
|
|
"callback_vars": {
|
|
"langfuse_public_key": "my-mock-public-key",
|
|
"langfuse_secret_key": "my-mock-secret-key",
|
|
"langfuse_host": "https://us.cloud.langfuse.com",
|
|
},
|
|
}
|
|
],
|
|
},
|
|
tpm_limit=None,
|
|
rpm_limit=None,
|
|
budget_duration=None,
|
|
budget_reset_at=None,
|
|
allowed_cache_controls=[],
|
|
permissions={},
|
|
model_spend={},
|
|
model_max_budget={},
|
|
soft_budget_cooldown=False,
|
|
litellm_budget_table=None,
|
|
org_id=None,
|
|
team_spend=None,
|
|
team_alias=None,
|
|
team_tpm_limit=None,
|
|
team_rpm_limit=None,
|
|
team_max_budget=None,
|
|
team_models=[],
|
|
team_blocked=False,
|
|
soft_budget=None,
|
|
team_model_aliases=None,
|
|
team_member_spend=None,
|
|
team_metadata=None,
|
|
end_user_id=None,
|
|
end_user_tpm_limit=None,
|
|
end_user_rpm_limit=None,
|
|
end_user_max_budget=None,
|
|
last_refreshed_at=None,
|
|
api_key=None,
|
|
user_role=None,
|
|
allowed_model_region=None,
|
|
parent_otel_span=None,
|
|
),
|
|
"proxy_config": proxy_config,
|
|
"general_settings": {},
|
|
"version": "0.0.0",
|
|
}
|
|
|
|
new_data = await add_litellm_data_to_request(**data)
|
|
print("NEW DATA: {}".format(new_data))
|
|
|
|
assert "langfuse_public_key" in new_data
|
|
assert new_data["langfuse_public_key"] == "my-mock-public-key"
|
|
assert "langfuse_secret_key" in new_data
|
|
assert new_data["langfuse_secret_key"] == "my-mock-secret-key"
|
|
|
|
if expected_success_callbacks:
|
|
assert "success_callback" in new_data
|
|
assert new_data["success_callback"] == expected_success_callbacks
|
|
|
|
if expected_failure_callbacks:
|
|
assert "failure_callback" in new_data
|
|
assert new_data["failure_callback"] == expected_failure_callbacks
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"disable_fallbacks_set",
|
|
[
|
|
True,
|
|
False,
|
|
],
|
|
)
|
|
async def test_disable_fallbacks_by_key(disable_fallbacks_set):
|
|
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
|
|
|
|
key_metadata = {"disable_fallbacks": disable_fallbacks_set}
|
|
existing_data = {
|
|
"model": "azure/gpt-4.1-mini",
|
|
"messages": [{"role": "user", "content": "write 1 sentence poem"}],
|
|
}
|
|
data = LiteLLMProxyRequestSetup.add_key_level_controls(
|
|
key_metadata=key_metadata,
|
|
data=existing_data,
|
|
_metadata_variable_name="metadata",
|
|
)
|
|
|
|
assert data["disable_fallbacks"] == disable_fallbacks_set
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"callback_type, expected_success_callbacks, expected_failure_callbacks",
|
|
[
|
|
("success", ["gcs_bucket"], []),
|
|
("failure", [], ["gcs_bucket"]),
|
|
("success_and_failure", ["gcs_bucket"], ["gcs_bucket"]),
|
|
],
|
|
)
|
|
async def test_add_callback_via_key_litellm_pre_call_utils_gcs_bucket(
|
|
mock_prisma_client,
|
|
callback_type,
|
|
expected_success_callbacks,
|
|
expected_failure_callbacks,
|
|
):
|
|
import json
|
|
|
|
from fastapi import HTTPException, Request, Response
|
|
from starlette.datastructures import URL
|
|
|
|
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
|
|
|
|
setattr(litellm.proxy.proxy_server, "prisma_client", mock_prisma_client)
|
|
setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
|
|
|
|
proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config")
|
|
|
|
request = Request(scope={"type": "http", "method": "POST", "headers": {}})
|
|
request._url = URL(url="/chat/completions")
|
|
|
|
test_data = {
|
|
"model": "azure/gpt-4.1-mini",
|
|
"messages": [
|
|
{"role": "user", "content": "write 1 sentence poem"},
|
|
],
|
|
"max_tokens": 10,
|
|
"mock_response": "Hello world",
|
|
"api_key": "my-fake-key",
|
|
}
|
|
|
|
json_bytes = json.dumps(test_data).encode("utf-8")
|
|
|
|
request._body = json_bytes
|
|
|
|
data = {
|
|
"data": {
|
|
"model": "azure/gpt-4.1-mini",
|
|
"messages": [{"role": "user", "content": "write 1 sentence poem"}],
|
|
"max_tokens": 10,
|
|
"mock_response": "Hello world",
|
|
"api_key": "my-fake-key",
|
|
},
|
|
"request": request,
|
|
"user_api_key_dict": UserAPIKeyAuth(
|
|
token=None,
|
|
key_name=None,
|
|
key_alias=None,
|
|
spend=0.0,
|
|
max_budget=None,
|
|
expires=None,
|
|
models=[],
|
|
aliases={},
|
|
config={},
|
|
user_id=None,
|
|
team_id=None,
|
|
max_parallel_requests=None,
|
|
metadata={
|
|
"allow_client_mock_response": True,
|
|
"logging": [
|
|
{
|
|
"callback_name": "gcs_bucket",
|
|
"callback_type": callback_type,
|
|
"callback_vars": {
|
|
"gcs_bucket_name": "key-logging-project1",
|
|
"gcs_path_service_account": "pathrise-convert-1606954137718-a956eef1a2a8.json",
|
|
},
|
|
}
|
|
],
|
|
},
|
|
tpm_limit=None,
|
|
rpm_limit=None,
|
|
budget_duration=None,
|
|
budget_reset_at=None,
|
|
allowed_cache_controls=[],
|
|
permissions={},
|
|
model_spend={},
|
|
model_max_budget={},
|
|
soft_budget_cooldown=False,
|
|
litellm_budget_table=None,
|
|
org_id=None,
|
|
team_spend=None,
|
|
team_alias=None,
|
|
team_tpm_limit=None,
|
|
team_rpm_limit=None,
|
|
team_max_budget=None,
|
|
team_models=[],
|
|
team_blocked=False,
|
|
soft_budget=None,
|
|
team_model_aliases=None,
|
|
team_member_spend=None,
|
|
team_metadata=None,
|
|
end_user_id=None,
|
|
end_user_tpm_limit=None,
|
|
end_user_rpm_limit=None,
|
|
end_user_max_budget=None,
|
|
last_refreshed_at=None,
|
|
api_key=None,
|
|
user_role=None,
|
|
allowed_model_region=None,
|
|
parent_otel_span=None,
|
|
),
|
|
"proxy_config": proxy_config,
|
|
"general_settings": {},
|
|
"version": "0.0.0",
|
|
}
|
|
|
|
new_data = await add_litellm_data_to_request(**data)
|
|
print("NEW DATA: {}".format(new_data))
|
|
|
|
assert "gcs_bucket_name" in new_data
|
|
assert new_data["gcs_bucket_name"] == "key-logging-project1"
|
|
assert "gcs_path_service_account" in new_data
|
|
assert (
|
|
new_data["gcs_path_service_account"]
|
|
== "pathrise-convert-1606954137718-a956eef1a2a8.json"
|
|
)
|
|
|
|
if expected_success_callbacks:
|
|
assert "success_callback" in new_data
|
|
assert new_data["success_callback"] == expected_success_callbacks
|
|
|
|
if expected_failure_callbacks:
|
|
assert "failure_callback" in new_data
|
|
assert new_data["failure_callback"] == expected_failure_callbacks
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"callback_type, expected_success_callbacks, expected_failure_callbacks",
|
|
[
|
|
("success", ["langsmith"], []),
|
|
("failure", [], ["langsmith"]),
|
|
("success_and_failure", ["langsmith"], ["langsmith"]),
|
|
],
|
|
)
|
|
async def test_add_callback_via_key_litellm_pre_call_utils_langsmith(
|
|
mock_prisma_client,
|
|
callback_type,
|
|
expected_success_callbacks,
|
|
expected_failure_callbacks,
|
|
):
|
|
import json
|
|
|
|
from fastapi import HTTPException, Request, Response
|
|
from starlette.datastructures import URL
|
|
|
|
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
|
|
|
|
setattr(litellm.proxy.proxy_server, "prisma_client", mock_prisma_client)
|
|
setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
|
|
|
|
proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config")
|
|
|
|
request = Request(scope={"type": "http", "method": "POST", "headers": {}})
|
|
request._url = URL(url="/chat/completions")
|
|
|
|
test_data = {
|
|
"model": "azure/gpt-4.1-mini",
|
|
"messages": [
|
|
{"role": "user", "content": "write 1 sentence poem"},
|
|
],
|
|
"max_tokens": 10,
|
|
"mock_response": "Hello world",
|
|
"api_key": "my-fake-key",
|
|
}
|
|
|
|
json_bytes = json.dumps(test_data).encode("utf-8")
|
|
|
|
request._body = json_bytes
|
|
|
|
data = {
|
|
"data": {
|
|
"model": "azure/gpt-4.1-mini",
|
|
"messages": [{"role": "user", "content": "write 1 sentence poem"}],
|
|
"max_tokens": 10,
|
|
"mock_response": "Hello world",
|
|
"api_key": "my-fake-key",
|
|
},
|
|
"request": request,
|
|
"user_api_key_dict": UserAPIKeyAuth(
|
|
token=None,
|
|
key_name=None,
|
|
key_alias=None,
|
|
spend=0.0,
|
|
max_budget=None,
|
|
expires=None,
|
|
models=[],
|
|
aliases={},
|
|
config={},
|
|
user_id=None,
|
|
team_id=None,
|
|
max_parallel_requests=None,
|
|
metadata={
|
|
"allow_client_mock_response": True,
|
|
"logging": [
|
|
{
|
|
"callback_name": "langsmith",
|
|
"callback_type": callback_type,
|
|
"callback_vars": {
|
|
"langsmith_api_key": "ls-1234",
|
|
"langsmith_project": "pr-brief-resemblance-72",
|
|
"langsmith_base_url": "https://api.smith.langchain.com",
|
|
},
|
|
}
|
|
],
|
|
},
|
|
tpm_limit=None,
|
|
rpm_limit=None,
|
|
budget_duration=None,
|
|
budget_reset_at=None,
|
|
allowed_cache_controls=[],
|
|
permissions={},
|
|
model_spend={},
|
|
model_max_budget={},
|
|
soft_budget_cooldown=False,
|
|
litellm_budget_table=None,
|
|
org_id=None,
|
|
team_spend=None,
|
|
team_alias=None,
|
|
team_tpm_limit=None,
|
|
team_rpm_limit=None,
|
|
team_max_budget=None,
|
|
team_models=[],
|
|
team_blocked=False,
|
|
soft_budget=None,
|
|
team_model_aliases=None,
|
|
team_member_spend=None,
|
|
team_metadata=None,
|
|
end_user_id=None,
|
|
end_user_tpm_limit=None,
|
|
end_user_rpm_limit=None,
|
|
end_user_max_budget=None,
|
|
last_refreshed_at=None,
|
|
api_key=None,
|
|
user_role=None,
|
|
allowed_model_region=None,
|
|
parent_otel_span=None,
|
|
),
|
|
"proxy_config": proxy_config,
|
|
"general_settings": {},
|
|
"version": "0.0.0",
|
|
}
|
|
|
|
new_data = await add_litellm_data_to_request(**data)
|
|
print("NEW DATA: {}".format(new_data))
|
|
|
|
assert "langsmith_api_key" in new_data
|
|
assert new_data["langsmith_api_key"] == "ls-1234"
|
|
assert "langsmith_project" in new_data
|
|
assert new_data["langsmith_project"] == "pr-brief-resemblance-72"
|
|
assert "langsmith_base_url" in new_data
|
|
assert new_data["langsmith_base_url"] == "https://api.smith.langchain.com"
|
|
|
|
if expected_success_callbacks:
|
|
assert "success_callback" in new_data
|
|
assert new_data["success_callback"] == expected_success_callbacks
|
|
|
|
if expected_failure_callbacks:
|
|
assert "failure_callback" in new_data
|
|
assert new_data["failure_callback"] == expected_failure_callbacks
|
|
|
|
|
|
@pytest.mark.parametrize("hidden", [True, False])
|
|
@pytest.mark.asyncio
|
|
async def test_model_info_alias_without_prisma(hidden):
|
|
from litellm.proxy.proxy_server import model_info_v1
|
|
|
|
_model_list = [
|
|
{
|
|
"model_name": "gpt-3.5-turbo",
|
|
"litellm_params": {"model": "gpt-3.5-turbo"},
|
|
}
|
|
]
|
|
|
|
model_alias = "gpt-4"
|
|
|
|
router = litellm.Router(
|
|
model_list=_model_list,
|
|
model_group_alias={
|
|
model_alias: {
|
|
"model": "gpt-3.5-turbo",
|
|
"hidden": hidden,
|
|
}
|
|
},
|
|
)
|
|
|
|
setattr(litellm.proxy.proxy_server, "llm_router", router)
|
|
setattr(litellm.proxy.proxy_server, "llm_model_list", _model_list)
|
|
setattr(litellm.proxy.proxy_server, "prisma_client", None)
|
|
|
|
resp = await model_info_v1(
|
|
user_api_key_dict=UserAPIKeyAuth(models=[]),
|
|
)
|
|
|
|
models = json.loads(resp.body)["data"]
|
|
|
|
alias_found = any(
|
|
m["model_name"] == model_alias
|
|
for m in models
|
|
)
|
|
|
|
assert alias_found is (not hidden)
|
|
|
|
|
|
@pytest.mark.parametrize("hidden", [True, False])
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).")
|
|
async def test_proxy_model_group_alias_checks(prisma_client, hidden): # noqa: F811 # pytest fixture, not a redefinition
|
|
"""
|
|
Check if model group alias is returned on
|
|
|
|
`/v1/models`
|
|
`/v1/model/info`
|
|
`/v1/model_group/info`
|
|
"""
|
|
import json
|
|
|
|
from fastapi import HTTPException, Request, Response
|
|
from starlette.datastructures import URL
|
|
|
|
from litellm.proxy.proxy_server import model_group_info, model_info_v1, model_list
|
|
|
|
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
|
|
setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
|
|
await litellm.proxy.proxy_server.prisma_client.connect()
|
|
|
|
proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config")
|
|
|
|
_model_list = [
|
|
{
|
|
"model_name": "gpt-3.5-turbo",
|
|
"litellm_params": {"model": "gpt-3.5-turbo"},
|
|
}
|
|
]
|
|
model_alias = "gpt-4"
|
|
router = litellm.Router(
|
|
model_list=_model_list,
|
|
model_group_alias={model_alias: {"model": "gpt-3.5-turbo", "hidden": hidden}},
|
|
)
|
|
setattr(litellm.proxy.proxy_server, "llm_router", router)
|
|
setattr(litellm.proxy.proxy_server, "llm_model_list", _model_list)
|
|
|
|
request = Request(scope={"type": "http", "method": "POST", "headers": {}})
|
|
request._url = URL(url="/v1/models")
|
|
|
|
resp = await model_list(
|
|
user_api_key_dict=UserAPIKeyAuth(models=[]),
|
|
)
|
|
|
|
if hidden:
|
|
assert len(resp["data"]) == 1
|
|
else:
|
|
assert len(resp["data"]) == 2
|
|
print(resp)
|
|
|
|
resp = await model_info_v1(
|
|
user_api_key_dict=UserAPIKeyAuth(models=[]),
|
|
)
|
|
models = json.loads(resp.body)["data"]
|
|
is_model_alias_in_list = False
|
|
for item in models:
|
|
if model_alias == item["model_name"]:
|
|
is_model_alias_in_list = True
|
|
|
|
if hidden:
|
|
assert is_model_alias_in_list is False
|
|
else:
|
|
assert is_model_alias_in_list
|
|
|
|
resp = await model_group_info(
|
|
user_api_key_dict=UserAPIKeyAuth(models=[]),
|
|
)
|
|
print(f"resp: {resp}")
|
|
models = resp["data"]
|
|
is_model_alias_in_list = False
|
|
print(f"model_alias: {model_alias}, models: {models}")
|
|
for item in models:
|
|
if model_alias == item.model_group:
|
|
is_model_alias_in_list = True
|
|
|
|
if hidden:
|
|
assert is_model_alias_in_list is False
|
|
else:
|
|
assert is_model_alias_in_list, f"models: {models}"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).")
|
|
async def test_proxy_model_group_info_rerank(prisma_client): # noqa: F811 # pytest fixture, not a redefinition
|
|
"""
|
|
Check if rerank model is returned on the following endpoints
|
|
|
|
`/v1/models`
|
|
`/v1/model/info`
|
|
`/v1/model_group/info`
|
|
"""
|
|
import json
|
|
|
|
from fastapi import HTTPException, Request, Response
|
|
from starlette.datastructures import URL
|
|
|
|
from litellm.proxy.proxy_server import model_group_info, model_info_v1, model_list
|
|
|
|
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
|
|
setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
|
|
await litellm.proxy.proxy_server.prisma_client.connect()
|
|
|
|
proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config")
|
|
|
|
_model_list = [
|
|
{
|
|
"model_name": "rerank-english-v3.0",
|
|
"litellm_params": {"model": "cohere/rerank-english-v3.0"},
|
|
"model_info": {
|
|
"mode": "rerank",
|
|
},
|
|
}
|
|
]
|
|
router = litellm.Router(model_list=_model_list)
|
|
setattr(litellm.proxy.proxy_server, "llm_router", router)
|
|
setattr(litellm.proxy.proxy_server, "llm_model_list", _model_list)
|
|
|
|
request = Request(scope={"type": "http", "method": "POST", "headers": {}})
|
|
request._url = URL(url="/v1/models")
|
|
|
|
resp = await model_list(
|
|
user_api_key_dict=UserAPIKeyAuth(models=[]),
|
|
)
|
|
|
|
assert len(resp["data"]) == 1
|
|
print(resp)
|
|
|
|
resp = await model_info_v1(
|
|
user_api_key_dict=UserAPIKeyAuth(models=[]),
|
|
)
|
|
models = json.loads(resp.body)["data"]
|
|
assert models[0]["model_info"]["mode"] == "rerank"
|
|
resp = await model_group_info(
|
|
user_api_key_dict=UserAPIKeyAuth(models=[]),
|
|
)
|
|
|
|
print(resp)
|
|
models = resp["data"]
|
|
assert models[0].mode == "rerank"
|
|
|
|
|
|
# @pytest.mark.asyncio
|
|
# async def test_proxy_team_member_add(prisma_client):
|
|
# """
|
|
# Add 10 people to a team. Confirm all 10 are added.
|
|
# """
|
|
# from litellm.proxy.management_endpoints.team_endpoints import (
|
|
# team_member_add,
|
|
# new_team,
|
|
# )
|
|
# from litellm.proxy._types import TeamMemberAddRequest, Member, NewTeamRequest
|
|
|
|
# setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
|
|
# setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
|
|
# try:
|
|
|
|
# async def test():
|
|
# await litellm.proxy.proxy_server.prisma_client.connect()
|
|
# from litellm.proxy.proxy_server import user_api_key_cache
|
|
|
|
# user_api_key_dict = UserAPIKeyAuth(
|
|
# user_role=LitellmUserRoles.PROXY_ADMIN,
|
|
# api_key=MASTER_KEY,
|
|
# user_id="1234",
|
|
# )
|
|
|
|
# new_team()
|
|
# for _ in range(10):
|
|
# request = TeamMemberAddRequest(
|
|
# team_id="1234",
|
|
# member=Member(
|
|
# user_id="1234",
|
|
# user_role=LitellmUserRoles.INTERNAL_USER,
|
|
# ),
|
|
# )
|
|
# key = await team_member_add(
|
|
# request, user_api_key_dict=user_api_key_dict
|
|
# )
|
|
|
|
# print(key)
|
|
# user_id = key.user_id
|
|
|
|
# # check /user/info to verify user_role was set correctly
|
|
# new_user_info = await user_info(
|
|
# user_id=user_id, user_api_key_dict=user_api_key_dict
|
|
# )
|
|
# new_user_info = new_user_info.user_info
|
|
# print("new_user_info=", new_user_info)
|
|
# assert new_user_info["user_role"] == LitellmUserRoles.INTERNAL_USER
|
|
# assert new_user_info["user_id"] == user_id
|
|
|
|
# generated_key = key.key
|
|
# bearer_token = "Bearer " + generated_key
|
|
|
|
# assert generated_key not in user_api_key_cache.in_memory_cache.cache_dict
|
|
|
|
# value_from_prisma = await prisma_client.get_data(
|
|
# token=generated_key,
|
|
# )
|
|
# print("token from prisma", value_from_prisma)
|
|
|
|
# request = Request(
|
|
# {
|
|
# "type": "http",
|
|
# "route": api_route,
|
|
# "path": api_route.path,
|
|
# "headers": [("Authorization", bearer_token)],
|
|
# }
|
|
# )
|
|
|
|
# # use generated key to auth in
|
|
# result = await user_api_key_auth(request=request, api_key=bearer_token)
|
|
# print("result from user auth with new key", result)
|
|
|
|
# asyncio.run(test())
|
|
# except Exception as e:
|
|
# pytest.fail(f"An exception occurred - {str(e)}")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_server_prisma_setup():
|
|
from litellm.proxy.proxy_server import ProxyStartupEvent
|
|
from litellm.proxy.utils import ProxyLogging
|
|
from litellm.caching import DualCache
|
|
|
|
user_api_key_cache = DualCache()
|
|
|
|
with patch.object(
|
|
litellm.proxy.proxy_server, "PrismaClient", new=MagicMock()
|
|
) as mock_prisma_client:
|
|
mock_client = mock_prisma_client.return_value # This is the mocked instance
|
|
mock_client.connect = AsyncMock() # Mock the connect method
|
|
mock_client.start_view_setup_task = MagicMock()
|
|
mock_client.health_check = AsyncMock() # Mock the health_check method
|
|
mock_client.start_db_health_watchdog_task = AsyncMock()
|
|
# Mock the db attribute with start_token_refresh_task for RDS IAM token refresh
|
|
mock_db = MagicMock()
|
|
mock_db.start_token_refresh_task = AsyncMock()
|
|
mock_client.db = mock_db
|
|
|
|
prisma_client = await ProxyStartupEvent._setup_prisma_client(
|
|
database_url="postgresql://user:pass@localhost:5432/litellm",
|
|
proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache),
|
|
user_api_key_cache=user_api_key_cache,
|
|
)
|
|
|
|
assert prisma_client is mock_client
|
|
mock_client.connect.assert_called_once()
|
|
mock_client.start_view_setup_task.assert_called_once()
|
|
|
|
# Note: This is REALLY IMPORTANT to check that the health check is called
|
|
# This is how we ensure the DB is ready before proceeding
|
|
mock_client.health_check.assert_called_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_server_prisma_setup_invalid_db(monkeypatch):
|
|
"""
|
|
PROD TEST: Test that proxy server startup fails when it's unable to connect to the database
|
|
|
|
Think 2-3 times before editing / deleting this test, it's important for PROD
|
|
"""
|
|
import httpx
|
|
|
|
from litellm.proxy.proxy_server import ProxyStartupEvent
|
|
from litellm.proxy.utils import ProxyLogging
|
|
from litellm.caching import DualCache
|
|
|
|
user_api_key_cache = DualCache()
|
|
invalid_db_url = "postgresql://invalid:invalid@localhost:5432/nonexistent"
|
|
|
|
monkeypatch.setenv("DATABASE_URL", invalid_db_url)
|
|
|
|
with pytest.raises(httpx.ConnectError):
|
|
await ProxyStartupEvent._setup_prisma_client(
|
|
database_url=invalid_db_url,
|
|
proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache),
|
|
user_api_key_cache=user_api_key_cache,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_ui_settings_spend_logs_threshold():
|
|
"""
|
|
Test that get_ui_settings correctly sets DISABLE_EXPENSIVE_DB_QUERIES based on spend_logs_row_count threshold
|
|
"""
|
|
from litellm.proxy.management_endpoints.ui_sso import get_ui_settings
|
|
from litellm.proxy.proxy_server import proxy_state
|
|
from fastapi import Request
|
|
from litellm.constants import MAX_SPENDLOG_ROWS_TO_QUERY
|
|
|
|
# Create a mock request
|
|
mock_request = Request(
|
|
scope={
|
|
"type": "http",
|
|
"headers": [],
|
|
"method": "GET",
|
|
"scheme": "http",
|
|
"server": ("testserver", 80),
|
|
"path": "/sso/get/ui_settings",
|
|
"query_string": b"",
|
|
}
|
|
)
|
|
|
|
# Test case 1: When spend_logs_row_count > MAX_SPENDLOG_ROWS_TO_QUERY
|
|
proxy_state.set_proxy_state_variable(
|
|
"spend_logs_row_count", MAX_SPENDLOG_ROWS_TO_QUERY + 1
|
|
)
|
|
response = await get_ui_settings(mock_request)
|
|
print("response from get_ui_settings", json.dumps(response, indent=4))
|
|
assert response["DISABLE_EXPENSIVE_DB_QUERIES"] is True
|
|
assert response["NUM_SPEND_LOGS_ROWS"] == MAX_SPENDLOG_ROWS_TO_QUERY + 1
|
|
|
|
# Test case 2: When spend_logs_row_count < MAX_SPENDLOG_ROWS_TO_QUERY
|
|
proxy_state.set_proxy_state_variable(
|
|
"spend_logs_row_count", MAX_SPENDLOG_ROWS_TO_QUERY - 1
|
|
)
|
|
response = await get_ui_settings(mock_request)
|
|
print("response from get_ui_settings", json.dumps(response, indent=4))
|
|
assert response["DISABLE_EXPENSIVE_DB_QUERIES"] is False
|
|
assert response["NUM_SPEND_LOGS_ROWS"] == MAX_SPENDLOG_ROWS_TO_QUERY - 1
|
|
|
|
# Test case 3: Edge case - exactly MAX_SPENDLOG_ROWS_TO_QUERY
|
|
proxy_state.set_proxy_state_variable(
|
|
"spend_logs_row_count", MAX_SPENDLOG_ROWS_TO_QUERY
|
|
)
|
|
response = await get_ui_settings(mock_request)
|
|
print("response from get_ui_settings", json.dumps(response, indent=4))
|
|
assert response["DISABLE_EXPENSIVE_DB_QUERIES"] is False
|
|
assert response["NUM_SPEND_LOGS_ROWS"] == MAX_SPENDLOG_ROWS_TO_QUERY
|
|
|
|
# Clean up
|
|
proxy_state.set_proxy_state_variable("spend_logs_row_count", 0)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_background_health_check_reflects_llm_model_list(monkeypatch):
|
|
"""
|
|
Test that _run_background_health_check reflects changes to llm_model_list in each health check iteration.
|
|
"""
|
|
import litellm.proxy.proxy_server as proxy_server
|
|
import copy
|
|
|
|
test_model_list_1 = [{"model_name": "model-a"}]
|
|
test_model_list_2 = [{"model_name": "model-b"}]
|
|
called_model_lists = []
|
|
|
|
async def fake_perform_health_check(model_list, details, max_concurrency=None):
|
|
called_model_lists.append(copy.deepcopy(model_list))
|
|
return (["healthy"], ["unhealthy"], {})
|
|
|
|
monkeypatch.setattr(proxy_server, "health_check_interval", 1)
|
|
monkeypatch.setattr(proxy_server, "health_check_details", None)
|
|
monkeypatch.setattr(
|
|
proxy_server, "llm_model_list", copy.deepcopy(test_model_list_1)
|
|
)
|
|
monkeypatch.setattr(proxy_server, "perform_health_check", fake_perform_health_check)
|
|
monkeypatch.setattr(proxy_server, "health_check_results", {})
|
|
|
|
async def fake_sleep(interval):
|
|
raise asyncio.CancelledError()
|
|
|
|
monkeypatch.setattr(asyncio, "sleep", fake_sleep)
|
|
|
|
try:
|
|
await proxy_server._run_background_health_check()
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
monkeypatch.setattr(
|
|
proxy_server, "llm_model_list", copy.deepcopy(test_model_list_2)
|
|
)
|
|
|
|
try:
|
|
await proxy_server._run_background_health_check()
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
assert len(called_model_lists) >= 2
|
|
assert called_model_lists[0] == test_model_list_1
|
|
assert called_model_lists[1] == test_model_list_2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_background_health_check_skip_disabled_models(monkeypatch):
|
|
"""Ensure models with disable_background_health_check are skipped."""
|
|
import litellm.proxy.proxy_server as proxy_server
|
|
import copy
|
|
|
|
test_model_list = [
|
|
{"model_name": "model-a"},
|
|
{
|
|
"model_name": "model-b",
|
|
"model_info": {"disable_background_health_check": True},
|
|
},
|
|
]
|
|
called_model_lists = []
|
|
|
|
async def fake_perform_health_check(
|
|
model_list, details, max_concurrency=None, **kwargs
|
|
):
|
|
called_model_lists.append(copy.deepcopy(model_list))
|
|
return (["healthy"], [], {})
|
|
|
|
monkeypatch.setattr(proxy_server, "health_check_interval", 1)
|
|
monkeypatch.setattr(proxy_server, "health_check_details", None)
|
|
monkeypatch.setattr(proxy_server, "llm_model_list", copy.deepcopy(test_model_list))
|
|
monkeypatch.setattr(proxy_server, "perform_health_check", fake_perform_health_check)
|
|
monkeypatch.setattr(proxy_server, "health_check_results", {})
|
|
|
|
async def fake_sleep(interval):
|
|
raise asyncio.CancelledError()
|
|
|
|
monkeypatch.setattr(asyncio, "sleep", fake_sleep)
|
|
|
|
try:
|
|
await proxy_server._run_background_health_check()
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
assert called_model_lists == [[{"model_name": "model-a"}]]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_direct_health_check_with_instrumentation_legacy_three_arg_stub(
|
|
monkeypatch,
|
|
):
|
|
"""Monkeypatched perform_health_check with only base kwargs should still run."""
|
|
import litellm.proxy.proxy_server as proxy_server
|
|
|
|
async def fake_perform_health_check(model_list, details, max_concurrency=None):
|
|
return ([], [], {})
|
|
|
|
monkeypatch.setattr(proxy_server, "perform_health_check", fake_perform_health_check)
|
|
result = await proxy_server._run_direct_health_check_with_instrumentation(
|
|
[{"model_name": "m"}],
|
|
True,
|
|
1,
|
|
{"enabled": True, "source": "test", "cycle_id": "c1"},
|
|
)
|
|
assert result == ([], [], {})
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_direct_health_check_with_instrumentation_accepts_instrumentation_only(
|
|
monkeypatch,
|
|
):
|
|
"""Stub that accepts instrumentation_context but not health_check filter kwargs."""
|
|
import litellm.proxy.proxy_server as proxy_server
|
|
|
|
seen: list = []
|
|
|
|
async def fake_perform_health_check(
|
|
model_list, details, max_concurrency=None, instrumentation_context=None
|
|
):
|
|
seen.append(instrumentation_context)
|
|
return ([], [], {})
|
|
|
|
monkeypatch.setattr(proxy_server, "perform_health_check", fake_perform_health_check)
|
|
await proxy_server._run_direct_health_check_with_instrumentation(
|
|
[],
|
|
False,
|
|
2,
|
|
{"enabled": True, "source": "test", "cycle_id": "c2"},
|
|
)
|
|
assert len(seen) == 1
|
|
assert seen[0]["cycle_id"] == "c2"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_direct_health_check_with_instrumentation_accepts_filter_only(
|
|
monkeypatch,
|
|
):
|
|
"""Stub that accepts health_check_skip_disabled_background_models but not instrumentation."""
|
|
import litellm.proxy.proxy_server as proxy_server
|
|
|
|
seen: list = []
|
|
|
|
async def fake_perform_health_check(
|
|
model_list,
|
|
details,
|
|
max_concurrency=None,
|
|
health_check_skip_disabled_background_models=False,
|
|
):
|
|
seen.append(health_check_skip_disabled_background_models)
|
|
return ([], [], {})
|
|
|
|
monkeypatch.setattr(proxy_server, "perform_health_check", fake_perform_health_check)
|
|
await proxy_server._run_direct_health_check_with_instrumentation(
|
|
[],
|
|
True,
|
|
None,
|
|
{"enabled": False},
|
|
)
|
|
assert len(seen) == 1
|
|
assert seen[0] is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_direct_health_check_drops_only_the_rejected_kwarg(monkeypatch):
|
|
"""A callee that predates `router` must still get the skip-disabled filter: dropping the
|
|
rejected argument alongside working ones would probe deployments the operator opted out."""
|
|
import litellm.proxy.proxy_server as proxy_server
|
|
|
|
seen: list[tuple[dict[str, str] | None, bool]] = []
|
|
|
|
async def fake_perform_health_check(
|
|
model_list,
|
|
details,
|
|
max_concurrency=None,
|
|
instrumentation_context=None,
|
|
health_check_skip_disabled_background_models=False,
|
|
):
|
|
seen.append((instrumentation_context, health_check_skip_disabled_background_models))
|
|
return ([], [], {})
|
|
|
|
monkeypatch.setattr(proxy_server, "perform_health_check", fake_perform_health_check)
|
|
monkeypatch.setattr(
|
|
proxy_server,
|
|
"general_settings",
|
|
{"health_check_skip_disabled_background_models": True},
|
|
)
|
|
await proxy_server._run_direct_health_check_with_instrumentation([], True, 1, {"cycle_id": "c3"})
|
|
|
|
assert seen == [({"cycle_id": "c3"}, True)]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_direct_health_check_with_instrumentation_non_kw_typeerror_reraises(
|
|
monkeypatch,
|
|
):
|
|
import litellm.proxy.proxy_server as proxy_server
|
|
|
|
async def fake_perform_health_check(**kwargs):
|
|
raise TypeError("unsupported operand type(s)")
|
|
|
|
monkeypatch.setattr(proxy_server, "perform_health_check", fake_perform_health_check)
|
|
with pytest.raises(TypeError, match="unsupported operand"):
|
|
await proxy_server._run_direct_health_check_with_instrumentation(
|
|
[],
|
|
True,
|
|
1,
|
|
{},
|
|
)
|
|
|
|
|
|
def test_get_timeout_from_request():
|
|
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
|
|
|
|
headers = {
|
|
"x-litellm-timeout": "90",
|
|
}
|
|
timeout = LiteLLMProxyRequestSetup._get_timeout_from_request(headers)
|
|
assert timeout == 90
|
|
|
|
headers = {
|
|
"x-litellm-timeout": "90.5",
|
|
}
|
|
timeout = LiteLLMProxyRequestSetup._get_timeout_from_request(headers)
|
|
assert timeout == 90.5
|
|
|
|
|
|
def test_add_litellm_data_for_backend_llm_call_marks_client_side_timeout():
|
|
"""A caller-supplied x-litellm-timeout must be marked with client_side_timeout=True,
|
|
so the router's fallback-cooldown trigger can tell it apart from a deployment
|
|
actually timing out (a caller could otherwise force every deployment in a fallback
|
|
chain to look unhealthy with a single near-zero timeout request)."""
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
|
|
|
|
user_api_key_dict = UserAPIKeyAuth(api_key="test_api_key")
|
|
|
|
data = LiteLLMProxyRequestSetup.add_litellm_data_for_backend_llm_call(
|
|
headers={"x-litellm-timeout": "0.001"},
|
|
request_data={},
|
|
user_api_key_dict=user_api_key_dict,
|
|
)
|
|
assert data["timeout"] == 0.001
|
|
assert data["client_side_timeout"] is True
|
|
|
|
data_without_header = LiteLLMProxyRequestSetup.add_litellm_data_for_backend_llm_call(
|
|
headers={},
|
|
request_data={},
|
|
user_api_key_dict=user_api_key_dict,
|
|
)
|
|
assert "client_side_timeout" not in data_without_header
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"request_data",
|
|
[
|
|
{"timeout": 0.001},
|
|
{"request_timeout": 0.001},
|
|
{"stream_timeout": 0.001},
|
|
],
|
|
)
|
|
def test_add_litellm_data_for_backend_llm_call_marks_client_side_timeout_from_body(
|
|
request_data,
|
|
):
|
|
"""Router._get_timeout resolves the effective timeout from kwargs["timeout"],
|
|
kwargs["request_timeout"], or kwargs["stream_timeout"], and a caller can supply any
|
|
of those directly in the request body, not just via the x-litellm-timeout header.
|
|
Missing this would let a caller force a 408 on every deployment in a fallback chain
|
|
without it being recognized as caller-controlled, cooling down deployments other
|
|
tenants rely on."""
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
|
|
|
|
user_api_key_dict = UserAPIKeyAuth(api_key="test_api_key")
|
|
|
|
data = LiteLLMProxyRequestSetup.add_litellm_data_for_backend_llm_call(
|
|
headers={},
|
|
request_data=request_data,
|
|
user_api_key_dict=user_api_key_dict,
|
|
)
|
|
assert data["client_side_timeout"] is True
|
|
|
|
|
|
def test_add_litellm_data_for_backend_llm_call_ignores_forged_client_side_timeout():
|
|
"""The caller-supplied client_side_timeout key itself must never be trusted verbatim:
|
|
the marker is always recomputed from the actual timeout sources, so a caller can't
|
|
forge client_side_timeout=True to dodge cooldown on a real deployment failure."""
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
|
|
|
|
user_api_key_dict = UserAPIKeyAuth(api_key="test_api_key")
|
|
|
|
data = LiteLLMProxyRequestSetup.add_litellm_data_for_backend_llm_call(
|
|
headers={},
|
|
request_data={"client_side_timeout": True},
|
|
user_api_key_dict=user_api_key_dict,
|
|
)
|
|
assert "client_side_timeout" not in data
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"ui_exists, ui_has_content",
|
|
[
|
|
(True, True), # UI path exists and has content
|
|
(True, False), # UI path exists but is empty
|
|
(False, False), # UI path doesn't exist
|
|
],
|
|
)
|
|
def test_non_root_ui_path_logic(monkeypatch, tmp_path, ui_exists, ui_has_content):
|
|
"""
|
|
Test the non-root Docker UI path detection logic.
|
|
|
|
Tests that when LITELLM_NON_ROOT is set to "true":
|
|
- If UI path exists and has content, it should be used
|
|
- If UI path doesn't exist or is empty, proper error logging occurs
|
|
"""
|
|
import tempfile
|
|
import shutil
|
|
from unittest.mock import MagicMock
|
|
|
|
# Create a temporary directory to act as /tmp/litellm_ui
|
|
test_ui_path = tmp_path / "litellm_ui"
|
|
|
|
if ui_exists:
|
|
test_ui_path.mkdir(parents=True, exist_ok=True)
|
|
if ui_has_content:
|
|
# Create some dummy files to simulate built UI
|
|
(test_ui_path / "index.html").write_text("<html></html>")
|
|
(test_ui_path / "app.js").write_text("console.log('test');")
|
|
|
|
# Mock the environment variable and os.path operations
|
|
monkeypatch.setenv("LITELLM_NON_ROOT", "true")
|
|
|
|
# Create a mock logger to capture log messages
|
|
mock_logger = MagicMock()
|
|
|
|
# We need to reimport or reload the relevant code section
|
|
# Since this is module-level code, we'll test the logic directly
|
|
ui_path = None
|
|
non_root_ui_path = str(test_ui_path)
|
|
|
|
# Simulate the logic from proxy_server.py lines 909-920
|
|
if os.getenv("LITELLM_NON_ROOT", "").lower() == "true":
|
|
if os.path.exists(non_root_ui_path) and os.listdir(non_root_ui_path):
|
|
mock_logger.info(
|
|
f"Using pre-built UI for non-root Docker: {non_root_ui_path}"
|
|
)
|
|
mock_logger.info(
|
|
f"UI files found: {len(os.listdir(non_root_ui_path))} items"
|
|
)
|
|
ui_path = non_root_ui_path
|
|
else:
|
|
mock_logger.error(
|
|
f"UI not found at {non_root_ui_path}. UI will not be available."
|
|
)
|
|
mock_logger.error(
|
|
f"Path exists: {os.path.exists(non_root_ui_path)}, Has content: {os.path.exists(non_root_ui_path) and bool(os.listdir(non_root_ui_path))}"
|
|
)
|
|
|
|
# Verify behavior based on test parameters
|
|
if ui_exists and ui_has_content:
|
|
# UI should be found and used
|
|
assert ui_path == non_root_ui_path
|
|
assert mock_logger.info.call_count == 2
|
|
mock_logger.info.assert_any_call(
|
|
f"Using pre-built UI for non-root Docker: {non_root_ui_path}"
|
|
)
|
|
# Verify the second info call mentions the number of items
|
|
info_calls = [call[0][0] for call in mock_logger.info.call_args_list]
|
|
assert any("UI files found:" in call and "items" in call for call in info_calls)
|
|
assert mock_logger.error.call_count == 0
|
|
else:
|
|
# UI should not be found, error should be logged
|
|
assert ui_path is None
|
|
assert mock_logger.error.call_count == 2
|
|
mock_logger.error.assert_any_call(
|
|
f"UI not found at {non_root_ui_path}. UI will not be available."
|
|
)
|
|
# Verify the second error call has path existence info
|
|
error_calls = [call[0][0] for call in mock_logger.error.call_args_list]
|
|
assert any("Path exists:" in call for call in error_calls)
|
|
assert mock_logger.info.call_count == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_config_callbacks_with_all_types(client_no_auth):
|
|
"""
|
|
Test that /get/config/callbacks returns all three callback types:
|
|
- success_callback with type="success"
|
|
- failure_callback with type="failure"
|
|
- callbacks (success_and_failure) with type="success_and_failure"
|
|
"""
|
|
from litellm.proxy.proxy_server import ProxyConfig
|
|
|
|
# Create a mock config with all three callback types
|
|
mock_config_data = {
|
|
"litellm_settings": {
|
|
"success_callback": ["langfuse", "braintrust"],
|
|
"failure_callback": ["sentry"],
|
|
"callbacks": ["otel", "langsmith"],
|
|
},
|
|
"environment_variables": {
|
|
"LANGFUSE_PUBLIC_KEY": "test-public-key",
|
|
"LANGFUSE_SECRET_KEY": "test-secret-key",
|
|
"LANGFUSE_HOST": "https://test.langfuse.com",
|
|
"BRAINTRUST_API_KEY": "test-braintrust-key",
|
|
"OTEL_EXPORTER": "otlp",
|
|
"OTEL_ENDPOINT": "http://localhost:4317",
|
|
"LANGSMITH_API_KEY": "test-langsmith-key",
|
|
},
|
|
"general_settings": {},
|
|
}
|
|
|
|
proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config")
|
|
|
|
with patch.object(
|
|
proxy_config, "get_config", new=AsyncMock(return_value=mock_config_data)
|
|
):
|
|
response = client_no_auth.get("/get/config/callbacks")
|
|
|
|
assert response.status_code == 200
|
|
result = response.json()
|
|
|
|
# Verify response structure
|
|
assert "status" in result
|
|
assert result["status"] == "success"
|
|
assert "callbacks" in result
|
|
|
|
callbacks = [cb for cb in result["callbacks"] if not cb.get("read_only", False)]
|
|
|
|
# Verify we have all 5 callbacks (2 success + 1 failure + 2 success_and_failure)
|
|
assert len(callbacks) == 5
|
|
|
|
# Group callbacks by type
|
|
success_callbacks = [cb for cb in callbacks if cb.get("type") == "success"]
|
|
failure_callbacks = [cb for cb in callbacks if cb.get("type") == "failure"]
|
|
success_and_failure_callbacks = [
|
|
cb for cb in callbacks if cb.get("type") == "success_and_failure"
|
|
]
|
|
|
|
# Verify all callbacks have required fields
|
|
for callback in callbacks:
|
|
assert "name" in callback
|
|
assert "variables" in callback
|
|
assert "type" in callback
|
|
assert callback["type"] in ["success", "failure", "success_and_failure"]
|
|
|
|
# Verify success callbacks
|
|
assert len(success_callbacks) == 2
|
|
success_names = [cb["name"] for cb in success_callbacks]
|
|
assert "langfuse" in success_names
|
|
assert "braintrust" in success_names
|
|
|
|
# Verify failure callbacks
|
|
assert len(failure_callbacks) == 1
|
|
assert failure_callbacks[0]["name"] == "sentry"
|
|
|
|
# Verify success_and_failure callbacks
|
|
assert len(success_and_failure_callbacks) == 2
|
|
success_and_failure_names = [cb["name"] for cb in success_and_failure_callbacks]
|
|
assert "otel" in success_and_failure_names
|
|
assert "langsmith" in success_and_failure_names
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_config_callbacks_environment_variables(client_no_auth):
|
|
"""
|
|
Test that /get/config/callbacks correctly includes environment variables
|
|
for each callback type. Under ``client_no_auth`` the resolved role is
|
|
not ``PROXY_ADMIN``, so values matched by the redaction helper come back
|
|
as ``"REDACTED"`` and other values pass through verbatim.
|
|
"""
|
|
from litellm.proxy.proxy_server import ProxyConfig
|
|
|
|
# Create a mock config with callbacks and their env vars
|
|
mock_config_data = {
|
|
"litellm_settings": {
|
|
"success_callback": ["langfuse"],
|
|
"failure_callback": [],
|
|
"callbacks": ["otel"],
|
|
},
|
|
"environment_variables": {
|
|
"LANGFUSE_PUBLIC_KEY": "test-public-key",
|
|
"LANGFUSE_SECRET_KEY": "test-secret-key",
|
|
"LANGFUSE_HOST": "https://cloud.langfuse.com",
|
|
"OTEL_EXPORTER": "otlp",
|
|
"OTEL_ENDPOINT": "http://localhost:4317",
|
|
"OTEL_HEADERS": "key=value",
|
|
},
|
|
"general_settings": {},
|
|
}
|
|
|
|
proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config")
|
|
|
|
with patch.object(
|
|
proxy_config, "get_config", new=AsyncMock(return_value=mock_config_data)
|
|
):
|
|
response = client_no_auth.get("/get/config/callbacks")
|
|
|
|
assert response.status_code == 200
|
|
result = response.json()
|
|
|
|
callbacks = result["callbacks"]
|
|
|
|
# Find langfuse callback (success type)
|
|
langfuse_callback = next(
|
|
(cb for cb in callbacks if cb["name"] == "langfuse"), None
|
|
)
|
|
assert langfuse_callback is not None
|
|
assert langfuse_callback["type"] == "success"
|
|
assert "variables" in langfuse_callback
|
|
|
|
langfuse_vars = langfuse_callback["variables"]
|
|
assert "LANGFUSE_PUBLIC_KEY" in langfuse_vars
|
|
assert langfuse_vars["LANGFUSE_PUBLIC_KEY"] == "REDACTED"
|
|
assert "LANGFUSE_SECRET_KEY" in langfuse_vars
|
|
assert langfuse_vars["LANGFUSE_SECRET_KEY"] == "REDACTED"
|
|
assert "LANGFUSE_HOST" in langfuse_vars
|
|
assert langfuse_vars["LANGFUSE_HOST"] == "https://cloud.langfuse.com"
|
|
|
|
# Find otel callback (success_and_failure type)
|
|
otel_callback = next((cb for cb in callbacks if cb["name"] == "otel"), None)
|
|
assert otel_callback is not None
|
|
assert otel_callback["type"] == "success_and_failure"
|
|
assert "variables" in otel_callback
|
|
|
|
otel_vars = otel_callback["variables"]
|
|
assert "OTEL_EXPORTER" in otel_vars
|
|
assert otel_vars["OTEL_EXPORTER"] == "otlp"
|
|
assert "OTEL_ENDPOINT" in otel_vars
|
|
assert otel_vars["OTEL_ENDPOINT"] == "http://localhost:4317"
|
|
assert "OTEL_HEADERS" in otel_vars
|
|
assert otel_vars["OTEL_HEADERS"] == "REDACTED"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_update_config_success_callback_normalization(monkeypatch):
|
|
"""
|
|
Ensure success_callback values are normalized to lowercase when updating config.
|
|
This prevents delete_callback (which searches lowercase) from failing on mixed case inputs like 'SQS'.
|
|
"""
|
|
import litellm.proxy.proxy_server as proxy_server
|
|
from litellm.proxy._types import ConfigYAML
|
|
|
|
monkeypatch.setattr(proxy_server, "proxy_logging_obj", MagicMock())
|
|
|
|
existing_litellm_settings = {"success_callback": ["langfuse"]}
|
|
|
|
class FakeRow:
|
|
def __init__(self, name, value):
|
|
self.param_name = name
|
|
self.param_value = value
|
|
|
|
upserted = {}
|
|
|
|
async def fake_find_first(where=None):
|
|
if where and where.get("param_name") == "litellm_settings":
|
|
return FakeRow("litellm_settings", existing_litellm_settings)
|
|
return None
|
|
|
|
async def fake_upsert(where=None, data=None):
|
|
upserted[where["param_name"]] = json.loads(data["update"]["param_value"])
|
|
|
|
class MockPrisma:
|
|
def __init__(self):
|
|
self.db = MagicMock()
|
|
self.db.litellm_config = MagicMock()
|
|
self.db.litellm_config.find_first = AsyncMock(side_effect=fake_find_first)
|
|
self.db.litellm_config.upsert = AsyncMock(side_effect=fake_upsert)
|
|
self.writer_db = self.db
|
|
|
|
monkeypatch.setattr(proxy_server, "prisma_client", MockPrisma())
|
|
|
|
class MockProxyConfig:
|
|
async def add_deployment(self, prisma_client=None, proxy_logging_obj=None): # noqa: F811 # pytest fixture, not a redefinition
|
|
return None
|
|
|
|
def reject_config_owned_writes(self, *, section_name, changed_keys):
|
|
return None
|
|
|
|
monkeypatch.setattr(proxy_server, "proxy_config", MockProxyConfig())
|
|
|
|
config_update = ConfigYAML(litellm_settings={"success_callback": ["SQS", "sQs"]})
|
|
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
|
|
|
admin_user = UserAPIKeyAuth(
|
|
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-test"
|
|
)
|
|
request = MagicMock()
|
|
request.json = AsyncMock(return_value={"litellm_settings": {"success_callback": ["SQS", "sQs"]}})
|
|
await proxy_server.update_config(config_update, request=request, user_api_key_dict=admin_user)
|
|
|
|
assert (
|
|
"litellm_settings" in upserted
|
|
), "litellm_config.upsert was not called for litellm_settings"
|
|
callbacks = upserted["litellm_settings"]["success_callback"]
|
|
|
|
# Deduped and normalized
|
|
assert "sqs" in callbacks
|
|
assert "SQS" not in callbacks
|
|
assert "sQs" not in callbacks
|
|
# Existing callback should still be present
|
|
assert "langfuse" in callbacks
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"data",
|
|
[
|
|
{
|
|
"model": {
|
|
"model_name": "azure/gpt-4.1-mini",
|
|
"litellm_params": {"model": "azure/gpt-4.1-mini"},
|
|
"model_info": {"base_model": "gpt-4.1-mini"},
|
|
},
|
|
"expected": "gpt-4.1-mini",
|
|
},
|
|
{
|
|
"model": {
|
|
"model_name": "openai/gpt-4.1-mini",
|
|
"litellm_params": {"model": "openai/gpt-4.1-mini"},
|
|
},
|
|
"expected": "openai/gpt-4.1-mini",
|
|
},
|
|
{
|
|
"model": {
|
|
"model_name": "openai/gpt-4.1-mini",
|
|
"litellm_params": {"model": "openai/gpt-4.1-mini"},
|
|
"model_info": {"base_model": "gpt-4.1-mini"},
|
|
},
|
|
"expected": "gpt-4.1-mini",
|
|
},
|
|
{
|
|
"model": {
|
|
"model_name": "claude-sonnet-4-5-20250929",
|
|
"litellm_params": {"model": "anthropic/claude-sonnet-4-5@20250929"},
|
|
"model_info": {"base_model": "anthropic/claude-sonnet-4-5-20250929"},
|
|
},
|
|
"expected": "anthropic/claude-sonnet-4-5-20250929",
|
|
},
|
|
{
|
|
"model": {
|
|
"model_name": "gemini-2.5-flash-001",
|
|
"litellm_params": {"model": "gemini/gemini-2.5-flash@001"},
|
|
"model_info": {"base_model": "gemini-2.5-flash-001"},
|
|
},
|
|
"expected": "gemini-2.5-flash-001",
|
|
},
|
|
],
|
|
)
|
|
def test_get_litellm_model_info(data):
|
|
from litellm.proxy.proxy_server import get_litellm_model_info
|
|
|
|
model = data["model"]
|
|
get_info_mock = MagicMock()
|
|
|
|
with mock.patch(
|
|
"litellm.get_model_info",
|
|
new=get_info_mock,
|
|
):
|
|
get_litellm_model_info(model=model)
|
|
get_info_mock.assert_called_once_with(data["expected"])
|
|
|
|
|
|
def test_live_session_accept_failure_releases_reservations(
|
|
client_no_auth: TestClient,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
from litellm.proxy import proxy_server
|
|
|
|
released: Final[list[str]] = []
|
|
|
|
async def release_budget(user_api_key_dict: object) -> None:
|
|
released.append("budget")
|
|
|
|
async def release_slot(user_api_key_dict: object) -> None:
|
|
released.append("slot")
|
|
|
|
async def receive() -> dict[str, object]:
|
|
return {"type": "websocket.connect"}
|
|
|
|
async def send_failing_accept(message: Mapping[str, object]) -> None:
|
|
if message["type"] == "websocket.accept":
|
|
raise RuntimeError("accept failed")
|
|
|
|
scope: Final = {
|
|
"type": "websocket",
|
|
"path": "/v1/live/sessions",
|
|
"raw_path": b"/v1/live/sessions",
|
|
"root_path": "",
|
|
"scheme": "ws",
|
|
"query_string": b"",
|
|
"headers": [(b"host", b"testserver"), (b"authorization", b"Bearer test-api-key")],
|
|
"client": ("testclient", 50000),
|
|
"server": ("testserver", 80),
|
|
"subprotocols": [],
|
|
}
|
|
route_call: Final = AsyncMock()
|
|
monkeypatch.setattr(proxy_server, "route_request", route_call)
|
|
monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", release_budget)
|
|
monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", release_slot)
|
|
|
|
with pytest.raises(RuntimeError, match="accept failed"):
|
|
asyncio.run(client_no_auth.app(scope, receive, send_failing_accept))
|
|
|
|
assert released == ["budget", "slot"]
|
|
route_call.assert_not_awaited()
|
|
|
|
|
|
def _override_live_session_authorizer(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
authorizer: Callable[[WebSocket, str], Awaitable[UserAPIKeyAuth]],
|
|
) -> None:
|
|
from litellm.proxy import proxy_server
|
|
|
|
monkeypatch.setitem(
|
|
proxy_server.app.dependency_overrides,
|
|
proxy_server.live_session_model_authorizer,
|
|
lambda: authorizer,
|
|
)
|
|
|
|
|
|
def _start_live_session_and_wait_for_reply(client: TestClient) -> None:
|
|
with client.websocket_connect("/v1/live/sessions", headers={"Authorization": "Bearer test-api-key"}) as websocket:
|
|
websocket.send_json({"type": "session.start", "session": {"model": "gpt-3.5-turbo"}})
|
|
websocket.receive_text()
|
|
|
|
|
|
def test_live_session_model_check_error_releases_reservations(
|
|
client_no_auth: TestClient,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
from litellm.proxy import proxy_server
|
|
|
|
released: Final[list[str]] = []
|
|
|
|
async def release_budget(user_api_key_dict: object) -> None:
|
|
released.append("budget")
|
|
|
|
async def release_slot(user_api_key_dict: object) -> None:
|
|
released.append("slot")
|
|
|
|
async def failing_authorizer(websocket: object, model: str) -> UserAPIKeyAuth:
|
|
raise RuntimeError("agent lookup failed")
|
|
|
|
route_call: Final = AsyncMock()
|
|
monkeypatch.setattr(proxy_server, "route_request", route_call)
|
|
_override_live_session_authorizer(monkeypatch, failing_authorizer)
|
|
monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", release_budget)
|
|
monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", release_slot)
|
|
|
|
with pytest.raises(RuntimeError, match="agent lookup failed"):
|
|
_start_live_session_and_wait_for_reply(client_no_auth)
|
|
|
|
assert released == ["budget", "slot"]
|
|
route_call.assert_not_awaited()
|
|
|
|
|
|
class _PreCallAborted(BaseException):
|
|
pass
|
|
|
|
|
|
def test_live_session_pre_call_abort_releases_reservations(
|
|
client_no_auth: TestClient,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
from litellm.proxy import proxy_server
|
|
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
|
|
|
frame_auth: Final = UserAPIKeyAuth(token="frame-authorized")
|
|
released_for_frame: Final[list[str]] = []
|
|
|
|
async def release_budget(user_api_key_dict: UserAPIKeyAuth) -> None:
|
|
if user_api_key_dict is frame_auth:
|
|
released_for_frame.append("budget")
|
|
|
|
async def release_slot(user_api_key_dict: UserAPIKeyAuth) -> None:
|
|
if user_api_key_dict is frame_auth:
|
|
released_for_frame.append("slot")
|
|
|
|
async def authorize(websocket: object, model: str) -> UserAPIKeyAuth:
|
|
return frame_auth
|
|
|
|
class _AbortingPreCallProcessing(ProxyBaseLLMRequestProcessing):
|
|
async def common_processing_pre_call_logic(self, *args: object, **kwargs: object) -> NoReturn:
|
|
raise _PreCallAborted
|
|
|
|
route_call: Final = AsyncMock()
|
|
monkeypatch.setattr(proxy_server, "route_request", route_call)
|
|
_override_live_session_authorizer(monkeypatch, authorize)
|
|
monkeypatch.setattr(proxy_server, "ProxyBaseLLMRequestProcessing", _AbortingPreCallProcessing)
|
|
monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", release_budget)
|
|
monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", release_slot)
|
|
|
|
with pytest.raises(_PreCallAborted):
|
|
_start_live_session_and_wait_for_reply(client_no_auth)
|
|
|
|
assert released_for_frame == ["budget", "slot"]
|
|
route_call.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.parametrize("path", ["/openai/v1/live/sessions", "/v1/live/sessions", "/live/sessions"])
|
|
def test_live_session_invalid_start_is_rejected_without_routing(
|
|
path: str,
|
|
client_no_auth: TestClient,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
from litellm.proxy import proxy_server
|
|
from starlette.websockets import WebSocketDisconnect
|
|
|
|
async def no_op(user_api_key_dict: object) -> None:
|
|
return None
|
|
|
|
route_call: Final = AsyncMock()
|
|
monkeypatch.setattr(proxy_server, "route_request", route_call)
|
|
monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", no_op)
|
|
monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", no_op)
|
|
|
|
with client_no_auth.websocket_connect(path, headers={"Authorization": "Bearer test-api-key"}) as websocket:
|
|
websocket.send_text('{"type":"session.start","session":{"model":""}}')
|
|
assert websocket.receive_json() == {
|
|
"type": "error",
|
|
"error": {
|
|
"type": "invalid_request_error",
|
|
"code": "invalid_session_start",
|
|
"message": "First message must be a session.start JSON object with a non-empty session.model.",
|
|
},
|
|
}
|
|
with pytest.raises(WebSocketDisconnect) as error:
|
|
websocket.receive_json()
|
|
assert error.value.code == 1008
|
|
|
|
route_call.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("denial_type", "expected_code"),
|
|
(("user_model_access_denied", "model_not_allowed"), ("budget_exceeded", "session_not_allowed")),
|
|
)
|
|
def test_live_session_frame_authorization_denial_closes_with_policy_code(
|
|
client_no_auth: TestClient,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
denial_type: str,
|
|
expected_code: str,
|
|
) -> None:
|
|
from litellm.proxy import proxy_server
|
|
from litellm.proxy._types import ProxyException
|
|
from starlette.websockets import WebSocketDisconnect
|
|
|
|
async def no_op(user_api_key_dict: object) -> None:
|
|
return None
|
|
|
|
authorized_models: Final[list[str]] = []
|
|
|
|
async def deny(websocket: object, model: str) -> UserAPIKeyAuth:
|
|
authorized_models.append(model)
|
|
raise ProxyException("model access denied for this user", denial_type, "model", 403)
|
|
|
|
route_call: Final = AsyncMock()
|
|
monkeypatch.setattr(proxy_server, "route_request", route_call)
|
|
monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", no_op)
|
|
monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", no_op)
|
|
_override_live_session_authorizer(monkeypatch, deny)
|
|
|
|
with client_no_auth.websocket_connect(
|
|
"/v1/live/sessions",
|
|
headers={"Authorization": "Bearer test-api-key"},
|
|
) as websocket:
|
|
websocket.send_json({"type": "session.start", "session": {"model": "gpt-4o"}})
|
|
assert websocket.receive_json() == {
|
|
"type": "error",
|
|
"error": {
|
|
"type": "invalid_request_error",
|
|
"code": expected_code,
|
|
"message": "model access denied for this user",
|
|
},
|
|
}
|
|
with pytest.raises(WebSocketDisconnect) as error:
|
|
websocket.receive_json()
|
|
assert error.value.code == 1008
|
|
|
|
assert authorized_models == ["gpt-4o"]
|
|
route_call.assert_not_awaited()
|
|
|
|
|
|
def test_live_session_valid_start_routes_with_model_and_frame(
|
|
client_no_auth: TestClient,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
from types import SimpleNamespace
|
|
|
|
from litellm.proxy import proxy_server
|
|
from starlette.websockets import WebSocket
|
|
|
|
async def no_op(user_api_key_dict: object) -> None:
|
|
return None
|
|
|
|
async def complete_call() -> None:
|
|
return None
|
|
|
|
captured: Final = SimpleNamespace(data=None, route_type=None)
|
|
|
|
async def route(
|
|
*,
|
|
data: dict[str, object],
|
|
route_type: str,
|
|
llm_router: object,
|
|
user_model: str | None,
|
|
) -> Coroutine[object, object, None]:
|
|
captured.data = data
|
|
captured.route_type = route_type
|
|
await cast(WebSocket, data["websocket"]).send_text("routed")
|
|
return complete_call()
|
|
|
|
frame_auth: Final = UserAPIKeyAuth(token="frame-authorized")
|
|
|
|
async def authorize(websocket: object, model: str) -> UserAPIKeyAuth:
|
|
return frame_auth
|
|
|
|
monkeypatch.setattr(proxy_server, "route_request", route)
|
|
_override_live_session_authorizer(monkeypatch, authorize)
|
|
monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", no_op)
|
|
monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", no_op)
|
|
start_frame: Final[dict[str, object]] = {
|
|
"type": "session.start",
|
|
"session": {"model": "gpt-3.5-turbo", "instructions": "Stay brief"},
|
|
}
|
|
|
|
with client_no_auth.websocket_connect(
|
|
"/v1/live/sessions",
|
|
headers={"Authorization": "Bearer test-api-key"},
|
|
) as websocket:
|
|
websocket.send_json(start_frame)
|
|
assert websocket.receive_text() == "routed"
|
|
|
|
routed_data: Final = cast(dict[str, object], captured.data)
|
|
assert captured.route_type == "_arealtime"
|
|
assert routed_data["model"] == "gpt-3.5-turbo"
|
|
assert routed_data["live_session_start"] == start_frame
|
|
assert routed_data["user_api_key_dict"] is frame_auth
|
|
|
|
|
|
def test_live_session_ignores_query_model_and_authorizes_frame_model(
|
|
client_no_auth: TestClient,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
from fastapi import Request
|
|
|
|
from litellm.proxy import proxy_server
|
|
from starlette.websockets import WebSocket
|
|
|
|
async def no_op(user_api_key_dict: object) -> None:
|
|
return None
|
|
|
|
async def complete_call() -> None:
|
|
return None
|
|
|
|
async def route(
|
|
*,
|
|
data: dict[str, object],
|
|
route_type: str,
|
|
llm_router: object,
|
|
user_model: str | None,
|
|
) -> Coroutine[object, object, None]:
|
|
await cast(WebSocket, data["websocket"]).send_text("routed")
|
|
return complete_call()
|
|
|
|
upgrade_auth: Final = UserAPIKeyAuth(token="upgrade", models=["gpt-live-1"])
|
|
|
|
async def authenticate(request: Request, api_key: str) -> UserAPIKeyAuth:
|
|
if await request.body() != b'{"model": ""}':
|
|
raise HTTPException(status_code=403, detail="Model is not allowed")
|
|
return upgrade_auth
|
|
|
|
auth_call: Final = AsyncMock(side_effect=authenticate)
|
|
monkeypatch.setattr("litellm.proxy.auth.user_api_key_auth.user_api_key_auth", auth_call)
|
|
frame_auth: Final = UserAPIKeyAuth(token="frame-authorized", models=["gpt-live-1"])
|
|
authorize: Final = AsyncMock(return_value=frame_auth)
|
|
monkeypatch.setattr(proxy_server, "route_request", route)
|
|
monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", no_op)
|
|
monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", no_op)
|
|
_override_live_session_authorizer(monkeypatch, authorize)
|
|
|
|
with client_no_auth.websocket_connect(
|
|
"/v1/live/sessions?model=gpt-realtime",
|
|
headers={"Authorization": "Bearer test-api-key"},
|
|
) as websocket:
|
|
websocket.send_json({"type": "session.start", "session": {"model": "gpt-live-1"}})
|
|
assert websocket.receive_text() == "routed"
|
|
|
|
request_arg: Final = auth_call.await_args.kwargs["request"]
|
|
assert json.loads(asyncio.run(request_arg.body())) == {"model": ""}
|
|
assert authorize.await_args.args[1] == "gpt-live-1"
|
|
|
|
|
|
def test_live_session_reauthorized_token_still_needs_the_model_on_the_key(
|
|
client_no_auth: TestClient,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
from starlette.websockets import WebSocketDisconnect
|
|
|
|
from litellm.proxy import proxy_server
|
|
|
|
released_tokens: Final[list[str | None]] = []
|
|
|
|
async def release_budget(user_api_key_dict: UserAPIKeyAuth) -> None:
|
|
released_tokens.append(user_api_key_dict.token)
|
|
|
|
async def no_op(user_api_key_dict: object) -> None:
|
|
return None
|
|
|
|
async def custom_auth_without_common_checks(websocket: object, model: str) -> UserAPIKeyAuth:
|
|
return UserAPIKeyAuth(token="custom-auth", models=["gpt-4o"])
|
|
|
|
route_call: Final = AsyncMock()
|
|
monkeypatch.setattr(proxy_server, "route_request", route_call)
|
|
monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", release_budget)
|
|
monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", no_op)
|
|
_override_live_session_authorizer(monkeypatch, custom_auth_without_common_checks)
|
|
|
|
with client_no_auth.websocket_connect(
|
|
"/v1/live/sessions",
|
|
headers={"Authorization": "Bearer test-api-key"},
|
|
) as websocket:
|
|
websocket.send_json({"type": "session.start", "session": {"model": "gpt-3.5-turbo"}})
|
|
error_event: Final = websocket.receive_json()
|
|
with pytest.raises(WebSocketDisconnect) as error:
|
|
websocket.receive_json()
|
|
|
|
assert error_event["error"]["code"] == "model_not_allowed", error_event
|
|
assert error.value.code == 1008
|
|
assert "custom-auth" in released_tokens
|
|
route_call.assert_not_awaited()
|
|
|
|
|
|
def test_live_session_releases_the_upgrade_reservation_before_the_frame_reserves(
|
|
client_no_auth: TestClient,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
from starlette.websockets import WebSocket
|
|
|
|
from litellm.proxy import proxy_server
|
|
|
|
released_reservations: Final[list[object]] = []
|
|
|
|
async def release_budget(user_api_key_dict: UserAPIKeyAuth) -> None:
|
|
released_reservations.append(user_api_key_dict.budget_reservation)
|
|
|
|
async def no_op(user_api_key_dict: object) -> None:
|
|
return None
|
|
|
|
async def complete_call() -> None:
|
|
return None
|
|
|
|
async def route(
|
|
*,
|
|
data: dict[str, object],
|
|
route_type: str,
|
|
llm_router: object,
|
|
user_model: str | None,
|
|
) -> Coroutine[object, object, None]:
|
|
await cast(WebSocket, data["websocket"]).send_text("routed")
|
|
return complete_call()
|
|
|
|
frame_auth: Final = UserAPIKeyAuth(token="frame-authorized", budget_reservation={"reserved_for": "frame model"})
|
|
|
|
async def upgrade_auth() -> UserAPIKeyAuth:
|
|
return UserAPIKeyAuth(token="upgrade", budget_reservation={"reserved_for": "query model"})
|
|
|
|
released_before_frame_auth: Final[list[tuple[object, ...]]] = []
|
|
|
|
async def authorize(websocket: object, model: str) -> UserAPIKeyAuth:
|
|
released_before_frame_auth.append(tuple(released_reservations))
|
|
return frame_auth
|
|
|
|
monkeypatch.setattr(proxy_server, "route_request", route)
|
|
monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", release_budget)
|
|
monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", no_op)
|
|
monkeypatch.setitem(proxy_server.app.dependency_overrides, proxy_server.user_api_key_auth_live_websocket, upgrade_auth)
|
|
_override_live_session_authorizer(monkeypatch, authorize)
|
|
|
|
with client_no_auth.websocket_connect(
|
|
"/v1/live/sessions?model=gpt-4o",
|
|
headers={"Authorization": "Bearer test-api-key"},
|
|
) as websocket:
|
|
websocket.send_json({"type": "session.start", "session": {"model": "gpt-3.5-turbo"}})
|
|
assert websocket.receive_text() == "routed"
|
|
|
|
assert released_before_frame_auth == [({"reserved_for": "query model"},)], released_reservations
|
|
|
|
|
|
def test_live_session_preserves_pretty_printed_start_frame(
|
|
client_no_auth: TestClient,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
from types import SimpleNamespace
|
|
|
|
from starlette.websockets import WebSocket
|
|
|
|
from litellm.proxy import proxy_server
|
|
|
|
async def no_op(user_api_key_dict: object) -> None:
|
|
return None
|
|
|
|
async def complete_call() -> None:
|
|
return None
|
|
|
|
captured: Final = SimpleNamespace(data=None, route_type=None)
|
|
|
|
async def route(
|
|
*,
|
|
data: dict[str, object],
|
|
route_type: str,
|
|
llm_router: object,
|
|
user_model: str | None,
|
|
) -> Coroutine[object, object, None]:
|
|
captured.data = data
|
|
captured.route_type = route_type
|
|
await cast(WebSocket, data["websocket"]).send_text("routed")
|
|
return complete_call()
|
|
|
|
frame_auth: Final = UserAPIKeyAuth(token="frame-authorized")
|
|
|
|
async def authorize(websocket: object, model: str) -> UserAPIKeyAuth:
|
|
return frame_auth
|
|
|
|
monkeypatch.setattr(proxy_server, "route_request", route)
|
|
_override_live_session_authorizer(monkeypatch, authorize)
|
|
monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", no_op)
|
|
monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", no_op)
|
|
|
|
start_frame: Final = "\r\n".join(
|
|
[
|
|
"{",
|
|
' "type": "session.start",',
|
|
' "session": {"model": "gpt-live-1", "instructions": "a\\nb"}',
|
|
"}",
|
|
]
|
|
)
|
|
expected_start: Final[dict[str, object]] = {
|
|
"type": "session.start",
|
|
"session": {"model": "gpt-live-1", "instructions": "a\nb"},
|
|
}
|
|
|
|
with client_no_auth.websocket_connect(
|
|
"/v1/live/sessions",
|
|
headers={"Authorization": "Bearer test-api-key"},
|
|
) as websocket:
|
|
websocket.send_text(start_frame)
|
|
assert websocket.receive_text() == "routed"
|
|
|
|
routed_data: Final = cast(dict[str, object], captured.data)
|
|
assert captured.route_type == "_arealtime"
|
|
assert routed_data["model"] == "gpt-live-1"
|
|
assert routed_data["live_session_start"] == expected_start
|