pycontextdb 0.1.0__py3-none-any.whl

contextdb/graphs/base.py

@@ -0,0 +1,46 @@
+"""Shared scaffolding for graph implementations.
+
+Each graph stores its edges in its own SQLite table (same file as the memory
+store) so the free-tier deployment stays zero-config. :class:`BaseGraph`
+defines the contract and supplies a small helper for table creation.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from contextdb.core.models import Edge
+    from contextdb.store.sqlite_store import SQLiteStore
+
+
+class BaseGraph(ABC):
+    """Common contract for all graphs over memory items."""
+
+    def __init__(self, store: SQLiteStore) -> None:
+        self.store = store
+
+    @abstractmethod
+    async def initialize(self) -> None:
+        """Create tables / indices. Idempotent."""
+
+    @abstractmethod
+    async def add_node(self, memory_id: str, data: dict[str, Any]) -> None: ...
+
+    @abstractmethod
+    async def add_edge(self, edge: Edge) -> None: ...
+
+    @abstractmethod
+    async def get_neighbors(
+        self,
+        memory_id: str,
+        depth: int = 1,
+        max_results: int = 20,
+    ) -> list[tuple[str, float]]: ...
+
+    @abstractmethod
+    async def remove_node(self, memory_id: str) -> None: ...
+
+    @abstractmethod
+    async def get_edges(self, memory_id: str) -> list[Edge]: ...

contextdb/graphs/causal.py

@@ -0,0 +1,224 @@
+"""Causal graph — LLM-inferred cause/effect links between memories.
+
+The causal relationship is inherently a judgement call, so we ask the LLM
+for a structured verdict over a pair of memories. Edges carry a confidence
+weight in ``[0, 1]`` and a free-form ``reasoning`` payload.
+"""
+
+from __future__ import annotations
+
+import json
+from datetime import datetime
+from typing import TYPE_CHECKING, Any
+
+from contextdb.core.models import Edge
+from contextdb.graphs.base import BaseGraph
+
+if TYPE_CHECKING:
+    from contextdb.store.sqlite_store import SQLiteStore
+    from contextdb.utils.llm import LLMProvider
+
+_SCHEMA = """
+CREATE TABLE IF NOT EXISTS causal_edges (
+    source_id TEXT NOT NULL,
+    target_id TEXT NOT NULL,
+    relation TEXT NOT NULL,
+    weight REAL NOT NULL,
+    reasoning TEXT DEFAULT '',
+    metadata TEXT DEFAULT '{}',
+    created_at TEXT NOT NULL,
+    PRIMARY KEY (source_id, target_id)
+);
+CREATE INDEX IF NOT EXISTS idx_causal_source ON causal_edges(source_id);
+CREATE INDEX IF NOT EXISTS idx_causal_target ON causal_edges(target_id);
+"""
+
+_INFER_PROMPT = """You are a reasoning engine. Decide whether memory A caused,
+enabled, or has no causal link to memory B. Return strict JSON.
+
+Schema:
+{"relation": "CAUSES|ENABLES|NONE", "confidence": 0.0-1.0, "reasoning": "string"}
+
+Memory A (earlier): "{a_content}"
+Memory B (later):   "{b_content}"
+"""
+
+
+def _safe_json(text: str) -> dict[str, Any]:
+    text = text.strip()
+    if text.startswith("```"):
+        lines = text.splitlines()
+        text = "\n".join(line for line in lines if not line.startswith("```"))
+    try:
+        loaded = json.loads(text)
+        return loaded if isinstance(loaded, dict) else {}
+    except json.JSONDecodeError:
+        start = text.find("{")
+        end = text.rfind("}")
+        if start != -1 and end != -1 and end > start:
+            try:
+                loaded = json.loads(text[start : end + 1])
+                return loaded if isinstance(loaded, dict) else {}
+            except json.JSONDecodeError:
+                return {}
+        return {}
+
+
+class CausalGraph(BaseGraph):
+    """LLM-inferred cause/effect edges over memory items."""
+
+    def __init__(
+        self,
+        store: SQLiteStore,
+        llm: LLMProvider,
+        confidence_threshold: float = 0.5,
+        candidate_window: int = 10,
+    ) -> None:
+        super().__init__(store)
+        self.llm = llm
+        self.confidence_threshold = confidence_threshold
+        self.candidate_window = candidate_window
+
+    async def initialize(self) -> None:
+        conn = self.store._require_conn()
+        await conn.executescript(_SCHEMA)
+        await conn.commit()
+
+    async def add_node(self, memory_id: str, data: dict[str, Any]) -> None:
+        content = data.get("content", "")
+        if not content:
+            return
+        candidates = await self._recent_candidates(memory_id)
+        for other_id, other_content in candidates:
+            verdict = await self._infer(other_content, content)
+            relation = str(verdict.get("relation", "NONE")).upper()
+            confidence = float(verdict.get("confidence", 0.0))
+            if relation == "NONE" or confidence < self.confidence_threshold:
+                continue
+            reasoning = str(verdict.get("reasoning", ""))
+            await self.add_edge(
+                Edge(
+                    source_id=other_id,
+                    target_id=memory_id,
+                    graph_type="causal",
+                    weight=confidence,
+                    metadata={"relation": relation, "reasoning": reasoning},
+                )
+            )
+
+    async def _infer(self, a_content: str, b_content: str) -> dict[str, Any]:
+        prompt = _INFER_PROMPT.replace("{a_content}", a_content).replace(
+            "{b_content}", b_content
+        )
+        response = await self.llm.generate(prompt, temperature=0.0, max_tokens=200)
+        return _safe_json(response)
+
+    async def _recent_candidates(self, memory_id: str) -> list[tuple[str, str]]:
+        conn = self.store._require_conn()
+        cursor = await conn.execute(
+            "SELECT id, content FROM memories WHERE id != ? AND status = 'ACTIVE' "
+            "ORDER BY created_at DESC LIMIT ?",
+            (memory_id, self.candidate_window),
+        )
+        rows = await cursor.fetchall()
+        return [(row["id"], row["content"]) for row in rows]
+
+    async def add_edge(self, edge: Edge) -> None:
+        conn = self.store._require_conn()
+        meta = edge.metadata or {}
+        await conn.execute(
+            "INSERT OR REPLACE INTO causal_edges "
+            "(source_id, target_id, relation, weight, reasoning, metadata, created_at) "
+            "VALUES (?,?,?,?,?,?,?)",
+            (
+                edge.source_id,
+                edge.target_id,
+                str(meta.get("relation", "CAUSES")),
+                edge.weight,
+                str(meta.get("reasoning", "")),
+                json.dumps(meta),
+                edge.created_at.isoformat(),
+            ),
+        )
+        await conn.commit()
+
+    async def remove_node(self, memory_id: str) -> None:
+        conn = self.store._require_conn()
+        await conn.execute(
+            "DELETE FROM causal_edges WHERE source_id = ? OR target_id = ?",
+            (memory_id, memory_id),
+        )
+        await conn.commit()
+
+    async def get_edges(self, memory_id: str) -> list[Edge]:
+        conn = self.store._require_conn()
+        cursor = await conn.execute(
+            "SELECT source_id, target_id, relation, weight, reasoning, metadata, created_at "
+            "FROM causal_edges WHERE source_id = ? OR target_id = ?",
+            (memory_id, memory_id),
+        )
+        rows = await cursor.fetchall()
+        out: list[Edge] = []
+        for row in rows:
+            meta = json.loads(row["metadata"] or "{}")
+            meta.setdefault("relation", row["relation"])
+            meta.setdefault("reasoning", row["reasoning"])
+            out.append(
+                Edge(
+                    source_id=row["source_id"],
+                    target_id=row["target_id"],
+                    graph_type="causal",
+                    weight=float(row["weight"]),
+                    metadata=meta,
+                    created_at=datetime.fromisoformat(row["created_at"]),
+                )
+            )
+        return out
+
+    async def get_neighbors(
+        self,
+        memory_id: str,
+        depth: int = 1,
+        max_results: int = 20,
+    ) -> list[tuple[str, float]]:
+        conn = self.store._require_conn()
+        visited: dict[str, float] = {memory_id: 0.0}
+        frontier: list[tuple[str, float]] = [(memory_id, 1.0)]
+        for _ in range(depth):
+            next_frontier: list[tuple[str, float]] = []
+            for node, cum in frontier:
+                cursor = await conn.execute(
+                    "SELECT target_id, weight FROM causal_edges WHERE source_id = ?",
+                    (node,),
+                )
+                rows = await cursor.fetchall()
+                for row in rows:
+                    tid = row["target_id"]
+                    w = cum * float(row["weight"])
+                    if tid in visited and visited[tid] >= w:
+                        continue
+                    visited[tid] = w
+                    next_frontier.append((tid, w))
+            frontier = next_frontier
+        visited.pop(memory_id, None)
+        return sorted(visited.items(), key=lambda kv: kv[1], reverse=True)[:max_results]
+
+    async def get_causal_chain(self, memory_id: str, max_depth: int = 5) -> list[str]:
+        """Return the longest cause→effect chain ending at ``memory_id``."""
+        conn = self.store._require_conn()
+        chain = [memory_id]
+        seen = {memory_id}
+        current = memory_id
+        for _ in range(max_depth):
+            cursor = await conn.execute(
+                "SELECT source_id FROM causal_edges "
+                "WHERE target_id = ? ORDER BY weight DESC LIMIT 1",
+                (current,),
+            )
+            row = await cursor.fetchone()
+            if row is None or row["source_id"] in seen:
+                break
+            current = row["source_id"]
+            chain.append(current)
+            seen.add(current)
+        return list(reversed(chain))

contextdb/graphs/entity.py

@@ -0,0 +1,251 @@
+"""Entity graph — memories and named entities linked bidirectionally.
+
+Entity extraction is delegated to the configured LLM. Responses are expected
+to be JSON; we tolerate minor formatting noise and deduplicate on lowercased
+entity names.
+"""
+
+from __future__ import annotations
+
+import json
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any
+from uuid import uuid4
+
+from contextdb.core.models import Edge, Entity
+from contextdb.graphs.base import BaseGraph
+
+if TYPE_CHECKING:
+    from contextdb.store.sqlite_store import SQLiteStore
+    from contextdb.utils.llm import LLMProvider
+
+_SCHEMA = """
+CREATE TABLE IF NOT EXISTS entities (
+    id TEXT PRIMARY KEY,
+    name TEXT NOT NULL,
+    entity_type TEXT NOT NULL,
+    attributes TEXT DEFAULT '{}',
+    created_at TEXT NOT NULL,
+    updated_at TEXT NOT NULL
+);
+CREATE UNIQUE INDEX IF NOT EXISTS idx_entities_name_lc ON entities(LOWER(name));
+
+CREATE TABLE IF NOT EXISTS memory_entity_edges (
+    memory_id TEXT NOT NULL,
+    entity_id TEXT NOT NULL,
+    relation TEXT DEFAULT 'MENTIONS',
+    created_at TEXT NOT NULL,
+    PRIMARY KEY (memory_id, entity_id)
+);
+CREATE INDEX IF NOT EXISTS idx_mee_memory ON memory_entity_edges(memory_id);
+CREATE INDEX IF NOT EXISTS idx_mee_entity ON memory_entity_edges(entity_id);
+"""
+
+_EXTRACT_PROMPT = """Extract all named entities from the text. Return strict JSON.
+
+Schema:
+{"entities": [
+  {"name": "string",
+   "type": "PERSON|ORG|PRODUCT|LOCATION|EVENT|OTHER",
+   "attributes": {}}
+]}
+
+Text:
+"{text}"
+"""
+
+
+def _safe_json(text: str) -> dict[str, Any]:
+    text = text.strip()
+    if text.startswith("```"):
+        # strip triple-backtick fences
+        lines = text.splitlines()
+        text = "\n".join(line for line in lines if not line.startswith("```"))
+    try:
+        loaded = json.loads(text)
+        return loaded if isinstance(loaded, dict) else {}
+    except json.JSONDecodeError:
+        start = text.find("{")
+        end = text.rfind("}")
+        if start != -1 and end != -1 and end > start:
+            try:
+                loaded = json.loads(text[start : end + 1])
+                return loaded if isinstance(loaded, dict) else {}
+            except json.JSONDecodeError:
+                return {}
+        return {}
+
+
+class EntityGraph(BaseGraph):
+    """LLM-extracted entity overlay over the memory store."""
+
+    def __init__(self, store: SQLiteStore, llm: LLMProvider) -> None:
+        super().__init__(store)
+        self.llm = llm
+
+    async def initialize(self) -> None:
+        conn = self.store._require_conn()
+        await conn.executescript(_SCHEMA)
+        await conn.commit()
+
+    async def extract_entities(self, content: str) -> list[Entity]:
+        response = await self.llm.generate(_EXTRACT_PROMPT.replace("{text}", content))
+        payload = _safe_json(response)
+        out: list[Entity] = []
+        for raw in payload.get("entities", []) or []:
+            name = str(raw.get("name", "")).strip()
+            if not name:
+                continue
+            out.append(
+                Entity(
+                    name=name,
+                    entity_type=str(raw.get("type", "OTHER")),
+                    attributes=raw.get("attributes") or {},
+                )
+            )
+        return out
+
+    async def _find_or_create(self, entity: Entity) -> str:
+        conn = self.store._require_conn()
+        cursor = await conn.execute(
+            "SELECT id, attributes FROM entities WHERE LOWER(name) = ?",
+            (entity.name.lower(),),
+        )
+        row = await cursor.fetchone()
+        if row is not None:
+            if entity.attributes:
+                existing = json.loads(row["attributes"] or "{}")
+                existing.update(entity.attributes)
+                await conn.execute(
+                    "UPDATE entities SET attributes = ?, updated_at = ? WHERE id = ?",
+                    (
+                        json.dumps(existing),
+                        datetime.now(tz=timezone.utc).isoformat(),
+                        row["id"],
+                    ),
+                )
+                await conn.commit()
+            return str(row["id"])
+        eid = str(uuid4())
+        now = datetime.now(tz=timezone.utc).isoformat()
+        await conn.execute(
+            "INSERT INTO entities (id, name, entity_type, attributes, created_at, updated_at) "
+            "VALUES (?,?,?,?,?,?)",
+            (eid, entity.name, entity.entity_type, json.dumps(entity.attributes), now, now),
+        )
+        await conn.commit()
+        return eid
+
+    async def _link_memory(self, memory_id: str, entity_id: str) -> None:
+        conn = self.store._require_conn()
+        await conn.execute(
+            "INSERT OR IGNORE INTO memory_entity_edges "
+            "(memory_id, entity_id, relation, created_at) VALUES (?,?,?,?)",
+            (memory_id, entity_id, "MENTIONS", datetime.now(tz=timezone.utc).isoformat()),
+        )
+        await conn.commit()
+
+    async def add_node(self, memory_id: str, data: dict[str, Any]) -> None:
+        content = data.get("content", "")
+        if not content:
+            return
+        extracted = await self.extract_entities(content)
+        if not extracted:
+            return
+        for ent in extracted:
+            eid = await self._find_or_create(ent)
+            await self._link_memory(memory_id, eid)
+
+    async def add_edge(self, edge: Edge) -> None:
+        # Entity graph's primary edges are memory↔entity; generic Edge is a no-op.
+        return None
+
+    async def remove_node(self, memory_id: str) -> None:
+        conn = self.store._require_conn()
+        await conn.execute(
+            "DELETE FROM memory_entity_edges WHERE memory_id = ?", (memory_id,)
+        )
+        await conn.commit()
+
+    async def get_edges(self, memory_id: str) -> list[Edge]:
+        conn = self.store._require_conn()
+        cursor = await conn.execute(
+            "SELECT entity_id, created_at FROM memory_entity_edges WHERE memory_id = ?",
+            (memory_id,),
+        )
+        rows = await cursor.fetchall()
+        return [
+            Edge(
+                source_id=memory_id,
+                target_id=row["entity_id"],
+                graph_type="entity",
+                weight=1.0,
+                created_at=datetime.fromisoformat(row["created_at"]),
+            )
+            for row in rows
+        ]
+
+    async def get_neighbors(
+        self,
+        memory_id: str,
+        depth: int = 1,
+        max_results: int = 20,
+    ) -> list[tuple[str, float]]:
+        conn = self.store._require_conn()
+        cursor = await conn.execute(
+            "SELECT entity_id FROM memory_entity_edges WHERE memory_id = ?", (memory_id,)
+        )
+        entity_ids = [row["entity_id"] for row in await cursor.fetchall()]
+        if not entity_ids:
+            return []
+        placeholders = ",".join(["?"] * len(entity_ids))
+        cursor = await conn.execute(
+            f"SELECT memory_id, COUNT(*) as shared "
+            f"FROM memory_entity_edges WHERE entity_id IN ({placeholders}) AND memory_id != ? "
+            "GROUP BY memory_id ORDER BY shared DESC LIMIT ?",
+            [*entity_ids, memory_id, max_results],
+        )
+        rows = await cursor.fetchall()
+        return [(row["memory_id"], float(row["shared"])) for row in rows]
+
+    async def get_entity_profile(self, name: str) -> dict[str, Any]:
+        conn = self.store._require_conn()
+        cursor = await conn.execute(
+            "SELECT id, name, entity_type, attributes FROM entities WHERE LOWER(name) = ?",
+            (name.lower(),),
+        )
+        row = await cursor.fetchone()
+        if row is None:
+            return {"name": name, "memories": [], "attributes": {}}
+        eid = row["id"]
+        mem_cursor = await conn.execute(
+            "SELECT memory_id FROM memory_entity_edges WHERE entity_id = ?", (eid,)
+        )
+        memory_ids = [r["memory_id"] for r in await mem_cursor.fetchall()]
+        return {
+            "id": eid,
+            "name": row["name"],
+            "entity_type": row["entity_type"],
+            "attributes": json.loads(row["attributes"] or "{}"),
+            "memories": memory_ids,
+        }
+
+    async def get_entity_memories(self, name: str) -> list[str]:
+        profile = await self.get_entity_profile(name)
+        return list(profile.get("memories", []))
+
+    async def list_entities(self) -> list[Entity]:
+        conn = self.store._require_conn()
+        cursor = await conn.execute(
+            "SELECT id, name, entity_type, attributes FROM entities ORDER BY name"
+        )
+        rows = await cursor.fetchall()
+        return [
+            Entity(
+                name=row["name"],
+                entity_type=row["entity_type"],
+                attributes=json.loads(row["attributes"] or "{}"),
+                memory_ids=await self.get_entity_memories(row["name"]),
+            )
+            for row in rows
+        ]

contextdb/graphs/semantic.py

@@ -0,0 +1,156 @@
+"""Semantic graph — edges between memories with high embedding similarity."""
+
+from __future__ import annotations
+
+import json
+from datetime import datetime
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+
+from contextdb.core.models import Edge, MemoryStatus
+from contextdb.graphs.base import BaseGraph
+
+if TYPE_CHECKING:
+    from contextdb.store.sqlite_store import SQLiteStore
+
+_SCHEMA = """
+CREATE TABLE IF NOT EXISTS semantic_edges (
+    source_id TEXT NOT NULL,
+    target_id TEXT NOT NULL,
+    weight REAL NOT NULL,
+    metadata TEXT DEFAULT '{}',
+    created_at TEXT NOT NULL,
+    PRIMARY KEY (source_id, target_id)
+);
+CREATE INDEX IF NOT EXISTS idx_semantic_source ON semantic_edges(source_id);
+CREATE INDEX IF NOT EXISTS idx_semantic_target ON semantic_edges(target_id);
+"""
+
+
+def _cosine(a: list[float], b: list[float]) -> float:
+    va = np.asarray(a, dtype=np.float32)
+    vb = np.asarray(b, dtype=np.float32)
+    na = float(np.linalg.norm(va))
+    nb = float(np.linalg.norm(vb))
+    if na == 0 or nb == 0:
+        return 0.0
+    return float(np.dot(va, vb) / (na * nb))
+
+
+class SemanticGraph(BaseGraph):
+    """Edges represent cosine similarity above a threshold."""
+
+    def __init__(self, store: SQLiteStore, threshold: float = 0.6) -> None:
+        super().__init__(store)
+        self.threshold = threshold
+
+    async def initialize(self) -> None:
+        conn = self.store._require_conn()
+        await conn.executescript(_SCHEMA)
+        await conn.commit()
+
+    async def add_node(self, memory_id: str, data: dict[str, Any]) -> None:
+        embedding = data.get("embedding")
+        if embedding is None:
+            return
+        similar = await self.store.search_by_embedding(embedding, top_k=50)
+        for item in similar:
+            if item.id == memory_id or item.embedding is None:
+                continue
+            if item.status != MemoryStatus.ACTIVE:
+                continue
+            sim = _cosine(embedding, item.embedding)
+            if sim >= self.threshold:
+                await self.add_edge(
+                    Edge(
+                        source_id=memory_id,
+                        target_id=item.id,
+                        graph_type="semantic",
+                        weight=sim,
+                    )
+                )
+
+    async def add_edge(self, edge: Edge) -> None:
+        conn = self.store._require_conn()
+        await conn.execute(
+            "INSERT OR REPLACE INTO semantic_edges "
+            "(source_id, target_id, weight, metadata, created_at) VALUES (?,?,?,?,?)",
+            (
+                edge.source_id,
+                edge.target_id,
+                edge.weight,
+                json.dumps(edge.metadata),
+                edge.created_at.isoformat(),
+            ),
+        )
+        # Bidirectional for undirected similarity semantics.
+        await conn.execute(
+            "INSERT OR REPLACE INTO semantic_edges "
+            "(source_id, target_id, weight, metadata, created_at) VALUES (?,?,?,?,?)",
+            (
+                edge.target_id,
+                edge.source_id,
+                edge.weight,
+                json.dumps(edge.metadata),
+                edge.created_at.isoformat(),
+            ),
+        )
+        await conn.commit()
+
+    async def get_neighbors(
+        self,
+        memory_id: str,
+        depth: int = 1,
+        max_results: int = 20,
+    ) -> list[tuple[str, float]]:
+        conn = self.store._require_conn()
+        visited: dict[str, float] = {memory_id: 0.0}
+        frontier: list[tuple[str, float]] = [(memory_id, 1.0)]
+        for _ in range(depth):
+            next_frontier: list[tuple[str, float]] = []
+            for node, cum in frontier:
+                cursor = await conn.execute(
+                    "SELECT target_id, weight FROM semantic_edges WHERE source_id = ?",
+                    (node,),
+                )
+                rows = await cursor.fetchall()
+                for row in rows:
+                    tid = row["target_id"]
+                    w = cum * float(row["weight"])
+                    if tid in visited and visited[tid] >= w:
+                        continue
+                    visited[tid] = w
+                    next_frontier.append((tid, w))
+            frontier = next_frontier
+        visited.pop(memory_id, None)
+        ranked = sorted(visited.items(), key=lambda kv: kv[1], reverse=True)
+        return ranked[:max_results]
+
+    async def remove_node(self, memory_id: str) -> None:
+        conn = self.store._require_conn()
+        await conn.execute(
+            "DELETE FROM semantic_edges WHERE source_id = ? OR target_id = ?",
+            (memory_id, memory_id),
+        )
+        await conn.commit()
+
+    async def get_edges(self, memory_id: str) -> list[Edge]:
+        conn = self.store._require_conn()
+        cursor = await conn.execute(
+            "SELECT source_id, target_id, weight, metadata, created_at "
+            "FROM semantic_edges WHERE source_id = ?",
+            (memory_id,),
+        )
+        rows = await cursor.fetchall()
+        return [
+            Edge(
+                source_id=row["source_id"],
+                target_id=row["target_id"],
+                graph_type="semantic",
+                weight=float(row["weight"]),
+                metadata=json.loads(row["metadata"] or "{}"),
+                created_at=datetime.fromisoformat(row["created_at"]),
+            )
+            for row in rows
+        ]