pycontextdb 0.1.0__py3-none-any.whl

contextdb/graphs/temporal.py

@@ -0,0 +1,173 @@
+"""Temporal graph — edges carry ordering and proximity in event-time."""
+
+from __future__ import annotations
+
+import json
+from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING, Any
+
+from contextdb.core.models import Edge, MemoryItem
+from contextdb.graphs.base import BaseGraph
+
+if TYPE_CHECKING:
+    from contextdb.store.sqlite_store import SQLiteStore
+
+_SCHEMA = """
+CREATE TABLE IF NOT EXISTS temporal_edges (
+    source_id TEXT NOT NULL,
+    target_id TEXT NOT NULL,
+    relation TEXT NOT NULL,
+    weight REAL NOT NULL,
+    time_diff_seconds REAL NOT NULL,
+    metadata TEXT DEFAULT '{}',
+    created_at TEXT NOT NULL,
+    PRIMARY KEY (source_id, target_id)
+);
+CREATE INDEX IF NOT EXISTS idx_temporal_source ON temporal_edges(source_id);
+"""
+
+
+class TemporalGraph(BaseGraph):
+    """Link memories that occur near each other in event-time."""
+
+    def __init__(
+        self,
+        store: SQLiteStore,
+        proximity_window: timedelta = timedelta(hours=24),
+    ) -> None:
+        super().__init__(store)
+        self.proximity_window = proximity_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:
+        event_time = data.get("event_time")
+        if event_time is None:
+            return
+        nearby = await self._find_temporally_nearby(memory_id, event_time)
+        for other_id, other_time in nearby:
+            diff = (event_time - other_time).total_seconds()
+            if abs(diff) < 300:
+                relation = "CONCURRENT"
+            elif diff > 0:
+                relation = "AFTER"
+            else:
+                relation = "BEFORE"
+            weight = 1.0 / (1.0 + abs(diff) / 3600.0)
+            await self.add_edge(
+                Edge(
+                    source_id=memory_id,
+                    target_id=other_id,
+                    graph_type="temporal",
+                    weight=weight,
+                    metadata={"relation": relation, "time_diff_seconds": diff},
+                )
+            )
+
+    async def _find_temporally_nearby(
+        self, memory_id: str, event_time: datetime
+    ) -> list[tuple[str, datetime]]:
+        conn = self.store._require_conn()
+        window_start = (event_time - self.proximity_window).isoformat()
+        window_end = (event_time + self.proximity_window).isoformat()
+        cursor = await conn.execute(
+            "SELECT id, event_time FROM memories "
+            "WHERE id != ? AND event_time IS NOT NULL "
+            "AND event_time BETWEEN ? AND ? AND status = 'ACTIVE'",
+            (memory_id, window_start, window_end),
+        )
+        rows = await cursor.fetchall()
+        return [(row["id"], datetime.fromisoformat(row["event_time"])) 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 temporal_edges "
+            "(source_id, target_id, relation, weight, time_diff_seconds, metadata, created_at) "
+            "VALUES (?,?,?,?,?,?,?)",
+            (
+                edge.source_id,
+                edge.target_id,
+                str(meta.get("relation", "ADJACENT")),
+                edge.weight,
+                float(meta.get("time_diff_seconds", 0.0)),
+                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 temporal_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, relation, "
+            "time_diff_seconds FROM temporal_edges WHERE source_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("time_diff_seconds", row["time_diff_seconds"])
+            out.append(
+                Edge(
+                    source_id=row["source_id"],
+                    target_id=row["target_id"],
+                    graph_type="temporal",
+                    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()
+        cursor = await conn.execute(
+            "SELECT target_id, weight FROM temporal_edges WHERE source_id = ? "
+            "ORDER BY weight DESC LIMIT ?",
+            (memory_id, max_results),
+        )
+        rows = await cursor.fetchall()
+        return [(row["target_id"], float(row["weight"])) for row in rows]
+
+    async def get_timeline(
+        self,
+        entity: str | None = None,
+        start: datetime | None = None,
+        end: datetime | None = None,
+    ) -> list[MemoryItem]:
+        store = self.store
+        items = await store.list_memories(limit=10000)
+        filtered = [m for m in items if m.event_time is not None]
+        if start is not None:
+            filtered = [m for m in filtered if m.event_time and m.event_time >= start]
+        if end is not None:
+            filtered = [m for m in filtered if m.event_time and m.event_time <= end]
+        if entity is not None:
+            ent = entity.lower()
+            filtered = [
+                m
+                for m in filtered
+                if ent in (e.lower() for e in m.entity_mentions) or ent in m.content.lower()
+            ]
+        filtered.sort(key=lambda m: m.event_time or datetime.min.replace(tzinfo=timezone.utc))
+        return filtered

contextdb/integrations/__init__.py

@@ -0,0 +1,10 @@
+"""Framework adapters: LangChain, OpenAI function-calling, CrewAI, AutoGen."""
+
+from __future__ import annotations
+
+__all__ = [
+    "autogen_memory",
+    "crewai_memory",
+    "langchain_memory",
+    "openai_tools",
+]

contextdb/integrations/autogen.py

@@ -0,0 +1,39 @@
+"""AutoGen memory adapter.
+
+AutoGen passes strings around as agent "messages" and expects memory to
+surface a list of relevant prior messages on demand. We wrap ContextDB to
+match that shape.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from contextdb.client import ContextDB
+
+
+class ContextDBAutoGenMemory:
+    """AutoGen-style memory facade over ContextDB."""
+
+    def __init__(self, client: ContextDB, top_k: int = 5) -> None:
+        self.client = client
+        self.top_k = top_k
+
+    async def add_message(self, role: str, content: str) -> str:
+        item = await self.client.add(
+            content=f"{role}: {content}",
+            metadata={"role": role},
+        )
+        return item.id
+
+    async def get_relevant(self, query: str) -> list[str]:
+        hits = await self.client.search(query, top_k=self.top_k)
+        return [m.content for m in hits]
+
+    async def clear(self) -> None:
+        await self.client._ensure_init()
+        store = self.client._require_store()
+        memories = await store.list_memories(user_id=self.client.user_id, limit=100000)
+        for memory in memories:
+            await store.delete(memory.id, hard=True)

contextdb/integrations/crewai.py

@@ -0,0 +1,41 @@
+"""CrewAI memory adapter.
+
+CrewAI's memory contract is intentionally small (``save`` / ``search`` /
+``reset``). We mirror those method names on a thin wrapper so that a
+``Crew`` can consume a :class:`ContextDBCrewMemory` without importing
+``crewai`` here.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from contextdb.client import ContextDB
+
+
+class ContextDBCrewMemory:
+    """CrewAI-compatible memory wrapper backed by ContextDB."""
+
+    def __init__(self, client: ContextDB, top_k: int = 5) -> None:
+        self.client = client
+        self.top_k = top_k
+
+    async def save(self, value: str, metadata: dict[str, Any] | None = None) -> str:
+        item = await self.client.add(content=value, metadata=metadata)
+        return item.id
+
+    async def search(self, query: str, limit: int | None = None) -> list[dict[str, Any]]:
+        top_k = limit if limit is not None else self.top_k
+        hits = await self.client.search(query, top_k=top_k)
+        return [
+            {"id": m.id, "content": m.content, "metadata": m.metadata}
+            for m in hits
+        ]
+
+    async def reset(self) -> None:
+        await self.client._ensure_init()
+        store = self.client._require_store()
+        memories = await store.list_memories(user_id=self.client.user_id, limit=100000)
+        for memory in memories:
+            await store.delete(memory.id, hard=True)

contextdb/integrations/langchain.py

@@ -0,0 +1,132 @@
+"""LangChain memory adapter.
+
+Exposes :class:`ContextDBMemory`, a small class that conforms to LangChain's
+conversational memory interface (``load_memory_variables`` /
+``save_context``) backed by :class:`~contextdb.client.ContextDB`. We do
+**not** import ``langchain`` here — keeping the dependency optional — and
+instead structure the adapter so that LangChain's duck-typed interface is
+met by method signatures alone.
+
+Both sync and async entry points are provided. LangChain code paths still
+exist that call memories synchronously, so the sync wrappers bridge to the
+async client via :func:`asyncio.run` / ``run_until_complete``. The async
+methods (``aload_memory_variables`` / ``asave_context`` / ``aclear``) are
+preferred whenever the caller is already inside an event loop.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable
+from typing import TYPE_CHECKING, Any, TypeVar
+
+if TYPE_CHECKING:
+    from contextdb.client import ContextDB
+
+_T = TypeVar("_T")
+
+
+def _run_sync(coro: Awaitable[_T]) -> _T:
+    """Execute an awaitable from synchronous code.
+
+    Uses ``asyncio.run`` when no loop is running. When called from inside an
+    already-running loop (e.g. a Jupyter cell) we fall back to creating a
+    dedicated loop on a worker thread — ``asyncio.run`` refuses to nest, and
+    ``run_until_complete`` on a running loop deadlocks. Callers inside a
+    loop should prefer the ``a*`` methods.
+    """
+    try:
+        running = asyncio.get_running_loop()
+    except RuntimeError:
+        running = None
+    if running is None:
+        return asyncio.run(coro)  # type: ignore[arg-type]
+
+    import threading
+
+    result: list[_T] = []
+    error: list[BaseException] = []
+
+    def _worker() -> None:
+        loop = asyncio.new_event_loop()
+        try:
+            result.append(loop.run_until_complete(coro))
+        except BaseException as exc:  # noqa: BLE001
+            error.append(exc)
+        finally:
+            loop.close()
+
+    thread = threading.Thread(target=_worker)
+    thread.start()
+    thread.join()
+    if error:
+        raise error[0]
+    return result[0]
+
+
+class ContextDBMemory:
+    """LangChain-compatible memory backed by ContextDB.
+
+    Usage:
+        >>> from contextdb import init
+        >>> db = init()
+        >>> memory = ContextDBMemory(db, session_id="chat-1")
+        >>> chain = ConversationChain(llm=llm, memory=memory)
+    """
+
+    memory_key: str = "history"
+
+    def __init__(
+        self,
+        client: ContextDB,
+        session_id: str,
+        max_tokens: int = 2000,
+        top_k: int = 5,
+    ) -> None:
+        self.client = client
+        self.session_id = session_id
+        self.max_tokens = max_tokens
+        self.top_k = top_k
+
+    @property
+    def memory_variables(self) -> list[str]:
+        return [self.memory_key]
+
+    # -- Async (preferred) ------------------------------------------------- #
+
+    async def aload_memory_variables(self, inputs: dict[str, Any]) -> dict[str, Any]:
+        query = str(inputs.get("input", "")) or str(next(iter(inputs.values()), ""))
+        hits = await self.client.search(query, top_k=self.top_k) if query else []
+        working = self.client.working(self.session_id, max_tokens=self.max_tokens)
+        window = await working.context_window()
+        history = "\n".join(m.content for m in hits)
+        return {self.memory_key: f"{history}\n{window}".strip()}
+
+    async def asave_context(
+        self, inputs: dict[str, Any], outputs: dict[str, Any]
+    ) -> None:
+        working = self.client.working(self.session_id, max_tokens=self.max_tokens)
+        user_text = str(inputs.get("input", "")).strip()
+        bot_text = str(outputs.get("output", "")).strip()
+        if user_text:
+            await working.push(f"User: {user_text}")
+        if bot_text:
+            await working.push(f"Assistant: {bot_text}")
+
+    async def aclear(self) -> None:
+        working = self.client.working(self.session_id, max_tokens=self.max_tokens)
+        await working.clear()
+
+    # -- Sync (LangChain legacy paths) ------------------------------------- #
+
+    def load_memory_variables(self, inputs: dict[str, Any]) -> dict[str, Any]:
+        """Synchronous wrapper over :meth:`aload_memory_variables`."""
+        return _run_sync(self.aload_memory_variables(inputs))
+
+    def save_context(self, inputs: dict[str, Any], outputs: dict[str, Any]) -> None:
+        """Synchronous wrapper over :meth:`asave_context`."""
+        _run_sync(self.asave_context(inputs, outputs))
+
+    def clear(self) -> None:
+        """Synchronous wrapper over :meth:`aclear`."""
+        _run_sync(self.aclear())

contextdb/integrations/openai_tools.py

@@ -0,0 +1,124 @@
+"""OpenAI function-calling tool definitions for ContextDB.
+
+Returns a list of JSON Schema tool specs that can be passed directly into
+``openai.chat.completions.create(..., tools=[...])``. Each tool name maps to
+a bound async callable returned by :func:`make_tool_handlers`.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from contextdb.client import ContextDB
+
+
+def tool_schemas() -> list[dict[str, Any]]:
+    """OpenAI tools JSON schema for ContextDB memory operations."""
+    return [
+        {
+            "type": "function",
+            "function": {
+                "name": "memory_add",
+                "description": "Store a new fact, observation, or experience in long-term memory.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "content": {"type": "string"},
+                        "memory_type": {
+                            "type": "string",
+                            "enum": ["FACTUAL", "EXPERIENTIAL", "WORKING"],
+                        },
+                        "entity_mentions": {
+                            "type": "array",
+                            "items": {"type": "string"},
+                        },
+                    },
+                    "required": ["content"],
+                },
+            },
+        },
+        {
+            "type": "function",
+            "function": {
+                "name": "memory_search",
+                "description": "Search long-term memory semantically and temporally.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "query": {"type": "string"},
+                        "top_k": {"type": "integer", "default": 5},
+                    },
+                    "required": ["query"],
+                },
+            },
+        },
+        {
+            "type": "function",
+            "function": {
+                "name": "memory_get_entity",
+                "description": "Retrieve the profile and associated memories for a named entity.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {"name": {"type": "string"}},
+                    "required": ["name"],
+                },
+            },
+        },
+        {
+            "type": "function",
+            "function": {
+                "name": "memory_forget",
+                "description": "Bulk-delete memories by entity or age.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "entity": {"type": "string"},
+                        "older_than_days": {"type": "integer"},
+                    },
+                },
+            },
+        },
+    ]
+
+
+ToolHandler = Callable[[dict[str, Any]], Awaitable[Any]]
+
+
+def make_tool_handlers(client: ContextDB) -> dict[str, ToolHandler]:
+    """Bind tool names to async callables executing against ``client``."""
+    from datetime import timedelta
+
+    from contextdb.core.models import MemoryType
+
+    async def _add(args: dict[str, Any]) -> dict[str, Any]:
+        mt_value = args.get("memory_type", "FACTUAL")
+        mt = MemoryType(mt_value if isinstance(mt_value, str) else "FACTUAL")
+        item = await client.add(
+            content=args["content"],
+            memory_type=mt,
+            entity_mentions=args.get("entity_mentions"),
+        )
+        return {"id": item.id}
+
+    async def _search(args: dict[str, Any]) -> list[dict[str, Any]]:
+        hits = await client.search(args["query"], top_k=int(args.get("top_k", 5)))
+        return [{"id": m.id, "content": m.content} for m in hits]
+
+    async def _get_entity(args: dict[str, Any]) -> dict[str, Any]:
+        return await client.get_entity(args["name"])
+
+    async def _forget(args: dict[str, Any]) -> dict[str, int]:
+        older_than: timedelta | None = None
+        if "older_than_days" in args:
+            older_than = timedelta(days=int(args["older_than_days"]))
+        deleted = await client.forget(entity=args.get("entity"), older_than=older_than)
+        return {"deleted": deleted}
+
+    return {
+        "memory_add": _add,
+        "memory_search": _search,
+        "memory_get_entity": _get_entity,
+        "memory_forget": _forget,
+    }

