diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index d58ba9ecb5f..d0757ebbd2c 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3699,6 +3699,7 @@ dependencies = [ "macro_rules_attribute", "rstest", "serde", + "serde_json", "sqlx", "thiserror 2.0.19", "tokio", diff --git a/litellm-rust/crates/db/.sqlx/query-b0eb48b795bdaf728c79e986d2027df60ff71d8fed05b53be4458b335af88147.json b/litellm-rust/crates/db/.sqlx/query-b0eb48b795bdaf728c79e986d2027df60ff71d8fed05b53be4458b335af88147.json new file mode 100644 index 00000000000..b333ec814a1 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-b0eb48b795bdaf728c79e986d2027df60ff71d8fed05b53be4458b335af88147.json @@ -0,0 +1,14 @@ +{ + "db_name": "PostgreSQL", + "query": "INSERT INTO \"LiteLLM_TeamMembership\" (user_id, team_id, spend, total_spend)\nSELECT member.user_id, member.team_id, member.cost, member.cost\nFROM jsonb_to_recordset($1::jsonb) AS member(user_id text, team_id text, cost float8)\nWHERE EXISTS (\n SELECT 1 FROM \"LiteLLM_TeamTable\" t\n WHERE t.team_id = member.team_id\n AND t.members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', member.user_id))\n)\n OR EXISTS (\n SELECT 1 FROM \"LiteLLM_TeamMembership\" m\n WHERE m.user_id = member.user_id AND m.team_id = member.team_id\n)\nON CONFLICT (user_id, team_id) DO UPDATE\nSET spend = \"LiteLLM_TeamMembership\".spend + EXCLUDED.spend,\n total_spend = \"LiteLLM_TeamMembership\".total_spend + EXCLUDED.total_spend\n", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "Jsonb" + ] + }, + "nullable": [] + }, + "hash": "b0eb48b795bdaf728c79e986d2027df60ff71d8fed05b53be4458b335af88147" +} diff --git a/litellm-rust/crates/db/Cargo.toml b/litellm-rust/crates/db/Cargo.toml index bb6d5655b60..8c1a35c7535 100644 --- a/litellm-rust/crates/db/Cargo.toml +++ b/litellm-rust/crates/db/Cargo.toml @@ -9,6 +9,7 @@ repository.workspace = true chrono.workspace = true macro_rules_attribute.workspace = true serde.workspace = true +serde_json.workspace = true sqlx.workspace = true thiserror.workspace = true diff --git a/litellm-rust/crates/db/queries/spend/upsert_team_member_spend.sql b/litellm-rust/crates/db/queries/spend/upsert_team_member_spend.sql new file mode 100644 index 00000000000..460c631cbbd --- /dev/null +++ b/litellm-rust/crates/db/queries/spend/upsert_team_member_spend.sql @@ -0,0 +1,15 @@ +INSERT INTO "LiteLLM_TeamMembership" (user_id, team_id, spend, total_spend) +SELECT member.user_id, member.team_id, member.cost, member.cost +FROM jsonb_to_recordset($1::jsonb) AS member(user_id text, team_id text, cost float8) +WHERE EXISTS ( + SELECT 1 FROM "LiteLLM_TeamTable" t + WHERE t.team_id = member.team_id + AND t.members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', member.user_id)) +) + OR EXISTS ( + SELECT 1 FROM "LiteLLM_TeamMembership" m + WHERE m.user_id = member.user_id AND m.team_id = member.team_id +) +ON CONFLICT (user_id, team_id) DO UPDATE +SET spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend, + total_spend = "LiteLLM_TeamMembership".total_spend + EXCLUDED.total_spend diff --git a/litellm-rust/crates/db/src/lib.rs b/litellm-rust/crates/db/src/lib.rs index 4f009b411fd..afa962cff0e 100644 --- a/litellm-rust/crates/db/src/lib.rs +++ b/litellm-rust/crates/db/src/lib.rs @@ -19,6 +19,6 @@ pub use end_user::{EndUser, RateLimits, find_end_user}; pub use error::Error; pub use failure::DbFailure; pub use spend::{ - EntitySpend, add_end_user_spend, add_key_spend, add_team_spend, add_user_spend, begin, - lock_team, + EntitySpend, TeamMemberSpend, add_end_user_spend, add_key_spend, add_team_member_spend, + add_team_spend, add_user_spend, begin, lock_team, }; diff --git a/litellm-rust/crates/db/src/spend/mod.rs b/litellm-rust/crates/db/src/spend/mod.rs index 0459e63aaf8..1acf629af4e 100644 --- a/litellm-rust/crates/db/src/spend/mod.rs +++ b/litellm-rust/crates/db/src/spend/mod.rs @@ -1,11 +1,13 @@ mod end_user_spend; mod key_spend; mod team_lock; +mod team_member_spend; mod team_spend; mod user_spend; pub use end_user_spend::add_end_user_spend; pub use key_spend::{EntitySpend, add_key_spend}; pub use team_lock::{begin, lock_team}; +pub use team_member_spend::{TeamMemberSpend, add_team_member_spend}; pub use team_spend::add_team_spend; pub use user_spend::add_user_spend; diff --git a/litellm-rust/crates/db/src/spend/team_member_spend.rs b/litellm-rust/crates/db/src/spend/team_member_spend.rs new file mode 100644 index 00000000000..965a39ada4f --- /dev/null +++ b/litellm-rust/crates/db/src/spend/team_member_spend.rs @@ -0,0 +1,162 @@ +use sqlx::{Postgres, Transaction}; + +use super::lock_team; +use crate::Error; + +#[macro_rules_attribute::apply(crate::db_input)] +pub struct TeamMemberSpend { + pub team_id: String, + pub user_id: String, + pub cost: f64, +} + +pub async fn add_team_member_spend( + transaction: &mut Transaction<'_, Postgres>, + members: &[TeamMemberSpend], +) -> Result<(), Error> { + if members.is_empty() { + return Ok(()); + } + let mut members = members.to_vec(); + members.sort_by(|left, right| { + (&left.team_id, &left.user_id).cmp(&(&right.team_id, &right.user_id)) + }); + let mut teams: Vec<&str> = members + .iter() + .map(|member| member.team_id.as_str()) + .collect(); + teams.dedup(); + for team in teams { + lock_team(transaction, team).await?; + } + let rows = serde_json::Value::Array( + members + .iter() + .map(|member| { + serde_json::json!({ + "user_id": member.user_id, + "team_id": member.team_id, + "cost": member.cost, + }) + }) + .collect(), + ); + sqlx::query_file!("queries/spend/upsert_team_member_spend.sql", rows) + .execute(&mut **transaction) + .await?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::error::Error as StdError; + + use super::super::begin; + use super::{TeamMemberSpend, add_team_member_spend}; + use crate::test_support::database_pool; + use rstest::rstest; + use sqlx::PgPool; + + fn members(cost: f64) -> [TeamMemberSpend; 3] { + [ + TeamMemberSpend { + team_id: "member-spend-z".into(), + user_id: "off-roster-existing".into(), + cost: 2.0 * cost, + }, + TeamMemberSpend { + team_id: "member-spend-a".into(), + user_id: "on-roster-missing".into(), + cost, + }, + TeamMemberSpend { + team_id: "member-spend-a".into(), + user_id: "off-roster-missing".into(), + cost: 50.0 * cost, + }, + ] + } + + #[rstest] + #[case::zero_then_quarter(0.25, 0.5, 3.0, 8.0)] + #[case::zero_then_eighth(0.125, 0.25, 2.5, 7.5)] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn membership_spend_preserves_roster_guard_and_fractional_costs( + #[case] cost: f64, + #[case] expected_created: f64, + #[case] expected_spend: f64, + #[case] expected_total: f64, + #[future(awt)] database_pool: Result>, + ) -> Result<(), Box> { + let pool = database_pool?; + let mut tx = begin(&pool).await?; + #[expect( + clippy::disallowed_methods, + reason = "Fixture writes are outside the compile-checked query set" + )] + let _teams = sqlx::query( + r#"INSERT INTO "LiteLLM_TeamTable" (team_id, members_with_roles) + VALUES ('member-spend-a', '[{"user_id":"on-roster-missing","role":"user"}]'), + ('member-spend-z', '[]')"#, + ) + .execute(&mut *tx) + .await?; + #[expect( + clippy::disallowed_methods, + reason = "Fixture writes are outside the compile-checked query set" + )] + let _members = sqlx::query( + r#"INSERT INTO "LiteLLM_TeamMembership" (team_id, user_id, spend, total_spend) + VALUES ('member-spend-z', 'off-roster-existing', 2, 7), + ('member-spend-a', 'untouched', 100, 200)"#, + ) + .execute(&mut *tx) + .await?; + add_team_member_spend(&mut tx, &[]).await?; + add_team_member_spend(&mut tx, &members(0.0)).await?; + add_team_member_spend(&mut tx, &members(cost)).await?; + add_team_member_spend(&mut tx, &members(cost)).await?; + #[expect( + clippy::disallowed_methods, + reason = "Test assertions are outside the compile-checked query set" + )] + let rows: Vec<(String, String, f64, f64)> = sqlx::query_as( + r#"SELECT team_id, user_id, spend, total_spend FROM "LiteLLM_TeamMembership" + WHERE team_id IN ('member-spend-a', 'member-spend-z') ORDER BY team_id, user_id"#, + ) + .fetch_all(&mut *tx) + .await?; + assert_eq!( + rows, + vec![ + ( + "member-spend-a".into(), + "on-roster-missing".into(), + expected_created, + expected_created + ), + ("member-spend-a".into(), "untouched".into(), 100.0, 200.0), + ( + "member-spend-z".into(), + "off-roster-existing".into(), + expected_spend, + expected_total + ), + ] + ); + tx.rollback().await?; + #[expect( + clippy::disallowed_methods, + reason = "Test assertions are outside the compile-checked query set" + )] + let remaining: i64 = sqlx::query_scalar( + r#"SELECT COUNT(*) FROM "LiteLLM_TeamMembership" + WHERE team_id IN ('member-spend-a', 'member-spend-z')"#, + ) + .fetch_one(&pool) + .await?; + assert_eq!(remaining, 0); + Ok(()) + } +}