| import asyncio |
| import json |
| import logging |
| from datetime import datetime |
| from typing import Any, Dict, List, Optional, Sequence |
| from pathlib import Path |
| import uuid |
|
|
| import chromadb |
| from chromadb.config import Settings |
| import numpy as np |
| from sentence_transformers import SentenceTransformer |
|
|
| |
| try: |
| from mcp.server import Server |
| from mcp.server.models import InitializationOptions |
| from mcp.server.stdio import stdio_server |
| from mcp.types import ( |
| Resource, |
| Tool, |
| TextContent, |
| ImageContent, |
| EmbeddedResource, |
| LoggingLevel |
| ) |
| MCP_AVAILABLE = True |
| except ImportError: |
| print("MCP library not available, running in demo mode only") |
| MCP_AVAILABLE = False |
| |
| |
| class Server: |
| def __init__(self, name): |
| self.name = name |
| def list_tools(self): return lambda: None |
| def call_tool(self): return lambda: None |
| |
| class TextContent: |
| def __init__(self, type, text): |
| self.type = type |
| self.text = text |
|
|
| |
| logging.basicConfig(level=logging.INFO) |
| logger = logging.getLogger("ltm-mcp-server") |
|
|
| class LongTermMemoryServer: |
| def __init__(self): |
| self.server = Server("long-term-memory") |
| self.db_path = "./memory_db" |
| Path(self.db_path).mkdir(exist_ok=True) |
| |
| |
| self.chroma_client = chromadb.PersistentClient( |
| path=self.db_path, |
| settings=Settings(anonymized_telemetry=False) |
| ) |
| |
| |
| try: |
| self.collection = self.chroma_client.get_collection("memories") |
| except: |
| self.collection = self.chroma_client.create_collection( |
| name="memories", |
| metadata={"description": "Long-term memory storage for conversations"} |
| ) |
| |
| |
| self.encoder = SentenceTransformer('all-MiniLM-L6-v2') |
| |
| self.setup_handlers() |
| |
| def setup_handlers(self): |
| @self.server.list_tools() |
| async def handle_list_tools() -> List[Tool]: |
| """List available tools.""" |
| return [ |
| Tool( |
| name="save_memory", |
| description="Save important insights, conclusions, or context from conversation to long-term memory", |
| inputSchema={ |
| "type": "object", |
| "properties": { |
| "content": { |
| "type": "string", |
| "description": "The content/insight to save to memory" |
| }, |
| "title": { |
| "type": "string", |
| "description": "A short title/summary for this memory" |
| }, |
| "tags": { |
| "type": "array", |
| "items": {"type": "string"}, |
| "description": "Optional tags to categorize this memory", |
| "default": [] |
| }, |
| "context": { |
| "type": "string", |
| "description": "Additional context about when/why this was saved", |
| "default": "" |
| } |
| }, |
| "required": ["content", "title"] |
| } |
| ), |
| Tool( |
| name="search_memory", |
| description="Search through long-term memory for relevant information", |
| inputSchema={ |
| "type": "object", |
| "properties": { |
| "query": { |
| "type": "string", |
| "description": "Search query to find relevant memories" |
| }, |
| "limit": { |
| "type": "integer", |
| "description": "Maximum number of results to return", |
| "default": 5 |
| }, |
| "threshold": { |
| "type": "number", |
| "description": "Similarity threshold (0-1, higher = more similar)", |
| "default": 0.3 |
| } |
| }, |
| "required": ["query"] |
| } |
| ), |
| Tool( |
| name="list_memories", |
| description="List all stored memories with basic info", |
| inputSchema={ |
| "type": "object", |
| "properties": { |
| "limit": { |
| "type": "integer", |
| "description": "Maximum number of memories to return", |
| "default": 10 |
| } |
| } |
| } |
| ), |
| Tool( |
| name="delete_memory", |
| description="Delete a specific memory by ID", |
| inputSchema={ |
| "type": "object", |
| "properties": { |
| "memory_id": { |
| "type": "string", |
| "description": "The ID of the memory to delete" |
| } |
| }, |
| "required": ["memory_id"] |
| } |
| ) |
| ] |
|
|
| @self.server.call_tool() |
| async def handle_call_tool(name: str, arguments: Dict[str, Any]) -> Sequence[TextContent]: |
| """Handle tool calls.""" |
| try: |
| if name == "save_memory": |
| return await self._save_memory(**arguments) |
| elif name == "search_memory": |
| return await self._search_memory(**arguments) |
| elif name == "list_memories": |
| return await self._list_memories(**arguments) |
| elif name == "delete_memory": |
| return await self._delete_memory(**arguments) |
| else: |
| raise ValueError(f"Unknown tool: {name}") |
| except Exception as e: |
| logger.error(f"Error in tool {name}: {e}") |
| return [TextContent(type="text", text=f"Error: {str(e)}")] |
|
|
| async def _save_memory(self, content: str, title: str, tags: List[str] = None, context: str = "") -> Sequence[TextContent]: |
| """Save content to long-term memory.""" |
| if tags is None: |
| tags = [] |
| |
| memory_id = str(uuid.uuid4()) |
| timestamp = datetime.now().isoformat() |
| |
| |
| embedding = self.encoder.encode(f"{title} {content}").tolist() |
| |
| |
| metadata = { |
| "title": title, |
| "timestamp": timestamp, |
| "tags": json.dumps(tags), |
| "context": context, |
| "content_length": len(content) |
| } |
| |
| |
| self.collection.add( |
| documents=[content], |
| embeddings=[embedding], |
| metadatas=[metadata], |
| ids=[memory_id] |
| ) |
| |
| logger.info(f"Saved memory: {title} (ID: {memory_id})") |
| |
| result = f"✅ Memory saved successfully!\n\n" |
| result += f"**ID**: {memory_id}\n" |
| result += f"**Title**: {title}\n" |
| result += f"**Timestamp**: {timestamp}\n" |
| if tags: |
| result += f"**Tags**: {', '.join(tags)}\n" |
| if context: |
| result += f"**Context**: {context}\n" |
| result += f"**Content Preview**: {content[:200]}{'...' if len(content) > 200 else ''}" |
| |
| return [TextContent(type="text", text=result)] |
|
|
| async def _search_memory(self, query: str, limit: int = 5, threshold: float = 0.3) -> Sequence[TextContent]: |
| """Search through memories.""" |
| if self.collection.count() == 0: |
| return [TextContent(type="text", text="No memories stored yet.")] |
| |
| |
| query_embedding = self.encoder.encode(query).tolist() |
| |
| |
| results = self.collection.query( |
| query_embeddings=[query_embedding], |
| n_results=min(limit, self.collection.count()) |
| ) |
| |
| if not results['documents'][0]: |
| return [TextContent(type="text", text="No relevant memories found.")] |
| |
| |
| response = f"🔍 **Search Results for**: \"{query}\"\n\n" |
| |
| found_relevant = False |
| for i, (doc, metadata, distance) in enumerate(zip( |
| results['documents'][0], |
| results['metadatas'][0], |
| results['distances'][0] |
| )): |
| similarity = 1 - distance |
| if similarity >= threshold: |
| found_relevant = True |
| tags = json.loads(metadata.get('tags', '[]')) |
| |
| response += f"**{i+1}. {metadata['title']}** (Similarity: {similarity:.2f})\n" |
| response += f"*Saved*: {metadata['timestamp']}\n" |
| if tags: |
| response += f"*Tags*: {', '.join(tags)}\n" |
| if metadata.get('context'): |
| response += f"*Context*: {metadata['context']}\n" |
| response += f"*Content*: {doc}\n\n" |
| response += "---\n\n" |
| |
| if not found_relevant: |
| response += f"No memories found above similarity threshold of {threshold:.2f}" |
| |
| return [TextContent(type="text", text=response)] |
|
|
| async def _list_memories(self, limit: int = 10) -> Sequence[TextContent]: |
| """List all memories.""" |
| if self.collection.count() == 0: |
| return [TextContent(type="text", text="No memories stored yet.")] |
| |
| |
| results = self.collection.get() |
| |
| if not results['documents']: |
| return [TextContent(type="text", text="No memories found.")] |
| |
| response = f"📚 **All Memories** (showing up to {limit})\n\n" |
| |
| |
| memories = list(zip(results['ids'], results['documents'], results['metadatas'])) |
| memories.sort(key=lambda x: x[2]['timestamp'], reverse=True) |
| |
| for i, (memory_id, doc, metadata) in enumerate(memories[:limit]): |
| tags = json.loads(metadata.get('tags', '[]')) |
| |
| response += f"**{i+1}. {metadata['title']}**\n" |
| response += f"*ID*: {memory_id}\n" |
| response += f"*Saved*: {metadata['timestamp']}\n" |
| if tags: |
| response += f"*Tags*: {', '.join(tags)}\n" |
| response += f"*Preview*: {doc[:150]}{'...' if len(doc) > 150 else ''}\n\n" |
| response += "---\n\n" |
| |
| if len(memories) > limit: |
| response += f"... and {len(memories) - limit} more memories" |
| |
| return [TextContent(type="text", text=response)] |
|
|
| async def _delete_memory(self, memory_id: str) -> Sequence[TextContent]: |
| """Delete a memory by ID.""" |
| try: |
| |
| result = self.collection.get(ids=[memory_id]) |
| if not result['documents']: |
| return [TextContent(type="text", text=f"❌ Memory with ID {memory_id} not found.")] |
| |
| |
| metadata = result['metadatas'][0] |
| title = metadata['title'] |
| |
| |
| self.collection.delete(ids=[memory_id]) |
| |
| logger.info(f"Deleted memory: {title} (ID: {memory_id})") |
| |
| return [TextContent(type="text", text=f"✅ Memory deleted successfully!\n\n**Title**: {title}\n**ID**: {memory_id}")] |
| |
| except Exception as e: |
| logger.error(f"Error deleting memory {memory_id}: {e}") |
| return [TextContent(type="text", text=f"❌ Error deleting memory: {str(e)}")] |
|
|
| async def run(self): |
| """Run the MCP server.""" |
| if not MCP_AVAILABLE: |
| print("MCP not available, cannot run server") |
| return |
| |
| async with stdio_server() as (read_stream, write_stream): |
| await self.server.run( |
| read_stream, |
| write_stream, |
| InitializationOptions( |
| server_name="long-term-memory", |
| server_version="1.0.0", |
| capabilities=self.server.get_capabilities( |
| notification_options=None, |
| experimental_capabilities=None, |
| ), |
| ), |
| ) |
|
|
| |
| async def main(): |
| if not MCP_AVAILABLE: |
| print("MCP server cannot run without MCP library") |
| return |
| |
| server = LongTermMemoryServer() |
| await server.run() |
|
|
| if __name__ == "__main__": |
| asyncio.run(main()) |
