mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
feat(db): add typed daily team spend upsert (#48236)
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
1b874ceb85
commit
a2545134f3
5 changed files with 559 additions and 1 deletions
|
|
@ -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"
|
||||
}
|
||||
|
|
@ -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')
|
||||
|
|
@ -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,
|
||||
};
|
||||
|
|
|
|||
485
litellm-rust/crates/db/src/spend/daily_team_spend.rs
Normal file
485
litellm-rust/crates/db/src/spend/daily_team_spend.rs
Normal file
|
|
@ -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::<Vec<_>>()
|
||||
};
|
||||
}
|
||||
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<String>],
|
||||
&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<PgPool, Box<dyn StdError>>,
|
||||
) -> Result<(), Box<dyn StdError>> {
|
||||
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<PgPool, Box<dyn StdError>>,
|
||||
) -> Result<(), Box<dyn StdError>> {
|
||||
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<DailyTeamRow> = 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(())
|
||||
}
|
||||
}
|
||||
|
|
@ -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};
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue