know-do-graph 0.1.0__py3-none-any.whl

core/retrieval/retrieval.py

@@ -0,0 +1,346 @@
+from __future__ import annotations
+
+from typing import Optional
+
+from sqlalchemy import or_
+from sqlalchemy.orm import Session
+
+from core.graph.graph import KnowDoGraph
+from core.retrieval import vector_store
+from core.retrieval.embedder import build_embedding_text, get_default_embedder
+from core.retrieval.fusion import reciprocal_rank_fusion, trust_multiplier, usage_bump
+from core.schemas.edge import Edge, EdgeRelation
+from core.schemas.entry import Entry, EntryType
+from core.storage.models import EdgeModel, EntryModel
+
+_STOP_WORDS = {
+    "a", "an", "the", "of", "in", "on", "at", "to", "for", "and", "or", "is", "are", "be",
+}
+
+# Bidirectional synonym groups — each list is a set of interchangeable terms.
+# When any member appears as a query token the others are added automatically.
+_SYNONYM_GROUPS: list[list[str]] = [
+    ["cnt", "nanotube", "carbon-nanotube", "carbon nanotube"],
+    ["cnt", "carbon tube", "tube"],
+    ["filled tube", "filled cnt", "filled nanotube", "confined"],
+    ["ase", "atomic simulation environment"],
+    ["slab", "surface slab", "surface"],
+    ["interface", "heterostructure", "film substrate"],
+    ["supercell", "super cell", "expansion"],
+    ["lattice matching", "zsl", "coherent interface"],
+    ["methane", "ch4"],
+    ["dft", "density functional theory"],
+    ["md", "molecular dynamics"],
+    ["mlip", "mace", "machine learning potential", "interatomic potential"],
+    ["crystal", "bulk crystal", "bulk structure"],
+    ["nanoparticle", "nano particle", "nanostructure"],
+]
+
+# Build a fast lookup: token → set of synonym tokens (excluding itself)
+_SYNONYM_MAP: dict[str, set[str]] = {}
+for _group in _SYNONYM_GROUPS:
+    for _term in _group:
+        others = {t for t in _group if t != _term}
+        _SYNONYM_MAP.setdefault(_term, set()).update(others)
+
+
+def _expand_tokens(tokens: list[str]) -> list[str]:
+    """Return *tokens* plus any synonyms, deduped, preserving order."""
+    seen: set[str] = set(tokens)
+    expanded = list(tokens)
+    for tok in tokens:
+        for syn in _SYNONYM_MAP.get(tok, ()):
+            if syn not in seen:
+                seen.add(syn)
+                expanded.append(syn)
+    return expanded
+
+
+class RetrievalEngine:
+    """Search and graph-traversal interface over the persisted graph."""
+
+    def __init__(self, db: Session, graph: KnowDoGraph) -> None:
+        self._db = db
+        self._graph = graph
+
+    # ------------------------------------------------------------------
+    # Entry lookups
+    # ------------------------------------------------------------------
+
+    def get_entry_by_id(self, entry_id: str) -> Optional[Entry]:
+        row = self._db.get(EntryModel, entry_id)
+        return Entry(**row.to_dict()) if row else None
+
+    def get_entry_by_slug(self, slug: str) -> Optional[Entry]:
+        row = self._db.query(EntryModel).filter_by(slug=slug).first()
+        return Entry(**row.to_dict()) if row else None
+
+    def get_entry_by_alias(self, alias: str) -> Optional[Entry]:
+        """Return the first entry whose aliases list contains *alias* (case-insensitive)."""
+        alias_lower = alias.lower()
+        rows = (
+            self._db.query(EntryModel)
+            .filter(EntryModel.aliases.ilike(f"%{alias_lower}%"))
+            .all()
+        )
+        for row in rows:
+            entry = Entry(**row.to_dict())
+            if any(a.lower() == alias_lower for a in entry.aliases):
+                return entry
+        return None
+
+    def resolve_identifier(self, identifier: str) -> Optional[Entry]:
+        """Try ID → slug → alias in order and return the first match."""
+        return (
+            self.get_entry_by_id(identifier)
+            or self.get_entry_by_slug(identifier)
+            or self.get_entry_by_alias(identifier)
+        )
+
+    def list_entries(self, limit: int = 50, offset: int = 0) -> list[Entry]:
+        rows = self._db.query(EntryModel).offset(offset).limit(limit).all()
+        return [Entry(**row.to_dict()) for row in rows]
+
+    # ------------------------------------------------------------------
+    # Hybrid search
+    # ------------------------------------------------------------------
+
+    def search_entries(
+        self,
+        query: Optional[str] = None,
+        tags: Optional[list[str]] = None,
+        entry_type: Optional[EntryType] = None,
+        limit: int = 20,
+        mode: str = "hybrid",
+    ) -> list[Entry]:
+        return [e for _, e in self._search_impl(query, tags, entry_type, limit, mode)]
+
+    def search_entries_scored(
+        self,
+        query: Optional[str] = None,
+        tags: Optional[list[str]] = None,
+        entry_type: Optional[EntryType] = None,
+        limit: int = 20,
+        mode: str = "hybrid",
+    ) -> list[tuple[Entry, float]]:
+        """Like search_entries but returns (entry, score) with scores normalized 0.0–1.0."""
+        raw = self._search_impl(query, tags, entry_type, limit, mode)
+        if not raw:
+            return []
+        max_score = max(s for s, _ in raw) or 1.0
+        return [(e, s / max_score) for s, e in raw]
+
+    def _search_impl(
+        self,
+        query: Optional[str],
+        tags: Optional[list[str]],
+        entry_type: Optional[EntryType],
+        limit: int,
+        mode: str = "hybrid",
+    ) -> list[tuple[Entry, float]]:
+        """Retrieval with three modes:
+          - "hybrid": keyword + vector ANN fused with RRF (default)
+          - "semantic": vector-only ANN (best for conceptually similar queries)
+          - "keyword": keyword-only (best for exact title/tag/acronym lookups)
+
+        Falls back gracefully:
+          - No query → filter-only listing, scores all 0.
+          - No embedder / no vec index → pure keyword path regardless of mode.
+        """
+        if not query:
+            return [(0.0, e) for e in self._filter_only(tags=tags, entry_type=entry_type, limit=limit)]
+
+        entries_by_id: dict[str, Entry] = {}
+
+        # Channel A — keyword scorer (skipped in semantic mode).
+        keyword_ranked: list[str] = []
+        if mode != "semantic":
+            keyword_hits = self._keyword_search(
+                query=query, tags=tags, entry_type=entry_type, limit=200
+            )
+            keyword_ranked = [eid for eid, _ in keyword_hits]
+            entries_by_id.update({eid: e for eid, _, e in self._with_entries(keyword_hits)})
+
+        # Channel B — vector ANN (skipped in keyword mode).
+        vector_ranked: list[str] = []
+        embedder = get_default_embedder()
+        if mode != "keyword" and embedder.available:
+            qvec = embedder.embed([query])[0]
+            for eid, _dist in vector_store.knn(self._db, qvec, k=50):
+                if eid in entries_by_id:
+                    vector_ranked.append(eid)
+                    continue
+                row = self._db.get(EntryModel, eid)
+                if row is None:
+                    continue
+                d = row.to_dict()
+                if entry_type and d.get("entry_type") != entry_type.value:
+                    continue
+                if tags and not any(t in d["tags"] for t in tags):
+                    continue
+                entries_by_id[eid] = Entry(**d)
+                vector_ranked.append(eid)
+
+        # Fuse + rerank.
+        fused = reciprocal_rank_fusion([keyword_ranked, vector_ranked])
+        if not fused:
+            return []
+
+        scored: list[tuple[float, Entry]] = []
+        for eid, base in fused.items():
+            entry = entries_by_id.get(eid)
+            if entry is None:
+                continue
+            meta = entry.metadata
+            mult = trust_multiplier(
+                meta.verification_status.value,
+                trust_score_override=meta.trust_score,
+            ) * usage_bump(meta.usage_count or 0)
+            scored.append((base * mult, entry))
+
+        scored.sort(key=lambda x: x[0], reverse=True)
+        return scored[:limit]
+
+    # ------------------------------------------------------------------
+    # Channel implementations
+    # ------------------------------------------------------------------
+
+    def _filter_only(
+        self,
+        tags: Optional[list[str]],
+        entry_type: Optional[EntryType],
+        limit: int,
+    ) -> list[Entry]:
+        q = self._db.query(EntryModel)
+        if entry_type:
+            q = q.filter(EntryModel.entry_type == entry_type.value)
+        rows = q.limit(500).all()
+        out: list[Entry] = []
+        for row in rows:
+            d = row.to_dict()
+            if tags and not any(t in d["tags"] for t in tags):
+                continue
+            out.append(Entry(**d))
+            if len(out) >= limit:
+                break
+        return out
+
+    def _keyword_search(
+        self,
+        query: str,
+        tags: Optional[list[str]],
+        entry_type: Optional[EntryType],
+        limit: int,
+    ) -> list[tuple[str, int]]:
+        """Returns ``[(entry_id, raw_score), ...]`` sorted by descending score.
+
+        Same scoring formula as before — title 10 / alias 5 / tag 3 / content 1.
+        """
+        q = self._db.query(EntryModel)
+        if entry_type:
+            q = q.filter(EntryModel.entry_type == entry_type.value)
+
+        tokens = [
+            t for t in query.lower().split()
+            if len(t) > 2 and t not in _STOP_WORDS
+        ]
+        if not tokens:
+            tokens = [query.lower()]
+
+        # Expand tokens with domain synonyms so e.g. "CNT" also finds "nanotube"
+        tokens = _expand_tokens(tokens)
+
+        token_filters = []
+        for token in tokens:
+            tl = f"%{token}%"
+            token_filters.append(
+                or_(
+                    EntryModel.title.ilike(tl),
+                    EntryModel.content.ilike(tl),
+                    EntryModel.aliases.ilike(tl),
+                )
+            )
+        q = q.filter(or_(*token_filters))
+
+        rows = q.limit(500).all()
+
+        scored: list[tuple[int, str]] = []
+        for row in rows:
+            d = row.to_dict()
+            if tags and not any(t in d["tags"] for t in tags):
+                continue
+            title = (d.get("title") or "").lower()
+            aliases = str(d.get("aliases") or "").lower()
+            tags_str = str(d.get("tags") or "").lower()
+            content = (d.get("content") or "").lower()
+            score = 0
+            for token in tokens:
+                if token in title:
+                    score += 10
+                if token in aliases:
+                    score += 5
+                if token in tags_str:
+                    score += 3
+                if token in content and token not in title:
+                    score += 1
+            scored.append((score, d["id"]))
+
+        scored.sort(key=lambda x: x[0], reverse=True)
+        return [(eid, s) for s, eid in scored[:limit]]
+
+    def _with_entries(
+        self, hits: list[tuple[str, int]]
+    ) -> list[tuple[str, int, Entry]]:
+        out: list[tuple[str, int, Entry]] = []
+        for eid, score in hits:
+            row = self._db.get(EntryModel, eid)
+            if row is None:
+                continue
+            out.append((eid, score, Entry(**row.to_dict())))
+        return out
+
+    # ------------------------------------------------------------------
+    # Edge lookups
+    # ------------------------------------------------------------------
+
+    def get_edges_for_entry(self, entry_id: str) -> list[Edge]:
+        rows = (
+            self._db.query(EdgeModel)
+            .filter(
+                (EdgeModel.source_id == entry_id)
+                | (EdgeModel.target_id == entry_id)
+            )
+            .all()
+        )
+        return [Edge(**row.to_dict()) for row in rows]
+
+    # ------------------------------------------------------------------
+    # Graph traversal
+    # ------------------------------------------------------------------
+
+    def get_related_entries(
+        self,
+        entry_id: str,
+        depth: int = 1,
+        relation: Optional[EdgeRelation] = None,
+    ) -> list[Entry]:
+        related_ids = self._graph.get_related_ids(entry_id, depth=depth, relation=relation)
+        entries: list[Entry] = []
+        for rid in related_ids:
+            entry = self.get_entry_by_id(rid)
+            if entry:
+                entries.append(entry)
+        return entries
+
+    # ------------------------------------------------------------------
+    # Embedding maintenance — used by the backfill CLI
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def embedding_text_for(entry: Entry) -> str:
+        return build_embedding_text(
+            title=entry.title,
+            aliases=entry.aliases,
+            tags=entry.tags,
+            content=entry.content,
+        )

