""" 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: 2000 characters of German prose are roughly 550 tokens. # With the reranker, 2000/400 beat 1500/300 on both Kafka goldsets (eval/matrix_2000.out). CHUNK_SIZE = 2000 CHUNK_OVERLAP = 400 # 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. # Needed for Docling (which drops running heads); neutral to slightly harmful for pypdf, # whose pages carry the author in the running head anyway. CHUNK_HEADERS = False # 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" # The corpus search is offered to the LLM as a tool; the LLM decides per message whether # to search (content questions) or to answer directly (greetings, thanks, questions about # the conversation). It also writes the search query itself, using the chat history. SEARCH_TOOL_NAME = "search_kafka_corpus" SEARCH_TOOL_DESCRIPTION = ( "Durchsucht die Kafka-Bibliothek (Kafkas Texte und die Forschungsliteratur) und liefert die " "passendsten Textauszüge mit Quelle und Seite. Die Suchanfrage sollte eine eigenständige, " "präzise Frage oder Stichwortfolge auf Deutsch sein." ) MAX_TOOL_ROUNDS = 3 # searches per message before the LLM must answer # System prompt for every chat turn (German, like the corpus and the questions). SYSTEM_PROMPT = f"""Du beantwortest Fragen zu Franz Kafka, seinen Werken, seinem Leben und der Forschung. Du hast ein Werkzeug, {SEARCH_TOOL_NAME}, das eine Bibliothek aus Kafkas Texten und Forschungsliteratur durchsucht. - Für jede inhaltliche Frage zu Kafka, seinen Werken, Figuren, seiner Biografie, Sprache oder zur Forschung: Rufe zuerst das Werkzeug auf und stütze die Antwort auf die gefundenen Auszüge. Antworte nicht aus dem Gedächtnis. - Bei Begrüßung, Dank, Smalltalk oder Fragen zum Gespräch selbst: Antworte direkt, ohne Werkzeug. - Formuliere die Suchanfrage so, dass sie ohne den Gesprächsverlauf verständlich ist. - Finden sich in den Auszügen keine Belege, sage das offen, statt zu raten. - Nenne bei Aussagen aus den Auszügen die Quelle (Autor, Seite).""" class LocalRAG: def __init__( self, vectorstore_path="./vectorstore", llm_provider="ollama", ollama_model="gemma4:12b", openai_model="gpt-6-luna", 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 _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) @staticmethod def _text(message): """Plain text of an LLM reply (some models return a list of content blocks).""" content = getattr(message, "content", message) if isinstance(content, list): content = "".join(part.get("text", "") if isinstance(part, dict) else str(part) for part in content) return (content or "").strip() def query_with_history(self, question, chat_history=None, k=None, system_prompt=None): """Answer a chat message. The LLM decides whether to search the corpus: the search is offered as a tool (SEARCH_TOOL_NAME), which the LLM calls with its own query, up to MAX_TOOL_ROUNDS times. system_prompt defaults to SYSTEM_PROMPT (see personas.py for others). Returns dict with 'answer', 'retrieved' (chunks the LLM saw, with content, source, page, score) and 'searches' (the queries it used).""" from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage from langchain_core.tools import StructuredTool k = k or RETRIEVAL_K if self.vectorstore is None: return {"answer": "Error: No documents loaded. Please add documents first.", "retrieved": [], "searches": []} retrieved, searches = [], [] def search_kafka_corpus(query: str) -> str: retrieval_query = query + (f"\n\n{QUERY_ADDITIONAL_INSTRUCTIONS}" if QUERY_ADDITIONAL_INSTRUCTIONS else "") hits = self.retrieve(retrieval_query, k=k) searches.append(query) retrieved.extend(self._docs_scores_to_retrieved(hits)) print(f"[RAG] Tool search {query!r}: {len(hits)} chunk(s)") if not hits: return "Keine passenden Textstellen gefunden." return "\n\n".join( f"[{i}] {docling_loader.source_label(Path(d.metadata.get('source', '')).name)}" f"{', S. ' + str(d.metadata['page'] + 1) if d.metadata.get('page') is not None else ''}\n" f"{d.page_content.strip()}" for i, (d, _) in enumerate(hits, start=1)) tool = StructuredTool.from_function(func=search_kafka_corpus, name=SEARCH_TOOL_NAME, description=SEARCH_TOOL_DESCRIPTION) llm_with_tools = self.llm.bind_tools([tool]) messages = [SystemMessage(system_prompt or SYSTEM_PROMPT)] for turn in chat_history or []: content = (turn.get("content") or "").strip() if turn.get("role") == "user": messages.append(HumanMessage(content)) elif turn.get("role") == "assistant": messages.append(AIMessage(content)) messages.append(HumanMessage(question)) print(f"[RAG] User question: {question!r}") answer = None for _ in range(MAX_TOOL_ROUNDS): reply = llm_with_tools.invoke(messages) messages.append(reply) if not reply.tool_calls: answer = self._text(reply) break for n, call in enumerate(reply.tool_calls): result = tool.invoke(call["args"]) messages.append(ToolMessage(result, tool_call_id=call.get("id") or f"call_{len(searches)}_{n}")) if answer is None: # Search budget used up: answer with what was found, without further tool calls answer = self._text(self.llm.invoke(messages + [HumanMessage( "Beantworte jetzt die Frage mit den gefundenen Auszügen, ohne weitere Suche.")])) print(f"[RAG] Searches: {searches or 'keine'}") return {"answer": answer, "retrieved": retrieved, "searches": searches} def main(): """Example usage""" print("=" * 60) print("Local RAG with LangChain, Ollama/OpenAI, and FAISS") print("=" * 60) # Initialize rag = LocalRAG(ollama_model="gemma4:12b") # 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()