contextdb/memory/__init__.py

@@ -0,0 +1,9 @@
+"""Typed memory surfaces: factual, experiential, working."""
+
+from __future__ import annotations
+
+from contextdb.memory.experiential import ExperientialMemory
+from contextdb.memory.factual import FactualMemory
+from contextdb.memory.working import WorkingMemory
+
+__all__ = ["ExperientialMemory", "FactualMemory", "WorkingMemory"]

contextdb/memory/experiential.py

@@ -0,0 +1,102 @@
+"""Experiential memory — trajectories, reflections, outcomes.
+
+Experiential memories capture *what happened* during an agent run: the
+action taken, the outcome observed, and an optional post-hoc reflection.
+They are the substrate behind "agent learns from its own rollouts" —
+Memory-R1 and MAGMA both derive their gains from a well-organized
+experiential store.
+
+Two memories are written per trajectory:
+
+* A :class:`~contextdb.core.models.MemoryType.EXPERIENTIAL` item holding the
+  structured trajectory (action / outcome / context).
+* An optional linked reflection memory that the agent writes after the fact.
+"""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any
+
+from contextdb.core.models import MemoryItem, MemoryType
+
+if TYPE_CHECKING:
+    from contextdb.client import ContextDB
+
+
+class ExperientialMemory:
+    """Trajectories (action+outcome) and reflections (post-hoc insights)."""
+
+    def __init__(self, client: ContextDB, user_id: str | None = None) -> None:
+        self.client = client
+        self.user_id = user_id
+
+    async def record_trajectory(
+        self,
+        action: str,
+        outcome: str,
+        context: dict[str, Any] | None = None,
+        success: bool | None = None,
+    ) -> MemoryItem:
+        content = f"Action: {action}\nOutcome: {outcome}"
+        meta: dict[str, Any] = {
+            "action": action,
+            "outcome": outcome,
+            "context": context or {},
+        }
+        if success is not None:
+            meta["success"] = success
+        return await self.client.add(
+            content=content,
+            memory_type=MemoryType.EXPERIENTIAL,
+            metadata=meta,
+            event_time=datetime.now(tz=timezone.utc),
+        )
+
+    async def add_reflection(
+        self,
+        trajectory_id: str,
+        insight: str,
+        metadata: dict[str, Any] | None = None,
+    ) -> MemoryItem:
+        meta = dict(metadata or {})
+        meta["reflection_on"] = trajectory_id
+        return await self.client.add(
+            content=insight,
+            memory_type=MemoryType.EXPERIENTIAL,
+            metadata=meta,
+        )
+
+    async def recall_similar(
+        self,
+        situation: str,
+        top_k: int = 5,
+    ) -> list[MemoryItem]:
+        return await self.client.search(
+            situation, top_k=top_k, memory_type=MemoryType.EXPERIENTIAL
+        )
+
+    async def list_trajectories(self, limit: int = 100) -> list[MemoryItem]:
+        await self.client._ensure_init()
+        store = self.client._require_store()
+        items = await store.list_memories(
+            user_id=self.user_id, memory_type=MemoryType.EXPERIENTIAL, limit=limit
+        )
+        return [m for m in items if "reflection_on" not in m.metadata]
+
+    async def list_reflections(
+        self,
+        trajectory_id: str | None = None,
+        limit: int = 100,
+    ) -> list[MemoryItem]:
+        await self.client._ensure_init()
+        store = self.client._require_store()
+        items = await store.list_memories(
+            user_id=self.user_id, memory_type=MemoryType.EXPERIENTIAL, limit=limit
+        )
+        reflections = [m for m in items if "reflection_on" in m.metadata]
+        if trajectory_id is not None:
+            reflections = [
+                m for m in reflections if m.metadata.get("reflection_on") == trajectory_id
+            ]
+        return reflections

