Implement embeddings and ChromaDB storage with query retrieval - #5
Conversation
…rieval for Milestone 5
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
|
Warning Review limit reached
More reviews will be available in 34 minutes and 32 seconds. Learn how PR review limits work. Your organization has run out of usage credits. Purchase more in the billing tab. ⌛ How to resolve this issue?After more reviews become available, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans include higher PR review limits than trial, open-source, and free plans. In all cases, reviews become available again over time. During sustained high-volume PR review activity, CodeRabbit may temporarily slow when the next review becomes available. Please see our Fair Usage Limits Policy for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (4)
✨ Finishing Touches🧪 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 |
There was a problem hiding this comment.
Pull request overview
Adds an embedding + ChromaDB storage layer (Milestone 5) that turns chunked PDF content into Gemini text-embedding-004 vectors, persists them in a per-ticker Chroma collection, and exposes a filtered top-k retrieval helper. A standalone end-to-end script exercises the pipeline against a real PDF, and a numpy/chromadb compatibility shim is added to vector_store.py.
Changes:
- New
backend/engines/embedder.pywithembed_text/embed_text_query,store_chunks_in_vector_db, andretrieve_chunks(model fallback + tenacity retries, cosine-distance → relevance score). backend/core/vector_store.pymonkey-patchesnp.float_ = np.float64to keep ChromaDB working under NumPy 2.x.- Adds
backend/test_milestone5.pyintegration script and ignoreschroma_db/andtest_doc.pdfin.gitignore.
Reviewed changes
Copilot reviewed 3 out of 4 changed files in this pull request and generated 7 comments.
| File | Description |
|---|---|
| backend/engines/embedder.py | New Gemini embedding + ChromaDB store/retrieve API with retries, model fallback, and metadata filtering. |
| backend/core/vector_store.py | Adds NumPy 2.x compatibility shim before importing chromadb. |
| backend/test_milestone5.py | End-to-end script: prepares a PDF, chunks it, stores embeddings, queries, and asserts on relevance scores. |
| .gitignore | Ignores the local chroma_db/ directory and root-level test_doc.pdf. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| import numpy as np | ||
| np.float_ = np.float64 |
| @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=2, min=4, max=60)) | ||
| async def embed_text(text: str) -> list[float]: | ||
| """Embed a single text string using Gemini text-embedding-004 with fallback.""" | ||
| for model_name in ["models/text-embedding-004", "models/gemini-embedding-2"]: | ||
| try: | ||
| result = genai.embed_content( | ||
| model=model_name, | ||
| content=text, | ||
| task_type="retrieval_document" | ||
| ) | ||
| return result["embedding"] | ||
| except Exception as e: | ||
| err_msg = str(e).lower() | ||
| if "not found" in err_msg or "404" in err_msg or "not supported" in err_msg: | ||
| logger.warning(f"Embedding model {model_name} not found or unsupported, trying fallback...") | ||
| continue | ||
| logger.error(f"Embedding failed for model {model_name}: {e}") | ||
| raise | ||
| raise Exception("All embedding models failed.") |
| @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=2, min=4, max=60)) | ||
| async def embed_text_query(text: str) -> list[float]: | ||
| """Embed a query string (different task_type from document embedding).""" | ||
| for model_name in ["models/text-embedding-004", "models/gemini-embedding-2"]: | ||
| try: | ||
| result = genai.embed_content( | ||
| model=model_name, | ||
| content=text, | ||
| task_type="retrieval_query" | ||
| ) | ||
| return result["embedding"] | ||
| except Exception as e: | ||
| err_msg = str(e).lower() | ||
| if "not found" in err_msg or "404" in err_msg or "not supported" in err_msg: | ||
| logger.warning(f"Embedding model {model_name} not found or unsupported, trying fallback...") | ||
| continue | ||
| logger.error(f"Query embedding failed for model {model_name}: {e}") | ||
| raise | ||
| raise Exception("All query embedding models failed.") |
| for chunk in chunks: | ||
| try: | ||
| chroma_id = f"{ticker}_{document_id}_{chunk.chunk_index if hasattr(chunk, 'chunk_index') else chunk.get('chunk_index', 0)}" | ||
| content = chunk.content if hasattr(chunk, 'content') else chunk.get('content', '') | ||
| section_type = chunk.section_type if hasattr(chunk, 'section_type') else chunk.get('section_type', 'GENERAL') | ||
|
|
||
| embedding = await embed_text(content) | ||
|
|
||
| collection.add( | ||
| ids=[chroma_id], | ||
| embeddings=[embedding], | ||
| documents=[content], | ||
| metadatas=[{ | ||
| "ticker": ticker, | ||
| "document_id": document_id, | ||
| "document_type": document_type, | ||
| "fiscal_year": fiscal_year or 0, | ||
| "section_type": section_type, | ||
| "section_title": chunk.section_title if hasattr(chunk, 'section_title') else chunk.get('section_title', ''), | ||
| "chunk_index": chunk.chunk_index if hasattr(chunk, 'chunk_index') else chunk.get('chunk_index', 0), | ||
| "token_estimate": chunk.token_estimate if hasattr(chunk, 'token_estimate') else chunk.get('token_estimate', 0), | ||
| }] | ||
| ) | ||
| stored_count += 1 | ||
|
|
||
| except Exception as e: | ||
| logger.warning(f"Failed to store chunk {chunk.chunk_index if hasattr(chunk, 'chunk_index') else chunk.get('chunk_index', '?') if isinstance(chunk, dict) else '?'}: {e}") | ||
|
|
||
| logger.info(f"Stored {stored_count}/{len(chunks)} chunks for {ticker} doc {document_id}") | ||
| return stored_count |
| async def retrieve_chunks( | ||
| ticker: str, | ||
| query: str, | ||
| n_results: int = 8, | ||
| section_type_filter: str = None, | ||
| document_type_filter: str = None, | ||
| fiscal_year_filter: int = None | ||
| ) -> list[dict]: | ||
| """ | ||
| Retrieve top-k relevant chunks for a query. | ||
| Apply optional filters for section type, document type, fiscal year. | ||
| """ | ||
| collection_name = f"stockiq_{ticker.lower()}" | ||
|
|
||
| try: | ||
| collection = get_or_create_collection(collection_name) | ||
| except Exception: | ||
| return [] | ||
|
|
||
| # Build where filter | ||
| where_conditions = [] | ||
| if section_type_filter: | ||
| where_conditions.append({"section_type": {"$eq": section_type_filter}}) | ||
| if document_type_filter: | ||
| where_conditions.append({"document_type": {"$eq": document_type_filter}}) | ||
| if fiscal_year_filter: | ||
| where_conditions.append({"fiscal_year": {"$eq": fiscal_year_filter}}) |
| # If no local PDF, download a sample | ||
| print("No local PDF found. Downloading a sample Reliance financial PDF...") | ||
| url = "https://www.rfil.co.in/pdf/financial/reliance-financial-2023.pdf" # Fallback sample URL | ||
| try: | ||
| urllib.request.urlretrieve(url, test_pdf) | ||
| print(f"Downloaded sample to {test_pdf}") | ||
| except Exception as e: | ||
| print(f"Download failed: {e}") | ||
| # Try a simpler, reliable PDF if the above fails | ||
| simple_url = "https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf" | ||
| urllib.request.urlretrieve(simple_url, test_pdf) | ||
| print(f"Downloaded dummy PDF to {test_pdf}") |
| collection_name = f"stockiq_{ticker.lower()}" | ||
|
|
||
| try: | ||
| collection = get_or_create_collection(collection_name) | ||
| except Exception: | ||
| return [] |
Implement embeddings and ChromaDB storage engine with query