graph-engine-memory 0.6.1__py3-none-any.whl

graph_engine/server.py

@@ -0,0 +1,239 @@
+"""FastAPI server on top of the Brain (private git repo as storage).
+
+Env control:
+  IG_BRAIN_PATH   — path to the brain clone (default: ~/graph-engine-brain)
+  IG_BRAIN_REMOTE — SSH/GitHub URL (only used for `git clone` on first
+                    setup; no personal default, existing clones keep
+                    their own origin)
+  IG_BRAIN_MODE   — "git" (real repo) or "local" (FS only, for tests)
+  IDEAGRAPH_EMBEDDER — "st" (sentence-transformers) or "hash" (demo/tests)
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from fastapi.concurrency import run_in_threadpool
+from fastapi.responses import FileResponse, JSONResponse
+from pydantic import BaseModel
+
+from . import runtime, access, namespaces
+
+# Shipped UI assets live INSIDE the package so a pip install serves them too
+# (repo-root docs/ would not exist in site-packages).
+DOCS_DIR = Path(__file__).resolve().parent / "web"
+
+app = FastAPI(title="GraphEngine")
+
+
+@app.middleware("http")
+async def access_control(request, call_next):
+    # Pin selection before any awaits: switching the CLI while an ingest runs
+    # must not send the old graph's payload to the new graph's subscribers.
+    graph_token = namespaces.request_graph.set(namespaces.selected())
+    identity_token = None
+    try:
+        if access.policy_path():
+            try:
+                identity = access.token_principal(request.headers.get("authorization", ""))
+            except PermissionError:
+                access.audit("http", "", "authenticate", False)
+                return JSONResponse({"error": "Authentication required"}, status_code=401)
+            identity_token = access.request_principal.set(identity)
+        return await call_next(request)
+    finally:
+        if identity_token is not None:
+            access.request_principal.reset(identity_token)
+        namespaces.request_graph.reset(graph_token)
+
+
+@app.exception_handler(PermissionError)
+async def permission_error_handler(_request, _exc):
+    return JSONResponse({"error": "Access denied"}, status_code=403)
+
+
+
+class ConnectionManager:
+    def __init__(self):
+        self.active: list[WebSocket] = []
+        self.graphs: dict[WebSocket, str] = {}
+
+    async def connect(self, ws: WebSocket):
+        await ws.accept()
+        self.active.append(ws)
+        self.graphs[ws] = runtime.brain_path()
+
+    def disconnect(self, ws: WebSocket):
+        if ws in self.active:
+            self.active.remove(ws)
+        self.graphs.pop(ws, None)
+
+    async def broadcast(self, message: dict):
+        if access.policy_path():
+            return  # protected realtime transport is intentionally disabled
+        for ws in list(self.active):
+            if self.graphs.get(ws) != runtime.brain_path():
+                continue
+            try:
+                await ws.send_json(message)
+            except Exception:
+                self.disconnect(ws)
+
+
+manager = ConnectionManager()
+
+
+NO_CACHE = {"Cache-Control": "no-store"}
+
+
+@app.get("/")
+def index():
+    return FileResponse(DOCS_DIR / "index.html", headers=NO_CACHE)
+
+
+@app.get("/app.js")
+def app_js():
+    return FileResponse(DOCS_DIR / "app.js", headers=NO_CACHE)
+
+
+@app.get("/review")
+def review():
+    return FileResponse(DOCS_DIR / "review.html", headers=NO_CACHE)
+
+
+@app.get("/review.js")
+def review_js():
+    return FileResponse(DOCS_DIR / "review.js", headers=NO_CACHE)
+
+
+@app.get("/api/graph")
+def graph():
+    return JSONResponse(runtime.make_brain().graph_state())
+
+
+@app.get("/report")
+def report_page():
+    return FileResponse(DOCS_DIR / "report.html", headers=NO_CACHE)
+
+
+@app.get("/report.js")
+def report_js():
+    return FileResponse(DOCS_DIR / "report.js", headers=NO_CACHE)
+
+
+@app.get("/api/report")
+def api_report(since: str | None = None, top: int = 5):
+    """Read-only BRAIN_REPORT digest (report #7). Never mutates the brain."""
+    from .report import report_data, render_report
+    brain = runtime.make_brain()
+    try:
+        data = report_data(brain, since=since, top=top)
+        data["markdown"] = render_report(brain, since=since, top=top)
+        return JSONResponse(data)
+    finally:
+        if hasattr(brain, "close"):
+            brain.close()
+
+
+class IngestBody(BaseModel):
+    text: str
+    source: str = "human"
+    tags: list[str] = []
+    allow_duplicates: bool = False
+
+
+@app.exception_handler(ValueError)
+async def value_error_handler(_req, exc: ValueError):
+    """Audit #57: engine ValueErrors (e.g. empty ingest text) are client errors,
+    not 500s — return a clean 400 with the message."""
+    from fastapi.responses import JSONResponse
+    return JSONResponse({"error": str(exc)}, status_code=400)
+
+
+@app.post("/api/ingest")
+async def ingest(body: IngestBody):
+    # Audit #16: git pull/push + embedding inference are blocking I/O — they
+    # run in the threadpool, not on the event loop (otherwise a slow ingest
+    # stalls all requests and WS broadcasts).
+    engine = runtime.make_engine()
+    node, edges, is_dup = await run_in_threadpool(
+        engine.ingest, body.text, body.source, body.tags,
+        body.allow_duplicates)
+    await manager.broadcast({
+        "type": "ingested",
+        "node": node.to_dict(),
+        "edges": [e.to_dict() for e in edges],
+        "duplicate": is_dup,
+    })
+    return {"node": node.to_dict(), "suggested": [e.to_dict() for e in edges],
+            "duplicate": is_dup}
+
+
+@app.post("/api/edge/{edge_id}/accept")
+async def accept_edge(edge_id: str):
+    edge = await run_in_threadpool(runtime.make_engine().resolve, edge_id, True)
+    if edge is None:
+        return JSONResponse({"error": "edge not found or not pending"}, status_code=404)
+    await manager.broadcast({"type": "edge_resolved", "edge": edge.to_dict(), "accepted": True})
+    return edge.to_dict()
+
+
+@app.post("/api/edge/{edge_id}/reject")
+async def reject_edge(edge_id: str):
+    edge = await run_in_threadpool(runtime.make_engine().resolve, edge_id, False)
+    if edge is None:
+        return JSONResponse({"error": "edge not found or not pending"}, status_code=404)
+    await manager.broadcast({"type": "edge_resolved", "edge": edge.to_dict(), "accepted": False})
+    return {"rejected": edge_id}
+
+
+class LinkBody(BaseModel):
+    source: str
+    target: str
+    kind: str = "same_as"
+
+
+@app.post("/api/edge")
+async def link_edge(body: LinkBody):
+    def _link():
+        try:
+            return runtime.make_engine().link(body.source, body.target, body.kind), None
+        except ValueError as exc:
+            return None, str(exc)
+    edge, err = await run_in_threadpool(_link)
+    if err is not None or edge is None:
+        return JSONResponse({"error": err or "link failed"}, status_code=400)
+    await manager.broadcast({"type": "edge_linked", "edge": edge.to_dict()})
+    return edge.to_dict()
+
+
+@app.websocket("/ws")
+async def ws_endpoint(ws: WebSocket):
+    if access.policy_path():
+        await ws.close(code=1008)
+        return
+    await manager.connect(ws)
+    try:
+        while True:
+            # Audit #29: catching only WebSocketDisconnect left zombies behind
+            # (binary frame → KeyError, TCP reset → RuntimeError) — the socket
+            # stayed in manager.active forever and was never closed.
+            await ws.receive_text()
+    except WebSocketDisconnect:
+        manager.disconnect(ws)
+    except Exception:
+        manager.disconnect(ws)
+        try:
+            await ws.close()
+        except Exception:
+            pass
+
+
+@app.post("/api/edge/{edge_id}/undo")
+async def undo_edge(edge_id: str):
+    edge = await run_in_threadpool(runtime.make_engine().undo, edge_id)
+    if edge is None:
+        return JSONResponse({"error": "edge not found or already resolved"}, status_code=409)
+    await manager.broadcast({"type": "edge_restored", "edge": edge.to_dict()})
+    return edge.to_dict()

graph_engine/similarity.py

@@ -0,0 +1,42 @@
+"""Similarity: cosine similarity + k nearest neighbors."""
+
+from __future__ import annotations
+
+import math
+
+
+def cosine(a: list[float], b: list[float]) -> float:
+    """Cosine similarity. Audit #8: mismatched dimensions are a data error
+    (e.g. 384-dim ST vector vs. 64-dim hash vector in one brain) — a silent
+    truncation to the shorter length produces plausible-looking scores and can
+    trigger a wrong auto-merge at the dedup threshold. Therefore: hard fail."""
+    if len(a) != len(b):
+        raise ValueError(
+            f"cosine: dimension mismatch ({len(a)} vs {len(b)}) — "
+            "inhomogeneous embedders in the same brain? Vectors are not comparable.")
+    dot = sum(x * y for x, y in zip(a, b))
+    na = math.sqrt(sum(x * x for x in a))
+    nb = math.sqrt(sum(x * x for x in b))
+    if na == 0.0 or nb == 0.0:
+        return 0.0
+    return dot / (na * nb)
+
+
+def knn(query: list[float], candidates: dict[str, list[float]], k: int = 3) -> list[tuple[str, float]]:
+    """Returns the k nearest neighbors as (id, similarity), sorted descending.
+
+    Audit #8 follow-up (fix wave 2): foreign-dimension candidates are
+    skipped instead of crashing the whole query — the same two-level
+    decision as in _find_duplicate: the primitive cosine() is strict,
+    the call sites are tolerant. A brain with a few legacy vectors of the
+    wrong dimension (e.g. after an embedder switch) degrades cleanly to the
+    compatible neighbors.
+    Audit #60: k<=0 returns [] instead of all items (k=0) or the last one
+    being dropped (k=-1) — a limit means limit."""
+    if k <= 0:
+        return []
+    query_dim = len(query)
+    scored = [(nid, cosine(query, vec)) for nid, vec in candidates.items()
+              if vec and len(vec) == query_dim]
+    scored.sort(key=lambda t: t[1], reverse=True)
+    return scored[:k]

graph_engine/store.py

@@ -0,0 +1,63 @@
+"""JSONL store: one JSON line per object. No DB overhead."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from .model import Node, Edge
+
+
+class Store:
+    """Holds nodes and edges as JSONL (separate files)."""
+
+    def __init__(self, path: str | Path = "data.jsonl"):
+        self.path = Path(path)
+        self.edges_path = self.path.with_suffix(".edges.jsonl")
+        self.nodes: list[Node] = []
+        self.edges: list[Edge] = []
+        self._load()
+
+    def _load(self) -> None:
+        if self.path.exists():
+            for line in self.path.read_text(encoding="utf-8").splitlines():
+                if line.strip():
+                    self.nodes.append(Node.from_dict(json.loads(line)))
+        if self.edges_path.exists():
+            for line in self.edges_path.read_text(encoding="utf-8").splitlines():
+                if line.strip():
+                    self.edges.append(Edge.from_dict(json.loads(line)))
+
+    def add_node(self, node: Node) -> None:
+        self.nodes.append(node)
+        with self.path.open("a", encoding="utf-8") as f:
+            f.write(json.dumps(node.to_dict(), ensure_ascii=False) + "\n")
+
+    def add_edge(self, edge: Edge) -> None:
+        self.edges.append(edge)
+        with self.edges_path.open("a", encoding="utf-8") as f:
+            f.write(json.dumps(edge.to_dict(), ensure_ascii=False) + "\n")
+
+    def resolve_edge(self, edge_id: str, accept: bool) -> Edge | None:
+        """Accept a pending edge (pending=False) or discard it (remove)."""
+        edge = next((e for e in self.edges if e.id == edge_id), None)
+        if edge is None or not edge.pending:
+            return None
+        if accept:
+            edge.pending = False
+            self._rewrite_edges()
+            return edge
+        self.edges.remove(edge)
+        self._rewrite_edges()
+        return edge
+
+    def _rewrite_edges(self) -> None:
+        with self.edges_path.open("w", encoding="utf-8") as f:
+            for e in self.edges:
+                f.write(json.dumps(e.to_dict(), ensure_ascii=False) + "\n")
+
+    def graph_state(self) -> dict:
+        return {
+            "nodes": [n.to_dict() for n in self.nodes],
+            "edges": [e.to_dict() for e in self.edges],
+        }

graph_engine/suggester.py

@@ -0,0 +1,63 @@
+"""Edge suggestions: the k nearest neighbors become typed suggestions.
+
+Rules (V2#3 — confidence bands instead of bare thresholds):
+- sim >= 0.95            → "similar", AUTO-ACCEPT (pending=False), confidence=sim
+- 0.75 <= sim < 0.95     → "similar", pending
+- 0.45 <= sim < 0.75     → "extends", pending
+- sim < 0.45             → no suggestion
+
+Every suggestion carries a confidence (the cosine similarity). The caller
+(BrainEngine) decides from the band + env override whether it stays pending.
+
+"same_as" is never suggested automatically — it is created manually
+(CLI `link`, cockpit) for translation/alias pairs.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from .similarity import knn
+
+THRESHOLD_SIMILAR = 0.75
+THRESHOLD_EXTEND = 0.45
+AUTO_ACCEPT_CONFIDENCE = 0.95
+
+
+@dataclass
+class Suggestion:
+    source: str
+    target: str
+    kind: str
+    confidence: float = 0.0
+
+
+def suggest(source_id: str, query_vec: list[float],
+            candidates: dict[str, list[float]], k: int = 3) -> list[Suggestion]:
+    out: list[Suggestion] = []
+    seen: set[str] = set()
+    for nid, sim in knn(query_vec, candidates, k):
+        if nid == source_id or nid in seen:
+            continue
+        if sim >= THRESHOLD_SIMILAR:
+            kind = "similar"
+        elif sim >= THRESHOLD_EXTEND:
+            kind = "extends"
+        else:
+            continue
+        seen.add(nid)
+        out.append(Suggestion(source=source_id, target=nid, kind=kind, confidence=sim))
+    return out
+
+
+def is_auto_accept(confidence: float) -> bool:
+    """Confidence band: >=0.95 is accepted without HITL."""
+    return confidence >= AUTO_ACCEPT_CONFIDENCE
+
+
+# Backwards-compatible for the JSONL engine (v0.0.1)
+def suggest_edges(source_id: str, query_vec: list[float],
+                  candidates: dict[str, list[float]], k: int = 3):
+    from .model import Edge
+    return [Edge(source=s.source, target=s.target, kind=s.kind, pending=True)
+            for s in suggest(source_id, query_vec, candidates, k)]

graph_engine/temporal.py

@@ -0,0 +1,216 @@
+"""Temporal reasoning (Phase 2): bi-temporal queries over the graph.
+
+The engine has bi-temporal edges (valid_from/valid_to) since the Zep/Graphiti
+lesson, but no way to QUERY them. This module adds the read side:
+
+  valid_at(brain, timestamp)  — the graph as it was at a point in time
+  history(brain, node_id)     — how a node's edges evolved over time
+  when(engine, query, at)     — retrieval restricted to what was known then
+
+All functions are read-only. They never modify the brain.
+"""
+
+from __future__ import annotations
+
+import datetime
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .brain import Brain, Edge, Node
+    from .brain_engine import BrainEngine
+
+
+def _parse_ts(ts: str) -> datetime.datetime:
+    """Parse ISO-8601 timestamp, tolerating 'Z' suffix."""
+    return datetime.datetime.fromisoformat(ts.replace("Z", "+00:00"))
+
+
+def _is_valid_at(edge: "Edge", ts: datetime.datetime) -> bool:
+    """Was this edge valid at the given timestamp?
+
+    An edge is valid at T when:
+    - valid_from <= T (it had been created by then)
+    - valid_to is None OR valid_to > T (its validity had not ended)
+    - invalidated_at is None OR invalidated_at > T
+    """
+    return edge.is_valid_at(ts)
+
+
+def valid_at(brain: "Brain", timestamp: str) -> dict:
+    """The graph as it was at a point in time.
+
+    Returns {nodes: [...], edges: [...]} — only nodes and edges that were
+    live at the given timestamp. Nodes are included if they existed (created
+    <= T) and were not yet tombstoned. Edges are included if they were valid
+    at T (see _is_valid_at).
+    """
+    try:
+        ts = _parse_ts(timestamp)
+    except (ValueError, TypeError):
+        raise ValueError(f"Invalid timestamp: {timestamp!r}")
+
+    nodes: list[Node] = []
+    for n in brain.read_nodes():
+        if n.status == "tombstone":
+            continue
+        try:
+            created = _parse_ts(n.created) if n.created else None
+        except (ValueError, TypeError):
+            created = None
+        if created and created > ts:
+            continue
+        nodes.append(n)
+
+    edges: list[Edge] = []
+    for e in brain.read_edges(include_rejected=True):
+        if e.pending or e.rejected:
+            continue
+        if _is_valid_at(e, ts):
+            edges.append(e)
+
+    return {"nodes": nodes, "edges": edges}
+
+
+def history(brain: "Brain", node_id: str) -> list[dict]:
+    """How a node's edges evolved over time.
+
+    Returns a chronological list of edge events (creation, expiry, invalidation)
+    involving this node, sorted by timestamp. Each entry has:
+    {timestamp, event, edge_id, kind, other_node, origin}
+    """
+    events: list[dict] = []
+    for e in brain.read_edges(include_rejected=True):
+        if node_id not in (e.source, e.target):
+            continue
+        other = e.target if e.source == node_id else e.source
+        # Creation event
+        events.append({
+            "timestamp": e.valid_from,
+            "event": "created",
+            "edge_id": e.id,
+            "kind": e.kind,
+            "other_node": other,
+            "origin": e.origin,
+            "rejected": e.rejected,
+        })
+        # Keep an extracted expiry separate from a later review invalidation.
+        endings = [(e.valid_to, "expired" if e.extracted_validity is not None else "invalidated")]
+        if e.extracted_validity is not None:
+            endings.append((e.invalidated_at, "invalidated"))
+        for timestamp, event in endings:
+            if not timestamp:
+                continue
+            events.append({
+                "timestamp": timestamp,
+                "event": event,
+                "edge_id": e.id,
+                "kind": e.kind,
+                "other_node": other,
+                "origin": e.origin,
+                "rejected": e.rejected,
+            })
+    events.sort(key=lambda x: x["timestamp"] or "")
+    return events
+
+
+def when(engine: "BrainEngine", query: str, at: str, k: int = 5) -> list[tuple[str, float]]:
+    """Retrieval restricted to what was known at a point in time (Phase 2).
+
+    Runs the normal hybrid retrieval, then filters the results to only
+    include nodes that existed at the given timestamp. This is the
+    "what did I know then?" query — useful for understanding how the
+    brain's knowledge has evolved.
+    """
+    from .retrieval import retrieve
+
+    # Parse timestamp FIRST — an invalid timestamp must raise even when
+    # there are no results to filter.
+    try:
+        ts = _parse_ts(at)
+    except (ValueError, TypeError):
+        raise ValueError(f"Invalid timestamp: {at!r}")
+
+    # Get the full ranking (no tracking — this is a historical query)
+    results = retrieve(engine, query, k=k * 3, track=False)
+    if not results:
+        return []
+
+    # Filter to nodes that existed at the timestamp
+    filtered: list[tuple[str, float]] = []
+    for node_id, score in results:
+        node = engine.brain.read_node(node_id)
+        if node is None:
+            continue
+        if node.status == "tombstone":
+            continue
+        try:
+            created = _parse_ts(node.created) if node.created else None
+        except (ValueError, TypeError):
+            created = None
+        if created and created > ts:
+            continue
+        filtered.append((node_id, score))
+        if len(filtered) >= k:
+            break
+
+    return filtered
+
+
+def edge_timeline(brain: "Brain", *, since: str | None = None,
+                  until: str | None = None) -> list[dict]:
+    """All edge events (created/expired/invalidated) in a time range.
+
+    Useful for "what changed in the brain recently?" — the temporal
+    equivalent of `ig report --since`.
+    """
+    events: list[dict] = []
+    for e in brain.read_edges(include_rejected=True):
+        if e.pending or e.rejected:
+            continue
+        try:
+            vf = _parse_ts(e.valid_from) if e.valid_from else None
+        except (ValueError, TypeError):
+            vf = None
+
+        if vf:
+            if since and vf < _parse_ts(since):
+                pass  # still include, filtered below
+            if until and vf > _parse_ts(until):
+                pass
+            events.append({
+                "timestamp": e.valid_from,
+                "event": "created",
+                "edge_id": e.id,
+                "kind": e.kind,
+                "source": e.source,
+                "target": e.target,
+                "origin": e.origin,
+            })
+        endings = [(e.valid_to, "expired" if e.extracted_validity is not None else "invalidated")]
+        if e.extracted_validity is not None:
+            endings.append((e.invalidated_at, "invalidated"))
+        for timestamp, event in endings:
+            if not timestamp:
+                continue
+            try:
+                _parse_ts(timestamp)
+            except (ValueError, TypeError):
+                continue
+            events.append({
+                "timestamp": timestamp,
+                "event": event,
+                "edge_id": e.id,
+                "kind": e.kind,
+                "source": e.source,
+                "target": e.target,
+                "origin": e.origin,
+            })
+
+    # Filter by range
+    if since:
+        events = [e for e in events if (e["timestamp"] or "") >= since]
+    if until:
+        events = [e for e in events if (e["timestamp"] or "") <= until]
+
+    events.sort(key=lambda x: x["timestamp"] or "")
+    return events