boring-agent-memory 0.2.2__py3-none-any.whl

boring_agent_memory/index.py

@@ -0,0 +1,449 @@
+from __future__ import annotations
+
+import hashlib
+import json
+import sqlite3
+from collections import defaultdict
+from collections.abc import Callable
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from pathlib import Path
+
+from .chunking import CHUNKER_VERSION, DEFAULT_CHUNK_SIZE
+from .ingest import IngestedRecord, ingest_file, iter_candidate_files
+from .privacy import DEFAULT_EXCLUDE_GLOBS, REDACTION_VERSION, SECRET_PATTERNS
+from .schema import (
+    SCHEMA_VERSION,
+    clear_index,
+    connect,
+    init_db,
+    require_index,
+    schema_version,
+)
+
+
+@dataclass(frozen=True)
+class _UpdateOperations:
+    delete_ids: tuple[str, ...]
+    insert_records: tuple[IngestedRecord, ...]
+    removal_paths: tuple[str, ...]
+
+
+def build_index(
+    db_path: Path | str,
+    includes: list[str],
+    excludes: list[str] | None = None,
+    workspace: Path | str | None = None,
+    source_type: str = "file",
+    max_bytes: int = 512 * 1024,
+    chunk_size: int = DEFAULT_CHUNK_SIZE,
+    identity_namespace: str | None = None,
+) -> dict[str, int]:
+    excludes = excludes or []
+    workspace_path = _workspace_path(workspace)
+    records, skipped = _ingest_candidates(
+        includes,
+        excludes,
+        workspace_path,
+        source_type,
+        max_bytes,
+        chunk_size,
+        identity_namespace,
+    )
+    config_json, fingerprint = index_configuration(
+        includes,
+        excludes,
+        workspace_path,
+        source_type,
+        max_bytes,
+        chunk_size,
+        identity_namespace,
+    )
+
+    conn = connect(db_path)
+    try:
+        init_db(conn)
+        # Hold the write lock first, so another build cannot fill the index between this check and the clear.
+        # init_db commits on a new index but leaves a transaction open on an existing one.
+        if not conn.in_transaction:
+            conn.execute("BEGIN IMMEDIATE")
+        with conn:
+            if not records and conn.execute("SELECT 1 FROM documents LIMIT 1").fetchone():
+                raise ValueError("no indexable files matched the include paths; refusing to replace a populated index")
+            clear_index(conn)
+            for record in records:
+                insert_record(conn, record)
+            _write_index_metadata(conn, config_json, fingerprint, chunk_size)
+            _validate_index(conn)
+    finally:
+        conn.close()
+    return {
+        "indexed": len(records),
+        "chunks": sum(len(record.chunks) for record in records),
+        "skipped": skipped,
+    }
+
+
+def update_index(
+    db_path: Path | str,
+    includes: list[str],
+    excludes: list[str] | None = None,
+    workspace: Path | str | None = None,
+    source_type: str = "file",
+    max_bytes: int = 512 * 1024,
+    chunk_size: int = DEFAULT_CHUNK_SIZE,
+    dry_run: bool = False,
+    before_commit: Callable[[sqlite3.Connection], None] | None = None,
+    identity_namespace: str | None = None,
+) -> dict[str, object]:
+    """Apply a hash-based index update in one transaction."""
+    path = require_index(db_path)
+
+    excludes = excludes or []
+    workspace_path = _workspace_path(workspace)
+    records, skipped = _ingest_candidates(
+        includes,
+        excludes,
+        workspace_path,
+        source_type,
+        max_bytes,
+        chunk_size,
+        identity_namespace,
+    )
+    config_json, fingerprint = index_configuration(
+        includes,
+        excludes,
+        workspace_path,
+        source_type,
+        max_bytes,
+        chunk_size,
+        identity_namespace,
+    )
+
+    if dry_run:
+        wal_path = Path(f"{path}-wal")
+        if wal_path.exists() and wal_path.stat().st_size:
+            raise ValueError(
+                "dry-run requires a checkpointed index with no pending WAL content"
+            )
+        conn = sqlite3.connect(
+            path.resolve().as_uri() + "?mode=ro&immutable=1",
+            uri=True,
+        )
+        conn.row_factory = sqlite3.Row
+        if schema_version(conn) != SCHEMA_VERSION:
+            conn.close()
+            raise ValueError("dry-run requires a current index schema; run bam build first")
+    else:
+        conn = connect(path)
+    try:
+        if not dry_run:
+            init_db(conn)
+        metadata = conn.execute("SELECT * FROM index_metadata WHERE singleton = 1").fetchone()
+        if metadata is None:
+            raise ValueError("index has no configuration fingerprint; run bam build once")
+        if metadata["config_fingerprint"] != fingerprint:
+            raise ValueError(
+                "index configuration changed; run bam build before using incremental update"
+            )
+
+        existing_rows = conn.execute(
+            """
+            SELECT
+              d.id,
+              d.source_path,
+              d.source_hash,
+              (SELECT count(*) FROM chunks c WHERE c.document_id = d.id) AS chunk_count
+            FROM documents d
+            WHERE d.workspace = ? AND d.source_type = ?
+            ORDER BY source_path
+            """,
+            (workspace_path.as_posix(), source_type),
+        ).fetchall()
+        report, operations = _plan_update(existing_rows, records, skipped)
+        _validate_removal_access(operations.removal_paths)
+        report["applied"] = not dry_run
+        report["dry_run"] = dry_run
+        if dry_run:
+            return report
+        if not operations.delete_ids and not operations.insert_records:
+            return report
+
+        with conn:
+            for document_id in operations.delete_ids:
+                delete_document(conn, document_id)
+            for record in operations.insert_records:
+                insert_record(conn, record)
+            _write_index_metadata(conn, config_json, fingerprint, chunk_size)
+            _validate_index(conn)
+            if before_commit:
+                before_commit(conn)
+        return report
+    finally:
+        conn.close()
+
+
+def insert_record(conn: sqlite3.Connection, record: IngestedRecord) -> None:
+    conn.execute(
+        """
+        INSERT INTO documents (
+          id, source_type, source_path, workspace, title, content,
+          source_hash, content_hash, metadata_json, updated_at
+        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+        """,
+        (
+            record.id,
+            record.source_type,
+            record.source_path,
+            record.workspace,
+            record.title,
+            record.content,
+            record.source_hash,
+            record.content_hash,
+            record.metadata_json,
+            record.updated_at,
+        ),
+    )
+    for chunk in record.chunks:
+        conn.execute(
+            """
+            INSERT INTO chunks (
+              id, document_id, heading, heading_key, ordinal,
+              start_line, end_line, content, content_hash
+            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
+            """,
+            (
+                chunk.id,
+                record.id,
+                chunk.heading,
+                chunk.heading_key,
+                chunk.ordinal,
+                chunk.start_line,
+                chunk.end_line,
+                chunk.content,
+                chunk.content_hash,
+            ),
+        )
+        conn.execute(
+            """
+            INSERT INTO chunks_fts (chunk_id, title, heading, content, source_path)
+            VALUES (?, ?, ?, ?, ?)
+            """,
+            (chunk.id, record.title, chunk.heading, chunk.content, record.source_path),
+        )
+
+
+def delete_document(conn: sqlite3.Connection, document_id: str) -> None:
+    conn.execute(
+        "DELETE FROM chunks_fts WHERE chunk_id IN "
+        "(SELECT id FROM chunks WHERE document_id = ?)",
+        (document_id,),
+    )
+    conn.execute("DELETE FROM documents WHERE id = ?", (document_id,))
+
+
+def index_configuration(
+    includes: list[str],
+    excludes: list[str],
+    workspace: Path,
+    source_type: str,
+    max_bytes: int,
+    chunk_size: int,
+    identity_namespace: str | None = None,
+) -> tuple[str, str]:
+    payload = {
+        "schema_version": SCHEMA_VERSION,
+        "chunker_version": CHUNKER_VERSION,
+        "chunk_size": chunk_size,
+        "includes": sorted(_normalize_include(value, workspace) for value in includes),
+        "excludes": sorted(excludes),
+        "workspace": workspace.as_posix(),
+        "source_type": source_type,
+        "max_bytes": max_bytes,
+        "document_id_version": 1,
+        "identity_namespace": identity_namespace,
+        "redaction_version": REDACTION_VERSION,
+        "tokenizer": "unicode61",
+        "privacy_policy": hashlib.sha256(
+            json.dumps(
+                {
+                    "default_excludes": DEFAULT_EXCLUDE_GLOBS,
+                    "secret_patterns": [
+                        {"pattern": pattern.pattern, "flags": pattern.flags}
+                        for pattern in SECRET_PATTERNS
+                    ],
+                },
+                sort_keys=True,
+            ).encode("utf-8")
+        ).hexdigest(),
+    }
+    rendered = json.dumps(payload, sort_keys=True, separators=(",", ":"))
+    return rendered, hashlib.sha256(rendered.encode("utf-8")).hexdigest()
+
+
+def _ingest_candidates(
+    includes: list[str],
+    excludes: list[str],
+    workspace: Path,
+    source_type: str,
+    max_bytes: int,
+    chunk_size: int,
+    identity_namespace: str | None,
+) -> tuple[list[IngestedRecord], int]:
+    records: list[IngestedRecord] = []
+    skipped = 0
+    try:
+        candidates = iter_candidate_files(includes, excludes, workspace)
+    except OSError as exc:
+        raise OSError(f"failed to scan candidate sources in {workspace}: {exc}") from exc
+    for file_path in candidates:
+        try:
+            record = ingest_file(
+                file_path,
+                workspace=workspace,
+                source_type=source_type,
+                max_bytes=max_bytes,
+                chunk_size=chunk_size,
+                identity_namespace=identity_namespace,
+            )
+        except OSError as exc:
+            raise OSError(f"failed to read candidate source {file_path}: {exc}") from exc
+        if record is None:
+            skipped += 1
+            continue
+        records.append(record)
+    return records, skipped
+
+
+def _plan_update(
+    existing_rows: list[sqlite3.Row],
+    records: list[IngestedRecord],
+    skipped: int,
+) -> tuple[dict[str, object], _UpdateOperations]:
+    existing_by_path = {row["source_path"]: row for row in existing_rows}
+    current_by_path = {record.source_path: record for record in records}
+
+    unchanged = sorted(
+        path
+        for path in existing_by_path.keys() & current_by_path.keys()
+        if existing_by_path[path]["source_hash"] == current_by_path[path].source_hash
+    )
+    modified = sorted(
+        path
+        for path in existing_by_path.keys() & current_by_path.keys()
+        if existing_by_path[path]["source_hash"] != current_by_path[path].source_hash
+    )
+    added = set(current_by_path.keys() - existing_by_path.keys())
+    removed = set(existing_by_path.keys() - current_by_path.keys())
+
+    removed_by_hash: dict[str, list[str]] = defaultdict(list)
+    added_by_hash: dict[str, list[str]] = defaultdict(list)
+    for path in removed:
+        removed_by_hash[existing_by_path[path]["source_hash"]].append(path)
+    for path in added:
+        added_by_hash[current_by_path[path].source_hash].append(path)
+
+    moves: list[dict[str, str]] = []
+    for content_hash in sorted(removed_by_hash.keys() & added_by_hash.keys()):
+        old_paths = sorted(removed_by_hash[content_hash])
+        new_paths = sorted(added_by_hash[content_hash])
+        if len(old_paths) == 1 and len(new_paths) == 1:
+            old_path = old_paths[0]
+            new_path = new_paths[0]
+            moves.append({"from": old_path, "to": new_path})
+            removed.remove(old_path)
+            added.remove(new_path)
+
+    delete_paths = sorted([*removed, *modified, *(move["from"] for move in moves)])
+    insert_paths = sorted([*added, *modified, *(move["to"] for move in moves)])
+    report: dict[str, object] = {
+        "added": len(added),
+        "modified": len(modified),
+        "moved": len(moves),
+        "removed": len(removed),
+        "unchanged": len(unchanged),
+        "skipped": skipped,
+        "chunks_added": sum(len(current_by_path[path].chunks) for path in insert_paths),
+        "chunks_removed": sum(
+            int(existing_by_path[path]["chunk_count"]) for path in delete_paths
+        ),
+        "changes": {
+            "added": sorted(added),
+            "modified": modified,
+            "moved": moves,
+            "removed": sorted(removed),
+        },
+    }
+    operations = _UpdateOperations(
+        delete_ids=tuple(existing_by_path[path]["id"] for path in delete_paths),
+        insert_records=tuple(current_by_path[path] for path in insert_paths),
+        removal_paths=tuple(sorted(removed)),
+    )
+    return report, operations
+
+
+def _write_index_metadata(
+    conn: sqlite3.Connection,
+    config_json: str,
+    fingerprint: str,
+    chunk_size: int,
+) -> None:
+    conn.execute(
+        """
+        INSERT INTO index_metadata (
+          singleton, config_fingerprint, config_json, built_at,
+          chunker_version, chunk_size
+        ) VALUES (1, ?, ?, ?, ?, ?)
+        ON CONFLICT(singleton) DO UPDATE SET
+          config_fingerprint = excluded.config_fingerprint,
+          config_json = excluded.config_json,
+          built_at = excluded.built_at,
+          chunker_version = excluded.chunker_version,
+          chunk_size = excluded.chunk_size
+        """,
+        (
+            fingerprint,
+            config_json,
+            datetime.now(UTC).isoformat(),
+            CHUNKER_VERSION,
+            chunk_size,
+        ),
+    )
+
+
+def _workspace_path(workspace: Path | str | None) -> Path:
+    return Path(workspace).expanduser().resolve() if workspace else Path.cwd().resolve()
+
+
+def _normalize_include(value: str, workspace: Path) -> str:
+    path = Path(value).expanduser()
+    if not path.is_absolute():
+        path = workspace / path
+    return path.resolve().as_posix()
+
+
+def _validate_index(conn: sqlite3.Connection) -> None:
+    foreign_key_error = conn.execute("PRAGMA foreign_key_check").fetchone()
+    if foreign_key_error is not None:
+        raise sqlite3.IntegrityError(f"index foreign-key violation: {tuple(foreign_key_error)}")
+    chunk_count = conn.execute("SELECT count(*) FROM chunks").fetchone()[0]
+    fts_count = conn.execute("SELECT count(*) FROM chunks_fts").fetchone()[0]
+    if chunk_count != fts_count:
+        raise sqlite3.IntegrityError(
+            f"chunk/FTS row mismatch: chunks={chunk_count} fts={fts_count}"
+        )
+
+
+def _validate_removal_access(paths: tuple[str, ...]) -> None:
+    """Distinguish real deletion from a source hidden by an access failure."""
+    for value in paths:
+        path = Path(value)
+        try:
+            with path.open("rb") as source:
+                source.read(1)
+        except FileNotFoundError:
+            continue
+        except OSError as exc:
+            raise OSError(
+                f"cannot confirm removal of indexed source {path}: {exc}"
+            ) from exc

