99 lines
4.0 KiB
Python
99 lines
4.0 KiB
Python
"""
|
|
Defines the abstract interface for Emma's memory systems and provides a
|
|
non-persistent, in-memory implementation for testing.
|
|
"""
|
|
from abc import ABC, abstractmethod
|
|
from typing import List, Dict, Any, Optional
|
|
from contracts.common import EmmaConversation, MemoryRecord, AuditEvent, CallEmmaRequest
|
|
|
|
class MemoryStore(ABC):
|
|
"""Abstract Base Class for all memory store implementations."""
|
|
|
|
@abstractmethod
|
|
async def create_conversation(self, owner_id: str, workspace_id: str, created_by: str) -> EmmaConversation:
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def save_conversation(self, conversation: EmmaConversation) -> None:
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def get_conversation(self, conversation_id: str, owner_id: str) -> Optional[EmmaConversation]:
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def append_message(self, conversation_id: str, message: Dict[str, Any]) -> None:
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def list_messages(self, conversation_id: str, limit: int = 50) -> List[Dict[str, Any]]:
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def store_memory(self, record: MemoryRecord) -> None:
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def recall_memory(self, embedding: List[float], owner_id: str, scopes: List[str]) -> List[MemoryRecord]:
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def record_audit_event(self, event: AuditEvent) -> str:
|
|
pass
|
|
|
|
class InMemoryMemoryStore(MemoryStore):
|
|
"""
|
|
A non-persistent, in-memory implementation of the MemoryStore for testing and
|
|
local development. DO NOT USE IN PRODUCTION.
|
|
"""
|
|
def __init__(self):
|
|
self._conversations: Dict[str, EmmaConversation] = {}
|
|
self._messages: Dict[str, List[Dict[str, Any]]] = {}
|
|
self._memories: List[MemoryRecord] = []
|
|
self._audit: List[AuditEvent] = []
|
|
print("WARNING: InMemoryMemoryStore is active. Data will not be persisted.")
|
|
|
|
async def create_conversation(self, owner_id: str, workspace_id: str, created_by: str) -> EmmaConversation:
|
|
convo = EmmaConversation(owner_id=owner_id, workspace_id=workspace_id, created_by=created_by)
|
|
self._conversations[convo.conversation_id] = convo
|
|
self._messages[convo.conversation_id] = []
|
|
return convo
|
|
|
|
async def save_conversation(self, conversation: EmmaConversation) -> None:
|
|
if not conversation.owner_id:
|
|
raise ValueError("Conversation must have an owner_id.")
|
|
# Prevent an existing conversation from being claimed by a new owner.
|
|
existing = self._conversations.get(conversation.conversation_id)
|
|
if existing and existing.owner_id != conversation.owner_id:
|
|
raise ValueError("Cannot change the owner of an existing conversation.")
|
|
|
|
self._conversations[conversation.conversation_id] = conversation
|
|
# Ensure the message list is initialized.
|
|
if conversation.conversation_id not in self._messages:
|
|
self._messages[conversation.conversation_id] = []
|
|
|
|
async def get_conversation(self, conversation_id: str, owner_id: str) -> Optional[EmmaConversation]:
|
|
convo = self._conversations.get(conversation_id)
|
|
if convo and convo.owner_id == owner_id:
|
|
return convo
|
|
return None
|
|
|
|
async def append_message(self, conversation_id: str, message: Dict[str, Any]) -> None:
|
|
if conversation_id in self._messages:
|
|
self._messages[conversation_id].append(message)
|
|
|
|
async def list_messages(self, conversation_id: str, limit: int = 50) -> List[Dict[str, Any]]:
|
|
return self._messages.get(conversation_id, [])[-limit:]
|
|
|
|
async def store_memory(self, record: MemoryRecord) -> None:
|
|
if record.owner_id:
|
|
self._memories.append(record)
|
|
|
|
async def recall_memory(self, embedding: List[float], owner_id: str, scopes: List[str]) -> List[MemoryRecord]:
|
|
recalled = [mem for mem in self._memories if mem.owner_id == owner_id]
|
|
return recalled[:5]
|
|
|
|
async def record_audit_event(self, event: AuditEvent) -> str:
|
|
self._audit.append(event)
|
|
return event.event_id
|