memorytrace 0.1.0__py3-none-any.whl

engram/extraction/regex_extractor.py

@@ -0,0 +1,188 @@
+"""Rule-based extractor — NER + fact extraction + relation extraction."""
+
+from __future__ import annotations
+
+import copy
+import re
+from typing import Optional
+
+from engram.models.entity import Entity
+from engram.models.fact import Fact
+from engram.models.relation import Relation
+from engram.models.source import Source
+from engram.extraction.ner.english import extract_english_entities
+from engram.extraction.ner.korean import extract_korean_entities
+from engram.extraction.ner.cjk import extract_cjk_entities
+
+
+# Fact extraction patterns: "Subject <verb> Object"
+_FACT_PATTERNS = [
+    # English: "X is/was the Y of Z"
+    re.compile(
+        r'([A-Z][a-z]+(?:\s+[A-Z][a-z]+)+)\s+'
+        r'(?:is|was|serves?\s+as|became|joined|founded|leads?|runs?|heads?)\s+'
+        r'(?:the\s+)?(.{1,200}?)(?:\.|,|;|$)',
+        re.MULTILINE,
+    ),
+    # English: "X, the Y of Z"
+    re.compile(
+        r'([A-Z][a-z]+(?:\s+[A-Z][a-z]+)+),\s+'
+        r'(?:the\s+)?(\w[\w\s]{0,200}?)(?:\.|,|;|$)',
+        re.MULTILINE,
+    ),
+]
+
+# Relation patterns
+_RELATION_PATTERNS = [
+    # "X is the CEO of Y" → X -[CEO_OF]-> Y
+    re.compile(
+        r'([A-Z][a-z]+(?:\s+[A-Z][a-z]+)+)\s+'
+        r'(?:is|was)\s+(?:the\s+)?'
+        r'(CEO|CTO|CFO|COO|founder|co-founder|president|director|head|member|partner)\s+'
+        r'(?:of|at)\s+'
+        r'([A-Z][\w\s]*?)(?:\.|,|;|$)',
+        re.IGNORECASE | re.MULTILINE,
+    ),
+    # "X works at/for Y"
+    re.compile(
+        r'([A-Z][a-z]+(?:\s+[A-Z][a-z]+)+)\s+'
+        r'(?:works?\s+(?:at|for)|joined|left)\s+'
+        r'([A-Z][\w\s]*?)(?:\.|,|;|$)',
+        re.MULTILINE,
+    ),
+]
+
+# Predicate normalization
+_ROLE_KEYWORDS = frozenset({
+    "ceo", "cto", "cfo", "coo", "founder", "co-founder",
+    "president", "director", "head", "member", "partner",
+    "manager", "lead", "engineer", "scientist", "analyst",
+})
+
+
+class RegexExtractor:
+    """Rule-based entity, fact, and relation extraction.
+
+    Improvements over MemKraft:
+    - Context-aware entity type classification
+    - Dictionary-based Korean NER (not "any 2-4 hangul chars")
+    - Deduplication at extraction time
+    - Confidence scoring per pattern
+    """
+
+    def __init__(self, default_source: Optional[Source] = None):
+        self.default_source = default_source or Source()
+
+    def extract_entities(self, text: str) -> list[Entity]:
+        """Extract named entities from text across all supported languages."""
+        entities: list[Entity] = []
+        seen: set[str] = set()
+
+        # English NER
+        for e in extract_english_entities(text):
+            key = e.name.lower()
+            if key not in seen:
+                seen.add(key)
+                entities.append(e)
+
+        # Korean NER
+        for e in extract_korean_entities(text):
+            key = e.name.lower()
+            if key not in seen:
+                seen.add(key)
+                entities.append(e)
+
+        # CJK NER
+        for e in extract_cjk_entities(text):
+            key = e.name.lower()
+            if key not in seen:
+                seen.add(key)
+                entities.append(e)
+
+        return entities
+
+    def extract_facts(self, text: str, entities: list[Entity]) -> list[Fact]:
+        """Extract facts from text, grounded against known entities."""
+        facts: list[Fact] = []
+        entity_names = {e.name.lower(): e for e in entities}
+
+        for pattern in _FACT_PATTERNS:
+            for m in pattern.finditer(text):
+                subject_text = m.group(1).strip()
+                object_text = m.group(2).strip()
+
+                # Ground subject to known entity
+                subject_key = subject_text.lower()
+                entity = entity_names.get(subject_key)
+                if not entity:
+                    continue
+
+                # Determine predicate
+                predicate = self._classify_predicate(object_text)
+
+                # Truncate overly long objects
+                if len(object_text) > 200:
+                    object_text = object_text[:200]
+
+                # Copy source to avoid shared mutable reference
+                fact_source = copy.copy(self.default_source)
+                facts.append(Fact(
+                    entity_id=entity.id,
+                    subject=entity.name,
+                    predicate=predicate,
+                    object=object_text,
+                    raw_text=m.group(0).strip(),
+                    source=fact_source,
+                    confidence=0.5,  # Will be recomputed by quality gate
+                ))
+
+        return facts
+
+    def extract_relations(self, text: str, entities: list[Entity]) -> list[Relation]:
+        """Extract directed relations between entities."""
+        relations: list[Relation] = []
+        entity_names = {e.name.lower(): e for e in entities}
+
+        for pattern in _RELATION_PATTERNS:
+            for m in pattern.finditer(text):
+                groups = m.groups()
+                if len(groups) == 3:
+                    # Pattern: X is ROLE of Y
+                    subject = groups[0].strip()
+                    role = groups[1].strip().upper()
+                    obj = groups[2].strip()
+                    relation_type = f"{role}_OF"
+                elif len(groups) == 2:
+                    # Pattern: X works at Y
+                    subject = groups[0].strip()
+                    obj = groups[1].strip()
+                    relation_type = "WORKS_AT"
+                else:
+                    continue
+
+                # Ground to known entities
+                from_entity = entity_names.get(subject.lower())
+                to_entity = entity_names.get(obj.lower())
+
+                if from_entity and to_entity:
+                    relations.append(Relation(
+                        from_entity_id=from_entity.id,
+                        to_entity_id=to_entity.id,
+                        relation_type=relation_type,
+                    ))
+
+        return relations
+
+    def _classify_predicate(self, object_text: str) -> str:
+        """Determine the predicate type from the object text."""
+        lower = object_text.lower()
+        for keyword in _ROLE_KEYWORDS:
+            if keyword in lower:
+                return "role"
+        if any(w in lower for w in ("founded", "started", "created", "launched")):
+            return "action"
+        if any(w in lower for w in ("based in", "located", "lives in", "moved to")):
+            return "location"
+        if any(w in lower for w in ("invested", "fund", "raised", "capital")):
+            return "investment"
+        return "attribute"

