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

litellm/litellm_core_utils/litellm_logging.py

@@ -1453,6 +1453,9 @@ def _response_cost_calculator(
             response_cost = litellm.response_cost_calculator(**response_cost_calculator_kwargs)
 
             verbose_logger.debug(f"response_cost: {response_cost}")
+            additional_response_cost: object = self.model_call_details.get("additional_response_cost")
+            if isinstance(additional_response_cost, (int, float)) and additional_response_cost > 0:
+                return (response_cost or 0.0) + additional_response_cost
             return response_cost
         except Exception as e:  # error calculating cost
             debug_info = StandardLoggingModelCostFailureDebugInformation(

litellm/proxy/rag_endpoints/endpoints.py

@@ -11,14 +11,16 @@
 
 import orjson
 from fastapi import APIRouter, Depends, HTTPException, Request, Response, status
-from fastapi.responses import ORJSONResponse
+from fastapi.responses import ORJSONResponse, StreamingResponse
 
 import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
+from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
 from litellm.proxy._types import *
 from litellm.proxy.auth.auth_utils import is_request_body_safe
 from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
+from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
 from litellm.proxy.common_utils.http_parsing_utils import (
     _read_request_body,
     _safe_get_request_headers,
@@ -604,6 +606,7 @@ async def rag_query(
         general_settings,
         llm_router,
         proxy_config,
+        select_data_generator,
         version,
     )
 
@@ -673,6 +676,31 @@ async def rag_query(
             **request_data,
         )
 
+        hidden_params = getattr(response, "_hidden_params", {}) or {}
+        custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
+            user_api_key_dict=user_api_key_dict,
+            call_id=hidden_params.get("litellm_call_id", None) or "",
+            model_id=hidden_params.get("model_id", None) or "",
+            cache_key=hidden_params.get("cache_key", None) or "",
+            api_base=hidden_params.get("api_base", None) or "",
+            version=version,
+            response_cost=hidden_params.get("response_cost", None),
+            request_data=request_data,
+        )
+
+        if isinstance(response, CustomStreamWrapper):
+            return StreamingResponse(
+                select_data_generator(
+                    response=response,
+                    user_api_key_dict=user_api_key_dict,
+                    request_data=request_data,
+                    request=request,
+                ),
+                media_type="text/event-stream",
+                headers=custom_headers,
+            )
+
+        fastapi_response.headers.update(custom_headers)
         return response
 
     except HTTPException:

litellm/rag/main.py

@@ -11,12 +11,14 @@
 
 import asyncio
 import contextvars
+from contextlib import contextmanager
 from functools import partial
 from typing import (
     TYPE_CHECKING,
     Any,
     Coroutine,
     Dict,
+    Iterator,
     List,
     Optional,
     Tuple,
@@ -27,6 +29,9 @@
 import httpx
 
 import litellm
+from litellm._internal_context import is_internal_call
+from litellm.cost_calculator import vector_store_search_cost
+from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion
 from litellm.rag.ingestion.bedrock_ingestion import BedrockRAGIngestion
 from litellm.rag.ingestion.gemini_ingestion import GeminiRAGIngestion
@@ -188,6 +193,25 @@ async def aingest(
         )
 
 
+@contextmanager
+def _suppressed_sub_call_billing() -> Iterator[None]:
+    """
+    Suppress a sub-call's own billing event so the parent aquery event bills it.
+
+    Every suppressed sub-call's cost must be folded into the parent event:
+    into the response's hidden response_cost on the non-streaming path, or via
+    the logging object's additional_response_cost on the streaming path (the
+    streamed cost is computed from assembled chunks after this pipeline
+    returns, so there is no response object to fold into here).
+    """
+    previous = is_internal_call.get()
+    is_internal_call.set(True)
+    try:
+        yield
+    finally:
+        is_internal_call.set(previous)
+
+
 async def _execute_query_pipeline(
     model: str,
     messages: List[Any],
@@ -209,56 +233,87 @@ async def _execute_query_pipeline(
         raise ValueError("No query found in messages for RAG query")
 
     # 2. Search vector store
-    search_response = await litellm.vector_stores.asearch(
-        vector_store_id=retrieval_config["vector_store_id"],
-        query=query_text,
-        max_num_results=retrieval_config.get("top_k", 10),
-        custom_llm_provider=retrieval_config.get("custom_llm_provider", "openai"),
-        **kwargs,
-    )
+    with _suppressed_sub_call_billing():
+        search_response = await litellm.vector_stores.asearch(
+            vector_store_id=retrieval_config["vector_store_id"],
+            query=query_text,
+            max_num_results=retrieval_config.get("top_k", 10),
+            custom_llm_provider=retrieval_config.get("custom_llm_provider", "openai"),
+            **kwargs,
+        )
+
+    search_provider = retrieval_config.get("custom_llm_provider", "openai")
+    try:
+        search_cost = sum(
+            vector_store_search_cost(
+                model=search_provider if "/" in search_provider else None,
+                custom_llm_provider=search_provider,
+                response=search_response,
+            )
+        )
+    except Exception:  # noqa: BLE001 - cost accounting must never break the query path
+        search_cost = 0.0
 
     rerank_response = None
+    rerank_cost = 0.0
     context_chunks = search_response.get("data", [])
 
     # 3. Optional rerank
     if rerank and rerank.get("enabled"):
         documents = RAGQuery.extract_documents_from_search(search_response)
         if documents:
-            rerank_response = await litellm.arerank(
-                model=rerank["model"],
-                query=query_text,
-                documents=documents,
-                top_n=rerank.get("top_n", 5),
-            )
+            with _suppressed_sub_call_billing():
+                rerank_response = await litellm.arerank(
+                    model=rerank["model"],
+                    query=query_text,
+                    documents=documents,
+                    top_n=rerank.get("top_n", 5),
+                )
+            rerank_hidden_params = getattr(rerank_response, "_hidden_params", None)
+            if isinstance(rerank_hidden_params, dict):
+                rerank_response_cost: float | None = rerank_hidden_params.get("response_cost")
+                rerank_cost = rerank_response_cost or 0.0
             context_chunks = RAGQuery.get_top_chunks_from_rerank(search_response, rerank_response)
 
     # 4. Build context message and call completion
     context_message = RAGQuery.build_context_message(context_chunks)
     modified_messages = messages[:-1] + [context_message] + [messages[-1]]
 
     # Use router if available to properly resolve virtual model names
-    if router is not None:
-        response = await router.acompletion(
-            model=model,
-            messages=modified_messages,
-            stream=stream,
-            **kwargs,
-        )
-    else:
-        response = await litellm.acompletion(
-            model=model,
-            messages=modified_messages,
-            stream=stream,
-            **kwargs,
-        )
+    with _suppressed_sub_call_billing():
+        if router is not None:
+            response = await router.acompletion(
+                model=model,
+                messages=modified_messages,
+                stream=stream,
+                **kwargs,
+            )
+        else:
+            response = await litellm.acompletion(
+                model=model,
+                messages=modified_messages,
+                stream=stream,
+                **kwargs,
+            )
 
     # 5. Attach search results to response
+    sub_call_cost = search_cost + rerank_cost
     if not stream and isinstance(response, ModelResponse):
         response = RAGQuery.add_search_results_to_response(
             response=response,
             search_results=search_response,
             rerank_results=rerank_response,
         )
+        if sub_call_cost > 0:
+            hidden_params = getattr(response, "_hidden_params", None)
+            if isinstance(hidden_params, dict):
+                completion_response_cost: float | None = hidden_params.get("response_cost")
+                if completion_response_cost is not None:
+                    hidden_params["response_cost"] = completion_response_cost + sub_call_cost
+    elif sub_call_cost > 0:
+        logging_obj: object = kwargs.get("litellm_logging_obj")
+        if isinstance(logging_obj, LiteLLMLoggingObj):
+            logging_obj.model_call_details["additional_response_cost"] = sub_call_cost
 
     return response  # type: ignore[return-value]
 

litellm/types/utils.py

@@ -402,6 +402,11 @@ class CallTypes(str, Enum):
     vector_store_search = "vector_store_search"
     avector_store_search = "avector_store_search"
 
+    ingest = "ingest"
+    aingest = "aingest"
+    query = "query"
+    aquery = "aquery"
+
     #########################################################
     # Container Call Types
     #########################################################

tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py

@@ -242,3 +242,88 @@ def test_clean_bedrock_ingest_options_not_rejected(self, client_internal_user):
         assert response.status_code != 400, (
             f"Clean Bedrock ingest_options should not be rejected: {response.json()}"
         )
+
+
+def test_rag_query_returns_response_cost_header(client_internal_user):
+    """
+    /v1/rag/query must surface the completion cost via the
+    x-litellm-response-cost response header, like /v1/chat/completions does.
+    """
+    from litellm.types.utils import ModelResponse
+
+    mock_response = ModelResponse(
+        id="chatcmpl-test",
+        choices=[
+            {
+                "index": 0,
+                "message": {"role": "assistant", "content": "The codename is AZURE-FALCON-42."},
+                "finish_reason": "stop",
+            }
+        ],
+        model="gpt-4o-mini",
+        usage={"prompt_tokens": 35, "completion_tokens": 14, "total_tokens": 49},
+    )
+    mock_response._hidden_params["response_cost"] = 3.45e-06
+
+    with patch(
+        "litellm.proxy.rag_endpoints.endpoints.litellm.aquery",
+        new_callable=AsyncMock,
+        return_value=mock_response,
+    ), patch("litellm.vector_store_registry", None), patch(
+        "litellm.proxy.proxy_server.prisma_client", None
+    ):
+        response = client_internal_user.post(
+            "/v1/rag/query",
+            json={
+                "model": "gpt-4o-mini",
+                "messages": [{"role": "user", "content": "What is the codename?"}],
+                "retrieval_config": {
+                    "vector_store_id": "vs_test_123",
+                    "custom_llm_provider": "openai",
+                },
+            },
+        )
+
+    assert response.status_code == 200, response.json()
+    assert response.headers.get("x-litellm-response-cost") == "3.45e-06"
+
+
+def test_rag_query_stream_returns_event_stream(client_internal_user):
+    """
+    A stream=true /v1/rag/query must return an SSE response. Returning the raw
+    stream wrapper makes FastAPI try to serialize it, which raises and turns
+    every streaming RAG query into a 500; the stream then never drains, so its
+    single billing event (which carries the folded sub-call costs) never fires.
+    """
+    import litellm as litellm_module
+
+    async def fake_aquery(**kwargs):
+        return await litellm_module.acompletion(
+            model="gpt-4o-mini",
+            messages=[{"role": "user", "content": "What is the codename?"}],
+            mock_response="The codename is AZURE-FALCON-42.",
+            stream=True,
+            api_key="test-key",
+        )
+
+    with patch(
+        "litellm.proxy.rag_endpoints.endpoints.litellm.aquery",
+        new=AsyncMock(side_effect=fake_aquery),
+    ), patch("litellm.vector_store_registry", None), patch("litellm.proxy.proxy_server.prisma_client", None):
+        response = client_internal_user.post(
+            "/v1/rag/query",
+            json={
+                "model": "gpt-4o-mini",
+                "messages": [{"role": "user", "content": "What is the codename?"}],
+                "retrieval_config": {
+                    "vector_store_id": "vs_test_123",
+                    "custom_llm_provider": "openai",
+                },
+                "stream": True,
+            },
+        )
+
+    assert response.status_code == 200, response.text
+    assert response.headers.get("content-type", "").startswith("text/event-stream")
+    assert '"object":"chat.completion.chunk"' in response.text
+    assert "data: [DONE]" in response.text