Expose exact LLM token usage for agentic queries

[mahikaw]

Description

• collect query-scoped provider token usage from the ReAct and selection stages
• expose normalized input, output, and total usage through Python and service APIs
• add retriever query --agentic --include-usage
• preserve existing agentic CLI and Python return formats by default

Why

Consumers need exact inner-agent token accounting. Parsing logs or estimating from text does not provide reliable provider-reported usage.

Implementation

• adds normalized usage totals with the original per-stage breakdown
• exposes operator-owned usage through atomic pop_query_usage() calls
• maps internal positional graph query IDs back to caller query IDs
• collects usage in finally blocks so backend state is cleared safely
• adds typed agentic workflow and service response models
• returns usage: null when the provider does not report usage

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[greptile-apps]

Greptile Summary

This PR collects provider-reported token usage from the ReAct and final-selection stages and exposes normalized totals through optional Python, CLI, REST, and MCP response metadata.

• Adds query-scoped usage collection and atomic cleanup across agent operators.
• Introduces typed agentic result and service response models.
• Adds retriever query --agentic --include-usage while preserving default output formats.
• Extends tests and documentation for usage-present and usage-absent behavior.

Confidence Score: 4/5

The PR appears safe to merge, with only non-blocking public API documentation gaps in the new usage interfaces.

Usage ownership, query attribution, normalization, and additive service compatibility remain coherent, but the newly exposed result types and methods do not fully document their fields, return contracts, and errors.

Files Needing Attention: nemo_retriever/src/nemo_retriever/query/workflow.py, nemo_retriever/src/nemo_retriever/query/agentic.py, and the two agent operator modules

Important Files Changed

Filename	Overview
nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/usage.py	Adds normalization across OpenAI-style and input/output token aliases while preserving exact stage breakdowns.
nemo_retriever/src/nemo_retriever/query/agentic.py	Adds per-query usage collection and positional-to-caller query ID mapping; the new public interfaces need complete contract documentation.
nemo_retriever/src/nemo_retriever/query/workflow.py	Adds the typed hits-plus-usage workflow result and cleanup path, but leaves the new public contract under-documented.
nemo_retriever/src/nemo_retriever/service/query_schema.py	Adds an additive agentic response subtype with optional normalized usage.
nemo_retriever/src/nemo_retriever/service/vectordb_app.py	Routes agentic requests through the new response model without changing classic response serialization.
nemo_retriever/src/nemo_retriever/cli/query/app.py	Adds the agentic-only include-usage flag while retaining the legacy hits-list output by default.

Sequence Diagram

sequenceDiagram
    participant Caller
    participant Workflow
    participant ReAct
    participant Selection
    Caller->>Workflow: Agentic query
    Workflow->>ReAct: Run query with positional ID
    ReAct-->>Workflow: Ranked candidates + stage usage
    Workflow->>Selection: Select final documents
    Selection-->>Workflow: Final ranking + stage usage
    Workflow->>Workflow: Pop, merge, and normalize usage
    Workflow-->>Caller: Hits + optional usage

Loading

Prompt To Fix All With AI

### Issue 1
nemo_retriever/src/nemo_retriever/query/workflow.py:30-33
**Incomplete public usage contracts**

The new public result classes and usage-returning methods omit field, parameter, return, and error documentation, leaving API consumers without a defined usage shape or ownership contract. The same incomplete pattern also appears on `AgenticRetrieveResult`, `retrieve_with_usage()`, and both operator `pop_query_usage()` methods.

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

_{Reviews (1): Last reviewed commit: "Expose exact token usage for agentic que..." | Re-trigger Greptile}

[greptile-apps]

Incomplete public usage contracts

The new public result classes and usage-returning methods omit field, parameter, return, and error documentation, leaving API consumers without a defined usage shape or ownership contract. The same incomplete pattern also appears on AgenticRetrieveResult, retrieve_with_usage(), and both operator pop_query_usage() methods.

Rule Used: Every public class and function in nemo_retriever ... (source)

Prompt To Fix With AI

This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/query/workflow.py
Line: 30-33

Comment:
**Incomplete public usage contracts**

The new public result classes and usage-returning methods omit field, parameter, return, and error documentation, leaving API consumers without a defined usage shape or ownership contract. The same incomplete pattern also appears on `AgenticRetrieveResult`, `retrieve_with_usage()`, and both operator `pop_query_usage()` methods.

**Rule Used:** Every public class and function in nemo_retriever ... ([source](https://github.com/nvidia/nemo-retriever/blob/a48f2b12cccbcdec8ad6465b03821d6f4cc9fc80/nemo_retriever/.greptile/config.json))

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!