contextdb/memory/factual.py

@@ -0,0 +1,58 @@
+"""Factual memory — durable statements of fact.
+
+The simplest of the three memory surfaces: a typed filter over
+:meth:`contextdb.client.ContextDB.add` / :meth:`~contextdb.client.ContextDB.search`
+that forces ``memory_type=FACTUAL``. Convenience wrappers add a few
+fact-specific affordances (``recall``, ``update_fact``, ``list_facts``).
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from contextdb.core.models import MemoryItem, MemoryType
+
+if TYPE_CHECKING:
+    from contextdb.client import ContextDB
+
+
+class FactualMemory:
+    """Thin typed layer over the general-purpose client."""
+
+    def __init__(self, client: ContextDB, user_id: str | None = None) -> None:
+        self.client = client
+        self.user_id = user_id
+
+    async def add(
+        self,
+        content: str,
+        metadata: dict[str, Any] | None = None,
+        entity_mentions: list[str] | None = None,
+        confidence: float = 1.0,
+    ) -> MemoryItem:
+        meta = dict(metadata or {})
+        meta.setdefault("confidence", confidence)
+        return await self.client.add(
+            content=content,
+            memory_type=MemoryType.FACTUAL,
+            metadata=meta,
+            entity_mentions=entity_mentions,
+        )
+
+    async def recall(self, query: str, top_k: int = 5) -> list[MemoryItem]:
+        return await self.client.search(query, top_k=top_k, memory_type=MemoryType.FACTUAL)
+
+    async def update_fact(
+        self,
+        memory_id: str,
+        content: str | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> MemoryItem:
+        return await self.client.update(memory_id, content=content, metadata=metadata)
+
+    async def list_facts(self, limit: int = 100) -> list[MemoryItem]:
+        await self.client._ensure_init()
+        store = self.client._require_store()
+        return await store.list_memories(
+            user_id=self.user_id, memory_type=MemoryType.FACTUAL, limit=limit
+        )