Expose agentic retrieval via MCP and HTTP service endpoints - #2416
Expose agentic retrieval via MCP and HTTP service endpoints#2416mahikaw wants to merge 1 commit into
Conversation
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| async with _query_semaphore: | ||
| return await asyncio.to_thread( | ||
| run_agentic_query, | ||
| req, |
There was a problem hiding this comment.
When an agentic request times out or disconnects before the ReAct workflow finishes, asyncio.to_thread leaves the synchronous work running while it retains the semaphore shared with plain queries, causing subsequent queries to block after all four slots are occupied.
Rule Used: Models and stateful components accessed from Ray a... (source)
Knowledge Base Used: Service: HTTP/MCP API surface
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/service/vectordb_app.py
Line: 513-516
Comment:
**Timed-out work retains query slots**
When an agentic request times out or disconnects before the ReAct workflow finishes, `asyncio.to_thread` leaves the synchronous work running while it retains the semaphore shared with plain queries, causing subsequent queries to block after all four slots are occupied.
**Rule Used:** Models and stateful components accessed from Ray a... ([source](nemo_retriever/.greptile))
**Knowledge Base Used:** [Service: HTTP/MCP API surface](https://app.greptile.com/nvidia-public-github/-/custom-context/knowledge-base/nvidia/nemo-retriever/-/docs/service-api.md)
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.| class AgenticQueryRequest(BaseModel): | ||
| """One query for the server-configured agentic retrieval pipeline.""" | ||
|
|
||
| query: str = Field(min_length=1) |
There was a problem hiding this comment.
The public request schema accepts arbitrarily large query strings and forwards them into the multi-step LLM workflow, allowing disproportionate prompt size, inference cost, and latency; add a defensible max_length constraint.
Knowledge Base Used: Service: HTTP/MCP API surface
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/service/query_schema.py
Line: 42
Comment:
**Agentic query length is unbounded**
The public request schema accepts arbitrarily large query strings and forwards them into the multi-step LLM workflow, allowing disproportionate prompt size, inference cost, and latency; add a defensible `max_length` constraint.
**Knowledge Base Used:** [Service: HTTP/MCP API surface](https://app.greptile.com/nvidia-public-github/-/custom-context/knowledge-base/nvidia/nemo-retriever/-/docs/service-api.md)
---
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!
Greptile SummaryAdds config-gated agentic retrieval across the VectorDB HTTP API, gateway proxy, MCP facade, and Helm deployment.
|
| Filename | Overview |
|---|---|
| nemo_retriever/src/nemo_retriever/service/vectordb_app.py | Adds the execution endpoint and CLI configuration, but non-cancellable threaded workflows can retain shared query capacity after callers time out. |
| nemo_retriever/src/nemo_retriever/query/workflow.py | Adds explicit embedding-key resolution for service mode but reuses that credential for the independently configured LLM endpoint. |
| nemo_retriever/src/nemo_retriever/service/agentic_query.py | Provides a thin request adapter but exposes no separate credential path for the remote agentic LLM. |
| nemo_retriever/src/nemo_retriever/service/query_schema.py | Defines slim agentic request and response models, with the query text lacking an upper size bound. |
| nemo_retriever/src/nemo_retriever/service/routers/ingest.py | Adds a config-gated gateway proxy with the configured agentic HTTP timeout. |
| nemo_retriever/src/nemo_retriever/service/mcp_server.py | Adds a config-gated MCP tool and dedicated client timeout without changing existing tools. |
| nemo_retriever/src/nemo_retriever/service/config.py | Adds validated server-owned agentic settings requiring a remote model and invoke URL when enabled. |
| nemo_retriever/helm/templates/deployment-vectordb.yaml | Wires agentic CLI arguments and the shared NVIDIA credential into the VectorDB deployment. |
| nemo_retriever/helm/templates/configmap.yaml | Renders the agentic configuration for the gateway service. |
| nemo_retriever/tests/test_service_agentic_query.py | Covers configuration, mapping, endpoint execution, top-k validation, and proxying but not timeout/disconnect lifecycle behavior. |
| nemo_retriever/tests/test_service_mcp.py | Covers the dedicated MCP request path, additive tool registration, and config-derived timeout. |
Sequence Diagram
sequenceDiagram
participant C as REST/MCP Client
participant G as Gateway Service
participant V as VectorDB App
participant E as Embedding Endpoint
participant L as Agentic LLM
C->>G: POST /v1/agentic/query
G->>V: Proxy query + top_k
V->>V: Acquire shared query semaphore
V->>E: Multi-hop retrieval embeddings
V->>L: ReAct and selection calls
L-->>V: Selected document IDs
V-->>G: rank/doc_id/result_source
G-->>C: AgenticQueryResponse
Comments Outside Diff (1)
-
nemo_retriever/src/nemo_retriever/query/workflow.py, line 166-184 (link)Embedding key overwrites LLM credential
When the remote embedding service and OpenAI-compatible LLM require different credentials, this code assigns the embedding API key to both clients, causing one upstream request path to fail authentication.
Knowledge Base Used:
Prompt To Fix With AI
This is a comment left during a code review. Path: nemo_retriever/src/nemo_retriever/query/workflow.py Line: 166-184 Comment: **Embedding key overwrites LLM credential** When the remote embedding service and OpenAI-compatible LLM require different credentials, this code assigns the embedding API key to both clients, causing one upstream request path to fail authentication. **Knowledge Base Used:** - [Query Pipeline](https://app.greptile.com/nvidia-public-github/-/custom-context/knowledge-base/nvidia/nemo-retriever/-/docs/query-pipeline.md) - [Repo Overview and Deployment](https://app.greptile.com/nvidia-public-github/-/custom-context/knowledge-base/nvidia/nemo-retriever/-/docs/repo-overview-and-deployment.md) --- For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.
Prompt To Fix All With AI
### Issue 1
nemo_retriever/src/nemo_retriever/service/vectordb_app.py:513-516
**Timed-out work retains query slots**
When an agentic request times out or disconnects before the ReAct workflow finishes, `asyncio.to_thread` leaves the synchronous work running while it retains the semaphore shared with plain queries, causing subsequent queries to block after all four slots are occupied.
### Issue 2
nemo_retriever/src/nemo_retriever/query/workflow.py:166-184
**Embedding key overwrites LLM credential**
When the remote embedding service and OpenAI-compatible LLM require different credentials, this code assigns the embedding API key to both clients, causing one upstream request path to fail authentication.
### Issue 3
nemo_retriever/src/nemo_retriever/service/query_schema.py:42
**Agentic query length is unbounded**
The public request schema accepts arbitrarily large query strings and forwards them into the multi-step LLM workflow, allowing disproportionate prompt size, inference cost, and latency; add a defensible `max_length` constraint.
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.Reviews (1): Last reviewed commit: "Expose agentic retrieval via MCP and HTT..." | Re-trigger Greptile
1 participant
Description
Add server-owned AgenticConfig and thin service adapters so existing agentic_query_documents() can run in service mode without changing AgenticRetriever internals.
Expose POST /v1/agentic/query on the VectorDB app (execution) and gateway proxy; add config-gated MCP tool agentic_query alongside unchanged plain query.
Wire Helm/retriever-service.yaml so gateway gets enablement + timeout from ConfigMap, and the VectorDB deployment gets --agentic* CLI flags + API key when a remote LLM URL is set.
Service mode requires remote invoke_url + llm_model (and remote embed); request body stays {query, top_k} with server-owned ReAct knobs. Hits remain slim (rank / doc_id / result_source).
Checklist