igel-qe-core 1.0.0

package/knowledge_base/retrieval/knowledge_graph.py

@@ -0,0 +1,267 @@
+"""GraphRAG helpers: query expansion over knowledge_graph_nodes / knowledge_graph_edges."""
+
+from __future__ import annotations
+
+import logging
+import re
+from dataclasses import dataclass
+
+from knowledge_base.db.client import get_conn
+from knowledge_base.db.kg_introspect import get_kg_layout
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class KGContext:
+    """Structured graph context fused into prompts."""
+
+    text_block: str
+    node_ids: list[str]
+    node_count: int
+    edge_count: int
+
+
+def classify_query_style(query: str) -> str:
+    if re.search(r"app\.[a-z0-9_.]+", query, re.IGNORECASE):
+        return "exact"
+    if re.search(r"\b(step|workflow|scenario|procedure)\b", query, re.IGNORECASE):
+        return "workflow"
+    return "semantic"
+
+
+def _q(col: str) -> str:
+    return '"' + col.replace('"', '""') + '"'
+
+
+def _search_terms(query: str) -> list[str]:
+    q = query.strip()
+    keys = re.findall(r"app\.[a-z0-9_.]+", q, flags=re.IGNORECASE)
+    tokens = [t for t in re.split(r"[^\w\-./]+", q) if len(t) >= 3]
+    out: list[str] = []
+    for x in keys + tokens[:12]:
+        lx = x.lower()
+        if lx not in {y.lower() for y in out}:
+            out.append(x)
+        if len(out) >= 16:
+            break
+    return out
+
+
+def _find_matching_node_ids(
+    *,
+    lay,
+    terms: list[str],
+    seed_parent_ids: list[str] | None,
+    limit: int,
+) -> list[str]:
+    if not lay.node_id or not lay.value_col:
+        return []
+
+    nid = _q(lay.node_id)
+    vcol = _q(lay.value_col)
+    tname = _q(lay.nodes_table)
+
+    clauses: list[str] = []
+    params: list[object] = []
+
+    pat_clauses: list[str] = []
+    for t in terms:
+        esc = "%" + t.replace("%", r"\%").replace("_", r"\_") + "%"
+        pat_clauses.append(f"{vcol} ILIKE %s")
+        params.append(esc)
+
+    if lay.normalized_col:
+        ncol = _q(lay.normalized_col)
+        for t in terms:
+            esc = "%" + t.lower() + "%"
+            pat_clauses.append(f"{ncol} ILIKE %s")
+            params.append(esc)
+
+    if pat_clauses:
+        clauses.append("(" + " OR ".join(pat_clauses) + ")")
+
+    if lay.parent_chunk_col and seed_parent_ids:
+        pcol = _q(lay.parent_chunk_col)
+        clauses.append(f"{pcol}::uuid = ANY(%s::uuid[])")
+        params.append(seed_parent_ids[:512])
+
+    if not clauses:
+        return []
+
+    where_sql = "(" + (" OR ".join(clauses)) + ")" if len(clauses) > 1 else clauses[0]
+    sql = f"SELECT {nid} FROM {tname} WHERE {where_sql} LIMIT %s"
+    params.append(limit)
+
+    with get_conn() as conn:
+        with conn.cursor() as cur:
+            cur.execute(sql, tuple(params))
+            return [str(r[0]) for r in cur.fetchall()]
+
+
+def _expand_neighbors(
+    layout,
+    frontier: list[str],
+    *,
+    exclude: set[str],
+) -> tuple[list[str], int]:
+    if (
+        not frontier
+        or not getattr(layout, "source_col", None)
+        or not getattr(layout, "target_col", None)
+    ):
+        return [], 0
+
+    src_q = _q(layout.source_col)
+    tgt_q = _q(layout.target_col)
+    tbl = _q(layout.edges_table)
+
+    sql = (
+        f"SELECT {src_q}::text, {tgt_q}::text FROM {tbl} WHERE "
+        f"{src_q}::uuid = ANY(%s::uuid[]) OR {tgt_q}::uuid = ANY(%s::uuid[])"
+    )
+    batch = frontier[:512]
+    with get_conn() as conn:
+        with conn.cursor() as cur:
+            cur.execute(sql, (batch, batch))
+            rows = cur.fetchall()
+
+    new_neighbors: list[str] = []
+    for a, b in rows:
+        for nid in (a, b):
+            if nid and nid not in exclude:
+                new_neighbors.append(str(nid))
+                exclude.add(str(nid))
+    return new_neighbors, len(rows)
+
+
+def retrieve_kg_context(
+    query: str,
+    max_hops: int = 2,
+    max_nodes: int = 48,
+    seed_parent_chunk_ids: list[str] | None = None,
+    feature_area: str | None = None,
+) -> KGContext:
+    """
+    Expand around nodes matching query terms (+ optional anchor parent chunks).
+    Routes to Neo4j when KG_BACKEND=neo4j, else PostgreSQL.
+    """
+    from knowledge_base.config import cfg
+
+    if cfg.KG_BACKEND == "neo4j":
+        from knowledge_base.graph_db.kg_traversal import retrieve_kg_context_neo4j
+
+        return retrieve_kg_context_neo4j(
+            query,
+            max_hops=max_hops,
+            max_nodes=max_nodes,
+            feature_area=feature_area,
+            seed_parent_chunk_ids=seed_parent_chunk_ids,
+        )
+
+    lay = get_kg_layout()
+    if lay is None or not lay.node_id or not lay.value_col:
+        return KGContext(text_block="", node_ids=[], node_count=0, edge_count=0)
+
+    hops = max(0, min(int(max_hops), 3))
+    terms = _search_terms(query)
+    seeds_parent = seed_parent_chunk_ids or []
+
+    matched = _find_matching_node_ids(
+        lay=lay,
+        terms=terms,
+        seed_parent_ids=seeds_parent or None,
+        limit=max_nodes * 4,
+    )
+
+    expanded_ids = list(dict.fromkeys(matched))
+    exclude = set(expanded_ids)
+    total_edges = 0
+
+    can_expand = (
+        hops > 0
+        and bool(lay.edges_table)
+        and lay.source_col
+        and lay.target_col
+    )
+
+    frontier = expanded_ids[:]
+    for _ in range(hops):
+        if not can_expand or not frontier:
+            break
+        new_nodes, ec = _expand_neighbors(lay, frontier, exclude=exclude)
+        total_edges += ec
+        frontier = []
+        for nid in new_nodes:
+            if nid not in expanded_ids and len(expanded_ids) < max_nodes * 4:
+                expanded_ids.append(nid)
+                frontier.append(nid)
+            if len(expanded_ids) >= max_nodes * 4:
+                break
+
+    capped = expanded_ids[: max_nodes * 2]
+    text_block = _format_nodes_block(lay, capped)
+
+    return KGContext(
+        text_block=text_block,
+        node_ids=capped,
+        node_count=len(expanded_ids),
+        edge_count=total_edges,
+    )
+
+
+def _format_nodes_block(layout, node_ids: list[str]) -> str:
+    """Format KG nodes grouped by type so the LLM can map them to executable steps."""
+    if not node_ids or not layout.node_id or not layout.value_col:
+        return ""
+
+    nid = _q(layout.node_id)
+    vcol = _q(layout.value_col)
+    tname = _q(layout.nodes_table)
+
+    select_list = [vcol]
+    type_col = layout.node_type_col
+    if type_col:
+        select_list.append(_q(type_col))
+
+    sql = (
+        "SELECT DISTINCT " + ", ".join(select_list) + " FROM " + tname + " WHERE "
+        f"{nid}::uuid = ANY(%s::uuid[])"
+    )
+    with get_conn() as conn:
+        with conn.cursor() as cur:
+            cur.execute(sql, (node_ids[:200],))
+            rows = cur.fetchall()
+
+    if not rows:
+        return ""
+
+    grouped: dict[str, list[str]] = {}
+    for row in rows:
+        value = (str(row[0]).strip() if row[0] else "")
+        ntype = (str(row[1]).strip() if len(row) > 1 and row[1] else "entity")
+        if value:
+            grouped.setdefault(ntype, []).append(value)
+
+    _TYPE_HEADINGS = {
+        "registry_key": "Registry keys (use with `get`/`set` on device shell)",
+        "ui_path": "TC Setup / UI navigation paths (exact menu paths for configuration)",
+        "feature": "Features & sections (cross-document relationships)",
+    }
+
+    lines = ["—— Knowledge graph entities (use these in test steps for accuracy) ——"]
+    for ntype in ("ui_path", "registry_key", "feature"):
+        items = grouped.pop(ntype, [])
+        if not items:
+            continue
+        heading = _TYPE_HEADINGS.get(ntype, ntype)
+        lines.append(f"\n[{heading}]")
+        for val in items[:30]:
+            lines.append(f"  • {val}")
+
+    for ntype, items in sorted(grouped.items()):
+        lines.append(f"\n[{ntype}]")
+        for val in items[:20]:
+            lines.append(f"  • {val}")
+
+    return "\n".join(lines)

