Skip to content

feat(rag): add native document retrieval service - #335

Merged
nvddr merged 5 commits into
mainfrom
agent/native-rag-service
Aug 4, 2026
Merged

feat(rag): add native document retrieval service#335
nvddr merged 5 commits into
mainfrom
agent/native-rag-service

Conversation

@nvddr

@nvddr nvddr commented Aug 3, 2026

Copy link
Copy Markdown
Collaborator

What

  • add an OpenAI-compatible embedding protocol, client, preset, and persistent vLLM server
  • add a standalone dense document RAG service with recursive ingestion, chunking, model-aware caching, and relevance filtering
  • expose retrieval and document discovery as the native typed xr_rag NAT function group
  • add the embedding service to shared model-server GPU profiles and document the service boundaries

Why

Samples need reusable retrieval without embedding FastMCP or hand-written model HTTP in each worker. This turns the earlier prototype into the current NAT-first architecture: model HTTP stays in xr-ai-models, retrieval state stays in a typed service, and agents compose ordinary NAT functions.

Impact

  • new persistent model endpoint: embedding server on port 8109
  • new private capability endpoint: RAG service on port 8340
  • no existing sample behavior changes until a sample opts into RAGFunctionsConfig

Checks

  • ruff check on all changed Python paths
  • uv run --project tests pytest -q tests/test_rag_service.py tests/test_models_config.py tests/test_models_openai_compat.py tests/test_models_protocols.py tests/test_models_deprecated_aliases.py tests/test_launcher_config.py tests/test_launcher_gpu.py
  • 103 tests passed
  • git diff --check

Comment thread tests/test_rag_service.py Fixed
@nvddr
nvddr marked this pull request as ready for review August 3, 2026 17:55
@wenxind-nvidia

Copy link
Copy Markdown
Collaborator

Review findings on current head 711ee6b:

Blockers

  1. Raise the supported vLLM minimum for the embedding model. ai-services/embedding-server/pyproject.toml:13 declares vllm>=0.12.0, while the NVIDIA model card requires vLLM 0.14.0 or newer. The default 26.04 Docker image contains vLLM 0.19.0, but the pip backend still permits unsupported 0.12/0.13 installations. Please raise the dependency floor to vllm>=0.14.0 and update DEPENDENCIES.md.

  2. Make RAG health reflect the required embedding backend. services/rag-service/rag_service/service.py:35-40 always returns ready: true, but every retrieval makes a live embedding request in DenseIndex.retrieve(). If the separately managed embedding server becomes unavailable after startup, RAGClient.health() reports healthy while retrieval fails. Please incorporate the embedder health into get_health and add a test that makes the fake embedder unhealthy.

Suggestions

  • Validate both cached and newly generated vectors against the complete expected shape (len(chunks), embedding_dim), and reject non-finite values before caching. Checking only the cached row count can defer an incompatible cache failure to an opaque NumPy matrix multiplication error.
  • Add cache-focused tests covering cache reuse, document-content invalidation, model or cache_key invalidation, dimension mismatch, and min_score filtering. The current RAG tests do not exercise the caching behavior described by the PR.
  • The current model card uses the model-configured pooling behavior and does not require --convert embed. Remove that forced conversion, or validate retrieval parity against the supported invocation on real hardware so the service does not accidentally replace the configured mean-pooling behavior.
  • Avoid logging the complete retrieval query at INFO. Logging sources and scores provides useful observability without persisting potentially sensitive user text.
  • Clean the remaining documentation drift: embedding_server.yaml still mentions rag-mcp-server; test_models_protocols.py still says there are four protocols; and some architecture/service lists omit EmbeddingService or the persistent embedding container.

@nvddr

nvddr commented Aug 3, 2026

Copy link
Copy Markdown
Collaborator Author

Addressed the review feedback in 4f521ed:

  • Raised the embedding server's pip backend requirement to vllm>=0.14.0 and updated the dependency documentation.
  • Made RAG health depend on the live embedding backend, with unhealthy-backend coverage.
  • Added exact vector-shape and finite-value validation for generated and cached embeddings; invalid derived caches are rebuilt.
  • Added coverage for cache reuse, document and cache-key invalidation, cached dimension mismatch, invalid generated vectors, and min_score filtering.
  • Removed the forced embedding conversion/runner flags and now rely on the model-configured mean-pooling behavior supported by the model card.
  • Removed complete query text from INFO logs; retrieval logs retain only sources and scores.
  • Corrected the remaining embedding/RAG protocol and service documentation drift.
  • Replaced the flagged cancelled-task cleanup expression with explicit asyncio.gather(..., return_exceptions=True).

Validation: 79 focused tests passed locally, Ruff passed, and all GitHub checks are green on the updated head.

@nvddr
nvddr force-pushed the agent/native-rag-service branch 2 times, most recently from 9ab7eda to 712a5da Compare August 4, 2026 21:42

@wenxind-nvidia wenxind-nvidia left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Requesting changes on current head 712a5da:

  • Rebase onto current main. The PR is currently unmergeable; merge analysis reports conflicts in DEPENDENCIES.md, README.md, agent-samples/model-servers/main.py, docs/changelog.md, and docs/source/getting_started/quickstart.md. The branch contains older render/model-server commits overlapping changes already merged through #337.
  • Please address the two inline findings before rerunning the required checks.

Verification: 8 RAG, 72 models, and 29 launcher tests passed in focused project environments. The combined test project cannot resolve on macOS because PyNvVideoCodec==2.2.0 has no macOS wheel. DCO and CodeQL pass on this head.

Comment thread services/rag-service/rag_service/index.py Outdated
Comment thread services/rag-service/rag_service/__main__.py
nvddr added 4 commits August 4, 2026 14:50
Signed-off-by: Devdeep Ray <devdeepr@Nvidia.com>
Signed-off-by: Devdeep Ray <devdeepr@Nvidia.com>
Signed-off-by: Devdeep Ray <devdeepr@Nvidia.com>
Signed-off-by: Devdeep Ray <devdeepr@Nvidia.com>
@nvddr
nvddr force-pushed the agent/native-rag-service branch from 712a5da to a2b6dca Compare August 4, 2026 21:51
Signed-off-by: Devdeep Ray <devdeepr@Nvidia.com>
@nvddr

nvddr commented Aug 4, 2026

Copy link
Copy Markdown
Collaborator Author

Addressed the requested changes on the current head a48f114: rebased onto current main after #337 merged, fixed the oversized-token chunking edge case with a regression test, and fixed the Ruff import-spacing finding. The focused suite passes (105 tests), and the full required CI—including DCO—is green.

@nvddr
nvddr merged commit b605035 into main Aug 4, 2026
14 checks passed
@nvddr
nvddr deleted the agent/native-rag-service branch August 4, 2026 22:05
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants