agent-memory-sdk 0.1.0__py3-none-any.whl

agent_memory/explain.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+from agent_memory.models import MemoryAction, MemoryDecision, RetrievalResult
+from agent_memory.policy import DecisionPolicy, DefaultPolicy
+
+
+def build_score_breakdown(
+    result: RetrievalResult,
+    policy: DecisionPolicy,
+) -> dict[str, float]:
+    if isinstance(policy, DefaultPolicy):
+        return policy.score_breakdown(result.entry, result.semantic_score, result.keyword_score)
+
+    return {
+        "semantic_score": result.semantic_score,
+        "keyword_score": result.keyword_score,
+        "recency_score": 0.0,
+        "confidence_score": result.entry.confidence,
+        "usage_score": 0.0,
+        "policy_score": result.final_score,
+        "final_score": result.decision_score,
+    }
+
+
+def build_reasons(result: RetrievalResult, policy: DecisionPolicy) -> list[str]:
+    reasons: list[str] = []
+    breakdown = build_score_breakdown(result, policy)
+
+    if result.semantic_score >= 0.85:
+        reasons.append("high semantic match")
+    elif result.semantic_score >= 0.70:
+        reasons.append("moderate semantic match")
+
+    if result.keyword_score >= 0.50:
+        reasons.append("keyword match")
+
+    if breakdown.get("recency_score", 0) >= 0.80:
+        reasons.append("recent memory")
+    elif result.entry.access_count > 0:
+        reasons.append("prior usage")
+
+    if result.entry.access_count >= 3:
+        reasons.append("frequently used")
+
+    if result.entry.confidence >= 0.90:
+        reasons.append("high confidence memory")
+
+    if result.entry.requires_verification:
+        reasons.append("requires verification")
+
+    if result.entry.type.value in {"fact", "tool_output", "workflow"}:
+        reasons.append("freshness-sensitive type")
+
+    return reasons
+
+
+def enrich_decision(
+    decision: MemoryDecision,
+    policy: DecisionPolicy,
+    *,
+    replay_threshold: float,
+    restore_threshold: float,
+) -> MemoryDecision:
+    """Attach matched IDs, reason tags, and score breakdown to a decision."""
+    if not decision.context:
+        decision.matched = []
+        decision.reasons = ["no matching memories"]
+        return decision
+
+    best = decision.context[0]
+    decision.matched = [r.entry.id for r in decision.context]
+    decision.reasons = build_reasons(best, policy)
+    decision.scores = build_score_breakdown(best, policy)
+
+    if decision.action == MemoryAction.REPLAY and best.decision_score >= replay_threshold:
+        if "high semantic match" not in decision.reasons and best.semantic_score >= 0.85:
+            decision.reasons.insert(0, "exact or near-exact query match")
+    elif decision.action == MemoryAction.NONE:
+        if best.decision_score < restore_threshold:
+            decision.reasons.append("below restore threshold")
+
+    return decision
+
+
+def format_explanation(decision: MemoryDecision) -> str:
+    if not decision.scores:
+        return "No score breakdown available (no memory matches)."
+
+    lines = [
+        f"action: {decision.action.value}",
+        f"confidence: {decision.confidence:.2f}",
+        f"matched: {', '.join(decision.matched) or 'none'}",
+        "",
+    ]
+
+    if decision.reasons:
+        lines.append("reasons:")
+        for reason in decision.reasons:
+            lines.append(f"  - {reason}")
+        lines.append("")
+
+    lines.append("scores:")
+    label_map = {
+        "semantic_score": "semantic_score",
+        "keyword_score": "keyword_score",
+        "recency_score": "recency_score",
+        "confidence_score": "confidence_score",
+        "usage_score": "usage_score",
+        "policy_score": "policy_score",
+        "final_score": "final_score",
+    }
+    for key, label in label_map.items():
+        if key in decision.scores:
+            lines.append(f"  {label}: {decision.scores[key]:.2f}")
+
+    return "\n".join(lines)

agent_memory/manager.py

