diff --git a/docs/docs/extraction/workflow-agentic-retrieval.md b/docs/docs/extraction/workflow-agentic-retrieval.md index fe7057c6e..4c72ee34d 100644 --- a/docs/docs/extraction/workflow-agentic-retrieval.md +++ b/docs/docs/extraction/workflow-agentic-retrieval.md @@ -187,6 +187,18 @@ curl -X POST http://localhost:7670/v1/query \ -d '{"query": "find documents about parser behavior", "top_k": 5, "agentic": true}' ``` +The service-query CLI sends the same request when you pass `--agentic`: + +```bash +retriever query service "find documents about parser behavior" \ + --service-url http://localhost:7670 \ + --agentic +``` + +This command uses the service's configured agentic LLM and embedding endpoints. +It does not start a local agent LLM or accept the local query command's +`--agentic-llm-model` and `--agentic-invoke-url` options. + When service auth is enabled, send `Authorization: Bearer ` (`NEMO_RETRIEVER_API_TOKEN`). Requests with `agentic: true` return HTTP `400` when agentic retrieval is not configured on the service. `top_k` cannot exceed the configured `agentic.backend_top_k` (default 20). Agentic queries are capped at 4,096 characters. diff --git a/nemo_retriever/docs/cli/README.md b/nemo_retriever/docs/cli/README.md index ed4ecb00b..c2ce852e3 100644 --- a/nemo_retriever/docs/cli/README.md +++ b/nemo_retriever/docs/cli/README.md @@ -119,6 +119,18 @@ retriever query service "What is in this corpus?" \ --service-url http://localhost:7670 ``` +To run the service's configured agentic retrieval workflow, add `--agentic`: + +```bash +retriever query service "What is in this corpus?" \ + --service-url http://localhost:7670 \ + --agentic +``` + +The service supplies the agent LLM and embedding configuration. This command +does not support local agent settings such as `--agentic-llm-model` or +`--agentic-invoke-url`. + ### Route ingest to hosted or self-hosted NIM endpoints ```bash diff --git a/nemo_retriever/src/nemo_retriever/cli/query/app.py b/nemo_retriever/src/nemo_retriever/cli/query/app.py index 6eef67091..6e98735c1 100644 --- a/nemo_retriever/src/nemo_retriever/cli/query/app.py +++ b/nemo_retriever/src/nemo_retriever/cli/query/app.py @@ -337,6 +337,7 @@ def _service_command( content_types: opts.ContentTypesOption = None, output_format: opts.OutputFormatOption = "hits", max_text_chars: opts.MaxTextCharsOption = None, + agentic: opts.AgenticOption = False, ) -> None: _validate_output_options(output_format, max_text_chars) silence_noisy_libraries() @@ -355,6 +356,7 @@ def _service_command( service_url=service_url, service_api_token=service_api_token, ), + agentic=agentic, ) ) except ROOT_CLI_ERRORS as exc: diff --git a/nemo_retriever/src/nemo_retriever/query/options.py b/nemo_retriever/src/nemo_retriever/query/options.py index 74def4e52..596a0e4d9 100644 --- a/nemo_retriever/src/nemo_retriever/query/options.py +++ b/nemo_retriever/src/nemo_retriever/query/options.py @@ -92,3 +92,4 @@ class ServiceQueryRequest: query: str retrieval: QueryRetrievalOptions = field(default_factory=QueryRetrievalOptions) service: QueryServiceOptions = field(default_factory=QueryServiceOptions) + agentic: bool = False diff --git a/nemo_retriever/src/nemo_retriever/query/service.py b/nemo_retriever/src/nemo_retriever/query/service.py index ef36f135e..3467101a1 100644 --- a/nemo_retriever/src/nemo_retriever/query/service.py +++ b/nemo_retriever/src/nemo_retriever/query/service.py @@ -27,7 +27,7 @@ def query_documents(request: ServiceQueryRequest) -> list[RetrievalHit]: base_url=request.service.service_url, api_token=request.service.service_api_token, ) - raw_result_sets = client.query(request.query, top_k=retrieval_top_k) + raw_result_sets = client.query(request.query, top_k=retrieval_top_k, agentic=request.agentic) raw_hits = raw_result_sets[0] return shape_query_hits( raw_hits, diff --git a/nemo_retriever/src/nemo_retriever/service/client.py b/nemo_retriever/src/nemo_retriever/service/client.py index 86d84698d..7b5e66745 100644 --- a/nemo_retriever/src/nemo_retriever/service/client.py +++ b/nemo_retriever/src/nemo_retriever/service/client.py @@ -553,6 +553,7 @@ def query( *, top_k: int, collection_name: str | None = None, + agentic: bool = False, ) -> list[list[dict[str, Any]]] | list[QueryHit]: """Search ingested documents through ``POST /v1/query``. @@ -560,7 +561,7 @@ def query( ``top_k`` is required here but defaults to 10 on :meth:`aquery`. That asymmetry is part of the released signature; do not unify it. """ - return self._run(self.aquery(query, top_k=top_k, collection_name=collection_name)) + return self._run(self.aquery(query, top_k=top_k, collection_name=collection_name, agentic=agentic)) def _query_hit(self, hit: dict[str, Any]) -> QueryHit: return self._model( @@ -580,12 +581,15 @@ async def aquery( *, top_k: int = 10, collection_name: str | None = None, + agentic: bool = False, ) -> list[list[dict[str, Any]]] | list[QueryHit]: """Asynchronously search through ``POST /v1/query``.""" payload: dict[str, Any] = {"query": query, "top_k": int(top_k)} if collection_name: payload["collection_name"] = collection_name + if agentic: + payload["agentic"] = True body = await self._arequest("POST", "/v1/query", json=payload) try: parsed = QueryResponse.model_validate(body).hits_by_query( diff --git a/nemo_retriever/tests/test_query_workflow_options.py b/nemo_retriever/tests/test_query_workflow_options.py index f17d274bd..e27419b46 100644 --- a/nemo_retriever/tests/test_query_workflow_options.py +++ b/nemo_retriever/tests/test_query_workflow_options.py @@ -189,8 +189,8 @@ class FakeServiceClient: def __init__(self, *, base_url: str, api_token: str | None = None, **_kwargs: Any) -> None: client_calls.append({"base_url": base_url, "api_token": api_token}) - def query(self, query: str, *, top_k: int) -> list[list[dict[str, Any]]]: - client_calls.append({"query": query, "top_k": top_k}) + def query(self, query: str, *, top_k: int, agentic: bool = False) -> list[list[dict[str, Any]]]: + client_calls.append({"query": query, "top_k": top_k, "agentic": agentic}) return [ [ {"text": "keep", "source": "doc.pdf", "page_number": 1, "metadata": {"type": "text"}}, @@ -217,7 +217,7 @@ def query(self, query: str, *, top_k: int) -> list[list[dict[str, Any]]]: ] assert client_calls == [ {"base_url": "http://svc:7670", "api_token": "secret"}, - {"query": "deployment?", "top_k": 3}, + {"query": "deployment?", "top_k": 3, "agentic": False}, ] diff --git a/nemo_retriever/tests/test_root_query_cli.py b/nemo_retriever/tests/test_root_query_cli.py index 27ce7960f..7d37de66e 100644 --- a/nemo_retriever/tests/test_root_query_cli.py +++ b/nemo_retriever/tests/test_root_query_cli.py @@ -938,6 +938,7 @@ def test_root_query_service_help_hides_local_only_options() -> None: assert "--content-types" in result.output assert "--format" in result.output assert "--max-text-chars" in result.output + assert "--agentic" in result.output assert "--run-mode" not in result.output assert "--lancedb-uri" not in result.output assert "--table-name" not in result.output @@ -991,11 +992,27 @@ def fake_query_documents(request: Any) -> list[dict[str, Any]]: assert request.retrieval.candidate_k == 5 assert request.retrieval.page_dedup is True assert request.retrieval.content_types == "text" + assert request.agentic is False assert json.loads(result.output) == [ {"modality": "text", "page_number": 3, "score": 0.2, "source": "doc.pdf", "text": "service passage"}, ] +def test_root_query_service_forwards_agentic_flag(monkeypatch) -> None: + requests: list[Any] = [] + + def fake_query_documents(request: Any) -> list[dict[str, Any]]: + requests.append(request) + return [] + + monkeypatch.setattr(query_cli_app, "query_service_documents", fake_query_documents) + + result = RUNNER.invoke(cli_main.app, ["query", "service", "deployment?", "--agentic"]) + + assert result.exit_code == 0 + assert requests[0].agentic is True + + def test_root_query_service_mode_rejects_local_storage_flags() -> None: result = RUNNER.invoke( cli_main.app, diff --git a/nemo_retriever/tests/test_service_query_client.py b/nemo_retriever/tests/test_service_query_client.py index 5232a6844..869d89a9e 100644 --- a/nemo_retriever/tests/test_service_query_client.py +++ b/nemo_retriever/tests/test_service_query_client.py @@ -73,6 +73,15 @@ def test_service_client_query_accepts_empty_hits(monkeypatch) -> None: assert RetrieverServiceClient(base_url="http://svc:7670").query("deployment?", top_k=2) == [[]] +def test_service_client_query_sends_agentic_flag(monkeypatch) -> None: + calls: list[dict[str, Any]] = [] + _install_query_response(monkeypatch, {"results": [{"hits": []}]}, calls) + + RetrieverServiceClient(base_url="http://svc:7670").query("deployment?", top_k=2, agentic=True) + + assert calls[1]["json"] == {"query": "deployment?", "top_k": 2, "agentic": True} + + @pytest.mark.parametrize( ("body", "match"), [