Expose exact LLM token usage for agentic queries

docs/docs/extraction/workflow-agentic-retrieval.md

@@ -69,6 +69,7 @@ The following options apply only with `--agentic`. For the full flag list, refer
 | `--agentic-local-tensor-parallel-size` | `1` | vLLM `tensor_parallel_size` for the in-process agent LLM. Set to `2` for local `super-49b`. Ignored when `--agentic-invoke-url` is set. |
 | `--agentic-react-max-steps` | `50` | Maximum ReAct loop iterations. |
 | `--agentic-reasoning-effort` | `high` | Forwarded on OpenAI-compatible agent LLM calls. Ignored by the local adapter. |
+| `--include-usage` | off | Print an object with `hits` and provider-reported LLM `usage` instead of the default hits list. |
 
 Embedding credentials use `NVIDIA_API_KEY` or `NGC_API_KEY` when you call a remote embedding endpoint. The CLI also reuses `--embed-invoke-url`, `--top-k`, `--lancedb-uri`, and `--table-name` from standard retrieval.
 
@@ -229,9 +230,41 @@ Every agentic hit carries the one-pass hit fields (`text`, `metadata`, `source`,
 - `rank` — the position in the final ranking.
 - `result_source` — `final_results`, `rrf`, or `selection_agent`, depending on which stage produced the ranked ID.
 
-CLI `retriever query --agentic` prints those hits as JSON objects.
+CLI `retriever query --agentic` keeps its default output as a JSON list of those
+hits. Add `--include-usage` to return a JSON object that contains `hits` and
+provider-reported LLM `usage`:
 
-Service `POST /v1/query` with `agentic=true` uses the same hits envelope as classic retrieval. Successful responses set `query_mode` to `"agentic"`. Classic dense or hybrid `/v1/query` (including `format=evidence`) sets `query_mode` to `"classic"`. For backward compatibility with the previous agentic service contract, service and MCP hits also copy `rank` and `result_source` under `metadata`; the top-level fields are authoritative and carry the same values.
+```bash
+retriever query "find documents about parser behavior" \
+  --agentic \
+  --include-usage
+```
+
+```json
+{
+  "hits": [
+    {
+      "doc_id": "parser-guide",
+      "rank": 1,
+      "result_source": "final_results"
+    }
+  ],
+  "usage": {
+    "input_tokens": 1250,
+    "output_tokens": 184,
+    "total_tokens": 1434
+  }
+}
+```
+
+The `usage` object reports the exact token counts returned by the LLM provider.
+It contains `input_tokens`, `output_tokens`, and `total_tokens`. When available,
+`stages` preserves the provider-reported breakdown for the ReAct and
+final-selection calls. If the provider does not report usage, the response
+sets `usage` to `null`. `--include-usage` applies only to agentic queries. Classic
+`retriever query` output is unchanged.
+
+Service `POST /v1/query` with `agentic=true` uses the same hits envelope as classic retrieval and can include the same optional `usage` object at the response root. Successful responses set `query_mode` to `"agentic"`. Classic dense or hybrid `/v1/query` (including `format=evidence`) sets `query_mode` to `"classic"` and does not add usage metadata. For backward compatibility with the previous agentic service contract, service and MCP hits also copy `rank` and `result_source` under `metadata`; the top-level fields are authoritative and carry the same values.
 
 An agent can name a document that no retrieval hop returned, which leaves nothing to rehydrate. Those hits report null one-pass fields, and `source` falls back to `doc_id`.
 

nemo_retriever/docs/cli/README.md

@@ -231,6 +231,38 @@ the agent named without retrieving it reports null hit fields. It reuses the sam
 `--embed-model-name` options as standard retrieval. Agentic retrieval uses the
 selected table's model automatically when `--embed-model-name` is omitted.
 
+The default `retriever query --agentic` output remains a JSON hits list. Add
+`--include-usage` to print a JSON object with `hits` and exact provider-reported
+LLM usage:
+
+```bash
+retriever query "how does the ingestion pipeline handle tables?" \
+  --agentic \
+  --include-usage
+```
+
+```json
+{
+  "hits": [
+    {
+      "doc_id": "ingestion-guide",
+      "rank": 1,
+      "result_source": "final_results"
+    }
+  ],
+  "usage": {
+    "input_tokens": 1250,
+    "output_tokens": 184,
+    "total_tokens": 1434
+  }
+}
+```
+
+The `usage` object can also include `stages`, which preserves the
+provider-reported breakdown for the ReAct and final-selection calls. The output
+sets `usage` to `null` when the LLM provider does not report it. This flag applies only
+with `--agentic`; classic query behavior and output are unchanged.
+
 **How it works.** Each agentic query runs `Query -> ReActAgentOperator -> (RRF
 fusion) -> SelectionAgentOperator -> ranked results`:
 
@@ -273,6 +305,8 @@ Agentic-only knobs (apply only with `--agentic`):
   calls; omit to use the endpoint/model default (`0.0` = greedy). Local and
   non-NVIDIA OpenAI-compatible endpoints allow up to `2.0`; NVIDIA-hosted
   endpoints allow up to `1.0`.
+- `--include-usage` (default: off) — replace the default hits-list output with
+  an object that contains `hits` and provider-reported LLM `usage`.
 
 <!-- --8<-- [end:quickstart] -->
 

nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/usage.py

@@ -218,6 +218,56 @@ def sum_usage_breakdown(usage_by_stage: Optional[Mapping[str, Any]]) -> Dict[str
     return total
 
 
+def normalize_usage_breakdown(usage_by_stage: Optional[Mapping[str, Any]]) -> Dict[str, Any]:
+    """Return stable token totals plus the exact provider stage breakdown.
+
+    Providers use both OpenAI-style ``prompt_tokens`` / ``completion_tokens``
+    and ``input_tokens`` / ``output_tokens`` names. Values are copied from the
+    provider aggregate; ``total_tokens`` is only filled by exact addition when
+    both component totals are present.
+    """
+    if not usage_by_stage:
+        return {}
+
+    stages = deepcopy(dict(usage_by_stage))
+
+    def _integer(usage: Mapping[str, Any], *keys: str) -> Optional[int]:
+        for key in keys:
+            value = usage.get(key)
+            if isinstance(value, int) and not isinstance(value, bool):
+                return int(value)
+        return None
+
+    input_by_stage: list[int | None] = []
+    output_by_stage: list[int | None] = []
+    total_by_stage: list[int | None] = []
+    for usage in usage_by_stage.values():
+        if not isinstance(usage, Mapping) or not usage:
+            continue
+        stage_input = _integer(usage, "input_tokens", "prompt_tokens")
+        stage_output = _integer(usage, "output_tokens", "completion_tokens")
+        stage_total = _integer(usage, "total_tokens")
+        if stage_total is None and stage_input is not None and stage_output is not None:
+            stage_total = stage_input + stage_output
+        input_by_stage.append(stage_input)
+        output_by_stage.append(stage_output)
+        total_by_stage.append(stage_total)
+
+    def _complete_sum(values: list[int | None]) -> Optional[int]:
+        return (
+            sum(value for value in values if value is not None)
+            if values and all(v is not None for v in values)
+            else None
+        )
+
+    return {
+        "input_tokens": _complete_sum(input_by_stage),
+        "output_tokens": _complete_sum(output_by_stage),
+        "total_tokens": _complete_sum(total_by_stage),
+        "stages": stages,
+    }
+
+
 def coerce_usage_to_dict(usage: Any) -> Optional[Dict[str, Any]]:
     """Best-effort conversion of a provider usage object to a plain dict.
 

nemo_retriever/src/nemo_retriever/cli/query/app.py

@@ -15,6 +15,9 @@
 from nemo_retriever.query.evidence import build_evidence_result
 from nemo_retriever.cli.query import options as opts
 from nemo_retriever.cli.query_workflow import agentic_query_documents as query_agentic_documents
+from nemo_retriever.cli.query_workflow import (
+    agentic_query_documents_with_metadata as query_agentic_documents_with_metadata,
+)
 from nemo_retriever.cli.query_workflow import query_documents_with_metadata as query_local_documents_with_metadata
 from nemo_retriever.query.agentic_options import (
     agentic_llm_client_error,
@@ -187,6 +190,7 @@ def _local_command(
     output_format: opts.OutputFormatOption = "hits",
     max_text_chars: opts.MaxTextCharsOption = None,
     agentic: opts.AgenticOption = False,
+    include_usage: opts.IncludeUsageOption = False,
     agentic_llm_model: opts.AgenticLlmModelOption = None,
     agentic_invoke_url: opts.AgenticInvokeUrlOption = None,
     agentic_reasoning_effort: opts.AgenticReasoningEffortOption = "high",
@@ -197,6 +201,9 @@ def _local_command(
     agentic_llm_client: opts.AgenticLlmClientOption = None,
 ) -> None:
     _validate_output_options(output_format, max_text_chars)
+    if include_usage and not agentic:
+        typer.echo("Error: --include-usage requires --agentic.", err=True)
+        raise typer.Exit(1)
     if reranker_invoke_url is None:
         reranker_invoke_url = os.environ.get("RERANKER_INVOKE_URL") or None
     if rerank is None:
@@ -283,8 +290,12 @@ def _local_command(
                 ),
             )
             with quiet_capture():
-                ranked = query_agentic_documents(request)
-            typer.echo(json.dumps(ranked, indent=2, sort_keys=True, default=str))
+                if include_usage:
+                    result = query_agentic_documents_with_metadata(request)
+                else:
+                    result = query_agentic_documents(request)
+            payload = {"hits": result.hits, "usage": result.usage or None} if include_usage else result
+            typer.echo(json.dumps(payload, indent=2, sort_keys=True, default=str))
             return
 
         def _request() -> QueryRequest:

nemo_retriever/src/nemo_retriever/cli/query/options.py

@@ -165,6 +165,13 @@
         help="Run an LLM-driven agentic (ReAct) retrieval loop instead of the default retrieval pass.",
     ),
 ]
+IncludeUsageOption = Annotated[
+    bool,
+    typer.Option(
+        "--include-usage",
+        help="With --agentic, emit a {hits, usage} JSON envelope containing provider-reported LLM token usage.",
+    ),
+]
 AgenticLlmModelOption = Annotated[
     str | None,
     typer.Option(

nemo_retriever/src/nemo_retriever/cli/query_workflow.py

@@ -7,8 +7,11 @@
 from typing import Any
 
 from nemo_retriever.query.options import QueryRequest
-from nemo_retriever.query.workflow import QueryDocumentsResult
+from nemo_retriever.query.workflow import AgenticQueryDocumentsResult, QueryDocumentsResult
 from nemo_retriever.query.workflow import agentic_query_documents as run_agentic_query_documents
+from nemo_retriever.query.workflow import (
+    agentic_query_documents_with_metadata as run_agentic_query_documents_with_metadata,
+)
 from nemo_retriever.query.workflow import query_documents as run_query_documents
 from nemo_retriever.query.workflow import query_documents_with_metadata as run_query_documents_with_metadata
 from nemo_retriever.common.vdb.records import RetrievalHit
@@ -27,3 +30,8 @@ def query_documents_with_metadata(request: QueryRequest) -> QueryDocumentsResult
 def agentic_query_documents(request: QueryRequest) -> list[dict[str, Any]]:
     """Run the typed root agentic (ReAct) query workflow."""
     return run_agentic_query_documents(request)
+
+
+def agentic_query_documents_with_metadata(request: QueryRequest) -> AgenticQueryDocumentsResult:
+    """Run the typed root agentic query workflow with LLM usage metadata."""
+    return run_agentic_query_documents_with_metadata(request)

nemo_retriever/src/nemo_retriever/operators/graph_ops/react_agent_operator.py

@@ -274,6 +274,12 @@ def _ensure_agent(self) -> Agent:
             )
         return self._agent
 
+    def pop_query_usage(self, query_id: str) -> Dict[str, Any]:
+        """Remove and return provider-reported LLM usage for one query."""
+        if self._agent is None:
+            return {}
+        return self._agent.llm.pop_query_usage(str(query_id))
+
     # ------------------------------------------------------------------
     # AbstractOperator interface
     # ------------------------------------------------------------------

nemo_retriever/src/nemo_retriever/operators/graph_ops/selection_agent_operator.py

@@ -216,6 +216,12 @@ def _ensure_agent(self) -> SelectionAgent:
             )
         return self._sel
 
+    def pop_query_usage(self, query_id: str) -> Dict[str, Any]:
+        """Remove and return provider-reported LLM usage for one query."""
+        if self._sel is None:
+            return {}
+        return self._sel.llm.pop_query_usage(str(query_id))
+
     # ------------------------------------------------------------------
     # AbstractOperator interface
     # ------------------------------------------------------------------