@@ -0,0 +1,348 @@
+from __future__ import annotations
+
+import asyncio
+from pathlib import Path
+
+from agent_memory.decision import DecisionEngine
+from agent_memory.models import MemoryDecision, MemoryEntry, MemoryScope, MemoryState, MemoryType
+from agent_memory.policy import DecisionPolicy, DefaultPolicy
+from agent_memory.retriever import MemoryRetriever
+from agent_memory.sqlite_store import SqliteMemoryStore
+from agent_memory.store import ChromaDBStore, MemoryStore
+from agent_memory.ttl import parse_ttl
+
+
+class Memory:
+    """High-level SDK for persistent agent memory with a decision layer."""
+
+    def __init__(
+        self,
+        persist_dir: str | Path = ".agent_memory",
+        collection_name: str = "agent_memories",
+        policy: DecisionPolicy | None = None,
+        replay_threshold: float = 0.85,
+        restore_threshold: float = 0.70,
+        verify_threshold: float = 0.80,
+        backend: str = "sqlite",  # "chromadb" or "sqlite"
+    ) -> None:
+        self.store: MemoryStore
+        if backend == "sqlite":
+            self.store = SqliteMemoryStore(persist_dir=persist_dir, collection_name=collection_name)
+        elif backend == "chromadb":
+            self.store = ChromaDBStore(persist_dir=persist_dir, collection_name=collection_name)
+        else:
+            raise ValueError(f"Unknown backend: {backend}. Use 'sqlite' or 'chromadb'")
+        self._policy = policy or DefaultPolicy()
+        self.retriever = MemoryRetriever(self.store, policy=self._policy)
+        self.decision_engine = DecisionEngine(
+            self.retriever,
+            policy=self._policy,
+            replay_threshold=replay_threshold,
+            restore_threshold=restore_threshold,
+            verify_threshold=verify_threshold,
+        )
+
+    def remember(
+        self,
+        query: str,
+        response: str,
+        *,
+        content: str | None = None,
+        type: MemoryType | str = MemoryType.CONVERSATION,
+        scope: MemoryScope | str = MemoryScope.USER,
+        metadata: dict | None = None,
+        tags: list[str] | None = None,
+        confidence: float = 1.0,
+        requires_verification: bool = False,
+        ttl: str | int | float | None = None,
+    ) -> MemoryEntry:
+        memory_type = MemoryType(type) if isinstance(type, str) else type
+        memory_scope = MemoryScope(scope) if isinstance(scope, str) else scope
+        entry = MemoryEntry(
+            query=query,
+            response=response,
+            content=content or response,
+            type=memory_type,
+            scope=memory_scope,
+            metadata=metadata or {},
+            tags=tags or [],
+            confidence=confidence,
+            requires_verification=requires_verification,
+            expires_at=parse_ttl(ttl),
+        )
+        entry.refresh_state()
+        return self.store.store(entry)
+
+    async def aremember(
+        self,
+        query: str,
+        response: str,
+        *,
+        content: str | None = None,
+        type: MemoryType | str = MemoryType.CONVERSATION,
+        scope: MemoryScope | str = MemoryScope.USER,
+        metadata: dict | None = None,
+        tags: list[str] | None = None,
+        confidence: float = 1.0,
+        requires_verification: bool = False,
+        ttl: str | int | float | None = None,
+    ) -> MemoryEntry:
+        """Async version of remember()."""
+        return await asyncio.to_thread(
+            self.remember,
+            query,
+            response,
+            content=content,
+            type=type,
+            scope=scope,
+            metadata=metadata,
+            tags=tags,
+            confidence=confidence,
+            requires_verification=requires_verification,
+            ttl=ttl,
+        )
+
+    def resolve(
+        self,
+        query: str,
+        *,
+        mode: str = "auto",
+        top_k: int = 3,
+        scope: list[MemoryScope | str] | None = None,
+        enable_verify: bool = True,
+    ) -> MemoryDecision:
+        scopes = None
+        if scope:
+            scopes = [MemoryScope(s) if isinstance(s, str) else s for s in scope]
+        return self.decision_engine.decide(
+            query,
+            mode=mode,
+            top_k=top_k,
+            scopes=scopes,
+            enable_verify=enable_verify,
+        )
+
+    async def aresolve(
+        self,
+        query: str,
+        *,
+        mode: str = "auto",
+        top_k: int = 3,
+        scope: list[MemoryScope | str] | None = None,
+        enable_verify: bool = True,
+    ) -> MemoryDecision:
+        """Async version of resolve()."""
+        return await asyncio.to_thread(
+            self.resolve,
+            query,
+            mode=mode,
+            top_k=top_k,
+            scope=scope,
+            enable_verify=enable_verify,
+        )
+
+    def list(
+        self,
+        limit: int = 100,
+        offset: int = 0,
+        *,
+        scope: list[MemoryScope | str] | None = None,
+        include_archived: bool = False,
+        type: MemoryType | str | None = None,
+    ) -> list[MemoryEntry]:
+        scopes = [MemoryScope(s) if isinstance(s, str) else s for s in scope] if scope else None
+        memory_type = MemoryType(type) if isinstance(type, str) else type
+        return self.store.list_all(
+            limit=limit,
+            offset=offset,
+            scopes=scopes,
+            include_archived=include_archived,
+            memory_type=memory_type,
+        )
+
+    async def alist(
+        self,
+        limit: int = 100,
+        offset: int = 0,
+        *,
+        scope: list[MemoryScope | str] | None = None,
+        include_archived: bool = False,
+        type: MemoryType | str | None = None,
+    ) -> list[MemoryEntry]:
+        """Async version of list()."""
+        return await asyncio.to_thread(
+            self.list,
+            limit=limit,
+            offset=offset,
+            scope=scope,
+            include_archived=include_archived,
+            type=type,
+        )
+
+    def get(self, memory_id: str) -> MemoryEntry | None:
+        return self.store.get(memory_id)
+
+    async def aget(self, memory_id: str) -> MemoryEntry | None:
+        """Async version of get()."""
+        return await asyncio.to_thread(self.get, memory_id)
+
+    def forget(self, memory_id: str) -> bool:
+        return self.store.delete(memory_id)
+
+    async def aforget(self, memory_id: str) -> bool:
+        """Async version of forget()."""
+        return await asyncio.to_thread(self.forget, memory_id)
+
+    def archive(self, memory_id: str) -> MemoryEntry | None:
+        entry = self.store.get(memory_id)
+        if not entry:
+            return None
+        entry.archived = True
+        entry.refresh_state()
+        return self.store.update(entry)
+
+    async def aarchive(self, memory_id: str) -> MemoryEntry | None:
+        """Async version of archive()."""
+        return await asyncio.to_thread(self.archive, memory_id)
+
+    def cleanup(self, *, delete: bool = False) -> dict[str, int]:
+        """
+        Mark expired memories as expired, optionally deleting them.
+
+        Returns counts: {"expired": N, "deleted": M}
+        """
+        entries = self.store.list_all(limit=10_000, include_archived=True, include_expired=True)
+        expired_count = 0
+        deleted_count = 0
+
+        for entry in entries:
+            entry.refresh_state()
+            if not entry.is_expired:
+                continue
+            if delete:
+                if self.store.delete(entry.id):
+                    deleted_count += 1
+            else:
+                if entry.state != MemoryState.EXPIRED:
+                    entry.state = MemoryState.EXPIRED
+                    self.store.update(entry)
+                expired_count += 1
+
+        return {"expired": expired_count, "deleted": deleted_count}
+
+    async def acleanup(self, *, delete: bool = False) -> dict[str, int]:
+        """Async version of cleanup()."""
+        return await asyncio.to_thread(self.cleanup, delete=delete)
+
+    def stats(self) -> dict:
+        """Return aggregate memory and usage statistics."""
+        entries = self.store.list_all(limit=10_000, include_archived=True, include_expired=True)
+        by_state: dict[str, int] = {}
+        by_type: dict[str, int] = {}
+        total_access = 0
+
+        for entry in entries:
+            entry.refresh_state()
+            by_state[entry.state.value] = by_state.get(entry.state.value, 0) + 1
+            by_type[entry.type.value] = by_type.get(entry.type.value, 0) + 1
+            total_access += entry.access_count
+
+        return {
+            "total": len(entries),
+            "by_state": by_state,
+            "by_type": by_type,
+            "total_access_count": total_access,
+        }
+
+    async def astats(self) -> dict:
+        """Async version of stats()."""
+        return await asyncio.to_thread(self.stats)
+
+    def consolidate(self, similarity_threshold: float = 0.95) -> list[MemoryEntry]:
+        """
+        Merge near-duplicate active memories into summary entries.
+        Returns newly created summary memories.
+        """
+        entries = self.store.list_all(limit=10_000)
+        created: list[MemoryEntry] = []
+        seen: set[str] = set()
+
+        for entry in entries:
+            if entry.id in seen or entry.archived or entry.type == MemoryType.SUMMARY:
+                continue
+
+            duplicates = [entry]
+            for other in entries:
+                if other.id == entry.id or other.id in seen or other.archived:
+                    continue
+                if entry.type != other.type or entry.scope != other.scope:
+                    continue
+                hits = self.store.search(entry.query, top_k=5)
+                for hit, score in hits:
+                    if hit.id == other.id and score >= similarity_threshold:
+                        duplicates.append(other)
+                        break
+
+            if len(duplicates) < 2:
+                continue
+
+            for dup in duplicates:
+                seen.add(dup.id)
+                dup.archived = True
+                self.store.update(dup)
+
+            merged_query = duplicates[0].query
+            merged_content = "\n".join(f"- {d.response}" for d in duplicates)
+            summary = self.remember(
+                query=merged_query,
+                response=merged_content,
+                content=merged_content,
+                type=MemoryType.SUMMARY,
+                scope=duplicates[0].scope,
+                tags=list({tag for d in duplicates for tag in d.tags}),
+                metadata={"consolidated_from": [d.id for d in duplicates]},
+            )
+            created.append(summary)
+
+        return created
+
+    async def aconsolidate(self, similarity_threshold: float = 0.95) -> list[MemoryEntry]:
+        """Async version of consolidate()."""
+        return await asyncio.to_thread(self.consolidate, similarity_threshold=similarity_threshold)
+
+    def format_restore_context(self, decision: MemoryDecision) -> str:
+        if not decision.context:
+            return ""
+
+        lines = ["## Retrieved Memory Context", ""]
+        for result in decision.context:
+            entry = result.entry
+            lines.append(
+                f"### Memory {result.rank} (score: {result.final_score:.2f}, type: {entry.type.value})"
+            )
+            lines.append(f"**Prior query:** {entry.query}")
+            lines.append(f"**Prior response:** {entry.response}")
+            if entry.tags:
+                lines.append(f"**Tags:** {', '.join(entry.tags)}")
+            lines.append("")
+        return "\n".join(lines)
+
+    def format_verify_context(self, decision: MemoryDecision) -> str:
+        if not decision.memory:
+            return ""
+        entry = decision.memory
+        return (
+            "## Memory Pending Verification\n\n"
+            f"**Query:** {entry.query}\n"
+            f"**Stored response:** {entry.response}\n"
+            f"**Type:** {entry.type.value}\n"
+            f"**Last updated:** {entry.updated_at.isoformat()}\n\n"
+            "Validate this memory with available tools before reusing or regenerating."
+        )
+
+    # Backward compatibility
+    def list_memories(self, limit: int = 100, offset: int = 0) -> list[MemoryEntry]:
+        return self.list(limit=limit, offset=offset)
+
+
+MemoryManager = Memory

