From a2545134f370696ae1dd2ba7102d0ac4da552968 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Mon, 26 Oct 2026 17:05:07 -0700 Subject: [PATCH] feat(db): add typed daily team spend upsert (#48236) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...74dca981ff954452b35ffbace3e022035b416.json | 31 ++ .../queries/spend/upsert_daily_team_spend.sql | 40 ++ litellm-rust/crates/db/src/lib.rs | 2 +- .../crates/db/src/spend/daily_team_spend.rs | 485 ++++++++++++++++++ litellm-rust/crates/db/src/spend/mod.rs | 2 + 5 files changed, 559 insertions(+), 1 deletion(-) create mode 100644 litellm-rust/crates/db/.sqlx/query-9c51401a8907f818e96da28050b74dca981ff954452b35ffbace3e022035b416.json create mode 100644 litellm-rust/crates/db/queries/spend/upsert_daily_team_spend.sql create mode 100644 litellm-rust/crates/db/src/spend/daily_team_spend.rs diff --git a/litellm-rust/crates/db/.sqlx/query-9c51401a8907f818e96da28050b74dca981ff954452b35ffbace3e022035b416.json b/litellm-rust/crates/db/.sqlx/query-9c51401a8907f818e96da28050b74dca981ff954452b35ffbace3e022035b416.json new file mode 100644 index 00000000000..9886a9505d8 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-9c51401a8907f818e96da28050b74dca981ff954452b35ffbace3e022035b416.json @@ -0,0 +1,31 @@ +{ + "db_name": "PostgreSQL", + "query": "INSERT INTO \"LiteLLM_DailyTeamSpend\" (\n id, team_id, date, api_key, model, custom_llm_provider,\n mcp_namespaced_tool_name, endpoint, model_group, prompt_tokens,\n completion_tokens, api_requests, successful_requests, failed_requests,\n cache_read_input_tokens, cache_creation_input_tokens, total_response_time_ms,\n timed_requests, spend, updated_at\n)\nSELECT gen_random_uuid()::text, r.entity_id, r.date, r.api_key, r.model,\n r.custom_llm_provider, r.mcp_namespaced_tool_name, r.endpoint, r.model_group, r.prompt_tokens,\n r.completion_tokens, r.api_requests, r.successful_requests,\n r.failed_requests, r.cache_read, r.cache_creation, r.response_time,\n r.timed_requests, r.spend, (NOW() AT TIME ZONE 'UTC')\nFROM UNNEST(\n $1::text[], $2::text[], $3::text[], $4::text[], $5::text[],\n $6::text[], $7::text[], $8::int8[], $9::int8[], $10::int8[],\n $11::int8[], $12::int8[], $13::int8[], $14::int8[], $15::int8[],\n $16::int8[], $17::float8[], $18::text[]\n) AS r(\n entity_id, date, api_key, model, custom_llm_provider, endpoint,\n model_group, prompt_tokens, completion_tokens, api_requests,\n successful_requests, failed_requests, cache_read, cache_creation,\n response_time, timed_requests, spend, mcp_namespaced_tool_name\n)\nORDER BY r.date, r.entity_id, r.api_key, r.model, r.custom_llm_provider,\n r.mcp_namespaced_tool_name, r.endpoint\nON CONFLICT (\n team_id, date, api_key, model, custom_llm_provider,\n mcp_namespaced_tool_name, endpoint\n) DO UPDATE SET\n prompt_tokens = \"LiteLLM_DailyTeamSpend\".prompt_tokens + EXCLUDED.prompt_tokens,\n completion_tokens = \"LiteLLM_DailyTeamSpend\".completion_tokens + EXCLUDED.completion_tokens,\n api_requests = \"LiteLLM_DailyTeamSpend\".api_requests + EXCLUDED.api_requests,\n successful_requests = \"LiteLLM_DailyTeamSpend\".successful_requests + EXCLUDED.successful_requests,\n failed_requests = \"LiteLLM_DailyTeamSpend\".failed_requests + EXCLUDED.failed_requests,\n cache_read_input_tokens = \"LiteLLM_DailyTeamSpend\".cache_read_input_tokens + EXCLUDED.cache_read_input_tokens,\n cache_creation_input_tokens = \"LiteLLM_DailyTeamSpend\".cache_creation_input_tokens + EXCLUDED.cache_creation_input_tokens,\n total_response_time_ms = \"LiteLLM_DailyTeamSpend\".total_response_time_ms + EXCLUDED.total_response_time_ms,\n timed_requests = \"LiteLLM_DailyTeamSpend\".timed_requests + EXCLUDED.timed_requests,\n spend = \"LiteLLM_DailyTeamSpend\".spend + EXCLUDED.spend,\n updated_at = (NOW() AT TIME ZONE 'UTC')\n", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Float8Array", + "TextArray" + ] + }, + "nullable": [] + }, + "hash": "9c51401a8907f818e96da28050b74dca981ff954452b35ffbace3e022035b416" +} diff --git a/litellm-rust/crates/db/queries/spend/upsert_daily_team_spend.sql b/litellm-rust/crates/db/queries/spend/upsert_daily_team_spend.sql new file mode 100644 index 00000000000..a9bdef34043 --- /dev/null +++ b/litellm-rust/crates/db/queries/spend/upsert_daily_team_spend.sql @@ -0,0 +1,40 @@ +INSERT INTO "LiteLLM_DailyTeamSpend" ( + id, team_id, date, api_key, model, custom_llm_provider, + mcp_namespaced_tool_name, endpoint, model_group, prompt_tokens, + completion_tokens, api_requests, successful_requests, failed_requests, + cache_read_input_tokens, cache_creation_input_tokens, total_response_time_ms, + timed_requests, spend, updated_at +) +SELECT gen_random_uuid()::text, r.entity_id, r.date, r.api_key, r.model, + r.custom_llm_provider, r.mcp_namespaced_tool_name, r.endpoint, r.model_group, r.prompt_tokens, + r.completion_tokens, r.api_requests, r.successful_requests, + r.failed_requests, r.cache_read, r.cache_creation, r.response_time, + r.timed_requests, r.spend, (NOW() AT TIME ZONE 'UTC') +FROM UNNEST( + $1::text[], $2::text[], $3::text[], $4::text[], $5::text[], + $6::text[], $7::text[], $8::int8[], $9::int8[], $10::int8[], + $11::int8[], $12::int8[], $13::int8[], $14::int8[], $15::int8[], + $16::int8[], $17::float8[], $18::text[] +) AS r( + entity_id, date, api_key, model, custom_llm_provider, endpoint, + model_group, prompt_tokens, completion_tokens, api_requests, + successful_requests, failed_requests, cache_read, cache_creation, + response_time, timed_requests, spend, mcp_namespaced_tool_name +) +ORDER BY r.date, r.entity_id, r.api_key, r.model, r.custom_llm_provider, + r.mcp_namespaced_tool_name, r.endpoint +ON CONFLICT ( + team_id, date, api_key, model, custom_llm_provider, + mcp_namespaced_tool_name, endpoint +) DO UPDATE SET + prompt_tokens = "LiteLLM_DailyTeamSpend".prompt_tokens + EXCLUDED.prompt_tokens, + completion_tokens = "LiteLLM_DailyTeamSpend".completion_tokens + EXCLUDED.completion_tokens, + api_requests = "LiteLLM_DailyTeamSpend".api_requests + EXCLUDED.api_requests, + successful_requests = "LiteLLM_DailyTeamSpend".successful_requests + EXCLUDED.successful_requests, + failed_requests = "LiteLLM_DailyTeamSpend".failed_requests + EXCLUDED.failed_requests, + cache_read_input_tokens = "LiteLLM_DailyTeamSpend".cache_read_input_tokens + EXCLUDED.cache_read_input_tokens, + cache_creation_input_tokens = "LiteLLM_DailyTeamSpend".cache_creation_input_tokens + EXCLUDED.cache_creation_input_tokens, + total_response_time_ms = "LiteLLM_DailyTeamSpend".total_response_time_ms + EXCLUDED.total_response_time_ms, + timed_requests = "LiteLLM_DailyTeamSpend".timed_requests + EXCLUDED.timed_requests, + spend = "LiteLLM_DailyTeamSpend".spend + EXCLUDED.spend, + updated_at = (NOW() AT TIME ZONE 'UTC') diff --git a/litellm-rust/crates/db/src/lib.rs b/litellm-rust/crates/db/src/lib.rs index 3bcd050440a..5b88bcabf32 100644 --- a/litellm-rust/crates/db/src/lib.rs +++ b/litellm-rust/crates/db/src/lib.rs @@ -21,5 +21,5 @@ pub use failure::DbFailure; pub use spend::{ DailySpend, EntitySpend, SpendLog, SpendLogRow, TeamMemberSpend, add_end_user_spend, add_key_spend, add_team_member_spend, add_team_spend, add_user_spend, begin, insert_spend_logs, - lock_team, upsert_daily_user_spend, + lock_team, upsert_daily_team_spend, upsert_daily_user_spend, }; diff --git a/litellm-rust/crates/db/src/spend/daily_team_spend.rs b/litellm-rust/crates/db/src/spend/daily_team_spend.rs new file mode 100644 index 00000000000..f2acd23f468 --- /dev/null +++ b/litellm-rust/crates/db/src/spend/daily_team_spend.rs @@ -0,0 +1,485 @@ +use sqlx::PgExecutor; + +use super::DailySpend; +use crate::Error; + +pub async fn upsert_daily_team_spend( + executor: impl PgExecutor<'_>, + rows: &[DailySpend], +) -> Result<(), Error> { + if rows.is_empty() { + return Ok(()); + } + macro_rules! column { + ($field:ident) => { + rows.iter() + .map(|row| row.$field.clone()) + .collect::>() + }; + } + sqlx::query_file!( + "queries/spend/upsert_daily_team_spend.sql", + &column!(entity_id), + &column!(date), + &column!(api_key), + &column!(model), + &column!(custom_llm_provider), + &column!(endpoint), + &column!(model_group) as &[Option], + &column!(prompt_tokens), + &column!(completion_tokens), + &column!(api_requests), + &column!(successful_requests), + &column!(failed_requests), + &column!(cache_read_input_tokens), + &column!(cache_creation_input_tokens), + &column!(total_response_time_ms), + &column!(timed_requests), + &column!(spend), + &column!(mcp_namespaced_tool_name), + ) + .execute(executor) + .await?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::error::Error as StdError; + + use super::{DailySpend, upsert_daily_team_spend}; + use crate::{DailyTeamRow, test_support::database_pool}; + use rstest::rstest; + use sqlx::{PgPool, types::chrono::NaiveDateTime}; + + #[rstest] + #[case::distinct_tools("server.fetch", "server.search")] + #[case::default_and_named_tool("", "server.search")] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn daily_team_spend_separates_tools_and_accumulates( + #[case] first_tool: &str, + #[case] second_tool: &str, + #[future(awt)] database_pool: Result>, + ) -> Result<(), Box> { + let pool = database_pool?; + let mut tx = pool.begin().await?; + let first = DailySpend { + entity_id: "daily-team-tools".to_owned(), + date: "2026-01-15".to_owned(), + api_key: "daily-key-tools".to_owned(), + model: "daily-model".to_owned(), + custom_llm_provider: "daily-provider".to_owned(), + endpoint: "/mcp".to_owned(), + mcp_namespaced_tool_name: first_tool.to_owned(), + api_requests: 1, + spend: 0.125, + ..Default::default() + }; + let second = DailySpend { + mcp_namespaced_tool_name: second_tool.to_owned(), + api_requests: 2, + spend: 0.25, + ..first.clone() + }; + upsert_daily_team_spend(&mut *tx, &[second.clone(), first.clone()]).await?; + upsert_daily_team_spend( + &mut *tx, + &[ + DailySpend { + api_requests: 3, + spend: 0.5, + ..first + }, + DailySpend { + api_requests: 5, + spend: 1.0, + ..second + }, + ], + ) + .await?; + + #[expect( + clippy::disallowed_methods, + reason = "The integration test reads stored tool attribution and counters" + )] + let rows = sqlx::query_as::<_, DailyTeamRow>( + r#"SELECT * FROM "LiteLLM_DailyTeamSpend" + WHERE team_id = $1 AND date = $2 + ORDER BY mcp_namespaced_tool_name"#, + ) + .bind("daily-team-tools") + .bind("2026-01-15") + .fetch_all(&mut *tx) + .await?; + + assert_eq!(rows.len(), 2); + assert_eq!( + rows[0].mcp_namespaced_tool_name.as_deref(), + Some(first_tool) + ); + assert_eq!( + rows[1].mcp_namespaced_tool_name.as_deref(), + Some(second_tool) + ); + assert_eq!((rows[0].api_requests, rows[0].spend), (4, 0.625)); + assert_eq!((rows[1].api_requests, rows[1].spend), (7, 1.25)); + Ok(()) + } + + #[rstest] + #[case::full_keys( + "daily-team-base", + "daily-key-base", + "daily-model", + "daily-provider", + "/chat/completions", + Some("first-group") + )] + #[case::initial_model_group_none( + "daily-team-base", + "daily-key-base", + "daily-model", + "daily-provider", + "/chat/completions", + None + )] + #[case::empty_optional_keys("", "", "", "", "", Some(""))] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn daily_team_spend_upsert_preserves_keys_and_accumulates( + #[case] team_id: &'static str, + #[case] api_key: &'static str, + #[case] model: &'static str, + #[case] custom_llm_provider: &'static str, + #[case] endpoint: &'static str, + #[case] first_model_group: Option<&'static str>, + #[future(awt)] database_pool: Result>, + ) -> Result<(), Box> { + let pool = database_pool?; + let mut tx = pool.begin().await?; + + #[expect( + clippy::disallowed_methods, + reason = "The integration test exercises UTC timestamp behavior under a non-UTC session" + )] + sqlx::query("SET LOCAL TIME ZONE 'America/Los_Angeles'") + .execute(&mut *tx) + .await?; + + let first_batch = [ + DailySpend { + entity_id: team_id.to_owned(), + date: "2026-01-15".to_owned(), + api_key: api_key.to_owned(), + model: model.to_owned(), + custom_llm_provider: custom_llm_provider.to_owned(), + mcp_namespaced_tool_name: String::new(), + endpoint: endpoint.to_owned(), + model_group: first_model_group.map(str::to_owned), + prompt_tokens: 1, + completion_tokens: 2, + api_requests: 3, + successful_requests: 4, + failed_requests: 5, + cache_read_input_tokens: 6, + cache_creation_input_tokens: 7, + total_response_time_ms: 8, + timed_requests: 9, + spend: 0.125, + }, + DailySpend { + entity_id: team_id.to_owned(), + date: "2026-01-15".to_owned(), + api_key: api_key.to_owned(), + model: model.to_owned(), + custom_llm_provider: custom_llm_provider.to_owned(), + mcp_namespaced_tool_name: String::new(), + endpoint: "endpoint-sentinel".to_owned(), + model_group: Some("sentinel-endpoint-group".to_owned()), + prompt_tokens: 101, + completion_tokens: 102, + api_requests: 103, + successful_requests: 104, + failed_requests: 105, + cache_read_input_tokens: 106, + cache_creation_input_tokens: 107, + total_response_time_ms: 108, + timed_requests: 109, + spend: 1.25, + }, + DailySpend { + entity_id: team_id.to_owned(), + date: "2026-01-15".to_owned(), + api_key: "api-key-sentinel".to_owned(), + model: model.to_owned(), + custom_llm_provider: custom_llm_provider.to_owned(), + mcp_namespaced_tool_name: String::new(), + endpoint: endpoint.to_owned(), + model_group: Some("sentinel-key-group".to_owned()), + prompt_tokens: 201, + completion_tokens: 202, + api_requests: 203, + successful_requests: 204, + failed_requests: 205, + cache_read_input_tokens: 206, + cache_creation_input_tokens: 207, + total_response_time_ms: 208, + timed_requests: 209, + spend: 2.5, + }, + DailySpend { + entity_id: "entity-sentinel".to_owned(), + date: "2026-01-15".to_owned(), + api_key: api_key.to_owned(), + model: model.to_owned(), + custom_llm_provider: custom_llm_provider.to_owned(), + mcp_namespaced_tool_name: String::new(), + endpoint: endpoint.to_owned(), + model_group: Some("sentinel-entity-group".to_owned()), + prompt_tokens: 301, + completion_tokens: 302, + api_requests: 303, + successful_requests: 304, + failed_requests: 305, + cache_read_input_tokens: 306, + cache_creation_input_tokens: 307, + total_response_time_ms: 308, + timed_requests: 309, + spend: 3.75, + }, + ]; + upsert_daily_team_spend(&mut *tx, &first_batch).await?; + + #[expect( + clippy::disallowed_methods, + reason = "The integration test captures the inserted row identity before conflict handling" + )] + let (initial_id, initial_created_at): (String, NaiveDateTime) = sqlx::query_as( + r#"SELECT id, created_at + FROM "LiteLLM_DailyTeamSpend" + WHERE team_id = $1 AND date = $2 AND api_key = $3 AND model = $4 + AND custom_llm_provider = $5 AND endpoint = $6"#, + ) + .bind(team_id) + .bind("2026-01-15") + .bind(api_key) + .bind(model) + .bind(custom_llm_provider) + .bind(endpoint) + .fetch_one(&mut *tx) + .await?; + assert!(!initial_id.is_empty()); + + #[expect( + clippy::disallowed_methods, + reason = "The integration test seeds updated_at before exercising the conflict update" + )] + let seeded_updated_at: NaiveDateTime = sqlx::query_scalar( + r#"UPDATE "LiteLLM_DailyTeamSpend" + SET updated_at = TIMESTAMP '2000-01-01' + WHERE id = $1 + RETURNING updated_at"#, + ) + .bind(&initial_id) + .fetch_one(&mut *tx) + .await?; + + let increment = [DailySpend { + entity_id: team_id.to_owned(), + date: "2026-01-15".to_owned(), + api_key: api_key.to_owned(), + model: model.to_owned(), + custom_llm_provider: custom_llm_provider.to_owned(), + mcp_namespaced_tool_name: String::new(), + endpoint: endpoint.to_owned(), + model_group: Some("second-group".to_owned()), + prompt_tokens: 10, + completion_tokens: 20, + api_requests: 30, + successful_requests: 40, + failed_requests: 50, + cache_read_input_tokens: 60, + cache_creation_input_tokens: 70, + total_response_time_ms: 80, + timed_requests: 90, + spend: 0.625, + }]; + upsert_daily_team_spend(&mut *tx, &increment).await?; + upsert_daily_team_spend(&mut *tx, &[]).await?; + + #[expect( + clippy::disallowed_methods, + reason = "The integration test captures the transaction timestamp for assertions" + )] + let expected_updated_at: NaiveDateTime = + sqlx::query_scalar("SELECT (NOW() AT TIME ZONE 'UTC')::timestamp(3)") + .fetch_one(&mut *tx) + .await?; + + #[expect( + clippy::disallowed_methods, + reason = "The integration test reads all rows to verify attribution and accumulation" + )] + let rows: Vec = sqlx::query_as( + r#"SELECT id, team_id, api_key, model, custom_llm_provider, endpoint, + model_group, mcp_namespaced_tool_name, prompt_tokens, + completion_tokens, api_requests, successful_requests, + failed_requests, cache_read_input_tokens, + cache_creation_input_tokens, total_response_time_ms, + timed_requests, spend, date, created_at, updated_at + FROM "LiteLLM_DailyTeamSpend" + WHERE date = $1 AND model = $2 AND custom_llm_provider = $3 + AND mcp_namespaced_tool_name = '' AND team_id = ANY($4) + ORDER BY CASE + WHEN team_id = $5 AND api_key = $6 AND endpoint = $7 THEN 0 + WHEN team_id = $5 AND api_key = $6 AND endpoint = $9 THEN 1 + WHEN team_id = $5 AND api_key = $8 AND endpoint = $7 THEN 2 + WHEN team_id = $10 AND api_key = $6 AND endpoint = $7 THEN 3 + ELSE 4 + END"#, + ) + .bind("2026-01-15") + .bind(model) + .bind(custom_llm_provider) + .bind(vec![team_id, "entity-sentinel"]) + .bind(team_id) + .bind(api_key) + .bind(endpoint) + .bind("api-key-sentinel") + .bind("endpoint-sentinel") + .bind("entity-sentinel") + .fetch_all(&mut *tx) + .await?; + + assert_eq!(rows.len(), 4); + assert_eq!(rows[0].id, initial_id); + assert_eq!(rows[0].team_id.as_deref(), Some(team_id)); + assert_eq!(rows[0].api_key, api_key); + assert_eq!(rows[0].model.as_deref(), Some(model)); + assert_eq!( + rows[0].custom_llm_provider.as_deref(), + Some(custom_llm_provider) + ); + assert_eq!(rows[0].endpoint.as_deref(), Some(endpoint)); + assert_eq!(rows[0].model_group.as_deref(), first_model_group); + assert_eq!(rows[0].mcp_namespaced_tool_name.as_deref(), Some("")); + assert_eq!( + ( + rows[0].prompt_tokens, + rows[0].completion_tokens, + rows[0].api_requests, + rows[0].successful_requests, + rows[0].failed_requests, + rows[0].cache_read_input_tokens, + rows[0].cache_creation_input_tokens, + rows[0].total_response_time_ms, + rows[0].timed_requests, + rows[0].spend + ), + (11, 22, 33, 44, 55, 66, 77, 88, 99, 0.75) + ); + assert_eq!(rows[0].date, "2026-01-15"); + assert_eq!(rows[0].created_at, initial_created_at); + assert!(rows[0].updated_at > seeded_updated_at); + assert_eq!(rows[0].updated_at, expected_updated_at); + + assert!(!rows[1].id.is_empty()); + assert_eq!(rows[1].team_id.as_deref(), Some(team_id)); + assert_eq!(rows[1].api_key, api_key); + assert_eq!(rows[1].model.as_deref(), Some(model)); + assert_eq!( + rows[1].custom_llm_provider.as_deref(), + Some(custom_llm_provider) + ); + assert_eq!(rows[1].endpoint.as_deref(), Some("endpoint-sentinel")); + assert_eq!( + rows[1].model_group.as_deref(), + Some("sentinel-endpoint-group") + ); + assert_eq!(rows[1].mcp_namespaced_tool_name.as_deref(), Some("")); + assert_eq!( + ( + rows[1].prompt_tokens, + rows[1].completion_tokens, + rows[1].api_requests, + rows[1].successful_requests, + rows[1].failed_requests, + rows[1].cache_read_input_tokens, + rows[1].cache_creation_input_tokens, + rows[1].total_response_time_ms, + rows[1].timed_requests, + rows[1].spend + ), + (101, 102, 103, 104, 105, 106, 107, 108, 109, 1.25) + ); + assert_eq!(rows[1].date, "2026-01-15"); + assert!(rows[1].updated_at > seeded_updated_at); + assert_eq!(rows[1].updated_at, expected_updated_at); + + assert!(!rows[2].id.is_empty()); + assert_eq!(rows[2].team_id.as_deref(), Some(team_id)); + assert_eq!(rows[2].api_key, "api-key-sentinel"); + assert_eq!(rows[2].model.as_deref(), Some(model)); + assert_eq!( + rows[2].custom_llm_provider.as_deref(), + Some(custom_llm_provider) + ); + assert_eq!(rows[2].endpoint.as_deref(), Some(endpoint)); + assert_eq!(rows[2].model_group.as_deref(), Some("sentinel-key-group")); + assert_eq!(rows[2].mcp_namespaced_tool_name.as_deref(), Some("")); + assert_eq!( + ( + rows[2].prompt_tokens, + rows[2].completion_tokens, + rows[2].api_requests, + rows[2].successful_requests, + rows[2].failed_requests, + rows[2].cache_read_input_tokens, + rows[2].cache_creation_input_tokens, + rows[2].total_response_time_ms, + rows[2].timed_requests, + rows[2].spend + ), + (201, 202, 203, 204, 205, 206, 207, 208, 209, 2.5) + ); + assert_eq!(rows[2].date, "2026-01-15"); + assert!(rows[2].updated_at > seeded_updated_at); + assert_eq!(rows[2].updated_at, expected_updated_at); + + assert!(!rows[3].id.is_empty()); + assert_eq!(rows[3].team_id.as_deref(), Some("entity-sentinel")); + assert_eq!(rows[3].api_key, api_key); + assert_eq!(rows[3].model.as_deref(), Some(model)); + assert_eq!( + rows[3].custom_llm_provider.as_deref(), + Some(custom_llm_provider) + ); + assert_eq!(rows[3].endpoint.as_deref(), Some(endpoint)); + assert_eq!( + rows[3].model_group.as_deref(), + Some("sentinel-entity-group") + ); + assert_eq!(rows[3].mcp_namespaced_tool_name.as_deref(), Some("")); + assert_eq!( + ( + rows[3].prompt_tokens, + rows[3].completion_tokens, + rows[3].api_requests, + rows[3].successful_requests, + rows[3].failed_requests, + rows[3].cache_read_input_tokens, + rows[3].cache_creation_input_tokens, + rows[3].total_response_time_ms, + rows[3].timed_requests, + rows[3].spend + ), + (301, 302, 303, 304, 305, 306, 307, 308, 309, 3.75) + ); + assert_eq!(rows[3].date, "2026-01-15"); + assert!(rows[3].updated_at > seeded_updated_at); + assert_eq!(rows[3].updated_at, expected_updated_at); + Ok(()) + } +} diff --git a/litellm-rust/crates/db/src/spend/mod.rs b/litellm-rust/crates/db/src/spend/mod.rs index f50440e06de..3f71ae4bf9f 100644 --- a/litellm-rust/crates/db/src/spend/mod.rs +++ b/litellm-rust/crates/db/src/spend/mod.rs @@ -1,3 +1,4 @@ +mod daily_team_spend; mod daily_user_spend; mod end_user_spend; mod key_spend; @@ -7,6 +8,7 @@ mod team_member_spend; mod team_spend; mod user_spend; +pub use daily_team_spend::upsert_daily_team_spend; pub use daily_user_spend::{DailySpend, upsert_daily_user_spend}; pub use end_user_spend::add_end_user_spend; pub use key_spend::{EntitySpend, add_key_spend};