""" Local RAG setup with LangChain, Ollama/OpenAI, and FAISS Minimal dependencies, simple code """ import json import os import re from pathlib import Path import numpy as np from langchain_community.document_loaders import PyPDFLoader, TextLoader from langchain_community.vectorstores import FAISS from langchain_huggingface import HuggingFaceEmbeddings from langchain_ollama import ChatOllama from langchain_openai import ChatOpenAI from langchain_text_splitters import RecursiveCharacterTextSplitter import docling_loader # Pipeline tuning: used when adding documents (rebuild index after changing). # Characters, not tokens: 1500 characters of German prose are roughly 400 tokens. CHUNK_SIZE = 1500 CHUNK_OVERLAP = 300 # Multilingual embedding model with a long input window (8192 tokens), so chunks # are never truncated. Changing it requires rebuilding the index. EMBEDDING_MODEL = "BAAI/bge-m3" # How PDFs are read: "pypdf" (PyPDFLoader, one unit per page) or "docling" # (structured output from prepare_corpus.py, cleaned by docling_loader.py). PDF_LOADER = "pypdf" # Where prepare_corpus.py put the Docling output for the PDFs being indexed. PROCESSED_DIR = "corpus_processed/data_split" # Prefix every chunk with author and title from the file name (Docling: plus section), # so questions like "according to Lehmann ..." can match. See docling_loader.chunk_header. CHUNK_HEADERS = True # Retrieval / query tuning RETRIEVAL_K = 5 # Hybrid search: besides the vector search, a keyword search (BM25 on German word stems) # finds names, titles and quotations that embeddings blur. # "dense" vector search alone # "hybrid" both lists merged with Reciprocal Rank Fusion, which also sets the order # "union" BM25 only adds candidates (BM25_CANDIDATES) to the vector candidates; # the order comes from the vector search or, if enabled, from the reranker RETRIEVAL_MODE = "hybrid" CANDIDATES_K = 50 # candidates per search before merging and reranking BM25_CANDIDATES = 20 # extra keyword candidates in "union" mode RRF_K = 60 # standard constant of Reciprocal Rank Fusion # Cross-encoder that rereads question and candidate together and re-sorts the candidates. # None switches reranking off. RERANKER_MODEL = "BAAI/bge-reranker-v2-m3" # Appended to the embedding query for FAISS (does not affect the answer-generation prompt). # Left empty on purpose: embedding models do not follow instructions, and on the Kafka # corpus the former text ("Do not return a list of references ...") lowered Recall@5 # from 0.67 to 0.53 (eval_retrieval.py --query-instructions). Measure before re-enabling. QUERY_ADDITIONAL_INSTRUCTIONS = "" # Written next to every index; an index built with another embedding model is not loaded. INDEX_INFO_FILE = "index_info.json" # Opening instructions for the final answer LLM call only. The rest (history, context, question) # is assembled in query_with_history. ANSWER_PROMPT = """You are an assistant for question-answering. Use the chat history (if any) and the retrieved context below to answer the current question. If you don't know the answer, say so. Keep the conversation coherent.""" class LocalRAG: def __init__( self, vectorstore_path="./vectorstore", llm_provider="ollama", ollama_model="gpt-oss:20b", openai_model="gpt-5-mini", ollama_base_url="http://localhost:11434", pdf_loader=PDF_LOADER, processed_dir=PROCESSED_DIR, retrieval_mode=RETRIEVAL_MODE, reranker_model=RERANKER_MODEL, chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP, chunk_headers=CHUNK_HEADERS, ): """Initialize local RAG system. llm_provider: 'ollama' or 'openai'. pdf_loader: 'pypdf' or 'docling' (see PDF_LOADER).""" self.vectorstore_path = vectorstore_path self.llm_provider = llm_provider self.pdf_loader = pdf_loader self.processed_dir = Path(processed_dir) self.chunk_size, self.chunk_overlap, self.chunk_headers = chunk_size, chunk_overlap, chunk_headers self.retrieval_mode = retrieval_mode self.reranker_model = reranker_model self._bm25 = None # built on first use from the documents in the index self._reranker = None # loaded on first use # Embeddings print("Loading embeddings model...") # Normalized vectors: FAISS's L2 distance then ranks exactly like cosine similarity. self.embeddings = HuggingFaceEmbeddings( model_name=EMBEDDING_MODEL, encode_kwargs={"normalize_embeddings": True}, query_encode_kwargs={"normalize_embeddings": True}, ) # Text splitter (used when adding documents; tune CHUNK_SIZE / CHUNK_OVERLAP above) self.text_splitter = RecursiveCharacterTextSplitter( chunk_size=chunk_size, chunk_overlap=chunk_overlap, ) print(f"Text splitter: chunk_size={chunk_size}, chunk_overlap={chunk_overlap}") # LLM (Ollama or OpenAI) if llm_provider == "openai": api_key = os.environ.get("OPENAI_API_KEY") if not api_key: raise ValueError( "OPENAI_API_KEY environment variable is required when llm_provider='openai'" ) print(f"Using OpenAI (model: {openai_model})...") self.llm = ChatOpenAI(model=openai_model, api_key=api_key) else: print(f"Using Ollama (model: {ollama_model})...") self.llm = ChatOllama( model=ollama_model, base_url=ollama_base_url ) # Vector store (load if exists, otherwise None) self.vectorstore = None self._load_vectorstore() def _index_info(self): """Settings that must match between the index on disk and this code.""" return { "embedding_model": EMBEDDING_MODEL, "chunk_size": self.chunk_size, "chunk_overlap": self.chunk_overlap, "pdf_loader": self.pdf_loader, "chunk_headers": self.chunk_headers, } def _index_matches(self, index_file): """False if the index was built with a different embedding model than EMBEDDING_MODEL.""" info_file = os.path.join(self.vectorstore_path, INDEX_INFO_FILE) if os.path.exists(info_file): with open(info_file, encoding="utf-8") as f: built_with = json.load(f).get("embedding_model") if built_with != EMBEDDING_MODEL: print(f"Index in {self.vectorstore_path} was built with {built_with}, " f"but EMBEDDING_MODEL is {EMBEDDING_MODEL}. Rebuild it (add_pdfs.py).") return False return True # Older index without info file: at least the vector dimensions must agree import faiss index_dim = faiss.read_index(index_file).d model_dim = len(self.embeddings.embed_query("dimension check")) if index_dim != model_dim: print(f"Index in {self.vectorstore_path} has {index_dim} dimensions, {EMBEDDING_MODEL} " f"produces {model_dim}. It was built with another model; rebuild it (add_pdfs.py).") return False return True def _load_vectorstore(self): """Load existing vector store if available""" index_file = os.path.join(self.vectorstore_path, "index.faiss") if os.path.exists(index_file) and self._index_matches(index_file): try: self.vectorstore = FAISS.load_local( self.vectorstore_path, self.embeddings, allow_dangerous_deserialization=True ) print(f"Loaded existing vector store from {self.vectorstore_path}") except Exception as e: print(f"Could not load vector store: {e}") self.vectorstore = None def add_documents(self, file_paths): """Add documents to the vector store""" print(f"\nLoading {len(file_paths)} document(s)...") all_docs = [] for file_path in file_paths: path = Path(file_path) if not path.exists(): print(f"Warning: {file_path} not found, skipping") continue # Docling: already structured and cleaned, the loader returns finished chunks if path.suffix.lower() == '.pdf' and self.pdf_loader == "docling": processed = self.processed_dir / f"{path.stem}.json" if processed.exists(): chunks = docling_loader.load(processed, str(path), self.text_splitter, header=self.chunk_headers) all_docs.extend(chunks) print(f" - {path.name}: {len(chunks)} chunks (docling)") continue print(f"Warning: no Docling output for {path.name} in {self.processed_dir}, using PyPDFLoader") # Load document if path.suffix.lower() == '.pdf': loader = PyPDFLoader(str(path)) elif path.suffix.lower() in ['.txt', '.md']: loader = TextLoader(str(path)) else: print(f"Warning: Unsupported file type {path.suffix}, skipping") continue docs = loader.load() chunks = self.text_splitter.split_documents(docs) if self.chunk_headers: header = docling_loader.chunk_header(str(path)) for chunk in chunks: chunk.page_content = header + chunk.page_content all_docs.extend(chunks) print(f" - {path.name}: {len(chunks)} chunks") if not all_docs: print("No documents loaded!") return # Create or update vector store print(f"\nCreating embeddings for {len(all_docs)} chunks...") if self.vectorstore is None: self.vectorstore = FAISS.from_documents(all_docs, self.embeddings) else: new_store = FAISS.from_documents(all_docs, self.embeddings) self.vectorstore.merge_from(new_store) # Save os.makedirs(self.vectorstore_path, exist_ok=True) self.vectorstore.save_local(self.vectorstore_path) with open(os.path.join(self.vectorstore_path, INDEX_INFO_FILE), "w", encoding="utf-8") as f: json.dump(self._index_info(), f, indent=2) print(f"Vector store saved to {self.vectorstore_path}") def list_documents(self): """List all documents in the vector store""" if self.vectorstore is None: print("No documents in vector store.") return [] # Get all documents from the vector store # We'll retrieve a large number to get all documents all_docs = self.vectorstore.similarity_search("", k=10000) # Large k to get all # Extract unique document sources from metadata documents = {} for doc in all_docs: source = doc.metadata.get('source', 'Unknown') if source not in documents: documents[source] = { 'source': source, 'chunks': 0, 'page': doc.metadata.get('page', None) } documents[source]['chunks'] += 1 # Convert to list and sort doc_list = list(documents.values()) doc_list.sort(key=lambda x: x['source']) print(f"\nDocuments in vector store ({len(doc_list)} unique documents):") print("-" * 60) for doc_info in doc_list: print(f" - {doc_info['source']}") print(f" Chunks: {doc_info['chunks']}") if doc_info['page'] is not None: print(f" Page: {doc_info['page']}") return doc_list def _format_history(self, chat_history): """Format chat history as a string for prompts.""" lines = [] for turn in chat_history or []: role = (turn.get("role") or "").lower() content = (turn.get("content") or "").strip() if role == "user": lines.append(f"User: {content}") elif role == "assistant": lines.append(f"Assistant: {content}") return "\n".join(lines) if lines else "" def _docs_to_retrieved(self, docs): """Convert document list to retrieved chunks format for API.""" return [ { "content": doc.page_content, "source": doc.metadata.get("source", ""), "page": doc.metadata.get("page"), } for doc in docs ] def _docs_scores_to_retrieved(self, docs_with_scores): """Convert (Document, score) list to retrieved chunks format with score. FAISS returns L2 distance (lower = more similar).""" return [ { "content": doc.page_content, "source": doc.metadata.get("source", ""), "page": doc.metadata.get("page"), "score": float(score), } for doc, score in docs_with_scores ] # ---------- retrieval: vector search, keyword search, reranking ---------- _stemmer = None @classmethod def _bm25_tokens(cls, text): """Lowercase German word stems, so "ohnmächtigen" matches "Ohnmacht".""" if cls._stemmer is None: import snowballstemmer cls._stemmer = snowballstemmer.stemmer("german") return cls._stemmer.stemWords(re.findall(r"[a-zäöüß0-9]{2,}", text.lower())) def _bm25_index(self): if self._bm25 is None: from rank_bm25 import BM25Okapi docs = list(self.vectorstore.docstore._dict.values()) self._bm25 = (BM25Okapi([self._bm25_tokens(d.page_content) for d in docs]), docs) return self._bm25 def retrieve(self, query, k=None): """The k best chunks for a query as [(Document, score)]. dense: vector search only (score = L2 distance, lower is better) hybrid: vector search + BM25, merged by Reciprocal Rank Fusion (higher is better) union: vector candidates plus BM25_CANDIDATES keyword candidates, vector order first With a reranker, the merged candidates are re-sorted by the cross-encoder (score = relevance, higher is better).""" k = k or RETRIEVAL_K pool = max(k, CANDIDATES_K) if (self.retrieval_mode != "dense" or self.reranker_model) else k dense = self.vectorstore.similarity_search_with_score(query, k=pool) key = lambda doc: (doc.metadata.get("source"), doc.metadata.get("page"), doc.page_content[:200]) if self.retrieval_mode in ("hybrid", "union"): bm25, docs = self._bm25_index() scores = bm25.get_scores(self._bm25_tokens(query)) keyword = [docs[i] for i in np.argsort(scores)[::-1][:pool]] if self.retrieval_mode == "union": seen = {key(d) for d, _ in dense} extra = [d for d in keyword if key(d) not in seen][:BM25_CANDIDATES] candidates = [[d, s] for d, s in dense] + [[d, float("nan")] for d in extra] elif self.retrieval_mode == "hybrid": fused = {} for ranking in ([d for d, _ in dense], keyword): for rank, doc in enumerate(ranking): entry = fused.setdefault(key(doc), [doc, 0.0]) entry[1] += 1 / (RRF_K + rank + 1) candidates = sorted(fused.values(), key=lambda e: e[1], reverse=True)[:pool] else: candidates = [[d, s] for d, s in dense] if self.reranker_model and candidates: if self._reranker is None: from sentence_transformers import CrossEncoder print(f"Loading reranker {self.reranker_model}...") self._reranker = CrossEncoder(self.reranker_model, max_length=1024) relevance = self._reranker.predict([(query, d.page_content) for d, _ in candidates]) candidates = [[d, float(r)] for (d, _), r in zip(candidates, relevance)] candidates.sort(key=lambda e: e[1], reverse=True) return [(d, s) for d, s in candidates[:k]] def query(self, question, k=None): """Query the RAG system (no conversation history). Returns dict with 'answer' and 'retrieved'.""" if k is None: k = RETRIEVAL_K return self.query_with_history(question, chat_history=[], k=k) def query_with_history(self, question, chat_history=None, k=None): """Query the RAG with conversation history: rephrase question using history for retrieval, then answer with full conversation + retrieved context in the prompt. Returns dict with 'answer' and 'retrieved' (list of chunks with content, source, page). """ if k is None: k = RETRIEVAL_K if self.vectorstore is None: return { "answer": "Error: No documents loaded. Please add documents first.", "retrieved": [], } history_str = self._format_history(chat_history) search_query = question print(f"[RAG] User question: {question!r}") # 1) If we have history, rephrase the question into a standalone query for better retrieval if history_str.strip(): rephrase_prompt = f"""Given this chat history and the latest user question, write a single standalone question that captures what the user is asking. Do not answer it; only output the standalone question. If the latest question is already clear on its own, output it unchanged. Chat history: {history_str} Latest user question: {question} Standalone question:""" rephrase_response = self.llm.invoke(rephrase_prompt) search_query = (rephrase_response.content if hasattr(rephrase_response, "content") else str(rephrase_response)).strip() or question print(f"[RAG] Standalone search query (rephrased): {search_query!r}") retrieval_query = search_query if QUERY_ADDITIONAL_INSTRUCTIONS: retrieval_query += f"\n\n{QUERY_ADDITIONAL_INSTRUCTIONS}" print(f"[RAG] Search query: {search_query!r}") print(f"[RAG] Retrieval query sent to vector store: {retrieval_query!r}") # 2) Retrieve documents with scores (see retrieve() for what the score means) docs_with_scores = self.retrieve(retrieval_query, k=k) docs = [doc for doc, _ in docs_with_scores] retrieved = self._docs_scores_to_retrieved(docs_with_scores) if docs_with_scores: scores = [f"{s:.3f}" for _, s in docs_with_scores] print(f"[RAG] Retrieved {len(docs)} chunk(s), scores (L2 dist): [{', '.join(scores)}]") context = "\n\n".join([doc.page_content for doc in docs]) # 3) Answer using conversation history + retrieved context history_block = f"Chat history:\n{history_str}\n\n" if history_str else "" answer_prompt = ( ANSWER_PROMPT + f""" {history_block}Relevant context from documents: {context} Current question: {question} Answer:""" ) print(f"[RAG] Composed answer prompt:\n{answer_prompt}") response = self.llm.invoke(answer_prompt) answer = response.content if hasattr(response, "content") else str(response) return {"answer": answer, "retrieved": retrieved} def main(): """Example usage""" print("=" * 60) print("Local RAG with LangChain, Ollama/OpenAI, and FAISS") print("=" * 60) # Initialize rag = LocalRAG(ollama_model="gpt-oss:20b") # Add documents (uncomment and add your file paths) # rag.add_documents([ # "data/dok1.pdf", # "data/dok2.pdf", # "data/dok3.pdf" # ]) # List documents rag.list_documents() # Query question = "What do you knowabout modality for perceived message perception?" result = rag.query(question) print(f"\nQuestion: {question}") print(f"Answer: {result['answer']}") if result.get("retrieved"): print(f"Retrieved {len(result['retrieved'])} chunks") # print("\nSetup complete! Uncomment the code above to add documents and query.") if __name__ == "__main__": main()