memorytrace 0.1.0__py3-none-any.whl

engram/config.py

@@ -0,0 +1,72 @@
+"""Engram configuration."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Optional
+
+
+@dataclass
+class EngramConfig:
+    """Central configuration for the Engram memory system."""
+
+    # Storage
+    base_dir: Path = field(default_factory=lambda: Path.home() / ".engram")
+    db_name: str = "memory.db"
+    enable_markdown_export: bool = True
+    export_dir_name: str = "readable"
+
+    # Search
+    search_backend: str = "fts5"  # "fts5" | "semantic" | "hybrid"
+    max_search_results: int = 20
+    default_token_budget: int = 500
+    fts5_tokenizer: str = "unicode61 remove_diacritics 2"
+
+    # Extraction
+    extractor_backend: str = "regex"  # "regex" | "llm"
+    llm_provider: Optional[str] = None  # "anthropic" | "openai"
+    llm_model: Optional[str] = None
+
+    # Quality
+    min_confidence: float = 0.3
+    auto_resolve_threshold: float = 0.3  # trust diff for auto-supersede
+    enable_pii_detection: bool = True
+
+    # Session
+    default_agent_id: str = "default"
+
+    # Maintenance
+    thin_page_threshold: int = 300  # bytes
+    decay_days: int = 90
+    compact_days: int = 180
+
+    @property
+    def db_path(self) -> Path:
+        return self.base_dir / self.db_name
+
+    @property
+    def export_dir(self) -> Path:
+        return self.base_dir / self.export_dir_name
+
+    def ensure_dirs(self) -> None:
+        """Create required directories if they don't exist."""
+        self.base_dir.mkdir(parents=True, exist_ok=True)
+        if self.enable_markdown_export:
+            self.export_dir.mkdir(parents=True, exist_ok=True)
+
+    @classmethod
+    def from_dict(cls, data: dict) -> EngramConfig:
+        config = cls()
+        if "base_dir" in data:
+            config.base_dir = Path(data["base_dir"]).expanduser().resolve()
+        for key in (
+            "db_name", "enable_markdown_export", "export_dir_name",
+            "search_backend", "max_search_results", "default_token_budget",
+            "extractor_backend", "llm_provider", "llm_model",
+            "min_confidence", "auto_resolve_threshold", "enable_pii_detection",
+            "default_agent_id", "thin_page_threshold", "decay_days", "compact_days",
+        ):
+            if key in data:
+                setattr(config, key, data[key])
+        return config

engram/engine.py

