monitor_app/mcp/ai_memory.py

0001 """
0002 AI Memory MCP tools for cross-session dialogue persistence.
0003
0004 Provides tools for recording and retrieving AI dialogue history,
0005 enabling context continuity across Claude Code sessions.
0006
0007 Opt-in via SWF_DIALOGUE_TURNS environment variable.
0008 """
0009
0010 import logging
0011 from django.utils import timezone
0012 from asgiref.sync import sync_to_async
0013
0014 from mcp_server import mcp_server as mcp
0015
0016 from ..models import AIMemory
0017
0018 logger = logging.getLogger(__name__)
0019
0020
0021 @mcp.tool()
0022 async def swf_record_ai_memory(
0023     username: str,
0024     session_id: str,
0025     role: str,
0026     content: str,
0027     namespace: str = None,
0028     project_path: str = None,
0029 ) -> dict:
0030     """
0031     Record a dialogue exchange for AI memory persistence.
0032
0033     Called by Claude Code hooks to store user prompts and assistant responses.
0034     Each exchange is stored as a separate record for retrieval across sessions.
0035
0036     Args:
0037         username: Developer username (required)
0038         session_id: Claude Code session ID (required)
0039         role: Either 'user' or 'assistant' (required)
0040         content: The message content (required)
0041         namespace: Testbed namespace if applicable
0042         project_path: Project directory path
0043
0044     Returns:
0045         Success/failure status with record ID
0046     """
0047     if role not in ('user', 'assistant'):
0048         return {"success": False, "error": f"Invalid role '{role}'. Must be 'user' or 'assistant'."}
0049
0050     if not username or not session_id or not content:
0051         return {"success": False, "error": "username, session_id, and content are required"}
0052
0053     @sync_to_async
0054     def do_record():
0055         try:
0056             record = AIMemory.objects.create(
0057                 username=username,
0058                 session_id=session_id,
0059                 role=role,
0060                 content=content,
0061                 namespace=namespace,
0062                 project_path=project_path,
0063             )
0064             logger.debug(
0065                 f"AI memory recorded: user={username} session={session_id[:8]}... "
0066                 f"role={role} len={len(content)}"
0067             )
0068             return {
0069                 "success": True,
0070                 "id": record.id,
0071                 "username": username,
0072                 "role": role,
0073                 "content_length": len(content),
0074             }
0075         except Exception as e:
0076             logger.error(f"Failed to record AI memory: {e}")
0077             return {"success": False, "error": str(e)}
0078
0079     return await do_record()
0080
0081
0082 @mcp.tool()
0083 async def swf_get_ai_memory(
0084     username: str,
0085     turns: int = 20,
0086     namespace: str = None,
0087 ) -> list:
0088     """
0089     Get recent dialogue history for session context.
0090
0091     Called at session start to load recent exchanges into the AI's context.
0092     Returns chronologically ordered messages (oldest first) for natural
0093     conversation flow.
0094
0095     Args:
0096         username: Developer username (required)
0097         turns: Number of conversation turns to retrieve (default: 20).
0098                Each turn = 1 user + 1 assistant message, so 20 turns = up to 40 messages.
0099         namespace: Filter to messages from this namespace (optional)
0100
0101     Returns:
0102         List of dialogue entries with: role, content, created_at, session_id
0103     """
0104     if not username:
0105         return {"error": "username is required"}
0106
0107     max_messages = turns * 2  # Each turn is user + assistant
0108
0109     @sync_to_async
0110     def fetch():
0111         qs = AIMemory.objects.filter(username=username)
0112
0113         if namespace:
0114             qs = qs.filter(namespace=namespace)
0115
0116         # Get most recent messages, then reverse for chronological order
0117         recent = qs.order_by('-created_at')[:max_messages]
0118         messages = list(reversed([
0119             {
0120                 "role": m.role,
0121                 "content": m.content,
0122                 "created_at": m.created_at.isoformat() if m.created_at else None,
0123                 "session_id": m.session_id,
0124                 "post_id": m.namespace or '',
0125                 "root_id": m.project_path or '',
0126             }
0127             for m in recent
0128         ]))
0129
0130         return {
0131             "items": messages,
0132             "count": len(messages),
0133             "username": username,
0134             "turns_requested": turns,
0135             "namespace": namespace,
0136         }
0137
0138     return await fetch()