memos_rest_client.py

"""
memOS REST API Client Service
Provides a clean interface for interacting with memOS REST API endpoints
"""

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

logger = logging.getLogger(__name__)

class MemoryEntity(BaseModel):
    """Represents a memory entity in memOS"""
    id: Optional[str] = None
    name: str
    entity_type: str = Field(alias="entityType")
    observations: List[str] = []
    created_at: Optional[datetime] = Field(None, alias="createdAt")
    updated_at: Optional[datetime] = Field(None, alias="updatedAt")

class MemoryRelation(BaseModel):
    """Represents a relation between memory entities"""
    from_entity: str = Field(alias="from")
    to_entity: str = Field(alias="to")
    relation_type: str = Field(alias="relationType")
    metadata: Dict[str, Any] = {}

class MemOSRestClient:
    """REST API client for memOS memory management system"""
    
    def __init__(self, base_url: str, api_key: Optional[str] = None):
        """
        Initialize memOS REST client
        
        Args:
            base_url: Base URL for memOS REST API
            api_key: Optional API key for authentication
        """
        self.base_url = base_url.rstrip('/')
        self.api_key = api_key
        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"""
        headers = {"Content-Type": "application/json"}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"
            
        self.session = aiohttp.ClientSession(
            base_url=self.base_url,
            headers=headers,
            timeout=aiohttp.ClientTimeout(total=30)
        )
        
    async def close(self):
        """Close aiohttp session"""
        if self.session:
            await self.session.close()
            self.session = None
            
    async def _make_request(self, method: str, endpoint: str, **kwargs) -> Dict[str, Any]:
        """Make HTTP request to memOS API"""
        if not self.session:
            await self.connect()
            
        url = f"{self.base_url}{endpoint}"
        
        try:
            async with self.session.request(method, endpoint, **kwargs) as response:
                response.raise_for_status()
                return await response.json()
        except aiohttp.ClientError as e:
            logger.error(f"memOS API request failed: {e}")
            raise
            
    async def create_entity(self, entity: MemoryEntity) -> MemoryEntity:
        """Create a new memory entity"""
        data = entity.model_dump(by_alias=True, exclude_none=True)
        response = await self._make_request("POST", "/api/v1/entities", json=data)
        return MemoryEntity(**response)
        
    async def get_entity(self, entity_id: str) -> MemoryEntity:
        """Get a memory entity by ID"""
        response = await self._make_request("GET", f"/api/v1/entities/{entity_id}")
        return MemoryEntity(**response)
        
    async def update_entity(self, entity_id: str, entity: MemoryEntity) -> MemoryEntity:
        """Update an existing memory entity"""
        data = entity.model_dump(by_alias=True, exclude_none=True)
        response = await self._make_request("PUT", f"/api/v1/entities/{entity_id}", json=data)
        return MemoryEntity(**response)
        
    async def delete_entity(self, entity_id: str) -> bool:
        """Delete a memory entity"""
        await self._make_request("DELETE", f"/api/v1/entities/{entity_id}")
        return True
        
    async def search_entities(self, query: str, entity_type: Optional[str] = None) -> List[MemoryEntity]:
        """Search for memory entities"""
        params = {"q": query}
        if entity_type:
            params["type"] = entity_type
            
        response = await self._make_request("GET", "/api/v1/entities/search", params=params)
        return [MemoryEntity(**entity) for entity in response.get("entities", [])]
        
    async def add_observations(self, entity_id: str, observations: List[str]) -> MemoryEntity:
        """Add observations to an existing entity"""
        data = {"observations": observations}
        response = await self._make_request(
            "POST", 
            f"/api/v1/entities/{entity_id}/observations", 
            json=data
        )
        return MemoryEntity(**response)
        
    async def create_relation(self, relation: MemoryRelation) -> Dict[str, Any]:
        """Create a relation between entities"""
        data = relation.model_dump(by_alias=True)
        return await self._make_request("POST", "/api/v1/relations", json=data)
        
    async def get_relations(self, entity_id: str) -> List[MemoryRelation]:
        """Get all relations for an entity"""
        response = await self._make_request("GET", f"/api/v1/entities/{entity_id}/relations")
        return [MemoryRelation(**rel) for rel in response.get("relations", [])]
        
    async def search_memories(self, query: str, limit: int = 10) -> List[Dict[str, Any]]:
        """Search across all memories"""
        params = {"q": query, "limit": limit}
        response = await self._make_request("GET", "/api/v1/search", params=params)
        return response.get("results", [])