scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-04 07:27:51 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 195

wip

Browse source
This commit is contained in:
Yujong Lee 2026-10-26 17:39:06 -07:00
parent ee2cb0b6e0
commit cebc7e6dd7
19 changed files with 391 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm-rust/Cargo.lock generated
View file

@ -4229,6 +4229,7 @@ dependencies = [
"litellm-cache-response",
"litellm-callbacks-legacy-python",
"litellm-core-utils",
"litellm-db",
"litellm-host",
"litellm-host-python",
"litellm-http",

26
litellm-rust/crates/db/.sqlx/query-136362927bf4b07894c7098e4f6ebd0f1c8d415ecdb3eb9f1a698e879dd270ae.json generated Normal file
View file

@ -0,0 +1,26 @@
{
"db_name": "PostgreSQL",
"query": "SELECT ui_settings\nFROM \"LiteLLM_UISettings\"\nWHERE id = 'user_banner'\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "ui_settings",
"type_info": "Jsonb",
"origin": {
"Table": {
"table": "\"LiteLLM_UISettings\"",
"name": "ui_settings"
}
}
}
],
"parameters": {
"Left": []
},
"nullable": [
false
]
},
"hash": "136362927bf4b07894c7098e4f6ebd0f1c8d415ecdb3eb9f1a698e879dd270ae"
}

14
litellm-rust/crates/db/.sqlx/query-c8da5047ab370f06f239ef8d5798dc13486a1f56d03400f54edf3f25be95cccc.json generated Normal file
View file

@ -0,0 +1,14 @@
{
"db_name": "PostgreSQL",
"query": "INSERT INTO \"LiteLLM_UISettings\" (id, ui_settings, updated_at)\nVALUES ('user_banner', $1, CURRENT_TIMESTAMP)\nON CONFLICT (id) DO UPDATE\nSET ui_settings = EXCLUDED.ui_settings,\n updated_at = EXCLUDED.updated_at\n",
"describe": {
"columns": [],
"parameters": {
"Left": [
"Jsonb"
]
},
"nullable": []
},
"hash": "c8da5047ab370f06f239ef8d5798dc13486a1f56d03400f54edf3f25be95cccc"
}

3
litellm-rust/crates/db/queries/user_banner/find_user_banner.sql Normal file
View file

@ -0,0 +1,3 @@
SELECT ui_settings
FROM "LiteLLM_UISettings"
WHERE id = 'user_banner'

5
litellm-rust/crates/db/queries/user_banner/save_user_banner.sql Normal file
View file

@ -0,0 +1,5 @@
INSERT INTO "LiteLLM_UISettings" (id, ui_settings, updated_at)
VALUES ('user_banner', $1, CURRENT_TIMESTAMP)
ON CONFLICT (id) DO UPDATE
SET ui_settings = EXCLUDED.ui_settings,
updated_at = EXCLUDED.updated_at

2
litellm-rust/crates/db/src/lib.rs
View file

@ -13,6 +13,7 @@ mod failure;
mod spend;
#[cfg(test)]
mod test_support;
mod user_banner;
pub use daily_spend::{DailyEndUserRow, DailyTeamRow, DailyUserRow};
pub use end_user::{EndUser, RateLimits, find_end_user};
@ -24,3 +25,4 @@ pub use spend::{
begin, insert_spend_logs, lock_team, upsert_daily_end_user_spend, upsert_daily_spend,
upsert_daily_team_spend, upsert_daily_user_spend,
};
pub use user_banner::{find_user_banner, save_user_banner};

114
litellm-rust/crates/db/src/user_banner.rs Normal file
View file

