agentdatabase 0.1.0__py3-none-any.whl

agentdb/core/store.py

@@ -0,0 +1,373 @@
+from __future__ import annotations
+
+import json
+import sqlite3
+from contextlib import contextmanager
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Optional
+
+from .models import ConflictRecord, MemoryRecord
+
+
+# ---------------------------------------------------------------------------
+# SQLite schema
+# ---------------------------------------------------------------------------
+
+_CREATE_MEMORIES = """
+CREATE TABLE IF NOT EXISTS memories (
+    id TEXT PRIMARY KEY,
+    key TEXT NOT NULL,
+    value TEXT NOT NULL,
+    origin TEXT NOT NULL,
+    trust_level TEXT NOT NULL,
+    agent_id TEXT NOT NULL,
+    confidence REAL NOT NULL DEFAULT 1.0,
+    version INTEGER NOT NULL DEFAULT 1,
+    archived INTEGER NOT NULL DEFAULT 0,
+    retrieval_count INTEGER NOT NULL DEFAULT 0,
+    created_at TEXT NOT NULL,
+    entities TEXT NOT NULL DEFAULT '[]',
+    last_retrieved TEXT,
+    conflict_id TEXT,
+    expiry_time TEXT,
+    inferred INTEGER NOT NULL DEFAULT 0,
+    working_dirs TEXT NOT NULL DEFAULT '[]'
+)
+"""
+
+_MEMORIES_MIGRATION_COLUMNS = (
+    ("last_retrieved", "TEXT"),
+    ("conflict_id", "TEXT"),
+    ("expiry_time", "TEXT"),
+    ("inferred", "INTEGER NOT NULL DEFAULT 0"),
+    ("working_dirs", "TEXT NOT NULL DEFAULT '[]'"),
+)
+
+_CREATE_CONFLICTS = """
+CREATE TABLE IF NOT EXISTS conflicts (
+    id          TEXT PRIMARY KEY,
+    record_a_id TEXT NOT NULL,
+    record_b_id TEXT NOT NULL,
+    conflict_type TEXT NOT NULL,
+    resolved    INTEGER NOT NULL DEFAULT 0,
+    auto_resolved INTEGER NOT NULL DEFAULT 0,
+    created_at  TEXT NOT NULL
+)
+"""
+
+
+def _utcnow() -> datetime:
+    return datetime.now(timezone.utc)
+
+
+# ---------------------------------------------------------------------------
+# In-memory store backed by SQLite
+# ---------------------------------------------------------------------------
+
+class MemoryStore:
+    """SQLite-backed store for MemoryRecord objects."""
+
+    def __init__(self, db_path: Path) -> None:
+        self._db_path = db_path
+        self._by_id: dict[str, MemoryRecord] = {}
+        db_path.parent.mkdir(parents=True, exist_ok=True)
+        with self._connect() as conn:
+            conn.execute(_CREATE_MEMORIES)
+            conn.execute(_CREATE_CONFLICTS)
+            for column, coltype in _MEMORIES_MIGRATION_COLUMNS:
+                try:
+                    conn.execute(f"ALTER TABLE memories ADD COLUMN {column} {coltype}")
+                except sqlite3.OperationalError:
+                    pass  # column already exists (pre-existing database)
+        # Load existing records into memory
+        self._reload()
+
+    @contextmanager
+    def _connect(self):
+        conn = sqlite3.connect(str(self._db_path))
+        conn.row_factory = sqlite3.Row
+        conn.execute("PRAGMA journal_mode=WAL")
+        try:
+            yield conn
+            conn.commit()
+        except Exception:
+            conn.rollback()
+            raise
+        finally:
+            conn.close()
+
+    def _reload(self) -> None:
+        with self._connect() as conn:
+            rows = conn.execute("SELECT * FROM memories").fetchall()
+        for row in rows:
+            record = self._row_to_record(row)
+            self._by_id[record.id] = record
+
+    def _row_to_record(self, row: sqlite3.Row) -> MemoryRecord:
+        return MemoryRecord(
+            id=row["id"],
+            key=row["key"],
+            value=json.loads(row["value"]),
+            origin=row["origin"],
+            trust_level=row["trust_level"],
+            agent_id=row["agent_id"],
+            confidence=row["confidence"],
+            version=row["version"],
+            archived=bool(row["archived"]),
+            retrieval_count=row["retrieval_count"],
+            created_at=datetime.fromisoformat(row["created_at"]),
+            entities=json.loads(row["entities"]),
+            last_retrieved=(
+                datetime.fromisoformat(row["last_retrieved"])
+                if row["last_retrieved"] else None
+            ),
+            conflict_id=row["conflict_id"],
+            expiry_time=(
+                datetime.fromisoformat(row["expiry_time"])
+                if row["expiry_time"] else None
+            ),
+            inferred=bool(row["inferred"]) if row["inferred"] is not None else False,
+            working_dirs=json.loads(row["working_dirs"] if "working_dirs" in row.keys() else "[]"),
+        )
+
+    def write(self, record: MemoryRecord) -> MemoryRecord:
+        with self._connect() as conn:
+            conn.execute(
+                """
+                INSERT OR REPLACE INTO memories
+                    (id, key, value, origin, trust_level, agent_id, confidence,
+                     version, archived, retrieval_count, created_at, entities,
+                     last_retrieved, conflict_id, expiry_time, inferred, working_dirs)
+                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+                """,
+                (
+                    record.id,
+                    record.key,
+                    json.dumps(record.value),
+                    str(record.origin),
+                    str(record.trust_level),
+                    record.agent_id,
+                    record.confidence,
+                    record.version,
+                    int(record.archived),
+                    record.retrieval_count,
+                    record.created_at.isoformat(),
+                    json.dumps(record.entities),
+                    record.last_retrieved.isoformat() if record.last_retrieved else None,
+                    record.conflict_id,
+                    record.expiry_time.isoformat() if record.expiry_time else None,
+                    int(record.inferred),
+                    json.dumps(record.working_dirs),
+                ),
+            )
+        self._by_id[record.id] = record
+        return record
+
+    def get_by_key(self, key: str) -> Optional[MemoryRecord]:
+        for record in self._by_id.values():
+            if record.key == key and not record.archived:
+                return record
+        return None
+
+    def get_by_id(self, record_id: str) -> Optional[MemoryRecord]:
+        return self._by_id.get(record_id)
+
+    def delete_by_key(self, key: str) -> None:
+        record = self.get_by_key(key)
+        if record is None:
+            return
+        record.archived = True
+        self._persist_field(record.id, "archived", 1)
+
+    def _persist_field(self, record_id: str, field_name: str, value: Any) -> None:
+        with self._connect() as conn:
+            conn.execute(
+                f"UPDATE memories SET {field_name} = ? WHERE id = ?",
+                (value, record_id),
+            )
+
+    def increment_retrieval_count(self, record_id: str) -> None:
+        record = self._by_id.get(record_id)
+        if record:
+            record.retrieval_count += 1
+            self._persist_field(record_id, "retrieval_count", record.retrieval_count)
+            record.last_retrieved = _utcnow()
+            self._persist_field(record_id, "last_retrieved", record.last_retrieved.isoformat())
+
+    def update_confidence(self, record_id: str, new_confidence: float) -> None:
+        record = self._by_id.get(record_id)
+        if record:
+            record.confidence = round(new_confidence, 10)
+            self._persist_field(record_id, "confidence", record.confidence)
+
+    def archive_record(self, record_id: str) -> None:
+        """Mark a record as archived by ID."""
+        record = self._by_id.get(record_id)
+        if record:
+            record.archived = True
+            self._persist_field(record_id, "archived", 1)
+
+    def list_all(self) -> list[MemoryRecord]:
+        return [r for r in self._by_id.values() if not r.archived]
+
+    def list_by_agent(self, agent_id: str) -> list[MemoryRecord]:
+        return [r for r in self._by_id.values() if r.agent_id == agent_id and not r.archived]
+
+    def search(self, query: str, limit: int = 5, agent_id: Optional[str] = None) -> list[MemoryRecord]:
+        """
+        Simple text search: find records whose key or value contains
+        any word from query (case-insensitive). agent_id, if given, is
+        applied during the scan (before the `limit` slice below) so an
+        isolated caller's own matches can't be truncated away by another
+        agent's larger matching set.
+        """
+        words = [w.lower() for w in query.split() if w]
+        if not words:
+            return []
+
+        results: list[MemoryRecord] = []
+        for record in self._by_id.values():
+            if record.archived:
+                continue
+            if agent_id is not None and record.agent_id != agent_id:
+                continue
+            key_lower = record.key.lower()
+            value_str = (
+                json.dumps(record.value) if isinstance(record.value, (dict, list)) else str(record.value)
+            ).lower()
+            for word in words:
+                if word in key_lower or word in value_str:
+                    results.append(record)
+                    break
+
+        return results[:limit]
+
+    def search_by_entities(self, entities: list[str], agent_id: Optional[str] = None) -> list[MemoryRecord]:
+        """
+        Find records where record.entities has overlap with the given entities.
+        Fallback: search key and value for entity name substrings.
+        agent_id, if given, is applied inline in BOTH passes below — not as a
+        filter on the final chosen list — so an isolated caller's fallback
+        search is scoped to their own records rather than being gated by
+        whether OTHER agents had entity-tag hits.
+        """
+        if not entities:
+            return []
+
+        entity_set = set(e.lower() for e in entities)
+
+        # First pass: exact entity match
+        matches: list[MemoryRecord] = []
+        for record in self._by_id.values():
+            if record.archived:
+                continue
+            if agent_id is not None and record.agent_id != agent_id:
+                continue
+            record_entities = set(e.lower() for e in (record.entities or []))
+            if record_entities & entity_set:
+                matches.append(record)
+
+        if matches:
+            return matches
+
+        # Fallback: substring search in key and value
+        fallback: list[MemoryRecord] = []
+        for record in self._by_id.values():
+            if record.archived:
+                continue
+            if agent_id is not None and record.agent_id != agent_id:
+                continue
+            key_lower = record.key.lower()
+            value_str = (
+                json.dumps(record.value) if isinstance(record.value, (dict, list)) else str(record.value)
+            ).lower()
+            for entity in entities:
+                entity_lower = entity.lower()
+                if entity_lower in key_lower or entity_lower in value_str:
+                    fallback.append(record)
+                    break
+
+        return fallback
+
+    # ------------------------------------------------------------------
+    # Conflict storage
+    # ------------------------------------------------------------------
+
+    def store_conflict(self, conflict: ConflictRecord) -> None:
+        now = _utcnow().isoformat()
+        with self._connect() as conn:
+            conn.execute(
+                """
+                INSERT OR REPLACE INTO conflicts
+                    (id, record_a_id, record_b_id, conflict_type, resolved, auto_resolved, created_at)
+                VALUES (?, ?, ?, ?, ?, ?, ?)
+                """,
+                (
+                    conflict.id,
+                    conflict.record_a_id,
+                    conflict.record_b_id,
+                    conflict.conflict_type,
+                    int(conflict.resolved),
+                    int(conflict.auto_resolved),
+                    now,
+                ),
+            )
+
+    def get_conflicts_for_records(self, record_ids: list[str]) -> list[ConflictRecord]:
+        if not record_ids:
+            return []
+        placeholders = ",".join("?" * len(record_ids))
+        with self._connect() as conn:
+            rows = conn.execute(
+                f"""
+                SELECT id, record_a_id, record_b_id, conflict_type, resolved, auto_resolved
+                FROM conflicts
+                WHERE record_a_id IN ({placeholders}) OR record_b_id IN ({placeholders})
+                """,
+                record_ids + record_ids,
+            ).fetchall()
+        return [
+            ConflictRecord(
+                id=row["id"],
+                record_a_id=row["record_a_id"],
+                record_b_id=row["record_b_id"],
+                conflict_type=row["conflict_type"],
+                resolved=bool(row["resolved"]),
+                auto_resolved=bool(row["auto_resolved"]),
+            )
+            for row in rows
+        ]
+
+    def list_conflicts(self, resolved: bool = False) -> list[ConflictRecord]:
+        """
+        Return conflict records from the conflicts table.
+
+        resolved=False (default) -> pending only (resolved = 0).
+        resolved=True -> all conflicts (pending + resolved).
+        """
+        query = (
+            "SELECT id, record_a_id, record_b_id, conflict_type, resolved, auto_resolved "
+            "FROM conflicts"
+        )
+        with self._connect() as conn:
+            if not resolved:
+                query += " WHERE resolved = 0"
+            rows = conn.execute(query).fetchall()
+        return [
+            ConflictRecord(
+                id=row["id"],
+                record_a_id=row["record_a_id"],
+                record_b_id=row["record_b_id"],
+                conflict_type=row["conflict_type"],
+                resolved=bool(row["resolved"]),
+                auto_resolved=bool(row["auto_resolved"]),
+            )
+            for row in rows
+        ]
+
+    def update_conflict_id(self, record_id: str, conflict_id: str) -> None:
+        self._persist_field(record_id, "conflict_id", conflict_id)
+        record = self._by_id.get(record_id)
+        if record is not None:
+            record.conflict_id = conflict_id

