diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index f91deaaa401..14e5064fa04 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4344,6 +4344,7 @@ dependencies = [ "litellm-cache-response", "litellm-callbacks-legacy-python", "litellm-core-utils", + "litellm-db", "litellm-error", "litellm-host", "litellm-host-python", diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index e3ef45f7876..8cea0850216 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -19,6 +19,8 @@ huggingface = ["litellm-token-counter/huggingface"] tiktoken = ["litellm-token-counter/tiktoken"] [dependencies] +litellm-db.workspace = true +sqlx.workspace = true litellm-secret-value.workspace = true fancy-regex.workspace = true litellm-tracing.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/db.rs b/litellm-rust/crates/python-bridge/src/db.rs index 03fb30b986f..c50ee72cd98 100644 --- a/litellm-rust/crates/python-bridge/src/db.rs +++ b/litellm-rust/crates/python-bridge/src/db.rs @@ -1,19 +1,56 @@ -use litellm_host_python::ready_future; -use pyo3::{prelude::*, types::PyNone}; +use std::str::FromStr; + +use litellm_db::{EntitySpend, add_key_spend}; +use litellm_host_python::run_async_value; +use pyo3::{exceptions::PyRuntimeError, prelude::*}; +use sqlx::postgres::{PgConnectOptions, PgPool, PgPoolOptions}; #[pyclass(frozen)] -pub(crate) struct NativeDatabase {} +pub(crate) struct NativeDatabase { + pool: PgPool, +} + +fn connect_options() -> PyResult { + let url = std::env::var("DATABASE_URL") + .map_err(|_| PyRuntimeError::new_err("DATABASE_URL is not set"))?; + let base = url.split('?').next().unwrap_or(&url); + PgConnectOptions::from_str(base) + .map(|options| options.application_name("litellm-native")) + .map_err(|error| PyRuntimeError::new_err(error.to_string())) +} -// TODO: take typed writer/reader settings and a token provider, then open a sqlx PgPool, once the first -// repository moves to native. The Prisma DATABASE_URL carries Prisma-only params and rotating IAM tokens #[pymethods] impl NativeDatabase { #[staticmethod] fn connect(py: Python<'_>) -> PyResult> { - ready_future(py, Bound::new(py, NativeDatabase {})?.as_any()) + let options = connect_options()?; + run_async_value(py, async move { + let pool = PgPoolOptions::new() + .max_connections(5) + .connect_lazy_with(options); + Ok(NativeDatabase { pool }) + }) + } + + fn add_key_spend<'py>( + &self, + py: Python<'py>, + ids: Vec, + costs: Vec, + ) -> PyResult> { + let pool = self.pool.clone(); + run_async_value(py, async move { + add_key_spend(&pool, &EntitySpend { ids, costs }) + .await + .map_err(|error| PyRuntimeError::new_err(error.to_string())) + }) } fn close<'py>(&self, py: Python<'py>) -> PyResult> { - ready_future(py, PyNone::get(py).as_any()) + let pool = self.pool.clone(); + run_async_value(py, async move { + pool.close().await; + Ok(()) + }) } } diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index d9fb236632b..9affad7ed98 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -61,6 +61,7 @@ from litellm.proxy.db.daily_spend_bulk_upsert import ( merge_by_conflict_key, ) from litellm.proxy.db.db_span import db_span +from litellm.repositories import database as database_runtime from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( DailySpendUpdateQueue, ) @@ -2196,19 +2197,11 @@ class DBSpendUpdateWriter: for i in range(n_retry_times + 1): start_time = time.time() try: - async with _spend_update_tx(prisma_client, "LiteLLM_VerificationToken") as transaction: - async with transaction.batch_() as batcher: - # Sort by token for consistent lock ordering across pods to prevent deadlocks. - for token, response_cost in sorted(key_list_transactions.items()): - spend_increment: _SpendIncrement = {"increment": response_cost} - batcher.litellm_verificationtoken.update_many( # 'update_many' prevents error from being raised if no row exists - where={"token": token}, - data={ - "spend": spend_increment, - "total_spend": spend_increment, - "last_active": datetime.now(timezone.utc), - }, - ) + native: Final = database_runtime.current_native + if native is None: + raise RuntimeError("native database is not open") + ordered: Final = sorted(key_list_transactions.items()) + await native.add_key_spend([t for t, _ in ordered], [c for _, c in ordered]) break except Exception as e: await self.handle_spend_update_failure( diff --git a/litellm/repositories/database.py b/litellm/repositories/database.py index d60a2099477..870295d44d7 100644 --- a/litellm/repositories/database.py +++ b/litellm/repositories/database.py @@ -10,6 +10,9 @@ if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient +current_native: NativeDatabase | None = None + + @dataclass(frozen=True, slots=True) class Database: prisma: "PrismaClient" @@ -22,6 +25,8 @@ class Database: except BaseException: await prisma.aclose() raise + global current_native # noqa: PLW0603 # fuzz port experiment: spend writer has no handle yet + current_native = native # rebind-ok: fuzz port experiment return cls(prisma=prisma, native=native) async def aclose(self) -> None: diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 8177e4a5cc9..ca7b67a106c 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -30,6 +30,7 @@ class NativeClickHouseSpendStorage: class NativeDatabase: @staticmethod def connect() -> Future[NativeDatabase]: ... + def add_key_spend(self, ids: Sequence[str], costs: Sequence[float]) -> Future[None]: ... def close(self) -> Future[None]: ... @final