455 lines
18 KiB
Python
455 lines
18 KiB
Python
"""Hybrid-Retrieval: BM25 (FTS5) + Dense (bge-m3) -> RRF-Fusion,
|
|
milde Stand-Aktualitätsgewichtung und kontrollierte cross_ref-Erweiterung.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import sqlite3
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
from .config import Config
|
|
from .normalize import fts_query
|
|
from .ollama_client import OllamaClient
|
|
from .query_planner import SubQuery, decision_support_plan
|
|
|
|
|
|
@dataclass
|
|
class ChunkResult:
|
|
chunk_id: int
|
|
entry_id: str
|
|
section: str
|
|
text: str
|
|
title: str
|
|
stand: str
|
|
work: str
|
|
chapter: str
|
|
topic: str
|
|
tags: list = field(default_factory=list)
|
|
legal_bases: list = field(default_factory=list)
|
|
cross_refs: list = field(default_factory=list)
|
|
batch: int = 0
|
|
score: float = 0.0
|
|
source: str = "fused" # bm25 | dense | fused | cross_ref
|
|
|
|
|
|
def _section_priority(section: str) -> int:
|
|
"""Kontext-Sektionen priorisieren: Inhalt vor Navigation.
|
|
|
|
„Verweise“-Sektionen sind Navigationslisten (KB-IDs) — sie tragen
|
|
Retrieval-Signal (Stichworte), sind aber als Kontextblock wertlos und
|
|
provozieren Fehlverweigerungen. BM25-Längennormalisierung rangiert sie
|
|
bevorzugt, daher wird pro Eintrag bewusst die beste Inhaltssektion
|
|
gewählt (Fix 2026-09-14, q-008).
|
|
"""
|
|
s = (section or "").casefold()
|
|
if s.startswith("zusammenfassung"):
|
|
return 0
|
|
if s.startswith("kernwerte"):
|
|
return 1
|
|
if s.startswith("rechtsgrundlagen"):
|
|
return 2
|
|
if s.startswith("payroll"):
|
|
return 3
|
|
if s.startswith("verweise"):
|
|
return 5
|
|
return 4
|
|
|
|
|
|
class Retriever:
|
|
def __init__(self, cfg: Config, db_path: str | None = None, client=None):
|
|
self.cfg = cfg
|
|
self.db_path = str(db_path or cfg.db_path)
|
|
if not Path(self.db_path).is_file():
|
|
raise RuntimeError(
|
|
f"index fehlt ({self.db_path}) — zuerst 'python -m agent.cli ingest' ausführen"
|
|
)
|
|
self._con = sqlite3.connect(self.db_path)
|
|
self._con.row_factory = sqlite3.Row
|
|
self._client = client
|
|
self._owns_client = client is None
|
|
self._mat: np.ndarray | None = None
|
|
self._mat_chunk_ids: list[int] | None = None
|
|
self._scope_cache: dict[str, set[int]] = {}
|
|
row = self._con.execute("SELECT MIN(stand), MAX(stand) FROM chunks").fetchone()
|
|
self._stand_min = int((row[0] or "2026-01").replace("-", ""))
|
|
self._stand_max = int((row[1] or "2026-01").replace("-", ""))
|
|
|
|
def close(self) -> None:
|
|
self._con.close()
|
|
if self._owns_client and self._client is not None:
|
|
self._client.close()
|
|
|
|
# -- Index-Kennzahlen ---------------------------------------------------
|
|
|
|
def stats(self) -> dict:
|
|
n_chunks = self._con.execute("SELECT COUNT(*) FROM chunks").fetchone()[0]
|
|
n_entries = self._con.execute(
|
|
"SELECT COUNT(DISTINCT entry_id) FROM chunks"
|
|
).fetchone()[0]
|
|
n_vec = self._con.execute(
|
|
"SELECT COUNT(*) FROM vectors WHERE model = ?",
|
|
(self.cfg.embed_model,),
|
|
).fetchone()[0]
|
|
meta = dict(self._con.execute("SELECT key, value FROM meta").fetchall())
|
|
return {
|
|
"n_entries": n_entries,
|
|
"n_chunks": n_chunks,
|
|
"n_vectors": n_vec,
|
|
"dense_available": n_vec > 0 and not self.cfg.embed_off,
|
|
"stand_min": str(self._stand_min),
|
|
"stand_max": str(self._stand_max),
|
|
"built_at": meta.get("built_at"),
|
|
}
|
|
|
|
# -- Einzelverfahren ----------------------------------------------------
|
|
|
|
def _scope_ids(self, scope: str | None) -> set[int] | None:
|
|
"""Chunk-ID-Menge je Sub-Query-Scope (Cache je Instanz): 'gesetz' =
|
|
alles außer Branchen-KV, 'kv' = nur Branchen-KV."""
|
|
if scope is None:
|
|
return None
|
|
if scope not in self._scope_cache:
|
|
if scope == "kv":
|
|
sql = "SELECT chunk_id FROM chunks WHERE entry_id LIKE 'kv-%'"
|
|
else: # gesetz
|
|
sql = "SELECT chunk_id FROM chunks WHERE entry_id NOT LIKE 'kv-%'"
|
|
self._scope_cache[scope] = {
|
|
row[0] for row in self._con.execute(sql).fetchall()
|
|
}
|
|
return self._scope_cache[scope]
|
|
|
|
def _bm25(self, question: str, limit: int, scope: str | None = None) -> dict[int, float]:
|
|
q = fts_query(question)
|
|
if not q:
|
|
return {}
|
|
fetch = limit * 4 if scope else limit
|
|
rows = self._con.execute(
|
|
"SELECT rowid, bm25(chunks_fts) AS rank FROM chunks_fts "
|
|
"WHERE chunks_fts MATCH ? ORDER BY rank LIMIT ?",
|
|
(q, fetch),
|
|
).fetchall()
|
|
scope_ids = self._scope_ids(scope)
|
|
if scope_ids is not None:
|
|
rows = [r for r in rows if r["rowid"] in scope_ids][:limit]
|
|
# bm25(): kleinere Werte = besser -> negieren für "größer = besser"
|
|
return {r["rowid"]: -float(r["rank"]) for r in rows}
|
|
|
|
def _dense(self, question: str, limit: int, scope: str | None = None) -> dict[int, float]:
|
|
if self.cfg.embed_off:
|
|
return {}
|
|
self._ensure_matrix()
|
|
if self._mat is None or len(self._mat) == 0:
|
|
return {}
|
|
if self._client is None:
|
|
self._client = OllamaClient(
|
|
self.cfg.ollama_url,
|
|
embed_timeout_s=self.cfg.embed_timeout_s,
|
|
chat_timeout_s=self.cfg.chat_timeout_s,
|
|
)
|
|
try:
|
|
qvec = np.asarray(
|
|
self._client.embed(self.cfg.embed_model, [question])[0],
|
|
dtype=np.float32,
|
|
)
|
|
except Exception:
|
|
return {} # Ollama nicht erreichbar -> BM25-only weiter
|
|
qn = np.linalg.norm(qvec)
|
|
if qn == 0:
|
|
return {}
|
|
sims = self._mat_norm @ (qvec / qn)
|
|
order = np.argsort(-sims)
|
|
scope_ids = self._scope_ids(scope)
|
|
out: dict[int, float] = {}
|
|
for i in order:
|
|
if len(out) >= limit:
|
|
break
|
|
cid = self._mat_chunk_ids[i]
|
|
if scope_ids is not None and cid not in scope_ids:
|
|
continue
|
|
out[cid] = float(sims[i])
|
|
return out
|
|
|
|
def _ensure_matrix(self) -> None:
|
|
if self._mat is not None:
|
|
return
|
|
rows = self._con.execute(
|
|
"SELECT c.chunk_id, v.vec, v.dim FROM chunks c "
|
|
"JOIN vectors v ON v.content_hash = c.content_hash AND v.model = ?",
|
|
(self.cfg.embed_model,),
|
|
).fetchall()
|
|
if not rows:
|
|
self._mat = np.zeros((0, 1), dtype=np.float32)
|
|
self._mat_chunk_ids = []
|
|
return
|
|
ids = [r[0] for r in rows]
|
|
mat = np.vstack(
|
|
[np.frombuffer(r[1], dtype=np.float32) for r in rows]
|
|
)
|
|
norms = np.linalg.norm(mat, axis=1, keepdims=True)
|
|
self._mat = mat
|
|
self._mat_norm = mat / np.where(norms == 0, 1.0, norms)
|
|
self._mat_chunk_ids = ids
|
|
|
|
# -- Metadaten & Fusion -------------------------------------------------
|
|
|
|
def _stand_factor(self, stand: str) -> float:
|
|
if self._stand_max <= self._stand_min:
|
|
return 0.0
|
|
try:
|
|
s = int(stand.replace("-", ""))
|
|
except (ValueError, AttributeError):
|
|
return 0.0
|
|
f = (s - self._stand_min) / (self._stand_max - self._stand_min)
|
|
return min(1.0, max(0.0, f))
|
|
|
|
def _fetch_chunks(self, chunk_ids: list[int]) -> dict[int, sqlite3.Row]:
|
|
out: dict[int, sqlite3.Row] = {}
|
|
for i in range(0, len(chunk_ids), 500):
|
|
part = chunk_ids[i:i + 500]
|
|
qm = ",".join("?" * len(part))
|
|
for r in self._con.execute(
|
|
f"SELECT * FROM chunks WHERE chunk_id IN ({qm})", part
|
|
).fetchall():
|
|
out[r["chunk_id"]] = r
|
|
return out
|
|
|
|
def _row_to_result(self, row: sqlite3.Row, score: float, source: str) -> ChunkResult:
|
|
return ChunkResult(
|
|
chunk_id=row["chunk_id"],
|
|
entry_id=row["entry_id"],
|
|
section=row["section"],
|
|
text=row["text"],
|
|
title=row["title"],
|
|
stand=row["stand"],
|
|
work=row["work"],
|
|
chapter=row["chapter"],
|
|
topic=row["topic"],
|
|
tags=json.loads(row["tags"]),
|
|
legal_bases=json.loads(row["legal_bases"]),
|
|
cross_refs=json.loads(row["cross_refs"]),
|
|
batch=row["batch"],
|
|
score=score,
|
|
source=source,
|
|
)
|
|
|
|
def _representative_chunk(
|
|
self, entry_id: str, ranked: list[ChunkResult]
|
|
) -> ChunkResult:
|
|
"""Beste Inhaltssektion des Eintrags als Kontextblock.
|
|
|
|
Bevorzugt die rangierte (gefundene) Sektion mit bester Priorität;
|
|
traf der Eintrag nur über „Verweise“, wird seine beste Inhalts-
|
|
sektion aus dem Index nachgeladen (source="section-swap").
|
|
"""
|
|
content = [c for c in ranked if _section_priority(c.section) < 5]
|
|
if content:
|
|
return min(content, key=lambda c: (_section_priority(c.section), -c.score))
|
|
rows = self._con.execute(
|
|
"SELECT * FROM chunks WHERE entry_id = ? AND section NOT LIKE 'Verweise%' "
|
|
"ORDER BY CASE WHEN section LIKE 'Zusammenfassung%' THEN 0 "
|
|
"WHEN section LIKE 'Kernwerte%' THEN 1 ELSE 2 END, chunk_id LIMIT 1",
|
|
(entry_id,),
|
|
).fetchall()
|
|
if rows:
|
|
return self._row_to_result(rows[0], 0.0, "section-swap")
|
|
return ranked[0] # Eintrag hat nur Verweise-Sektionen
|
|
|
|
def _best_chunk_of_entry(self, entry_id: str) -> ChunkResult | None:
|
|
rows = self._con.execute(
|
|
"SELECT * FROM chunks WHERE entry_id = ? "
|
|
"ORDER BY CASE WHEN section LIKE 'Zusammenfassung%' THEN 0 ELSE 1 END, "
|
|
"chunk_id LIMIT 1",
|
|
(entry_id,),
|
|
).fetchall()
|
|
if not rows:
|
|
return None
|
|
return self._row_to_result(rows[0], 0.0, "cross_ref")
|
|
|
|
# -- öffentliche Suche --------------------------------------------------
|
|
|
|
def _search_one(
|
|
self, sq: SubQuery, pool: int
|
|
) -> tuple[dict[int, float], dict[int, float], dict[int, float]]:
|
|
"""BM25+Dense für EINE Sub-Query inkl. Scope-Filter (mit Fallback auf
|
|
unscoped, wenn die gefilterte Suche leer bleibt) und Query-RRF."""
|
|
try:
|
|
bm = self._bm25(sq.text, pool, scope=sq.scope)
|
|
except Exception:
|
|
bm = {}
|
|
try:
|
|
dn = self._dense(sq.text, pool, scope=sq.scope)
|
|
except Exception:
|
|
dn = {}
|
|
if sq.scope and not bm and not dn:
|
|
try:
|
|
bm = self._bm25(sq.text, pool)
|
|
dn = self._dense(sq.text, pool)
|
|
except Exception:
|
|
pass
|
|
fused_q: dict[int, float] = {}
|
|
for ranking, weight in ((bm, 1.0), (dn, self.cfg.dense_weight)):
|
|
ordered = sorted(ranking.items(), key=lambda kv: -kv[1])
|
|
for rank, (cid, _) in enumerate(ordered):
|
|
fused_q[cid] = fused_q.get(cid, 0.0) + weight / (
|
|
self.cfg.rrf_k + rank
|
|
)
|
|
return bm, dn, fused_q
|
|
|
|
def _fuse_queries(
|
|
self, sub_queries: list, pool: int
|
|
) -> tuple[dict[int, float], set[int], set[int]]:
|
|
"""RRF-Fusion über BM25+Dense je Sub-Query (Beiträge summieren);
|
|
liefert (fused, bm25-Chunk-IDs, dense-Chunk-IDs) für die Quelle-
|
|
Markierung. Bei einer einzelnen Sub-Query identisch zum bisherigen
|
|
Verhalten."""
|
|
fused: dict[int, float] = {}
|
|
bm_all: set[int] = set()
|
|
dn_all: set[int] = set()
|
|
for sq in sub_queries:
|
|
bm, dn, fused_q = self._search_one(sq, pool)
|
|
bm_all.update(bm)
|
|
dn_all.update(dn)
|
|
for cid, score in fused_q.items():
|
|
fused[cid] = fused.get(cid, 0.0) + score
|
|
return fused, bm_all, dn_all
|
|
|
|
def _context_from_fused(
|
|
self,
|
|
fused: dict[int, float],
|
|
n: int,
|
|
bm_all: set[int],
|
|
dn_all: set[int],
|
|
stand_years: set[str] | None = None,
|
|
reserved: list[str] | None = None,
|
|
) -> list[ChunkResult]:
|
|
"""Fusion -> Score (+Recency +Temporal-Boost) -> Entry-Dedup ->
|
|
Vertreter-Chunk -> cross_ref-Erweiterung."""
|
|
if not fused:
|
|
return []
|
|
rows = self._fetch_chunks(list(fused))
|
|
results: list[ChunkResult] = []
|
|
for cid, score in fused.items():
|
|
if cid not in rows:
|
|
continue
|
|
source = (
|
|
"fused" if (cid in bm_all and cid in dn_all)
|
|
else ("bm25" if cid in bm_all else "dense")
|
|
)
|
|
score += self.cfg.recency_boost * self._stand_factor(rows[cid]["stand"])
|
|
if stand_years and rows[cid]["entry_id"].startswith("kv-"):
|
|
if rows[cid]["stand"][:4] in stand_years:
|
|
score += self.cfg.temporal_boost
|
|
results.append(self._row_to_result(rows[cid], score, source))
|
|
results.sort(key=lambda r: -r.score)
|
|
|
|
# Bester Chunk je Eintrag -> Kontext (Entry-Level-Dedup).
|
|
# Der Vertreter-Chunk ist die beste Inhaltssektion des Eintrags,
|
|
# nicht die Rangfolge-Beste (vgl. _section_priority).
|
|
# Per-Query-Slots (Multi-Query): reservierte Einträge zuerst —
|
|
# jeder Sub-Query-Aspekt hält garantiert seine Top-Quelle(n).
|
|
main: list[ChunkResult] = []
|
|
by_entry: dict[str, list[ChunkResult]] = {}
|
|
for r in results:
|
|
by_entry.setdefault(r.entry_id, []).append(r)
|
|
seen: set[str] = set()
|
|
reserved = reserved or []
|
|
for entry_id in reserved:
|
|
if len(main) >= n or entry_id not in by_entry:
|
|
continue
|
|
seen.add(entry_id)
|
|
main.append(self._representative_chunk(entry_id, by_entry[entry_id]))
|
|
for entry_id, chunks in by_entry.items():
|
|
if len(main) >= n:
|
|
break
|
|
if entry_id in seen:
|
|
continue
|
|
seen.add(entry_id)
|
|
main.append(self._representative_chunk(entry_id, chunks))
|
|
|
|
# cross_ref-Erweiterung (kontrolliert, markiert, begrenzt)
|
|
extra: list[ChunkResult] = []
|
|
budget = self.cfg.cross_ref_max_extra
|
|
for r in main[: self.cfg.cross_ref_expand]:
|
|
for ref in r.cross_refs:
|
|
if budget <= 0:
|
|
break
|
|
if ref in seen:
|
|
continue
|
|
er = self._best_chunk_of_entry(ref)
|
|
if er is not None:
|
|
extra.append(er)
|
|
seen.add(ref)
|
|
budget -= 1
|
|
return main + extra
|
|
|
|
def search(self, question: str, n_entries: int | None = None) -> list[ChunkResult]:
|
|
"""Liefert die Top-Kontextblöcke (Hauptretrieval + cross_ref-Erweiterung).
|
|
|
|
Gestaltungsfragen zu zusätzlichen Arbeitnehmerleistungen werden
|
|
deterministisch in Direktzahlung und Alternativen zerlegt. Das gilt
|
|
auch für den Offline-Retrieval-Eval, der keinen LLM-Planer aufruft.
|
|
"""
|
|
n = n_entries or self.cfg.context_blocks
|
|
latest_year = str(self._stand_max)[:4]
|
|
deterministic = decision_support_plan(question, default_year=latest_year)
|
|
if deterministic:
|
|
sub_queries, _qtype = deterministic
|
|
return self.search_multi(sub_queries, n_entries=n)
|
|
fused, bm_all, dn_all = self._fuse_queries(
|
|
[SubQuery(text=question)], self.cfg.candidate_pool
|
|
)
|
|
return self._context_from_fused(fused, n, bm_all, dn_all)
|
|
|
|
def search_multi(
|
|
self, sub_queries: list, n_entries: int | None = None
|
|
) -> list[ChunkResult]:
|
|
"""Multi-Query-Retrieval (Stufe 1): je Sub-Query BM25+Dense, RRF-
|
|
Beiträge summieren. Per-Query-Slots: jede Sub-Query sichert ihre
|
|
Top-Einträge im Kontext (Multi-Hop: jeder Frageaspekt kommt mit
|
|
seiner besten Quelle hinein — sonst dominieren Einträge, die in
|
|
mehreren Sub-Queries mittelgut matchen). Temporal-Intent: kv-
|
|
Einträge im gefragten Geltungsjahr erhalten temporal_boost."""
|
|
n = n_entries or self.cfg.context_blocks
|
|
if (
|
|
len(sub_queries) == 1
|
|
and sub_queries[0].scope is None
|
|
and sub_queries[0].stand_year is None
|
|
):
|
|
return self.search(sub_queries[0].text, n_entries=n)
|
|
pool = self.cfg.candidate_pool
|
|
fused_total: dict[int, float] = {}
|
|
bm_all: set[int] = set()
|
|
dn_all: set[int] = set()
|
|
per_query_entries: list[list[str]] = []
|
|
for sq in sub_queries:
|
|
bm, dn, fused_q = self._search_one(sq, pool)
|
|
bm_all.update(bm)
|
|
dn_all.update(dn)
|
|
for cid, score in fused_q.items():
|
|
fused_total[cid] = fused_total.get(cid, 0.0) + score
|
|
rows_q = self._fetch_chunks(list(fused_q))
|
|
entry_best: dict[str, float] = {}
|
|
for cid, score in fused_q.items():
|
|
row = rows_q.get(cid)
|
|
if row is not None:
|
|
eid = row["entry_id"]
|
|
entry_best[eid] = max(entry_best.get(eid, 0.0), score)
|
|
per_query_entries.append(
|
|
[e for e, _ in sorted(entry_best.items(), key=lambda kv: -kv[1])]
|
|
)
|
|
reserved: list[str] = []
|
|
for entries in per_query_entries:
|
|
taken = 0
|
|
for eid in entries:
|
|
if taken >= self.cfg.per_query_slots:
|
|
break
|
|
if eid not in reserved:
|
|
reserved.append(eid)
|
|
taken += 1
|
|
years = {sq.stand_year for sq in sub_queries if sq.stand_year}
|
|
return self._context_from_fused(
|
|
fused_total, n, bm_all, dn_all, stand_years=years, reserved=reserved
|
|
) |