fix(rag): track LLM completion usage and spend for /v1/rag/query

[yassin-berriai]

Relevant issues

Linear ticket

Resolves LIT-4187

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• I have added meaningful tests
• My PR passes all CI/CD checks (e.g., lint, format, unit tests)
• My PR's scope is as isolated as possible; it only solves 1 specific problem
• I have received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment @greptileai to re-request a review after pushing changes)

Screenshots / Proof of Fix

All runs below are against a live proxy (python litellm/proxy/proxy_cli.py --config config.yaml --port 4187) backed by a real Postgres, with a real OpenAI vector store and real gpt-4o-mini calls costing real money. Setup used for both runs:

KEY=$(curl -s -X POST http://127.0.0.1:4187/key/generate \
  -H "Authorization: Bearer sk-123" -H "Content-Type: application/json" \
  -d '{"key_alias": "rag-spend-repro"}' | jq -r .key)
printf 'The secret project codename is AZURE-FALCON-42.\n' > /tmp/doc.txt
curl -s -X POST http://127.0.0.1:4187/v1/rag/ingest \
  -H "Authorization: Bearer $KEY" \
  -F file="@/tmp/doc.txt" \
  -F 'request={"ingest_options": {"vector_store": {"custom_llm_provider": "openai"}}}'
# {"id":"ingest_d29041ab-...","status":"completed","vector_store_id":"vs_6a4dedbeec08819194f088568d8c4ac5",...}

curl -s -D - -X POST http://127.0.0.1:4187/v1/rag/query \
  -H "Authorization: Bearer $KEY" -H "Content-Type: application/json" \
  -d '{"model": "gpt-4o-mini",
       "messages": [{"role": "user", "content": "What is the secret project codename?"}],
       "retrieval_config": {"vector_store_id": "vs_6a4dedbeec08819194f088568d8c4ac5", "custom_llm_provider": "openai", "top_k": 3},
       "max_tokens": 200}'

Before (commit d6cbf6e, branch base)

The query succeeds and the response body carries real usage ("total_tokens": 49), but there is no x-litellm-response-cost header, the SpendLogs row for the request is zeroed, and key spend only reflects the /v1/chat/completions control request made on the same key:

                 request_id                  |  call_type  |       model        |  spend   | prompt_tokens | completion_tokens | total_tokens
---------------------------------------------+-------------+--------------------+----------+---------------+-------------------+--------------
 ingest_d29041ab-418f-4681-a8c7-e20cf14f2e2e | aingest     |                    |        0 |             0 |                 0 |            0
 e0685e2a-0515-444f-b1bd-fa90251b63aa        | aquery      | openai/gpt-4o-mini |        0 |             0 |                 0 |            0
 chatcmpl-DzG303GDCo7RKnElJyezAzF1BLanl      | acompletion | openai/gpt-4o-mini | 2.55e-06 |             9 |                 2 |           11

       key_alias        |  spend
------------------------+----------
 rag-spend-repro-before | 2.55e-06

After (commit 2124e20)

Same curl. The cost header is returned and the request is tracked with its real tokens and cost, attributed to the key:

x-litellm-response-cost: 1.365e-05
x-litellm-key-spend: 1.365e-05

                 request_id                  |  call_type  |       model        |   spend   | prompt_tokens | completion_tokens | total_tokens
---------------------------------------------+-------------+--------------------+-----------+---------------+-------------------+--------------
 chatcmpl-DzGQ0IItLeeoGz1Vmd400qu1Erhbz      | aquery      | openai/gpt-4o-mini | 1.365e-05 |            35 |                14 |           49

       key_alias        |   spend   | max_budget
------------------------+-----------+------------
 rag-spend-repro-after  | 1.365e-05 |          1

Exactly one SpendLogs row is written per RAG query (the vector store search and the underlying completion do not double-bill), and the /v1/chat/completions control on the same proxy still tracks as before

Budget enforcement (commit 2124e20)