@ -0,0 +1,114 @@
use serde_json::Value;
use sqlx::PgExecutor;
use crate::Error;
pub async fn find_user_banner(executor: impl PgExecutor<'_>) -> Result<Option<Value>, Error> {
Ok(
sqlx::query_file_scalar!("queries/user_banner/find_user_banner.sql")
.fetch_optional(executor)
.await?,
)
}
pub async fn save_user_banner(
executor: impl PgExecutor<'_>,
settings: &Value,
) -> Result<(), Error> {
sqlx::query_file!("queries/user_banner/save_user_banner.sql", settings)
.execute(executor)
.await?;
Ok(())
}
#[cfg(test)]
mod tests {
use std::error::Error as StdError;
use rstest::rstest;
use serde_json::{Value, json};
use sqlx::PgPool;
use sqlx::types::chrono::NaiveDateTime;
use super::{find_user_banner, save_user_banner};
use crate::test_support::database_pool;
async fn clear(tx: &mut sqlx::PgConnection) -> Result<(), sqlx::Error> {
#[expect(
clippy::disallowed_methods,
reason = "Fixture cleanup is outside the compile-checked query set"
)]
sqlx::query(r#"DELETE FROM "LiteLLM_UISettings" WHERE id = 'user_banner'"#)
.execute(tx)
.await?;
Ok(())
}
async fn stored_row(
tx: &mut sqlx::PgConnection,
) -> Result<(Value, NaiveDateTime, NaiveDateTime), sqlx::Error> {
#[expect(
clippy::disallowed_methods,
reason = "Fixture reads are outside the compile-checked query set"
)]
sqlx::query_as(
r#"SELECT ui_settings, created_at, updated_at FROM "LiteLLM_UISettings" WHERE id = 'user_banner'"#,
)
.fetch_one(tx)
.await
}
#[rstest]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn missing_banner_reads_as_none(
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let mut tx = database_pool?.begin().await?;
clear(&mut tx).await?;
assert_eq!(find_user_banner(&mut *tx).await?, None);
Ok(())
}
#[rstest]
#[case::first_publish(None)]
#[case::republish(Some(json!({"enabled": false, "message": "old", "severity": "info", "revision": "a"})))]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn save_replaces_the_stored_banner_and_stamps_updated_at(
#[case] previous: Option<Value>,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let mut tx = database_pool?.begin().await?;
clear(&mut tx).await?;
if let Some(previous) = &previous {
save_user_banner(&mut *tx, previous).await?;
}
let banner = json!({"enabled": true, "message": "maintenance at 5pm", "severity": "warning", "revision": "b"});
save_user_banner(&mut *tx, &banner).await?;
assert_eq!(find_user_banner(&mut *tx).await?, Some(banner.clone()));
let (stored, created_at, updated_at) = stored_row(&mut tx).await?;
assert_eq!(stored, banner);
assert!(updated_at >= created_at);
Ok(())
}
#[rstest]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn banner_text_is_stored_as_data_not_sql(
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let mut tx = database_pool?.begin().await?;
clear(&mut tx).await?;
let banner = json!({"message": r#"'); DROP TABLE "LiteLLM_UISettings"; --"#});
save_user_banner(&mut *tx, &banner).await?;
assert_eq!(find_user_banner(&mut *tx).await?, Some(banner));
Ok(())
}
}

13
litellm-rust/crates/python-bridge/Cargo.toml
View file

@ -19,18 +19,14 @@ huggingface = ["litellm-token-counter/huggingface"]
tiktoken = ["litellm-token-counter/tiktoken"]
[dependencies]
fancy-regex.workspace = true
litellm-tracing.workspace = true
litellm-spend-clickhouse.workspace = true
litellm-storage-clickhouse.workspace = true
litellm-host.workspace = true
bytes.workspace = true
futures-util.workspace = true
litellm-cache.workspace = true
litellm-cache-memory.workspace = true
litellm-cache-redis.workspace = true
litellm-cache-response.workspace = true
serde.workspace = true
litellm-auth.workspace = true
litellm-auth-aws.workspace = true
litellm-callbacks-legacy-python.workspace = true
@ -41,6 +37,7 @@ litellm-inference-messages.workspace = true
litellm-inference-chat.workspace = true
litellm-inference-ocr.workspace = true
litellm-core-utils.workspace = true
litellm-db.workspace = true
litellm-http.workspace = true
litellm-llms.workspace = true
litellm-secrets = { workspace = true, features = ["aws", "azure", "google", "hashicorp", "cyberark"] }
@ -49,6 +46,11 @@ litellm-llms-types.workspace = true
litellm-router-types.workspace = true
litellm-host-python.workspace = true
litellm-token-counter = { path = "../token-counter", default-features = false }
serde.workspace = true
futures-util.workspace = true
bytes.workspace = true
fancy-regex.workspace = true
pyo3.workspace = true
prost.workspace = true
pyo3-async-runtimes.workspace = true
@ -61,9 +63,10 @@ thiserror.workspace = true
tokio = { workspace = true, features = ["rt", "sync"] }
[dev-dependencies]
tracing.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }
litellm-secrets-aws.workspace = true
tracing.workspace = true
serde.workspace = true
serde_with.workspace = true
criterion.workspace = true

