pmock/localRAG
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Primary-text index for personas, faster retrieval, OpenAI model options

Browse Source 
- Contemporary and Kafka personas search only "Die Verwandlung" (own index and tool description); per-persona subtitle
- New contemporary portrait and avatar alt text
- ChatGPT: gpt-6.1-sol, reasoning_effort low via Responses API; default in the model menu
- GPT-Realtime-2.1 Mini as third option (realtime_chat.py, WebSocket, text only)
- Load .env for OPENAI_API_KEY
- Reranker in fp16 on GPU with 25 candidates; reranker, embeddings and BM25 shared and built at server start
- BM25 stem cache (first search ~31s -> ~1.7s)
- Timing output per LLM call and search

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Philipp Mock
2026-10-26 07:04:59 +01:00
co-authored by Claude Opus 5.5
parent f68d7a6e25
commit e4b7f0ce34
6 changed files with 243 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
local_rag.py
+100 -23
View File
@@ -5,6 +5,7 @@ Minimal dependencies, simple code
import json
import os
import re
import time
from pathlib import Path
import numpy as np
@@ -49,7 +50,9 @@ RETRIEVAL_K = 5
# "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
# 25 instead of 50 halves the reranker time and was not worse on either Kafka goldset
# (pypdf 2000/400, R@5 human .70 -> .73, synthetic .91 -> .98).
CANDIDATES_K = 25 # 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
@@ -87,6 +90,43 @@ Du hast ein Tool, {SEARCH_TOOL_NAME}, das eine Bibliothek aus Kafkas Texten und
- Nenne bei Aussagen aus den Auszügen die Quelle (Autor, Seite)."""
# Cross-encoders by model name, shared by all LocalRAG instances in a process (one copy in memory)
_RERANKERS = {}
# Embedding models, shared the same way
_EMBEDDINGS = {}
# BM25 keyword indexes by (vector store path, number of chunks), shared by instances on the same index
_BM25_INDEXES = {}
def load_embeddings(model_name):
"""Load an embedding model once per process and return the shared instance."""
if model_name not in _EMBEDDINGS:
print("Loading embeddings model...")
# Normalized vectors: FAISS's L2 distance then ranks exactly like cosine similarity.
_EMBEDDINGS[model_name] = HuggingFaceEmbeddings(
model_name=model_name,
encode_kwargs={"normalize_embeddings": True},
query_encode_kwargs={"normalize_embeddings": True},
)
return _EMBEDDINGS[model_name]
def load_reranker(model_name):
"""Load a cross-encoder once per process and return the shared instance."""
if model_name not in _RERANKERS:
import torch
from sentence_transformers import CrossEncoder
# Half precision on the GPU (Apple MPS or CUDA): about 15-25 % faster, same rankings
# on both Kafka goldsets. On the CPU fp16 is slow, so keep fp32 there.
gpu = torch.backends.mps.is_available() or torch.cuda.is_available()
model_kwargs = {"torch_dtype": torch.float16} if gpu else {}
print(f"Loading reranker {model_name} ({'fp16' if gpu else 'fp32'})...")
reranker = CrossEncoder(model_name, max_length=1024, model_kwargs=model_kwargs)
reranker.predict([("Probe", "Probe")]) # warm-up, so the first real query is not slower
_RERANKERS[model_name] = reranker
return _RERANKERS[model_name]
class LocalRAG:
def __init__(
self,
@@ -94,6 +134,7 @@ class LocalRAG:
llm_provider="ollama",
ollama_model="gemma4:12b",
openai_model="gpt-6-luna",
realtime_model="gpt-realtime-2.1-mini",
ollama_base_url="http://localhost:11434",
pdf_loader=PDF_LOADER,
processed_dir=PROCESSED_DIR,
@@ -103,7 +144,8 @@ class LocalRAG:
chunk_overlap=CHUNK_OVERLAP,
chunk_headers=CHUNK_HEADERS,
):
"""Initialize local RAG system. llm_provider: 'ollama' or 'openai'.
"""Initialize local RAG system. llm_provider: 'ollama', 'openai' or 'realtime'
(OpenAI Realtime API, see realtime_chat.py).
pdf_loader: 'pypdf' or 'docling' (see PDF_LOADER)."""
self.vectorstore_path = vectorstore_path
self.llm_provider = llm_provider
@@ -113,16 +155,11 @@ class LocalRAG:
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
self._reranker = None # loaded on first use, or up front via preload()
self.last_rerank_seconds = 0.0 # time of the last reranking, for the timing trace
# 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},
)
# Embeddings (one shared copy per process)
self.embeddings = load_embeddings(EMBEDDING_MODEL)
# Text splitter (used when adding documents; tune CHUNK_SIZE / CHUNK_OVERLAP above)
self.text_splitter = RecursiveCharacterTextSplitter(
@@ -131,15 +168,23 @@ class LocalRAG:
)
print(f"Text splitter: chunk_size={chunk_size}, chunk_overlap={chunk_overlap}")
# LLM (Ollama or OpenAI)
if llm_provider == "openai":
# LLM (Ollama, OpenAI or OpenAI Realtime)
if llm_provider in ("openai", "realtime"):
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'"
f"OPENAI_API_KEY environment variable is required when llm_provider='{llm_provider}'"
)
if llm_provider == "realtime":
from realtime_chat import ChatOpenAIRealtime
print(f"Using OpenAI Realtime (model: {realtime_model})...")
self.llm = ChatOpenAIRealtime(model=realtime_model, api_key=api_key)
elif llm_provider == "openai":
print(f"Using OpenAI (model: {openai_model})...")
self.llm = ChatOpenAI(model=openai_model, api_key=api_key)
# Responses API: GPT-6 models accept function tools on Chat Completions only with
# reasoning_effort="none"; with "low" the search tool needs the Responses API.
self.llm = ChatOpenAI(model=openai_model, api_key=api_key, reasoning_effort="low",
use_responses_api=True)
else:
print(f"Using Ollama (model: {ollama_model})...")
self.llm = ChatOllama(
@@ -318,6 +363,7 @@ class LocalRAG:
# ---------- retrieval: vector search, keyword search, reranking ----------
_stemmer = None
_stems = {} # word -> stem; the corpus has far fewer distinct words than word occurrences
@classmethod
def _bm25_tokens(cls, text):
@@ -325,13 +371,26 @@ class LocalRAG:
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()))
stems = cls._stems
tokens = []
for word in re.findall(r"[a-zäöüß0-9]{2,}", text.lower()):
stem = stems.get(word)
if stem is None:
stem = stems[word] = cls._stemmer.stemWord(word)
tokens.append(stem)
return tokens
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)
key = (self.vectorstore_path, self.vectorstore.index.ntotal)
if key not in _BM25_INDEXES:
from rank_bm25 import BM25Okapi
started = time.perf_counter()
docs = list(self.vectorstore.docstore._dict.values())
_BM25_INDEXES[key] = (BM25Okapi([self._bm25_tokens(d.page_content) for d in docs]), docs)
print(f"BM25 index for {self.vectorstore_path}: {len(docs)} chunks, "
f"{time.perf_counter() - started:.1f}s")
self._bm25 = _BM25_INDEXES[key]
return self._bm25
def retrieve(self, query, k=None):
@@ -365,14 +424,22 @@ class LocalRAG:
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)
self._reranker = load_reranker(self.reranker_model)
started = time.perf_counter()
relevance = self._reranker.predict([(query, d.page_content) for d, _ in candidates])
self.last_rerank_seconds = time.perf_counter() - started
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 preload(self):
"""Load the reranker and build the BM25 index now instead of on the first search
(e.g. at server start)."""
if self.reranker_model and self._reranker is None:
self._reranker = load_reranker(self.reranker_model)
if self.retrieval_mode != "dense" and self.vectorstore is not None:
self._bm25_index()
def query(self, question, k=None):
"""Query the RAG system (no conversation history). Returns dict with 'answer' and 'retrieved'."""
if k is None:
@@ -407,10 +474,13 @@ class LocalRAG:
def search_kafka_corpus(query: str) -> str:
retrieval_query = query + (f"\n\n{QUERY_ADDITIONAL_INSTRUCTIONS}" if QUERY_ADDITIONAL_INSTRUCTIONS else "")
started = time.perf_counter()
hits = self.retrieve(retrieval_query, k=k)
seconds = time.perf_counter() - started
searches.append(query)
retrieved.extend(self._docs_scores_to_retrieved(hits))
print(f"[RAG] Tool search {query!r}: {len(hits)} chunk(s)")
print(f"[Zeit] Suche {seconds:.2f}s (davon Reranker {self.last_rerank_seconds:.2f}s)")
if not hits:
return "Keine passenden Textstellen gefunden."
return "\n\n".join(
@@ -433,9 +503,13 @@ class LocalRAG:
messages.append(HumanMessage(question))
print(f"[RAG] User question: {question!r}")
started_total = time.perf_counter()
answer = None
for _ in range(MAX_TOOL_ROUNDS):
for round_no in range(1, MAX_TOOL_ROUNDS + 1):
started = time.perf_counter()
reply = llm_with_tools.invoke(messages)
print(f"[Zeit] LLM-Aufruf {round_no}: {time.perf_counter() - started:.2f}s "
f"({'Suche angefordert' if reply.tool_calls else 'Antwort'})")
messages.append(reply)
if not reply.tool_calls:
answer = self._text(reply)
@@ -445,8 +519,11 @@ class LocalRAG:
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
started = time.perf_counter()
answer = self._text(self.llm.invoke(messages + [HumanMessage(
"Beantworte jetzt die Frage mit den gefundenen Auszügen, ohne weitere Suche.")]))
print(f"[Zeit] LLM-Aufruf (Antwort erzwungen): {time.perf_counter() - started:.2f}s")
print(f"[Zeit] Gesamt {time.perf_counter() - started_total:.2f}s")
print(f"[RAG] Searches: {searches or 'keine'}")
return {"answer": answer, "retrieved": retrieved, "searches": searches}