memoryagent-lib 0.1.1__py3-none-any.whl

memoryagent/indexers.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+from typing import List
+
+from memoryagent.models import MemoryItem, MemoryType, StorageTier
+from memoryagent.storage.base import FeatureStore, GraphStore, VectorIndex
+from memoryagent.utils import tokenize
+
+
+class EpisodicIndexer:
+    """Indexes episodic content into the vector index."""
+
+    def __init__(self, vector_index: VectorIndex) -> None:
+        self.vector_index = vector_index
+
+    async def index_hot(self, item: MemoryItem) -> None:
+        await self.vector_index.upsert(
+            item.id,
+            text=item.text(),
+            metadata={"owner": item.owner, "tier": StorageTier.HOT.value, "type": item.type.value, "item": item},
+        )
+
+    async def index_archive(self, item: MemoryItem) -> None:
+        await self.vector_index.upsert(
+            item.id,
+            text=item.summary,
+            metadata={"owner": item.owner, "tier": StorageTier.ARCHIVE_INDEX.value, "type": item.type.value, "item": item},
+        )
+
+
+class SemanticGraphIndexer:
+    """Extracts simple fact-like triples from tags for demo use."""
+
+    def __init__(self, graph_store: GraphStore) -> None:
+        self.graph_store = graph_store
+
+    async def index(self, item: MemoryItem) -> None:
+        if item.type != MemoryType.SEMANTIC:
+            return
+        if len(item.tags) < 2:
+            return
+        subject = item.tags[0]
+        for tag in item.tags[1:]:
+            await self.graph_store.upsert_fact(item.owner, subject, "related_to", tag)
+
+
+class PerceptualIndexer:
+    """Summarizes perceptual inputs into feature store entries."""
+
+    def __init__(self, feature_store: FeatureStore) -> None:
+        self.feature_store = feature_store
+
+    async def index(self, item: MemoryItem) -> None:
+        if item.type != MemoryType.PERCEPTUAL:
+            return
+        payload = {
+            "summary": item.summary,
+            "tags": item.tags,
+            "confidence": item.confidence,
+        }
+        await self.feature_store.write_feature(item.owner, payload)

memoryagent/models.py

@@ -0,0 +1,156 @@
+from __future__ import annotations
+
+from datetime import datetime, timedelta, timezone
+from enum import Enum
+from typing import Any, Dict, List, Optional
+from uuid import UUID, uuid4
+
+from pydantic import BaseModel, Field
+
+
+def utc_now() -> datetime:
+    return datetime.now(timezone.utc)
+
+
+class MemoryType(str, Enum):
+    WORKING = "working"
+    EPISODIC = "episodic"
+    SEMANTIC = "semantic"
+    PERCEPTUAL = "perceptual"
+
+
+class StorageTier(str, Enum):
+    HOT = "hot"
+    COLD = "cold"
+    ARCHIVE_INDEX = "archive_index"
+
+
+class MemoryItem(BaseModel):
+    """Canonical memory item stored in metadata storage."""
+
+    id: UUID = Field(default_factory=uuid4)
+    type: MemoryType
+    owner: str
+    summary: str
+    created_at: datetime = Field(default_factory=utc_now)
+    updated_at: datetime = Field(default_factory=utc_now)
+    last_accessed: Optional[datetime] = None
+    tier: StorageTier = StorageTier.HOT
+    pointer: Dict[str, Any] = Field(default_factory=dict)
+    content: Optional[Any] = None
+    tags: List[str] = Field(default_factory=list)
+    ttl_seconds: Optional[int] = None
+    confidence: float = 0.5
+    authority: float = 0.5
+    stability: float = 0.5
+
+    def is_expired(self, now: Optional[datetime] = None) -> bool:
+        if self.ttl_seconds is None:
+            return False
+        now = now or utc_now()
+        return self.created_at + timedelta(seconds=self.ttl_seconds) <= now
+
+    def text(self) -> str:
+        if isinstance(self.content, str):
+            return self.content
+        return self.summary
+
+
+class MemoryEvent(BaseModel):
+    """Developer-facing input for memory writes."""
+
+    content: Any
+    type: MemoryType = MemoryType.WORKING
+    owner: str
+    summary: Optional[str] = None
+    tags: List[str] = Field(default_factory=list)
+    ttl_seconds: Optional[int] = None
+    confidence: float = 0.5
+    authority: float = 0.5
+    stability: float = 0.5
+    pointer: Dict[str, Any] = Field(default_factory=dict)
+
+    def to_item(self) -> MemoryItem:
+        summary = self.summary or (self.content if isinstance(self.content, str) else str(self.content))
+        return MemoryItem(
+            type=self.type,
+            owner=self.owner,
+            summary=summary,
+            content=self.content,
+            tags=self.tags,
+            ttl_seconds=self.ttl_seconds,
+            confidence=self.confidence,
+            authority=self.authority,
+            stability=self.stability,
+            pointer=self.pointer,
+        )
+
+
+class MemoryQuery(BaseModel):
+    text: str
+    owner: str
+    types: Optional[List[MemoryType]] = None
+    top_k: int = 10
+    time_range_seconds: Optional[int] = None
+
+
+class RetrievalPlan(BaseModel):
+    """Routing + budgets + thresholds for retrieval pipeline."""
+
+    hot_top_k: int = 30
+    archive_top_k: int = 30
+    cold_fetch_limit: int = 20
+    cold_fetch_min_score: float = 0.25
+    hot_confidence: float = 0.62
+    archive_confidence: float = 0.72
+    cold_fetch_confidence: float = 0.58
+    max_results: int = 50
+    max_context_tokens: int = 600
+
+
+class ScoredMemory(BaseModel):
+    item: MemoryItem
+    score: float
+    tier: StorageTier
+    explanation: Optional[str] = None
+
+
+class ConfidenceReport(BaseModel):
+    total: float
+    semantic_relevance: float
+    coverage: float
+    temporal_fit: float
+    authority: float
+    consistency: float
+    recommendation: str
+
+
+class MemoryBlock(BaseModel):
+    text: str
+    item_id: UUID
+    memory_type: MemoryType
+    tier: StorageTier
+    score: float
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+
+
+class RetrievalTrace(BaseModel):
+    steps: List[str] = Field(default_factory=list)
+    escalations: List[str] = Field(default_factory=list)
+    sources: List[str] = Field(default_factory=list)
+
+    def add_step(self, text: str) -> None:
+        self.steps.append(text)
+
+    def add_escalation(self, text: str) -> None:
+        self.escalations.append(text)
+
+
+class MemoryBundle(BaseModel):
+    query: str
+    results: List[ScoredMemory]
+    blocks: List[MemoryBlock]
+    confidence: ConfidenceReport
+    used_tiers: List[StorageTier]
+    trace: RetrievalTrace
+    warnings: List[str] = Field(default_factory=list)

memoryagent/policy.py