engram/integrations/__init__.py

@@ -0,0 +1 @@
+"""Integrations for Engram (MCP, SDK)."""

engram/integrations/mcp_server.py

@@ -0,0 +1,213 @@
+"""MCP server for Engram — Claude Code / Codex integration.
+
+Requires: pip install engram[mcp]
+
+Usage:
+    python -m engram.integrations.mcp_server
+
+    Or via CLI:
+    engram serve
+"""
+
+from __future__ import annotations
+
+import atexit
+import json
+import sys
+from typing import Optional
+
+from engram.engine import MemoryEngine
+from engram.models.search import SearchOptions
+from engram.models.source import Source, SourceType
+
+# Lazy engine instance (created on first tool call)
+_engine: Optional[MemoryEngine] = None
+
+
+def _get_engine() -> MemoryEngine:
+    global _engine
+    if _engine is None:
+        _engine = MemoryEngine()
+        atexit.register(_engine.close)
+    return _engine
+
+
+def _safe_source_type(value: str) -> SourceType:
+    """Parse source type string, falling back to agent_inference."""
+    try:
+        return SourceType(value)
+    except ValueError:
+        return SourceType.AGENT_INFERENCE
+
+
+def run_server(transport: str = "stdio", port: int = 8080) -> None:
+    """Start the MCP server."""
+    try:
+        from mcp.server import Server
+        from mcp.server.stdio import stdio_server
+    except ImportError:
+        print("MCP package not installed. Run: pip install engram[mcp]", file=sys.stderr)
+        sys.exit(1)
+
+    app = Server("engram")
+
+    @app.tool()
+    async def memory_search(
+        query: str,
+        max_results: int = 5,
+        max_tokens: int = 500,
+        min_confidence: float = 0.0,
+    ) -> str:
+        """Search persistent memory for relevant context about people, organizations, or past decisions.
+
+        Use when the user mentions someone/something discussed before,
+        or references "last time", "before", "we discussed".
+
+        Args:
+            query: Natural language search query
+            max_results: Maximum entities to return
+            max_tokens: Token budget for response
+            min_confidence: Minimum fact confidence (0.0-1.0)
+        """
+        engine = _get_engine()
+        options = SearchOptions(
+            query=query,
+            max_results=max_results,
+            max_tokens=max_tokens,
+            min_confidence=min_confidence,
+        )
+        result = engine.search(query, options)
+        return result.to_agent_context(max_tokens=max_tokens)
+
+    @app.tool()
+    async def memory_store(
+        content: str,
+        source_type: str = "agent_inference",
+        confidence: float = 0.7,
+    ) -> str:
+        """Store new factual information in persistent memory.
+
+        Use when new facts about people/orgs/projects are shared.
+        Do NOT store: greetings, opinions, or speculative statements.
+
+        Args:
+            content: Factual content to store
+            source_type: One of: direct_speech, document, api, web, agent_inference, user_input
+            confidence: How confident (0.0-1.0)
+        """
+        engine = _get_engine()
+        source = Source(
+            type=_safe_source_type(source_type),
+            confidence=confidence,
+            channel="mcp",
+        )
+        result = engine.store(content, source=source)
+        return json.dumps(result.to_dict(), ensure_ascii=False)
+
+    @app.tool()
+    async def memory_get_entity(name: str) -> str:
+        """Get all known information about a specific entity.
+
+        Args:
+            name: Entity name to look up
+        """
+        engine = _get_engine()
+        entity = engine.get_entity(name)
+        if not entity:
+            return f"No entity found with name '{name}'."
+        facts = engine.get_facts(name)
+        lines = [
+            f"Name: {entity.name}",
+            f"Type: {entity.entity_type}",
+            f"Tier: {entity.tier.value}",
+        ]
+        if entity.summary:
+            lines.append(f"Summary: {entity.summary}")
+        if entity.state.role:
+            lines.append(f"Role: {entity.state.role}")
+        if entity.state.affiliation:
+            lines.append(f"Affiliation: {entity.state.affiliation}")
+        if facts:
+            lines.append(f"Facts ({len(facts)}):")
+            for f in facts[:10]:
+                lines.append(f"  - {f.raw_text} [{f.confidence:.0%}]")
+        return "\n".join(lines)
+
+    @app.tool()
+    async def memory_list_entities(
+        entity_type: str = "",
+        tier: str = "",
+        limit: int = 20,
+    ) -> str:
+        """List known entities with optional filters.
+
+        Args:
+            entity_type: Filter by type (person, organization, project, concept)
+            tier: Filter by tier (core, recall, archival)
+            limit: Max entries to return
+        """
+        engine = _get_engine()
+        from engram.models.entity import Tier
+        entities = engine.list_entities(
+            tier=Tier(tier) if tier else None,
+            entity_type=entity_type or None,
+            limit=limit,
+        )
+        if not entities:
+            return "No entities found."
+        lines = []
+        for e in entities:
+            lines.append(f"- {e.name} ({e.entity_type}, {e.tier.value})")
+        return "\n".join(lines)
+
+    @app.tool()
+    async def memory_session_start(agent_id: str = "claude-code") -> str:
+        """Start a new memory session. Call at the beginning of a conversation.
+
+        Args:
+            agent_id: Identifier for the calling agent
+        """
+        engine = _get_engine()
+        session = engine.start_session(agent_id)
+        context = engine.get_session_context(agent_id)
+        return json.dumps({
+            "session_id": session.session_id,
+            "entity_count": context["entity_count"],
+            "previous_summary": context.get("previous_summary", ""),
+        }, ensure_ascii=False)
+
+    @app.tool()
+    async def memory_session_end(session_id: str, summary: str = "") -> str:
+        """End the current memory session.
+
+        Args:
+            session_id: Session ID to end
+            summary: Brief summary of what was discussed
+        """
+        engine = _get_engine()
+        session = engine.end_session(session_id, summary=summary or None)
+        return f"Session ended. Duration: {session.duration_minutes or 0}min."
+
+    @app.tool()
+    async def memory_resolve_conflict(conflict_id: str, resolution: str) -> str:
+        """Resolve a data conflict in memory.
+
+        Args:
+            conflict_id: Conflict ID to resolve
+            resolution: One of: accept_new, keep_old, merge
+        """
+        engine = _get_engine()
+        engine.resolve_conflict(conflict_id, resolution)
+        return f"Conflict {conflict_id[:8]} resolved: {resolution}"
+
+    # Run server
+    if transport == "stdio":
+        import asyncio
+        asyncio.run(stdio_server(app))
+    else:
+        print(f"Transport '{transport}' not yet supported.", file=sys.stderr)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    run_server()

