Report LLM token usage from API side

deploy/compose/nemoguardrails/config-store/nemoguard_cloud/config.yml

@@ -17,5 +17,7 @@ rails:
       - content safety check input $model=content_safety
       - topic safety check input $model=topic_control
   output:
+    streaming:
+      enabled: true
     flows:
       - content safety check output $model=content_safety

src/nvidia_rag/rag_server/main.py

@@ -33,6 +33,7 @@
 
 """
 
+import asyncio
 import json
 import logging
 import math
@@ -41,6 +42,7 @@
 from collections.abc import AsyncGenerator, Generator
 from concurrent.futures import ThreadPoolExecutor
 from traceback import print_exc
+from types import SimpleNamespace
 from typing import Any
 
 import requests
@@ -64,6 +66,7 @@
     Citations,
     ErrorCodeMapping,
     RAGResponse,
+    Usage,
     generate_answer_async,
     prepare_citations,
     prepare_llm_request,
@@ -93,6 +96,7 @@
     get_llm,
     get_prompts,
     get_streaming_filter_think_parser_async,
+    streaming_filter_think,
 )
 from nvidia_rag.utils.observability.otel_metrics import OtelMetrics
 from nvidia_rag.utils.reranker import get_ranking_model
@@ -106,6 +110,43 @@ async def _async_iter(items) -> AsyncGenerator[Any, None]:
         yield item
 
 
+def _extract_token_usage_from_llm_message(message) -> Usage | None:
+    """Extract token_usage from an LLM invoke response (AIMessage).
+
+    Supports usage_metadata (NVIDIA), response_metadata.token_usage / usage, and
+    prompt_tokens_details is ignored; only prompt_tokens, completion_tokens, total_tokens are used.
+    """
+    prompt_tokens = None
+    completion_tokens = None
+    total_tokens = None
+
+    # Prefer usage_metadata (e.g. NVIDIA AIMessage)
+    um = getattr(message, "usage_metadata", None) or {}
+    if um:
+        prompt_tokens = um.get("prompt_tokens") or um.get("input_tokens")
+        completion_tokens = um.get("completion_tokens") or um.get("output_tokens")
+        total_tokens = um.get("total_tokens")
+
+    # Fallback: response_metadata (token_usage or usage)
+    if prompt_tokens is None or completion_tokens is None:
+        meta = getattr(message, "response_metadata", None) or {}
+        tu = meta.get("token_usage") or meta.get("usage") or {}
+        if tu:
+            prompt_tokens = prompt_tokens or tu.get("prompt_tokens") or tu.get("input_tokens")
+            completion_tokens = completion_tokens or tu.get("completion_tokens") or tu.get("output_tokens")
+            total_tokens = total_tokens or tu.get("total_tokens")
+
+    if prompt_tokens is None and completion_tokens is None:
+        return None
+    if total_tokens is None and (prompt_tokens is not None or completion_tokens is not None):
+        total_tokens = (prompt_tokens or 0) + (completion_tokens or 0)
+    return Usage(
+        prompt_tokens=prompt_tokens or 0,
+        completion_tokens=completion_tokens or 0,
+        total_tokens=total_tokens or 0,
+    )
+
+
 logger = logging.getLogger(__name__)
 
 MAX_COLLECTION_NAMES = 5
@@ -1534,9 +1575,22 @@ async def _llm_chain(
             stream_gen = chain.astream(
                 {"question": query_text}, config={"run_name": "llm-stream"}
             )
-            # Eagerly fetch first chunk to trigger any errors before returning response
-            prefetched_stream = await self._eager_prefetch_astream(stream_gen)
-
+            chain_for_usage = prompt_template | llm
+            loop = asyncio.get_event_loop()
+
+            prefetched_stream, usage_message = await asyncio.gather(
+                # Eagerly fetch first chunk to trigger any errors before returning response
+                self._eager_prefetch_astream(stream_gen),
+                loop.run_in_executor(
+                    None,
+                    lambda: chain_for_usage.invoke(
+                        {"question": query_text},
+                        config={"run_name": "llm-invoke-usage"},
+                    ),
+                ),
+            )
+            token_usage = _extract_token_usage_from_llm_message(usage_message)
+
             logger.info("LLM stream initiated successfully (first chunk received)")
             logger.info("-" * 80)
 
@@ -1548,6 +1602,7 @@ async def _llm_chain(
                     collection_name="",
                     enable_citations=enable_citations,
                     otel_metrics_client=metrics,
+                    token_usage=token_usage,
                 ),
                 status_code=ErrorCodeMapping.SUCCESS,
             )
@@ -2914,8 +2969,25 @@ def generate_filter_for_collection(collection_name):
                 response_reflection_counter = ReflectionCounter(
                     self.config.reflection.max_loops
                 )
-                initial_response = await chain.ainvoke(
-                    {"question": query, "context": docs}
+                chain_invoke = prompt | llm
+                response_message = await chain_invoke.ainvoke(
+                    {"question": query, "context": docs},
+                    config={"run_name": "llm-invoke-reflection"},
+                )
+                token_usage = _extract_token_usage_from_llm_message(response_message)
+                raw_content = response_message.content
+                if isinstance(raw_content, str):
+                    content_str = raw_content
+                elif isinstance(raw_content, list) and raw_content:
+                    content_str = (
+                        raw_content[0].get("text", "")
+                        if isinstance(raw_content[0], dict)
+                        else str(raw_content[0])
+                    )
+                else:
+                    content_str = ""
+                initial_response = "".join(
+                    streaming_filter_think(iter([SimpleNamespace(content=content_str)]))
                 )
                 logger.info("Initial LLM response generated, checking groundedness...")
                 try:
@@ -2970,6 +3042,7 @@ def generate_filter_for_collection(collection_name):
                         retrieval_time_ms=retrieval_time_ms,
                         rag_start_time_sec=rag_start_time_sec,
                         otel_metrics_client=metrics,
+                        token_usage=token_usage,
                     ),
                     status_code=ErrorCodeMapping.SUCCESS,
                 )
@@ -2980,9 +3053,22 @@ def generate_filter_for_collection(collection_name):
                     {"question": query, "context": docs},
                     config={"run_name": "llm-stream"},
                 )
-                # Eagerly fetch first chunk to trigger any errors before returning response
-                prefetched_stream = await self._eager_prefetch_astream(stream_gen)
-
+                chain_for_usage = prompt | llm
+                loop = asyncio.get_event_loop()
+
+                prefetched_stream, usage_message = await asyncio.gather(
+                    # Eagerly fetch first chunk to trigger any errors before returning response
+                    self._eager_prefetch_astream(stream_gen),
+                    loop.run_in_executor(
+                        None,
+                        lambda: chain_for_usage.invoke(
+                            {"question": query, "context": docs},
+                            config={"run_name": "llm-invoke-usage-rag"},
+                        ),
+                    ),
+                )
+                token_usage = _extract_token_usage_from_llm_message(usage_message)
+
                 logger.info("LLM stream initiated successfully (first chunk received)")
                 logger.info("-" * 80)
                 logger.info("=" * 80)
@@ -3000,6 +3086,7 @@ def generate_filter_for_collection(collection_name):
                         retrieval_time_ms=retrieval_time_ms,
                         rag_start_time_sec=rag_start_time_sec,
                         otel_metrics_client=metrics,
+                        token_usage=token_usage,
                     ),
                     status_code=ErrorCodeMapping.SUCCESS,
                 )

src/nvidia_rag/rag_server/response_generator.py

@@ -419,6 +419,7 @@ def generate_answer(
     retrieval_time_ms: float | None = None,
     rag_start_time_sec: float | None = None,
     otel_metrics_client: OtelMetrics | None = None,
+    token_usage: Usage | None = None,
 ):
     """Generate and stream the response to the provided prompt.
 