A key with "max_budget": 0.00001 makes one successful RAG query (cost 1.365e-05), then the second query is rejected. Before this fix, RAG spend never accrued so this limit was unenforceable:

curl -s -w "HTTP %{http_code}\n" -X POST http://127.0.0.1:4187/v1/rag/query ... # same body as above
# HTTP 429
# {"error":{"message":"Budget has been exceeded! Key=rag-budget-enforce (sk-...Qpgg) Current cost: 1.3649999999999998e-05, Max budget: 1e-05","type":"budget_exceeded","param":null,"code":"429"}}

Streaming (commit b71998b)

On the previous head, the same streamed query against a live proxy returned HTTP 500 (PydanticSerializationError from FastAPI trying to serialize the raw stream wrapper) and wrote no SpendLogs row. After this round, against a live proxy on port 4244 backed by a real Postgres, with the same real OpenAI vector store flow and a real gpt-5.5 completion:

curl -s -N -D - -X POST http://127.0.0.1:4244/v1/rag/query \
  -H "Authorization: Bearer $KEY" -H "Content-Type: application/json" \
  -d '{"model": "gpt-5.5",
       "messages": [{"role": "user", "content": "What is the secret project codename?"}],
       "retrieval_config": {"vector_store_id": "vs_6a5a651693948191ae72fff8a8b8f9f5", "custom_llm_provider": "openai", "top_k": 3},
       "stream": true,
       "max_tokens": 200}'
# HTTP/1.1 200 OK
# content-type: text/event-stream; charset=utf-8
# data: {"id":"chatcmpl-E2gbxkwle7BT21FEvkkU2JGRqmD4P","object":"chat.completion.chunk",...,"choices":[{"index":0,"delta":{"content":"AZ"}}]}
# ...
# data: [DONE]

Once the stream drains, the request bills exactly one aquery row with the completion's real usage, attributed to the key:

                 request_id                  | call_type |     model      |  spend  | prompt_tokens | completion_tokens | total_tokens
---------------------------------------------+-----------+----------------+---------+---------------+-------------------+--------------
 chatcmpl-E2gbxkwle7BT21FEvkkU2JGRqmD4P      | aquery    | openai/gpt-5.5 | 0.00113 |            34 |                32 |           66

    key_alias     |  spend
------------------+---------
 ragstream-verify | 0.00113

OpenAI prices vector store search at 0 per query, so the folded sub-call cost is 0 here; the fold arithmetic for priced providers (e.g. Vertex AI search API) and for rerank is pinned by the unit tests

Independent e2e verification

Reproduced independently on this branch (commit 2124e20) against a live proxy on port 4000 backed by a real Postgres, with a real OpenAI vector store and real gpt-4o-mini calls. A virtual key with max_budget 1.0 ingested a small text doc, then issued one /v1/rag/query; the response carried both x-litellm-response-cost and x-litellm-key-spend (3.87e-05), the aquery SpendLogs row recorded spend 3.87e-05 with 122/34/156 prompt/completion/total tokens as the only billing row, and the key spend incremented to match. A second key with max_budget 0.00001 succeeded once then returned HTTP 429 budget_exceeded on an identical second query after the spend flushed, confirming budget enforcement now applies to RAG traffic

Full terminal walkthrough:

Response headers carrying the cost:

SpendLogs aquery row and matching key spend:

Budget enforcement returning HTTP 429:

Type

🐛 Bug Fix

Changes

The RAG query pipeline (litellm/rag/main.py) forwards its kwargs, including the parent litellm_logging_obj injected by the @client wrapper, into its sub-calls (vector_stores.asearch, then router.acompletion/litellm.acompletion). All three calls therefore share one logging object, and should_run_logging() allows only one async_success event per object. The vector store search finishes first and consumes the slot, so the completion's usage and cost never reach _PROXY_track_cost_callback and SpendLogs records the request with zero tokens and zero spend, bypassing max_budget