agentdb/core/system.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+import sqlite3
+from contextlib import contextmanager
+from pathlib import Path
+
+
+class SystemStats:
+    def __init__(self, db_path: Path) -> None:
+        self._db_path = db_path
+
+    @contextmanager
+    def _connect(self):
+        conn = sqlite3.connect(str(self._db_path))
+        conn.row_factory = sqlite3.Row
+        conn.execute("PRAGMA journal_mode=WAL")
+        try:
+            yield conn
+            conn.commit()
+        except Exception:
+            conn.rollback()
+            raise
+        finally:
+            conn.close()
+
+    def agentdb_stats(self) -> dict:
+        with self._connect() as conn:
+            # Total memories (non-archived)
+            total = conn.execute(
+                "SELECT COUNT(*) FROM memories WHERE archived = 0"
+            ).fetchone()[0]
+
+            # Average confidence
+            avg_conf_row = conn.execute(
+                "SELECT AVG(confidence) FROM memories WHERE archived = 0"
+            ).fetchone()
+            avg_confidence = avg_conf_row[0] if avg_conf_row[0] is not None else 0.0
+
+            # Low confidence count (< 0.5)
+            low_conf = conn.execute(
+                "SELECT COUNT(*) FROM memories WHERE archived = 0 AND confidence < 0.5"
+            ).fetchone()[0]
+
+            # Retrieval count total (sum of retrieval_count across all memories)
+            retrieval_total_row = conn.execute(
+                "SELECT SUM(retrieval_count) FROM memories WHERE archived = 0"
+            ).fetchone()
+            retrieval_total = retrieval_total_row[0] if retrieval_total_row[0] is not None else 0
+
+            # Conflict pending count - from audit log
+            conflict_pending = 0  # We don't track ConflictRecord table; use 0
+
+            # Outcome positive rate
+            outcome_rate = None
+            try:
+                outcomes_rows = conn.execute(
+                    "SELECT outcome_value FROM outcomes"
+                ).fetchall()
+                if len(outcomes_rows) >= 10:
+                    positive = sum(1 for r in outcomes_rows if r["outcome_value"] >= 0.5)
+                    outcome_rate = positive / len(outcomes_rows)
+            except Exception:
+                pass
+
+            # External conflict rate
+            external_conflict_rate = None
+            try:
+                ext_rows = conn.execute(
+                    "SELECT had_external_context, had_conflict FROM external_retrieves"
+                ).fetchall()
+                external_retrieves = [r for r in ext_rows if r["had_external_context"]]
+                if external_retrieves:
+                    conflicts = sum(1 for r in external_retrieves if r["had_conflict"])
+                    external_conflict_rate = conflicts / len(external_retrieves)
+            except Exception:
+                pass
+
+        return {
+            "total_memories": total,
+            "avg_confidence": avg_confidence or 0.0,
+            "low_confidence_count": low_conf,
+            "conflict_pending_count": conflict_pending,
+            "outcome_positive_rate": outcome_rate,
+            "retrieval_count_total": retrieval_total,
+            "external_conflict_rate": external_conflict_rate,
+        }