boring_agent_memory/ingest.py

@@ -0,0 +1,175 @@
+from __future__ import annotations
+
+import hashlib
+import json
+from collections.abc import Iterable
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from glob import glob
+from pathlib import Path
+
+from .chunking import DEFAULT_CHUNK_SIZE, TextChunk, chunk_text
+from .privacy import redact_secrets, should_exclude_path
+
+TEXT_EXTENSIONS = {
+    ".adoc",
+    ".cfg",
+    ".conf",
+    ".csv",
+    ".ini",
+    ".json",
+    ".jsonl",
+    ".log",
+    ".md",
+    ".mdx",
+    ".py",
+    ".rst",
+    ".sql",
+    ".toml",
+    ".txt",
+    ".yaml",
+    ".yml",
+}
+
+
+@dataclass(frozen=True)
+class IngestedRecord:
+    id: str
+    source_type: str
+    source_path: str
+    workspace: str
+    title: str
+    content: str
+    source_hash: str
+    content_hash: str
+    metadata_json: str
+    updated_at: str
+    chunks: tuple[TextChunk, ...]
+
+
+def iter_candidate_files(
+    includes: Iterable[str],
+    excludes: Iterable[str] = (),
+    workspace: Path | None = None,
+) -> list[Path]:
+    root = (workspace or Path.cwd()).expanduser().resolve()
+    files: list[Path] = []
+    seen: set[Path] = set()
+
+    for include in includes:
+        candidates = _expand_include(include, root)
+        for candidate in candidates:
+            resolved = candidate.resolve()
+            if resolved in seen:
+                continue
+            if should_exclude_path(resolved, root=root, exclude_globs=tuple(excludes)):
+                continue
+            if not looks_like_text_file(resolved):
+                continue
+            seen.add(resolved)
+            files.append(resolved)
+
+    return sorted(files, key=lambda p: p.as_posix())
+
+
+def _expand_include(include: str, root: Path) -> list[Path]:
+    path = Path(include).expanduser()
+    if not path.is_absolute():
+        path = root / path
+
+    pattern = path.as_posix()
+    if any(char in pattern for char in "*?[]"):
+        matches = [Path(match) for match in glob(pattern, recursive=True)]
+    elif path.exists():
+        matches = [path]
+    else:
+        raise FileNotFoundError(f"include path does not exist: {include}")
+
+    files: list[Path] = []
+    for match in matches:
+        if match.is_file():
+            files.append(match)
+        elif match.is_dir():
+            files.extend(p for p in match.rglob("*") if p.is_file())
+    return files
+
+
+def ingest_file(
+    path: Path,
+    workspace: Path | None = None,
+    source_type: str = "file",
+    max_bytes: int = 512 * 1024,
+    chunk_size: int = DEFAULT_CHUNK_SIZE,
+    identity_namespace: str | None = None,
+) -> IngestedRecord | None:
+    stat = path.stat()
+    if stat.st_size > max_bytes:
+        return None
+
+    raw = path.read_bytes()
+    if b"\x00" in raw:
+        return None
+
+    try:
+        content = raw.decode("utf-8")
+    except UnicodeDecodeError:
+        content = raw.decode("utf-8", errors="replace")
+
+    source_hash = hashlib.sha256(raw).hexdigest()
+    content, redactions = redact_secrets(content)
+    content_hash = hashlib.sha256(content.encode("utf-8")).hexdigest()
+    resolved = path.expanduser().resolve()
+    root = (workspace or Path.cwd()).expanduser().resolve()
+    try:
+        relative_path = resolved.relative_to(root).as_posix()
+    except ValueError:
+        relative_path = resolved.as_posix()
+    identity_scope = (
+        f"namespace:{identity_namespace}"
+        if identity_namespace is not None
+        else root.as_posix()
+    )
+    identity = f"{source_type}\0{identity_scope}\0{relative_path}"
+    record_id = hashlib.sha256(identity.encode("utf-8")).hexdigest()
+    title = title_for_path(resolved)
+    metadata = {
+        "bytes": stat.st_size,
+        "mtime": datetime.fromtimestamp(stat.st_mtime, tz=UTC).isoformat(),
+        "extension": resolved.suffix.lower(),
+        "redactions": redactions,
+        "relative_path": relative_path,
+    }
+    chunks = chunk_text(record_id, content, resolved.suffix.lower(), max_chars=chunk_size)
+
+    return IngestedRecord(
+        id=record_id,
+        source_type=source_type,
+        source_path=resolved.as_posix(),
+        workspace=root.as_posix(),
+        title=title,
+        content=content,
+        source_hash=source_hash,
+        content_hash=content_hash,
+        metadata_json=json.dumps(metadata, sort_keys=True),
+        updated_at=datetime.now(UTC).isoformat(),
+        chunks=chunks,
+    )
+
+
+def looks_like_text_file(path: Path) -> bool:
+    if path.suffix.lower() in TEXT_EXTENSIONS:
+        return True
+    with path.open("rb") as source:
+        chunk = source.read(2048)
+    if b"\x00" in chunk:
+        return False
+    if not chunk:
+        return True
+    printable = sum(1 for byte in chunk if byte in b"\n\r\t" or 32 <= byte <= 126)
+    return printable / len(chunk) > 0.85
+
+
+def title_for_path(path: Path) -> str:
+    if path.name.lower() in {"readme.md", "index.md"} and path.parent.name:
+        return path.parent.name
+    return path.name