The fix wraps the search and completion sub-calls in the is_internal_call context, the same pattern the emulated file-search handler uses for its nested calls, so the parent aquery event bills exactly once with the final completion response. Additionally, query/aquery/ingest/aingest are registered in the CallTypes enum (they were the only @client entry points missing from it), and the /v1/rag/query route now sets the standard response headers so x-litellm-response-cost is returned like other endpoints. schema.d.ts is regenerated for the enum change

Review follow-ups: the pipeline now also folds sub-call costs into that single billing event. Vector store search cost (per-query pricing, e.g. Vertex AI search API; 0 for providers without search pricing) is computed via the existing vector_store_search_cost helper, and the optional rerank sub-call runs under the same internal-call context with its already-computed response cost added in, instead of firing a standalone billing event that carried no proxy key metadata. The aquery SpendLogs row therefore carries completion plus search plus rerank cost, all attributed to the calling key

Streaming is handled by carrying the accumulated sub-call cost into the stream's single final billing event, as suggested in review. On the non-streaming path the fold mutates the response's hidden response_cost as before; on the streaming path there is no response object to fold into (the cost is computed from the assembled chunks after the pipeline returns), so the pipeline stores the search plus rerank cost on the parent logging object as additional_response_cost and Logging._response_cost_calculator adds it when it prices the assembled stream. All sub-calls are therefore suppressed on both paths and their cost lands on the one aquery event, which also restores proxy key attribution for rerank spend on streamed queries (its previous standalone event carried no key metadata)

While verifying this on a live proxy it turned out stream=true against /v1/rag/query failed outright with a 500: the route returned the raw stream wrapper, which FastAPI cannot serialize, so the stream never drained and its billing event never fired at all. The route now returns a proper text/event-stream StreamingResponse via select_data_generator, the same shape the other LLM routes use, which is also what lets the stream drain and bill

Tests: tests/test_litellm/rag/test_main.py drives litellm.aquery through the real @client wrappers (mock transport only) and asserts the single billing event carries the completion's usage, cost, and aquery call type; it fails on the unfixed code because the event then carries the vector store search response. A streaming test asserts the one streamed aquery billing event includes the priced search and rerank cost, and fails without the additional_response_cost carry. tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py gains regression tests for the cost header and for stream=true returning an SSE response instead of a 500

Link to Devin session: https://app.devin.ai/sessions/7481bf2805db40948c206bf32d349ffc

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[greptile-apps]

Greptile Summary

This PR fixes zero-spend tracking for /v1/rag/query by ensuring the parent aquery billing event fires with the real completion usage and cost, rather than being consumed by the vector store search sub-call. It also fixes streaming RAG queries, which previously returned HTTP 500 because the raw CustomStreamWrapper was returned directly to FastAPI.

• Billing fix (rag/main.py): Wraps vector store search, optional rerank, and the completion call in a _suppressed_sub_call_billing context manager that sets is_internal_call=True, preventing each sub-call from firing its own standalone billing event. Sub-call costs are folded into the single parent aquery event — via _hidden_params["response_cost"] mutation on the non-streaming path and via a new additional_response_cost field on the logging object for the streaming path.
• Streaming response fix (endpoints.py): Detects CustomStreamWrapper results and routes them through select_data_generator wrapped in a StreamingResponse, matching the pattern used by other LLM routes, and adds the standard x-litellm-response-cost / x-litellm-key-spend headers to non-streaming RAG responses.
• CallTypes registration (types/utils.py): Adds ingest, aingest, query, and aquery to the enum so deployment hooks and call-type-driven logic are no longer silently no-oped for RAG entry points.

Confidence Score: 5/5

Safe to merge — the changes are well-scoped to the RAG pipeline, backed by end-to-end proof with real billing data, and do not touch the core authentication or routing paths.