@@ -0,0 +1,171 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Iterable, List, Optional
+
+from memoryagent.models import MemoryEvent, MemoryItem, MemoryType, StorageTier
+from memoryagent.utils import tokenize
+
+
+@dataclass
+class MemoryDecision:
+    store: bool
+    memory_type: MemoryType
+    summary: Optional[str] = None
+    tags: Optional[List[str]] = None
+    reasons: Optional[List[str]] = None
+
+
+class ConversationMemoryPolicy:
+    """Decides if a turn should be persisted to memory."""
+
+    def should_store(
+        self,
+        owner: str,
+        history: List[str],
+        user_message: str,
+        assistant_message: str,
+    ) -> MemoryDecision:
+        raise NotImplementedError
+
+    def to_event(self, owner: str, decision: MemoryDecision) -> Optional[MemoryEvent]:
+        if not decision.store or not decision.summary:
+            return None
+        return MemoryEvent(
+            content=decision.summary,
+            type=decision.memory_type,
+            owner=owner,
+            tags=decision.tags or [],
+        )
+
+
+class HeuristicMemoryPolicy(ConversationMemoryPolicy):
+    """Default heuristic policy for deciding what to store."""
+
+    def __init__(
+        self,
+        min_tokens: int = 24,
+        novelty_threshold: float = 0.65,
+        short_turn_min_novelty: float = 0.8,
+        preference_keywords: Optional[Iterable[str]] = None,
+    ) -> None:
+        self.min_tokens = min_tokens
+        self.novelty_threshold = novelty_threshold
+        self.short_turn_min_novelty = short_turn_min_novelty
+        self.preference_keywords = set(
+            k.lower()
+            for k in (preference_keywords or ["prefer", "always", "never", "likes", "dislikes"])
+        )
+
+    def should_store(
+        self,
+        owner: str,
+        history: List[str],
+        user_message: str,
+        assistant_message: str,
+    ) -> MemoryDecision:
+        combined = f"{user_message} {assistant_message}"
+        tokens = tokenize(combined)
+        reasons: List[str] = []
+        memory_type = MemoryType.EPISODIC
+
+        is_preference = any(word in combined.lower() for word in self.preference_keywords)
+        if len(tokens) < self.min_tokens:
+            reasons.append("short_turn")
+        if is_preference:
+            memory_type = MemoryType.SEMANTIC
+            reasons.append("preference_signal")
+
+        if history:
+            recent_entries = history[-3:]
+            recent_text = " ".join(self._history_entry_text(entry) for entry in recent_entries)
+            novelty = 1.0 - self._overlap_ratio(tokens, tokenize(recent_text))
+            novelty_floor = self.short_turn_min_novelty if len(tokens) < self.min_tokens else self.novelty_threshold
+            if novelty < novelty_floor:
+                reasons.append("low_novelty")
+
+        if is_preference:
+            store = True
+        else:
+            store = "short_turn" not in reasons and "low_novelty" not in reasons
+        summary = self._summarize(user_message, assistant_message, memory_type)
+        tags = ["conversation", memory_type.value]
+        return MemoryDecision(store=store, memory_type=memory_type, summary=summary, tags=tags, reasons=reasons)
+
+    def _overlap_ratio(self, tokens_a: List[str], tokens_b: List[str]) -> float:
+        if not tokens_a or not tokens_b:
+            return 0.0
+        set_a = set(tokens_a)
+        set_b = set(tokens_b)
+        return len(set_a & set_b) / max(1, len(set_a | set_b))
+
+    def _summarize(self, user_message: str, assistant_message: str, memory_type: MemoryType) -> str:
+        if memory_type == MemoryType.SEMANTIC:
+            return f"User preference: {user_message.strip()}"
+        return f"User asked: {user_message.strip()} | Assistant replied: {assistant_message.strip()}"
+
+    def _history_entry_text(self, entry) -> str:
+        if isinstance(entry, str):
+            return entry
+        if isinstance(entry, dict):
+            if "user" in entry and "assistant" in entry:
+                return f"User: {entry['user']} Assistant: {entry['assistant']}"
+            if "role" in entry and "text" in entry:
+                return f"{entry['role']}: {entry['text']}"
+        return str(entry)
+
+
+@dataclass
+class RoutingDecision:
+    write_hot: bool
+    write_vector: bool
+    write_features: bool
+    archive_cold: bool
+    reasons: List[str]
+
+
+class MemoryRoutingPolicy:
+    """Concrete policy for routing memory writes across tiers and indexes."""
+
+    def __init__(
+        self,
+        hot_min_confidence: float = 0.4,
+        cold_min_confidence: float = 0.55,
+        vector_min_confidence: float = 0.5,
+        feature_min_confidence: float = 0.45,
+    ) -> None:
+        self.hot_min_confidence = hot_min_confidence
+        self.cold_min_confidence = cold_min_confidence
+        self.vector_min_confidence = vector_min_confidence
+        self.feature_min_confidence = feature_min_confidence
+
+    def route(self, item: MemoryItem) -> RoutingDecision:
+        reasons: List[str] = []
+        confidence = item.confidence
+
+        write_hot = confidence >= self.hot_min_confidence
+        if not write_hot:
+            reasons.append("low_confidence_hot")
+
+        write_vector = confidence >= self.vector_min_confidence and item.type != MemoryType.WORKING
+        if not write_vector:
+            reasons.append("skip_vector")
+
+        write_features = item.type == MemoryType.PERCEPTUAL and confidence >= self.feature_min_confidence
+        if not write_features and item.type == MemoryType.PERCEPTUAL:
+            reasons.append("skip_features")
+
+        archive_cold = (
+            item.type in {MemoryType.EPISODIC, MemoryType.SEMANTIC, MemoryType.PERCEPTUAL}
+            and confidence >= self.cold_min_confidence
+        )
+        if not archive_cold:
+            reasons.append("skip_cold")
+
+        return RoutingDecision(
+            write_hot=write_hot,
+            write_vector=write_vector,
+            write_features=write_features,
+            archive_cold=archive_cold,
+            reasons=reasons,
+        )