@@ -0,0 +1,612 @@
+"""MemoryEngine — the main orchestrator that composes all modules.
+
+This is the single public entry point for the Engram memory system.
+No print() calls — all output is via structured return values.
+"""
+
+from __future__ import annotations
+
+import copy
+import json
+from dataclasses import dataclass, field
+from typing import Optional
+
+from engram.config import EngramConfig
+from engram.exceptions import EntityAlreadyExistsError, EntityNotFoundError
+from engram.models.entity import Entity, Tier
+from engram.models.fact import Fact, FactStatus
+from engram.models.quality import Action, ValidationResult
+from engram.models.relation import Relation
+from engram.models.search import SearchOptions, SearchResult
+from engram.models.session import Session
+from engram.models.source import Source, SourceType
+from engram.quality.gate import QualityGate
+from engram.search.fts5_search import FTS5Search
+from engram.session.context import build_session_context
+from engram.session.manager import SessionManager
+from engram.session.working_memory import WorkingMemory
+from engram.storage.sqlite_store import SQLiteStorage
+
+
+@dataclass
+class StoreResult:
+    """Result of a store operation."""
+    entities_created: list[Entity] = field(default_factory=list)
+    entities_updated: list[Entity] = field(default_factory=list)
+    facts_added: list[Fact] = field(default_factory=list)
+    facts_quarantined: list[Fact] = field(default_factory=list)
+    facts_conflicted: list[Fact] = field(default_factory=list)
+    relations_added: list[Relation] = field(default_factory=list)
+    validations: list[ValidationResult] = field(default_factory=list)
+
+    def to_dict(self) -> dict:
+        return {
+            "entities_created": [{"name": e.name, "type": e.entity_type} for e in self.entities_created],
+            "entities_updated": [{"name": e.name, "type": e.entity_type} for e in self.entities_updated],
+            "facts_added": [
+                {"subject": f.subject, "predicate": f.predicate, "object": f.object}
+                for f in self.facts_added
+            ],
+            "facts_quarantined": len(self.facts_quarantined),
+            "facts_conflicted": len(self.facts_conflicted),
+            "relations_added": [
+                {"type": r.relation_type} for r in self.relations_added
+            ],
+        }
+
+    def to_json(self) -> str:
+        return json.dumps(self.to_dict(), ensure_ascii=False, indent=2)
+
+
+class MemoryEngine:
+    """Main entry point. Composes storage, search, extraction, quality, sessions.
+
+    Design principles:
+    - No print() — all output via structured return values
+    - All writes go through quality gate
+    - Session tracking for every operation
+    """
+
+    def __init__(self, config: Optional[EngramConfig] = None):
+        self.config = config or EngramConfig()
+        self.config.ensure_dirs()
+
+        # Core components
+        self.storage = SQLiteStorage(self.config.db_path)
+        self.storage.initialize()
+        self.search_engine = FTS5Search(self.config.db_path)
+        self.quality_gate = QualityGate(self.config, self.storage)
+        self.session_mgr = SessionManager(self.storage)
+
+        # Optional markdown exporter
+        self._md_exporter = None
+        if self.config.enable_markdown_export:
+            from engram.storage.markdown_export import MarkdownExporter
+            self._md_exporter = MarkdownExporter(self.config.export_dir)
+
+        # Extractor (lazy-loaded based on config)
+        self._extractor = None
+
+        # Working memory per session
+        self._working_memories: dict[str, WorkingMemory] = {}
+
+    @property
+    def extractor(self):
+        if self._extractor is None:
+            self._extractor = self._build_extractor()
+        return self._extractor
+
+    def close(self) -> None:
+        """Close all connections."""
+        self.storage.close()
+        self.search_engine.close()
+
+    # ── Read Operations ──
+
+    def search(
+        self,
+        query: str,
+        options: Optional[SearchOptions] = None,
+        session_id: Optional[str] = None,
+    ) -> SearchResult:
+        """Search memory with BM25 ranking. Boosts entities from working memory."""
+        if options is None:
+            options = SearchOptions(query=query)
+        else:
+            options.query = query
+        result = self.search_engine.search(options)
+
+        # Boost entities currently in working memory (session context awareness)
+        wm = self._working_memories.get(session_id) if session_id else None
+        if wm and wm.active_entities and result.hits:
+            for hit in result.hits:
+                if hit.entity.id in wm.active_entities:
+                    hit.relevance_score *= 1.5
+            result.hits.sort(key=lambda h: h.relevance_score, reverse=True)
+            wm.add_query(query)
+
+        return result
+
+    def get_entity(self, name: str) -> Optional[Entity]:
+        """Look up an entity by name. Increments access_count."""
+        entity = self.storage.get_entity_by_name(name)
+        if entity:
+            self._touch_entity(entity)
+        return entity
+
+    def _touch_entity(self, entity: Entity) -> None:
+        """Increment access_count without triggering FTS reindex."""
+        from datetime import datetime
+        self.storage._conn.execute(
+            "UPDATE entities SET access_count = access_count + 1, last_accessed = ? WHERE id = ?",
+            (datetime.now().isoformat(), entity.id),
+        )
+        self.storage._conn.commit()
+        entity.access_count += 1
+        entity.last_accessed = datetime.now()
+
+    def get_entity_by_id(self, entity_id: str) -> Optional[Entity]:
+        """Look up an entity by ID."""
+        return self.storage.get_entity(entity_id)
+
+    def get_facts(self, entity_name: str, current_only: bool = True) -> list[Fact]:
+        """Get facts for an entity by name."""
+        entity = self.storage.get_entity_by_name(entity_name)
+        if not entity:
+            return []
+        if current_only:
+            return self.storage.get_current_facts(entity.id)
+        return self.storage.get_facts(entity.id)
+
+    def list_entities(
+        self,
+        tier: Optional[Tier] = None,
+        entity_type: Optional[str] = None,
+        limit: int = 100,
+    ) -> list[Entity]:
+        """List entities with optional filters."""
+        return self.storage.list_entities(tier=tier, entity_type=entity_type, limit=limit)
+
+    def get_relations(self, entity_name: str) -> list[Relation]:
+        """Get relations for an entity."""
+        entity = self.storage.get_entity_by_name(entity_name)
+        if not entity:
+            return []
+        return self.storage.get_relations(entity.id)
+
+    # ── Write Operations (through quality gate) ──
+
+    def store(
+        self,
+        text: str,
+        source: Optional[Source] = None,
+        session_id: Optional[str] = None,
+    ) -> StoreResult:
+        """Extract entities, facts, and relations from text and store them.
+
+        All facts pass through the quality gate before storage.
+        """
+        if not text or not text.strip():
+            return StoreResult()
+
+        source = source or Source(type=SourceType.USER_INPUT)
+        result = StoreResult()
+
+        # Step 1: Extract entities
+        ext = self.extractor
+        extracted_entities = ext.extract_entities(text)
+
+        # Step 2: Create or update entities in storage
+        all_entities: list[Entity] = []
+        for entity in extracted_entities:
+            existing = self.storage.get_entity_by_name(entity.name)
+            if existing:
+                result.entities_updated.append(existing)
+                entity.id = existing.id
+                all_entities.append(existing)
+                if session_id:
+                    self.session_mgr.record_entity_modified(session_id, existing.id)
+            else:
+                try:
+                    self.storage.create_entity(entity)
+                    result.entities_created.append(entity)
+                    all_entities.append(entity)
+                    if session_id:
+                        self.session_mgr.record_entity_modified(session_id, entity.id)
+                except EntityAlreadyExistsError:
+                    # TOCTOU: another thread created it between our check and create
+                    existing = self.storage.get_entity_by_name(entity.name)
+                    if existing:
+                        result.entities_updated.append(existing)
+                        entity.id = existing.id
+                        all_entities.append(existing)
+                    else:
+                        all_entities.append(entity)
+
+        # Step 3: Extract and validate facts (with deferred reindex)
+        extracted_facts = ext.extract_facts(text, all_entities)
+        entities_to_reindex: set[str] = set()
+
+        for fact in extracted_facts:
+            fact_source = copy.copy(source)
+            if fact_source.session_id is None and session_id:
+                fact_source.session_id = session_id
+            fact.source = fact_source
+
+            validation = self.quality_gate.validate(fact, extraction_method="regex")
+            result.validations.append(validation)
+
+            if validation.action == Action.REJECT:
+                continue
+
+            if validation.action == Action.ACCEPT:
+                self.storage.add_fact(fact, reindex=False)
+                result.facts_added.append(fact)
+                entities_to_reindex.add(fact.entity_id)
+                if session_id:
+                    self.session_mgr.record_fact_added(session_id, fact.id)
+            elif validation.action == Action.QUARANTINE:
+                self.storage.add_fact(fact, reindex=False)
+                result.facts_quarantined.append(fact)
+                entities_to_reindex.add(fact.entity_id)
+            elif validation.action == Action.FLAG_CONFLICT:
+                self.storage.add_fact(fact, reindex=False)
+                result.facts_conflicted.append(fact)
+                entities_to_reindex.add(fact.entity_id)
+                for conflict in validation.conflicts:
+                    self.storage.add_conflict({
+                        "id": conflict.conflict_id,
+                        "existing_fact_id": conflict.existing_fact.id if conflict.existing_fact else None,
+                        "new_fact_id": fact.id,
+                        "conflict_type": conflict.conflict_type,
+                        "suggested_resolution": conflict.suggested_resolution,
+                    })
+
+        # Step 3.5: Sync entity state from accepted facts
+        for fact in result.facts_added:
+            entity = next((e for e in all_entities if e.id == fact.entity_id), None)
+            if not entity:
+                continue
+            state_updated = False
+            if fact.predicate == "role" and not entity.state.role:
+                entity.state.role = fact.object
+                state_updated = True
+            elif fact.predicate == "location" and not entity.state.location:
+                entity.state.location = fact.object
+                state_updated = True
+            if state_updated:
+                self.storage.update_entity(entity)
+
+        # Step 3.6: Deferred reindex (once per entity, not per fact)
+        for eid in entities_to_reindex:
+            entity = self.storage.get_entity(eid)
+            if entity:
+                self.storage._reindex_entity(entity)
+        if entities_to_reindex:
+            self.storage._conn.commit()
+
+        # Step 4: Extract and store relations
+        extracted_relations = ext.extract_relations(text, all_entities)
+        for relation in extracted_relations:
+            self.storage.add_relation(relation)
+            result.relations_added.append(relation)
+
+        # Step 4.5: Update working memory
+        wm = self._working_memories.get(session_id) if session_id else None
+        if wm:
+            for entity in all_entities:
+                wm.touch_entity(entity.id)
+
+        # Step 5: Export to Markdown if enabled
+        if self._md_exporter:
+            for entity in result.entities_created + result.entities_updated:
+                e = self.storage.get_entity_by_name(entity.name) or entity
+                facts = self.storage.get_facts(e.id)
+                rels = self.storage.get_relations(e.id)
+                self._md_exporter.export_entity(e, facts, rels)
+
+        return result
+
+    def add_fact(
+        self,
+        entity_name: str,
+        fact_text: str,
+        predicate: str = "attribute",
+        source: Optional[Source] = None,
+        session_id: Optional[str] = None,
+    ) -> ValidationResult:
+        """Add a single fact to an entity, through quality gate."""
+        entity = self.storage.get_entity_by_name(entity_name)
+        if not entity:
+            raise EntityNotFoundError(f"Entity '{entity_name}' not found")
+
+        source = source or Source(type=SourceType.USER_INPUT)
+        fact = Fact(
+            entity_id=entity.id,
+            subject=entity.name,
+            predicate=predicate,
+            object=fact_text,
+            raw_text=fact_text,
+            source=source,
+        )
+
+        validation = self.quality_gate.validate(fact, extraction_method="manual")
+
+        if validation.action in (Action.ACCEPT, Action.QUARANTINE, Action.FLAG_CONFLICT):
+            self.storage.add_fact(fact)
+            if session_id:
+                self.session_mgr.record_fact_added(session_id, fact.id)
+                self.session_mgr.record_entity_modified(session_id, entity.id)
+
+        if validation.action == Action.FLAG_CONFLICT:
+            for conflict in validation.conflicts:
+                self.storage.add_conflict({
+                    "id": conflict.conflict_id,
+                    "existing_fact_id": conflict.existing_fact.id if conflict.existing_fact else None,
+                    "new_fact_id": fact.id,
+                    "conflict_type": conflict.conflict_type,
+                    "suggested_resolution": conflict.suggested_resolution,
+                })
+
+        return validation
+
+    def create_entity(
+        self,
+        name: str,
+        entity_type: str = "person",
+        tier: Tier = Tier.RECALL,
+        summary: str = "",
+    ) -> Entity:
+        """Manually create an entity."""
+        entity = Entity(
+            name=name,
+            entity_type=entity_type,
+            tier=tier,
+            summary=summary,
+        )
+        self.storage.create_entity(entity)
+        return entity
+
+    def merge_entities(self, primary_name: str, secondary_name: str) -> Entity:
+        """Merge secondary entity into primary. All facts, relations, aliases transfer."""
+        primary = self.storage.get_entity_by_name(primary_name)
+        secondary = self.storage.get_entity_by_name(secondary_name)
+        if not primary:
+            raise EntityNotFoundError(f"Entity '{primary_name}' not found")
+        if not secondary:
+            raise EntityNotFoundError(f"Entity '{secondary_name}' not found")
+        if primary.id == secondary.id:
+            return primary
+
+        # Transfer facts
+        self.storage._conn.execute(
+            "UPDATE facts SET entity_id = ?, subject = ? WHERE entity_id = ?",
+            (primary.id, primary.name, secondary.id),
+        )
+        # Transfer relations
+        self.storage._conn.execute(
+            "UPDATE relations SET from_entity_id = ? WHERE from_entity_id = ?",
+            (primary.id, secondary.id),
+        )
+        self.storage._conn.execute(
+            "UPDATE relations SET to_entity_id = ? WHERE to_entity_id = ?",
+            (primary.id, secondary.id),
+        )
+        # Add aliases
+        if secondary.name not in primary.aliases:
+            primary.aliases.append(secondary.name)
+        for alias in secondary.aliases:
+            if alias not in primary.aliases:
+                primary.aliases.append(alias)
+        # Merge state (secondary fills gaps)
+        if not primary.state.role and secondary.state.role:
+            primary.state.role = secondary.state.role
+        if not primary.state.affiliation and secondary.state.affiliation:
+            primary.state.affiliation = secondary.state.affiliation
+        if not primary.state.location and secondary.state.location:
+            primary.state.location = secondary.state.location
+
+        self.storage.update_entity(primary)
+        self.storage.delete_entity(secondary.id)
+        return primary
+
+    def resolve_conflict(self, conflict_id: str, resolution: str, resolved_by: str = "user") -> None:
+        """Resolve a pending conflict and update the losing fact's status."""
+        # Get conflict details before resolving
+        pending = self.storage.get_pending_conflicts()
+        conflict = None
+        for c in pending:
+            if c["id"] == conflict_id:
+                conflict = c
+                break
+
+        self.storage.resolve_conflict(conflict_id, resolution, resolved_by)
+
+        # Update fact statuses based on resolution
+        if conflict:
+            if resolution == "accept_new" and conflict.get("existing_fact_id"):
+                old_facts = self.storage._conn.execute(
+                    "SELECT id FROM facts WHERE id = ?", (conflict["existing_fact_id"],)
+                ).fetchone()
+                if old_facts:
+                    self.storage._conn.execute(
+                        "UPDATE facts SET status = 'retracted', superseded_by = ? WHERE id = ?",
+                        (conflict.get("new_fact_id"), conflict["existing_fact_id"]),
+                    )
+                    self.storage._conn.commit()
+            elif resolution == "keep_old" and conflict.get("new_fact_id"):
+                self.storage._conn.execute(
+                    "UPDATE facts SET status = 'retracted' WHERE id = ?",
+                    (conflict["new_fact_id"],),
+                )
+                self.storage._conn.commit()
+
+    # ── Session Operations ──
+
+    def start_session(self, agent_id: str = "default") -> Session:
+        """Start a new memory session."""
+        session = self.session_mgr.start_session(agent_id)
+        self._working_memories[session.session_id] = WorkingMemory(session)
+        return session
+
+    def end_session(self, session_id: str, summary: Optional[str] = None) -> Session:
+        """End a session. Runs light maintenance on modified entities."""
+        wm = self._working_memories.pop(session_id, None)
+        session = self.session_mgr.end_session(session_id, summary)
+
+        # Light maintenance: compact entities that were heavily modified this session
+        if session.entities_modified:
+            for eid in session.entities_modified[:10]:
+                try:
+                    current_facts = self.storage.get_current_facts(eid)
+                    entity = self.storage.get_entity(eid)
+                    if entity and len(current_facts) > 15:
+                        self.compact_entity(entity.name)
+                except Exception:
+                    pass  # Don't fail session end for maintenance errors
+
+        return session
+
+    def get_session_context(self, agent_id: str) -> dict:
+        """Get context for starting a new session."""
+        return build_session_context(self.storage, agent_id)
+
+    # ── Maintenance ──
+
+    def health_check(self) -> dict:
+        """Run health checks on the memory system."""
+        pending = self.storage.get_pending_conflicts()
+        return {
+            "entity_count": self.storage.entity_count(),
+            "fact_count": self.storage.fact_count(),
+            "pending_conflicts": len(pending),
+            "status": "healthy" if len(pending) == 0 else "needs_attention",
+        }
+
+    def export_markdown(self) -> int:
+        """Export all entities to Markdown."""
+        if not self._md_exporter:
+            from engram.storage.markdown_export import MarkdownExporter
+            self._md_exporter = MarkdownExporter(self.config.export_dir)
+
+        count = 0
+        offset = 0
+        batch_size = 500
+        while True:
+            batch = self.storage.list_entities(limit=batch_size, offset=offset)
+            if not batch:
+                break
+            data = []
+            for entity in batch:
+                facts = self.storage.get_facts(entity.id)
+                rels = self.storage.get_relations(entity.id)
+                data.append((entity, facts, rels))
+            count += self._md_exporter.export_all(data)
+            offset += batch_size
+        return count
+
+    def reindex(self) -> int:
+        """Rebuild the search index."""
+        return self.storage.reindex_all()
+
+    def compact_entity(self, entity_name: str) -> dict:
+        """Compact an entity's facts: keep best per predicate, update summary.
+
+        This is the core "learning" operation — transforms accumulated facts
+        into synthesized knowledge.
+        """
+        entity = self.storage.get_entity_by_name(entity_name)
+        if not entity:
+            raise EntityNotFoundError(f"Entity '{entity_name}' not found")
+
+        all_facts = self.storage.get_facts(entity.id)
+        current = [f for f in all_facts if f.is_current]
+
+        # Group by predicate, keep highest-confidence per predicate
+        from collections import defaultdict
+        by_predicate: dict[str, list[Fact]] = defaultdict(list)
+        for f in current:
+            by_predicate[f.predicate].append(f)
+
+        kept = 0
+        superseded = 0
+        for predicate, facts in by_predicate.items():
+            if len(facts) <= 1:
+                kept += 1
+                continue
+            # Sort by confidence desc, then recency
+            ranked = sorted(facts, key=lambda f: (f.confidence, f.created_at.isoformat()), reverse=True)
+            kept += 1  # keep the best
+            for old_fact in ranked[1:]:
+                old_fact.supersede(ranked[0].id)
+                self.storage.update_fact(old_fact)
+                superseded += 1
+
+        # Auto-generate summary from remaining current facts
+        remaining = self.storage.get_current_facts(entity.id)
+        summary_parts = []
+        for f in remaining[:10]:
+            summary_parts.append(f.raw_text)
+        if summary_parts:
+            entity.summary = ". ".join(summary_parts[:5])
+            if not entity.summary.endswith("."):
+                entity.summary += "."
+            self.storage.update_entity(entity)
+
+        return {
+            "entity": entity_name,
+            "facts_kept": kept,
+            "facts_superseded": superseded,
+            "new_summary": entity.summary,
+        }
+
+    def maintenance(self) -> dict:
+        """Run maintenance: decay stale facts, compact heavy entities, clean old sessions."""
+        from engram.quality.decay import find_stale_facts
+
+        results: dict = {"stale_facts": 0, "compacted_entities": 0, "cleaned_sessions": 0}
+
+        # 1. Find and flag stale facts
+        entities = self.storage.list_entities(limit=10000)
+        for entity in entities:
+            facts = self.storage.get_current_facts(entity.id)
+            stale = find_stale_facts(facts, days=self.config.decay_days)
+            for f in stale:
+                f.status = FactStatus.EXPIRED
+                self.storage.update_fact(f)
+                results["stale_facts"] += 1
+
+        # 2. Compact entities with too many current facts
+        for entity in entities:
+            current = self.storage.get_current_facts(entity.id)
+            if len(current) > 20:  # Threshold for compaction
+                self.compact_entity(entity.name)
+                results["compacted_entities"] += 1
+
+        # 3. Auto-resolve old conflicts (> 30 days)
+        from datetime import datetime, timedelta
+        pending = self.storage.get_pending_conflicts()
+        for conflict in pending:
+            # Auto-resolve based on suggested_resolution
+            if conflict.get("suggested_resolution") == "supersede":
+                self.storage.resolve_conflict(
+                    conflict["id"], "accept_new", "auto_maintenance"
+                )
+
+        return results
+
+    # ── Private ──
+
+    def _build_extractor(self):
+        """Build extractor based on config."""
+        if self.config.extractor_backend == "llm":
+            try:
+                from engram.extraction.llm_extractor import LLMExtractor
+                return LLMExtractor(
+                    provider=self.config.llm_provider or "anthropic",
+                    model=self.config.llm_model,
+                )
+            except ImportError:
+                pass  # Fall back to regex
+        from engram.extraction.regex_extractor import RegexExtractor
+        return RegexExtractor()