Feat/phase 14 nemo retriever#15
Merged
Merged
Conversation
…essing Implement core/nemo_retriever.py wrapping NVIDIA's NeMo Retriever capabilities into a single interface for HECTOR's ingestion pipeline. NemoRetrieverProvider provides: - OCR processing via Nemotron OCR API (nvidia/nemotron-ocr-v1) * ocr_page(): process raw image bytes to text * ocr_page_from_pdf(): render PDF page then OCR * Returns NemoOCRResult with text, markdown, confidence, timing - Embedding generation via Nemotron embedding API * embed_documents(): batch embed document chunks (2048d) * embed_query(): single query embedding for search - Reranking via Nemotron reranker API * rerank(): score and sort documents by query relevance * Returns NemoRerankResult with index, score, text, reasons - Full document processing pipeline * process_document(): OCR → chunk → embed in one call * Intelligent chunking with sentence boundary detection * Automatic fallback from pypdf extraction to OCR Data classes: - NemoOCRResult: OCR output with confidence and timing - NemoChunk: processed document chunk with metadata and embedding - NemoRerankResult: reranked document with score and reasons Factory function: - get_nemo_retriever(): creates provider if HECTOR_NEMO_RETRIEVER_ENABLED=1 * Checks NVIDIA_API_KEY availability * Performs health check before returning * Returns None if not enabled or unreachable Env vars: - HECTOR_NEMO_RETRIEVER_ENABLED: '1' to enable (default: '0') - NVIDIA_API_KEY: Required for all NeMo operations - HECTOR_NEMO_OCR_MODEL: OCR model override - HECTOR_NEMO_EMBED_MODEL: Embedding model override - HECTOR_NEMO_RERANK_MODEL: Rerank model override
Update utils/enhanced_ingestor.py to use NemoRetrieverProvider when HECTOR_NEMO_RETRIEVER_ENABLED=1 and NVIDIA_API_KEY is set. Changes: - __init__(): initialize nemo_retriever via get_nemo_retriever() * Stores as self.nemo_retriever for use in OCR fallback * Logs availability status for debugging - _nvidia_ocr_fallback(): try NeMo Retriever first * If nemo_retriever.is_available, use ocr_page_from_pdf() * Falls back to legacy direct API call on failure * Maintains backward compatibility when NeMo is disabled The integration is opt-in and transparent: - Default behavior unchanged (HECTOR_NEMO_RETRIEVER_ENABLED=0) - When enabled, uses NeMo for OCR with automatic fallback - Legacy OCR path preserved for environments without NeMo access Env vars: - HECTOR_NEMO_RETRIEVER_ENABLED: '1' to enable (default: '0') - NVIDIA_API_KEY: Required for NeMo operations
Tests for the unified NemoRetriever provider covering all capabilities: Data classes: - NemoOCRResult: stores text, markdown, confidence, timing, model - NemoChunk: stores text, metadata, optional embedding vector - NemoRerankResult: stores index, score, text, reasons Provider initialization: - Default models (nemotron-ocr-v1, nemotron-embed-4b-v1, rerank-1b-v2) - Custom model override via constructor arguments - Environment variable configuration (NVIDIA_API_KEY, HECTOR_NEMO_*) Health check: - Returns False without API key - Returns True on 200/401/405 responses - Returns False on timeout/exception - Result is cached after first call OCR processing: - Raises ValueError without API key - Returns NemoOCRResult with text and confidence on success - Propagates API errors Embedding: - Raises ValueError without API key - Single document returns 2048-dimensional vector - Batch embedding returns correct count - Query embedding returns single vector Reranking: - Raises ValueError without API key - Empty document list returns empty result - Returns sorted NemoRerankResult list with scores Document pipeline: - Short text returns single chunk - Empty/whitespace text returns empty list - Long text produces multiple non-empty chunks Factory: - Returns None when disabled (default) - Returns provider when enabled with valid API - Returns None without API key - Returns None when API unreachable Integration: - Source code verification for ingestor integration All tests use explicit api_key=None + manual assignment to avoid env var leakage from concurrent test execution (966 tests total).
Update utils/diagnostics.py to include a 4th diagnostic test for the unified NeMo Retriever provider. Changes: - test_nemo_retriever(): new diagnostic method [4/4] * Imports get_nemo_retriever from core.nemo_retriever * Checks if HECTOR_NEMO_RETRIEVER_ENABLED=1 * If enabled, verifies API reachability via is_available * Reports SKIPPED/SUCCESS/FAILED status - Updated test counters from [1/3]-[3/3] to [1/4]-[4/4] - Updated run_diagnostics() to include NeMo test in results The NeMo Retriever test is non-blocking: returns True (pass) when skipped, allowing the diagnostic suite to complete even when NeMo is not configured.
|
Caution Review failedThe pull request is closed. ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (4)
📝 WalkthroughWalkthroughA new ChangesNeMo Retriever Provider Integration
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Poem
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit
New Features
Bug Fixes