@@ -429,6 +430,7 @@ def generate_answer(
         collection_name: Name of the collection used for retrieval
         enable_citations: Whether to enable citations in the response
         otel_metrics_client: Optional OpenTelemetry metrics client for updating latency histograms
+        token_usage: Optional token usage metrics (prompt_tokens, completion_tokens, total_tokens)
     """
 
     try:
@@ -537,6 +539,8 @@ def generate_answer(
             # Create response first, then attach metrics for clarity
             chain_response = ChainResponse()
             chain_response.metrics = final_metrics
+            if token_usage is not None and token_usage.total_tokens > 0:
+                chain_response.usage = token_usage
 
             # [DONE] indicate end of response from server
             response_choice = ChainResponseChoices(
@@ -585,6 +589,7 @@ async def generate_answer_async(
     retrieval_time_ms: float | None = None,
     rag_start_time_sec: float | None = None,
     otel_metrics_client: OtelMetrics | None = None,
+    token_usage: Usage | None = None,
 ):
     """Generate and stream the response to the provided prompt asynchronously.
 
@@ -595,6 +600,7 @@ async def generate_answer_async(
         collection_name: Name of the collection used for retrieval
         enable_citations: Whether to enable citations in the response
         otel_metrics_client: Optional OpenTelemetry metrics client for updating latency histograms
+        token_usage: Optional token usage (prompt_tokens, completion_tokens, total_tokens)
     """
 
     try:
@@ -703,6 +709,8 @@ async def generate_answer_async(
             # Create response first, then attach metrics for clarity
             chain_response = ChainResponse()
             chain_response.metrics = final_metrics
+            if token_usage is not None and token_usage.total_tokens > 0:
+                chain_response.usage = token_usage
 
             # [DONE] indicate end of response from server
             response_choice = ChainResponseChoices(

src/nvidia_rag/utils/observability/langchain_callback_handler.py

@@ -40,6 +40,17 @@
 
 from .otel_metrics import OtelMetrics
 
+# Hardcoded attribute keys (replacing deprecated SpanAttributes constants)
+GEN_AI_PROMPTS = "gen_ai.prompt"
+GEN_AI_COMPLETIONS = "gen_ai.completion"
+LLM_REQUEST_MODEL = "gen_ai.request.model"
+LLM_RESPONSE_MODEL = "gen_ai.response.model"
+# Missing in opentelemetry.semconv_ai SpanAttributes (use llm.* to match existing semconv)
+LLM_REQUEST_MAX_TOKENS = "llm.request.max_tokens"
+LLM_REQUEST_TEMPERATURE = "llm.request.temperature"
+LLM_REQUEST_TOP_P = "llm.request.top_p"
+LLM_SYSTEM = "llm.system"
+
 
 class Config:
     exception_logger = None
@@ -137,9 +148,9 @@ def _set_request_params(span, kwargs, span_holder: SpanHolder):
     else:
         model = "unknown"
 
-    span.set_attribute(SpanAttributes.LLM_REQUEST_MODEL, model)
+    span.set_attribute(LLM_REQUEST_MODEL, model)
     # response is not available for LLM requests (as opposed to chat)
-    span.set_attribute(SpanAttributes.LLM_RESPONSE_MODEL, model)
+    span.set_attribute(LLM_RESPONSE_MODEL, model)
 
     if "invocation_params" in kwargs:
         params = (
@@ -150,13 +161,11 @@ def _set_request_params(span, kwargs, span_holder: SpanHolder):
 
     _set_span_attribute(
         span,
-        SpanAttributes.LLM_REQUEST_MAX_TOKENS,
+        LLM_REQUEST_MAX_TOKENS,
         params.get("max_tokens") or params.get("max_new_tokens"),
     )
-    _set_span_attribute(
-        span, SpanAttributes.LLM_REQUEST_TEMPERATURE, params.get("temperature")
-    )
-    _set_span_attribute(span, SpanAttributes.LLM_REQUEST_TOP_P, params.get("top_p"))
+    _set_span_attribute(span, LLM_REQUEST_TEMPERATURE, params.get("temperature"))
+    _set_span_attribute(span, LLM_REQUEST_TOP_P, params.get("top_p"))
 
 
 def _set_llm_request(
@@ -171,11 +180,11 @@ def _set_llm_request(
     if should_send_prompts():
         for i, msg in enumerate(prompts):
             span.set_attribute(
-                f"{SpanAttributes.LLM_PROMPTS}.{i}.role",
+                f"{GEN_AI_PROMPTS}.{i}.role",
                 "user",
             )
             span.set_attribute(
-                f"{SpanAttributes.LLM_PROMPTS}.{i}.content",
+                f"{GEN_AI_PROMPTS}.{i}.content",
                 msg,
             )
 
@@ -207,18 +216,18 @@ def _set_chat_request(
         for message in messages:
             for msg in message:
                 span.set_attribute(
-                    f"{SpanAttributes.LLM_PROMPTS}.{i}.role",
+                    f"{GEN_AI_PROMPTS}.{i}.role",
                     _message_type_to_role(msg.type),
                 )
                 # if msg.content is string
                 if isinstance(msg.content, str):
                     span.set_attribute(
-                        f"{SpanAttributes.LLM_PROMPTS}.{i}.content",
+                        f"{GEN_AI_PROMPTS}.{i}.content",
                         msg.content,
                     )
                 else:
                     span.set_attribute(
-                        f"{SpanAttributes.LLM_PROMPTS}.{i}.content",
+                        f"{GEN_AI_PROMPTS}.{i}.content",
                         json.dumps(msg.content, cls=CallbackFilteredJSONEncoder),
                     )
                 i += 1
@@ -252,7 +261,7 @@ def _set_chat_response(span: Span, response: LLMResult) -> None:
                 )
                 total_tokens = input_tokens + output_tokens
 
-            prefix = f"{SpanAttributes.LLM_COMPLETIONS}.{i}"
+            prefix = f"{GEN_AI_COMPLETIONS}.{i}"
             if hasattr(generation, "text") and generation.text != "":
                 span.set_attribute(
                     f"{prefix}.content",
@@ -317,11 +326,11 @@ def _set_chat_response(span: Span, response: LLMResult) -> None:
 
     if input_tokens > 0 or output_tokens > 0 or total_tokens > 0:
         span.set_attribute(
-            SpanAttributes.LLM_USAGE_PROMPT_TOKENS,
+            "gen_ai.usage.input_tokens",
             input_tokens,
         )
         span.set_attribute(
-            SpanAttributes.LLM_USAGE_COMPLETION_TOKENS,
+            "gen_ai.usage.output_tokens",
             output_tokens,
         )
         span.set_attribute(
@@ -462,7 +471,7 @@ def _create_llm_span(
             entity_path=entity_path,
             metadata=metadata,
         )
-        span.set_attribute(SpanAttributes.LLM_SYSTEM, "Langchain")
+        span.set_attribute(LLM_SYSTEM, "Langchain")
         span.set_attribute(SpanAttributes.LLM_REQUEST_TYPE, request_type.value)
 
         return span
@@ -650,10 +659,10 @@ def on_llm_end(
                 "model_name"
             ) or response.llm_output.get("model_id")
             if model_name is not None:
-                span.set_attribute(SpanAttributes.LLM_RESPONSE_MODEL, model_name)
+                span.set_attribute(LLM_RESPONSE_MODEL, model_name)
 
                 if self.spans[run_id].request_model is None:
-                    span.set_attribute(SpanAttributes.LLM_REQUEST_MODEL, model_name)
+                    span.set_attribute(LLM_REQUEST_MODEL, model_name)
 
         token_usage = (response.llm_output or {}).get("token_usage") or (
             response.llm_output or {}
@@ -673,12 +682,8 @@ def on_llm_end(
                 prompt_tokens + completion_tokens
             )
 
-            _set_span_attribute(
-                span, SpanAttributes.LLM_USAGE_PROMPT_TOKENS, prompt_tokens
-            )
-            _set_span_attribute(
-                span, SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, completion_tokens
-            )
+            _set_span_attribute(span, "gen_ai.usage.input_tokens", prompt_tokens)
+            _set_span_attribute(span, "gen_ai.usage.output_tokens", completion_tokens)
             _set_span_attribute(
                 span, SpanAttributes.LLM_USAGE_TOTAL_TOKENS, total_tokens
             )