pycontextdb 0.1.0__py3-none-any.whl

contextdb/__init__.py

@@ -0,0 +1,72 @@
+"""ContextDB — The unified context layer for AI agents."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from contextdb.client import ContextDB
+from contextdb.core.config import ContextDBConfig
+from contextdb.core.exceptions import (
+    ConfigError,
+    ContextDBError,
+    MemoryNotFoundError,
+    PrivacyError,
+    StorageError,
+)
+from contextdb.core.models import (
+    Edge,
+    Entity,
+    MemoryItem,
+    MemoryStatus,
+    MemoryType,
+    PIIAnnotation,
+    PIIType,
+    RetentionPolicy,
+)
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "ConfigError",
+    "ContextDB",
+    "ContextDBConfig",
+    "ContextDBError",
+    "Edge",
+    "Entity",
+    "MemoryItem",
+    "MemoryNotFoundError",
+    "MemoryStatus",
+    "MemoryType",
+    "PIIAnnotation",
+    "PIIType",
+    "PrivacyError",
+    "RetentionPolicy",
+    "StorageError",
+    "__version__",
+    "init",
+]
+
+
+def init(
+    user_id: str | None = None,
+    config: ContextDBConfig | None = None,
+    **kwargs: Any,
+) -> ContextDB:
+    """Create a :class:`ContextDB` client.
+
+    The client is lazy — resources are provisioned on the first ``await`` on
+    any I/O method, so ``init()`` itself does not touch the disk or network.
+
+    Args:
+        user_id: Optional user scope. Every write carries this as a filter.
+        config: Pre-built configuration. When ``None`` one is constructed
+            from ``kwargs`` and environment variables prefixed with
+            ``CONTEXTDB_``.
+        **kwargs: Forwarded to :class:`ContextDBConfig` if ``config`` is
+            not provided.
+
+    Returns:
+        A fully-configured :class:`ContextDB` client.
+    """
+    resolved = config or ContextDBConfig(**kwargs)
+    return ContextDB(resolved, user_id=user_id)

contextdb/agents/__init__.py

@@ -0,0 +1,8 @@
+"""Multi-agent primitives: memory bus, RL-driven memory manager."""
+
+from __future__ import annotations
+
+from contextdb.agents.memory_bus import MemoryBus
+from contextdb.agents.rl_manager import RLMemoryManager
+
+__all__ = ["MemoryBus", "RLMemoryManager"]

contextdb/agents/memory_bus.py

@@ -0,0 +1,80 @@
+"""In-process pub/sub bus for multi-agent memory sharing.
+
+The bus is deliberately minimal: ``publish`` + ``subscribe`` with optional
+filters. Agents that share a :class:`MemoryBus` can broadcast new memories,
+events, or reflections to peers without coupling directly to each other.
+
+For cross-process sharing, swap the backing queue for Redis or NATS — the
+:class:`MemoryBus` contract stays identical. That's out of scope for v0.1.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass, field
+from typing import Any
+from uuid import uuid4
+
+
+@dataclass
+class _Subscription:
+    id: str
+    topic: str
+    callback: Callable[[dict[str, Any]], Awaitable[None]]
+    filters: dict[str, Any] = field(default_factory=dict)
+
+
+class MemoryBus:
+    """Async, in-process pub/sub across cooperating agents."""
+
+    def __init__(self) -> None:
+        self._subscriptions: dict[str, list[_Subscription]] = {}
+        self._lock = asyncio.Lock()
+
+    async def publish(self, topic: str, payload: dict[str, Any]) -> int:
+        """Deliver ``payload`` to all matching subscribers. Returns delivery count."""
+        async with self._lock:
+            subs = list(self._subscriptions.get(topic, []))
+            wildcard = list(self._subscriptions.get("*", []))
+        delivered = 0
+        for sub in subs + wildcard:
+            if not _matches_filters(payload, sub.filters):
+                continue
+            try:
+                await sub.callback(payload)
+                delivered += 1
+            except Exception:  # noqa: BLE001
+                # Subscriber failure must not break the bus.
+                continue
+        return delivered
+
+    async def subscribe(
+        self,
+        topic: str,
+        callback: Callable[[dict[str, Any]], Awaitable[None]],
+        filters: dict[str, Any] | None = None,
+    ) -> str:
+        sub = _Subscription(
+            id=str(uuid4()), topic=topic, callback=callback, filters=filters or {}
+        )
+        async with self._lock:
+            self._subscriptions.setdefault(topic, []).append(sub)
+        return sub.id
+
+    async def unsubscribe(self, subscription_id: str) -> bool:
+        async with self._lock:
+            for topic, subs in self._subscriptions.items():
+                for i, sub in enumerate(subs):
+                    if sub.id == subscription_id:
+                        del self._subscriptions[topic][i]
+                        return True
+        return False
+
+    async def topics(self) -> list[str]:
+        async with self._lock:
+            return [t for t, subs in self._subscriptions.items() if subs]
+
+
+def _matches_filters(payload: dict[str, Any], filters: dict[str, Any]) -> bool:
+    return all(payload.get(key) == expected for key, expected in filters.items())

contextdb/agents/rl_manager.py

