pycontextdb 0.1.0__py3-none-any.whl

contextdb/memory/working.py

@@ -0,0 +1,90 @@
+"""Working memory — short-lived, token-budgeted session scratchpad.
+
+Working memories are the agent's "what am I doing right now" buffer. They are
+scoped to a session, bounded by a token budget (characters ÷ 4 as a coarse
+proxy for tokens), and trimmed FIFO when the budget is exceeded.
+
+Items are stored with ``memory_type=WORKING`` and a ``session_id`` metadata
+tag. The retention policy short-circuits them to a 24-hour TTL by default.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from contextdb.core.models import MemoryItem, MemoryStatus, MemoryType
+
+if TYPE_CHECKING:
+    from contextdb.client import ContextDB
+
+
+def _approx_tokens(text: str) -> int:
+    """Rough token estimate: 4 chars ≈ 1 token (GPT-style average)."""
+    return max(1, len(text) // 4)
+
+
+class WorkingMemory:
+    """Token-budgeted FIFO scratchpad scoped to a session."""
+
+    def __init__(
+        self,
+        client: ContextDB,
+        session_id: str,
+        max_tokens: int = 4000,
+    ) -> None:
+        self.client = client
+        self.session_id = session_id
+        self.max_tokens = max_tokens
+
+    async def push(
+        self,
+        content: str,
+        metadata: dict[str, Any] | None = None,
+    ) -> MemoryItem:
+        meta = dict(metadata or {})
+        meta["session_id"] = self.session_id
+        item = await self.client.add(
+            content=content,
+            memory_type=MemoryType.WORKING,
+            metadata=meta,
+        )
+        await self._evict_if_over_budget()
+        return item
+
+    async def list_session(self) -> list[MemoryItem]:
+        await self.client._ensure_init()
+        store = self.client._require_store()
+        items = await store.list_memories(memory_type=MemoryType.WORKING, limit=10000)
+        return [
+            m
+            for m in items
+            if m.metadata.get("session_id") == self.session_id
+            and m.status == MemoryStatus.ACTIVE
+        ]
+
+    async def context_window(self) -> str:
+        items = await self.list_session()
+        items.sort(key=lambda m: m.created_at)
+        return "\n".join(m.content for m in items)
+
+    async def clear(self) -> int:
+        items = await self.list_session()
+        await self.client._ensure_init()
+        store = self.client._require_store()
+        for item in items:
+            await store.delete(item.id, hard=True)
+        return len(items)
+
+    async def _evict_if_over_budget(self) -> None:
+        items = await self.list_session()
+        items.sort(key=lambda m: m.created_at)
+        total = sum(_approx_tokens(m.content) for m in items)
+        if total <= self.max_tokens:
+            return
+        await self.client._ensure_init()
+        store = self.client._require_store()
+        for item in items:
+            if total <= self.max_tokens:
+                break
+            await store.delete(item.id, hard=True)
+            total -= _approx_tokens(item.content)

contextdb/privacy/__init__.py

@@ -0,0 +1,9 @@
+"""Privacy primitives: PII detection, retention enforcement, audit trail."""
+
+from __future__ import annotations
+
+from contextdb.privacy.audit import AuditEntry, AuditLogger
+from contextdb.privacy.pii_detector import PIIDetector
+from contextdb.privacy.retention import RetentionManager
+
+__all__ = ["AuditEntry", "AuditLogger", "PIIDetector", "RetentionManager"]

contextdb/privacy/audit.py

@@ -0,0 +1,199 @@
+"""Hash-chained audit trail for every memory-touching operation.
+
+Each :class:`AuditEntry` carries the SHA-256 hash of the previous entry, so
+tampering with any record invalidates the chain downstream. This gives us
+tamper-evidence without a signing authority — enough for internal compliance
+and most external audits short of regulated industries (which should add a
+signing key on top).
+
+Verifying the chain is O(n) in entry count but reads are sequential and
+bounded by storage. For production workloads, consider periodic anchoring
+to an external log (e.g., a ledger) — not in scope for v0.1.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any
+from uuid import uuid4
+
+from pydantic import BaseModel, Field
+
+if TYPE_CHECKING:
+    from contextdb.store.sqlite_store import SQLiteStore
+
+_SCHEMA = """
+CREATE TABLE IF NOT EXISTS audit_log (
+    id TEXT PRIMARY KEY,
+    sequence INTEGER NOT NULL,
+    operation TEXT NOT NULL,
+    memory_id TEXT,
+    user_id TEXT,
+    details TEXT DEFAULT '{}',
+    previous_hash TEXT NOT NULL,
+    entry_hash TEXT NOT NULL,
+    timestamp TEXT NOT NULL
+);
+CREATE INDEX IF NOT EXISTS idx_audit_sequence ON audit_log(sequence);
+CREATE INDEX IF NOT EXISTS idx_audit_memory ON audit_log(memory_id);
+CREATE INDEX IF NOT EXISTS idx_audit_user ON audit_log(user_id);
+"""
+
+_GENESIS_HASH = "0" * 64
+
+
+class AuditEntry(BaseModel):
+    """A single record in the hash-chained audit log."""
+
+    id: str = Field(default_factory=lambda: str(uuid4()))
+    sequence: int
+    operation: str
+    memory_id: str | None = None
+    user_id: str | None = None
+    details: dict[str, Any] = Field(default_factory=dict)
+    previous_hash: str
+    entry_hash: str
+    timestamp: datetime = Field(default_factory=lambda: datetime.now(tz=timezone.utc))
+
+    def canonical_payload(self) -> str:
+        return json.dumps(
+            {
+                "id": self.id,
+                "sequence": self.sequence,
+                "operation": self.operation,
+                "memory_id": self.memory_id,
+                "user_id": self.user_id,
+                "details": self.details,
+                "previous_hash": self.previous_hash,
+                "timestamp": self.timestamp.isoformat(),
+            },
+            sort_keys=True,
+            separators=(",", ":"),
+        )
+
+
+def _compute_hash(payload: str) -> str:
+    return hashlib.sha256(payload.encode("utf-8")).hexdigest()
+
+
+class AuditLogger:
+    """Append-only audit log with per-entry SHA-256 chaining."""
+
+    def __init__(self, store: SQLiteStore) -> None:
+        self.store = store
+
+    async def initialize(self) -> None:
+        conn = self.store._require_conn()
+        await conn.executescript(_SCHEMA)
+        await conn.commit()
+
+    async def log(
+        self,
+        operation: str,
+        memory_id: str | None = None,
+        user_id: str | None = None,
+        details: dict[str, Any] | None = None,
+    ) -> AuditEntry:
+        conn = self.store._require_conn()
+        cursor = await conn.execute(
+            "SELECT sequence, entry_hash FROM audit_log ORDER BY sequence DESC LIMIT 1"
+        )
+        row = await cursor.fetchone()
+        sequence = (row["sequence"] + 1) if row else 1
+        previous_hash = row["entry_hash"] if row else _GENESIS_HASH
+
+        entry = AuditEntry(
+            sequence=sequence,
+            operation=operation,
+            memory_id=memory_id,
+            user_id=user_id,
+            details=details or {},
+            previous_hash=previous_hash,
+            entry_hash="",
+        )
+        entry.entry_hash = _compute_hash(entry.canonical_payload())
+
+        await conn.execute(
+            "INSERT INTO audit_log "
+            "(id, sequence, operation, memory_id, user_id, details, "
+            "previous_hash, entry_hash, timestamp) VALUES (?,?,?,?,?,?,?,?,?)",
+            (
+                entry.id,
+                entry.sequence,
+                entry.operation,
+                entry.memory_id,
+                entry.user_id,
+                json.dumps(entry.details),
+                entry.previous_hash,
+                entry.entry_hash,
+                entry.timestamp.isoformat(),
+            ),
+        )
+        await conn.commit()
+        return entry
+
+    async def get_history(
+        self,
+        memory_id: str | None = None,
+        user_id: str | None = None,
+        limit: int = 1000,
+    ) -> list[AuditEntry]:
+        conn = self.store._require_conn()
+        clauses: list[str] = []
+        params: list[Any] = []
+        if memory_id is not None:
+            clauses.append("memory_id = ?")
+            params.append(memory_id)
+        if user_id is not None:
+            clauses.append("user_id = ?")
+            params.append(user_id)
+        where = f"WHERE {' AND '.join(clauses)}" if clauses else ""
+        params.append(limit)
+        cursor = await conn.execute(
+            f"SELECT * FROM audit_log {where} ORDER BY sequence ASC LIMIT ?",
+            params,
+        )
+        rows = await cursor.fetchall()
+        return [
+            AuditEntry(
+                id=row["id"],
+                sequence=int(row["sequence"]),
+                operation=row["operation"],
+                memory_id=row["memory_id"],
+                user_id=row["user_id"],
+                details=json.loads(row["details"] or "{}"),
+                previous_hash=row["previous_hash"],
+                entry_hash=row["entry_hash"],
+                timestamp=datetime.fromisoformat(row["timestamp"]),
+            )
+            for row in rows
+        ]
+
+    async def verify_chain(self) -> bool:
+        conn = self.store._require_conn()
+        cursor = await conn.execute(
+            "SELECT * FROM audit_log ORDER BY sequence ASC"
+        )
+        rows = await cursor.fetchall()
+        expected_prev = _GENESIS_HASH
+        for row in rows:
+            if row["previous_hash"] != expected_prev:
+                return False
+            entry = AuditEntry(
+                id=row["id"],
+                sequence=int(row["sequence"]),
+                operation=row["operation"],
+                memory_id=row["memory_id"],
+                user_id=row["user_id"],
+                details=json.loads(row["details"] or "{}"),
+                previous_hash=row["previous_hash"],
+                entry_hash=row["entry_hash"],
+                timestamp=datetime.fromisoformat(row["timestamp"]),
+            )
+            expected_hash = _compute_hash(entry.canonical_payload())
+            if expected_hash != row["entry_hash"]:
+                return False
+            expected_prev = row["entry_hash"]
+        return True

