- 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>
563 lines
26 KiB
Python
563 lines
26 KiB
Python
"""
|
|
Local RAG setup with LangChain, Ollama/OpenAI, and FAISS
|
|
Minimal dependencies, simple code
|
|
"""
|
|
import json
|
|
import os
|
|
import re
|
|
import time
|
|
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"
|
|
# 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
|
|
|
|
# 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 Tool, {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 Tool.
|
|
- 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)."""
|
|
|
|
|
|
# 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,
|
|
vectorstore_path="./vectorstore",
|
|
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,
|
|
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', '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
|
|
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, or up front via preload()
|
|
self.last_rerank_seconds = 0.0 # time of the last reranking, for the timing trace
|
|
|
|
# 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(
|
|
chunk_size=chunk_size,
|
|
chunk_overlap=chunk_overlap,
|
|
)
|
|
print(f"Text splitter: chunk_size={chunk_size}, chunk_overlap={chunk_overlap}")
|
|
|
|
# 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(
|
|
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})...")
|
|
# 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(
|
|
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
|
|
_stems = {} # word -> stem; the corpus has far fewer distinct words than word occurrences
|
|
|
|
@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")
|
|
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:
|
|
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):
|
|
"""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:
|
|
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:
|
|
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,
|
|
tool_description=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 and tool_description to
|
|
SEARCH_TOOL_DESCRIPTION (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 "")
|
|
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(
|
|
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=tool_description or 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}")
|
|
|
|
started_total = time.perf_counter()
|
|
answer = None
|
|
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)
|
|
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
|
|
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}
|
|
|
|
|
|
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()
|