@@ -0,0 +1,89 @@
+"""RL-guided memory manager — decide ADD / UPDATE / DELETE / NOOP per write.
+
+In v0.1 we ship the inference side of the RL manager: a policy prompt that
+asks the LLM to pick one of four actions for a new piece of content given
+the current top-20 candidate memories. This mirrors Memory-R1's action
+space (Yu et al., 2024) and gives us a drop-in surface for swapping in a
+trained model later without reshuffling the client.
+
+Training (PPO over trajectories of ADD/UPDATE/DELETE/NOOP outcomes) lives
+in a separate research harness — not shipped with the library.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from contextdb.core.models import MemoryItem
+    from contextdb.utils.llm import LLMProvider
+
+
+_POLICY_PROMPT = """You are a memory manager. Given new content and the top
+candidate existing memories, choose exactly one action and return strict JSON.
+
+Actions:
+- "ADD"    — store the content as a new memory.
+- "UPDATE" — merge into an existing memory; must include target_memory_id and merged content.
+- "DELETE" — discard an existing memory superseded by this content; must include target_memory_id.
+- "NOOP"   — the content is redundant; do nothing.
+
+Schema:
+{"action": "ADD|UPDATE|DELETE|NOOP",
+ "target_memory_id": "string|null",
+ "content": "string|null",
+ "reasoning": "string"}
+
+New content: "{content}"
+
+Candidate memories (id :: content):
+{candidates}
+"""
+
+
+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 RLMemoryManager:
+    """LLM-driven inference-time policy over ADD/UPDATE/DELETE/NOOP."""
+
+    def __init__(self, llm: LLMProvider, max_candidates: int = 10) -> None:
+        self.llm = llm
+        self.max_candidates = max_candidates
+
+    async def decide(
+        self,
+        content: str,
+        candidates: list[MemoryItem],
+    ) -> dict[str, Any]:
+        cand_snippets = "\n".join(
+            f"- {m.id} :: {m.content[:200]}" for m in candidates[: self.max_candidates]
+        ) or "(none)"
+        prompt = _POLICY_PROMPT.replace("{content}", content).replace(
+            "{candidates}", cand_snippets
+        )
+        response = await self.llm.generate(prompt, temperature=0.0, max_tokens=400)
+        decision = _safe_json(response)
+        action = str(decision.get("action", "ADD")).upper()
+        if action not in {"ADD", "UPDATE", "DELETE", "NOOP"}:
+            action = "ADD"
+        decision["action"] = action
+        return decision

contextdb/cli.py

@@ -0,0 +1,107 @@
+"""Tiny command-line interface for ContextDB.
+
+The CLI is intentionally minimal — enough to add/search/stats/export/import
+from a shell, but not a replacement for the Python SDK. We use argparse
+rather than click/typer so the library doesn't grow another dependency.
+
+Usage:
+    contextdb add "My birthday is March 5"
+    contextdb search "when is my birthday"
+    contextdb stats
+    contextdb export dump.json
+    contextdb import dump.json
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import json
+import sys
+from pathlib import Path
+from typing import Any
+
+from contextdb import init
+from contextdb.core.models import MemoryType
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(prog="contextdb")
+    parser.add_argument("--storage-url", default=None, help="Override storage URL.")
+    parser.add_argument("--user-id", default=None, help="Optional user_id scope.")
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    p_add = sub.add_parser("add", help="Add a memory.")
+    p_add.add_argument("content")
+    p_add.add_argument(
+        "--type",
+        choices=[t.value for t in MemoryType],
+        default=MemoryType.FACTUAL.value,
+    )
+    p_add.add_argument("--source", default="")
+
+    p_search = sub.add_parser("search", help="Search memories.")
+    p_search.add_argument("query")
+    p_search.add_argument("--top-k", type=int, default=5)
+
+    sub.add_parser("stats", help="Print store statistics.")
+
+    p_export = sub.add_parser("export", help="Export the store to JSON.")
+    p_export.add_argument("path")
+
+    p_import = sub.add_parser("import", help="Import a ContextDB JSON dump.")
+    p_import.add_argument("path")
+
+    return parser
+
+
+async def _run(args: argparse.Namespace) -> int:
+    kwargs: dict[str, Any] = {}
+    if args.storage_url:
+        kwargs["storage_url"] = args.storage_url
+    client = init(user_id=args.user_id, **kwargs)
+    try:
+        if args.command == "add":
+            item = await client.add(
+                content=args.content,
+                memory_type=MemoryType(args.type),
+                source=args.source,
+            )
+            print(json.dumps({"id": item.id, "content": item.content}, indent=2))
+        elif args.command == "search":
+            hits = await client.search(args.query, top_k=args.top_k)
+            print(
+                json.dumps(
+                    [{"id": m.id, "content": m.content} for m in hits], indent=2
+                )
+            )
+        elif args.command == "stats":
+            stats = await client.stats()
+            print(json.dumps(stats, indent=2, default=str))
+        elif args.command == "export":
+            from contextdb.utils.migrations import JSONExporter
+
+            exporter = JSONExporter(client)
+            count = await exporter.export(Path(args.path))
+            print(f"Exported {count} memories to {args.path}")
+        elif args.command == "import":
+            from contextdb.utils.migrations import JSONImporter
+
+            importer = JSONImporter(client)
+            count = await importer.import_path(Path(args.path))
+            print(f"Imported {count} memories from {args.path}")
+        else:
+            return 1
+        return 0
+    finally:
+        await client.close()
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = _build_parser()
+    args = parser.parse_args(argv)
+    return asyncio.run(_run(args))
+
+
+if __name__ == "__main__":  # pragma: no cover
+    sys.exit(main())