agent_memory/models.py

@@ -0,0 +1,172 @@
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from enum import Enum
+from typing import Any
+from uuid import uuid4
+
+from pydantic import BaseModel, Field, model_validator
+
+# Memory types that default to verify-before-reuse when similarity is moderate.
+VERIFY_TYPES: frozenset[str] = frozenset({"fact", "tool_output", "workflow"})
+
+
+class MemoryAction(str, Enum):
+    """How the agent should use a retrieved memory."""
+
+    REPLAY = "replay"
+    RESTORE = "restore"
+    VERIFY = "verify"
+    NONE = "none"
+
+
+class MemoryType(str, Enum):
+    CONVERSATION = "conversation"
+    FACT = "fact"
+    WORKFLOW = "workflow"
+    TOOL_OUTPUT = "tool_output"
+    DOCUMENT = "document"
+    SUMMARY = "summary"
+    CODE = "code"
+    PREFERENCE = "preference"
+
+
+class MemoryScope(str, Enum):
+    SESSION = "session"
+    USER = "user"
+    PROJECT = "project"
+    WORKSPACE = "workspace"
+    TEAM = "team"
+    GLOBAL = "global"
+
+
+class MemoryState(str, Enum):
+    ACTIVE = "active"
+    ARCHIVED = "archived"
+    EXPIRED = "expired"
+    DELETED = "deleted"
+
+
+class MemoryEntry(BaseModel):
+    """A stored memory with structured metadata."""
+
+    id: str = Field(default_factory=lambda: str(uuid4()))
+    query: str
+    response: str
+    content: str = ""
+    type: MemoryType = MemoryType.CONVERSATION
+    scope: MemoryScope = MemoryScope.USER
+    metadata: dict[str, Any] = Field(default_factory=dict)
+    tags: list[str] = Field(default_factory=list)
+    confidence: float = Field(default=1.0, ge=0.0, le=1.0)
+    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+    updated_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+    access_count: int = 0
+    last_accessed_at: datetime | None = None
+    archived: bool = False
+    requires_verification: bool = False
+    expires_at: datetime | None = None
+    state: MemoryState = MemoryState.ACTIVE
+
+    @property
+    def is_expired(self) -> bool:
+        if self.state == MemoryState.EXPIRED:
+            return True
+        if self.expires_at is None:
+            return False
+        now = datetime.now(timezone.utc)
+        expires = self.expires_at
+        if expires.tzinfo is None:
+            expires = expires.replace(tzinfo=timezone.utc)
+        return now >= expires
+
+    def refresh_state(self) -> None:
+        if self.archived:
+            self.state = MemoryState.ARCHIVED
+        elif self.is_expired:
+            self.state = MemoryState.EXPIRED
+        else:
+            self.state = MemoryState.ACTIVE
+
+    @model_validator(mode="after")
+    def default_content(self) -> MemoryEntry:
+        if not self.content:
+            self.content = self.response
+        return self
+
+    def touch(self) -> None:
+        self.access_count += 1
+        self.last_accessed_at = datetime.now(timezone.utc)
+        self.updated_at = datetime.now(timezone.utc)
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "id": self.id,
+            "content": self.content,
+            "query": self.query,
+            "response": self.response,
+            "type": self.type.value,
+            "scope": self.scope.value,
+            "created_at": self.created_at.isoformat(),
+            "updated_at": self.updated_at.isoformat(),
+            "last_accessed": self.last_accessed_at.isoformat() if self.last_accessed_at else None,
+            "access_count": self.access_count,
+            "confidence": self.confidence,
+            "tags": self.tags,
+            "archived": self.archived,
+            "state": self.state.value,
+            "expires_at": self.expires_at.isoformat() if self.expires_at else None,
+        }
+
+
+class RetrievalResult(BaseModel):
+    """A memory match from hybrid retrieval."""
+
+    entry: MemoryEntry
+    semantic_score: float
+    keyword_score: float = 0.0
+    final_score: float = 0.0
+    rank: int = 0
+
+    @property
+    def similarity(self) -> float:
+        """Backward-compatible alias for decision score."""
+        return self.decision_score
+
+    @property
+    def decision_score(self) -> float:
+        """Score used for action thresholds — never below raw retrieval signals."""
+        retrieval = max(self.semantic_score, 0.7 * self.semantic_score + 0.3 * self.keyword_score)
+        return max(self.final_score, retrieval)
+
+
+class MemoryDecision(BaseModel):
+    """Agent-facing decision on how to answer a query."""
+
+    action: MemoryAction
+    query: str
+    confidence: float
+    response: str | None = None
+    memory: MemoryEntry | None = None
+    context: list[RetrievalResult] = Field(default_factory=list)
+    reason: str = ""
+    matched: list[str] = Field(default_factory=list)
+    reasons: list[str] = Field(default_factory=list)
+    scores: dict[str, float] = Field(default_factory=dict)
+
+    def explain(self) -> str:
+        from agent_memory.explain import format_explanation
+
+        return format_explanation(self)
+
+    def __repr__(self) -> str:
+        matched = ", ".join(self.matched[:3])
+        if len(self.matched) > 3:
+            matched += ", ..."
+        return (
+            f"Decision(action={self.action.value!r}, confidence={self.confidence:.2f}, "
+            f"matched=[{matched}], reasons={self.reasons!r})"
+        )
+
+    def __str__(self) -> str:
+        return self.__repr__()