diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 9b60fb5520e..5da517b88b3 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3752,10 +3752,12 @@ dependencies = [ "base64 0.22.1", "envy", "futures-util", + "http 1.4.2", "http-body-util", "litellm-auth-types", "litellm-config", "litellm-db", + "litellm-db-testing", "litellm-gateway-auth", "litellm-gateway-hooks", "litellm-gateway-inference", @@ -3768,10 +3770,12 @@ dependencies = [ "litellm-router", "litellm-secrets", "litellm-tracing", + "rand 0.8.7", "redis", "rstest", "serde", "serde_json", + "sqlx", "strum", "tempfile", "thiserror 2.0.19", diff --git a/litellm-rust/crates/db/.sqlx/query-07328686221e3bfe19c205d68fe3420df984f9ae9460ceefa9a16e61bae0e3e2.json b/litellm-rust/crates/db/.sqlx/query-07328686221e3bfe19c205d68fe3420df984f9ae9460ceefa9a16e61bae0e3e2.json new file mode 100644 index 00000000000..876eb6b645c --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-07328686221e3bfe19c205d68fe3420df984f9ae9460ceefa9a16e61bae0e3e2.json @@ -0,0 +1,15 @@ +{ + "db_name": "PostgreSQL", + "query": "WITH added AS (SELECT * FROM UNNEST($1::text[], $2::float8[]) AS a(id, cost)),\n locked AS (SELECT t.token FROM \"LiteLLM_VerificationToken\" t\n WHERE t.token IN (SELECT id FROM added) ORDER BY t.token FOR UPDATE)\n UPDATE \"LiteLLM_VerificationToken\" t\n SET spend = t.spend + added.cost,\n total_spend = t.total_spend + added.cost,\n last_active = (NOW() AT TIME ZONE 'UTC'),\n updated_at = (NOW() AT TIME ZONE 'UTC')\n FROM added\n WHERE t.token = added.id AND t.token IN (SELECT token FROM locked)", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "TextArray", + "Float8Array" + ] + }, + "nullable": [] + }, + "hash": "07328686221e3bfe19c205d68fe3420df984f9ae9460ceefa9a16e61bae0e3e2" +} diff --git a/litellm-rust/crates/db/.sqlx/query-5c74f47a2b84574133c520d235192bf33e16fea1dbb0c32cf89a376a874927f4.json b/litellm-rust/crates/db/.sqlx/query-5c74f47a2b84574133c520d235192bf33e16fea1dbb0c32cf89a376a874927f4.json new file mode 100644 index 00000000000..7eef5278cb5 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-5c74f47a2b84574133c520d235192bf33e16fea1dbb0c32cf89a376a874927f4.json @@ -0,0 +1,21 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT set_config('lock_timeout', '5000ms', true)", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "set_config", + "type_info": "Text", + "origin": "Expression" + } + ], + "parameters": { + "Left": [] + }, + "nullable": [ + null + ] + }, + "hash": "5c74f47a2b84574133c520d235192bf33e16fea1dbb0c32cf89a376a874927f4" +} diff --git a/litellm-rust/crates/db/.sqlx/query-8337c5f41713afadeff72900c4cd48cba84f072116aca0fcedf3566f81fa14b6.json b/litellm-rust/crates/db/.sqlx/query-8337c5f41713afadeff72900c4cd48cba84f072116aca0fcedf3566f81fa14b6.json new file mode 100644 index 00000000000..071a2a834ef --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-8337c5f41713afadeff72900c4cd48cba84f072116aca0fcedf3566f81fa14b6.json @@ -0,0 +1,14 @@ +{ + "db_name": "PostgreSQL", + "query": "INSERT INTO \"LiteLLM_TeamMembership\" (user_id, team_id, spend, total_spend)\n SELECT member.user_id, member.team_id, member.cost, member.cost\n FROM jsonb_to_recordset($1::jsonb) AS member(user_id text, team_id text, cost float8)\n WHERE EXISTS (SELECT 1 FROM \"LiteLLM_TeamTable\" t 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 OR EXISTS (SELECT 1 FROM \"LiteLLM_TeamMembership\" m\n WHERE m.user_id = member.user_id AND m.team_id = member.team_id)\n ON CONFLICT (user_id, team_id) DO UPDATE\n SET spend = \"LiteLLM_TeamMembership\".spend + EXCLUDED.spend,\n total_spend = \"LiteLLM_TeamMembership\".total_spend + EXCLUDED.total_spend", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "Jsonb" + ] + }, + "nullable": [] + }, + "hash": "8337c5f41713afadeff72900c4cd48cba84f072116aca0fcedf3566f81fa14b6" +} diff --git a/litellm-rust/crates/db/.sqlx/query-83be2dd1f5892e3581ecee3d2d270fa5f0c55f253258f26bcae44fe2be097272.json b/litellm-rust/crates/db/.sqlx/query-83be2dd1f5892e3581ecee3d2d270fa5f0c55f253258f26bcae44fe2be097272.json new file mode 100644 index 00000000000..97ebe2f8fda --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-83be2dd1f5892e3581ecee3d2d270fa5f0c55f253258f26bcae44fe2be097272.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS \"locked!\"", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "locked!", + "type_info": "Bool", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "Text" + ] + }, + "nullable": [ + null + ] + }, + "hash": "83be2dd1f5892e3581ecee3d2d270fa5f0c55f253258f26bcae44fe2be097272" +} diff --git a/litellm-rust/crates/db/.sqlx/query-8abe8b069d1b8bb22f4ff49046e5cd10df9db5c4508b859068d800ca8be02b93.json b/litellm-rust/crates/db/.sqlx/query-8abe8b069d1b8bb22f4ff49046e5cd10df9db5c4508b859068d800ca8be02b93.json new file mode 100644 index 00000000000..db44ac1c3c6 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-8abe8b069d1b8bb22f4ff49046e5cd10df9db5c4508b859068d800ca8be02b93.json @@ -0,0 +1,40 @@ +{ + "db_name": "PostgreSQL", + "query": "INSERT INTO \"LiteLLM_SpendLogs\" (\n request_id, call_type, api_key, spend, total_tokens, prompt_tokens,\n completion_tokens, \"startTime\", \"endTime\", \"completionStartTime\",\n request_duration_ms, model, model_id, model_group, custom_llm_provider, api_base,\n \"user\", metadata, cache_hit, cache_key, request_tags, team_id, organization_id,\n end_user, requester_ip_address, messages, response, proxy_server_request,\n session_id, status, litellm_call_id)\n SELECT l.request_id, l.call_type, l.api_key, l.spend, l.total_tokens, l.prompt_tokens,\n l.completion_tokens, l.start_time::timestamp(3), l.end_time::timestamp(3),\n l.completion_start_time::timestamp(3), l.request_duration_ms, l.model,\n l.model_id, l.model_group, l.custom_llm_provider, l.api_base, l.user_id,\n l.metadata, l.cache_hit, l.cache_key, '[]'::jsonb, l.team_id,\n l.organization_id, l.end_user, l.requester_ip_address, '{}'::jsonb,\n '{}'::jsonb, '{}'::jsonb, l.session_id, l.status, l.litellm_call_id\n FROM UNNEST($1::text[], $2::text[], $3::text[], $4::float8[], $5::int4[], $6::int4[],\n $7::int4[], $8::text[], $9::text[], $10::text[], $11::int4[], $12::text[],\n $13::text[], $14::text[], $15::text[], $16::text[], $17::text[],\n $18::jsonb[], $19::text[], $20::text[], $21::text[], $22::text[],\n $23::text[], $24::text[], $25::text[], $26::text[], $27::text[])\n AS l(request_id, call_type, api_key, spend, total_tokens, prompt_tokens,\n completion_tokens, start_time, end_time, completion_start_time,\n request_duration_ms, model, model_id, model_group, custom_llm_provider,\n api_base, user_id, metadata, cache_hit, cache_key, team_id, organization_id,\n end_user, requester_ip_address, session_id, status, litellm_call_id)\n ON CONFLICT (request_id) DO NOTHING", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "TextArray", + "TextArray", + "TextArray", + "Float8Array", + "Int4Array", + "Int4Array", + "Int4Array", + "TextArray", + "TextArray", + "TextArray", + "Int4Array", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "JsonbArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray" + ] + }, + "nullable": [] + }, + "hash": "8abe8b069d1b8bb22f4ff49046e5cd10df9db5c4508b859068d800ca8be02b93" +} diff --git a/litellm-rust/crates/db/.sqlx/query-94187798c95011cae91b61417c8855a7324d296ae3f15ccee839a29007693f87.json b/litellm-rust/crates/db/.sqlx/query-94187798c95011cae91b61417c8855a7324d296ae3f15ccee839a29007693f87.json new file mode 100644 index 00000000000..9eb6a78c06e --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-94187798c95011cae91b61417c8855a7324d296ae3f15ccee839a29007693f87.json @@ -0,0 +1,30 @@ +{ + "db_name": "PostgreSQL", + "query": "INSERT INTO \"LiteLLM_DailyUserSpend\" (id, user_id, date, api_key, model,\n custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group,\n prompt_tokens, completion_tokens, api_requests, successful_requests,\n failed_requests, cache_read_input_tokens, cache_creation_input_tokens,\n total_response_time_ms, timed_requests, spend, updated_at)\n SELECT gen_random_uuid()::text, r.entity_id, r.date, r.api_key, r.model,\n r.custom_llm_provider, '', 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')\n FROM UNNEST($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[])\n AS r(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)\n ORDER BY r.date, r.entity_id, r.api_key, r.model, r.custom_llm_provider\n ON CONFLICT (user_id, date, api_key, model, custom_llm_provider,\n mcp_namespaced_tool_name, endpoint) DO UPDATE SET\n prompt_tokens = \"LiteLLM_DailyUserSpend\".prompt_tokens + EXCLUDED.prompt_tokens,\n completion_tokens = \"LiteLLM_DailyUserSpend\".completion_tokens + EXCLUDED.completion_tokens,\n api_requests = \"LiteLLM_DailyUserSpend\".api_requests + EXCLUDED.api_requests,\n successful_requests = \"LiteLLM_DailyUserSpend\".successful_requests + EXCLUDED.successful_requests,\n failed_requests = \"LiteLLM_DailyUserSpend\".failed_requests + EXCLUDED.failed_requests,\n cache_read_input_tokens = \"LiteLLM_DailyUserSpend\".cache_read_input_tokens + EXCLUDED.cache_read_input_tokens,\n cache_creation_input_tokens = \"LiteLLM_DailyUserSpend\".cache_creation_input_tokens + EXCLUDED.cache_creation_input_tokens,\n total_response_time_ms = \"LiteLLM_DailyUserSpend\".total_response_time_ms + EXCLUDED.total_response_time_ms,\n timed_requests = \"LiteLLM_DailyUserSpend\".timed_requests + EXCLUDED.timed_requests,\n spend = \"LiteLLM_DailyUserSpend\".spend + EXCLUDED.spend,\n updated_at = (NOW() AT TIME ZONE 'UTC')", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Float8Array" + ] + }, + "nullable": [] + }, + "hash": "94187798c95011cae91b61417c8855a7324d296ae3f15ccee839a29007693f87" +} diff --git a/litellm-rust/crates/db/.sqlx/query-94e8e0b4897318bd005c9d4cd50fd1abc967056d615940b0804915a3ebed9a2b.json b/litellm-rust/crates/db/.sqlx/query-94e8e0b4897318bd005c9d4cd50fd1abc967056d615940b0804915a3ebed9a2b.json new file mode 100644 index 00000000000..f2d7ba555b9 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-94e8e0b4897318bd005c9d4cd50fd1abc967056d615940b0804915a3ebed9a2b.json @@ -0,0 +1,15 @@ +{ + "db_name": "PostgreSQL", + "query": "WITH added AS (SELECT * FROM UNNEST($1::text[], $2::float8[]) AS a(id, cost)),\n locked AS (SELECT t.team_id FROM \"LiteLLM_TeamTable\" t\n WHERE t.team_id IN (SELECT id FROM added) ORDER BY t.team_id FOR UPDATE)\n UPDATE \"LiteLLM_TeamTable\" t\n SET spend = t.spend + added.cost, updated_at = (NOW() AT TIME ZONE 'UTC')\n FROM added\n WHERE t.team_id = added.id AND t.team_id IN (SELECT team_id FROM locked)", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "TextArray", + "Float8Array" + ] + }, + "nullable": [] + }, + "hash": "94e8e0b4897318bd005c9d4cd50fd1abc967056d615940b0804915a3ebed9a2b" +} diff --git a/litellm-rust/crates/db/.sqlx/query-9e6d4c7704bfbb38ce6d12b647376fbde114f184f7cc6efe70d49d4c78e8f036.json b/litellm-rust/crates/db/.sqlx/query-9e6d4c7704bfbb38ce6d12b647376fbde114f184f7cc6efe70d49d4c78e8f036.json new file mode 100644 index 00000000000..c2d058c96ac --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-9e6d4c7704bfbb38ce6d12b647376fbde114f184f7cc6efe70d49d4c78e8f036.json @@ -0,0 +1,30 @@ +{ + "db_name": "PostgreSQL", + "query": "INSERT INTO \"LiteLLM_DailyTeamSpend\" (id, team_id, date, api_key, model,\n custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group,\n prompt_tokens, completion_tokens, api_requests, successful_requests,\n failed_requests, cache_read_input_tokens, cache_creation_input_tokens,\n total_response_time_ms, timed_requests, spend, updated_at)\n SELECT gen_random_uuid()::text, r.entity_id, r.date, r.api_key, r.model,\n r.custom_llm_provider, '', 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')\n FROM UNNEST($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[])\n AS r(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)\n ORDER BY r.date, r.entity_id, r.api_key, r.model, r.custom_llm_provider\n ON CONFLICT (team_id, date, api_key, model, custom_llm_provider,\n mcp_namespaced_tool_name, endpoint) 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')", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Float8Array" + ] + }, + "nullable": [] + }, + "hash": "9e6d4c7704bfbb38ce6d12b647376fbde114f184f7cc6efe70d49d4c78e8f036" +} diff --git a/litellm-rust/crates/db/.sqlx/query-b0e3f762f1dd5efdd4672953f7745ad5e942f4896cac72c0f002f3c544e25094.json b/litellm-rust/crates/db/.sqlx/query-b0e3f762f1dd5efdd4672953f7745ad5e942f4896cac72c0f002f3c544e25094.json new file mode 100644 index 00000000000..90d6f426bb6 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-b0e3f762f1dd5efdd4672953f7745ad5e942f4896cac72c0f002f3c544e25094.json @@ -0,0 +1,15 @@ +{ + "db_name": "PostgreSQL", + "query": "WITH added AS (SELECT * FROM UNNEST($1::text[], $2::float8[]) AS a(id, cost)),\n locked AS (SELECT u.user_id FROM \"LiteLLM_UserTable\" u\n WHERE u.user_id IN (SELECT id FROM added) ORDER BY u.user_id FOR UPDATE)\n UPDATE \"LiteLLM_UserTable\" u\n SET spend = u.spend + added.cost, updated_at = (NOW() AT TIME ZONE 'UTC')\n FROM added\n WHERE u.user_id = added.id AND u.user_id IN (SELECT user_id FROM locked)", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "TextArray", + "Float8Array" + ] + }, + "nullable": [] + }, + "hash": "b0e3f762f1dd5efdd4672953f7745ad5e942f4896cac72c0f002f3c544e25094" +} diff --git a/litellm-rust/crates/db/.sqlx/query-bb412ba2896341351e1f6ca96761fd9c45b29330c73c66cdbee80e259dbb391b.json b/litellm-rust/crates/db/.sqlx/query-bb412ba2896341351e1f6ca96761fd9c45b29330c73c66cdbee80e259dbb391b.json new file mode 100644 index 00000000000..7e046d3f04f --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-bb412ba2896341351e1f6ca96761fd9c45b29330c73c66cdbee80e259dbb391b.json @@ -0,0 +1,15 @@ +{ + "db_name": "PostgreSQL", + "query": "INSERT INTO \"LiteLLM_EndUserTable\" (user_id, spend, blocked)\n SELECT id, cost, false FROM UNNEST($1::text[], $2::float8[]) AS a(id, cost) ORDER BY id\n ON CONFLICT (user_id) DO UPDATE\n SET spend = \"LiteLLM_EndUserTable\".spend + EXCLUDED.spend", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "TextArray", + "Float8Array" + ] + }, + "nullable": [] + }, + "hash": "bb412ba2896341351e1f6ca96761fd9c45b29330c73c66cdbee80e259dbb391b" +} diff --git a/litellm-rust/crates/db/.sqlx/query-f56e12df3f7b46178c50a58013ecb0b2cba889274ab089938b1fe011c7c94df2.json b/litellm-rust/crates/db/.sqlx/query-f56e12df3f7b46178c50a58013ecb0b2cba889274ab089938b1fe011c7c94df2.json new file mode 100644 index 00000000000..5824982661f --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-f56e12df3f7b46178c50a58013ecb0b2cba889274ab089938b1fe011c7c94df2.json @@ -0,0 +1,30 @@ +{ + "db_name": "PostgreSQL", + "query": "INSERT INTO \"LiteLLM_DailyEndUserSpend\" (id, end_user_id, date, api_key, model,\n custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group,\n prompt_tokens, completion_tokens, api_requests, successful_requests,\n failed_requests, cache_read_input_tokens, cache_creation_input_tokens,\n total_response_time_ms, timed_requests, spend, updated_at)\n SELECT gen_random_uuid()::text, r.entity_id, r.date, r.api_key, r.model,\n r.custom_llm_provider, '', 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')\n FROM UNNEST($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[])\n AS r(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)\n ORDER BY r.date, r.entity_id, r.api_key, r.model, r.custom_llm_provider\n ON CONFLICT (end_user_id, date, api_key, model, custom_llm_provider,\n mcp_namespaced_tool_name, endpoint) DO UPDATE SET\n prompt_tokens = \"LiteLLM_DailyEndUserSpend\".prompt_tokens + EXCLUDED.prompt_tokens,\n completion_tokens = \"LiteLLM_DailyEndUserSpend\".completion_tokens + EXCLUDED.completion_tokens,\n api_requests = \"LiteLLM_DailyEndUserSpend\".api_requests + EXCLUDED.api_requests,\n successful_requests = \"LiteLLM_DailyEndUserSpend\".successful_requests + EXCLUDED.successful_requests,\n failed_requests = \"LiteLLM_DailyEndUserSpend\".failed_requests + EXCLUDED.failed_requests,\n cache_read_input_tokens = \"LiteLLM_DailyEndUserSpend\".cache_read_input_tokens + EXCLUDED.cache_read_input_tokens,\n cache_creation_input_tokens = \"LiteLLM_DailyEndUserSpend\".cache_creation_input_tokens + EXCLUDED.cache_creation_input_tokens,\n total_response_time_ms = \"LiteLLM_DailyEndUserSpend\".total_response_time_ms + EXCLUDED.total_response_time_ms,\n timed_requests = \"LiteLLM_DailyEndUserSpend\".timed_requests + EXCLUDED.timed_requests,\n spend = \"LiteLLM_DailyEndUserSpend\".spend + EXCLUDED.spend,\n updated_at = (NOW() AT TIME ZONE 'UTC')", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Float8Array" + ] + }, + "nullable": [] + }, + "hash": "f56e12df3f7b46178c50a58013ecb0b2cba889274ab089938b1fe011c7c94df2" +} diff --git a/litellm-rust/crates/db/src/lib.rs b/litellm-rust/crates/db/src/lib.rs index c099edb0003..39ef58241b2 100644 --- a/litellm-rust/crates/db/src/lib.rs +++ b/litellm-rust/crates/db/src/lib.rs @@ -1,11 +1,17 @@ mod error; mod failure; mod pool; +mod spend; mod verification_token; pub use error::Error; pub use failure::DbFailure; pub use pool::{PoolSettings, connect_lazy}; +pub use spend::{ + DailySpend, DailyTable, EntitySpend, SpendLog, TeamMemberSpend, add_end_user_spend, + add_key_spend, add_team_member_spend, add_team_spend, add_user_spend, insert_spend_logs, + upsert_daily_spend, +}; pub use sqlx::PgPool; pub use verification_token::{ Restrictions, Team, TeamMembership, User, VerificationToken, find_verification_token, diff --git a/litellm-rust/crates/db/src/spend.rs b/litellm-rust/crates/db/src/spend.rs new file mode 100644 index 00000000000..67d11e4c7d3 --- /dev/null +++ b/litellm-rust/crates/db/src/spend.rs @@ -0,0 +1,451 @@ +//! The proxy's spend writes: additive spend on the key, user, team, team membership and end user +//! rows, rows in `LiteLLM_SpendLogs`, and the daily activity tables. Every statement takes a +//! batch, so a flush costs one round trip per table. + +use sqlx::{PgPool, Postgres, Transaction}; + +use crate::Error; + +/// Spend to add per id. Ids are locked in sorted order, as the proxy does, so writers on other +/// replicas cannot deadlock with this one. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct EntitySpend { + pub ids: Vec, + pub costs: Vec, +} + +impl EntitySpend { + pub fn is_empty(&self) -> bool { + self.ids.is_empty() + } +} + +impl FromIterator<(String, f64)> for EntitySpend { + fn from_iter>(entries: I) -> Self { + let mut entries: Vec<_> = entries.into_iter().collect(); + entries.sort_by(|left, right| left.0.cmp(&right.0)); + let (ids, costs) = entries.into_iter().unzip(); + Self { ids, costs } + } +} + +#[derive(Clone, Debug, PartialEq)] +pub struct TeamMemberSpend { + pub team_id: String, + pub user_id: String, + pub cost: f64, +} + +/// Starts a transaction that waits at most 5 s for a row lock, the proxy's +/// `SPEND_ROLLUP_LOCK_TIMEOUT_MS`. +async fn begin(pool: &PgPool) -> Result, Error> { + let mut transaction = pool.begin().await?; + sqlx::query!("SELECT set_config('lock_timeout', '5000ms', true)") + .fetch_one(&mut *transaction) + .await?; + Ok(transaction) +} + +/// `LiteLLM_VerificationToken.spend` and `total_spend`, and `last_active` set to now. +pub async fn add_key_spend(pool: &PgPool, spend: &EntitySpend) -> Result<(), Error> { + if spend.is_empty() { + return Ok(()); + } + let mut transaction = begin(pool).await?; + sqlx::query!( + r#"WITH added AS (SELECT * FROM UNNEST($1::text[], $2::float8[]) AS a(id, cost)), + locked AS (SELECT t.token FROM "LiteLLM_VerificationToken" t + WHERE t.token IN (SELECT id FROM added) ORDER BY t.token FOR UPDATE) + UPDATE "LiteLLM_VerificationToken" t + SET spend = t.spend + added.cost, + total_spend = t.total_spend + added.cost, + last_active = (NOW() AT TIME ZONE 'UTC'), + updated_at = (NOW() AT TIME ZONE 'UTC') + FROM added + WHERE t.token = added.id AND t.token IN (SELECT token FROM locked)"#, + &spend.ids, + &spend.costs, + ) + .execute(&mut *transaction) + .await?; + transaction.commit().await?; + Ok(()) +} + +/// `LiteLLM_UserTable.spend`. A user without a row is skipped, as the proxy skips it. +pub async fn add_user_spend(pool: &PgPool, spend: &EntitySpend) -> Result<(), Error> { + if spend.is_empty() { + return Ok(()); + } + let mut transaction = begin(pool).await?; + sqlx::query!( + r#"WITH added AS (SELECT * FROM UNNEST($1::text[], $2::float8[]) AS a(id, cost)), + locked AS (SELECT u.user_id FROM "LiteLLM_UserTable" u + WHERE u.user_id IN (SELECT id FROM added) ORDER BY u.user_id FOR UPDATE) + UPDATE "LiteLLM_UserTable" u + SET spend = u.spend + added.cost, updated_at = (NOW() AT TIME ZONE 'UTC') + FROM added + WHERE u.user_id = added.id AND u.user_id IN (SELECT user_id FROM locked)"#, + &spend.ids, + &spend.costs, + ) + .execute(&mut *transaction) + .await?; + transaction.commit().await?; + Ok(()) +} + +/// `LiteLLM_TeamTable.spend`. +pub async fn add_team_spend(pool: &PgPool, spend: &EntitySpend) -> Result<(), Error> { + if spend.is_empty() { + return Ok(()); + } + let mut transaction = begin(pool).await?; + sqlx::query!( + r#"WITH added AS (SELECT * FROM UNNEST($1::text[], $2::float8[]) AS a(id, cost)), + locked AS (SELECT t.team_id FROM "LiteLLM_TeamTable" t + WHERE t.team_id IN (SELECT id FROM added) ORDER BY t.team_id FOR UPDATE) + UPDATE "LiteLLM_TeamTable" t + SET spend = t.spend + added.cost, updated_at = (NOW() AT TIME ZONE 'UTC') + FROM added + WHERE t.team_id = added.id AND t.team_id IN (SELECT team_id FROM locked)"#, + &spend.ids, + &spend.costs, + ) + .execute(&mut *transaction) + .await?; + transaction.commit().await?; + Ok(()) +} + +/// `LiteLLM_EndUserTable.spend`, creating the end user on first spend as the proxy does. +pub async fn add_end_user_spend(pool: &PgPool, spend: &EntitySpend) -> Result<(), Error> { + if spend.is_empty() { + return Ok(()); + } + let mut transaction = begin(pool).await?; + sqlx::query!( + r#"INSERT INTO "LiteLLM_EndUserTable" (user_id, spend, blocked) + SELECT id, cost, false FROM UNNEST($1::text[], $2::float8[]) AS a(id, cost) ORDER BY id + ON CONFLICT (user_id) DO UPDATE + SET spend = "LiteLLM_EndUserTable".spend + EXCLUDED.spend"#, + &spend.ids, + &spend.costs, + ) + .execute(&mut *transaction) + .await?; + transaction.commit().await?; + Ok(()) +} + +/// `LiteLLM_TeamMembership.spend` and `total_spend`. A membership row is created only for a +/// user on the team's roster, under a per-team advisory lock, as the proxy writes it. +pub async fn add_team_member_spend( + pool: &PgPool, + 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 transaction = begin(pool).await?; + let mut teams: Vec<&str> = members + .iter() + .map(|member| member.team_id.as_str()) + .collect(); + teams.dedup(); + for team in teams { + sqlx::query!( + r#"SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS "locked!""#, + team + ) + .fetch_one(&mut *transaction) + .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!( + r#"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"#, + rows, + ) + .execute(&mut *transaction) + .await?; + transaction.commit().await?; + Ok(()) +} + +/// One `LiteLLM_SpendLogs` row. Times are UTC, `YYYY-MM-DD HH:MM:SS.fff`. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct SpendLog { + pub request_id: String, + pub call_type: String, + pub api_key: String, + pub spend: f64, + pub total_tokens: i32, + pub prompt_tokens: i32, + pub completion_tokens: i32, + pub start_time: String, + pub end_time: String, + pub completion_start_time: Option, + pub request_duration_ms: Option, + pub model: String, + pub model_id: String, + pub model_group: String, + pub custom_llm_provider: String, + pub api_base: String, + pub user: String, + pub metadata: serde_json::Value, + pub cache_hit: String, + pub cache_key: String, + pub team_id: String, + pub organization_id: String, + pub end_user: String, + pub requester_ip_address: Option, + pub session_id: Option, + pub status: String, + pub litellm_call_id: Option, +} + +/// Inserts spend logs, skipping request ids already written, as the proxy's +/// `create_many(skip_duplicates=True)` does. +pub async fn insert_spend_logs(pool: &PgPool, logs: &[SpendLog]) -> Result<(), Error> { + if logs.is_empty() { + return Ok(()); + } + macro_rules! column { + ($field:ident) => { + logs.iter() + .map(|log| log.$field.clone()) + .collect::>() + }; + } + sqlx::query!( + r#"INSERT INTO "LiteLLM_SpendLogs" ( + request_id, call_type, api_key, spend, total_tokens, prompt_tokens, + completion_tokens, "startTime", "endTime", "completionStartTime", + request_duration_ms, model, model_id, model_group, custom_llm_provider, api_base, + "user", metadata, cache_hit, cache_key, request_tags, team_id, organization_id, + end_user, requester_ip_address, messages, response, proxy_server_request, + session_id, status, litellm_call_id) + SELECT l.request_id, l.call_type, l.api_key, l.spend, l.total_tokens, l.prompt_tokens, + l.completion_tokens, l.start_time::timestamp(3), l.end_time::timestamp(3), + l.completion_start_time::timestamp(3), l.request_duration_ms, l.model, + l.model_id, l.model_group, l.custom_llm_provider, l.api_base, l.user_id, + l.metadata, l.cache_hit, l.cache_key, '[]'::jsonb, l.team_id, + l.organization_id, l.end_user, l.requester_ip_address, '{}'::jsonb, + '{}'::jsonb, '{}'::jsonb, l.session_id, l.status, l.litellm_call_id + FROM UNNEST($1::text[], $2::text[], $3::text[], $4::float8[], $5::int4[], $6::int4[], + $7::int4[], $8::text[], $9::text[], $10::text[], $11::int4[], $12::text[], + $13::text[], $14::text[], $15::text[], $16::text[], $17::text[], + $18::jsonb[], $19::text[], $20::text[], $21::text[], $22::text[], + $23::text[], $24::text[], $25::text[], $26::text[], $27::text[]) + AS l(request_id, call_type, api_key, spend, total_tokens, prompt_tokens, + completion_tokens, start_time, end_time, completion_start_time, + request_duration_ms, model, model_id, model_group, custom_llm_provider, + api_base, user_id, metadata, cache_hit, cache_key, team_id, organization_id, + end_user, requester_ip_address, session_id, status, litellm_call_id) + ON CONFLICT (request_id) DO NOTHING"#, + &column!(request_id), + &column!(call_type), + &column!(api_key), + &column!(spend), + &column!(total_tokens), + &column!(prompt_tokens), + &column!(completion_tokens), + &column!(start_time), + &column!(end_time), + &column!(completion_start_time) as &[Option], + &column!(request_duration_ms) as &[Option], + &column!(model), + &column!(model_id), + &column!(model_group), + &column!(custom_llm_provider), + &column!(api_base), + &column!(user), + &column!(metadata), + &column!(cache_hit), + &column!(cache_key), + &column!(team_id), + &column!(organization_id), + &column!(end_user), + &column!(requester_ip_address) as &[Option], + &column!(session_id) as &[Option], + &column!(status), + &column!(litellm_call_id) as &[Option], + ) + .execute(pool) + .await?; + Ok(()) +} + +/// The daily activity tables, keyed by user, team or end user. +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +pub enum DailyTable { + User, + Team, + EndUser, +} + +/// One row of a daily activity table. Key columns are `''` rather than NULL, as the proxy +/// writes them, so its unique index matches. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct DailySpend { + pub entity_id: String, + /// `YYYY-MM-DD`, UTC. + pub date: String, + pub api_key: String, + pub model: String, + pub custom_llm_provider: String, + pub endpoint: String, + /// Written on insert only. + pub model_group: Option, + pub prompt_tokens: i64, + pub completion_tokens: i64, + pub api_requests: i64, + pub successful_requests: i64, + pub failed_requests: i64, + pub cache_read_input_tokens: i64, + pub cache_creation_input_tokens: i64, + pub total_response_time_ms: i64, + pub timed_requests: i64, + pub spend: f64, +} + +/// Adds `rows` to a daily table. Rows must be unique by their key columns. +pub async fn upsert_daily_spend( + pool: &PgPool, + table: DailyTable, + rows: &[DailySpend], +) -> Result<(), Error> { + if rows.is_empty() { + return Ok(()); + } + macro_rules! column { + ($field:ident) => { + rows.iter() + .map(|row| row.$field.clone()) + .collect::>() + }; + } + let entity_id = column!(entity_id); + let date = column!(date); + let api_key = column!(api_key); + let model = column!(model); + let custom_llm_provider = column!(custom_llm_provider); + let endpoint = column!(endpoint); + let model_group: Vec> = column!(model_group); + let prompt_tokens = column!(prompt_tokens); + let completion_tokens = column!(completion_tokens); + let api_requests = column!(api_requests); + let successful_requests = column!(successful_requests); + let failed_requests = column!(failed_requests); + let cache_read = column!(cache_read_input_tokens); + let cache_creation = column!(cache_creation_input_tokens); + let response_time = column!(total_response_time_ms); + let timed_requests = column!(timed_requests); + let spend = column!(spend); + let mut transaction = begin(pool).await?; + // One statement per table: a table name cannot be a parameter. + macro_rules! upsert { + ($table:literal, $entity:literal) => { + sqlx::query!( + r#"INSERT INTO "# + $table + r#" (id, "# + $entity + r#", 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.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[]) + 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) + ORDER BY r.date, r.entity_id, r.api_key, r.model, r.custom_llm_provider + ON CONFLICT ("# + $entity + r#", date, api_key, model, custom_llm_provider, + mcp_namespaced_tool_name, endpoint) DO UPDATE SET + prompt_tokens = "# + $table + r#".prompt_tokens + EXCLUDED.prompt_tokens, + completion_tokens = "# + $table + r#".completion_tokens + EXCLUDED.completion_tokens, + api_requests = "# + $table + r#".api_requests + EXCLUDED.api_requests, + successful_requests = "# + $table + r#".successful_requests + EXCLUDED.successful_requests, + failed_requests = "# + $table + r#".failed_requests + EXCLUDED.failed_requests, + cache_read_input_tokens = "# + $table + r#".cache_read_input_tokens + EXCLUDED.cache_read_input_tokens, + cache_creation_input_tokens = "# + $table + r#".cache_creation_input_tokens + EXCLUDED.cache_creation_input_tokens, + total_response_time_ms = "# + $table + r#".total_response_time_ms + EXCLUDED.total_response_time_ms, + timed_requests = "# + $table + r#".timed_requests + EXCLUDED.timed_requests, + spend = "# + $table + r#".spend + EXCLUDED.spend, + updated_at = (NOW() AT TIME ZONE 'UTC')"#, + &entity_id, + &date, + &api_key, + &model, + &custom_llm_provider, + &endpoint, + &model_group as &[Option], + &prompt_tokens, + &completion_tokens, + &api_requests, + &successful_requests, + &failed_requests, + &cache_read, + &cache_creation, + &response_time, + &timed_requests, + &spend, + ) + .execute(&mut *transaction) + .await? + }; + } + match table { + DailyTable::User => upsert!(r#""LiteLLM_DailyUserSpend""#, "user_id"), + DailyTable::Team => upsert!(r#""LiteLLM_DailyTeamSpend""#, "team_id"), + DailyTable::EndUser => upsert!(r#""LiteLLM_DailyEndUserSpend""#, "end_user_id"), + }; + transaction.commit().await?; + Ok(()) +} + +impl Error { + /// Whether retrying the same write can succeed: a deadlock, a lock wait timeout, a + /// serialization failure, or a lost connection. + pub fn is_retryable(&self) -> bool { + match self { + Self::Query(sqlx::Error::Database(error)) => { + matches!(error.code().as_deref(), Some("40P01" | "55P03" | "40001")) + } + Self::Query( + sqlx::Error::Io(_) | sqlx::Error::PoolTimedOut | sqlx::Error::PoolClosed, + ) => true, + Self::Query(_) | Self::Url(_) => false, + } + } +} diff --git a/litellm-rust/crates/db/tests/spend.rs b/litellm-rust/crates/db/tests/spend.rs new file mode 100644 index 00000000000..b01ef55a68c --- /dev/null +++ b/litellm-rust/crates/db/tests/spend.rs @@ -0,0 +1,281 @@ +use litellm_db::{ + DailySpend, DailyTable, EntitySpend, SpendLog, TeamMemberSpend, add_end_user_spend, + add_key_spend, add_team_member_spend, add_team_spend, add_user_spend, insert_spend_logs, + upsert_daily_spend, +}; +use litellm_db_testing::MigratedPostgres; +use rstest::rstest; +use serde_json::{Value, json}; +use sqlx::PgPool; + +const FIXTURES: &str = r#" +INSERT INTO "LiteLLM_TeamTable" (team_id, members_with_roles, created_at, updated_at) + VALUES ('team-1', '[{"user_id": "user-1", "role": "user"}]', now(), now()); +INSERT INTO "LiteLLM_UserTable" (user_id, spend) VALUES ('user-1', 1.0); +INSERT INTO "LiteLLM_VerificationToken" (token, spend, total_spend, team_id, user_id) + VALUES ('hash-1', 2.0, 2.0, 'team-1', 'user-1'); +"#; + +#[expect( + clippy::disallowed_methods, + reason = "test fixture rows and assertions; the writes under test use the checked macro" +)] +async fn setup(pool: &PgPool) { + sqlx::raw_sql(FIXTURES).execute(pool).await.unwrap(); +} + +#[expect( + clippy::disallowed_methods, + reason = "test assertions read rows back; the writes under test use the checked macro" +)] +async fn read(pool: &PgPool, sql: &str) -> Value { + sqlx::query_scalar::<_, Value>(sqlx::AssertSqlSafe(format!( + "SELECT to_jsonb(r) FROM ({sql}) r" + ))) + .fetch_one(pool) + .await + .unwrap() +} + +#[expect( + clippy::disallowed_methods, + reason = "test assertions count rows; the writes under test use the checked macro" +)] +async fn count(pool: &PgPool, sql: &str) -> i64 { + sqlx::query_scalar::<_, i64>(sqlx::AssertSqlSafe(sql.to_owned())) + .fetch_one(pool) + .await + .unwrap() +} + +fn spend(entries: &[(&str, f64)]) -> EntitySpend { + entries + .iter() + .map(|(id, cost)| ((*id).to_owned(), *cost)) + .collect() +} + +#[rstest] +#[tokio::test] +async fn entity_spend_adds_up_and_creates_only_what_the_proxy_creates() { + let postgres = MigratedPostgres::start().await.unwrap(); + let pool = postgres.pool(); + setup(pool).await; + + for _ in 0..2 { + add_key_spend(pool, &spend(&[("hash-1", 0.25), ("missing", 1.0)])) + .await + .unwrap(); + add_user_spend(pool, &spend(&[("user-1", 0.5), ("missing", 1.0)])) + .await + .unwrap(); + add_team_spend(pool, &spend(&[("team-1", 0.125)])) + .await + .unwrap(); + add_end_user_spend(pool, &spend(&[("end-1", 0.75)])) + .await + .unwrap(); + add_team_member_spend( + pool, + &[ + TeamMemberSpend { + team_id: "team-1".into(), + user_id: "user-1".into(), + cost: 0.5, + }, + TeamMemberSpend { + team_id: "team-1".into(), + user_id: "stranger".into(), + cost: 9.0, + }, + ], + ) + .await + .unwrap(); + } + + let key = read( + pool, + r#"SELECT spend, total_spend, last_active IS NOT NULL AS active + FROM "LiteLLM_VerificationToken" WHERE token = 'hash-1'"#, + ) + .await; + assert_eq!( + key, + json!({"spend": 2.5, "total_spend": 2.5, "active": true}) + ); + let user = read( + pool, + r#"SELECT spend FROM "LiteLLM_UserTable" WHERE user_id = 'user-1'"#, + ) + .await; + assert_eq!(user, json!({"spend": 2})); + let team = read( + pool, + r#"SELECT spend FROM "LiteLLM_TeamTable" WHERE team_id = 'team-1'"#, + ) + .await; + assert_eq!(team, json!({"spend": 0.25})); + let end_user = read( + pool, + r#"SELECT spend, blocked FROM "LiteLLM_EndUserTable" WHERE user_id = 'end-1'"#, + ) + .await; + assert_eq!(end_user, json!({"spend": 1.5, "blocked": false})); + let member = read( + pool, + r#"SELECT spend, total_spend FROM "LiteLLM_TeamMembership" + WHERE team_id = 'team-1' AND user_id = 'user-1'"#, + ) + .await; + assert_eq!(member, json!({"spend": 1, "total_spend": 1})); + assert_eq!( + count(pool, r#"SELECT count(*) FROM "LiteLLM_UserTable""#).await, + 1 + ); + assert_eq!( + count( + pool, + r#"SELECT count(*) FROM "LiteLLM_TeamMembership" WHERE user_id = 'stranger'"# + ) + .await, + 0 + ); +} + +fn log(request_id: &str) -> SpendLog { + SpendLog { + request_id: request_id.into(), + call_type: "anthropic_messages".into(), + api_key: "hash-1".into(), + spend: 0.5, + total_tokens: 30, + prompt_tokens: 10, + completion_tokens: 20, + start_time: "2026-10-17 01:02:03.456".into(), + end_time: "2026-10-17 01:02:04.000".into(), + completion_start_time: Some("2026-10-17 01:02:03.500".into()), + request_duration_ms: Some(544), + model: "claude-x".into(), + model_group: "claude".into(), + custom_llm_provider: "anthropic".into(), + metadata: json!({"user_api_key": "hash-1", "status": null}), + cache_hit: "None".into(), + cache_key: "Cache OFF".into(), + status: "success".into(), + ..SpendLog::default() + } +} + +#[rstest] +#[tokio::test] +async fn spend_logs_are_written_once() { + let postgres = MigratedPostgres::start().await.unwrap(); + let pool = postgres.pool(); + insert_spend_logs(pool, &[log("a"), log("b")]) + .await + .unwrap(); + insert_spend_logs(pool, &[log("a")]).await.unwrap(); + + assert_eq!( + count(pool, r#"SELECT count(*) FROM "LiteLLM_SpendLogs""#).await, + 2 + ); + let row = read( + pool, + r#"SELECT "startTime", "completionStartTime", request_duration_ms, team_id, end_user, + metadata, request_tags, messages, response, cache_hit, status + FROM "LiteLLM_SpendLogs" WHERE request_id = 'a'"#, + ) + .await; + assert_eq!( + row, + json!({ + "startTime": "2026-10-17T01:02:03.456", + "completionStartTime": "2026-10-17T01:02:03.5", + "request_duration_ms": 544, + "team_id": "", + "end_user": "", + "metadata": {"user_api_key": "hash-1", "status": null}, + "request_tags": [], + "messages": {}, + "response": {}, + "cache_hit": "None", + "status": "success", + }) + ); +} + +fn daily(entity: &str, model_group: &str) -> DailySpend { + DailySpend { + entity_id: entity.into(), + date: "2026-10-17".into(), + api_key: "hash-1".into(), + model: "claude-x".into(), + custom_llm_provider: "anthropic".into(), + endpoint: String::new(), + model_group: Some(model_group.into()), + prompt_tokens: 10, + completion_tokens: 20, + api_requests: 1, + successful_requests: 1, + cache_read_input_tokens: 3, + total_response_time_ms: 100, + timed_requests: 1, + spend: 0.5, + ..DailySpend::default() + } +} + +#[rstest] +#[case::user(DailyTable::User, "LiteLLM_DailyUserSpend", "user_id")] +#[case::team(DailyTable::Team, "LiteLLM_DailyTeamSpend", "team_id")] +#[case::end_user(DailyTable::EndUser, "LiteLLM_DailyEndUserSpend", "end_user_id")] +#[tokio::test] +async fn daily_rows_add_up_and_keep_their_first_model_group( + #[case] table: DailyTable, + #[case] name: &str, + #[case] entity: &str, +) { + let postgres = MigratedPostgres::start().await.unwrap(); + let pool = postgres.pool(); + upsert_daily_spend(pool, table, &[daily("e-1", "first"), daily("e-2", "first")]) + .await + .unwrap(); + upsert_daily_spend(pool, table, &[daily("e-1", "second")]) + .await + .unwrap(); + + let row = read( + pool, + &format!( + r#"SELECT model_group, prompt_tokens, completion_tokens, api_requests, + successful_requests, failed_requests, cache_read_input_tokens, + total_response_time_ms, timed_requests, spend, endpoint, + mcp_namespaced_tool_name + FROM "{name}" WHERE {entity} = 'e-1'"# + ), + ) + .await; + assert_eq!( + row, + json!({ + "model_group": "first", + "prompt_tokens": 20, + "completion_tokens": 40, + "api_requests": 2, + "successful_requests": 2, + "failed_requests": 0, + "cache_read_input_tokens": 6, + "total_response_time_ms": 200, + "timed_requests": 2, + "spend": 1, + "endpoint": "", + "mcp_namespaced_tool_name": "", + }) + ); + assert_eq!( + count(pool, &format!(r#"SELECT count(*) FROM "{name}""#)).await, + 2 + ); +} diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml index 7d5ccf6069d..d7e2e59bb53 100644 --- a/litellm-rust/crates/gateway/Cargo.toml +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -15,6 +15,8 @@ litellm-gateway-mcp.workspace = true tokio-util = "0.7" thiserror.workspace = true futures-util.workspace = true +http.workspace = true +rand.workspace = true litellm-gateway-auth.workspace = true litellm-gateway-hooks.workspace = true litellm-model-catalog.workspace = true @@ -38,6 +40,8 @@ tower-sessions-moka-store = "0.15.0" [dev-dependencies] base64.workspace = true +litellm-db-testing.workspace = true +sqlx.workspace = true rstest.workspace = true tempfile.workspace = true tokio = { workspace = true, features = ["sync", "io-util"] } diff --git a/litellm-rust/crates/gateway/src/keys.rs b/litellm-rust/crates/gateway/src/keys.rs index e4a92953e84..fdf52ff90e0 100644 --- a/litellm-rust/crates/gateway/src/keys.rs +++ b/litellm-rust/crates/gateway/src/keys.rs @@ -39,24 +39,16 @@ impl KeyLookup for PostgresKeys { } } -/// The key store, when a database is configured: `general_settings.database_url` (which may be +/// The proxy's database, when one is configured: `general_settings.database_url` (which may be /// an `os.environ/` reference) or else `DATABASE_URL`, as the Python proxy resolves it. The pool /// connects on first use, so the gateway starts, and the master key works, with the database down. -pub async fn store( +pub async fn database( config: &Config, secrets: &dyn SecretSource, -) -> Result>, Error> { +) -> Result, Error> { let Some(url) = database_url(config, secrets).await? else { return Ok(None); }; - // Inference does not run global guardrails or policies, so a virtual key would get past - // them; refuse to start rather than serve keys the operator expects to be guarded. - if !config.guardrails.is_empty() - || !config.policies.is_empty() - || !config.policy_attachments.is_empty() - { - return Err(Error::GuardedVirtualKeys); - } let general = &config.general_settings; let defaults = PoolSettings::default(); let settings = PoolSettings { @@ -73,11 +65,23 @@ pub async fn store( .and_then(|seconds| Duration::try_from_secs_f64(seconds).ok()) .or(defaults.idle_timeout), }; - let pool = litellm_db::connect_lazy(url.expose(), &settings)?; - Ok(Some(Arc::new(PostgresKeys { + Ok(Some(litellm_db::connect_lazy(url.expose(), &settings)?)) +} + +/// The key store over the proxy's database. +pub fn store(config: &Config, pool: PgPool) -> Result, Error> { + // Inference does not run global guardrails or policies, so a virtual key would get past + // them; refuse to start rather than serve keys the operator expects to be guarded. + if !config.guardrails.is_empty() + || !config.policies.is_empty() + || !config.policy_attachments.is_empty() + { + return Err(Error::GuardedVirtualKeys); + } + Ok(Arc::new(PostgresKeys { pool, personal_keys_use_teams: personal_keys_use_teams(config), - }))) + })) } /// `litellm_settings.personal_key_model_access_from_teams`. diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index 55fd19bb941..7e5e50fcab1 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -8,6 +8,7 @@ mod key_status; mod keys; pub mod metrics; mod secrets; +pub mod spend; pub use error::Error; use std::{sync::Arc, time::Instant}; @@ -135,7 +136,26 @@ pub async fn build_keys( config: &Config, secrets: &dyn litellm_secrets::source::SecretSource, ) -> Result>, Error> { - keys::store(config, secrets).await + build_database(config, secrets) + .await? + .map(|pool| build_keys_on(config, pool)) + .transpose() +} + +/// The proxy's database, if one is configured (see [`keys::database`]). +pub async fn build_database( + config: &Config, + secrets: &dyn litellm_secrets::source::SecretSource, +) -> Result, Error> { + keys::database(config, secrets).await +} + +/// The virtual key store over an already built database pool. +pub fn build_keys_on( + config: &Config, + pool: litellm_db::PgPool, +) -> Result, Error> { + keys::store(config, pool) } pub fn router( diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs index ed4f293b4de..fccfa6cfc66 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -91,8 +91,11 @@ async fn main() -> Result<(), Box> { }, ) }); - let keys = litellm_gateway::build_keys(&config, inference.secrets.as_ref()) - .await? + let database = litellm_gateway::build_database(&config, inference.secrets.as_ref()).await?; + let keys = database + .clone() + .map(|pool| litellm_gateway::build_keys_on(&config, pool)) + .transpose()? .map(|keys| litellm_gateway::key_cache(keys, &config)); let virtual_keys = keys.is_some(); let coordination = @@ -109,7 +112,13 @@ async fn main() -> Result<(), Box> { coordination.start(inference.clone(), keys.clone(), shutdown.clone()); } // Built-in hooks and observers are added here as the gateway gains them. - litellm_gateway::hooks::install(&inference, &config, Vec::new(), Vec::new()).await?; + let mut observers: Vec> = Vec::new(); + let spend = database.zip(litellm_gateway::spend::SpendSettings::from_config(&config)); + let spend_tracking = spend.is_some(); + if let Some((pool, settings)) = spend { + observers.push(litellm_gateway::spend::SpendWriter::start(pool, settings)); + } + litellm_gateway::hooks::install(&inference, &config, Vec::new(), observers).await?; let listener = tokio::net::TcpListener::bind((settings.host.as_str(), settings.port)).await?; let address = listener.local_addr()?; // Without TCP_NODELAY, Nagle holds small SSE chunks until the client's delayed ACK (~40 ms per stream). @@ -120,7 +129,7 @@ async fn main() -> Result<(), Box> { }); let shared_state = coordination.is_some(); - tracing::info!(%address, models = config.model_list.len(), virtual_keys, shared_state, log_level = %level, "gateway listening"); + tracing::info!(%address, models = config.model_list.len(), virtual_keys, shared_state, spend_tracking, log_level = %level, "gateway listening"); let draining = litellm_gateway::health::Draining::default(); let metrics = litellm_gateway::metrics::Metrics::new(&config, draining.flag()); diff --git a/litellm-rust/crates/gateway/src/spend.rs b/litellm-rust/crates/gateway/src/spend.rs new file mode 100644 index 00000000000..fc373307bec --- /dev/null +++ b/litellm-rust/crates/gateway/src/spend.rs @@ -0,0 +1,708 @@ +//! The spend writer: an observer that records every request in the proxy's Postgres tables, as +//! the proxy's own spend writer does, so the proxy's dashboards, spend reports and budget checks +//! see the gateway's traffic. +//! +//! Requests are added up in memory and written in batches: spend logs every couple of seconds, +//! and the key, user, team, team membership, end user and daily activity rows every +//! `proxy_batch_write_at` seconds. Every write adds to what is stored, so replicas never +//! overwrite each other. + +use std::{ + collections::HashMap, + mem, + sync::{Arc, Mutex, Weak}, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use futures_util::future::BoxFuture; +use litellm_config::Config; +use litellm_db::{ + DailySpend, DailyTable, EntitySpend, PgPool, SpendLog, TeamMemberSpend, add_end_user_spend, + add_key_spend, add_team_member_spend, add_team_spend, add_user_spend, insert_spend_logs, + upsert_daily_spend, +}; +use litellm_gateway_hooks::{Facts, Observer, Record, Route, Usage}; +use serde_json::{Map, Value, json}; + +/// How often spend logs are written, as the proxy's spend-log queue monitor polls. +const LOG_INTERVAL: Duration = Duration::from_secs(2); +/// Rows per spend-log statement, the proxy's `SPEND_LOG_WRITE_BATCH_MAX_ROWS`. +const LOG_BATCH: usize = 100; +/// Rows per daily-table statement. +const DAILY_BATCH: usize = 100; +/// Spend logs held while the database is unreachable; older ones are dropped past this. +const MAX_PENDING_LOGS: usize = 10_000; +/// Retries of a write that failed on a deadlock, a lock timeout or a lost connection. +const RETRIES: u32 = 3; + +/// The keys of the proxy's spend-log `metadata`, every one present and null when unset. +const METADATA_KEYS: [&str; 49] = [ + "actor_agent_id", + "additional_usage_values", + "agent_execution_mode", + "applied_guardrails", + "attempted_fallbacks", + "attempted_retries", + "autorouter_baseline_observation", + "autorouter_savings", + "autorouter_savings_estimate", + "azure_spillover", + "batch_failed_requests", + "batch_models", + "batch_successful_requests", + "billing_agent_id", + "cold_storage_object_key", + "compression_savings", + "cost_breakdown", + "error_information", + "eval_information", + "guardrail_information", + "internal_call_origin", + "litellm_call_id", + "litellm_gateway_injected_cache", + "litellm_overhead_time_ms", + "litellm_roi_estimator", + "max_retries", + "mcp_tool_call_metadata", + "model_map_information", + "original_model_group", + "proxy_server_request", + "requester_ip_address", + "router_metadata", + "routing_decision", + "spend_logs_metadata", + "status", + "target_agent_id", + "usage_object", + "used_client_oauth_token", + "user_agent", + "user_api_key", + "user_api_key_alias", + "user_api_key_org_id", + "user_api_key_project_alias", + "user_api_key_project_id", + "user_api_key_team_alias", + "user_api_key_team_id", + "user_api_key_user_id", + "vector_store_request_metadata", + "verified_human_user_id", +]; + +/// The annotation hooks set to add to a request's `spend_logs_metadata`. +pub const SPEND_LOGS_METADATA: &str = "spend_logs_metadata"; + +#[derive(Clone, Debug, PartialEq)] +pub struct SpendSettings { + /// `general_settings.proxy_batch_write_at`, plus the proxy's 0-5 s of start-up jitter. + pub interval: Duration, + /// `general_settings.disable_spend_logs`: everything but the spend-log rows is written. + pub disable_spend_logs: bool, + /// `litellm_settings.disable_end_user_cost_tracking`. + pub disable_end_user_tracking: bool, +} + +impl SpendSettings { + /// The writer's settings, or `None` when `general_settings.disable_spend_updates` turns + /// spend tracking off. + pub fn from_config(config: &Config) -> Option { + let general = &config.general_settings.additional_fields; + if flag(general.get("disable_spend_updates")) { + return None; + } + let seconds = match general.get("proxy_batch_write_at") { + Some(litellm_config::Value::Number(seconds)) => seconds.as_f64(), + Some(litellm_config::Value::String(seconds)) => seconds.trim().parse().ok(), + _ => std::env::var("PROXY_BATCH_WRITE_AT") + .ok() + .and_then(|seconds| seconds.trim().parse().ok()), + } + .filter(|seconds: &f64| seconds.is_finite() && *seconds > 0.0) + .unwrap_or(10.0); + let jitter = f64::from(rand::random::() % 6); + Some(Self { + interval: Duration::from_secs_f64(seconds + jitter), + disable_spend_logs: flag(general.get("disable_spend_logs")), + disable_end_user_tracking: flag( + config + .litellm_settings + .additional_fields + .get("disable_end_user_cost_tracking"), + ), + }) + } +} + +fn flag(value: Option<&litellm_config::Value>) -> bool { + match value { + Some(litellm_config::Value::Bool(flag)) => *flag, + Some(litellm_config::Value::String(text)) => text.eq_ignore_ascii_case("true"), + _ => false, + } +} + +type DailyKey = (DailyTable, String, String, String, String, String, String); + +/// Spend added up since the last write. +#[derive(Default)] +struct Pending { + keys: HashMap, + users: HashMap, + teams: HashMap, + end_users: HashMap, + members: HashMap<(String, String), f64>, + daily: HashMap, +} + +impl Pending { + fn add_daily(&mut self, table: DailyTable, row: DailySpend) { + let key = ( + table, + row.entity_id.clone(), + row.date.clone(), + row.api_key.clone(), + row.model.clone(), + row.custom_llm_provider.clone(), + row.endpoint.clone(), + ); + match self.daily.get_mut(&key) { + Some(total) => { + total.prompt_tokens += row.prompt_tokens; + total.completion_tokens += row.completion_tokens; + total.api_requests += row.api_requests; + total.successful_requests += row.successful_requests; + total.failed_requests += row.failed_requests; + total.cache_read_input_tokens += row.cache_read_input_tokens; + total.cache_creation_input_tokens += row.cache_creation_input_tokens; + total.total_response_time_ms += row.total_response_time_ms; + total.timed_requests += row.timed_requests; + total.spend += row.spend; + } + None => { + self.daily.insert(key, row); + } + } + } +} + +pub struct SpendWriter { + pool: PgPool, + settings: SpendSettings, + pending: Mutex, + logs: Mutex>, + /// Serializes flushes, so the periodic one and the one at shutdown never interleave. + flushing: tokio::sync::Mutex<()>, +} + +impl SpendWriter { + /// Starts the writer and its timers. Must be called inside a tokio runtime. + pub fn start(pool: PgPool, settings: SpendSettings) -> Arc { + let writer = Arc::new(Self { + pool, + settings, + pending: Mutex::default(), + logs: Mutex::default(), + flushing: tokio::sync::Mutex::new(()), + }); + tokio::spawn(run(Arc::downgrade(&writer))); + writer + } + + fn add(&self, record: &Record) { + let request = &record.request; + let key = record.key(); + let cost = record.outcome.cost.unwrap_or(0.0); + let end_user = request + .end_user + .as_deref() + .filter(|_| !self.settings.disable_end_user_tracking); + let hash = key.map(|key| key.hash.as_str()).unwrap_or_default(); + let user_id = key.and_then(|key| key.user_id.as_deref()); + let team_id = key.and_then(|key| key.team_id.as_deref()); + if hash.is_empty() && user_id.is_none() && team_id.is_none() && end_user.is_none() { + return; + } + let log = spend_log(record, cost, end_user); + { + let mut pending = self.pending.lock().unwrap_or_else(|e| e.into_inner()); + if !hash.is_empty() { + *pending.keys.entry(hash.to_owned()).or_default() += cost; + } + if let Some(user_id) = user_id { + *pending.users.entry(user_id.to_owned()).or_default() += cost; + } + if let Some(team_id) = team_id { + *pending.teams.entry(team_id.to_owned()).or_default() += cost; + if let Some(user_id) = user_id { + *pending + .members + .entry((team_id.to_owned(), user_id.to_owned())) + .or_default() += cost; + } + } + if let Some(end_user) = end_user { + *pending.end_users.entry(end_user.to_owned()).or_default() += cost; + } + let row = daily_row(record, &log); + pending.add_daily( + DailyTable::User, + DailySpend { + entity_id: log.user.clone(), + ..row.clone() + }, + ); + pending.add_daily( + DailyTable::Team, + DailySpend { + entity_id: log.team_id.clone(), + ..row.clone() + }, + ); + if let Some(end_user) = end_user { + pending.add_daily( + DailyTable::EndUser, + DailySpend { + entity_id: end_user.to_owned(), + ..row + }, + ); + } + } + if !self.settings.disable_spend_logs { + let mut logs = self.logs.lock().unwrap_or_else(|e| e.into_inner()); + logs.push(log); + if logs.len() > MAX_PENDING_LOGS { + let excess = logs.len() - MAX_PENDING_LOGS; + logs.drain(..excess); + tracing::warn!(dropped = excess, "spend log queue full; dropped the oldest"); + } + } + } + + /// Writes the spend logs held so far. + pub async fn flush_logs(&self) { + let _flushing = self.flushing.lock().await; + self.write_logs().await; + } + + /// Writes everything held so far. + pub async fn flush_all(&self) { + let _flushing = self.flushing.lock().await; + self.write_logs().await; + self.write_spend().await; + } + + async fn write_logs(&self) { + let logs = mem::take(&mut *self.logs.lock().unwrap_or_else(|e| e.into_inner())); + let mut failed = Vec::new(); + for batch in logs.chunks(LOG_BATCH) { + if let Err(error) = retry(|| insert_spend_logs(&self.pool, batch)).await { + tracing::error!(%error, rows = batch.len(), "writing spend logs failed"); + if error.is_retryable() { + failed.extend_from_slice(batch); + } + } + } + if !failed.is_empty() { + let mut logs = self.logs.lock().unwrap_or_else(|e| e.into_inner()); + failed.append(&mut logs); + *logs = failed; + } + } + + async fn write_spend(&self) { + let pending = mem::take(&mut *self.pending.lock().unwrap_or_else(|e| e.into_inner())); + let entity = |spend: HashMap| spend.into_iter().collect::(); + // The proxy's order: user, end user, key, team, team membership, then the daily tables. + // Entity spend that cannot be written is dropped, as the proxy drops it. + let users = entity(pending.users); + log_failure("user", retry(|| add_user_spend(&self.pool, &users)).await); + let end_users = entity(pending.end_users); + log_failure( + "end user", + retry(|| add_end_user_spend(&self.pool, &end_users)).await, + ); + let keys = entity(pending.keys); + log_failure("key", retry(|| add_key_spend(&self.pool, &keys)).await); + let teams = entity(pending.teams); + log_failure("team", retry(|| add_team_spend(&self.pool, &teams)).await); + let members: Vec<_> = pending + .members + .into_iter() + .map(|((team_id, user_id), cost)| TeamMemberSpend { + team_id, + user_id, + cost, + }) + .collect(); + log_failure( + "team member", + retry(|| add_team_member_spend(&self.pool, &members)).await, + ); + + let mut tables: HashMap> = HashMap::new(); + for ((table, ..), row) in pending.daily { + tables.entry(table).or_default().push(row); + } + let mut requeue = Vec::new(); + for table in [DailyTable::User, DailyTable::Team, DailyTable::EndUser] { + let rows = tables.remove(&table).unwrap_or_default(); + for batch in rows.chunks(DAILY_BATCH) { + if let Err(error) = retry(|| upsert_daily_spend(&self.pool, table, batch)).await { + tracing::error!(%error, ?table, rows = batch.len(), "writing daily spend failed"); + if error.is_retryable() { + requeue.extend(batch.iter().cloned().map(|row| (table, row))); + } + } + } + } + if !requeue.is_empty() { + let mut pending = self.pending.lock().unwrap_or_else(|e| e.into_inner()); + for (table, row) in requeue { + pending.add_daily(table, row); + } + } + } +} + +impl Observer for SpendWriter { + fn name(&self) -> &'static str { + "spend_writer" + } + + fn reads(&self) -> Facts { + Facts::KEY + | Facts::END_USER + | Facts::CLIENT_HEADERS + | Facts::DEPLOYMENT + | Facts::USAGE + | Facts::COST + | Facts::TIMINGS + | Facts::ANNOTATIONS + } + + fn observe<'a>(&'a self, record: &'a Record) -> BoxFuture<'a, ()> { + self.add(record); + Box::pin(async {}) + } + + fn flush(&self) -> BoxFuture<'_, ()> { + Box::pin(self.flush_all()) + } +} + +/// Writes spend logs every [`LOG_INTERVAL`] and everything else every batch interval, until the +/// writer is dropped. +async fn run(writer: Weak) { + let interval = match writer.upgrade() { + Some(writer) => writer.settings.interval, + None => return, + }; + let mut since_spend = Duration::ZERO; + loop { + tokio::time::sleep(LOG_INTERVAL).await; + let Some(writer) = writer.upgrade() else { + return; + }; + since_spend += LOG_INTERVAL; + if since_spend >= interval { + since_spend = Duration::ZERO; + writer.flush_all().await; + } else { + writer.flush_logs().await; + } + } +} + +async fn retry(mut write: F) -> Result<(), litellm_db::Error> +where + F: FnMut() -> Fut, + Fut: Future>, +{ + let mut attempt = 0; + loop { + match write().await { + Err(error) if error.is_retryable() && attempt < RETRIES => { + let base = 1u64 << (attempt + 1); + let jitter = u64::from(rand::random::()) * base / u64::from(u16::MAX); + tokio::time::sleep(Duration::from_millis((base * 1000 + jitter * 1000) / 2)).await; + attempt += 1; + } + result => return result, + } + } +} + +fn log_failure(table: &str, result: Result<(), litellm_db::Error>) { + if let Err(error) = result { + tracing::error!(%error, table, "writing spend failed; the spend is dropped"); + } +} + +fn status(record: &Record) -> &'static str { + match (200..300).contains(&record.outcome.status) { + true => "success", + false => "failure", + } +} + +/// The model as the proxy logs it: the deployment's model without its provider route. +fn logged_model(record: &Record) -> String { + match &record.outcome.deployment { + Some(deployment) => match (deployment.provider(), deployment.model.split_once('/')) { + (Some(provider), Some((prefix, rest))) if prefix == provider => rest.to_owned(), + _ => deployment.model.clone(), + }, + None => record.request.requested_model.to_string(), + } +} + +fn spend_log(record: &Record, cost: f64, end_user: Option<&str>) -> SpendLog { + let request = &record.request; + let outcome = &record.outcome; + let key = record.key(); + let deployment = outcome.deployment.as_ref(); + let usage = outcome.usage.unwrap_or_default(); + let duration = record.duration(); + let completion_start = match (request.stream, outcome.time_to_first_byte) { + (true, Some(first_byte)) => request.started_at + first_byte, + _ => outcome.ended_at, + }; + let tokens = |count: u64| i32::try_from(count).unwrap_or(i32::MAX); + SpendLog { + request_id: request.request_id.to_string(), + call_type: request.route.call_type().into(), + api_key: key.map(|key| key.hash.clone()).unwrap_or_default(), + spend: cost, + total_tokens: tokens(usage.total_tokens()), + prompt_tokens: tokens(usage.prompt_tokens()), + completion_tokens: tokens(usage.output_tokens), + start_time: timestamp(request.started_at), + end_time: timestamp(outcome.ended_at), + completion_start_time: Some(timestamp(completion_start)), + request_duration_ms: Some(i32::try_from(duration.as_millis()).unwrap_or(i32::MAX)), + model: logged_model(record), + model_id: deployment.map(|d| d.id.to_string()).unwrap_or_default(), + model_group: deployment + .map(|d| d.model_group.to_string()) + .unwrap_or_else(|| request.requested_model.to_string()), + custom_llm_provider: deployment + .and_then(|d| d.provider()) + .unwrap_or_default() + .to_owned(), + api_base: deployment + .and_then(|d| d.api_base.clone()) + .unwrap_or_default(), + user: key.and_then(|key| key.user_id.clone()).unwrap_or_default(), + metadata: metadata(record), + cache_hit: if request.stream { "False" } else { "None" }.into(), + cache_key: "Cache OFF".into(), + team_id: key.and_then(|key| key.team_id.clone()).unwrap_or_default(), + organization_id: key.and_then(|key| key.org_id.clone()).unwrap_or_default(), + end_user: end_user.unwrap_or_default().to_owned(), + requester_ip_address: None, + session_id: Some(request.request_id.to_string()), + status: status(record).into(), + litellm_call_id: Some(request.request_id.to_string()), + } +} + +fn daily_row(record: &Record, log: &SpendLog) -> DailySpend { + let usage = record.outcome.usage.unwrap_or_default(); + let success = status(record) == "success"; + let count = |count: u64| i64::try_from(count).unwrap_or(i64::MAX); + DailySpend { + entity_id: String::new(), + date: date(record.request.started_at), + api_key: log.api_key.clone(), + model: log.model.clone(), + custom_llm_provider: log.custom_llm_provider.clone(), + // The proxy maps only chat completions to an endpoint here. + endpoint: match record.request.route { + Route::ChatCompletions => "/chat/completions".into(), + Route::Messages => String::new(), + }, + model_group: Some(log.model_group.clone()), + prompt_tokens: count(usage.prompt_tokens()), + completion_tokens: count(usage.output_tokens), + api_requests: 1, + successful_requests: i64::from(success), + failed_requests: i64::from(!success), + cache_read_input_tokens: count(usage.cache_read_tokens), + cache_creation_input_tokens: count(usage.cache_write_tokens), + total_response_time_ms: match success { + true => i64::from(log.request_duration_ms.unwrap_or(0)), + false => 0, + }, + timed_requests: i64::from(success), + spend: log.spend, + } +} + +/// The spend log's `metadata`, with every key the proxy writes. +fn metadata(record: &Record) -> Value { + let mut metadata: Map = METADATA_KEYS + .iter() + .map(|key| ((*key).to_owned(), Value::Null)) + .collect(); + let mut set = |key: &str, value: Value| { + metadata.insert(key.to_owned(), value); + }; + let text = + |value: Option<&String>| value.map_or(Value::Null, |value| Value::from(value.as_str())); + let request = &record.request; + if let Some(key) = record.key() { + set("user_api_key", Value::from(key.hash.as_str())); + set("user_api_key_alias", text(key.alias.as_ref())); + set("user_api_key_team_id", text(key.team_id.as_ref())); + set("user_api_key_user_id", text(key.user_id.as_ref())); + set("user_api_key_team_alias", text(key.team_alias.as_ref())); + set("user_api_key_org_id", text(key.org_id.as_ref())); + } + set("spend_logs_metadata", spend_logs_metadata(record)); + set( + "user_agent", + request + .client_headers + .get(http::header::USER_AGENT) + .and_then(|value| value.to_str().ok()) + .map_or(Value::Null, Value::from), + ); + set("litellm_call_id", Value::from(&*request.request_id)); + set("litellm_roi_estimator", Value::Bool(false)); + if let Some(usage) = record.outcome.usage { + let (usage_object, additional) = usage_values(&usage); + set("usage_object", usage_object); + set("additional_usage_values", additional); + } + if status(record) == "failure" { + set("status", Value::from("failure")); + set( + "error_information", + json!({ + "error_code": record.outcome.status.to_string(), + "error_class": "", + "llm_provider": record + .outcome + .deployment + .as_ref() + .and_then(|deployment| deployment.provider()) + .unwrap_or_default(), + "traceback": "", + "error_message": "", + }), + ); + } + Value::Object(metadata) +} + +/// `spend_logs_metadata` as the proxy merges it: the team's, then the key's, then what hooks +/// attached for the request, later keys winning. +fn spend_logs_metadata(record: &Record) -> Value { + let mut merged = Map::new(); + if let Some(key) = record.key() { + for metadata in [&key.team_metadata, &key.metadata] { + if let Some(Value::Object(fields)) = metadata.get("spend_logs_metadata") { + merged.extend(fields.clone()); + } + } + } + if let Some(Value::Object(fields)) = record.annotations.get(SPEND_LOGS_METADATA) { + merged.extend(fields.clone()); + } + match merged.is_empty() { + true => Value::Null, + false => Value::Object(merged), + } +} + +/// The proxy's usage object for a response, and its `additional_usage_values`: everything but +/// the three totals. +fn usage_values(usage: &Usage) -> (Value, Value) { + let mut additional = Map::new(); + additional.insert( + "prompt_tokens_details".into(), + json!({ + "cached_tokens": usage.cache_read_tokens, + "cache_creation_tokens": usage.cache_write_tokens, + }), + ); + additional.insert( + "cache_read_input_tokens".into(), + Value::from(usage.cache_read_tokens), + ); + additional.insert( + "cache_creation_input_tokens".into(), + Value::from(usage.cache_write_tokens), + ); + if usage.web_search_requests > 0 { + additional.insert( + "server_tool_use".into(), + json!({"web_search_requests": usage.web_search_requests}), + ); + } + let mut object = additional.clone(); + object.insert("prompt_tokens".into(), Value::from(usage.prompt_tokens())); + object.insert("completion_tokens".into(), Value::from(usage.output_tokens)); + object.insert("total_tokens".into(), Value::from(usage.total_tokens())); + (Value::Object(object), Value::Object(additional)) +} + +/// Days since 1970-01-01 as a UTC calendar date. +fn civil(days: i64) -> (i64, u32, u32) { + let days = days + 719_468; + let era = days.div_euclid(146_097); + let day_of_era = days.rem_euclid(146_097); + let year_of_era = + (day_of_era - day_of_era / 1460 + day_of_era / 36_524 - day_of_era / 146_096) / 365; + let day_of_year = day_of_era - (365 * year_of_era + year_of_era / 4 - year_of_era / 100); + let month_index = (5 * day_of_year + 2) / 153; + let day = (day_of_year - (153 * month_index + 2) / 5 + 1) as u32; + let month = if month_index < 10 { + month_index + 3 + } else { + month_index - 9 + } as u32; + let year = year_of_era + era * 400 + i64::from(month <= 2); + (year, month, day) +} + +fn since_epoch(time: SystemTime) -> Duration { + time.duration_since(UNIX_EPOCH).unwrap_or_default() +} + +/// `YYYY-MM-DD`, UTC. +fn date(time: SystemTime) -> String { + let (year, month, day) = civil((since_epoch(time).as_secs() / 86_400) as i64); + format!("{year:04}-{month:02}-{day:02}") +} + +/// `YYYY-MM-DD HH:MM:SS.fff`, UTC, as the proxy stores `timestamp(3)` columns. +fn timestamp(time: SystemTime) -> String { + let elapsed = since_epoch(time); + let seconds = elapsed.as_secs() % 86_400; + format!( + "{} {:02}:{:02}:{:02}.{:03}", + date(time), + seconds / 3600, + seconds / 60 % 60, + seconds % 60, + elapsed.subsec_millis() + ) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::epoch(0, "1970-01-01 00:00:00.000")] + #[case::leap_day(1_709_208_000_123, "2024-02-29 12:00:00.123")] + #[case::year_end(1_798_761_599_999, "2026-12-31 23:59:59.999")] + fn timestamps_are_utc(#[case] millis: u64, #[case] expected: &str) { + let time = UNIX_EPOCH + Duration::from_millis(millis); + assert_eq!(timestamp(time), expected); + assert_eq!(date(time), expected[..10]); + } +} diff --git a/litellm-rust/crates/gateway/tests/spend.rs b/litellm-rust/crates/gateway/tests/spend.rs new file mode 100644 index 00000000000..ce08d645a05 --- /dev/null +++ b/litellm-rust/crates/gateway/tests/spend.rs @@ -0,0 +1,304 @@ +//! The spend writer against a real proxy database: requests through the gateway end up in the +//! proxy's spend logs, entity spend and daily activity tables. + +use axum::{ + Router, + body::{Body, to_bytes}, + http::{Request, header::CONTENT_TYPE}, + response::IntoResponse, +}; +use litellm_auth_types::SecretValue; +use litellm_config::Config; +use litellm_db::PgPool; +use litellm_db_testing::MigratedPostgres; +use litellm_gateway::spend::{SpendSettings, SpendWriter}; +use litellm_gateway_auth::keys::KeyHash; +use rstest::rstest; +use serde_json::{Value, json}; +use tokio::net::TcpListener; +use tower::ServiceExt; + +const MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-x","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":3,"cache_read_input_tokens":2,"output_tokens":4}}"#; +const COMPLETION: &str = r#"{"id":"chatcmpl-1","object":"chat.completion","created":1,"model":"gpt-x","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":5,"completion_tokens":6,"total_tokens":11}}"#; + +async fn upstream() -> String { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let app = Router::new().fallback(|request: Request| async move { + let payload = match request.uri().path().ends_with("/chat/completions") { + true => COMPLETION, + false => MESSAGE, + }; + ([(CONTENT_TYPE, "application/json")], payload).into_response() + }); + tokio::spawn(async move { axum::serve(listener, app).await }); + format!("http://{address}") +} + +fn hash(token: &str) -> String { + KeyHash::from_token(&SecretValue::new(token)) + .as_str() + .to_owned() +} + +#[expect( + clippy::disallowed_methods, + reason = "test fixture rows; the writes under test use the checked macros" +)] +async fn fixtures(pool: &PgPool) { + let sql = format!( + r#" +INSERT INTO "LiteLLM_TeamTable" (team_id, team_alias, members_with_roles, metadata, created_at, updated_at) + VALUES ('team-1', 'research', '[{{"user_id": "user-1", "role": "user"}}]', + '{{"spend_logs_metadata": {{"project": "team", "cost_center": "42"}}}}', now(), now()); +INSERT INTO "LiteLLM_UserTable" (user_id, spend) VALUES ('user-1', 0); +INSERT INTO "LiteLLM_VerificationToken" (token, key_alias, spend, team_id, user_id, metadata) + VALUES ('{}', 'team-key', 0, 'team-1', 'user-1', + '{{"spend_logs_metadata": {{"project": "alpha"}}}}'); +"#, + hash("sk-team") + ); + sqlx::raw_sql(sqlx::AssertSqlSafe(sql)) + .execute(pool) + .await + .unwrap(); +} + +#[expect( + clippy::disallowed_methods, + reason = "test assertions read rows back; the writes under test use the checked macros" +)] +async fn rows(pool: &PgPool, sql: &str) -> Vec { + sqlx::query_scalar::<_, Value>(sqlx::AssertSqlSafe(format!( + "SELECT to_jsonb(r) FROM ({sql}) r" + ))) + .fetch_all(pool) + .await + .unwrap() +} + +async fn send(app: &Router, path: &str, key: &str, body: Value) { + let response = app + .clone() + .oneshot( + Request::post(path) + .header("content-type", "application/json") + .header("authorization", format!("Bearer {key}")) + .header("user-agent", "spend-test/1.0") + .body(Body::from(body.to_string())) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = to_bytes(response.into_body(), 1 << 20).await.unwrap(); + assert_eq!(status, 200, "{body:?}"); +} + +fn close(left: &Value, right: f64) -> bool { + (left.as_f64().unwrap() - right).abs() < 1e-9 +} + +#[rstest] +#[tokio::test] +async fn requests_are_written_as_the_proxy_writes_them() { + let postgres = MigratedPostgres::start().await.unwrap(); + let pool = postgres.pool(); + fixtures(pool).await; + let upstream = upstream().await; + let config = Config::from_yaml(&format!( + "model_list: + - model_name: claude + litellm_params: + model: anthropic/claude-x + api_base: {upstream} + api_key: sk-upstream + input_cost_per_token: 0.001 + output_cost_per_token: 0.002 + cache_read_input_token_cost: 0.0001 + model_info: + id: claude-1 + - model_name: gpt + litellm_params: + model: openai/gpt-x + api_base: {upstream}/v1 + api_key: sk-upstream + input_cost_per_token: 1 + output_cost_per_token: 2 + model_info: + id: gpt-1 +router_settings: + num_retries: 0 +general_settings: + master_key: sk-master + database_url: '{}' +", + postgres.url() + )) + .unwrap(); + let inference = litellm_gateway::build_inference(&config).unwrap(); + let database = litellm_gateway::build_database(&config, inference.secrets.as_ref()) + .await + .unwrap() + .unwrap(); + let keys = litellm_gateway::build_keys_on(&config, database.clone()).unwrap(); + let settings = SpendSettings::from_config(&config).unwrap(); + litellm_gateway::hooks::install( + &inference, + &config, + Vec::new(), + vec![SpendWriter::start(database, settings)], + ) + .await + .unwrap(); + let app = litellm_gateway::router_with_keys(inference.clone(), &config, None, None, Some(keys)); + + let message = json!({ + "model": "claude", + "max_tokens": 8, + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"user_id": "customer-7"}, + }); + send(&app, "/v1/messages", "sk-team", message.clone()).await; + send(&app, "/v1/messages", "sk-team", message).await; + let chat = json!({"model": "gpt", "messages": [{"role": "user", "content": "hi"}]}); + send(&app, "/v1/chat/completions", "sk-master", chat).await; + litellm_gateway::hooks::close(&inference).await; + + let message_cost = 3.0 * 0.001 + 2.0 * 0.0001 + 4.0 * 0.002; + let logs = rows( + pool, + r#"SELECT call_type, api_key, spend, prompt_tokens, completion_tokens, total_tokens, + model, model_id, model_group, custom_llm_provider, "user", team_id, end_user, + cache_hit, cache_key, status, metadata + FROM "LiteLLM_SpendLogs" ORDER BY call_type DESC, "startTime""#, + ) + .await; + assert_eq!(logs.len(), 3, "{logs:#?}"); + let [message_log, _, chat_log] = logs.as_slice() else { + unreachable!() + }; + assert!( + close(&message_log["spend"], message_cost), + "{message_log:#}" + ); + let mut message_log = message_log.clone(); + let metadata = message_log + .as_object_mut() + .unwrap() + .remove("metadata") + .unwrap(); + message_log.as_object_mut().unwrap().remove("spend"); + assert_eq!( + message_log, + json!({ + "call_type": "anthropic_messages", + "api_key": hash("sk-team"), + "prompt_tokens": 5, + "completion_tokens": 4, + "total_tokens": 9, + "model": "claude-x", + "model_id": "claude-1", + "model_group": "claude", + "custom_llm_provider": "anthropic", + "user": "user-1", + "team_id": "team-1", + "end_user": "customer-7", + "cache_hit": "None", + "cache_key": "Cache OFF", + "status": "success", + }) + ); + assert_eq!(metadata.as_object().unwrap().len(), 49, "{metadata:#}"); + assert_eq!(metadata["user_api_key_alias"], "team-key"); + assert_eq!(metadata["user_api_key_team_alias"], "research"); + assert_eq!(metadata["user_agent"], "spend-test/1.0"); + assert_eq!(metadata["status"], Value::Null); + assert_eq!( + metadata["spend_logs_metadata"], + json!({"project": "alpha", "cost_center": "42"}) + ); + assert_eq!(chat_log["call_type"], "acompletion"); + assert_eq!(chat_log["user"], "default_user_id"); + assert_eq!(chat_log["spend"], 17.0); + + let key = rows( + pool, + &format!( + r#"SELECT spend FROM "LiteLLM_VerificationToken" WHERE token = '{}'"#, + hash("sk-team") + ), + ) + .await; + assert!(close(&key[0]["spend"], 2.0 * message_cost), "{key:?}"); + for sql in [ + r#"SELECT spend FROM "LiteLLM_UserTable" WHERE user_id = 'user-1'"#, + r#"SELECT spend FROM "LiteLLM_TeamTable" WHERE team_id = 'team-1'"#, + r#"SELECT spend FROM "LiteLLM_TeamMembership" WHERE team_id = 'team-1'"#, + r#"SELECT spend FROM "LiteLLM_EndUserTable" WHERE user_id = 'customer-7'"#, + ] { + let row = rows(pool, sql).await; + assert!( + close(&row[0]["spend"], 2.0 * message_cost), + "{sql}: {row:?}" + ); + } + + let daily = rows( + pool, + r#"SELECT user_id, api_key, model, model_group, endpoint, api_requests, + successful_requests, prompt_tokens, completion_tokens, cache_read_input_tokens + FROM "LiteLLM_DailyUserSpend" ORDER BY user_id"#, + ) + .await; + assert_eq!( + daily, + [ + json!({ + "user_id": "default_user_id", + "api_key": hash("sk-master"), + "model": "gpt-x", + "model_group": "gpt", + "endpoint": "/chat/completions", + "api_requests": 1, + "successful_requests": 1, + "prompt_tokens": 5, + "completion_tokens": 6, + "cache_read_input_tokens": 0, + }), + json!({ + "user_id": "user-1", + "api_key": hash("sk-team"), + "model": "claude-x", + "model_group": "claude", + "endpoint": "", + "api_requests": 2, + "successful_requests": 2, + "prompt_tokens": 10, + "completion_tokens": 8, + "cache_read_input_tokens": 4, + }), + ] + ); + let teams = rows( + pool, + r#"SELECT team_id, api_requests FROM "LiteLLM_DailyTeamSpend" ORDER BY team_id"#, + ) + .await; + assert_eq!( + teams, + [ + json!({"team_id": "", "api_requests": 1}), + json!({"team_id": "team-1", "api_requests": 2}), + ] + ); + let end_users = rows( + pool, + r#"SELECT end_user_id, api_requests FROM "LiteLLM_DailyEndUserSpend""#, + ) + .await; + assert_eq!( + end_users, + [json!({"end_user_id": "customer-7", "api_requests": 2})] + ); +} diff --git a/litellm-rust/docs/customer-proxy/spend-writer.svg b/litellm-rust/docs/customer-proxy/spend-writer.svg new file mode 100644 index 00000000000..53bb0e03fdf --- /dev/null +++ b/litellm-rust/docs/customer-proxy/spend-writer.svg @@ -0,0 +1,52 @@ + + + + + + Spend writer: gateway requests land in the proxy's own tables + + + + Request record (observer queue) + key, user, team, end user, deployment, + usage, cost, timings, user agent, + spend_logs_metadata (team, key, hooks) + off the request path; never blocks it + + + + Added up in memory + spend per key, user, team, member, + end user; daily rows per entity, day, + key, model, provider, endpoint; + spend log rows queued (cap 10k) + + + + + Timers (as the proxy) + spend logs: every 2 s + everything else: proxy_batch_write_at + (default 10 s) plus 0-5 s jitter + shutdown: drain the queue, then flush + + + + + Proxy Postgres, same statements and order as the Python writer + LiteLLM_SpendLogs: batches of 100, INSERT ... ON CONFLICT DO NOTHING; columns and 49-key metadata as the proxy writes them + user, end user, key (spend, total_spend, last_active), team, team membership: spend += cost, ids sorted, lock_timeout 5 s + update-only for users, keys and teams; end users created unblocked; membership only for users on the team roster + LiteLLM_DailyUserSpend / DailyTeamSpend / DailyEndUserSpend: upsert that adds tokens, requests, spend; model_group kept from first insert + every write adds to what is stored, so gateway replicas and Python proxies can write side by side + deadlock, lock timeout, lost connection: 3 retries with backoff; then daily rows and spend logs are kept for the next flush, + entity spend is dropped (as the proxy drops it) + + + + + Read by what already reads these tables + proxy UI usage pages, /spend and /global/activity endpoints, key and team spend shown in the admin UI, + and the proxy's own budget checks on its next key lookup + +