68
litellm-rust/crates/python-bridge/src/database.rs Normal file
View file

@ -0,0 +1,68 @@
use std::error::Error as StdError;
use litellm_db::Error;
use pyo3::{exceptions::PyRuntimeError, prelude::*};
use serde_json::Value;
use sqlx::{PgPool, postgres::PgPoolOptions};
fn map_error(error: Error) -> PyErr {
let chain: Vec<String> =
std::iter::successors(Some(&error as &dyn StdError), |error| error.source())
.map(ToString::to_string)
.collect();
PyRuntimeError::new_err(chain.join(": "))
}
#[pyclass(frozen)]
pub struct NativeDatabase {
pool: PgPool,
}
#[pymethods]
impl NativeDatabase {
#[staticmethod]
fn connect(py: Python<'_>, url: String, max_connections: u32) -> PyResult<Bound<'_, PyAny>> {
crate::execution::run_async_value(py, async move {
let pool = PgPoolOptions::new()
.max_connections(max_connections)
.connect(&url)
.await
.map_err(|error| map_error(error.into()))?;
Ok(Self { pool })
})
}
fn find_user_banner<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let pool = self.pool.clone();
crate::execution::run_async(
py,
async move { litellm_db::find_user_banner(&pool).await },
map_error,
)
}
fn save_user_banner<'py>(
&self,
py: Python<'py>,
#[pyo3(from_py_with = litellm_host_python::from_py_argument)] settings: Value,
) -> PyResult<Bound<'py, PyAny>> {
let pool = self.pool.clone();
crate::execution::run_async(
py,
async move { litellm_db::save_user_banner(&pool, &settings).await },
map_error,
)
}
fn close<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let pool = self.pool.clone();
crate::execution::run_async(
py,
async move {
pool.close().await;
Ok::<(), Error>(())
},
map_error,
)
}
}

4
litellm-rust/crates/python-bridge/src/lib.rs
View file

@ -2,6 +2,7 @@ mod cache;
mod callable;
mod coercion;
mod credentials;
mod database;
mod diagnostics;
mod errors;
mod execution;
@ -20,6 +21,8 @@ mod _native {
use litellm_host_python::{ForkedAfterNativeRuntimeStarted, ProcessReservedForForking};
use pyo3::{prelude::*, types::PyModule};
#[pymodule_export]
use crate::database::NativeDatabase;
#[cfg(feature = "panic-test")]
#[pymodule_export]
use crate::diagnostics::_panic_for_test;
@ -87,6 +90,7 @@ mod tests {
Python::attach(|py| {
let mut expected = vec![
"NativeCacheHandle",
"NativeDatabase",
"RustBridgeDeclined",
"RustUpstreamError",
"ForkedAfterNativeRuntimeStarted",

8
litellm/_db.py
View file

@ -5,14 +5,20 @@ from typing import TYPE_CHECKING
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
from litellm.rust_bridge.database import NativeDatabase
@dataclass(frozen=True, slots=True)
class Database:
prisma: "PrismaClient"
native: "NativeDatabase | None" = None
async def aclose(self) -> None:
await self.prisma.disconnect()
try:
if self.native is not None:
await self.native.close()
finally:
await self.prisma.disconnect()
@asynccontextmanager

1
litellm/constants.py
View file

@ -2350,3 +2350,4 @@ HARNESS_SNAPSHOT_SKIP_DIRS: Final = frozenset(
)
DEFAULT_TOOL_LOOP_MAX_ROUNDS: Final = 20
NATIVE_DATABASE_MAX_CONNECTIONS: Final = 5

17
litellm/proxy/proxy_server.py
View file

@ -86,6 +86,7 @@ from litellm.constants import (
LITELLM_SETTINGS_SAFE_DB_OVERRIDES,
LITELLM_UI_ALLOW_HEADERS,
LITELLM_UI_SESSION_DURATION,
NATIVE_DATABASE_MAX_CONNECTIONS,
OPENAI_LIVE_SESSION_START_TIMEOUT_SECONDS,
RUNTIME_UPDATABLE_ROUTER_SETTINGS,
)
@ -214,6 +215,7 @@ if TYPE_CHECKING:
from litellm.integrations.opentelemetry import OpenTelemetry
from litellm.proxy.health_check_utils.shared_health_check_manager import SharedHealthCheckManager
from litellm.rust_bridge.database import NativeDatabase
Span = _Span | Any
else:
@ -1420,6 +1422,15 @@ async def _connect_to_count_stored_values() -> SupportsRawQueries:
return client.writer_db
async def _connect_native_database() -> "NativeDatabase | None":
database_url: Final = get_secret_str("DATABASE_URL")
if database_url is None or get_secret_bool("LITELLM_NATIVE_DATABASE") is not True:
return None
from litellm.rust_bridge.database import connect_native_database
return await connect_native_database(database_url, NATIVE_DATABASE_MAX_CONNECTIONS)
@asynccontextmanager
async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState, None]:
async with AsyncExitStack() as resources, _proxy_lifespan(app, resources) as state:
@ -1552,7 +1563,11 @@ async def _proxy_lifespan(app: FastAPI, resources: AsyncExitStack) -> AsyncGener
)
database: Final = await resources.enter_async_context(
_db.lifespan(None if prisma_client is None else _db.Database(prisma=prisma_client))
_db.lifespan(
None
if prisma_client is None
else _db.Database(prisma=prisma_client, native=await _connect_native_database())
)
)
await migrate_if_requested(

13
litellm/proxy/ui_crud_endpoints/user_banner_dependencies.py
View file

@ -6,7 +6,12 @@ from fastapi import Depends
from litellm.models.user_banner import UserBanner
from litellm.proxy._types import LitellmTableNames, UserAPIKeyAuth
from litellm.proxy.deps import DatabaseDependency
from litellm.repositories.user_banner_repository import USER_BANNER_ROW_ID, UserBannerRepository, UserBannerStore
from litellm.repositories.user_banner_repository import (
USER_BANNER_ROW_ID,
NativeUserBannerRepository,
UserBannerRepository,
UserBannerStore,
)
class UserBannerAudit(Protocol):
@ -14,7 +19,11 @@ class UserBannerAudit(Protocol):
async def get_user_banner_repository(database: DatabaseDependency) -> UserBannerStore | None:
return None if database is None else UserBannerRepository(database.prisma)
if database is None:
return None
if database.native is not None:
return NativeUserBannerRepository(database.native)
return UserBannerRepository(database.prisma)
async def get_user_banner_audit() -> UserBannerAudit:

14
litellm/repositories/user_banner_repository.py
View file

@ -1,4 +1,5 @@
import json
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final, Protocol
from pydantic import ValidationError
@ -9,6 +10,8 @@ from litellm.repositories.table_repositories import PrismaTableRepository
if TYPE_CHECKING:
from prisma import models as prisma_models # noqa: F401 # resolved only from the quoted base-class subscript below
from litellm.rust_bridge.database import NativeDatabase
USER_BANNER_ROW_ID: Final = "user_banner"
@ -28,6 +31,17 @@ def parse_user_banner(raw_settings: object) -> UserBanner:
return UserBanner()
@dataclass(frozen=True, slots=True)
class NativeUserBannerRepository:
native: "NativeDatabase"
async def get_banner(self) -> UserBanner:
return parse_user_banner(await self.native.find_user_banner())
async def save_banner(self, banner: UserBanner) -> None:
await self.native.save_user_banner(banner.model_dump())
class UserBannerRepository(PrismaTableRepository["prisma_models.LiteLLM_UISettings"]):
table_name = "litellm_uisettings"

9
litellm/rust_bridge/_native.pyi
View file

@ -25,6 +25,14 @@ class NativeClickHouseSpendStorage:
def ensure_schema(self) -> Future[None]: ...
def insert_rows(self, rows: Sequence[Mapping[str, object]]) -> Future[None]: ...
@final
class NativeDatabase:
@staticmethod
def connect(url: str, max_connections: int) -> Future[NativeDatabase]: ...
def close(self) -> Future[None]: ...
def find_user_banner(self) -> Future[JsonValue]: ...
def save_user_banner(self, settings: Mapping[str, JsonValue]) -> Future[None]: ...
@final
class NativeDiagnosticProcessor:
def __new__(cls, minimum_custom_key_length: int) -> NativeDiagnosticProcessor: ...
@ -207,6 +215,7 @@ __all__ = [
"HuggingFaceEncoding",
"NativeClickHouseSpendConfig",
"NativeClickHouseSpendStorage",
"NativeDatabase",
"NativeDiagnosticProcessor",
"ProcessReservedForForking",
"ResponsesWebSocketConnection",

33
litellm/rust_bridge/database.py Normal file
View file

@ -0,0 +1,33 @@
from collections.abc import Awaitable, Mapping
from typing import Final, Protocol, runtime_checkable
from pydantic import ConfigDict, TypeAdapter
from litellm.rust_bridge.loader import get_native_bridge
class NativeDatabase(Protocol):
def find_user_banner(self) -> Awaitable[object]: ...
def save_user_banner(self, settings: Mapping[str, object]) -> Awaitable[None]: ...
def close(self) -> Awaitable[None]: ...
class NativeDatabaseFactory(Protocol):
def connect(self, url: str, max_connections: int) -> Awaitable[NativeDatabase]: ...
@runtime_checkable
class NativeDatabaseBridge(Protocol):
NativeDatabase: NativeDatabaseFactory
_NATIVE: Final[TypeAdapter[NativeDatabaseBridge]] = TypeAdapter(
NativeDatabaseBridge, config=ConfigDict(arbitrary_types_allowed=True)
)
async def connect_native_database(url: str, max_connections: int) -> NativeDatabase:
bridge: Final = get_native_bridge()
if bridge is None:
raise RuntimeError("The native database requires the Rust extension")
return await _NATIVE.validate_python(bridge).NativeDatabase.connect(url, max_connections)

41
tests/unit/proxy/ui_crud_endpoints/test_user_banner_endpoints.py
View file

@ -229,3 +229,44 @@ class TestUpdateUserBanner:
response = client.patch("/update/user_banner", json=payload)
assert response.status_code == 422
mock_prisma.db.litellm_uisettings.upsert.assert_not_awaited()
class _NativeDatabase:
def __init__(self, stored=None):
self.stored = stored
async def find_user_banner(self):
return self.stored
async def save_user_banner(self, settings):
self.stored = settings
class TestNativeUserBannerBackend:
@pytest.fixture
def backends(self, monkeypatch):
prisma = MagicMock()
banners = _NativeDatabase()
monkeypatch.setitem(app.dependency_overrides, get_database, lambda: _db.Database(prisma=prisma, native=banners))
return prisma, banners
def test_published_banner_round_trips_through_the_native_backend(self, admin_auth, mock_audit_log, backends):
prisma, banners = backends
published = client.patch("/update/user_banner", json=PUBLISH_BODY).json()["banner"]
assert banners.stored == published
assert client.get("/get/user_banner").json() == published
assert prisma.mock_calls == []
@pytest.mark.parametrize(
"stored", [json.dumps(PUBLISHED_BANNER), PUBLISHED_BANNER], ids=["prisma-string", "object"]
)
def test_native_backend_reads_rows_written_by_either_backend(self, admin_auth, backends, stored):
_, banners = backends
banners.stored = stored
assert client.get("/get/user_banner").json() == PUBLISHED_BANNER
def test_native_backend_treats_invalid_stored_banner_as_disabled(self, admin_auth, backends):
_, banners = backends
banners.stored = {"enabled": True, "message": ""}
assert client.get("/get/user_banner").json() == DISABLED_BANNER

14
tests/unit/test__db.py
View file

@ -34,3 +34,17 @@ async def test_lifespan_without_database_preserves_owner_error():
with pytest.raises(RuntimeError, match="owner failed"):
await fail()
@pytest.mark.asyncio
async def test_closing_disconnects_prisma_even_when_the_native_pool_fails_to_close():
prisma: Final = MagicMock()
prisma.disconnect = AsyncMock()
native: Final = MagicMock()
native.close = AsyncMock(side_effect=RuntimeError("native close failed"))
with pytest.raises(RuntimeError, match="native close failed"):
async with _db.lifespan(_db.Database(prisma=prisma, native=native)):
pass
native.close.assert_awaited_once()
prisma.disconnect.assert_awaited_once()