From 2d2ca95481994dd00a0ba50669240b9356efe5db Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Mon, 19 Oct 2026 09:58:43 -0700 Subject: [PATCH] feat(gateway): write spend to the proxy's database from a native observer (#47041) Adds a spend writer observer that records every gateway request in the proxy's own Postgres tables with the same statements the Python proxy uses: spend logs, key/user/team/team member/end user spend, and the daily user, team and end-user activity tables. Writes are batched on the proxy's timers and are additive, so replicas and Python proxies can write side by side. The database pool is now shared between virtual keys and the writer. Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5 Co-authored-by: Claude Opus 5.5 (1M context) --- litellm-rust/Cargo.lock | 4 + ...3420df984f9ae9460ceefa9a16e61bae0e3e2.json | 15 + ...92bf33e16fea1dbb0c32cf89a376a874927f4.json | 21 + ...d48cba84f072116aca0fcedf3566f81fa14b6.json | 14 + ...70fa5f0c55f253258f26bcae44fe2be097272.json | 23 + ...5cd10df9db5c4508b859068d800ca8be02b93.json | 40 + ...855a7324d296ae3f15ccee839a29007693f87.json | 30 + ...fd1abc967056d615940b0804915a3ebed9a2b.json | 15 + ...76fbde114f184f7cc6efe70d49d4c78e8f036.json | 30 + ...45ad5e942f4896cac72c0f002f3c544e25094.json | 15 + ...1fd9c45b29330c73c66cdbee80e259dbb391b.json | 15 + ...cb0b2cba889274ab089938b1fe011c7c94df2.json | 30 + litellm-rust/crates/db/src/lib.rs | 6 + litellm-rust/crates/db/src/spend.rs | 451 +++++++++++ litellm-rust/crates/db/tests/spend.rs | 281 +++++++ litellm-rust/crates/gateway/Cargo.toml | 4 + litellm-rust/crates/gateway/src/keys.rs | 32 +- litellm-rust/crates/gateway/src/lib.rs | 22 +- litellm-rust/crates/gateway/src/main.rs | 17 +- litellm-rust/crates/gateway/src/spend.rs | 708 ++++++++++++++++++ litellm-rust/crates/gateway/tests/spend.rs | 304 ++++++++ .../docs/customer-proxy/spend-writer.svg | 52 ++ 22 files changed, 2110 insertions(+), 19 deletions(-) create mode 100644 litellm-rust/crates/db/.sqlx/query-07328686221e3bfe19c205d68fe3420df984f9ae9460ceefa9a16e61bae0e3e2.json create mode 100644 litellm-rust/crates/db/.sqlx/query-5c74f47a2b84574133c520d235192bf33e16fea1dbb0c32cf89a376a874927f4.json create mode 100644 litellm-rust/crates/db/.sqlx/query-8337c5f41713afadeff72900c4cd48cba84f072116aca0fcedf3566f81fa14b6.json create mode 100644 litellm-rust/crates/db/.sqlx/query-83be2dd1f5892e3581ecee3d2d270fa5f0c55f253258f26bcae44fe2be097272.json create mode 100644 litellm-rust/crates/db/.sqlx/query-8abe8b069d1b8bb22f4ff49046e5cd10df9db5c4508b859068d800ca8be02b93.json create mode 100644 litellm-rust/crates/db/.sqlx/query-94187798c95011cae91b61417c8855a7324d296ae3f15ccee839a29007693f87.json create mode 100644 litellm-rust/crates/db/.sqlx/query-94e8e0b4897318bd005c9d4cd50fd1abc967056d615940b0804915a3ebed9a2b.json create mode 100644 litellm-rust/crates/db/.sqlx/query-9e6d4c7704bfbb38ce6d12b647376fbde114f184f7cc6efe70d49d4c78e8f036.json create mode 100644 litellm-rust/crates/db/.sqlx/query-b0e3f762f1dd5efdd4672953f7745ad5e942f4896cac72c0f002f3c544e25094.json create mode 100644 litellm-rust/crates/db/.sqlx/query-bb412ba2896341351e1f6ca96761fd9c45b29330c73c66cdbee80e259dbb391b.json create mode 100644 litellm-rust/crates/db/.sqlx/query-f56e12df3f7b46178c50a58013ecb0b2cba889274ab089938b1fe011c7c94df2.json create mode 100644 litellm-rust/crates/db/src/spend.rs create mode 100644 litellm-rust/crates/db/tests/spend.rs create mode 100644 litellm-rust/crates/gateway/src/spend.rs create mode 100644 litellm-rust/crates/gateway/tests/spend.rs create mode 100644 litellm-rust/docs/customer-proxy/spend-writer.svg 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 + +