ltcai 9.9.4 → 9.9.9

package/lattice_brain/graph/rerank.py

@@ -0,0 +1,160 @@
+"""Optional cross-encoder rerank for hybrid retrieval (v9.9.5).
+
+Default path is identity (fused score preserved) — no model download, no
+latency tax, no unearned claims. Opt in with::
+
+    LATTICEAI_CROSS_ENCODER_RERANK=1
+    # optional model id (sentence-transformers CrossEncoder):
+    LATTICEAI_CROSS_ENCODER_MODEL=cross-encoder/ms-marco-MiniLM-L-6-v2
+
+When the env kill-switch is off, or ``sentence_transformers`` / the model is
+unavailable, :func:`rerank_matches` returns the candidates unchanged and
+reports ``mode="identity"``. Failures never raise into the search path.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any, Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+CROSS_ENCODER_RERANK_ENV = "LATTICEAI_CROSS_ENCODER_RERANK"
+CROSS_ENCODER_MODEL_ENV = "LATTICEAI_CROSS_ENCODER_MODEL"
+DEFAULT_CROSS_ENCODER_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2"
+
+_model_cache: Dict[str, Any] = {}
+
+
+def _rerank_enabled() -> bool:
+    raw = os.getenv(CROSS_ENCODER_RERANK_ENV, "").strip().lower()
+    return raw in {"1", "true", "yes", "on"}
+
+
+def _model_id() -> str:
+    raw = os.getenv(CROSS_ENCODER_MODEL_ENV, "").strip()
+    return raw or DEFAULT_CROSS_ENCODER_MODEL
+
+
+def _candidate_text(match: Dict[str, Any]) -> str:
+    parts = [
+        str(match.get("title") or ""),
+        str(match.get("summary") or ""),
+        str((match.get("metadata") or {}).get("snippet") or ""),
+    ]
+    return " ".join(p for p in parts if p).strip() or str(match.get("node_id") or "")
+
+
+def _load_cross_encoder(model_id: str) -> Any:
+    if model_id in _model_cache:
+        return _model_cache[model_id]
+    from sentence_transformers import CrossEncoder  # type: ignore
+
+    model = CrossEncoder(model_id)
+    _model_cache[model_id] = model
+    return model
+
+
+def identity_rerank(
+    query: str,
+    candidates: List[Dict[str, Any]],
+    *,
+    top_k: Optional[int] = None,
+) -> Dict[str, Any]:
+    """Preserve fused ordering; stamp identity scores for an honest contract."""
+    del query  # identity path does not use the query text
+    ranked = list(candidates)
+    for item in ranked:
+        item["rerank_score"] = float(item.get("score") or item.get("fused_score") or 0.0)
+    if top_k is not None:
+        ranked = ranked[: max(1, int(top_k))]
+    return {
+        "matches": ranked,
+        "mode": "identity",
+        "model": None,
+        "detail": None,
+    }
+
+
+def cross_encoder_rerank(
+    query: str,
+    candidates: List[Dict[str, Any]],
+    *,
+    top_k: Optional[int] = None,
+    model_id: Optional[str] = None,
+) -> Dict[str, Any]:
+    """Score (query, candidate) pairs with a CrossEncoder when available."""
+    if not candidates:
+        return {
+            "matches": [],
+            "mode": "cross_encoder",
+            "model": model_id or _model_id(),
+            "detail": None,
+        }
+    mid = model_id or _model_id()
+    try:
+        model = _load_cross_encoder(mid)
+    except Exception as exc:  # noqa: BLE001 — never break search
+        logger.info("cross-encoder unavailable (%s); falling back to identity", exc)
+        result = identity_rerank(query, candidates, top_k=top_k)
+        result["detail"] = f"cross_encoder_unavailable: {exc}"
+        return result
+
+    pairs = [[str(query or ""), _candidate_text(c)] for c in candidates]
+    try:
+        scores = model.predict(pairs)
+    except Exception as exc:  # noqa: BLE001
+        logger.warning("cross-encoder predict failed: %s", exc)
+        result = identity_rerank(query, candidates, top_k=top_k)
+        result["detail"] = f"cross_encoder_predict_failed: {exc}"
+        return result
+
+    ranked = list(candidates)
+    for item, score in zip(ranked, scores):
+        item["rerank_score"] = float(score)
+        # Surface the rerank score as the primary ranking key while keeping
+        # the pre-rerank fused score under scores.fused for audit.
+        scores_map = item.setdefault("scores", {})
+        if isinstance(scores_map, dict):
+            scores_map.setdefault("fused", float(item.get("score") or 0.0))
+            scores_map["rerank"] = float(score)
+        item["score"] = float(score)
+    ranked.sort(key=lambda m: (-float(m.get("rerank_score") or 0.0), str(m.get("node_id") or "")))
+    if top_k is not None:
+        ranked = ranked[: max(1, int(top_k))]
+    for rank, match in enumerate(ranked, start=1):
+        match["rank"] = rank
+    return {
+        "matches": ranked,
+        "mode": "cross_encoder",
+        "model": mid,
+        "detail": None,
+    }
+
+
+def rerank_matches(
+    query: str,
+    candidates: List[Dict[str, Any]],
+    *,
+    top_k: Optional[int] = None,
+    force: Optional[bool] = None,
+) -> Dict[str, Any]:
+    """Public entry: cross-encoder when enabled, else identity.
+
+    ``force=True/False`` overrides the env kill-switch (tests only).
+    """
+    enabled = _rerank_enabled() if force is None else bool(force)
+    if not enabled:
+        return identity_rerank(query, candidates, top_k=top_k)
+    return cross_encoder_rerank(query, candidates, top_k=top_k)
+
+
+__all__ = [
+    "CROSS_ENCODER_MODEL_ENV",
+    "CROSS_ENCODER_RERANK_ENV",
+    "DEFAULT_CROSS_ENCODER_MODEL",
+    "cross_encoder_rerank",
+    "identity_rerank",
+    "rerank_matches",
+]

package/lattice_brain/graph/retrieval.py

@@ -4,6 +4,14 @@ from __future__ import annotations
 
 from ._kg_common import *  # noqa: F403,F401
 
+# --- Compat seam (v9.9.5 decomposition) -------------------------------------
+# The non-search read surface (list_documents / workspaces_of /
+# filter_scoped_nodes / neighbors / get_node / relationship_search /
+# traverse / stats) moved byte-identically to .retrieval_reads as
+# KnowledgeGraphReadsMixin. Re-exported here so any legacy
+# ``from lattice_brain.graph.retrieval import ...`` site keeps resolving.
+from .retrieval_reads import KnowledgeGraphReadsMixin  # noqa: F401
+
 
 def context_quality_signal(
     mode: str,
@@ -68,119 +76,6 @@ class KnowledgeGraphRetrievalMixin:
         "Agent",  # 에이전트
     )
 
-    def list_documents(self, limit: int = 200) -> Dict[str, Any]:
-        """List ingested ``Document`` nodes with their ingest + index state.
-
-        Powers the Files view: every accepted upload and every indexed local
-        document becomes a ``Document`` node. A document is reported ``indexed``
-        once its retrieval chunks exist (searchable in Chat / Hybrid Search).
-        """
-        limit = max(1, min(int(limit or 200), 1000))
-        nt, _ = self._read_tables()
-        documents: List[Dict[str, Any]] = []
-        with self._connect() as conn:
-            rows = conn.execute(
-                f"SELECT id, title, summary, metadata_json, created_at, updated_at "
-                f"FROM {nt} WHERE type='Document' ORDER BY updated_at DESC, id ASC LIMIT ?",
-                (limit,),
-            ).fetchall()
-            for row in rows:
-                meta = _safe_loads(row["metadata_json"]) or {}
-                extracted = meta.get("extracted") or {}
-                node_id = row["id"]
-                chunk_count = conn.execute(
-                    "SELECT COUNT(*) AS c FROM chunks WHERE source_node=?",
-                    (node_id,),
-                ).fetchone()["c"]
-                if not chunk_count:
-                    # Legacy projections represented chunks as graph nodes and
-                    # linked them only through metadata_json. Keep read
-                    # compatibility without making the fragile LIKE path the
-                    # primary query.
-                    chunk_count = conn.execute(
-                        f"SELECT COUNT(*) AS c FROM {nt} WHERE type='Chunk' AND metadata_json LIKE ?",
-                        (f"%{node_id}%",),
-                    ).fetchone()["c"]
-                documents.append(
-                    {
-                        "id": node_id,
-                        "filename": meta.get("filename") or row["title"],
-                        "ext": meta.get("ext"),
-                        "mime_type": meta.get("mime_type"),
-                        "bytes": meta.get("bytes"),
-                        "sha256": meta.get("sha256"),
-                        "uploader": meta.get("uploader"),
-                        "chars": extracted.get("chars"),
-                        "chunks": int(chunk_count or 0),
-                        "indexed": int(chunk_count or 0) > 0,
-                        "ingest_state": "indexed"
-                        if int(chunk_count or 0) > 0
-                        else "ingested",
-                        "created_at": row["created_at"],
-                        "updated_at": row["updated_at"],
-                    }
-                )
-        return {
-            "documents": documents,
-            "total": len(documents),
-            "generated_at": datetime.now().isoformat(timespec="seconds"),
-        }
-
-    def workspaces_of(self, node_ids) -> Dict[str, Optional[str]]:
-        """Map known node ids to their workspace scope.
-
-        ``None`` is returned only for a row that is explicitly present in the
-        authoritative v2 projection with a NULL workspace.  Missing ids remain
-        missing, and projection/query failures propagate so callers can fail
-        closed instead of mistaking every candidate for legacy-global data.
-        """
-        ids = [str(i) for i in node_ids if i]
-        if not ids:
-            return {}
-        placeholders = ",".join("?" for _ in ids)
-        with self._connect() as conn:
-            return {
-                row["id"]: row["workspace_id"]
-                for row in conn.execute(
-                    f"SELECT id, workspace_id FROM nodes_v2 WHERE id IN ({placeholders})",
-                    ids,
-                ).fetchall()
-            }
-
-    def filter_scoped_nodes(
-        self,
-        items,
-        allowed_workspaces,
-        *,
-        id_key: str = "id",
-        include_legacy_global: bool = False,
-    ):
-        """Drop items scoped to a workspace the caller is not a member of.
-
-        ``allowed_workspaces=None`` means no scoping (single-user / no-auth
-        mode). In scoped/multi-user mode, unknown ids are private and
-        legacy-global rows require the explicit ``include_legacy_global=True``
-        compatibility opt-in.
-        """
-        candidates = list(items)
-        if allowed_workspaces is None:
-            return candidates
-        allowed = {str(workspace_id) for workspace_id in allowed_workspaces if workspace_id}
-        scopes = self.workspaces_of([item.get(id_key) for item in candidates])
-        visible = []
-        for item in candidates:
-            node_id = str(item.get(id_key) or "")
-            if not node_id or node_id not in scopes:
-                # Unknown/unprojected rows are never treated as public.
-                continue
-            workspace_id = scopes[node_id]
-            if workspace_id is None:
-                if include_legacy_global:
-                    visible.append(item)
-            elif str(workspace_id) in allowed:
-                visible.append(item)
-        return visible
-
     def graph(
         self,
         limit: int = 300,
@@ -679,7 +574,25 @@ class KnowledgeGraphRetrievalMixin:
                 match["score"] = round(float(match["score"]) * multiplier, 6)
 
         matches.sort(key=lambda item: (-item["score"], item["node_id"]))
-        matches = matches[:top_k]
+        # Optional cross-encoder rerank (v9.9.5). Off by default; when the
+        # env kill-switch is set and the model loads, pair scores reorder the
+        # fused list. Failures degrade to identity and never break search.
+        rerank_meta: Dict[str, Any]
+        try:
+            from .rerank import rerank_matches
+
+            # Rerank a slightly wider window, then cut to top_k.
+            window = matches[: max(top_k * 2, top_k)]
+            reranked = rerank_matches(search_query, window, top_k=top_k)
+            matches = list(reranked.get("matches") or matches[:top_k])
+            rerank_meta = {
+                "mode": reranked.get("mode") or "identity",
+                "model": reranked.get("model"),
+                "detail": reranked.get("detail"),
+            }
+        except Exception as exc:  # noqa: BLE001 — rerank must never break search
+            matches = matches[:top_k]
+            rerank_meta = {"mode": "identity", "model": None, "detail": str(exc)}
         for rank, match in enumerate(matches, start=1):
             match["rank"] = rank
         result = {
@@ -691,6 +604,7 @@ class KnowledgeGraphRetrievalMixin:
             "sources": {"lexical": len(lexical_matches), "vector": len(vector_matches)},
             "matches": matches,
             "policy": {"search_query": search_query, "rewrite_rules": rewrite_rules},
+            "rerank": rerank_meta,
             "detail": detail,
         }
         if vector_degraded is not None:
@@ -847,293 +761,6 @@ class KnowledgeGraphRetrievalMixin:
             with_meta=True,
         )
 
-    def neighbors(
-        self,
-        node_id: str,
-        *,
-        allowed_workspaces=None,
-        include_legacy_global: bool = False,
-    ) -> Dict[str, Any]:
-        """Return direct neighbors (1-hop) of a node."""
-        if allowed_workspaces is not None and not self.filter_scoped_nodes(
-            [{"id": node_id}],
-            allowed_workspaces,
-            include_legacy_global=include_legacy_global,
-        ):
-            raise ValueError(f"graph node not found: {node_id}")
-        nt, et = self._read_tables()
-        with self._connect() as conn:
-            edge_rows = conn.execute(
-                f"SELECT from_node, to_node, type, weight FROM {et} WHERE from_node=? OR to_node=? ORDER BY id ASC",
-                (node_id, node_id),
-            ).fetchall()
-            neighbor_ids: set = set()
-            edges = []
-            for row in edge_rows:
-                neighbor_ids.add(row["from_node"])
-                neighbor_ids.add(row["to_node"])
-                edges.append(
-                    {
-                        "from": row["from_node"],
-                        "to": row["to_node"],
-                        "type": row["type"],
-                        "weight": row["weight"],
-                    }
-                )
-            neighbor_ids.discard(node_id)
-            nodes = []
-            if neighbor_ids:
-                placeholders = ",".join("?" * len(neighbor_ids))
-                nodes = [
-                    {
-                        "id": row["id"],
-                        "type": row["type"],
-                        "title": row["title"],
-                        "summary": row["summary"],
-                        "metadata": _safe_loads(row["metadata_json"]),
-                    }
-                    for row in conn.execute(
-                        f"SELECT id, type, title, summary, metadata_json FROM {nt} WHERE id IN ({placeholders}) ORDER BY id ASC",
-                        list(neighbor_ids),
-                    )
-                ]
-        if allowed_workspaces is not None:
-            nodes = self.filter_scoped_nodes(
-                nodes,
-                allowed_workspaces,
-                include_legacy_global=include_legacy_global,
-            )
-            kept = {node.get("id") for node in nodes}
-            edges = [
-                edge for edge in edges
-                if (edge.get("from") == node_id or edge.get("from") in kept)
-                and (edge.get("to") == node_id or edge.get("to") in kept)
-            ]
-        return {"node_id": node_id, "neighbors": nodes, "edges": edges}
-
-    def get_node(
-        self,
-        node_id: str,
-        *,
-        allowed_workspaces=None,
-        include_legacy_global: bool = False,
-    ) -> Dict[str, Any]:
-        node_id = str(node_id or "").strip()
-        if not node_id:
-            raise ValueError("node_id required")
-        nt, et = self._read_tables()
-        with self._connect() as conn:
-            row = conn.execute(
-                f"""
-                    SELECT id, type, title, summary, metadata_json, updated_at
-                    FROM {nt}
-                    WHERE id=?
-                    """,
-                (node_id,),
-            ).fetchone()
-            if not row:
-                raise ValueError(f"graph node not found: {node_id}")
-            degree = conn.execute(
-                f"SELECT COUNT(*) AS c FROM {et} WHERE from_node=? OR to_node=?",
-                (node_id, node_id),
-            ).fetchone()["c"]
-        node = {
-            "id": row["id"],
-            "type": row["type"],
-            "title": row["title"],
-            "summary": row["summary"],
-            "metadata": _safe_loads(row["metadata_json"]),
-            "updated_at": row["updated_at"],
-            "degree": degree,
-        }
-        if allowed_workspaces is not None and not self.filter_scoped_nodes(
-            [node],
-            allowed_workspaces,
-            include_legacy_global=include_legacy_global,
-        ):
-            raise ValueError(f"graph node not found: {node_id}")
-        return node
-
-    def relationship_search(
-        self,
-        *,
-        query: str = "",
-        node_id: str = "",
-        relationship_type: str = "",
-        limit: int = 30,
-        allowed_workspaces=None,
-        include_legacy_global: bool = False,
-    ) -> Dict[str, Any]:
-        query = str(query or "").strip()
-        node_id = str(node_id or "").strip()
-        relationship_type = str(relationship_type or "").strip()
-        limit = max(1, min(int(limit or 30), 200))
-        nt, et = self._read_tables()
-        where = []
-        params: List[Any] = []
-        if node_id:
-            where.append("(e.from_node=? OR e.to_node=?)")
-            params.extend([node_id, node_id])
-        if relationship_type:
-            where.append("e.type LIKE ?")
-            params.append(f"%{relationship_type}%")
-        if query:
-            where.append(
-                "(e.type LIKE ? OR e.metadata_json LIKE ? OR src.title LIKE ? OR dst.title LIKE ? OR src.summary LIKE ? OR dst.summary LIKE ?)"
-            )
-            params.extend([f"%{query}%"] * 6)
-        where_sql = "WHERE " + " AND ".join(where) if where else ""
-        with self._connect() as conn:
-            rows = conn.execute(
-                f"""
-                    SELECT
-                      e.id, e.from_node, e.to_node, e.type, e.weight, e.metadata_json, e.created_at,
-                      src.type AS source_type, src.title AS source_title, src.summary AS source_summary,
-                      src.metadata_json AS source_metadata,
-                      dst.type AS target_type, dst.title AS target_title, dst.summary AS target_summary,
-                      dst.metadata_json AS target_metadata
-                    FROM {et} e
-                    JOIN {nt} src ON src.id=e.from_node
-                    JOIN {nt} dst ON dst.id=e.to_node
-                    {where_sql}
-                    ORDER BY e.weight DESC, e.created_at DESC, e.id ASC
-                    LIMIT ?
-                    """,
-                (*params, limit),
-            ).fetchall()
-        relationships = [
-                {
-                    "id": row["id"],
-                    "type": row["type"],
-                    "weight": row["weight"],
-                    "metadata": _safe_loads(row["metadata_json"]),
-                    "created_at": row["created_at"],
-                    "source": {
-                        "id": row["from_node"],
-                        "type": row["source_type"],
-                        "title": row["source_title"],
-                        "summary": row["source_summary"],
-                        "metadata": _safe_loads(row["source_metadata"]),
-                    },
-                    "target": {
-                        "id": row["to_node"],
-                        "type": row["target_type"],
-                        "title": row["target_title"],
-                        "summary": row["target_summary"],
-                        "metadata": _safe_loads(row["target_metadata"]),
-                    },
-                }
-                for row in rows
-            ]
-        if allowed_workspaces is not None:
-            kept = []
-            for rel in relationships:
-                endpoints = [
-                    {"id": (rel.get("source") or {}).get("id")},
-                    {"id": (rel.get("target") or {}).get("id")},
-                ]
-                if len(
-                    self.filter_scoped_nodes(
-                        endpoints,
-                        allowed_workspaces,
-                        include_legacy_global=include_legacy_global,
-                    )
-                ) == 2:
-                    kept.append(rel)
-            relationships = kept
-        return {
-            "query": query,
-            "node_id": node_id,
-            "relationship_type": relationship_type,
-            "relationships": relationships,
-        }
-
-    def traverse(
-        self,
-        node_id: str,
-        *,
-        depth: int = 1,
-        limit: int = 100,
-        allowed_workspaces=None,
-        include_legacy_global: bool = False,
-    ) -> Dict[str, Any]:
-        node_id = str(node_id or "").strip()
-        if not node_id:
-            raise ValueError("node_id required")
-        if allowed_workspaces is not None and not self.filter_scoped_nodes(
-            [{"id": node_id}],
-            allowed_workspaces,
-            include_legacy_global=include_legacy_global,
-        ):
-            raise ValueError(f"graph node not found: {node_id}")
-        depth = max(0, min(int(depth or 1), 4))
-        limit = max(1, min(int(limit or 100), 500))
-        nt, et = self._read_tables()
-        visited = {node_id}
-        frontier = {node_id}
-        edges_by_id: Dict[str, Dict[str, Any]] = {}
-        with self._connect() as conn:
-            for _ in range(depth):
-                if not frontier or len(visited) >= limit:
-                    break
-                placeholders = ",".join("?" * len(frontier))
-                rows = conn.execute(
-                    f"""
-                        SELECT id, from_node, to_node, type, weight, metadata_json
-                        FROM {et}
-                        WHERE from_node IN ({placeholders}) OR to_node IN ({placeholders})
-                        ORDER BY weight DESC, id ASC
-                        LIMIT ?
-                        """,
-                    (*frontier, *frontier, limit * 3),
-                ).fetchall()
-                next_frontier = set()
-                for row in rows:
-                    edges_by_id[row["id"]] = {
-                        "id": row["id"],
-                        "from": row["from_node"],
-                        "to": row["to_node"],
-                        "type": row["type"],
-                        "weight": row["weight"],
-                        "metadata": _safe_loads(row["metadata_json"]),
-                    }
-                    for candidate in (row["from_node"], row["to_node"]):
-                        if candidate not in visited and len(visited) < limit:
-                            visited.add(candidate)
-                            next_frontier.add(candidate)
-                frontier = next_frontier
-            placeholders = ",".join("?" * len(visited))
-            node_rows = conn.execute(
-                f"""
-                    SELECT id, type, title, summary, metadata_json, updated_at
-                    FROM {nt}
-                    WHERE id IN ({placeholders})
-                    ORDER BY updated_at DESC, id ASC
-                    """,
-                list(visited),
-            ).fetchall()
-        nodes = [
-                {
-                    "id": row["id"],
-                    "type": row["type"],
-                    "title": row["title"],
-                    "summary": row["summary"],
-                    "metadata": _safe_loads(row["metadata_json"]),
-                    "updated_at": row["updated_at"],
-                }
-                for row in node_rows
-            ]
-        edges = list(edges_by_id.values())
-        if allowed_workspaces is not None:
-            nodes = self.filter_scoped_nodes(
-                nodes,
-                allowed_workspaces,
-                include_legacy_global=include_legacy_global,
-            )
-            kept = {node.get("id") for node in nodes}
-            edges = [edge for edge in edges if edge.get("from") in kept and edge.get("to") in kept]
-        return {"root": node_id, "depth": depth, "nodes": nodes, "edges": edges}
-
     def delete_conversation(self, conversation_id: str) -> Dict[str, Any]:
         conversation_id = str(conversation_id or "").strip()
         if not conversation_id:
@@ -1228,44 +855,3 @@ class KnowledgeGraphRetrievalMixin:
             shutil.rmtree(self.blob_dir, ignore_errors=True)
             self.blob_dir.mkdir(parents=True, exist_ok=True)
         return {"status": "ok", "removed": counts}
-
-    def stats(self) -> Dict[str, Any]:
-        nt, et = self._read_tables()
-        with self._connect() as conn:
-            node_counts = {
-                row["type"]: row["count"]
-                for row in conn.execute(
-                    f"SELECT type, COUNT(*) AS count FROM {nt} GROUP BY type"
-                )
-            }
-            edge_counts = {
-                row["type"]: row["count"]
-                for row in conn.execute(
-                    f"SELECT type, COUNT(*) AS count FROM {et} GROUP BY type"
-                )
-            }
-            local_sources = conn.execute(
-                "SELECT COUNT(*) AS c FROM knowledge_sources"
-            ).fetchone()["c"]
-            local_file_status = {
-                row["status"]: row["count"]
-                for row in conn.execute(
-                    "SELECT status, COUNT(*) AS count FROM local_file_index GROUP BY status"
-                )
-            }
-        v2 = None
-        if KGStoreV2 is not None:
-            try:
-                v2 = KGStoreV2(self.db_path).stats()
-            except Exception as e:
-                v2 = {"available": False, "error": str(e)}
-        return {
-            "db_path": str(self.db_path),
-            "schema_version": GRAPH_SCHEMA_VERSION,
-            "v2_schema_available": KGStoreV2 is not None,
-            "nodes": node_counts,
-            "edges": edge_counts,
-            "local_sources": local_sources,
-            "local_file_status": local_file_status,
-            "v2": v2,
-        }