diff --git a/README.md b/README.md index 32159fc..d5a3716 100644 --- a/README.md +++ b/README.md @@ -76,6 +76,7 @@ The goal of this project is to demonstrate a production-style RAG architecture r RAG/ │ ├── app.py +├── pyproject.toml ├── requirements.txt ├── README.md ├── .env.example @@ -83,17 +84,17 @@ RAG/ ├── data/ │ └── pdf/ │ -├── faiss_store/ +├── faiss_store/ ← created at runtime │ ├── faiss.index │ └── metadata.pkl │ └── src/ + ├── __init__.py ├── rag.py ├── vector_store.py ├── synchronizer.py - ├── embedding.py - ├── document_generator.py - └── search.py + ├── embedding_generator.py + └── document_generator.py ``` --- diff --git a/app.py b/app.py index 77ddc7b..2cf3a31 100644 --- a/app.py +++ b/app.py @@ -1,13 +1,20 @@ from src.rag import RAG -def main(): - print("welcome!!!") -if __name__ == "__main__": +def main() -> None: rag = RAG() - query = "" + print("RAG system ready. Type 'exit' to quit.\n") + + while True: + query = input("Enter your query: ").strip() + if query.lower() == "exit": + print("Goodbye!") + break + if not query: + continue + answer = rag.search_and_summarize(query) + print(f"\nAnswer:\n{answer}\n") - while query != "exit": - query = input("Enter your RAG query:") - res = rag.search_and_summarize(query) - print(res) + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index dd15632..e829b60 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,13 +1,11 @@ [project] name = "rag" version = "0.1.0" -description = "Add your description here" +description = "Production-style RAG system with incremental FAISS indexing" readme = "README.md" requires-python = ">=3.10" dependencies = [ - "chromadb>=1.5.9", "faiss-cpu>=1.14.3", - "ipykernel>=7.3.0", "langchain>=1.3.14", "langchain-community>=0.4.2", "langchain-core>=1.4.9", @@ -17,3 +15,8 @@ dependencies = [ "python-dotenv>=1.2.2", "sentence-transformers>=5.6.0", ] + +[project.optional-dependencies] +dev = [ + "ipykernel>=7.3.0", +] diff --git a/requirements.txt b/requirements.txt index a51a2e2..535b0d9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,10 +1,9 @@ +faiss-cpu langchain langchain-core langchain-community -pypdf +langchain-groq pymupdf -faiss-cpu +pypdf +python-dotenv sentence-transformers -chromadb -langchain-groq -python-dotenv \ No newline at end of file diff --git a/src/document_generator.py b/src/document_generator.py index e2baf8f..0782884 100644 --- a/src/document_generator.py +++ b/src/document_generator.py @@ -1,25 +1,40 @@ -from langchain_community.document_loaders import PyPDFLoader,PyMuPDFLoader +from langchain_community.document_loaders import PyMuPDFLoader from pathlib import Path -from typing import List, Any -import os +from typing import List -def generate_documents(path,pdf_files:List[str]): - ##for loading pdf - all_documents=[] - for file in pdf_files: + +def generate_documents(path: Path, pdf_files: List[str]) -> List: + """Load and return LangChain documents from a list of PDF files. + + Args: + path: Directory that contains the PDF files. + pdf_files: List of PDF filenames (not full paths) to load. + + Returns: + A flat list of LangChain Document objects with ``source_name`` + and ``format_type`` set in their metadata. + """ + all_documents = [] + for file in pdf_files: try: - documents = PyMuPDFLoader(str(path.joinpath(file))).load() + documents = PyMuPDFLoader(str(path / file)).load() for doc in documents: - doc.metadata["source_name"] = str(file) - doc.metadata["format_type"] = 'pdf' + doc.metadata["source_name"] = file + doc.metadata["format_type"] = "pdf" all_documents.extend(documents) - print(f"Total {len(all_documents)} loaded from {file}") - + print(f"[INFO] Loaded {len(documents)} pages from '{file}' " + f"(total so far: {len(all_documents)})") except Exception as e: - print(f"[ERROR] Failed to load PDF {file}: {e}") + print(f"[ERROR] Failed to load PDF '{file}': {e}") continue - + return all_documents + if __name__ == "__main__": - generate_documents("../data/pdf") + # Resolve the data directory relative to this file so the script works + # regardless of the working directory. + pdf_dir = Path(__file__).parent.parent / "data" / "pdf" + pdf_files = [f.name for f in pdf_dir.glob("*.pdf")] + docs = generate_documents(pdf_dir, pdf_files) + print(f"Total documents loaded: {len(docs)}") diff --git a/src/embedding_generator.py b/src/embedding_generator.py index 05388cd..9d4c2ff 100644 --- a/src/embedding_generator.py +++ b/src/embedding_generator.py @@ -1,40 +1,56 @@ -from .document_generator import generate_documents +from pathlib import Path +from typing import List, Any + import numpy as np -from sentence_transformers import SentenceTransformer -from typing import List,Tuple,Dict,Any from langchain_text_splitters import RecursiveCharacterTextSplitter +from sentence_transformers import SentenceTransformer + +from .document_generator import generate_documents + class Generate_Embedding: - def __init__(self,model_name: str = "all-MiniLM-L6-v2",chunk_size=1000, chunk_overlap=200,): + """Splits documents into chunks and produces vector embeddings. + + A single :class:`SentenceTransformer` instance is created once and + reused for all calls, avoiding redundant model loads. + """ + + def __init__( + self, + model_name: str = "all-MiniLM-L6-v2", + chunk_size: int = 1000, + chunk_overlap: int = 200, + ): self.chunk_size = chunk_size self.chunk_overlap = chunk_overlap + self.model_name = model_name self.model = SentenceTransformer(model_name) - def convert_chunks(self,docs:List): + def convert_chunks(self, docs: List[Any]) -> List[Any]: + """Split a list of LangChain documents into smaller chunks.""" text_splitter = RecursiveCharacterTextSplitter( - chunk_size=self.chunk_size, - chunk_overlap=self.chunk_overlap, - length_function=len, - separators=["\n\n", "\n", " ", ""] + chunk_size=self.chunk_size, + chunk_overlap=self.chunk_overlap, + length_function=len, + separators=["\n\n", "\n", " ", ""], ) - chunks = text_splitter.split_documents(docs) - print(f"Split into {len(chunks)} chunks.") + print(f"[INFO] Split {len(docs)} documents into {len(chunks)} chunks.") return chunks - def embedd_chunks(self,texts: List[Any]): - text = [text.page_content for text in texts] - if self.model is None: - raise ValueError("Model is not loaded. Cannot generate embeddings.") - print(f"Generating embeddings for {len(text)} texts.") - embeddings = self.model.encode(text, show_progress_bar=True) - return embeddings + def embed_chunks(self, chunks: List[Any]) -> np.ndarray: + """Return a float32 embedding matrix for the given document chunks.""" + texts = [chunk.page_content for chunk in chunks] + print(f"[INFO] Generating embeddings for {len(texts)} chunks.") + embeddings = self.model.encode(texts, show_progress_bar=True) + return np.array(embeddings, dtype="float32") if __name__ == "__main__": - docs = generate_documents("../data/pdf") + pdf_dir = Path(__file__).parent.parent / "data" / "pdf" + pdf_files = [f.name for f in pdf_dir.glob("*.pdf")] + docs = generate_documents(pdf_dir, pdf_files) pipeline = Generate_Embedding() chunks = pipeline.convert_chunks(docs) - embeded_chunks = pipeline.embedd_chunks(chunks) - print(embeded_chunks) - \ No newline at end of file + embeddings = pipeline.embed_chunks(chunks) + print(f"Embedding matrix shape: {embeddings.shape}") diff --git a/src/rag.py b/src/rag.py index 7d94ec3..740c0a6 100644 --- a/src/rag.py +++ b/src/rag.py @@ -1,73 +1,152 @@ -from .vector_store import Vector_Store -from .document_generator import generate_documents import os -from langchain_groq import ChatGroq +from pathlib import Path +from typing import Optional + from dotenv import load_dotenv +from langchain_groq import ChatGroq +from langchain_core.messages import SystemMessage, HumanMessage + +from .document_generator import generate_documents from .synchronizer import Synchronizer -from pathlib import Path +from .vector_store import Vector_Store + load_dotenv() +# --------------------------------------------------------------------------- +# Default paths (resolved relative to this file so the app works from any +# working directory) +# --------------------------------------------------------------------------- +_PROJECT_ROOT = Path(__file__).parent.parent +_DEFAULT_PDF_DIR = _PROJECT_ROOT / "data" / "pdf" +_DEFAULT_PERSIST_DIR = str(_PROJECT_ROOT / "faiss_store") + + class RAG: - def __init__(self,persistent_str:str="faiss_store",model_name: str = "all-MiniLM-L6-v2",llm_model:str="llama-3.3-70b-versatile"): - self.docs = None - self.llm = None + """Retrieval-Augmented Generation pipeline. + + On first run (no persisted index) the system indexes every PDF found + in *pdf_dir*. On subsequent runs it loads the saved index and syncs + it against the current contents of *pdf_dir*. + + Args: + persist_dir: Directory for the FAISS index and metadata pickle. + pdf_dir: Directory that contains source PDF files. + model_name: Sentence-transformers model used for embeddings. + llm_model: Groq model identifier for answer generation. + distance_threshold: Maximum L2 distance for retrieved chunks to + be considered relevant. + """ + + def __init__( + self, + persist_dir: str = _DEFAULT_PERSIST_DIR, + pdf_dir: Optional[Path] = None, + model_name: str = "all-MiniLM-L6-v2", + llm_model: str = "llama-3.3-70b-versatile", + distance_threshold: float = 1.5, + ): self.llm_model = llm_model - self.persist_dir = persistent_str - self.store = Vector_Store(model_name=model_name) - self.synchronizer = Synchronizer(self.store) + self.pdf_dir = pdf_dir if pdf_dir is not None else _DEFAULT_PDF_DIR + + # Validate API key early so the error is clear and immediate. self.groq_api_key = os.getenv("GROQ_API_KEY") + if not self.groq_api_key: + raise EnvironmentError( + "GROQ_API_KEY is not set. " + "Add it to your .env file or export it as an environment variable." + ) - faiss_path = os.path.join(self.persist_dir, "faiss.index") - meta_path = os.path.join(self.persist_dir, "metadata.pkl") - - if not os.path.exists(faiss_path) or not os.path.exists(meta_path): + self.store = Vector_Store( + persist_dir=persist_dir, + model_name=model_name, + distance_threshold=distance_threshold, + ) + self.synchronizer = Synchronizer(self.store, pdf_dir=self.pdf_dir) + + # Lazy-initialised on first query. + self._llm: Optional[ChatGroq] = None + + faiss_path = os.path.join(persist_dir, "faiss.index") + meta_path = os.path.join(persist_dir, "metadata.pkl") - path = Path("data/pdf").resolve() - pdf_files = [file for file in os.listdir(path) if file.endswith(".pdf") ] - - self.docs = generate_documents(path,pdf_files) - self.store.building_from_doc(self.docs) + if not os.path.exists(faiss_path) or not os.path.exists(meta_path): + print("[INFO] No existing index found — building from scratch…") + pdf_files = [f.name for f in self.pdf_dir.glob("*.pdf")] + if not pdf_files: + print(f"[WARN] No PDF files found in '{self.pdf_dir}'.") + else: + docs = generate_documents(self.pdf_dir, pdf_files) + self.store.building_from_doc(docs) else: self.store.load() self.synchronizer.sync() - def search_and_summarize(self, query_text: str, top_k: int = 5) -> str: + # ------------------------------------------------------------------ + # LLM (lazy init) + # ------------------------------------------------------------------ - if self.llm is None: - self.llm = ChatGroq( + @property + def llm(self) -> ChatGroq: + """Return the ChatGroq client, creating it on first access.""" + if self._llm is None: + print(f"[INFO] Initialising Groq LLM: {self.llm_model}") + self._llm = ChatGroq( api_key=self.groq_api_key, model_name=self.llm_model, - temperature=0.1 + temperature=0.1, ) + return self._llm + + # ------------------------------------------------------------------ + # Public interface + # ------------------------------------------------------------------ + + def search_and_summarize(self, query_text: str, top_k: int = 5) -> str: + """Retrieve relevant chunks and generate an answer with the LLM. - print(f"[INFO] Groq LLM initialized: {self.llm_model}") + Args: + query_text: The user's question. + top_k: Number of chunks to retrieve from the vector store. + Returns: + A string answer grounded in the retrieved context. + """ results = self.store.query(query_text, top_k=top_k) - texts = [r.get("metadata", {}).get("text_content", "") for r in results] - context = "\n\n".join(texts) - if not context: - return "No relevant documents found." - - prompt=f""" - - If the answer is not explicitly stated in the context, respond exactly: - - I don't know based on the provided context. - - Do not infer, assume, or use outside knowledge. - - Context: - {context} - - Question: {query_text} - - Answer: - - """ - response = self.llm.invoke([prompt]) + + if not results: + return "I don't know based on the provided context." + + context_parts = [ + r["metadata"]["text_content"] + for r in results + if r.get("metadata", {}).get("text_content") + ] + context = "\n\n---\n\n".join(context_parts) + + if not context.strip(): + return "I don't know based on the provided context." + + system_prompt = ( + "You are a helpful assistant that answers questions strictly based on " + "the provided context. If the answer is not explicitly stated in the " + "context, respond with exactly: " + "'I don't know based on the provided context.' " + "Do not infer, assume, or use outside knowledge." + ) + + user_prompt = ( + f"Context:\n{context}\n\n" + f"Question: {query_text}\n\n" + "Answer:" + ) + + response = self.llm.invoke( + [SystemMessage(content=system_prompt), HumanMessage(content=user_prompt)] + ) return response.content + if __name__ == "__main__": rag = RAG() - #res = rag.search_and_summarize("what is attention") - #print(res) + result = rag.search_and_summarize("what is attention") + print(result) diff --git a/src/synchronizer.py b/src/synchronizer.py index 616904a..5f01528 100644 --- a/src/synchronizer.py +++ b/src/synchronizer.py @@ -1,60 +1,91 @@ -from .vector_store import Vector_Store -from .document_generator import generate_documents -import os -import pickle from pathlib import Path +from typing import Optional + +from .document_generator import generate_documents +from .vector_store import Vector_Store + class Synchronizer: - def __init__(self,store:Vector_Store,persistent_str:str="faiss_store"): + """Keeps the FAISS vector store in sync with the PDF directory. + + On each :meth:`sync` call the synchronizer: + + 1. Detects PDFs that have been added and incrementally indexes them. + 2. Detects PDFs that have been deleted and rebuilds the full index + (required because ``IndexFlatL2`` does not support per-vector + deletion efficiently). + + Args: + store: The :class:`Vector_Store` instance to synchronize. + pdf_dir: Directory that contains the source PDF files. Defaults + to ``/data/pdf``. + """ + + def __init__( + self, + store: Vector_Store, + pdf_dir: Optional[Path] = None, + ): self.store = store - self.docs = None - self.new_pdfs = [] - self.metadatas = [] - self.persist_dir = persistent_str - - def sync(self): - folder_data = set([file for file in os.listdir("data/pdf") if file.endswith(".pdf")]) - existing_metadata = set([f['source'] for f in self.store.metadata]) - print("folder_data:",folder_data) - print("existing_metadata:",existing_metadata) - pdfs = folder_data-existing_metadata - - if pdfs: - self.new_pdfs.extend(pdfs) - print(f"found {len(self.new_pdfs)} pdf has be added") - print(self.new_pdfs) - path = Path("data/pdf").resolve() - self.docs = generate_documents(path,self.new_pdfs) - self.store.building_from_doc(self.docs) - self.check_delete() - else: - print("found no new pdfs in folder") - self.check_delete() + # Fall back to the conventional location relative to *this* file. + self.pdf_dir: Path = ( + pdf_dir + if pdf_dir is not None + else Path(__file__).parent.parent / "data" / "pdf" + ) + + # ------------------------------------------------------------------ + # Public interface + # ------------------------------------------------------------------ + + def sync(self) -> None: + """Detect additions and deletions and update the store.""" + folder_pdfs = self._get_folder_pdfs() + indexed_pdfs = self._get_indexed_pdfs() + print(f"[INFO] PDFs on disk : {folder_pdfs}") + print(f"[INFO] PDFs indexed : {indexed_pdfs}") - def check_delete(self): - folder_data = set([file for file in os.listdir("data/pdf") if file.endswith(".pdf")]) - existing_metadata = set([f['source'] for f in self.store.metadata]) - print("folder_data:",folder_data) - print("existing_metadata:",existing_metadata) - deleted_pdfs = existing_metadata-folder_data + new_pdfs = folder_pdfs - indexed_pdfs + deleted_pdfs = indexed_pdfs - folder_pdfs + + if new_pdfs: + print(f"[INFO] {len(new_pdfs)} new PDF(s) detected: {new_pdfs}") + docs = generate_documents(self.pdf_dir, list(new_pdfs)) + self.store.building_from_doc(docs) if deleted_pdfs: - self.store.index=None - self.store.metadata=[] - print(f"found {len(deleted_pdfs)} pdf has been deleted") - print(f"current vector index and metadata has been flushed......") - path = Path("data/pdf").resolve() - pdf_files = [file for file in os.listdir(path) if file.endswith(".pdf") ] - - self.docs = generate_documents(path,pdf_files) - self.store.building_from_doc(self.docs) + self._handle_deletions(deleted_pdfs, folder_pdfs) + elif not new_pdfs: + print("[INFO] Vector store is already up to date.") + + # ------------------------------------------------------------------ + # Private helpers + # ------------------------------------------------------------------ + + def _get_folder_pdfs(self): + return {f.name for f in self.pdf_dir.glob("*.pdf")} + + def _get_indexed_pdfs(self): + return {entry["source"] for entry in self.store.metadata} + + def _handle_deletions(self, deleted_pdfs, remaining_folder_pdfs) -> None: + """Flush the index and rebuild from the surviving PDFs.""" + print(f"[INFO] {len(deleted_pdfs)} deleted PDF(s) detected: {deleted_pdfs}") + print("[INFO] Flushing index and metadata for full rebuild…") + + self.store.index = None + self.store.metadata = [] + + if remaining_folder_pdfs: + docs = generate_documents(self.pdf_dir, list(remaining_folder_pdfs)) + self.store.building_from_doc(docs) else: - print("found no deleted pdfs in folder") - self.store.load() + print("[INFO] No PDFs remaining — vector store is now empty.") if __name__ == "__main__": - synchronizer = Synchronizer(Vector_Store) + store = Vector_Store() + store.load() + synchronizer = Synchronizer(store) synchronizer.sync() - \ No newline at end of file diff --git a/src/vector_store.py b/src/vector_store.py index a93ee57..58d8a30 100644 --- a/src/vector_store.py +++ b/src/vector_store.py @@ -1,78 +1,157 @@ +import hashlib import os +import pickle +from pathlib import Path +from typing import Any, Dict, List, Optional + import faiss import numpy as np -import pickle -import hashlib -from typing import List, Any from sentence_transformers import SentenceTransformer + from .document_generator import generate_documents from .embedding_generator import Generate_Embedding + class Vector_Store: - def __init__(self,persistent_str:str="faiss_store",model_name: str = "all-MiniLM-L6-v2",chunk_size=1000, chunk_overlap=200,): + """Persistent FAISS vector store with incremental indexing support. + + A single :class:`SentenceTransformer` model instance is shared + between :class:`Generate_Embedding` (for building) and query-time + encoding to avoid loading the model twice. + + Args: + persist_dir: Directory where ``faiss.index`` and ``metadata.pkl`` + are stored. + model_name: HuggingFace sentence-transformers model identifier. + chunk_size: Token/character chunk size for text splitting. + chunk_overlap: Overlap between consecutive chunks. + distance_threshold: Maximum L2 distance for a result to be + considered relevant. Results beyond this threshold are + filtered out. Set to ``None`` to disable filtering. + """ + + def __init__( + self, + persist_dir: str = "faiss_store", + model_name: str = "all-MiniLM-L6-v2", + chunk_size: int = 1000, + chunk_overlap: int = 200, + distance_threshold: Optional[float] = 1.5, + ): self.chunk_size = chunk_size self.chunk_overlap = chunk_overlap - self.model = None - self.persist_dir = persistent_str + self.distance_threshold = distance_threshold + self.persist_dir = persist_dir os.makedirs(self.persist_dir, exist_ok=True) - self.index = None - self.metadata = [] - def building_from_doc(self,docs:List[Any]): - print(f"[INFO] Building vector store from {len(docs)} raw documents...") - pipeline = Generate_Embedding() + self.index: Optional[faiss.Index] = None + self.metadata: List[Dict[str, Any]] = [] + + # One shared model instance used for both building and querying. + self.model_name = model_name + self._model: Optional[SentenceTransformer] = None + + @property + def model(self) -> SentenceTransformer: + """Lazy-load the embedding model (initialised at most once).""" + if self._model is None: + print(f"[INFO] Loading embedding model '{self.model_name}'…") + self._model = SentenceTransformer(self.model_name) + return self._model + + # ------------------------------------------------------------------ + # Building / updating + # ------------------------------------------------------------------ + + def building_from_doc(self, docs: List[Any]) -> None: + """Chunk *docs*, embed them, and add to the FAISS index.""" + print(f"[INFO] Building vector store from {len(docs)} raw documents…") + pipeline = Generate_Embedding( + model_name=self.model_name, + chunk_size=self.chunk_size, + chunk_overlap=self.chunk_overlap, + ) + # Reuse the already-loaded model so there is only one copy in memory. + pipeline.model = self.model + chunks = pipeline.convert_chunks(docs) - embeded_chunks = pipeline.embedd_chunks(chunks) - metadatas = [{"text_content":doc.page_content,"id":hashlib.sha256(doc.page_content.encode()).hexdigest(), - "source":doc.metadata["source_name"]} for doc in chunks] - ids = [m["id"] for m in metadatas] - - data_exist = self.add_to_memory(np.array(embeded_chunks).astype("float32"),metadatas) - if data_exist: + embedded_chunks = pipeline.embed_chunks(chunks) + + metadatas = [ + { + "text_content": chunk.page_content, + "id": hashlib.sha256(chunk.page_content.encode()).hexdigest(), + "source": chunk.metadata["source_name"], + } + for chunk in chunks + ] + + added = self.add_to_memory(embedded_chunks, metadatas) + if added: self.save() - print(f"[INFO] Vector store built and saved to {self.persist_dir}") + print(f"[INFO] Vector store built and saved to '{self.persist_dir}'.") else: - print("No new files to save!!!") - - def add_to_memory(self,embeddings:np.ndarray,metadatas): - new_embeddings =[] - new_metadata = [] - dup = [] - existing_ids = [doc['id'] for doc in self.metadata] + print("[INFO] No new vectors to save.") + + def add_to_memory( + self, embeddings: np.ndarray, metadatas: List[Dict[str, Any]] + ) -> bool: + """Add non-duplicate vectors to the in-memory FAISS index. + + Returns: + ``True`` if at least one new vector was added, ``False`` + otherwise. + """ + existing_ids = {doc["id"] for doc in self.metadata} dim = embeddings.shape[1] + if self.index is None: self.index = faiss.IndexFlatL2(dim) - for embedding,metadata in zip(embeddings,metadatas): + + new_embeddings: List[np.ndarray] = [] + new_metadata: List[Dict[str, Any]] = [] + duplicates = 0 + + for embedding, metadata in zip(embeddings, metadatas): if metadata["id"] in existing_ids: - ##print(f"skipping duplicate id: {metadata['id']}") - dup.append(metadata['id']) + duplicates += 1 continue new_embeddings.append(embedding) new_metadata.append(metadata) - print(f"found total {len(dup)} duplicate files") + + print(f"[INFO] Skipped {duplicates} duplicate chunk(s).") if not new_embeddings: print("[INFO] No new documents to add.") return False - - new_embeddings = np.array(new_embeddings, dtype="float32") - self.index.add(new_embeddings) - if new_metadata: - self.metadata.extend(new_metadata) - print(f"[INFO] Added {new_embeddings.shape[0]} vectors to Faiss index.") - print(f"Total number of docs in FAISS: {self.index.ntotal}") - print(f"Total dnumber of docs in metadata: {len(self.metadata)}") + + batch = np.array(new_embeddings, dtype="float32") + self.index.add(batch) + self.metadata.extend(new_metadata) + + print(f"[INFO] Added {batch.shape[0]} vector(s). " + f"Index total: {self.index.ntotal} | Metadata total: {len(self.metadata)}") return True - def save(self): + # ------------------------------------------------------------------ + # Persistence + # ------------------------------------------------------------------ + + def save(self) -> None: + """Write the FAISS index and metadata to disk.""" faiss_path = os.path.join(self.persist_dir, "faiss.index") meta_path = os.path.join(self.persist_dir, "metadata.pkl") faiss.write_index(self.index, faiss_path) with open(meta_path, "wb") as f: pickle.dump(self.metadata, f) - print(f"[INFO] Saved Faiss index and metadata to {self.persist_dir}") + print(f"[INFO] Saved FAISS index and metadata to '{self.persist_dir}'.") + + def load(self) -> bool: + """Load the FAISS index and metadata from disk. - def load(self): + Returns: + ``True`` if files were found and loaded, ``False`` otherwise. + """ faiss_path = os.path.join(self.persist_dir, "faiss.index") meta_path = os.path.join(self.persist_dir, "metadata.pkl") @@ -81,44 +160,61 @@ def load(self): self.index = None self.metadata = [] return False - + self.index = faiss.read_index(faiss_path) with open(meta_path, "rb") as f: self.metadata = pickle.load(f) - print(f"[INFO] Loaded Faiss index and metadata from {self.persist_dir}") + print(f"[INFO] Loaded FAISS index ({self.index.ntotal} vectors) " + f"and {len(self.metadata)} metadata entries from '{self.persist_dir}'.") return True - def search(self, query_embedding: np.ndarray, top_k: int = 5): - D, I = self.index.search(query_embedding, top_k) - # print("Distances:", D) - # print("Indices:", I) + # ------------------------------------------------------------------ + # Retrieval + # ------------------------------------------------------------------ + + def search( + self, query_embedding: np.ndarray, top_k: int = 5 + ) -> List[Dict[str, Any]]: + """Run a nearest-neighbour search and return filtered results. + + Results whose L2 distance exceeds ``self.distance_threshold`` + are excluded to avoid returning irrelevant chunks. + """ + if self.index is None or self.index.ntotal == 0: + print("[WARN] FAISS index is empty — no results returned.") + return [] + + distances, indices = self.index.search(query_embedding, top_k) results = [] - for idx, dist in zip(I[0], D[0]): - #if dist < 0.6: - meta = self.metadata[idx] if idx < len(self.metadata) else None - results.append({"index": idx, "distance": dist, "metadata": meta}) + for idx, dist in zip(indices[0], distances[0]): + if idx < 0 or idx >= len(self.metadata): + continue + if self.distance_threshold is not None and dist > self.distance_threshold: + continue + results.append( + {"index": int(idx), "distance": float(dist), "metadata": self.metadata[idx]} + ) return results - def query(self,query_text: str, top_k: int = 3,model_name: str = "all-MiniLM-L6-v2",): - if self.model is None: - self.model = SentenceTransformer(model_name) - - print(f"[INFO] Querying vector store for: '{query_text}'") - query_emb = self.model.encode([query_text]).astype('float32') + def query(self, query_text: str, top_k: int = 3) -> List[Dict[str, Any]]: + """Encode *query_text* and return the top-k nearest chunks.""" + print(f"[INFO] Querying vector store: '{query_text}'") + query_emb = self.model.encode([query_text]).astype("float32") return self.search(query_emb, top_k=top_k) - if __name__ == "__main__": - docs = generate_documents("../data/pdf") - store = Vector_Store() + pdf_dir = Path(__file__).parent.parent / "data" / "pdf" + pdf_files = [f.name for f in pdf_dir.glob("*.pdf")] + docs = generate_documents(pdf_dir, pdf_files) + store = Vector_Store() if store.load(): print("Existing vector store loaded.") else: print("Creating a new vector store.") + store.building_from_doc(docs) - store.building_from_doc(docs) - - print(store.query("what is python")) - + results = store.query("what is python") + for r in results: + print(r)