memos_chat_client.py

"""
meMOS Chat API Client
Simplified client for interacting with meMOS chat endpoint
"""

import asyncio
import json
import logging
from typing import Dict, List, Optional, Any
from datetime import datetime
import aiohttp

logger = logging.getLogger(__name__)

class MemOSChatClient:
    """Chat-focused client for meMOS memory management system"""
    
    def __init__(self, base_url: str, user_id: str = "echo_mind_user"):
        """
        Initialize meMOS chat client
        
        Args:
            base_url: Base URL for meMOS REST API
            user_id: User ID for memory isolation
        """
        self.base_url = base_url.rstrip('/')
        self.user_id = user_id
        self.session = None
        
    async def __aenter__(self):
        """Async context manager entry"""
        await self.connect()
        return self
        
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Async context manager exit"""
        await self.close()
        
    async def connect(self):
        """Create aiohttp session"""
        self.session = aiohttp.ClientSession(
            headers={"Content-Type": "application/json"},
            timeout=aiohttp.ClientTimeout(total=30)
        )
        
    async def close(self):
        """Close aiohttp session"""
        if self.session:
            await self.session.close()
            self.session = None
            
    async def chat(self, query: str) -> Dict[str, Any]:
        """
        Send a chat query to meMOS
        
        Args:
            query: The user's message/query
            
        Returns:
            Response from meMOS including the AI response
        """
        if not self.session:
            await self.connect()
            
        url = f"{self.base_url}/chat"
        data = {
            "query": query,
            "user_id": self.user_id
        }
        
        try:
            async with self.session.post(url, json=data) as response:
                result = await response.json()
                
                if response.status == 200 and result.get("code") == 200:
                    return {
                        "success": True,
                        "response": result.get("data", ""),
                        "timestamp": datetime.now().isoformat()
                    }
                else:
                    logger.error(f"meMOS chat error: {result}")
                    return {
                        "success": False,
                        "error": result.get("message", "Unknown error"),
                        "timestamp": datetime.now().isoformat()
                    }
                    
        except aiohttp.ClientError as e:
            logger.error(f"meMOS API request failed: {e}")
            return {
                "success": False,
                "error": str(e),
                "timestamp": datetime.now().isoformat()
            }
            
    async def search_memories(self, query: str, limit: int = 5) -> List[Dict[str, Any]]:
        """
        Search memories through chat interface
        
        Args:
            query: Search query
            limit: Maximum results to return
            
        Returns:
            List of relevant memories
        """
        # Use chat endpoint to search memories by asking directly
        search_query = f"Based on our previous conversations, tell me what you remember about: {query}"
        result = await self.chat(search_query)
        
        if result.get("success"):
            # Parse the response to extract memory-like information
            return [{
                "content": result.get("response", ""),
                "timestamp": result.get("timestamp"),
                "relevance": 1.0
            }]
        
        return []
        
    async def store_memory(self, user_input: str, assistant_response: str) -> bool:
        """
        Store a conversation in memory via chat
        
        Args:
            user_input: The user's message
            assistant_response: The assistant's response
            
        Returns:
            Success status
        """
        # The chat endpoint automatically stores conversations
        # We just need to send the user input and get a response
        result = await self.chat(user_input)
        return result.get("success", False)