Upload 17 files

memory/vector.py

@@ -51,6 +51,10 @@ class _SentenceTransformerEmbedder:
         embeddings = self.model.encode(input, show_progress_bar=False)
         return embeddings.tolist()
 
+    def name(self) -> str:
+        """Required by ChromaDB EmbeddingFunction protocol."""
+        return f"sentence-transformers_{self.model_name}"
+
 
 class SemanticMemory:
     """ChromaDB-backed vector store with Markdown file mirror."""

memory_server.py

@@ -0,0 +1,572 @@
+#!/usr/bin/env python3
+"""
+Memory System MCP Server
+=========================
+A three-tier memory architecture exposed as MCP tools for AI agents.
+
+Tiers
+-----
+1. **Session** (short-term)  – conversation context, auto-expiring
+2. **Episodic** (mid-term)   – past tasks & events, searchable timeline
+3. **Semantic** (long-term)  – vector-backed RAG knowledge base
+
+Every entry is also persisted as a human-readable Markdown file.
+
+Usage
+-----
+    python memory_server.py                # stdio transport (for MCP clients)
+    python memory_server.py --sse 8765     # SSE transport on port 8765
+
+Transport is auto-detected via MCP protocol when run from an MCP host.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import sys
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+from mcp.server.fastmcp import FastMCP
+
+# ── local imports ────────────────────────────────────────────
+from memory.session import SessionMemory
+from memory.events import EpisodicMemory
+from memory.vector import SemanticMemory
+from memory.models import MemoryEntry, MemoryTier
+
+# ── logging ──────────────────────────────────────────────────
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s  %(levelname)-8s  %(name)s  %(message)s",
+)
+logger = logging.getLogger("memory-mcp")
+
+# ── resolve data root ───────────────────────────────────────
+DATA_ROOT = Path(os.environ.get("MEMORY_DATA_ROOT", Path(__file__).parent / "data"))
+DATA_ROOT.mkdir(parents=True, exist_ok=True)
+
+EMBEDDING_MODEL = os.environ.get("MEMORY_EMBEDDING_MODEL", "all-MiniLM-L6-v2")
+SESSION_TTL = int(os.environ.get("MEMORY_SESSION_TTL", "3600"))
+
+# ── initialise stores ───────────────────────────────────────
+session_store = SessionMemory(
+    base_dir=str(DATA_ROOT / "session"),
+    ttl=SESSION_TTL,
+)
+episodic_store = EpisodicMemory(
+    base_dir=str(DATA_ROOT / "events"),
+)
+semantic_store = SemanticMemory(
+    vector_dir=str(DATA_ROOT / "vector"),
+    md_dir=str(DATA_ROOT / "vector" / "docs"),
+    model_name=EMBEDDING_MODEL,
+)
+
+logger.info("🧠 Memory stores initialised  –  data_root=%s", DATA_ROOT)
+
+# ── MCP server ───────────────────────────────────────────────
+mcp = FastMCP("memory")
+
+
+# =====================================================================
+#  RESOURCES  – browse memory state
+# =====================================================================
+
+@mcp.resource("memory://status")
+def memory_status() -> str:
+    """Overview of the memory system."""
+    return json.dumps({
+        "session": {
+            "sessions": session_store.list_sessions(),
+            "total_entries": sum(
+                len(session_store.list_entries(sid))
+                for sid in session_store.list_sessions()
+            ),
+        },
+        "episodic": {
+            "total_entries": episodic_store.count(),
+        },
+        "semantic": {
+            "total_entries": semantic_store.count(),
+            "embedding_model": EMBEDDING_MODEL,
+        },
+    }, indent=2)
+
+
+@mcp.resource("memory://session/{session_id}")
+def session_resource(session_id: str) -> str:
+    """List entries in a session."""
+    entries = session_store.list_entries(session_id)
+    return json.dumps([e.to_dict() for e in entries], indent=2)
+
+
+@mcp.resource("memory://events/recent")
+def recent_events_resource() -> str:
+    """The 20 most recent episodic events."""
+    entries = episodic_store.recent(20)
+    return json.dumps([e.to_dict() for e in entries], indent=2)
+
+
+# =====================================================================
+#  PROMPTS
+# =====================================================================
+
+@mcp.prompt()
+def memory_context_prompt(query: str = "", session_id: str = "default") -> str:
+    """Build a comprehensive memory context for an LLM prompt."""
+    parts: List[str] = ["# Agent Memory Context\n"]
+
+    # session context
+    session_entries = session_store.list_entries(session_id)
+    if session_entries:
+        parts.append("## Recent Conversation (Session)")
+        for e in session_entries[-5:]:
+            parts.append(f"- [{e.created_at}] {e.title}: {e.content[:200]}")
+        parts.append("")
+
+    # episodic
+    recent = episodic_store.recent(5)
+    if recent:
+        parts.append("## Recent Tasks (Episodic)")
+        for e in recent:
+            parts.append(f"- [{e.created_at}] {e.title}: {e.content[:200]}")
+        parts.append("")
+
+    # semantic / RAG
+    if query:
+        hits = semantic_store.search(query, limit=3)
+        if hits:
+            parts.append("## Relevant Knowledge (Semantic / RAG)")
+            for h in hits:
+                parts.append(f"- [score={h.score:.2f}] {h.entry.title}: {h.entry.content[:300]}")
+            parts.append("")
+
+    return "\n".join(parts)
+
+
+# =====================================================================
+#  TOOLS  – full CRUD for each tier
+# =====================================================================
+
+# ─── Session (short-term) ───────────────────────────────────
+
+@mcp.tool()
+def session_create(
+    content: str,
+    title: str = "",
+    tags: str = "",
+    session_id: str = "default",
+    importance: float = 0.5,
+) -> Dict[str, Any]:
+    """
+    Create a new short-term / session memory entry.
+
+    Stores conversation context that auto-expires after the configured TTL.
+    Persisted as a Markdown file under data/session/<session_id>/.
+    """
+    entry = MemoryEntry(
+        content=content,
+        title=title or content[:60],
+        tags=[t.strip() for t in tags.split(",") if t.strip()] if tags else [],
+        importance=importance,
+    )
+    result = session_store.create(entry, session_id=session_id)
+    return {"status": "created", "entry": result.to_dict()}
+
+
+@mcp.tool()
+def session_read(entry_id: str, session_id: str = "default") -> Dict[str, Any]:
+    """Read a single session memory entry by ID."""
+    entry = session_store.read(entry_id, session_id)
+    if not entry:
+        return {"status": "not_found", "entry_id": entry_id}
+    return {"status": "ok", "entry": entry.to_dict()}
+
+
+@mcp.tool()
+def session_update(
+    entry_id: str,
+    session_id: str = "default",
+    content: str = "",
+    title: str = "",
+    tags: str = "",
+    importance: float = -1,
+) -> Dict[str, Any]:
+    """Update a session memory entry. Only provided fields are changed."""
+    kwargs: Dict[str, Any] = {}
+    if content:
+        kwargs["content"] = content
+    if title:
+        kwargs["title"] = title
+    if tags:
+        kwargs["tags"] = [t.strip() for t in tags.split(",") if t.strip()]
+    if importance >= 0:
+        kwargs["importance"] = importance
+    entry = session_store.update(entry_id, session_id, **kwargs)
+    if not entry:
+        return {"status": "not_found", "entry_id": entry_id}
+    return {"status": "updated", "entry": entry.to_dict()}
+
+
+@mcp.tool()
+def session_delete(entry_id: str, session_id: str = "default") -> Dict[str, Any]:
+    """Delete a session memory entry."""
+    ok = session_store.delete(entry_id, session_id)
+    return {"status": "deleted" if ok else "not_found", "entry_id": entry_id}
+
+
+@mcp.tool()
+def session_list(session_id: str = "default", tag: str = "") -> Dict[str, Any]:
+    """List all entries in a session, optionally filtered by tag."""
+    entries = session_store.list_entries(session_id, tag=tag or None)
+    return {"count": len(entries), "entries": [e.to_dict() for e in entries]}
+
+
+@mcp.tool()
+def session_search(query: str, session_id: str = "", limit: int = 10) -> Dict[str, Any]:
+    """Keyword search across session memories."""
+    results = session_store.search(query, session_id=session_id or None, limit=limit)
+    return {"count": len(results), "entries": [e.to_dict() for e in results]}
+
+
+@mcp.tool()
+def session_clear(session_id: str = "default") -> Dict[str, Any]:
+    """Clear all entries from a session."""
+    count = session_store.clear_session(session_id)
+    return {"status": "cleared", "session_id": session_id, "deleted": count}
+
+
+@mcp.tool()
+def session_gc() -> Dict[str, Any]:
+    """Garbage-collect expired session entries across all sessions."""
+    removed = session_store.gc()
+    return {"status": "gc_complete", "removed": removed}
+
+
+# ─── Episodic (mid-term) ────────────────────────────────────
+
+@mcp.tool()
+def episodic_create(
+    content: str,
+    title: str = "",
+    tags: str = "",
+    importance: float = 0.5,
+    source: str = "",
+) -> Dict[str, Any]:
+    """
+    Record a new episodic memory (task completion, event, interaction).
+
+    Stored as a timestamped Markdown file under data/events/.
+    """
+    tag_list = [t.strip() for t in tags.split(",") if t.strip()] if tags else []
+    entry = episodic_store.create(
+        content=content,
+        title=title,
+        tags=tag_list,
+        importance=importance,
+        source=source,
+    )
+    return {"status": "created", "entry": entry.to_dict()}
+
+
+@mcp.tool()
+def episodic_read(entry_id: str) -> Dict[str, Any]:
+    """Read a single episodic memory by ID."""
+    entry = episodic_store.read(entry_id)
+    if not entry:
+        return {"status": "not_found", "entry_id": entry_id}
+    return {"status": "ok", "entry": entry.to_dict()}
+
+
+@mcp.tool()
+def episodic_update(
+    entry_id: str,
+    content: str = "",
+    title: str = "",
+    tags: str = "",
+    importance: float = -1,
+) -> Dict[str, Any]:
+    """Update an episodic memory entry."""
+    kwargs: Dict[str, Any] = {}
+    if content:
+        kwargs["content"] = content
+    if title:
+        kwargs["title"] = title
+    if tags:
+        kwargs["tags"] = [t.strip() for t in tags.split(",") if t.strip()]
+    if importance >= 0:
+        kwargs["importance"] = importance
+    entry = episodic_store.update(entry_id, **kwargs)
+    if not entry:
+        return {"status": "not_found", "entry_id": entry_id}
+    return {"status": "updated", "entry": entry.to_dict()}
+
+
+@mcp.tool()
+def episodic_delete(entry_id: str) -> Dict[str, Any]:
+    """Delete an episodic memory entry."""
+    ok = episodic_store.delete(entry_id)
+    return {"status": "deleted" if ok else "not_found", "entry_id": entry_id}
+
+
+@mcp.tool()
+def episodic_list(
+    tag: str = "",
+    since: str = "",
+    until: str = "",
+    limit: int = 50,
+) -> Dict[str, Any]:
+    """List episodic memories, optionally filtered by tag and/or time range (ISO format)."""
+    entries = episodic_store.list_entries(
+        tag=tag or None,
+        since=since or None,
+        until=until or None,
+        limit=limit,
+    )
+    return {"count": len(entries), "entries": [e.to_dict() for e in entries]}
+
+
+@mcp.tool()
+def episodic_search(query: str, limit: int = 10) -> Dict[str, Any]:
+    """Keyword search across episodic memories."""
+    results = episodic_store.search(query, limit=limit)
+    return {"count": len(results), "entries": [e.to_dict() for e in results]}
+
+
+@mcp.tool()
+def episodic_recent(n: int = 10) -> Dict[str, Any]:
+    """Get the N most recent episodic events."""
+    entries = episodic_store.recent(n)
+    return {"count": len(entries), "entries": [e.to_dict() for e in entries]}
+
+
+# ─── Semantic / RAG (long-term) ─────────────────────────────
+
+@mcp.tool()
+def semantic_create(
+    content: str,
+    title: str = "",
+    tags: str = "",
+    importance: float = 0.5,
+    source: str = "",
+) -> Dict[str, Any]:
+    """
+    Add a document to the semantic / RAG knowledge base.
+
+    The content is embedded via sentence-transformers and stored in ChromaDB
+    for similarity search. Also persisted as a Markdown file.
+    """
+    tag_list = [t.strip() for t in tags.split(",") if t.strip()] if tags else []
+    entry = semantic_store.create(
+        content=content,
+        title=title,
+        tags=tag_list,
+        importance=importance,
+        source=source,
+    )
+    return {"status": "created", "entry": entry.to_dict()}
+
+
+@mcp.tool()
+def semantic_read(entry_id: str) -> Dict[str, Any]:
+    """Read a single semantic memory by ID."""
+    entry = semantic_store.read(entry_id)
+    if not entry:
+        return {"status": "not_found", "entry_id": entry_id}
+    return {"status": "ok", "entry": entry.to_dict()}
+
+
+@mcp.tool()
+def semantic_update(
+    entry_id: str,
+    content: str = "",
+    title: str = "",
+    tags: str = "",
+    importance: float = -1,
+) -> Dict[str, Any]:
+    """Update a semantic memory entry. Re-embeds automatically if content changes."""
+    kwargs: Dict[str, Any] = {}
+    if content:
+        kwargs["content"] = content
+    if title:
+        kwargs["title"] = title
+    if tags:
+        kwargs["tags"] = [t.strip() for t in tags.split(",") if t.strip()]
+    if importance >= 0:
+        kwargs["importance"] = importance
+    entry = semantic_store.update(entry_id, **kwargs)
+    if not entry:
+        return {"status": "not_found", "entry_id": entry_id}
+    return {"status": "updated", "entry": entry.to_dict()}
+
+
+@mcp.tool()
+def semantic_delete(entry_id: str) -> Dict[str, Any]:
+    """Delete a semantic memory entry from vector store and disk."""
+    ok = semantic_store.delete(entry_id)
+    return {"status": "deleted" if ok else "not_found", "entry_id": entry_id}
+
+
+@mcp.tool()
+def semantic_search(query: str, limit: int = 5) -> Dict[str, Any]:
+    """
+    Semantic similarity search (RAG retrieval).
+
+    Finds the most relevant documents in the knowledge base using
+    vector cosine similarity. This is the primary RAG endpoint.
+    """
+    results = semantic_store.search(query, limit=limit)
+    return {
+        "count": len(results),
+        "results": [
+            {
+                "score": round(r.score, 4),
+                "distance": round(r.distance, 4),
+                "entry": r.entry.to_dict(),
+            }
+            for r in results
+        ],
+    }
+
+
+@mcp.tool()
+def semantic_list(limit: int = 100, tag: str = "") -> Dict[str, Any]:
+    """List all entries in the semantic knowledge base."""
+    entries = semantic_store.list_entries(limit=limit, tag=tag or None)
+    return {"count": len(entries), "entries": [e.to_dict() for e in entries]}
+
+
+# ─── Cross-tier utilities ───────────────────────────────────
+
+@mcp.tool()
+def memory_search_all(query: str, limit: int = 5) -> Dict[str, Any]:
+    """
+    Search across ALL memory tiers (session + episodic + semantic).
+
+    Combines keyword search from session & episodic with
+    semantic vector search. Returns unified results sorted by relevance.
+    """
+    results: Dict[str, Any] = {}
+
+    # session
+    s_hits = session_store.search(query, limit=limit)
+    results["session"] = [e.to_dict() for e in s_hits]
+
+    # episodic
+    e_hits = episodic_store.search(query, limit=limit)
+    results["episodic"] = [e.to_dict() for e in e_hits]
+
+    # semantic (RAG)
+    v_hits = semantic_store.search(query, limit=limit)
+    results["semantic"] = [
+        {"score": round(r.score, 4), "entry": r.entry.to_dict()}
+        for r in v_hits
+    ]
+
+    results["total"] = len(s_hits) + len(e_hits) + len(v_hits)
+    return results
+
+
+@mcp.tool()
+def memory_promote(entry_id: str, from_tier: str, to_tier: str) -> Dict[str, Any]:
+    """
+    Promote a memory entry from one tier to another.
+
+    E.g. promote a session memory to episodic, or episodic to semantic.
+    The entry is copied to the target tier (source is kept).
+    """
+    # read from source
+    source_entry: Optional[MemoryEntry] = None
+    if from_tier == "session":
+        source_entry = session_store.read(entry_id)
+    elif from_tier == "episodic":
+        source_entry = episodic_store.read(entry_id)
+    elif from_tier == "semantic":
+        source_entry = semantic_store.read(entry_id)
+
+    if not source_entry:
+        return {"status": "not_found", "entry_id": entry_id, "tier": from_tier}
+
+    # write to target
+    if to_tier == "session":
+        new_entry = MemoryEntry(
+            content=source_entry.content,
+            title=source_entry.title,
+            tags=source_entry.tags,
+            importance=source_entry.importance,
+            metadata=source_entry.metadata,
+            source=f"promoted from {from_tier}:{entry_id}",
+        )
+        result = session_store.create(new_entry)
+    elif to_tier == "episodic":
+        result = episodic_store.create(
+            content=source_entry.content,
+            title=source_entry.title,
+            tags=source_entry.tags,
+            importance=source_entry.importance,
+            metadata=source_entry.metadata,
+            source=f"promoted from {from_tier}:{entry_id}",
+        )
+    elif to_tier == "semantic":
+        result = semantic_store.create(
+            content=source_entry.content,
+            title=source_entry.title,
+            tags=source_entry.tags,
+            importance=source_entry.importance,
+            metadata=source_entry.metadata,
+            source=f"promoted from {from_tier}:{entry_id}",
+        )
+    else:
+        return {"status": "error", "message": f"Unknown target tier: {to_tier}"}
+
+    return {
+        "status": "promoted",
+        "from": from_tier,
+        "to": to_tier,
+        "original_id": entry_id,
+        "new_entry": result.to_dict(),
+    }
+
+
+@mcp.tool()
+def memory_stats() -> Dict[str, Any]:
+    """Get statistics about all memory tiers."""
+    sessions = session_store.list_sessions()
+    session_total = sum(len(session_store.list_entries(sid)) for sid in sessions)
+    return {
+        "session": {
+            "sessions": len(sessions),
+            "total_entries": session_total,
+            "ttl_seconds": SESSION_TTL,
+        },
+        "episodic": {
+            "total_entries": episodic_store.count(),
+        },
+        "semantic": {
+            "total_entries": semantic_store.count(),
+            "embedding_model": EMBEDDING_MODEL,
+        },
+        "data_root": str(DATA_ROOT),
+    }
+
+
+# =====================================================================
+#  ENTRY POINT
+# =====================================================================
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Memory System MCP Server")
+    parser.add_argument("--sse", type=int, default=0, help="Run SSE transport on this port")
+    args = parser.parse_args()
+
+    if args.sse:
+        logger.info("🚀 Starting Memory MCP server (SSE) on port %d", args.sse)
+        mcp.run(transport="sse", sse_params={"port": args.sse})
+    else:
+        logger.info("🚀 Starting Memory MCP server (stdio)")
+        mcp.run(transport="stdio")