diff --git a/.gitignore b/.gitignore index fad2fff..db50ed6 100644 --- a/.gitignore +++ b/.gitignore @@ -16,3 +16,6 @@ documents/ # PDF test files backend/test_doc.pdf +chroma_db/ +test_doc.pdf + diff --git a/backend/core/vector_store.py b/backend/core/vector_store.py index ec4a733..05be863 100644 --- a/backend/core/vector_store.py +++ b/backend/core/vector_store.py @@ -1,3 +1,5 @@ +import numpy as np +np.float_ = np.float64 import chromadb from chromadb.config import Settings as ChromaSettings from loguru import logger diff --git a/backend/engines/embedder.py b/backend/engines/embedder.py new file mode 100644 index 0000000..1018d49 --- /dev/null +++ b/backend/engines/embedder.py @@ -0,0 +1,171 @@ +""" +Embedding + ChromaDB Storage Engine + +Uses Google Gemini's text-embedding-004 model. +Stores chunks in ChromaDB with metadata for filtered retrieval. + +Collection naming: stockiq_{ticker_lowercase} + e.g. stockiq_reliance, stockiq_tcs +""" + +import google.generativeai as genai +from loguru import logger +from tenacity import retry, stop_after_attempt, wait_exponential +from ..core.vector_store import get_or_create_collection +from ..core.config import settings + +genai.configure(api_key=settings.GEMINI_API_KEY) + +@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.") + +async def store_chunks_in_vector_db( + ticker: str, + document_id: int, + document_type: str, + fiscal_year: int | None, + chunks: list +) -> int: + """ + Embed all chunks and store in ChromaDB. + Returns number of chunks stored. + """ + collection_name = f"stockiq_{ticker.lower()}" + collection = get_or_create_collection(collection_name) + + stored_count = 0 + + 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}}) + + where = None + if len(where_conditions) == 1: + where = where_conditions[0] + elif len(where_conditions) > 1: + where = {"$and": where_conditions} + + # Embed the query + try: + query_embedding = await embed_text_query(query) + except Exception as e: + logger.error(f"Query embedding failed: {e}") + return [] + + count = collection.count() + if count == 0: + return [] + + kwargs = { + "query_embeddings": [query_embedding], + "n_results": min(n_results, count), + "include": ["documents", "metadatas", "distances"] + } + if where: + kwargs["where"] = where + + results = collection.query(**kwargs) + + chunks = [] + if results and "documents" in results and results["documents"]: + for i, doc in enumerate(results["documents"][0]): + chunks.append({ + "content": doc, + "metadata": results["metadatas"][0][i], + "distance": results["distances"][0][i], + "relevance_score": 1 - results["distances"][0][i] + }) + + return chunks + +@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.") diff --git a/backend/test_milestone5.py b/backend/test_milestone5.py new file mode 100644 index 0000000..3b1f9f7 --- /dev/null +++ b/backend/test_milestone5.py @@ -0,0 +1,110 @@ +# backend/test_milestone5.py +import os +import shutil +import urllib.request +import asyncio +from pathlib import Path +from backend.engines.pdf_extractor import extract_pdf +from backend.engines.semantic_chunker import chunk_document +from backend.engines.embedder import store_chunks_in_vector_db, retrieve_chunks +from backend.core.vector_store import get_chroma_client + + +def prepare_test_pdf(): + test_pdf = Path("test_doc.pdf") + if test_pdf.exists(): + print(f"Using existing {test_pdf}") + return + + # Check if test_doc.pdf is in the backend directory + backend_pdf = Path("backend/test_doc.pdf") + if backend_pdf.exists(): + shutil.copy(backend_pdf, test_pdf) + print(f"Copied {backend_pdf} to {test_pdf}") + return + + # Check local document directories + possible_dirs = [Path("documents/TCS"), Path("../documents/TCS"), Path("backend/documents/TCS")] + for d in possible_dirs: + if d.exists(): + pdfs = list(d.glob("*.pdf")) + if pdfs: + # Use the smallest one for speed + pdfs.sort(key=lambda p: p.stat().st_size) + shutil.copy(pdfs[0], test_pdf) + print(f"Copied {pdfs[0]} to {test_pdf}") + return + + + # 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}") + +async def run_milestone5_test(): + # 1. Prepare PDF + prepare_test_pdf() + + # 2. Extract PDF + result = extract_pdf("test_doc.pdf") + print(f"Pages extracted: {len(result['pages'])}") + print(f"Tables found: {len(result['tables'])}") + + # 3. Chunk Document + chunks = chunk_document(result["pages"], "QUARTERLY_RESULTS") + print(f"Chunks created: {len(chunks)}") + + # 4. Clean up any existing Chroma collection for TCS to ensure clean test + client = get_chroma_client() + try: + client.delete_collection("stockiq_tcs") + print("Cleared existing stockiq_tcs collection") + except Exception: + pass + + # 5. Store chunks in Vector DB + stored = await store_chunks_in_vector_db( + ticker="TCS", + document_id=1, + document_type="QUARTERLY_RESULTS", + fiscal_year=2024, + chunks=chunks + ) + print(f"Stored {stored} chunks in ChromaDB") + + # 6. Retrieve and verify + query = "revenue growth and EBITDA margin trends" + print(f"Retrieving chunks for query: '{query}'") + results = await retrieve_chunks( + ticker="TCS", + query=query, + n_results=5 + ) + + print(f"Retrieved {len(results)} chunks") + high_relevance_count = 0 + for i, r in enumerate(results): + score = r['relevance_score'] + section_type = r['metadata']['section_type'] + preview = r['content'][:150].replace('\n', ' ') + print(f" {i+1}. Section: {section_type} | Relevance Score: {score:.4f}") + print(f" Preview: {preview}...") + if score > 0.4: + high_relevance_count += 1 + + # 7. Assertions + assert len(results) >= 3, f"Expected at least 3 retrieved chunks, got {len(results)}" + assert high_relevance_count >= 3, f"Expected at least 3 chunks with relevance score > 0.4, got {high_relevance_count}" + + print("\nMilestone 5: ALL TESTS PASSED") + +if __name__ == "__main__": + asyncio.run(run_milestone5_test())