boring_agent_memory/privacy.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+import fnmatch
+import re
+from pathlib import Path
+
+REDACTION_VERSION = 2
+
+
+DEFAULT_EXCLUDE_GLOBS = (
+    "**/.git/**",
+    "**/__pycache__/**",
+    "**/.pytest_cache/**",
+    "**/.mypy_cache/**",
+    "**/.ruff_cache/**",
+    "**/node_modules/**",
+    "**/.venv/**",
+    "**/.env",
+    "**/.env.*",
+    "**/secrets/**",
+    "**/*secret*/**",
+    "**/*.pem",
+    "**/*.key",
+    "**/id_rsa",
+    "**/id_ed25519",
+    "**/.DS_Store",
+)
+
+
+SECRET_PATTERNS = (
+    re.compile(r"sk-[A-Za-z0-9_-]{20,}"),
+    re.compile(r"github_pat_[A-Za-z0-9_]{20,}"),
+    re.compile(r"ghp_[A-Za-z0-9_]{20,}"),
+    re.compile(r"AKIA[0-9A-Z]{16}"),
+    re.compile(r"(?i)bearer\s+[A-Za-z0-9._~+/=-]{20,}"),
+    re.compile(
+        r"-----BEGIN [A-Z ]*PRIVATE KEY-----.*?-----END [A-Z ]*PRIVATE KEY-----",
+        re.DOTALL,
+    ),
+    re.compile(
+        r"(?im)^([A-Z0-9_]*(?:SECRET|TOKEN|PASSWORD|API_KEY|PRIVATE_KEY)[A-Z0-9_]*\s*=\s*).+$"
+    ),
+)
+
+
+def normalized_path(path: Path) -> str:
+    return path.expanduser().resolve().as_posix()
+
+
+def should_exclude_path(
+    path: Path,
+    root: Path | None = None,
+    exclude_globs: list[str] | tuple[str, ...] = (),
+) -> bool:
+    absolute = normalized_path(path)
+    candidates = [absolute, path.as_posix()]
+    if root is not None:
+        try:
+            candidates.append(path.resolve().relative_to(root.resolve()).as_posix())
+        except ValueError:
+            pass
+
+    for pattern in (*DEFAULT_EXCLUDE_GLOBS, *exclude_globs):
+        expanded = str(Path(pattern).expanduser()) if pattern.startswith("~") else pattern
+        if any(fnmatch.fnmatch(candidate, expanded) for candidate in candidates):
+            return True
+    return False
+
+
+def redact_secrets(content: str) -> tuple[str, int]:
+    redactions = 0
+    redacted = content
+    for pattern in SECRET_PATTERNS:
+        redacted, count = pattern.subn(_replacement, redacted)
+        redactions += count
+    return redacted, redactions
+
+
+def _replacement(match: re.Match[str]) -> str:
+    if match.lastindex:
+        replacement = f"{match.group(1)}[REDACTED]"
+    else:
+        replacement = "[REDACTED]"
+    return replacement + ("\n" * match.group(0).count("\n"))