engram/integrations/sdk.py

@@ -0,0 +1,194 @@
+"""Programmatic Python SDK — clean wrapper around MemoryEngine.
+
+Usage:
+    from engram.integrations.sdk import EngramSDK
+
+    sdk = EngramSDK()
+    sdk.store("Simon Kim is the CEO of Hashed.")
+    results = sdk.search("CEO")
+    sdk.close()
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+from engram.config import EngramConfig
+from engram.engine import MemoryEngine, StoreResult
+from engram.models.entity import Entity, Tier
+from engram.models.quality import ValidationResult
+from engram.models.search import SearchOptions, SearchResult
+from engram.models.session import Session
+from engram.models.source import Source, SourceType
+
+
+class EngramSDK:
+    """High-level SDK for programmatic access to Engram.
+
+    Provides a clean API without CLI concerns.
+    All methods return structured objects — no print().
+    """
+
+    def __init__(self, config: Optional[EngramConfig] = None):
+        self.engine = MemoryEngine(config)
+        self._current_session: Optional[Session] = None
+
+    def close(self) -> None:
+        try:
+            if self._current_session:
+                self.end_session()
+        finally:
+            self.engine.close()
+
+    def __enter__(self) -> EngramSDK:
+        return self
+
+    def __exit__(self, *args) -> None:
+        self.close()
+
+    # ── Session ──
+
+    def start_session(self, agent_id: str = "sdk") -> Session:
+        self._current_session = self.engine.start_session(agent_id)
+        return self._current_session
+
+    def end_session(self, summary: Optional[str] = None) -> Optional[Session]:
+        if self._current_session:
+            session = self.engine.end_session(self._current_session.session_id, summary)
+            self._current_session = None
+            return session
+        return None
+
+    @property
+    def session_id(self) -> Optional[str]:
+        return self._current_session.session_id if self._current_session else None
+
+    # ── Store & Retrieve ──
+
+    def store(
+        self,
+        text: str,
+        source_type: str = "user_input",
+        confidence: float = 1.0,
+        author: str = "",
+    ) -> StoreResult:
+        source = Source(
+            type=SourceType(source_type),
+            confidence=confidence,
+            author=author,
+            channel="sdk",
+        )
+        return self.engine.store(text, source=source, session_id=self.session_id)
+
+    def search(
+        self,
+        query: str,
+        max_results: int = 10,
+        max_tokens: int = 500,
+        min_confidence: float = 0.0,
+    ) -> SearchResult:
+        options = SearchOptions(
+            query=query,
+            max_results=max_results,
+            max_tokens=max_tokens,
+            min_confidence=min_confidence,
+        )
+        return self.engine.search(query, options)
+
+    def get_entity(self, name: str) -> Optional[Entity]:
+        return self.engine.get_entity(name)
+
+    def get_facts(self, entity_name: str) -> list:
+        return self.engine.get_facts(entity_name)
+
+    def add_fact(
+        self,
+        entity_name: str,
+        fact_text: str,
+        predicate: str = "attribute",
+        source_type: str = "user_input",
+        confidence: float = 1.0,
+    ) -> ValidationResult:
+        source = Source(type=SourceType(source_type), confidence=confidence, channel="sdk")
+        return self.engine.add_fact(
+            entity_name, fact_text, predicate=predicate,
+            source=source, session_id=self.session_id,
+        )
+
+    def create_entity(
+        self,
+        name: str,
+        entity_type: str = "person",
+        tier: str = "recall",
+        summary: str = "",
+    ) -> Entity:
+        return self.engine.create_entity(name, entity_type, Tier(tier), summary)
+
+    def list_entities(self, entity_type: str = "", tier: str = "", limit: int = 100) -> list[Entity]:
+        return self.engine.list_entities(
+            tier=Tier(tier) if tier else None,
+            entity_type=entity_type or None,
+            limit=limit,
+        )
+
+    # ── Context ──
+
+    def get_context(self, agent_id: str = "sdk") -> dict:
+        return self.engine.get_session_context(agent_id)
+
+    def health(self) -> dict:
+        return self.engine.health_check()
+
+    # ── Tool Schemas (for OpenAI function calling) ──
+
+    @staticmethod
+    def get_tool_schemas() -> list[dict]:
+        """Return OpenAI-format function calling schemas for all tools."""
+        return [
+            {
+                "type": "function",
+                "function": {
+                    "name": "memory_search",
+                    "description": "Search persistent memory for relevant context",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "query": {"type": "string", "description": "Search query"},
+                            "max_results": {"type": "integer", "default": 5},
+                            "max_tokens": {"type": "integer", "default": 500},
+                        },
+                        "required": ["query"],
+                    },
+                },
+            },
+            {
+                "type": "function",
+                "function": {
+                    "name": "memory_store",
+                    "description": "Store new factual information in memory",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "content": {"type": "string", "description": "Content to store"},
+                            "source_type": {"type": "string", "default": "agent_inference"},
+                            "confidence": {"type": "number", "default": 0.7},
+                        },
+                        "required": ["content"],
+                    },
+                },
+            },
+            {
+                "type": "function",
+                "function": {
+                    "name": "memory_get_entity",
+                    "description": "Get all info about a specific entity",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "name": {"type": "string", "description": "Entity name"},
+                        },
+                        "required": ["name"],
+                    },
+                },
+            },
+        ]

engram/models/__init__.py

@@ -0,0 +1,19 @@
+"""Engram data models — pure dataclasses with no business logic."""
+
+from engram.models.entity import Entity, EntityState, Tier
+from engram.models.fact import Fact, FactStatus
+from engram.models.source import Source, SourceType
+from engram.models.relation import Relation
+from engram.models.session import Session, SessionEvent
+from engram.models.search import SearchResult, SearchHit, SearchOptions
+from engram.models.quality import ValidationResult, ConflictInfo, PIIMatch, Action
+
+__all__ = [
+    "Entity", "EntityState", "Tier",
+    "Fact", "FactStatus",
+    "Source", "SourceType",
+    "Relation",
+    "Session", "SessionEvent",
+    "SearchResult", "SearchHit", "SearchOptions",
+    "ValidationResult", "ConflictInfo", "PIIMatch", "Action",
+]

engram/models/entity.py

@@ -0,0 +1,72 @@
+"""Entity data models."""
+
+from __future__ import annotations
+
+import uuid
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum
+from typing import Optional
+
+
+class Tier(str, Enum):
+    CORE = "core"
+    RECALL = "recall"
+    ARCHIVAL = "archival"
+
+
+@dataclass
+class EntityState:
+    """Structured key-value attributes of an entity."""
+
+    role: Optional[str] = None
+    affiliation: Optional[str] = None
+    location: Optional[str] = None
+    email: Optional[str] = None
+    custom: dict[str, str] = field(default_factory=dict)
+
+    def to_dict(self) -> dict:
+        result: dict = {}
+        if self.role is not None:
+            result["role"] = self.role
+        if self.affiliation is not None:
+            result["affiliation"] = self.affiliation
+        if self.location is not None:
+            result["location"] = self.location
+        if self.email is not None:
+            result["email"] = self.email
+        if self.custom:
+            result["custom"] = self.custom
+        return result
+
+    @classmethod
+    def from_dict(cls, data: dict) -> EntityState:
+        return cls(
+            role=data.get("role"),
+            affiliation=data.get("affiliation"),
+            location=data.get("location"),
+            email=data.get("email"),
+            custom=data.get("custom", {}),
+        )
+
+
+@dataclass
+class Entity:
+    """A tracked entity: person, organization, project, or concept."""
+
+    id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    name: str = ""
+    entity_type: str = "person"
+    state: EntityState = field(default_factory=EntityState)
+    tier: Tier = Tier.RECALL
+    summary: str = ""
+    aliases: list[str] = field(default_factory=list)
+    created_at: datetime = field(default_factory=datetime.now)
+    updated_at: datetime = field(default_factory=datetime.now)
+    access_count: int = 0
+    last_accessed: Optional[datetime] = None
+
+    def touch(self) -> None:
+        """Record an access."""
+        self.access_count += 1
+        self.last_accessed = datetime.now()