128 lines
5.4 KiB
Python
128 lines
5.4 KiB
Python
"""
|
|
Firestore-backed implementation of the MemoryStore interface.
|
|
"""
|
|
from typing import List, Dict, Any, Optional
|
|
from datetime import datetime, timezone
|
|
|
|
from google.cloud import firestore
|
|
from google.cloud.firestore_v1.async_transaction import async_transactional
|
|
|
|
from contracts.common import EmmaConversation, MemoryRecord, AuditEvent
|
|
from .memory_store import MemoryStore
|
|
|
|
class FirestoreMemoryStore(MemoryStore):
|
|
"""
|
|
Implements the MemoryStore interface using Google Cloud Firestore.
|
|
"""
|
|
def __init__(self, project_id: str, client: Optional[firestore.AsyncClient] = None):
|
|
self._project_id = project_id
|
|
# Client can be injected for testing, otherwise it will be created on demand.
|
|
self._client = client
|
|
|
|
def _get_client(self) -> firestore.AsyncClient:
|
|
if not self._client:
|
|
self._client = firestore.AsyncClient(project=self._project_id)
|
|
return self._client
|
|
|
|
async def create_conversation(self, owner_id: str, workspace_id: str, created_by: str) -> EmmaConversation:
|
|
db = self._get_client()
|
|
convo = EmmaConversation(owner_id=owner_id, workspace_id=workspace_id, created_by=created_by)
|
|
doc_ref = db.collection("emma_conversations").document(convo.conversation_id)
|
|
await doc_ref.set(convo.model_dump(mode='json'))
|
|
return convo
|
|
|
|
async def save_conversation(self, conversation: EmmaConversation) -> None:
|
|
if not conversation.owner_id:
|
|
raise ValueError("Conversation must have an owner_id.")
|
|
|
|
db = self._get_client()
|
|
doc_ref = db.collection("emma_conversations").document(
|
|
conversation.conversation_id
|
|
)
|
|
conversation_data = conversation.model_dump(mode="json")
|
|
|
|
@async_transactional
|
|
async def _save(transaction) -> None:
|
|
snapshot = await doc_ref.get(transaction=transaction)
|
|
|
|
if not snapshot.exists:
|
|
transaction.create(doc_ref, conversation_data)
|
|
return
|
|
|
|
existing_data = snapshot.to_dict() or {}
|
|
|
|
if existing_data.get("owner_id") != conversation.owner_id:
|
|
raise PermissionError(
|
|
"Cannot modify a conversation owned by another owner."
|
|
)
|
|
|
|
if existing_data.get("workspace_id") != conversation.workspace_id:
|
|
raise PermissionError(
|
|
"Cannot move a conversation to another workspace."
|
|
)
|
|
|
|
conversation_data["owner_id"] = existing_data["owner_id"]
|
|
conversation_data["workspace_id"] = existing_data["workspace_id"]
|
|
|
|
if existing_data.get("created_at") is not None:
|
|
conversation_data["created_at"] = existing_data["created_at"]
|
|
|
|
transaction.update(doc_ref, conversation_data)
|
|
|
|
await _save(db.transaction())
|
|
|
|
async def get_conversation(self, conversation_id: str, owner_id: str) -> Optional[EmmaConversation]:
|
|
db = self._get_client()
|
|
doc_ref = db.collection("emma_conversations").document(conversation_id)
|
|
doc = await doc_ref.get()
|
|
if not doc.exists:
|
|
return None
|
|
|
|
convo = EmmaConversation(**doc.to_dict())
|
|
# Security: Enforce ownership
|
|
if convo.owner_id != owner_id:
|
|
return None # Or raise an exception
|
|
return convo
|
|
|
|
async def append_message(self, conversation_id: str, message: Dict[str, Any]) -> None:
|
|
db = self._get_client()
|
|
# Note: message_id should be auto-generated by Firestore for ordering
|
|
messages_ref = db.collection("emma_conversations").document(conversation_id).collection("messages")
|
|
await messages_ref.add(message)
|
|
# Update last_updated_at on parent conversation document
|
|
convo_ref = db.collection("emma_conversations").document(conversation_id)
|
|
await convo_ref.update({"last_updated_at": datetime.now(timezone.utc)})
|
|
|
|
async def list_messages(self, conversation_id: str, limit: int = 50) -> List[Dict[str, Any]]:
|
|
db = self._get_client()
|
|
messages_ref = db.collection("emma_conversations").document(conversation_id).collection("messages")
|
|
docs = messages_ref.order_by("created_at", direction=firestore.Query.DESCENDING).limit(limit).stream()
|
|
return [doc.to_dict() async for doc in docs][::-1] # Reverse to get chronological order
|
|
|
|
async def record_audit_event(self, event: AuditEvent) -> str:
|
|
db = self._get_client()
|
|
doc_ref = db.collection("emma_audit_events").document(event.event_id)
|
|
# Use create() to ensure the document does not already exist, guaranteeing append-only.
|
|
await doc_ref.create(event.model_dump(mode='json'))
|
|
return event.event_id
|
|
|
|
# --- Placeholder methods not fully implemented in Phase 2B ---
|
|
|
|
async def save_conversation_summary(self, conversation_id: str, summary: str) -> None:
|
|
# In a real implementation, this would update a summary field
|
|
pass
|
|
|
|
async def store_memory(self, record: MemoryRecord) -> None:
|
|
# This would write to the emma_memories collection
|
|
pass
|
|
|
|
async def recall_memory(self, embedding: List[float], owner_id: str, scopes: List[str]) -> List[MemoryRecord]:
|
|
# This would perform a query against a vector database, which is out of scope.
|
|
return []
|
|
|
|
async def get_profile(self, owner_id: str) -> Optional[Dict[str, Any]]:
|
|
pass
|
|
|
|
async def update_profile(self, owner_id: str, profile_data: Dict[str, Any]) -> None:
|
|
pass
|