The core billing fix follows an established pattern already used by the file-search emulation handler, and the non-streaming/streaming cost-folding paths are logically disjoint (hidden_params vs. model_call_details) with no double-counting possible. All new tests use try/finally teardown for global state and mock at the transport layer only. The single non-streaming edge case where sub_call_cost can be lost (completion cost absent from hidden_params) is a cost-accounting gap in a failure path, not a correctness regression on the happy path.

litellm/rag/main.py — the non-streaming sub_call_cost folding silently drops search/rerank cost when the completion's response_cost is absent from hidden_params (failure-path gap).

Important Files Changed

Filename	Overview
litellm/rag/main.py	Core fix: wraps vector store search, rerank, and completion sub-calls in _suppressed_sub_call_billing so only the parent aquery event fires; folds sub-call costs via hidden_params on non-streaming path and additional_response_cost on streaming path. Minor gap: sub_call_cost is silently dropped on the non-streaming path when the completion's response_cost is absent from hidden_params.
litellm/litellm_core_utils/litellm_logging.py	Adds additional_response_cost accumulation in _response_cost_calculator — reads the key from model_call_details, guards with isinstance and > 0, and adds it on top of the computed completion cost. Correctly scoped to the streaming path since non-streaming cost is folded via hidden_params["response_cost"] instead.
litellm/proxy/rag_endpoints/endpoints.py	Adds x-litellm-response-cost and related headers to the RAG query response, and routes streaming CustomStreamWrapper results through select_data_generator wrapped in StreamingResponse — fixing the HTTP 500 that previously occurred when FastAPI tried to serialize the raw stream wrapper.
litellm/types/utils.py	Registers ingest, aingest, query, and aquery in the CallTypes enum — previously missing entries for these @client-decorated entry points, which caused deployment hooks and call-type-driven logic to silently no-op for RAG calls.
tests/test_litellm/rag/test_main.py	New unit tests covering: single billing event with completion usage, hidden_params cost, priced vector store cost folding, rerank cost folding, and streaming additional_response_cost. All tests use try/finally to restore litellm.callbacks, addressing the previous teardown concern. The polling pattern (50 × 0.1s + 0.5s sleep) remains but is a pre-existing pattern.
tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py	Adds regression tests for the cost header on non-streaming responses and for stream=true returning an SSE response (content-type, chunk structure, [DONE] sentinel). Both tests mock litellm.aquery at the endpoint layer — no real network calls.
ui/litellm-dashboard/src/lib/http/schema.d.ts	Generated TypeScript schema update reflecting the four new CallTypes enum values (ingest, aingest, query, aquery).

_{Reviews (5): Last reviewed commit: "fix(rag): track LLM completion usage and..." | Re-trigger Greptile}

[codecov]

Codecov Report

❌ Patch coverage is 36.36364% with 35 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
litellm/rag/main.py	15.00%	34 Missing ⚠️
litellm/litellm_core_utils/litellm_logging.py	66.66%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[yassin-berriai]

Addressed the test cleanup in f87dcd5: the aquery billing test now restores litellm.callbacks in a finally block.

On the P1 about the rerank step: leaving arerank outside the is_internal_call guard is deliberate. The search and completion sub-calls forward **kwargs, so they share the parent litellm_logging_obj and would consume its single billing slot; that is what the guard prevents. The arerank call does not forward kwargs (litellm/rag/main.py, step 3), so it creates its own logging object and bills itself independently, exactly as it did before this PR. Wrapping it in is_internal_call would not consolidate its cost into the aquery record; it would suppress the rerank's only billing event while the parent aquery event still prices only the completion response, so the rerank spend would be dropped entirely. Consolidating rerank cost into the parent record would require summing costs across logging objects, which is a behavior change beyond this fix. The emulated file-search handler has the same structure: only the nested calls that share the parent logging object are marked internal, while its vector store search bills itself. Rerank spend attribution (the standalone rerank event carries no proxy key metadata) is a pre-existing gap that deserves its own ticket

@greptileai please review the current head f87dcd5