contextdb/privacy/pii_detector.py

@@ -0,0 +1,173 @@
+"""Regex-based PII detection, redaction, and encryption.
+
+Covers the common categories (email, phone, SSN, credit card) with well-known
+patterns. Name and address detection is heuristic here; richer recognizers
+can be plugged in later (spaCy NER, Presidio, etc.).
+
+Offsets in returned :class:`PIIAnnotation` objects index the **original**
+content, which lets callers redact from the tail forward without corrupting
+indices.
+
+Three actions are supported:
+
+* ``redact`` — replace each PII span with ``[<TYPE>]`` and keep plaintext
+  in :attr:`PIIAnnotation.original` (useful for operator audit on a trusted
+  store).
+* ``encrypt`` — replace each PII span with ``[<TYPE>]`` but store a Fernet
+  ciphertext in :attr:`PIIAnnotation.original`. Reversible via :meth:`decrypt`
+  when the caller holds the key.
+* ``flag`` / ``allow`` — leave text intact; only annotations are produced.
+"""
+
+from __future__ import annotations
+
+import base64
+import hashlib
+import logging
+import os
+import re
+from typing import Literal
+
+from cryptography.fernet import Fernet, InvalidToken
+
+from contextdb.core.models import PIIAnnotation, PIIType
+
+PIIAction = Literal["redact", "encrypt", "flag", "allow"]
+
+_logger = logging.getLogger(__name__)
+
+# Patterns applied in order. The credit-card pattern comes before phone so
+# 16-digit card numbers without dashes don't get misclassified when the phone
+# pattern is tightened.
+_PATTERNS: list[tuple[PIIType, re.Pattern[str]]] = [
+    (
+        PIIType.EMAIL,
+        re.compile(r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b"),
+    ),
+    (
+        PIIType.SSN,
+        re.compile(r"\b\d{3}-\d{2}-\d{4}\b"),
+    ),
+    (
+        PIIType.CREDIT_CARD,
+        # 13-16 digit sequences optionally grouped by hyphens or spaces.
+        re.compile(r"\b(?:\d{4}[-\s]){3}\d{3,4}\b|\b\d{13,16}\b"),
+    ),
+    (
+        PIIType.PHONE,
+        re.compile(
+            r"\b(?:\+?1[-.\s]?)?\(?\d{3}\)?[-.\s]\d{3}[-.\s]\d{4}\b"
+        ),
+    ),
+]
+
+
+def _derive_fernet_key(raw: str) -> bytes:
+    """Derive a 32-byte urlsafe-base64 Fernet key from any string.
+
+    We SHA-256 the input so users can supply a human-readable secret without
+    having to know Fernet's key format. This is intentionally not PBKDF2 —
+    the input is already expected to be a high-entropy operator secret, not
+    a user-memorable password.
+    """
+    digest = hashlib.sha256(raw.encode("utf-8")).digest()
+    return base64.urlsafe_b64encode(digest)
+
+
+class PIIDetector:
+    """Detect and apply a policy to PII in free-form text."""
+
+    def __init__(
+        self,
+        action: PIIAction = "redact",
+        encryption_key: str | None = None,
+    ) -> None:
+        self.action: PIIAction = action
+        self._fernet: Fernet | None = None
+        if action == "encrypt":
+            key = encryption_key or os.environ.get("CONTEXTDB_PII_KEY")
+            if key:
+                self._fernet = Fernet(_derive_fernet_key(key))
+            else:
+                _logger.warning(
+                    "PII action is 'encrypt' but no key is configured "
+                    "(pass encryption_key= or set CONTEXTDB_PII_KEY). "
+                    "Falling back to redact; originals will NOT be recoverable."
+                )
+
+    def detect(self, text: str) -> list[PIIAnnotation]:
+        """Return non-overlapping PII spans sorted by start offset."""
+        found: list[PIIAnnotation] = []
+        taken: list[tuple[int, int]] = []
+        for pii_type, pattern in _PATTERNS:
+            for match in pattern.finditer(text):
+                start, end = match.start(), match.end()
+                if any(_overlaps(start, end, s, e) for s, e in taken):
+                    continue
+                taken.append((start, end))
+                original = match.group()
+                stored_original = (
+                    self._fernet.encrypt(original.encode("utf-8")).decode("ascii")
+                    if self._fernet is not None
+                    else original
+                )
+                found.append(
+                    PIIAnnotation(
+                        pii_type=pii_type,
+                        start=start,
+                        end=end,
+                        original=stored_original,
+                        redacted=f"[{pii_type.value}]",
+                    )
+                )
+        found.sort(key=lambda a: a.start)
+        return found
+
+    def redact(
+        self,
+        text: str,
+        annotations: list[PIIAnnotation] | None = None,
+    ) -> str:
+        """Replace each PII span with its typed placeholder (e.g. ``[EMAIL]``)."""
+        if annotations is None:
+            annotations = self.detect(text)
+        result = text
+        for ann in sorted(annotations, key=lambda a: a.start, reverse=True):
+            result = result[: ann.start] + ann.redacted + result[ann.end :]
+        return result
+
+    def process(self, text: str) -> tuple[str, list[PIIAnnotation]]:
+        """Detect PII and apply the configured action.
+
+        Returns ``(processed_text, annotations)``. For ``allow``/``flag`` the
+        text is returned unchanged; for ``redact``/``encrypt`` the text is
+        returned with placeholders substituted. When action is ``encrypt``
+        the :attr:`PIIAnnotation.original` field holds a Fernet ciphertext
+        that can be round-tripped via :meth:`decrypt`.
+        """
+        annotations = self.detect(text)
+        if self.action in {"redact", "encrypt"}:
+            return self.redact(text, annotations), annotations
+        return text, annotations
+
+    def decrypt(self, annotation: PIIAnnotation) -> str:
+        """Recover the plaintext original from an encrypted annotation.
+
+        Raises :class:`ValueError` if the detector is not configured for
+        encryption or if the ciphertext is tampered / wrong key.
+        """
+        if self._fernet is None:
+            raise ValueError(
+                "PIIDetector is not configured for encryption. Pass "
+                "encryption_key= or set CONTEXTDB_PII_KEY when constructing."
+            )
+        try:
+            return self._fernet.decrypt(annotation.original.encode("ascii")).decode(
+                "utf-8"
+            )
+        except InvalidToken as exc:
+            raise ValueError("Invalid or tampered PII ciphertext.") from exc
+
+
+def _overlaps(a_start: int, a_end: int, b_start: int, b_end: int) -> bool:
+    return a_start < b_end and b_start < a_end

contextdb/privacy/retention.py

@@ -0,0 +1,99 @@
+"""Retention policy enforcement.
+
+Applies the configured :class:`~contextdb.core.models.RetentionPolicy` to
+the store: anything older than the per-type TTL gets archived (soft) or
+erased (hard, when ``right_to_erasure`` is honored and the user asks).
+
+The enforcer is explicit — nothing is deleted on a schedule unless
+:meth:`RetentionManager.enforce` is called. Operators can wire it to a
+cron / periodic task, but the library does not start background threads.
+"""
+
+from __future__ import annotations
+
+from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING
+
+from contextdb.core.models import MemoryStatus, MemoryType, RetentionPolicy
+
+if TYPE_CHECKING:
+    from contextdb.privacy.audit import AuditLogger
+    from contextdb.store.sqlite_store import SQLiteStore
+
+
+_TYPE_TO_TTL: dict[MemoryType, str] = {
+    MemoryType.FACTUAL: "factual_ttl",
+    MemoryType.EXPERIENTIAL: "experiential_ttl",
+    MemoryType.WORKING: "working_ttl",
+}
+
+
+class RetentionManager:
+    """Apply retention TTLs and honor right-to-erasure requests."""
+
+    def __init__(
+        self,
+        store: SQLiteStore,
+        audit: AuditLogger | None,
+        policy: RetentionPolicy,
+    ) -> None:
+        self.store = store
+        self.audit = audit
+        self.policy = policy
+
+    def _ttl_for(self, memory_type: MemoryType) -> timedelta | None:
+        field = _TYPE_TO_TTL[memory_type]
+        ttl = getattr(self.policy, field, None)
+        if ttl is None:
+            return self.policy.default_ttl
+        assert isinstance(ttl, timedelta)
+        return ttl
+
+    async def enforce(self, hard: bool = False) -> int:
+        """Archive (or hard-delete) memories past their TTL.
+
+        Returns the number of affected rows.
+        """
+        now = datetime.now(tz=timezone.utc)
+        affected = 0
+        for memory_type in MemoryType:
+            ttl = self._ttl_for(memory_type)
+            if ttl is None:
+                continue
+            cutoff = now - ttl
+            memories = await self.store.list_memories(memory_type=memory_type, limit=100000)
+            for memory in memories:
+                if memory.created_at > cutoff:
+                    continue
+                if hard:
+                    await self.store.delete(memory.id, hard=True)
+                    operation = "ERASE"
+                else:
+                    await self.store.update(memory.id, status=MemoryStatus.ARCHIVED)
+                    operation = "ARCHIVE"
+                affected += 1
+                if self.audit is not None:
+                    await self.audit.log(
+                        operation=operation,
+                        memory_id=memory.id,
+                        details={"reason": "retention_ttl", "type": memory_type.value},
+                    )
+        return affected
+
+    async def erase_user(self, user_id: str) -> int:
+        """Honor a right-to-erasure request for a specific user."""
+        if not self.policy.right_to_erasure:
+            return 0
+        memories = await self.store.list_memories(user_id=user_id, limit=100000)
+        deleted = 0
+        for memory in memories:
+            await self.store.delete(memory.id, hard=True)
+            deleted += 1
+            if self.audit is not None:
+                await self.audit.log(
+                    operation="ERASE",
+                    memory_id=memory.id,
+                    user_id=user_id,
+                    details={"reason": "right_to_erasure"},
+                )
+        return deleted

contextdb/store/__init__.py

@@ -0,0 +1,15 @@
+"""Storage backends and vector indices."""
+
+from __future__ import annotations
+
+from contextdb.store.base import BaseStore
+from contextdb.store.sqlite_store import SQLiteStore
+from contextdb.store.vector_index import NumpyIndex, VectorIndex, get_vector_index
+
+__all__ = [
+    "BaseStore",
+    "NumpyIndex",
+    "SQLiteStore",
+    "VectorIndex",
+    "get_vector_index",
+]

contextdb/store/base.py

@@ -0,0 +1,67 @@
+"""Abstract storage backend contract."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+from contextdb.core.models import MemoryStatus
+
+if TYPE_CHECKING:
+    from contextdb.core.models import MemoryItem, MemoryType
+
+
+class BaseStore(ABC):
+    """Abstract base for all persistent memory stores.
+
+    Implementations must be async and safe to call concurrently from one
+    event loop. They are not required to be safe across processes.
+    """
+
+    @abstractmethod
+    async def initialize(self) -> None:
+        """Create schema / open connections. Idempotent."""
+
+    @abstractmethod
+    async def add(self, item: MemoryItem) -> MemoryItem:
+        """Persist a memory and return the stored copy."""
+
+    @abstractmethod
+    async def get(self, memory_id: str) -> MemoryItem | None:
+        """Fetch one memory by id; increments access counters."""
+
+    @abstractmethod
+    async def update(self, memory_id: str, **kwargs: object) -> MemoryItem:
+        """Partial update; unknown keys raise ValueError."""
+
+    @abstractmethod
+    async def delete(self, memory_id: str, hard: bool = False) -> None:
+        """Soft delete by default (status=DELETED); ``hard=True`` removes the row."""
+
+    @abstractmethod
+    async def search_by_embedding(
+        self,
+        embedding: list[float],
+        top_k: int = 10,
+        filters: dict[str, object] | None = None,
+    ) -> list[MemoryItem]:
+        """Return top-k most similar memories by cosine similarity."""
+
+    @abstractmethod
+    async def list_memories(
+        self,
+        user_id: str | None = None,
+        memory_type: MemoryType | None = None,
+        status: MemoryStatus = MemoryStatus.ACTIVE,
+        limit: int = 100,
+        offset: int = 0,
+    ) -> list[MemoryItem]:
+        """List memories with optional filters."""
+
+    @abstractmethod
+    async def count(self, user_id: str | None = None) -> int:
+        """Count active memories, optionally scoped to a user."""
+
+    @abstractmethod
+    async def close(self) -> None:
+        """Release resources."""