loop-memory 0.4.0__py3-none-any.whl

loop_memory/graph/build.py

@@ -0,0 +1,259 @@
+"""Build the knowledge graph from the memories table.
+
+``KnowledgeGraph.rebuild()`` walks every memory row, runs the
+zero-dep entity extractor on its text, and:
+
+1. Upserts each entity row in ``entities`` (name, kind, weight, count).
+2. Computes pairwise co-occurrence within a sliding window per memory
+   and stores them as ``co_occurs_with`` relations in ``relations``.
+3. Records which memory id each relation is evidenced by so the UI
+   can highlight the underlying chunks.
+
+Plug in an LLM extractor for higher quality by replacing the ``extract``
+function via the ``KnowledgeGraph(extractor=...)`` constructor — the
+function signature is ``(text: str) -> list[tuple[name, kind]]``.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from collections.abc import Callable
+from dataclasses import dataclass
+
+from ..storage.sqlite_store import MemoryStore
+from .extract import extract_entities
+
+log = logging.getLogger(__name__)
+
+
+ExtractorFn = Callable[[str], list[tuple[str, str]]]
+
+
+@dataclass
+class BuildReport:
+    entities: int = 0
+    relations: int = 0
+    memories_scanned: int = 0
+    elapsed_ms: float = 0.0
+
+
+class KnowledgeGraph:
+    def __init__(
+        self,
+        store: MemoryStore,
+        extractor: ExtractorFn | None = None,
+        window: int = 6,
+        min_mention: int = 1,
+    ) -> None:
+        self.store = store
+        self.extractor = extractor or extract_entities
+        self.window = window
+        self.min_mention = min_mention
+
+    def rebuild(
+        self,
+        *,
+        clear: bool = False,
+        limit: int | None = None,
+    ) -> BuildReport:
+        t0 = time.time()
+        report = BuildReport()
+        if clear:
+            removed = self.store.delete_graph()
+            log.info("cleared graph: removed %s entities", removed)
+
+        rows = self.store.list_memories(limit=limit or 100_000)
+        report.memories_scanned = len(rows)
+        for row in rows:
+            ents = self.extractor(row.text or "")
+            ents = [(n, k) for (n, k) in ents if len(n) <= 32]
+            for name, kind in ents:
+                self.store.upsert_entity(name, kind, bump_weight=0.02)
+            # co-occurrence per memory text
+            if ents:
+                names = [n for (n, _) in ents]
+                pairs = self._pairs_in_window(names, window=self.window)
+                for a, b in pairs:
+                    self.store.upsert_relation(
+                        a, b, kind="co_occurs_with",
+                        weight=0.5, evidence_id=row.id,
+                    )
+                    report.relations += 1
+        stats = self.store.graph_stats()
+        report.entities = stats["entities"]
+        report.relations = stats["relations"]
+        report.elapsed_ms = (time.time() - t0) * 1000
+        return report
+
+    def _pairs_in_window(self, names, *, window: int) -> list[tuple[str, str]]:
+        # Local co-occurrence within the same memory text.
+        seen = set()
+        out: list[tuple[str, str]] = []
+        for i, a in enumerate(names):
+            for b in names[i + 1 : i + window]:
+                if a == b:
+                    continue
+                key = tuple(sorted([a, b]))
+                if key in seen:
+                    continue
+                seen.add(key)
+                out.append((a, b))
+        return out
+
+    # ------------------------------------------------------------------
+    # Wiki-based graph: build the knowledge graph from the *distilled*
+    # wiki pages rather than the raw memories. The result is denser,
+    # cleaner, and easier to navigate because every node represents a
+    # topic the user has already validated through consolidation.
+    # ------------------------------------------------------------------
+
+    def rebuild_from_wiki(
+        self,
+        *,
+        clear: bool = False,
+        limit: int | None = None,
+    ) -> BuildReport:
+        """Build the graph from ``wiki_pages`` instead of raw memories.
+
+        Nodes created (with ``kind`` so the UI can color them):
+
+          * ``wiki:<slug>``  — one node per wiki page
+          * ``tag:<name>``   — one node per tag used across pages
+          * ``concept:<name>`` — entities extracted from the page body
+
+        Edges created:
+
+          * ``wiki --tagged_with--> tag``            (per tag in page)
+          * ``wiki --mentions--> concept``           (entities in body)
+          * ``wiki --related_to--> wiki``            (pages sharing a tag or concept)
+
+        Each relation carries the page id (and any source-memory ids)
+        as evidence so the UI can drill back to the original chunks.
+        """
+        t0 = time.time()
+        report = BuildReport()
+        if clear:
+            removed = self.store.delete_graph()
+            log.info("cleared graph: removed %s entities", removed)
+
+        pages = self.store.list_wiki_pages(limit=limit or 1000)
+        report.memories_scanned = len(pages)
+        if not pages:
+            log.info("rebuild_from_wiki: no wiki pages; skipping")
+            report.elapsed_ms = (time.time() - t0) * 1000
+            return report
+
+        # 1) Page nodes
+        page_node_ids: Dict[str, str] = {}
+        for page in pages:
+            slug = (page.get("slug") or "").strip()
+            if not slug:
+                continue
+            importance = float(page.get("importance") or 0.5)
+            # Wiki nodes start at importance and accumulate tiny weight
+            # from shared tags/concepts — they should dominate visually.
+            ent = self.store.upsert_entity(
+                f"wiki:{slug}", kind="wiki_page",
+                bump_weight=max(0.0, importance - 0.5),
+            )
+            page_node_ids[slug] = ent.name
+            # Bump by mention_count too so a page with N tags/concepts
+            # is more prominent than a lonely one.
+            for _ in range(min(8, len(page.get("tags") or []) + 1)):
+                self.store.upsert_entity(
+                    f"wiki:{slug}", kind="wiki_page", bump_weight=0.01,
+                )
+
+        # 2) Tag nodes + tag edges
+        tag_to_pages: Dict[str, list[str]] = {}
+        for page in pages:
+            tags = page.get("tags") or []
+            slug = page.get("slug") or ""
+            if not slug:
+                continue
+            for raw in tags:
+                tag = (raw or "").strip().lower()
+                if not tag:
+                    continue
+                self.store.upsert_entity(
+                    f"tag:{tag}", kind="tag", bump_weight=0.05,
+                )
+                self.store.upsert_relation(
+                    f"wiki:{slug}", f"tag:{tag}",
+                    kind="tagged_with",
+                    weight=0.7,
+                    evidence_id=page.get("id"),
+                )
+                report.relations += 1
+                tag_to_pages.setdefault(tag, []).append(slug)
+
+        # 3) Concept entities from page body + title + summary
+        concept_to_pages: Dict[str, list[str]] = {}
+        for page in pages:
+            slug = page.get("slug") or ""
+            if not slug:
+                continue
+            text_chunks = [
+                page.get("title") or "",
+                page.get("summary") or "",
+                page.get("body") or "",
+            ]
+            text = chr(10).join(t for t in text_chunks if t)
+            ents = self.extractor(text)
+            seen_here: set = set()
+            for name, _kind in ents:
+                name = (name or "").strip()
+                if not name or len(name) > 32:
+                    continue
+                if name in seen_here:
+                    continue
+                seen_here.add(name)
+                self.store.upsert_entity(
+                    f"concept:{name}", kind="concept", bump_weight=0.02,
+                )
+                self.store.upsert_relation(
+                    f"wiki:{slug}", f"concept:{name}",
+                    kind="mentions",
+                    weight=0.4,
+                    evidence_id=page.get("id"),
+                )
+                report.relations += 1
+                concept_to_pages.setdefault(name, []).append(slug)
+
+        # 4) Wiki --related_to--> wiki when pages share a tag or concept
+        def _relate(a: str, b: str, evidence_id: str | None) -> None:
+            if a == b:
+                return
+            # Insert in both directions so the UI can highlight a↔b
+            # without having to do its own lookup.
+            for x, y in ((a, b), (b, a)):
+                self.store.upsert_relation(
+                    f"wiki:{x}", f"wiki:{y}",
+                    kind="related_to",
+                    weight=0.3,
+                    evidence_id=evidence_id,
+                )
+                report.relations += 1
+
+        for tag, slugs in tag_to_pages.items():
+            slugs = list(dict.fromkeys(slugs))
+            for i, a in enumerate(slugs):
+                for b in slugs[i + 1:]:
+                    _relate(a, b, evidence_id=None)
+        for _concept, slugs in concept_to_pages.items():
+            slugs = list(dict.fromkeys(slugs))
+            if len(slugs) <= 8:
+                for i, a in enumerate(slugs):
+                    for b in slugs[i + 1:]:
+                        _relate(a, b, evidence_id=None)
+
+        stats = self.store.graph_stats()
+        report.entities = stats["entities"]
+        report.relations = stats["relations"]
+        report.elapsed_ms = (time.time() - t0) * 1000
+        log.info(
+            "rebuild_from_wiki: %d pages -> %d entities, %d relations in %.1fms",
+            len(pages), report.entities, report.relations, report.elapsed_ms,
+        )
+        return report

loop_memory/graph/extract.py

@@ -0,0 +1,197 @@
+"""Lightweight entity + co-occurrence extractor (zero-dep, heuristic).
+
+We don't pull in spaCy here. This module produces two things:
+
+* ``extract_entities(text)`` → list of (name, kind) candidates, deduped
+  and filtered by stopwords.  Recognised entities include Capitalised
+  Latin tokens, ``#hashtag`` style tokens, file paths / URLs, CamelCase
+  product names, and short CJK noun-like runs (after stopword filter).
+
+* ``pair_cooccurrence(texts, window=...)`` → list of (a, b, count)
+  relations where ``a`` and ``b`` appear within ``window`` tokens of
+  each other in any of the input texts.
+
+This is intentionally tunable — for higher quality call out to an
+LLM-backed extractor and feed the output through the same API.
+"""
+
+from __future__ import annotations
+
+import re
+from collections import Counter
+from collections.abc import Iterable
+
+# --- tokenisation -----------------------------------------------------------
+
+_TOKEN_RE = re.compile(
+    r"[A-Za-z][A-Za-z0-9_]+"                  # Latin word
+    r"|\#[\w\u4e00-\u9fff]+"                  # #hashtag (English / CJK mix)
+    r"|https?://[^\s]+"                       # URL
+    r"|[A-Za-z0-9_./-]+\.[A-Za-z0-9]{2,}"      # file.ext or domain.tld
+    r"|[\u4e00-\u9fff]{2,8}"                  # CJK noun runs (2–8 chars)
+)
+
+
+def tokenize(text: str) -> list[str]:
+    return _TOKEN_RE.findall(text or "")
+
+
+# --- stopword lists ---------------------------------------------------------
+
+_EN_STOPWORDS = {
+    "the", "and", "for", "with", "from", "into", "this", "that", "these",
+    "those", "you", "your", "are", "was", "were", "have", "has", "had",
+    "but", "not", "any", "all", "can", "could", "would", "should", "will",
+    "shall", "may", "might", "must", "what", "why", "how", "who", "where",
+    "when", "then", "now", "also", "still", "just", "very", "more",
+    "less", "most", "least", "out", "off", "per", "via", "i", "me",
+    "my", "we", "us", "our", "they", "their", "them", "he", "she", "his",
+    "her", "is", "am", "be", "been", "being", "do", "does", "did", "done",
+    "doing", "of", "in", "on", "at", "to", "by", "as", "an", "or", "if",
+    "no", "yes", "so", "it", "its", "about", "than", "there", "here",
+    "above", "below", "under", "over", "again", "once", "each",
+    "something", "anything", "everything", "nothing", "some",
+    "which", "such", "after", "before",
+}
+
+# Common short CJK fragments that aren't real concepts.
+_CJK_STOPWORDS = {
+    "的", "了", "和", "是", "在", "我", "你", "他", "她", "它",
+    "我们", "你们", "他们", "这个", "那个", "什么", "怎么", "为什么",
+    "可以", "可能", "也许", "应该", "已经", "现在", "之前", "之后",
+    "因为", "所以", "如果", "但是", "不过", "然后", "可是", "而且",
+    "或者", "还有", "也", "都", "就", "才", "只", "再", "又", "很",
+    "非常", "比较", "一点", "一下", "一直", "顺便", "帮我", "我用",
+    "你用", "对她", "我对", "是不是",
+    "用", "打", "做", "搞", "弄", "给", "让", "把", "被", "由",
+    "从", "到", "向", "对", "跟", "比", "如", "若", "虽", "除非",
+    "将", "会", "能", "须", "必", "得", "地", "着", "过", "如何", "怎样", "为啥", "的工具", "的项目", "的系统", "的代码", "的内容", "的功能",
+    "一个", "一些", "这些", "那些", "今天",
+    "明天", "昨天", "谢谢", "感谢", "麻烦", "请帮",
+}
+
+_STOPWORDS = _EN_STOPWORDS
+
+
+# --- detection --------------------------------------------------------------
+
+
+
+
+# Extra stopwords: tokens that come up very frequently as filler in
+# extracted LLM-style text — they tend to drown out useful entities.
+_GENERIC_TERMS = {
+    "User", "Users", "Assistant", "Outcome", "Task", "Description",
+    "Issues", "Issue", "Files", "File", "Strengths", "Added",
+    "Required", "Implemented", "Complete", "Check", "Review",
+    "Code", "Title", "Body", "Read", "Note", "Source", "Target", "Message", "Context", "Result", "Output", "Input", "Step", "Steps", "List", "Section",
+    "Project", "Repository", "Repo", "Doc", "Documentation",
+    "Docs", "URL", "Path", "Line", "Lines", "Type", "Method",
+    "Function", "Class", "Module", "Package", "Import", "Export",
+    "Example", "Examples", "Sample", "Demo", "Use", "Using", "Used",
+    "Make", "Made", "Run", "Running", "Create", "Creates",
+    "Created", "Add", "Adding", "Remove", "Removed",
+    "Removing", "Update", "Updated", "Updating", "Change", "Changes",
+    "Changed", "Show", "Showing", "Found", "Founding", "Missing",
+    "First", "Second", "Third",
+    "Test", "Tests", "Tested", "Testing", "Pass", "Passes",
+    "Status", "OK", "Ok", "Notes", "Comments", "Comment",
+    "Done", "Completed", "Completion",
+}
+
+def _looks_like_proper_noun(tok: str) -> bool:
+    """True for tokens that may be useful named-entity material."""
+    if len(tok) < 2:
+        return False
+    if tok.isupper() and len(tok) <= 5:
+        return tok.isalpha()  # ACRONYM
+    if tok[0].isupper() and any(c.islower() for c in tok[1:]):
+        return True
+    if any(c.isupper() for c in tok[1:]) and any(c.islower() for c in tok):
+        return True  # CamelCase
+    return False
+
+
+def _kind_for(tok: str) -> str:
+    if tok.startswith("#"):
+        return "tag"
+    if tok.startswith("http"):
+        return "url"
+    if "." in tok and "/" in tok:
+        return "path"
+    if tok.isupper():
+        return "acronym"
+    if any(ord(c) > 127 for c in tok):
+        return "cjk"
+    return "concept"
+
+
+# Heuristic disqualifiers for CJK tokens: any of these characters in
+# any position is a strong signal the n-gram is not a real concept.
+_CJK_FRAGMENT_MARKERS = set("的了着过得把被让给向将从跟比和或而但所以因为")     | set("是吗呀啊嘛呢哦嗯哈呀嘛啊")
+
+
+def _is_useful(tok: str) -> bool:
+    if not tok:
+        return False
+    lower = tok.lower()
+    if lower in _STOPWORDS:
+        return False
+    if tok in _CJK_STOPWORDS:
+        return False
+    if tok in _GENERIC_TERMS:
+        return False
+    # Latin: only proper-nounish tokens qualify.
+    if ord(tok[0]) < 128:
+        if not _looks_like_proper_noun(tok):
+            return False
+        return True
+    # CJK token refinement
+    if any(c in _CJK_FRAGMENT_MARKERS for c in tok):
+        return False
+    return True
+
+
+def extract_entities(text: str, *, min_count: int = 1) -> list[tuple[str, str]]:
+    if not text:
+        return []
+    counts: Counter = Counter()
+    for tok in tokenize(text):
+        tok = tok.strip("'\"`")
+        if _is_useful(tok):
+            counts[tok] += 1
+    return [
+        (tok, _kind_for(tok))
+        for tok, n in counts.items()
+        if n >= min_count
+    ]
+
+
+def pair_cooccurrence(
+    texts: Iterable[str],
+    *,
+    window: int = 6,
+    min_count: int = 1,
+) -> list[tuple[str, str, int]]:
+    pair_counts: Counter = Counter()
+    for text in texts:
+        toks = [tok for tok in tokenize(text) if _is_useful(tok)]
+        seen_local: set = set()
+        for i, tok in enumerate(toks):
+            for other in toks[i + 1 : i + window]:
+                if other == tok:
+                    continue
+                key = tuple(sorted([tok, other]))
+                if key in seen_local:
+                    continue
+                seen_local.add(key)
+                pair_counts[key] += 1
+    return [(a, b, c) for (a, b), c in pair_counts.items() if c >= min_count]
+
+
+def canonical(name: str) -> str:
+    if not name:
+        return name
+    if ord(name[0]) >= 128:
+        return name
+    return name.strip()