core/retrieval/vector_store.py

@@ -0,0 +1,91 @@
+"""Thin wrapper around the sqlite-vec `entry_embeddings` virtual table.
+
+All operations are no-ops when sqlite-vec is not loaded — callers can use
+the wrapper unconditionally and let it degrade silently. The virtual table
+itself is created in ``core.storage.database.init_db``.
+"""
+
+from __future__ import annotations
+
+import logging
+import struct
+from typing import Optional
+
+from sqlalchemy import text
+from sqlalchemy.orm import Session
+
+logger = logging.getLogger(__name__)
+
+
+def _pack(vec: list[float]) -> bytes:
+    """sqlite-vec accepts vectors as little-endian float32 blobs."""
+    return struct.pack(f"<{len(vec)}f", *vec)
+
+
+def _table_available(db: Session) -> bool:
+    """Cheap probe — returns True iff the vec0 virtual table is queryable."""
+    try:
+        db.execute(text("SELECT entry_id FROM entry_embeddings LIMIT 0"))
+        return True
+    except Exception:
+        return False
+
+
+def upsert(db: Session, entry_id: str, vec: list[float]) -> bool:
+    if not vec or not _table_available(db):
+        return False
+    try:
+        blob = _pack(vec)
+        db.execute(text("DELETE FROM entry_embeddings WHERE entry_id = :id"), {"id": entry_id})
+        db.execute(
+            text("INSERT INTO entry_embeddings (entry_id, embedding) VALUES (:id, :v)"),
+            {"id": entry_id, "v": blob},
+        )
+        db.commit()
+        return True
+    except Exception as exc:
+        logger.warning("vector upsert failed for %s: %s", entry_id, exc)
+        db.rollback()
+        return False
+
+
+def delete(db: Session, entry_id: str) -> None:
+    if not _table_available(db):
+        return
+    try:
+        db.execute(text("DELETE FROM entry_embeddings WHERE entry_id = :id"), {"id": entry_id})
+        db.commit()
+    except Exception:
+        db.rollback()
+
+
+def knn(db: Session, vec: list[float], k: int = 50) -> list[tuple[str, float]]:
+    """Return ``[(entry_id, distance), ...]`` ordered by ascending distance.
+
+    Empty list if the index is unavailable or empty.
+    """
+    if not vec or not _table_available(db):
+        return []
+    try:
+        blob = _pack(vec)
+        rows = db.execute(
+            text(
+                "SELECT entry_id, distance FROM entry_embeddings "
+                "WHERE embedding MATCH :v AND k = :k "
+                "ORDER BY distance"
+            ),
+            {"v": blob, "k": int(k)},
+        ).all()
+        return [(r[0], float(r[1])) for r in rows]
+    except Exception as exc:
+        logger.warning("vector knn query failed: %s", exc)
+        return []
+
+
+def count(db: Session) -> Optional[int]:
+    if not _table_available(db):
+        return None
+    try:
+        return int(db.execute(text("SELECT COUNT(*) FROM entry_embeddings")).scalar() or 0)
+    except Exception:
+        return None

