"""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 )