agentdb/embeddings/__init__.py

@@ -0,0 +1,7 @@
+from .provider import EmbeddingProvider, FunctionEmbeddingProvider, LocalEmbeddingProvider, _cosine
+from .store import EmbeddingStore
+
+__all__ = [
+    "EmbeddingProvider", "FunctionEmbeddingProvider", "LocalEmbeddingProvider",
+    "_cosine", "EmbeddingStore",
+]

agentdb/embeddings/provider.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+import math
+import threading
+from typing import Callable, Protocol, runtime_checkable
+
+
+@runtime_checkable
+class EmbeddingProvider(Protocol):
+    @property
+    def model_name(self) -> str: ...
+    @property
+    def dimensions(self) -> int: ...
+    def embed(self, text: str) -> list[float]: ...
+    def embed_batch(self, texts: list[str]) -> list[list[float]]: ...
+
+
+def _cosine(a: list[float], b: list[float]) -> float:
+    """Cosine similarity between two vectors. Safe for zero vectors."""
+    dot = sum(x * y for x, y in zip(a, b))
+    norm_a = math.sqrt(sum(x * x for x in a))
+    norm_b = math.sqrt(sum(x * x for x in b))
+    if norm_a == 0.0 or norm_b == 0.0:
+        return 0.0
+    return dot / (norm_a * norm_b)
+
+
+class FunctionEmbeddingProvider:
+    def __init__(self, fn: Callable[[str], list[float]], model_name: str = "custom") -> None:
+        self._fn = fn
+        self._model_name = model_name
+        self._dimensions: int | None = None
+
+    @property
+    def model_name(self) -> str:
+        return self._model_name
+
+    @property
+    def dimensions(self) -> int:
+        if self._dimensions is None:
+            try:
+                probe = self._fn("probe")
+                self._dimensions = len(probe)
+            except Exception:
+                self._dimensions = 0
+        return self._dimensions
+
+    def embed(self, text: str) -> list[float]:
+        return self._fn(text)
+
+    def embed_batch(self, texts: list[str]) -> list[list[float]]:
+        return [self._fn(t) for t in texts]
+
+
+class LocalEmbeddingProvider:
+    _MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
+    _model = None
+    _lock = threading.Lock()
+
+    @property
+    def model_name(self) -> str:
+        return self._MODEL_NAME
+
+    @property
+    def dimensions(self) -> int:
+        return 384
+
+    def _load_model(self) -> None:
+        if LocalEmbeddingProvider._model is None:
+            try:
+                from sentence_transformers import SentenceTransformer  # type: ignore
+            except ImportError as e:
+                raise RuntimeError(
+                    "sentence-transformers is required for LocalEmbeddingProvider. "
+                    "Install it with: pip install 'agentdb[local-embed]'"
+                ) from e
+            import logging
+            logging.getLogger("agentdb").info(
+                "agentdb: loading local embedding model (first use)"
+            )
+            LocalEmbeddingProvider._model = SentenceTransformer(self._MODEL_NAME)
+
+    def embed(self, text: str) -> list[float]:
+        with LocalEmbeddingProvider._lock:
+            self._load_model()
+            vecs = LocalEmbeddingProvider._model.encode(
+                [text], normalize_embeddings=True
+            )
+        return vecs[0].tolist()
+
+    def embed_batch(self, texts: list[str]) -> list[list[float]]:
+        if not texts:
+            return []
+        with LocalEmbeddingProvider._lock:
+            self._load_model()
+            vecs = LocalEmbeddingProvider._model.encode(
+                texts, normalize_embeddings=True
+            )
+        return [v.tolist() for v in vecs]