[codspeed-hq]

Merging this PR will not alter performance

✅ 31 untouched benchmarks

---

_{Comparing litellm_rag_query_spend_tracking (b71998b) with litellm_internal_staging (561b679)¹}

Footnotes

1. No successful run was found on litellm_internal_staging (b0a0f11) during the generation of this report, so 561b679 was used instead as the comparison base. There might be some changes unrelated to this pull request in this report. ↩

[veria-ai]

PR overview

This pull request updates the RAG query path for /v1/rag/query to track LLM completion usage and spend, with changes in the RAG request handling and streaming response flow.

The PR has made progress, with two issues already addressed, but one billing gap remains open. An authenticated streaming caller can disconnect before end-of-stream finalization and avoid spend tracking for completed vector search and rerank sub-calls. The remaining impact is limited to usage and cost accounting rather than direct data access or code execution.

Open issues (1)

• Medium: Client disconnects bypass sub-call billing — litellm/rag/main.py:236

Fixed/addressed: 2 · PR risk: 5/10

[yassin-berriai]

@greptileai please review the current head ce96b0a (two test-only commits since your last review: router-branch coverage for the aquery billing test, and pinning the vector store globals so the cost header test is order-independent)

[yassin-berriai]

@greptileai please review the current head 72ffe4b. Since your last review: the vector store search cost is now folded into the aquery billing event (review feedback from the security pass), and the rerank concern from your review is resolved for real; the rerank sub-call now runs under the internal-call context and its cost is folded into the same single aquery event, so spend is no longer split across two rows

[yassin-berriai]

@greptileai please review the current head 10e9af4

[veria-ai]

Medium: Client disconnects bypass sub-call billing

The vector search, and the rerank below, finish before the streaming response is returned, but their own billing events are suppressed. An authenticated caller can request stream=true and disconnect before end-of-stream processing; CustomStreamWrapper.aclose() closes the provider stream without emitting the parent success event that consumes additional_response_cost, so the completed search and rerank calls never enter spend tracking. Keep these sub-calls on separate billing events, or explicitly finalize the accumulated sub-call cost during stream cancellation.

[yassin-berriai]

This is the proxy-wide streaming disconnect behavior rather than a gap this PR introduces. When a streaming client disconnects mid-stream, neither the success nor the failure logging callback fires for the request (documented on _release_max_parallel_requests_on_disconnect in litellm/proxy/utils.py, which exists precisely because of that), so the completion tokens already streamed are not billed either; that holds for /chat/completions and every streamed route today. The sub-call cost rides the same single billing event as the completion, so it bills exactly when the platform bills. Verified live on this branch: a fully drained stream=true query writes the aquery SpendLogs row with completion plus sub-call cost, and a mid-stream disconnect writes no row at all, completion tokens included, which is the platform contract this PR inherits rather than creates

Splitting the sub-calls back onto standalone events would reintroduce the two defects fixed earlier in this review: a direct vector_stores.asearch billing event prices through the web-search call_type branch at 0 (the vector store branch is keyed on avector_store_search, which only the router factory sets), and a standalone rerank event carries no key metadata, so both would under-bill or orphan spend on every streamed request instead of only on early disconnects. Finalizing only the sub-call cost at cancellation would bill a request's search while dropping its typically larger completion tokens, which is inconsistent partial accounting. Disconnect-time finalization of streamed spend is platform infrastructure that should cover all streamed routes in one change; happy to see it tracked separately, but it is out of scope for this fix

[veria-ai]

Thanks for the detail. I can't automatically confirm this is safe to dismiss, so I'm leaving the thread open for a maintainer to make the call.

[yassin-berriai]

Follow-up: the platform-wide disconnect billing gap this finding pointed at is now fixed in #33736, which finalizes partial streamed spend (completion tokens plus the folded sub-call cost) at disconnect time in the shared streaming cleanup, covering every streamed route rather than only this pipeline