core/schemas/edge.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+import uuid
+from datetime import datetime, timezone
+from enum import Enum
+
+from pydantic import BaseModel, Field
+
+
+class EdgeRelation(str, Enum):
+    dependency = "dependency"
+    compatible_with = "compatible_with"
+    alternative_to = "alternative_to"
+    related_workflow = "related_workflow"
+    generated_from = "generated_from"
+    memory_of = "memory_of"
+    refinement_of = "refinement_of"
+    derived_from = "derived_from"
+    warning_about = "warning_about"
+    cited_by = "cited_by"
+    wikilink = "wikilink"
+    prerequisite = "prerequisite"
+    replacement = "replacement"
+    execution_pathway = "execution_pathway"
+    transformation = "transformation"
+    provenance = "provenance"
+    compatibility = "compatibility"
+    implements = "implements"
+    uses = "uses"
+    documents = "documents"
+    # Hierarchical-memory edges (progressive disclosure).
+    # Direction convention is always *from child detail → parent skill*,
+    # so that traversal "out of" a planner-level node yields its details.
+    decomposes_to = "decomposes_to"   # L1 → L2  (capability decomposed into procedure)
+    heuristic_for = "heuristic_for"   # L3 → L1/L2  (heuristic attached to a skill)
+    constraint_on = "constraint_on"   # L4 → L1/L2  (failure-mode attached to a skill)
+
+
+class Edge(BaseModel):
+    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    source_id: str
+    target_id: str
+    relation: EdgeRelation = EdgeRelation.wikilink
+    weight: float = 1.0
+    metadata: dict = Field(default_factory=dict)
+    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))