memoryagent/retrieval.py

@@ -0,0 +1,154 @@
+from __future__ import annotations
+
+from typing import List, Sequence
+
+from memoryagent.confidence import evaluate_confidence
+from memoryagent.models import (
+    MemoryBlock,
+    MemoryBundle,
+    MemoryQuery,
+    MemoryType,
+    RetrievalPlan,
+    RetrievalTrace,
+    ScoredMemory,
+    StorageTier,
+)
+from memoryagent.storage.base import MetadataStore, ObjectStore, VectorIndex
+from memoryagent.utils import clamp
+
+
+class RetrievalOrchestrator:
+    def __init__(
+        self,
+        metadata_store: MetadataStore,
+        vector_index: VectorIndex,
+        object_store: ObjectStore,
+        plan: RetrievalPlan,
+    ) -> None:
+        self.metadata_store = metadata_store
+        self.vector_index = vector_index
+        self.object_store = object_store
+        self.plan = plan
+
+    async def retrieve(self, query: MemoryQuery) -> MemoryBundle:
+        used_tiers: List[StorageTier] = []
+        warnings: List[str] = []
+        trace = RetrievalTrace()
+
+        trace.add_step("hot search per type")
+        hot_results: List[ScoredMemory] = []
+        types = query.types or [MemoryType.WORKING, MemoryType.EPISODIC, MemoryType.SEMANTIC, MemoryType.PERCEPTUAL]
+        per_type_limit = max(1, self.plan.hot_top_k // max(1, len(types)))
+        for mem_type in types:
+            hot_results.extend(
+                await self.vector_index.query(
+                    query,
+                    filters={
+                        "owner": query.owner,
+                        "tier": StorageTier.HOT.value,
+                        "types": [mem_type],
+                    },
+                    limit=per_type_limit,
+                )
+            )
+        used_tiers.append(StorageTier.HOT)
+        confidence = evaluate_confidence(query, hot_results)
+
+        results = list(hot_results)
+
+        if confidence.total < self.plan.hot_confidence:
+            trace.add_escalation("hot confidence below threshold; searching archive")
+            archive_results = await self.vector_index.query(
+                query,
+                filters={"owner": query.owner, "tier": StorageTier.ARCHIVE_INDEX.value, "types": query.types},
+                limit=self.plan.archive_top_k,
+            )
+            if archive_results:
+                results.extend(archive_results)
+                used_tiers.append(StorageTier.ARCHIVE_INDEX)
+                confidence = evaluate_confidence(query, results)
+
+            if confidence.total < self.plan.cold_fetch_confidence:
+                trace.add_escalation("archive confidence low; fetching cold payloads")
+                cold_candidates = [
+                    item for item in archive_results if item.score >= self.plan.cold_fetch_min_score
+                ][: self.plan.cold_fetch_limit]
+                for item in cold_candidates:
+                    pointer = item.item.pointer.get("object_key")
+                    if not pointer:
+                        continue
+                    payload = await self.object_store.get(pointer)
+                    if payload is None:
+                        warnings.append(f"Missing cold object: {pointer}")
+                        continue
+                    if isinstance(payload, list):
+                        payload = next((p for p in payload if p.get("id") == str(item.item.id)), None)
+                        if payload is None:
+                            warnings.append(f"Missing id {item.item.id} in daily notes: {pointer}")
+                            continue
+                    hydrated = item.item.model_copy(update={"content": payload, "tier": StorageTier.COLD})
+                    results.append(
+                        ScoredMemory(item=hydrated, score=item.score, tier=StorageTier.COLD, explanation="cold hydrate")
+                    )
+                if cold_candidates:
+                    used_tiers.append(StorageTier.COLD)
+                    confidence = evaluate_confidence(query, results)
+
+        hydrated = await self._hydrate(results)
+        reranked = self._rerank(self._dedupe(hydrated))
+        blocks = self._to_blocks(reranked)
+        trace.sources = [f"{item.item.type}:{item.tier}" for item in reranked[:10]]
+
+        return MemoryBundle(
+            query=query.text,
+            results=reranked,
+            blocks=blocks,
+            confidence=confidence,
+            used_tiers=used_tiers,
+            trace=trace,
+            warnings=warnings,
+        )
+
+    def _rerank(self, results: Sequence[ScoredMemory]) -> List[ScoredMemory]:
+        def score(item: ScoredMemory) -> float:
+            return clamp(0.75 * item.score + 0.25 * item.item.confidence)
+
+        reranked = sorted(results, key=score, reverse=True)
+        return reranked[: self.plan.max_results]
+
+    def _to_blocks(self, results: Sequence[ScoredMemory]) -> List[MemoryBlock]:
+        blocks: List[MemoryBlock] = []
+        for item in results:
+            text = item.item.text()
+            blocks.append(
+                MemoryBlock(
+                    text=text,
+                    item_id=item.item.id,
+                    memory_type=item.item.type,
+                    tier=item.tier,
+                    score=item.score,
+                    metadata={"owner": item.item.owner, "tags": item.item.tags},
+                )
+            )
+        return blocks
+
+    async def _hydrate(self, results: Sequence[ScoredMemory]) -> List[ScoredMemory]:
+        hydrated: List[ScoredMemory] = []
+        for item in results:
+            if item.item.content is not None and item.item.tags:
+                hydrated.append(item)
+                continue
+            full_item = await self.metadata_store.get(item.item.id)
+            if full_item is None:
+                hydrated.append(item)
+                continue
+            hydrated.append(item.model_copy(update={"item": full_item}))
+        return hydrated
+
+    def _dedupe(self, results: Sequence[ScoredMemory]) -> List[ScoredMemory]:
+        best = {}
+        for item in results:
+            key = str(item.item.id)
+            if key not in best or item.score > best[key].score:
+                best[key] = item
+        return list(best.values())

memoryagent/storage/base.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Iterable, List, Optional
+
+from memoryagent.models import MemoryItem, MemoryQuery, ScoredMemory
+
+
+class MetadataStore(ABC):
+    """Stores canonical MemoryItem metadata."""
+
+    @abstractmethod
+    async def upsert(self, item: MemoryItem) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    async def get(self, item_id) -> Optional[MemoryItem]:
+        raise NotImplementedError
+
+    @abstractmethod
+    async def delete(self, item_id) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    async def list_by_owner(self, owner: str) -> List[MemoryItem]:
+        raise NotImplementedError
+
+    @abstractmethod
+    async def list_by_owner_and_type(self, owner: str, types: Iterable[str]) -> List[MemoryItem]:
+        raise NotImplementedError
+
+    @abstractmethod
+    async def update_access(self, item_id) -> None:
+        raise NotImplementedError
+
+
+class VectorIndex(ABC):
+    """Vector or lexical index supporting similarity search."""
+
+    @abstractmethod
+    async def upsert(self, item_id, text: str, metadata: dict) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    async def delete(self, item_id) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    async def query(self, query: MemoryQuery, filters: dict, limit: int) -> List[ScoredMemory]:
+        raise NotImplementedError
+
+
+class GraphStore(ABC):
+    """Stores semantic relationships (facts, preferences, rules)."""
+
+    @abstractmethod
+    async def upsert_fact(self, owner: str, subject: str, predicate: str, obj: str) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    async def query_related(self, owner: str, subject: str, limit: int) -> List[str]:
+        raise NotImplementedError
+
+
+class ObjectStore(ABC):
+    """Stores cold memory payloads."""
+
+    @abstractmethod
+    async def put(self, key: str, payload: dict) -> str:
+        raise NotImplementedError
+
+    @abstractmethod
+    async def get(self, key: str) -> Optional[dict]:
+        raise NotImplementedError
+
+
+class FeatureStore(ABC):
+    """Stores perceptual aggregates (time-series or feature logs)."""
+
+    @abstractmethod
+    async def write_feature(self, owner: str, payload: dict) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    async def query_features(self, owner: str, limit: int) -> List[dict]:
+        raise NotImplementedError