diff --git a/docs/docs/extraction/workflow-agentic-retrieval.md b/docs/docs/extraction/workflow-agentic-retrieval.md index fe7057c6e..12085c98a 100644 --- a/docs/docs/extraction/workflow-agentic-retrieval.md +++ b/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`. diff --git a/nemo_retriever/docs/cli/README.md b/nemo_retriever/docs/cli/README.md index ed4ecb00b..06261624a 100644 --- a/nemo_retriever/docs/cli/README.md +++ b/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`. diff --git a/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/usage.py b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/usage.py index da3a287bc..ceea064a4 100644 --- a/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/usage.py +++ b/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. diff --git a/nemo_retriever/src/nemo_retriever/cli/query/app.py b/nemo_retriever/src/nemo_retriever/cli/query/app.py index 6eef67091..bb814debf 100644 --- a/nemo_retriever/src/nemo_retriever/cli/query/app.py +++ b/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: diff --git a/nemo_retriever/src/nemo_retriever/cli/query/options.py b/nemo_retriever/src/nemo_retriever/cli/query/options.py index 215a0bca5..9f99312ea 100644 --- a/nemo_retriever/src/nemo_retriever/cli/query/options.py +++ b/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( diff --git a/nemo_retriever/src/nemo_retriever/cli/query_workflow.py b/nemo_retriever/src/nemo_retriever/cli/query_workflow.py index ce98e9abb..a2b284289 100644 --- a/nemo_retriever/src/nemo_retriever/cli/query_workflow.py +++ b/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) diff --git a/nemo_retriever/src/nemo_retriever/operators/graph_ops/react_agent_operator.py b/nemo_retriever/src/nemo_retriever/operators/graph_ops/react_agent_operator.py index 4c50c8fc9..6325e906e 100644 --- a/nemo_retriever/src/nemo_retriever/operators/graph_ops/react_agent_operator.py +++ b/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 # ------------------------------------------------------------------ diff --git a/nemo_retriever/src/nemo_retriever/operators/graph_ops/selection_agent_operator.py b/nemo_retriever/src/nemo_retriever/operators/graph_ops/selection_agent_operator.py index 24efb1259..fe67cd382 100644 --- a/nemo_retriever/src/nemo_retriever/operators/graph_ops/selection_agent_operator.py +++ b/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 # ------------------------------------------------------------------ diff --git a/nemo_retriever/src/nemo_retriever/query/agentic.py b/nemo_retriever/src/nemo_retriever/query/agentic.py index e6b66944a..7ab0972bf 100644 --- a/nemo_retriever/src/nemo_retriever/query/agentic.py +++ b/nemo_retriever/src/nemo_retriever/query/agentic.py @@ -19,6 +19,7 @@ import pandas as pd +from nemo_retriever._agentic.nemo_agent.llm.usage import deep_merge_usage_breakdown from nemo_retriever.common.params import build_embed_option_kwargs from nemo_retriever.operators.abstract_operator import AbstractOperator from nemo_retriever.models import VL_RERANK_MODEL @@ -73,6 +74,14 @@ _RETRIEVED_ONLY_RESULT_SOURCES = frozenset({"rrf", "selection_agent"}) +@dataclass(frozen=True) +class AgenticRetrieveResult: + """Ranked agentic documents and provider-reported usage by query ID.""" + + documents: pd.DataFrame + usage: dict[str, dict[str, Any]] + + class AgenticQueryInputOperator(AbstractOperator): """Adapt ``Retriever(graph=...)`` input DataFrames to agentic query schema.""" @@ -447,9 +456,18 @@ def retrieve(self, query_ids: Sequence[str], query_texts: Sequence[str]) -> pd.D when no hop returned it. Use :func:`rehydrated_agentic_hit` to layer the agentic annotations onto that hit. """ + return self.retrieve_with_usage(query_ids, query_texts).documents + + def retrieve_with_usage( + self, + query_ids: Sequence[str], + query_texts: Sequence[str], + ) -> AgenticRetrieveResult: + """Return ranked documents and exact per-query agent LLM usage.""" if len(query_ids) != len(query_texts): raise ValueError("query_ids and query_texts must have the same length.") + caller_query_ids = [str(query_id) for query_id in query_ids] with self._hit_cache_lock: self._hit_cache.clear() @@ -465,39 +483,41 @@ def retrieve(self, query_ids: Sequence[str], query_texts: Sequence[str]) -> pd.D per_hop_top_k = max(AGENTIC_RETRIEVER_TOP_K, target_top_k) chat_completion_fn = self._get_chat_completion_fn() + react_operator = ReActAgentOperator( + invoke_url=_none_if_empty(self._cfg.invoke_url), + llm_model=str(self._cfg.llm_model), + retriever_fn=self._retrieve_for_agent, + retriever_fn_accepts_query_id=True, + retriever_top_k=per_hop_top_k, + target_top_k=target_top_k, + max_steps=int(self._cfg.react_max_steps), + api_key=_none_if_empty(self._cfg.api_key), + parallel_tool_calls=AGENTIC_PARALLEL_TOOL_CALLS, + num_concurrent=int(self._cfg.num_concurrent), + reasoning_effort=self._cfg.reasoning_effort, + temperature=self._cfg.temperature, + backend=self._cfg.llm_client, + max_tokens=self._cfg.max_tokens, + chat_completion_fn=chat_completion_fn, + ) + selection_operator = SelectionAgentOperator( + invoke_url=_none_if_empty(self._cfg.invoke_url), + llm_model=str(self._cfg.llm_model), + top_k=target_top_k, + api_key=_none_if_empty(self._cfg.api_key), + parallel_tool_calls=AGENTIC_PARALLEL_TOOL_CALLS, + text_truncation=int(self._cfg.text_truncation), + reasoning_effort=self._cfg.reasoning_effort, + temperature=self._cfg.temperature, + backend=self._cfg.llm_client, + max_tokens=self._cfg.max_tokens, + chat_completion_fn=chat_completion_fn, + ) pipeline = ( AgenticQueryInputOperator() - >> ReActAgentOperator( - invoke_url=_none_if_empty(self._cfg.invoke_url), - llm_model=str(self._cfg.llm_model), - retriever_fn=self._retrieve_for_agent, - retriever_fn_accepts_query_id=True, - retriever_top_k=per_hop_top_k, - target_top_k=target_top_k, - max_steps=int(self._cfg.react_max_steps), - api_key=_none_if_empty(self._cfg.api_key), - parallel_tool_calls=AGENTIC_PARALLEL_TOOL_CALLS, - num_concurrent=int(self._cfg.num_concurrent), - reasoning_effort=self._cfg.reasoning_effort, - temperature=self._cfg.temperature, - backend=self._cfg.llm_client, - max_tokens=self._cfg.max_tokens, - chat_completion_fn=chat_completion_fn, - ) + >> react_operator >> RRFAggregatorOperator(k=AGENTIC_RRF_K) - >> SelectionAgentOperator( - invoke_url=_none_if_empty(self._cfg.invoke_url), - llm_model=str(self._cfg.llm_model), - top_k=target_top_k, - api_key=_none_if_empty(self._cfg.api_key), - parallel_tool_calls=AGENTIC_PARALLEL_TOOL_CALLS, - text_truncation=int(self._cfg.text_truncation), - reasoning_effort=self._cfg.reasoning_effort, - temperature=self._cfg.temperature, - backend=self._cfg.llm_client, - max_tokens=self._cfg.max_tokens, - chat_completion_fn=chat_completion_fn, - ) + >> selection_operator >> AgenticSelectionOutputOperator() ) graph_retriever = Retriever( @@ -505,14 +525,30 @@ def retrieve(self, query_ids: Sequence[str], query_texts: Sequence[str]) -> pd.D top_k=target_top_k, embed_kwargs={"text_column": "query_text"}, ) - raw_hits = graph_retriever.queries( - [str(query_text) for query_text in query_texts], - top_k=target_top_k, - ) - result = _raw_hits_to_agentic_result([str(query_id) for query_id in query_ids], raw_hits) + usage_by_query: dict[str, dict[str, Any]] = {} + try: + raw_hits = graph_retriever.queries( + [str(query_text) for query_text in query_texts], + top_k=target_top_k, + ) + finally: + # Retriever.queries assigns positional graph query IDs ("0", "1", + # ...). Pop both operator-owned backends after all query workers + # finish, then restore the caller's IDs. + for position, caller_query_id in enumerate(caller_query_ids): + breakdown: dict[str, Any] = {} + deep_merge_usage_breakdown(breakdown, react_operator.pop_query_usage(str(position))) + deep_merge_usage_breakdown(breakdown, selection_operator.pop_query_usage(str(position))) + if breakdown: + usage_by_query[caller_query_id] = breakdown + + result = _raw_hits_to_agentic_result(caller_query_ids, raw_hits) with self._hit_cache_lock: hit_cache = dict(self._hit_cache) - return _rehydrate_selected_hits(result, hit_cache) + return AgenticRetrieveResult( + documents=_rehydrate_selected_hits(result, hit_cache), + usage=usage_by_query, + ) def _retrieve_for_agent(self, query_text: str, top_k: int, *, query_id: str = "") -> list[dict[str, Any]]: """Retriever callback used by ``ReActAgentOperator``. diff --git a/nemo_retriever/src/nemo_retriever/query/workflow.py b/nemo_retriever/src/nemo_retriever/query/workflow.py index 46ea159a4..4d02e506f 100644 --- a/nemo_retriever/src/nemo_retriever/query/workflow.py +++ b/nemo_retriever/src/nemo_retriever/query/workflow.py @@ -27,6 +27,12 @@ class QueryDocumentsResult: strategies: list[str] +@dataclass(frozen=True) +class AgenticQueryDocumentsResult: + hits: list[dict[str, Any]] + usage: dict[str, Any] + + def _strategies_for_retrieval_mode(mode: LanceRetrievalMode | None) -> list[str]: if mode == "hybrid": return ["semantic", "lexical"] @@ -228,30 +234,50 @@ def agentic_query_documents(request: QueryRequest) -> list[dict[str, Any]]: ``Retriever`` that backs the agent's ``retrieve`` tool. Reranking therefore applies per agent retrieval hop. """ - from nemo_retriever.query.agentic import rehydrated_agentic_hit - retriever = build_agentic_retriever(request) try: result = retriever.retrieve(["0"], [str(request.query)]) - if "rank" in result.columns: - result = result.sort_values("rank") - ranked: list[dict[str, Any]] = [] - for _, row in result.iterrows(): - doc_id = str(row.get("doc_id", "")).strip() - if not doc_id: - continue - ranked.append( - rehydrated_agentic_hit( - row.get("hit"), - doc_id=doc_id, - rank=int(row.get("rank", len(ranked) + 1)), - result_source=str(row.get("result_source", "")), - ) - ) - if len(ranked) >= request.retrieval.top_k: - break - return ranked + return _agentic_rows_to_hits(result, top_k=request.retrieval.top_k) finally: # One-shot CLI/Python entry: tear down cached local vLLM EngineCore so the # process can exit instead of hanging on a live child worker. retriever.unload() + + +def agentic_query_documents_with_metadata(request: QueryRequest) -> AgenticQueryDocumentsResult: + """Run one agentic query and return ranked hits plus exact LLM usage.""" + from nemo_retriever._agentic.nemo_agent.llm.usage import normalize_usage_breakdown + + retriever = build_agentic_retriever(request) + try: + result = retriever.retrieve_with_usage(["0"], [str(request.query)]) + return AgenticQueryDocumentsResult( + hits=_agentic_rows_to_hits(result.documents, top_k=request.retrieval.top_k), + usage=normalize_usage_breakdown(result.usage.get("0")), + ) + finally: + retriever.unload() + + +def _agentic_rows_to_hits(result: Any, *, top_k: int) -> list[dict[str, Any]]: + """Convert an agentic result DataFrame to ranked, rehydrated hits.""" + from nemo_retriever.query.agentic import rehydrated_agentic_hit + + if "rank" in result.columns: + result = result.sort_values("rank") + ranked: list[dict[str, Any]] = [] + for _, row in result.iterrows(): + doc_id = str(row.get("doc_id", "")).strip() + if not doc_id: + continue + ranked.append( + rehydrated_agentic_hit( + row.get("hit"), + doc_id=doc_id, + rank=int(row.get("rank", len(ranked) + 1)), + result_source=str(row.get("result_source", "")), + ) + ) + if len(ranked) >= top_k: + break + return ranked diff --git a/nemo_retriever/src/nemo_retriever/service/agentic_query.py b/nemo_retriever/src/nemo_retriever/service/agentic_query.py index 316cfed67..5e7bb218c 100644 --- a/nemo_retriever/src/nemo_retriever/service/agentic_query.py +++ b/nemo_retriever/src/nemo_retriever/service/agentic_query.py @@ -15,9 +15,9 @@ QueryRetrievalOptions, QueryStorageOptions, ) -from nemo_retriever.query.workflow import agentic_query_documents +from nemo_retriever.query.workflow import agentic_query_documents_with_metadata from nemo_retriever.service.config import AgenticConfig -from nemo_retriever.service.query_schema import QueryResponse, QueryResult +from nemo_retriever.service.query_schema import AgenticQueryResponse, QueryResult #: Annotations the agentic workflow layers on top of the classic hit fields. @@ -140,8 +140,8 @@ def run_agentic_query( embed_model: str, embed_model_provider_prefix: str | None, embed_api_key: str, -) -> QueryResponse: - """Execute one agentic retrieval query and return ``QueryResponse`` hits.""" +) -> AgenticQueryResponse: + """Execute one agentic retrieval query and return hits plus LLM usage.""" query_request = build_agentic_query_request( query=query, top_k=top_k, @@ -153,8 +153,9 @@ def run_agentic_query( embed_model_provider_prefix=embed_model_provider_prefix, embed_api_key=embed_api_key, ) - ranked = agentic_query_documents(query_request) - return QueryResponse( - results=[QueryResult(hits=agentic_ranked_to_hits(ranked))], + result = agentic_query_documents_with_metadata(query_request) + return AgenticQueryResponse( + results=[QueryResult(hits=agentic_ranked_to_hits(result.hits))], query_mode="agentic", + usage=result.usage or None, ) diff --git a/nemo_retriever/src/nemo_retriever/service/query_schema.py b/nemo_retriever/src/nemo_retriever/service/query_schema.py index 40011e335..e3f71a93b 100644 --- a/nemo_retriever/src/nemo_retriever/service/query_schema.py +++ b/nemo_retriever/src/nemo_retriever/service/query_schema.py @@ -105,6 +105,15 @@ class QueryResult(BaseModel): hits: list[dict[str, Any]] +class AgenticTokenUsage(BaseModel): + """Provider-reported LLM usage for one agentic retrieval query.""" + + input_tokens: int | None = Field(default=None, ge=0) + output_tokens: int | None = Field(default=None, ge=0) + total_tokens: int | None = Field(default=None, ge=0) + stages: dict[str, dict[str, Any]] = Field(default_factory=dict) + + class QueryResponse(BaseModel): results: list[QueryResult] query_mode: QueryMode = Field( @@ -121,6 +130,16 @@ def hits_by_query(self, *, expected_results: int | None = None) -> list[list[dic return [result.hits for result in self.results] +class AgenticQueryResponse(QueryResponse): + """Agentic query response with exact provider-reported LLM usage.""" + + query_mode: Literal["agentic"] = "agentic" + usage: AgenticTokenUsage | None = Field( + default=None, + description="Exact provider-reported LLM usage; present for instrumented agentic queries.", + ) + + class Locator(BaseModel): """Where an evidence item lives in its source (page / segment / timestamp / bbox).""" diff --git a/nemo_retriever/src/nemo_retriever/service/vectordb_app.py b/nemo_retriever/src/nemo_retriever/service/vectordb_app.py index a59258961..89115f3e7 100644 --- a/nemo_retriever/src/nemo_retriever/service/vectordb_app.py +++ b/nemo_retriever/src/nemo_retriever/service/vectordb_app.py @@ -56,6 +56,7 @@ from nemo_retriever.service.agentic_query import run_agentic_query from nemo_retriever.service.config import AgenticConfig from nemo_retriever.service.query_schema import ( + AgenticQueryResponse, EvidenceQueryResponse, EvidenceResult, QueryRequest, @@ -667,13 +668,13 @@ async def delete_document( @app.post( "/v1/query", - response_model=Union[QueryResponse, EvidenceQueryResponse], + response_model=Union[AgenticQueryResponse, QueryResponse, EvidenceQueryResponse], tags=["query"], ) async def query( req: QueryRequest, x_nrl_scope: str | None = Header(None), - ) -> QueryResponse | EvidenceQueryResponse: + ) -> AgenticQueryResponse | QueryResponse | EvidenceQueryResponse: current = require_state() if req.agentic: if req.collection_name is not None: @@ -731,7 +732,7 @@ async def query( ) return QueryResponse(results=[QueryResult(hits=hits) for hits in hits_per_query]) - async def _run_agentic_query(req: QueryRequest) -> QueryResponse: + async def _run_agentic_query(req: QueryRequest) -> AgenticQueryResponse: """Run the blocking agentic workflow without consuming plain-query workers.""" current = require_state() if not agentic_config.enabled: diff --git a/nemo_retriever/tests/test_agentic_eval.py b/nemo_retriever/tests/test_agentic_eval.py index 4e85ba1bb..1fd15d499 100644 --- a/nemo_retriever/tests/test_agentic_eval.py +++ b/nemo_retriever/tests/test_agentic_eval.py @@ -12,8 +12,13 @@ import pytest -def _make_tool_call_response(fn_name: str, fn_args: dict, tc_id: str = "call_1") -> dict: - return { +def _make_tool_call_response( + fn_name: str, + fn_args: dict, + tc_id: str = "call_1", + usage: dict | None = None, +) -> dict: + response = { "choices": [ { "message": { @@ -30,6 +35,9 @@ def _make_tool_call_response(fn_name: str, fn_args: dict, tc_id: str = "call_1") } ] } + if usage is not None: + response["usage"] = usage + return response class FakeRetriever: @@ -133,6 +141,55 @@ def test_agentic_retriever_runs_graph_with_wrapped_retriever(): assert result["rank"].tolist() == list(range(1, 11)) +@patch("nemo_retriever.query.agentic.Retriever", FakeRetriever) +def test_agentic_retriever_returns_and_pops_query_usage(): + from nemo_retriever.query.agentic import AgenticRetrievalConfig, AgenticRetriever + + usage = {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15} + response = _make_tool_call_response( + "final_results", + { + "doc_ids": ["doc_1"] + [f"extra_{i}" for i in range(9)], + "message": "done", + "search_successful": "true", + }, + usage=usage, + ) + cfg = AgenticRetrievalConfig(llm_model="nemotron-8b") + with patch("nemo_retriever.query.agentic._build_agent_chat_completion_fn", return_value=lambda **_: response): + result = AgenticRetriever(cfg, match_mode="pdf_page").retrieve_with_usage(["customer-q"], ["find doc"]) + + assert result.usage == {"customer-q": {"main_agent": usage}} + assert result.documents["query_id"].tolist() == ["customer-q"] * 10 + + +@patch("nemo_retriever.query.agentic.Retriever", FakeRetriever) +def test_agentic_retriever_isolates_usage_for_concurrent_queries(): + from nemo_retriever.query.agentic import AgenticRetrievalConfig, AgenticRetriever + + usage = {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10} + response = _make_tool_call_response( + "final_results", + { + "doc_ids": ["doc_1"] + [f"extra_{i}" for i in range(9)], + "message": "done", + "search_successful": "true", + }, + usage=usage, + ) + cfg = AgenticRetrievalConfig(llm_model="nemotron-8b", num_concurrent=2) + with patch("nemo_retriever.query.agentic._build_agent_chat_completion_fn", return_value=lambda **_: response): + result = AgenticRetriever(cfg, match_mode="pdf_page").retrieve_with_usage( + ["customer-a", "customer-b"], + ["find a", "find b"], + ) + + assert result.usage == { + "customer-a": {"main_agent": usage}, + "customer-b": {"main_agent": usage}, + } + + @patch("nemo_retriever.query.agentic.Retriever", FakeRetriever) def test_agentic_retriever_rehydrates_only_retrieved_documents(): from nemo_retriever.query.agentic import AgenticRetrievalConfig, AgenticRetriever @@ -290,6 +347,45 @@ def test_agentic_query_documents_returns_classic_hit_fields_with_annotations(): retriever.unload.assert_called_once() +def test_agentic_query_documents_with_metadata_normalizes_usage(): + from nemo_retriever.query.agentic import AgenticRetrieveResult + from nemo_retriever.query.options import QueryAgenticOptions, QueryRequest, QueryRetrievalOptions + from nemo_retriever.query.workflow import agentic_query_documents_with_metadata + + retriever = MagicMock() + retriever.retrieve_with_usage.return_value = AgenticRetrieveResult( + documents=pd.DataFrame( + { + "query_id": ["0"], + "doc_id": ["doc_1"], + "rank": [1], + "result_source": ["final_results"], + "hit": [{"text": "body"}], + } + ), + usage={ + "0": { + "main_agent": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}, + "top1_agent": {"input_tokens": 3, "output_tokens": 2}, + } + }, + ) + request = QueryRequest( + query="find doc", + retrieval=QueryRetrievalOptions(top_k=1), + agentic=QueryAgenticOptions(enabled=True, llm_model="m", invoke_url=_REMOTE_URL), + ) + + with patch("nemo_retriever.query.workflow.build_agentic_retriever", return_value=retriever): + result = agentic_query_documents_with_metadata(request) + + assert result.usage["input_tokens"] == 14 + assert result.usage["output_tokens"] == 6 + assert result.usage["total_tokens"] == 20 + assert set(result.usage["stages"]) == {"main_agent", "top1_agent"} + retriever.unload.assert_called_once() + + @patch("nemo_retriever.query.agentic.Retriever", FakeRetriever) def test_agentic_retriever_honors_top_k(): """cfg.top_k drives the pipeline output count, not the hardcoded default of 10.""" diff --git a/nemo_retriever/tests/test_agentic_operators.py b/nemo_retriever/tests/test_agentic_operators.py index 37597320d..a9c995b03 100644 --- a/nemo_retriever/tests/test_agentic_operators.py +++ b/nemo_retriever/tests/test_agentic_operators.py @@ -188,6 +188,15 @@ def _op(self, **kwargs): def _input(self): return pd.DataFrame({"query_id": ["q1"], "query_text": ["What causes inflation?"]}) + def test_pop_query_usage_delegates_without_building_agent(self): + op = self._op() + assert op.pop_query_usage("q1") == {} + + op._agent = MagicMock() + op._agent.llm.pop_query_usage.return_value = {"main_agent": {"prompt_tokens": 3}} + assert op.pop_query_usage("q1") == {"main_agent": {"prompt_tokens": 3}} + op._agent.llm.pop_query_usage.assert_called_once_with("q1") + def test_retrieve_adapter_renames_and_coerces(self): op = self._op( retriever_fn=lambda q, k: [ @@ -380,6 +389,15 @@ def _rrf_frame(self, react_final_rank): } ) + def test_pop_query_usage_delegates_without_building_agent(self): + op = self._op() + assert op.pop_query_usage("q1") == {} + + op._sel = MagicMock() + op._sel.llm.pop_query_usage.return_value = {"top2_agent": {"completion_tokens": 2}} + assert op.pop_query_usage("q1") == {"top2_agent": {"completion_tokens": 2}} + op._sel.llm.pop_query_usage.assert_called_once_with("q1") + def test_final_results_passthrough(self): """Tier 1: a ReAct final list passes through; the selection agent is not run.""" from nemo_retriever.operators.graph_ops.selection_agent_operator import SelectionAgentOperator diff --git a/nemo_retriever/tests/test_root_query_cli.py b/nemo_retriever/tests/test_root_query_cli.py index 27ce7960f..ecfb4da24 100644 --- a/nemo_retriever/tests/test_root_query_cli.py +++ b/nemo_retriever/tests/test_root_query_cli.py @@ -455,6 +455,52 @@ def unload(self) -> None: ] +def test_root_query_agentic_include_usage_emits_envelope(monkeypatch) -> None: + from nemo_retriever.query.workflow import AgenticQueryDocumentsResult + + usage = { + "input_tokens": 12, + "output_tokens": 5, + "total_tokens": 17, + "stages": {"main_agent": {"prompt_tokens": 12, "completion_tokens": 5, "total_tokens": 17}}, + } + monkeypatch.setattr( + query_cli_app, + "query_agentic_documents_with_metadata", + lambda _request: AgenticQueryDocumentsResult( + hits=[{"doc_id": "a.pdf", "rank": 1, "result_source": "final_results"}], + usage=usage, + ), + ) + + result = RUNNER.invoke( + cli_main.app, + [ + "query", + "how does ingest work?", + "--agentic", + "--include-usage", + "--agentic-llm-model", + "model", + "--agentic-invoke-url", + "https://llm.example/v1/chat/completions", + ], + ) + + assert result.exit_code == 0 + assert json.loads(result.output) == { + "hits": [{"doc_id": "a.pdf", "rank": 1, "result_source": "final_results"}], + "usage": usage, + } + + +def test_root_query_include_usage_requires_agentic() -> None: + result = RUNNER.invoke(cli_main.app, ["query", "hello", "--include-usage"]) + + assert result.exit_code == 1 + assert "--include-usage requires --agentic" in result.output + + def test_root_query_agentic_rejects_candidate_k_below_top_k() -> None: result = RUNNER.invoke( cli_main.app, diff --git a/nemo_retriever/tests/test_service_agentic_query.py b/nemo_retriever/tests/test_service_agentic_query.py index 82e7771cb..51058d806 100644 --- a/nemo_retriever/tests/test_service_agentic_query.py +++ b/nemo_retriever/tests/test_service_agentic_query.py @@ -16,6 +16,7 @@ from nemo_retriever.service.agentic_query import ( agentic_ranked_to_hits, build_agentic_query_request, + run_agentic_query, ) from nemo_retriever.service.config import ( AgenticConfig, @@ -26,9 +27,9 @@ VectorDbConfig, ) from nemo_retriever.service.query_schema import ( + AgenticQueryResponse, MAX_AGENTIC_QUERY_CHARS, QueryRequest, - QueryResponse, QueryResult, ) from nemo_retriever.service.vectordb_app import VectorDBState, create_vectordb_app @@ -87,6 +88,44 @@ def test_build_agentic_query_request_maps_server_owned_configuration() -> None: assert request.agentic.react_max_steps == 7 +def test_run_agentic_query_includes_provider_usage() -> None: + from nemo_retriever.query.workflow import AgenticQueryDocumentsResult + + usage = { + "input_tokens": 12, + "output_tokens": 5, + "total_tokens": 17, + "stages": {"main_agent": {"prompt_tokens": 12, "completion_tokens": 5, "total_tokens": 17}}, + } + workflow_result = AgenticQueryDocumentsResult( + hits=[{"doc_id": "report_7", "rank": 1, "result_source": "final_results"}], + usage=usage, + ) + + with patch( + "nemo_retriever.service.agentic_query.agentic_query_documents_with_metadata", + return_value=workflow_result, + ): + response = run_agentic_query( + query="revenue trend", + top_k=1, + config=AgenticConfig( + enabled=True, + llm_model="model", + invoke_url="https://llm.example/v1/chat/completions", + ), + lancedb_uri="/indexes/finance", + table_name="finance", + embed_endpoint="https://embed.example/v1/embeddings", + embed_model="embed-model", + embed_model_provider_prefix=None, + embed_api_key="", + ) + + assert response.usage is not None + assert response.usage.model_dump() == usage + + def test_agentic_ranked_to_hits_keeps_rehydrated_classic_fields() -> None: hits = agentic_ranked_to_hits( [ @@ -187,7 +226,7 @@ def test_agentic_true_runs_react_workflow_on_v1_query(tmp_path) -> None: invoke_url="https://llm.example/v1/chat/completions", ), ) - expected = QueryResponse( + expected = AgenticQueryResponse( results=[ QueryResult( hits=agentic_ranked_to_hits( @@ -206,6 +245,12 @@ def test_agentic_true_runs_react_workflow_on_v1_query(tmp_path) -> None: ) ], query_mode="agentic", + usage={ + "input_tokens": 120, + "output_tokens": 30, + "total_tokens": 150, + "stages": {}, + }, ) with ( @@ -236,6 +281,12 @@ def test_agentic_true_runs_react_workflow_on_v1_query(tmp_path) -> None: } ], "query_mode": "agentic", + "usage": { + "input_tokens": 120, + "output_tokens": 30, + "total_tokens": 150, + "stages": {}, + }, } assert run_query.call_args.kwargs["query"] == "revenue trend" assert run_query.call_args.kwargs["top_k"] == 3 @@ -307,7 +358,7 @@ def test_agentic_query_slots_are_bounded_and_released_by_the_worker(tmp_path) -> invoke_url="https://llm.example/v1/chat/completions", ), ) - expected = QueryResponse(results=[QueryResult(hits=[])], query_mode="agentic") + expected = AgenticQueryResponse(results=[QueryResult(hits=[])]) with ( patch.object(VectorDBState, "table_exists", new_callable=PropertyMock, return_value=True), diff --git a/nemo_retriever/tests/test_service_mcp.py b/nemo_retriever/tests/test_service_mcp.py index bd4d45d73..38debd461 100644 --- a/nemo_retriever/tests/test_service_mcp.py +++ b/nemo_retriever/tests/test_service_mcp.py @@ -132,7 +132,13 @@ def _handler(request: httpx.Request) -> httpx.Response: } ] } - ] + ], + "usage": { + "input_tokens": 120, + "output_tokens": 30, + "total_tokens": 150, + "stages": {}, + }, }, ) @@ -157,6 +163,7 @@ def _handler(request: httpx.Request) -> httpx.Response: }, } assert result["results"][0]["hits"][0]["source"] == "report.pdf" + assert result["usage"]["total_tokens"] == 150 def test_query_methods_gate_mcp_retrieval_tools() -> None: