agentkernel-cli 0.1.0__py3-none-any.whl

agentkernel/redaction.py

@@ -0,0 +1,77 @@
+"""Secret redaction for tool output (design §18.1).
+
+Tool results — `bash` stdout, `read_file` contents, web/MCP output — can carry
+API keys and tokens. This module scrubs well-known secret formats from that text
+*before* it enters the context window and the trace, so a leaked credential in a
+command's output doesn't get memorialized in the conversation or logged.
+
+It is deliberately conservative: it matches high-signal token shapes (provider
+key prefixes, PEM private-key blocks, `Authorization` headers, and labelled
+`secret = …` assignments) rather than guessing at entropy, to keep false
+positives low. Structured `ToolResult.data` is never touched — only text content.
+"""
+
+from __future__ import annotations
+
+import re
+
+_PLACEHOLDER = "[REDACTED]"
+
+# PEM private-key blocks (multi-line) — replace the whole block.
+_PRIVATE_KEY = re.compile(
+    r"-----BEGIN [A-Z ]*PRIVATE KEY-----.*?-----END [A-Z ]*PRIVATE KEY-----",
+    re.DOTALL,
+)
+
+# Standalone, high-signal token formats; the whole match is replaced. Anthropic's
+# `sk-ant-` must come before the generic `sk-` so it wins.
+_TOKEN_PATTERNS = (
+    re.compile(r"sk-ant-[A-Za-z0-9_\-]{20,}"),
+    re.compile(r"sk-(?:proj-)?[A-Za-z0-9]{20,}"),
+    re.compile(r"\bgh[posru]_[A-Za-z0-9]{30,}"),
+    re.compile(r"\bgithub_pat_[A-Za-z0-9_]{20,}"),
+    re.compile(r"\bglpat-[A-Za-z0-9_\-]{20,}"),
+    re.compile(r"\bxox[baprs]-[A-Za-z0-9-]{10,}"),
+    re.compile(r"\bAKIA[0-9A-Z]{16}\b"),
+    re.compile(r"\bAIza[A-Za-z0-9_\-]{35}\b"),
+    re.compile(r"\bhf_[A-Za-z0-9]{30,}"),
+    re.compile(r"\b(?:sk|rk)_live_[A-Za-z0-9]{20,}"),
+)
+
+# Labelled assignments — keep the label + operator, redact the value. The value
+# must be a contiguous 12+ char run (no spaces), which keeps prose like
+# "token: the next word" from matching.
+_ASSIGNMENT = re.compile(
+    r"(?i)\b(api[_-]?key|secret|token|password|passwd|pwd|access[_-]?key)"
+    r"(\s*[=:]\s*)"
+    r"['\"]?([A-Za-z0-9._\-/+=]{12,})['\"]?"
+)
+
+# `Authorization: Bearer <token>` — keep the header prefix, redact the token.
+_AUTH_HEADER = re.compile(
+    r"(?i)(authorization\s*:\s*(?:bearer\s+)?)([A-Za-z0-9._\-]{16,})"
+)
+
+
+def redact_secrets(text: str) -> tuple[str, int]:
+    """Return ``(scrubbed_text, count)`` with known secret formats replaced.
+
+    ``count`` is how many secrets were redacted (0 means the text was clean).
+    """
+    if not text:
+        return text, 0
+
+    count = 0
+
+    def _apply(pattern: re.Pattern[str], repl, s: str) -> str:
+        nonlocal count
+        s, n = pattern.subn(repl, s)
+        count += n
+        return s
+
+    out = _apply(_PRIVATE_KEY, _PLACEHOLDER, text)
+    for pattern in _TOKEN_PATTERNS:
+        out = _apply(pattern, _PLACEHOLDER, out)
+    out = _apply(_ASSIGNMENT, lambda m: f"{m.group(1)}{m.group(2)}{_PLACEHOLDER}", out)
+    out = _apply(_AUTH_HEADER, lambda m: f"{m.group(1)}{_PLACEHOLDER}", out)
+    return out, count

agentkernel/semantic_index.py

@@ -0,0 +1,139 @@
+"""Approximate nearest-neighbor helpers for semantic memory.
+
+These use only the Python standard library so the kernel keeps its stdlib-only
+constraint. The default path is a brute-force cosine scan; when scale demands it,
+a small random-projection LSH index can prune the candidate set.
+"""
+
+from __future__ import annotations
+
+import json
+import random
+import sqlite3
+from collections.abc import Callable
+
+
+class LSHIndex:
+    """Random-projection locality-sensitive hash index for dense vectors.
+
+    Each vector is projected onto ``num_bits`` random hyperplanes; the sign of
+    each projection becomes one bit of an integer bucket key. Queries fetch the
+    exact bucket plus all buckets one bit away, which dramatically improves
+    recall without a full linear scan.
+
+    The hyperplanes are persisted in the same SQLite database as the vectors so
+    the index is stable across process restarts.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_bits: int,
+        conn: Callable[[], sqlite3.Connection],
+        *,
+        seed: int = 0,
+    ) -> None:
+        self.dim = dim
+        self.num_bits = num_bits
+        self._conn = conn
+        self._seed = seed
+        self._hyperplanes = self._ensure_hyperplanes()
+
+    def _ensure_hyperplanes(self) -> list[list[float]]:
+        conn = self._conn()
+        conn.execute(
+            "CREATE TABLE IF NOT EXISTS lsh_meta (key TEXT PRIMARY KEY, value TEXT)"
+        )
+        conn.execute(
+            "CREATE TABLE IF NOT EXISTS lsh_buckets ("
+            "note_id INTEGER PRIMARY KEY, bucket INTEGER NOT NULL"
+            ")"
+        )
+
+        bits_row = conn.execute(
+            'SELECT value FROM lsh_meta WHERE key = "bits"'
+        ).fetchone()
+        seed_row = conn.execute(
+            'SELECT value FROM lsh_meta WHERE key = "seed"'
+        ).fetchone()
+        planes_row = conn.execute(
+            'SELECT value FROM lsh_meta WHERE key = "hyperplanes"'
+        ).fetchone()
+
+        if bits_row and seed_row and planes_row:
+            stored_bits = int(bits_row["value"])
+            stored_seed = int(seed_row["value"])
+            planes = json.loads(planes_row["value"])
+            if (
+                stored_bits == self.num_bits
+                and stored_seed == self._seed
+                and len(planes) == self.num_bits
+                and all(len(p) == self.dim for p in planes)
+            ):
+                return planes
+
+        rng = random.Random(self._seed)
+        planes = [
+            [rng.gauss(0.0, 1.0) for _ in range(self.dim)]
+            for _ in range(self.num_bits)
+        ]
+        with conn:
+            conn.execute(
+                "INSERT OR REPLACE INTO lsh_meta (key, value) VALUES (?, ?)",
+                ("bits", str(self.num_bits)),
+            )
+            conn.execute(
+                "INSERT OR REPLACE INTO lsh_meta (key, value) VALUES (?, ?)",
+                ("seed", str(self._seed)),
+            )
+            conn.execute(
+                "INSERT OR REPLACE INTO lsh_meta (key, value) VALUES (?, ?)",
+                ("hyperplanes", json.dumps(planes)),
+            )
+            conn.execute("DELETE FROM lsh_buckets")
+        return planes
+
+    def hash(self, vector: list[float]) -> int:
+        """Return the integer bucket for ``vector``."""
+        bucket = 0
+        for bit, plane in enumerate(self._hyperplanes):
+            dot = sum(v * p for v, p in zip(vector, plane, strict=True))
+            if dot >= 0:
+                bucket |= 1 << bit
+        return bucket
+
+    def query_buckets(self, vector: list[float]) -> list[int]:
+        """Return the query bucket and all one-bit neighbors."""
+        base = self.hash(vector)
+        buckets = [base]
+        for bit in range(self.num_bits):
+            buckets.append(base ^ (1 << bit))
+        return buckets
+
+    def upsert(self, note_id: int, vector: list[float]) -> None:
+        """Store/update the bucket for ``note_id``."""
+        with self._conn():
+            self._conn().execute(
+                "INSERT OR REPLACE INTO lsh_buckets (note_id, bucket) VALUES (?, ?)",
+                (note_id, self.hash(vector)),
+            )
+
+    def remove(self, note_id: int) -> None:
+        with self._conn():
+            self._conn().execute(
+                "DELETE FROM lsh_buckets WHERE note_id = ?", (note_id,)
+            )
+
+    def candidate_ids(self, buckets: list[int]) -> list[int]:
+        """Return note ids whose bucket is in ``buckets``."""
+        if not buckets:
+            return []
+        placeholders = ",".join("?" for _ in buckets)
+        rows = self._conn().execute(
+            f"""
+            SELECT note_id FROM lsh_buckets
+            WHERE bucket IN ({placeholders})
+            """,
+            tuple(buckets),
+        ).fetchall()
+        return [row["note_id"] for row in rows]

agentkernel/semantic_memory.py

@@ -0,0 +1,253 @@
+"""Semantic note search over SQLite notebooks.
+
+This is a thin subclass of ``SqliteNoteStore`` so the JSONL / SQLite split and
+all existing memory-tool behavior stay intact. When an ``EmbeddingProvider`` is
+configured, each note stores a dense vector and recall is re-ranked by cosine
+similarity rather than only token overlap.
+
+For large notebooks the optional LSH index in ``semantic_index`` prunes the
+candidate set before the dense comparison, avoiding a full linear scan.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Sequence
+from pathlib import Path
+
+from agentkernel.embeddings import EmbeddingProvider, cosine_similarity
+from agentkernel.memory import MemoryNote, SqliteNoteStore
+from agentkernel.semantic_index import LSHIndex
+
+
+class SemanticSqliteNoteStore(SqliteNoteStore):
+    """SQLite notebook that also stores dense embeddings for semantic ranking.
+
+    Keyword and full-text search still retrieve candidates; dense similarity
+    refines their order. Notes created before the provider existed can be
+    backfilled with ``reindex_embeddings()``.
+    """
+
+    def __init__(
+        self,
+        path: str | Path,
+        *,
+        embedding_provider: EmbeddingProvider,
+        lsh_bits: int | None = None,
+    ) -> None:
+        self._embedding_provider = embedding_provider
+        self._lsh_bits = lsh_bits
+        self._lsh_index: LSHIndex | None = None
+        # Parent creates the notes table and optional FTS5 index.
+        super().__init__(path)
+        self._ensure_embedding_schema()
+
+    def _ensure_embedding_schema(self) -> None:
+        conn = self._connection()
+        conn.execute(
+            """
+            CREATE TABLE IF NOT EXISTS note_embeddings (
+                note_id INTEGER PRIMARY KEY,
+                embedding_json TEXT NOT NULL
+            )
+            """
+        )
+        conn.commit()
+
+    def _ensure_lsh_index(self, sample_vector: list[float] | None = None) -> None:
+        """Create the LSH index once we know the vector dimension."""
+        if self._lsh_index is not None or not self._lsh_bits:
+            return
+
+        dim = len(sample_vector) if sample_vector else None
+        if dim is None:
+            rows = self._connection().execute(
+                "SELECT embedding_json FROM note_embeddings LIMIT 1"
+            ).fetchall()
+            if rows:
+                dim = len(json.loads(rows[0]["embedding_json"]))
+        if dim is None:
+            return  # no embeddings yet; will initialize on first note
+
+        self._lsh_index = LSHIndex(
+            dim=dim,
+            num_bits=self._lsh_bits,
+            conn=self._connection,
+            seed=0,
+        )
+        # Backfill buckets for any existing embeddings (e.g. after a schema
+        # change or when opening an older notebook).
+        count = self._connection().execute(
+            "SELECT COUNT(*) FROM lsh_buckets"
+        ).fetchone()[0]
+        if count == 0:
+            rows = self._connection().execute(
+                "SELECT note_id, embedding_json FROM note_embeddings"
+            ).fetchall()
+            for row in rows:
+                self._lsh_index.upsert(row["note_id"], json.loads(row["embedding_json"]))
+
+    def _upsert_embedding(self, note_id: int, text: str) -> None:
+        if not text or self._embedding_provider is None:
+            return
+        vector = self._embedding_provider.embed([text])[0]
+        if not vector:
+            return
+        self._ensure_lsh_index(vector)
+        with self._connection():
+            self._connection().execute(
+                """
+                INSERT INTO note_embeddings (note_id, embedding_json)
+                VALUES (?, ?)
+                ON CONFLICT(note_id) DO UPDATE SET
+                    embedding_json = excluded.embedding_json
+                """,
+                (note_id, json.dumps(vector)),
+            )
+        if self._lsh_index is not None:
+            self._lsh_index.upsert(note_id, vector)
+
+    def add(self, text: str, *, tags: Sequence[str] | None = None) -> MemoryNote:
+        note = super().add(text, tags=tags)
+        self._upsert_embedding(note.note_id, note.text)
+        return note
+
+    def update(
+        self,
+        note_id: int,
+        text: str,
+        *,
+        tags: Sequence[str] | None = None,
+    ) -> MemoryNote | None:
+        note = super().update(note_id, text, tags=tags)
+        if note is not None:
+            self._upsert_embedding(note.note_id, note.text)
+        return note
+
+    def search(self, query: str, *, limit: int = 5) -> list[MemoryNote]:
+        query = query.strip()
+        if not query:
+            return super().search(query, limit=limit)
+        if self._embedding_provider is None:
+            return super().search(query, limit=limit)
+
+        query_vec = self._embedding_provider.embed([query])[0]
+        candidates = self._candidates(query_vec, limit=limit)
+        if not candidates:
+            return []
+
+        ids = [note.note_id for note in candidates]
+        vectors = self._load_embeddings(ids)
+
+        scored: list[tuple[float, int, MemoryNote]] = []
+        for note in candidates:
+            vec = vectors.get(note.note_id)
+            if vec and query_vec:
+                similarity = cosine_similarity(query_vec, vec)
+            else:
+                # Notes without embeddings fall below any scored note.
+                similarity = -1.0 if vectors else 0.0
+            scored.append((similarity, note.note_id, note))
+
+        # Highest similarity first; tie-break by note_id for stability.
+        scored.sort(key=lambda item: (item[0], item[1]), reverse=True)
+        ranked = [note for _, _, note in scored[:limit]]
+        for note in ranked:
+            self._touch(note)
+        return ranked
+
+    def _candidates(
+        self, query_vec: list[float], *, limit: int
+    ) -> list[MemoryNote]:
+        """Return notes to score for ``query_vec``.
+
+        If an LSH index is active, use it to narrow the set; otherwise fall
+        back to scanning every note. When the pruned set is too small we also
+        fall back to avoid missing neighbors due to hash collisions.
+        """
+        all_notes = self.all()
+        if not self._lsh_bits or not all_notes:
+            return all_notes
+
+        self._ensure_lsh_index()
+        if self._lsh_index is None:
+            return all_notes
+
+        buckets = self._lsh_index.query_buckets(query_vec)
+        candidate_ids = self._lsh_index.candidate_ids(buckets)
+        # LSH is only a speedup; if the bucket is empty or tiny, scan the table
+        # so accuracy does not suffer on small notebooks or unlucky hashes.
+        if len(candidate_ids) < max(limit * 2, 8):
+            return all_notes
+
+        placeholders = ",".join("?" for _ in candidate_ids)
+        rows = self._connection().execute(
+            f"""
+            SELECT * FROM notes
+            WHERE note_id IN ({placeholders})
+            ORDER BY note_id
+            """,
+            tuple(candidate_ids),
+        ).fetchall()
+        return [self._row_to_note(r) for r in rows]
+
+    def _load_embeddings(self, note_ids: list[int]) -> dict[int, list[float]]:
+        if not note_ids:
+            return {}
+        placeholders = ",".join("?" for _ in note_ids)
+        rows = self._connection().execute(
+            f"""
+            SELECT note_id, embedding_json
+            FROM note_embeddings
+            WHERE note_id IN ({placeholders})
+            """,
+            tuple(note_ids),
+        ).fetchall()
+        return {row["note_id"]: json.loads(row["embedding_json"]) for row in rows}
+
+    def reindex_embeddings(self) -> int:
+        """Compute and store embeddings for all notes that do not have one yet."""
+        rows = self._connection().execute(
+            """
+            SELECT n.note_id, n.text
+            FROM notes n
+            LEFT JOIN note_embeddings e ON n.note_id = e.note_id
+            WHERE e.note_id IS NULL AND n.text IS NOT NULL AND n.text != ''
+            """
+        ).fetchall()
+        if not rows:
+            return 0
+        vectors = self._embedding_provider.embed([r["text"] for r in rows])
+        count = 0
+        with self._connection():
+            for row, vec in zip(rows, vectors, strict=True):
+                if vec:
+                    self._ensure_lsh_index(vec)
+                    self._connection().execute(
+                        "INSERT INTO note_embeddings (note_id, embedding_json) VALUES (?, ?)",
+                        (row["note_id"], json.dumps(vec)),
+                    )
+                    if self._lsh_index is not None:
+                        self._lsh_index.upsert(row["note_id"], vec)
+                    count += 1
+        return count
+
+    def forget(
+        self,
+        *,
+        note_id: int | None = None,
+        text_prefix: str | None = None,
+    ) -> list[MemoryNote]:
+        removed = super().forget(note_id=note_id, text_prefix=text_prefix)
+        if removed:
+            ids = [note.note_id for note in removed]
+            placeholders = ",".join("?" for _ in ids)
+            with self._connection():
+                self._connection().execute(
+                    f"DELETE FROM note_embeddings WHERE note_id IN ({placeholders})",
+                    tuple(ids),
+                )
+                if self._lsh_index is not None:
+                    for nid in ids:
+                        self._lsh_index.remove(nid)
+        return removed