pmock/localRAG
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
30fd55da21e475fdb4ea0c9d8b35e387febab2be
localRAG/local_rag.py
T
Philipp Mock 30fd55da21 - eval mit gpt-6-luna 
- new chat interface with personas
2026-10-08 10:06:32 +02:00

484 lines
22 KiB
Python
Raw Blame History

"""
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()
Reference in New Issue View Git Blame Copy Permalink