package/knowledge_base/retrieval/reranker.py

@@ -0,0 +1,308 @@
+"""
+Pluggable cross-encoder reranking for retrieval (after RRF, before budget assembly).
+
+Backends:
+  local  — sentence-transformers CrossEncoder (OSS, CPU-optimized by default)
+  cohere — Cohere Rerank API (optional hosted)
+  none   — passthrough (scores inferred from rank order)
+
+CPU defaults (192.168.204.65 server — no GPU):
+  Model:      cross-encoder/ms-marco-MiniLM-L-6-v2  (~80MB, ~2-5s for 16 passages)
+  top_n:      16
+  max_length: 256 tokens per query-passage pair
+  batch_size: 8
+
+Multilingual OSS alternative (slower on CPU, ~30-60s):
+  RERANKER_MODEL=BAAI/bge-reranker-base
+"""
+from __future__ import annotations
+
+import logging
+import os
+from dataclasses import dataclass, field
+from time import perf_counter
+
+from knowledge_base.config import cfg
+
+logger = logging.getLogger(__name__)
+
+_CROSS_ENCODER = None
+_CROSS_ENCODER_KEY: str | None = None
+_DEPS_PREPARED = False
+
+
+def _prepare_reranker_environment() -> None:
+    """Fix dependency conflicts and tune PyTorch for CPU inference."""
+    global _DEPS_PREPARED
+    if _DEPS_PREPARED:
+        return
+    try:
+        import markupsafe
+
+        if not hasattr(markupsafe, "soft_unicode"):
+            markupsafe.soft_unicode = getattr(markupsafe, "soft_str", str)
+    except ImportError:
+        pass
+
+    device = _resolve_device()
+    if device == "cpu":
+        threads = max(1, int(getattr(cfg, "RERANKER_TORCH_THREADS", 4)))
+        os.environ.setdefault("OMP_NUM_THREADS", str(threads))
+        os.environ.setdefault("MKL_NUM_THREADS", str(threads))
+        try:
+            import torch
+
+            torch.set_num_threads(threads)
+            torch.set_grad_enabled(False)
+        except ImportError:
+            pass
+
+    _DEPS_PREPARED = True
+
+
+def _resolve_device() -> str:
+    explicit = (getattr(cfg, "RERANKER_DEVICE", None) or "cpu").strip().lower()
+    if explicit not in ("cpu", "cuda", "mps"):
+        explicit = "cpu"
+    if explicit == "cuda":
+        try:
+            import torch
+
+            if not torch.cuda.is_available():
+                logger.info("RERANKER_DEVICE=cuda requested but no GPU — using cpu")
+                return "cpu"
+        except ImportError:
+            return "cpu"
+    return explicit
+
+
+def _load_cross_encoder(model_name: str):
+    """Load sentence-transformers CrossEncoder with CPU-optimized settings."""
+    global _CROSS_ENCODER, _CROSS_ENCODER_KEY
+    _prepare_reranker_environment()
+
+    device = _resolve_device()
+    max_length = max(64, int(getattr(cfg, "RERANKER_MAX_LENGTH", 256)))
+    cache_key = f"{model_name}|{device}|{max_length}"
+
+    if _CROSS_ENCODER is not None and _CROSS_ENCODER_KEY == cache_key:
+        return _CROSS_ENCODER
+
+    try:
+        from sentence_transformers import CrossEncoder
+    except ImportError as e:
+        raise ImportError(
+            "sentence-transformers not available. Run: "
+            "bash knowledge_base/deploy/install_reranker.sh"
+        ) from e
+
+    logger.info(
+        "Loading OSS CrossEncoder: model=%s device=%s max_length=%d",
+        model_name,
+        device,
+        max_length,
+    )
+    t0 = perf_counter()
+    _CROSS_ENCODER = CrossEncoder(model_name, max_length=max_length, device=device)
+    _CROSS_ENCODER_KEY = cache_key
+    logger.info("Reranker ready in %.1fs", perf_counter() - t0)
+    return _CROSS_ENCODER
+
+
+@dataclass
+class RerankItem:
+    child_chunk_id: str
+    passage: str
+    parent_passage: str = ""
+    section_title: str = ""
+    source_url: str = ""
+    metadata: dict = field(default_factory=dict)
+
+
+def build_rerank_passage(item: RerankItem) -> str:
+    """Build a compact passage for cross-encoder input (CPU-friendly length)."""
+    passage_limit = int(getattr(cfg, "RERANKER_PASSAGE_CHARS", 2000))
+    parent_limit = int(getattr(cfg, "RERANKER_PARENT_CHARS", 1200))
+
+    parts: list[str] = []
+    if item.section_title:
+        parts.append(f"Section: {item.section_title[:200]}")
+    if item.parent_passage and item.parent_passage.strip() != item.passage.strip():
+        parts.append(f"Context:\n{item.parent_passage[:parent_limit]}")
+    parts.append(f"Passage:\n{item.passage[:passage_limit]}")
+    combined = "\n\n".join(parts)
+    return combined[: passage_limit + parent_limit + 250]
+
+
+def rerank_passages(query: str, items: list[RerankItem]) -> list[tuple[str, float]]:
+    """
+    Return list of (child_chunk_id, score) sorted by descending relevance.
+    Scores normalized roughly to [0, 1] within the batch where applicable.
+    """
+    if not items:
+        return []
+    backend = (cfg.RERANKER_BACKEND or "none").lower().strip()
+
+    if backend == "none":
+        return _score_by_rank_fallback(items)
+
+    if backend == "cohere":
+        return _rerank_cohere(query, items)
+
+    if backend == "local":
+        return _rerank_local(query, items)
+
+    logger.warning("Unknown RERANKER_BACKEND=%r; using rank fallback", backend)
+    return _score_by_rank_fallback(items)
+
+
+def apply_rerank_threshold(
+    ranked: list[tuple[str, float]],
+    threshold: float | None = None,
+) -> list[tuple[str, float]]:
+    """Filter reranked results below RERANKER_SCORE_THRESHOLD."""
+    cutoff = threshold if threshold is not None else getattr(cfg, "RERANKER_SCORE_THRESHOLD", None)
+    if cutoff is None or cutoff <= 0:
+        return ranked
+    filtered = [(cid, sc) for cid, sc in ranked if sc >= cutoff]
+    if len(filtered) < len(ranked):
+        logger.debug(
+            "rerank threshold %.2f removed %d/%d items",
+            cutoff,
+            len(ranked) - len(filtered),
+            len(ranked),
+        )
+    return filtered or ranked[: max(1, len(ranked) // 3)]
+
+
+def _score_by_rank_fallback(items: list[RerankItem]) -> list[tuple[str, float]]:
+    n = len(items)
+    if n == 0:
+        return []
+    return [(it.child_chunk_id, 1.0 - (i / max(n, 1))) for i, it in enumerate(items)]
+
+
+def _normalize_scores(pairs: list[tuple[str, float]]) -> list[tuple[str, float]]:
+    if not pairs:
+        return []
+    vals = [s for _, s in pairs]
+    lo, hi = min(vals), max(vals)
+    if hi <= lo:
+        return [(cid, 0.5) for cid, _ in pairs]
+    return [(cid, (s - lo) / (hi - lo)) for cid, s in pairs]
+
+
+def _rerank_local(query: str, items: list[RerankItem]) -> list[tuple[str, float]]:
+    model_name = cfg.RERANKER_MODEL or "cross-encoder/ms-marco-MiniLM-L-6-v2"
+    try:
+        encoder = _load_cross_encoder(model_name)
+    except ImportError as e:
+        logger.error("%s", e)
+        return _score_by_rank_fallback(items)
+    except Exception as e:
+        logger.error("Failed to load reranker model %s: %s", model_name, e)
+        return _score_by_rank_fallback(items)
+
+    # Cap items — retriever should already limit, but enforce here for safety
+    top_n = max(1, int(getattr(cfg, "RERANKER_TOP_N", 16)))
+    items = items[:top_n]
+
+    pairs = [(query, build_rerank_passage(it)) for it in items]
+    batch_size = max(1, int(getattr(cfg, "RERANKER_BATCH_SIZE", 8)))
+
+    t0 = perf_counter()
+    raw = encoder.predict(pairs, show_progress_bar=False, batch_size=batch_size)
+    elapsed = perf_counter() - t0
+    logger.info(
+        "Reranked %d passages in %.2fs (model=%s, device=%s)",
+        len(items),
+        elapsed,
+        model_name,
+        _resolve_device(),
+    )
+
+    scored = [(items[i].child_chunk_id, float(raw[i])) for i in range(len(items))]
+    scored.sort(key=lambda x: x[1], reverse=True)
+    normalized = _normalize_scores(scored)
+    return apply_rerank_threshold(normalized)
+
+
+def _rerank_cohere(query: str, items: list[RerankItem]) -> list[tuple[str, float]]:
+    try:
+        import cohere
+    except ImportError as e:
+        logger.error("cohere package not installed: %s", e)
+        return _score_by_rank_fallback(items)
+
+    key = getattr(cfg, "COHERE_API_KEY", "") or ""
+    if not key:
+        logger.error("COHERE_API_KEY not set — cannot use cohere reranker")
+        return _score_by_rank_fallback(items)
+
+    client = cohere.Client(key)
+    docs = [build_rerank_passage(it) for it in items]
+    model = getattr(cfg, "COHERE_RERANK_MODEL", None) or "rerank-multilingual-v3.0"
+    top_n = min(len(docs), max(cfg.RERANKER_TOP_N, len(docs)))
+    rsp = client.rerank(model=model, query=query, documents=docs, top_n=top_n)
+
+    out: list[tuple[str, float]] = []
+    for r in rsp.results:
+        out.append((items[r.index].child_chunk_id, float(r.relevance_score)))
+    out.sort(key=lambda x: x[1], reverse=True)
+    normalized = _normalize_scores(out)
+    return apply_rerank_threshold(normalized)
+
+
+def rerank_web_pages(
+    query: str,
+    pages: list[tuple[str, str, str]],
+) -> list[tuple[str, float]]:
+    """
+    Rerank web pages by pure query-passage relevancy (url, title, content).
+    Uses the same cross-encoder as local RAG but without score threshold filtering.
+    """
+    if not pages:
+        return []
+
+    items = [
+        RerankItem(
+            child_chunk_id=url,
+            passage=content,
+            section_title=title,
+            source_url=url,
+        )
+        for url, title, content in pages
+    ]
+    backend = (cfg.RERANKER_BACKEND or "none").lower().strip()
+    if backend == "local":
+        model_name = cfg.RERANKER_MODEL or "cross-encoder/ms-marco-MiniLM-L-6-v2"
+        try:
+            encoder = _load_cross_encoder(model_name)
+        except Exception as e:
+            logger.warning("web reranker unavailable: %s", e)
+            return [(url, 0.5 - i * 0.01) for i, (url, _, _) in enumerate(pages)]
+
+        pairs = [(query, build_rerank_passage(it)) for it in items]
+        batch_size = max(1, int(getattr(cfg, "RERANKER_BATCH_SIZE", 8)))
+        raw = encoder.predict(pairs, show_progress_bar=False, batch_size=batch_size)
+        scored = [(pages[i][0], float(raw[i])) for i in range(len(pages))]
+        scored.sort(key=lambda x: x[1], reverse=True)
+        return _normalize_scores(scored)
+
+    if backend == "cohere":
+        ranked = _rerank_cohere(query, items)
+        return ranked
+
+    return [(url, 0.5 - i * 0.01) for i, (url, _, _) in enumerate(pages)]
+
+
+def warm_start_reranker() -> None:
+    """Pre-load cross-encoder model at startup."""
+    if cfg.RERANKER_BACKEND == "local":
+        _load_cross_encoder()
+
+
+def is_reranker_loaded() -> bool:
+    if cfg.RERANKER_BACKEND != "local":
+        return True
+    return _CROSS_ENCODER is not None

package/knowledge_base/retrieval/retrieval_debug.py

@@ -0,0 +1,69 @@
+"""Retrieval observability: scores, ranks, inclusion reasons for tuning GraphRAG."""
+
+from __future__ import annotations
+
+import json
+import logging
+from dataclasses import dataclass, field, asdict
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ChunkScoreTrace:
+    child_chunk_id: str
+    parent_chunk_id: str
+    source_file: str
+    section_title: str
+    vector_rank: int | None
+    fts_rank: int | None
+    rrf_score: float
+    rerank_score: float | None
+    included: bool
+    exclusion_reason: str  # "" when included
+
+
+@dataclass
+class RetrievalDebugReport:
+    query: str
+    query_classification: str
+    total_vector_candidates: int
+    total_fts_candidates: int
+    rrf_merged_count: int
+    reranker_applied: bool
+    post_rerank_count: int
+    kg_nodes_found: int
+    kg_edges_traversed: int
+    final_chunk_count: int
+    total_tokens_used: int
+    token_budget: int
+    traces: list[ChunkScoreTrace] = field(default_factory=list)
+    timing_ms: dict[str, float] = field(default_factory=dict)
+
+    def to_json(self) -> str:
+        return json.dumps(asdict(self), indent=2)
+
+    def log_summary(self) -> None:
+        logger.info(
+            "retrieval_debug: class=%s vec=%d fts=%d rrf_merged=%d rerank=%s "
+            "final=%d kg_nodes=%d kg_edges=%d tokens~%d/%d",
+            self.query_classification,
+            self.total_vector_candidates,
+            self.total_fts_candidates,
+            self.rrf_merged_count,
+            self.reranker_applied,
+            self.final_chunk_count,
+            self.kg_nodes_found,
+            self.kg_edges_traversed,
+            self.total_tokens_used,
+            self.token_budget,
+        )
+
+
+def merge_traces_into_report(
+    base: RetrievalDebugReport,
+    traces: list[ChunkScoreTrace],
+) -> RetrievalDebugReport:
+    base.traces = traces
+    return base