""" Memory API Service Provides external access to memory read and write operations through API Key authentication. This service validates inputs and delegates to MemoryAgentService for core memory operations. """ import uuid from typing import Any, Dict, Optional from sqlalchemy.orm import Session from app.celery_task_scheduler import scheduler from app.core.error_codes import BizCode from app.core.exceptions import BusinessException, ResourceNotFoundException from app.core.logging_config import get_logger from app.models.app_model import App from app.models.end_user_model import EndUser from app.schemas.memory_config_schema import ConfigurationError from app.schemas.memory_agent_schema import WriteMemoryRequest from app.services.memory_agent_service import MemoryAgentService logger = get_logger(__name__) class MemoryAPIService: """Service for memory API operations with validation and delegation to MemoryAgentService. This service provides a thin layer that: 1. Validates end_user exists and belongs to the authorized workspace 2. Maps end_user_id to end_user_id for memory operations 3. Delegates to MemoryAgentService for actual memory read/write operations """ def __init__(self, db: Session): """Initialize MemoryAPIService. Args: db: SQLAlchemy database session """ self.db = db def validate_end_user( self, end_user_id: str, workspace_id: uuid.UUID ) -> EndUser: """Validate that end_user exists and belongs to the workspace. Args: end_user_id: End user ID to validate workspace_id: Workspace ID from API key authorization Returns: EndUser object if valid Raises: ResourceNotFoundException: If end_user not found BusinessException: If end_user not in authorized workspace """ logger.info(f"Validating end_user: {end_user_id} for workspace: {workspace_id}") # Query end_user by ID try: end_user_uuid = uuid.UUID(end_user_id) except ValueError: logger.warning(f"Invalid end_user_id format: {end_user_id}") raise BusinessException( message=f"Invalid end_user_id format: {end_user_id}", code=BizCode.INVALID_PARAMETER ) end_user = self.db.query(EndUser).filter(EndUser.id == end_user_uuid).first() if not end_user: logger.warning(f"End user not found: {end_user_id}") raise ResourceNotFoundException( resource_type="EndUser", resource_id=end_user_id ) # Verify end_user belongs to the workspace via App relationship app = self.db.query(App).filter( App.id == end_user.app_id, App.is_active.is_(True) ).first() if not app: logger.warning(f"App not found for end_user: {end_user_id}") # raise ResourceNotFoundException( # resource_type="App", # resource_id=str(end_user.app_id) # ) # temporally allow any workspace to access # if end_user.workspace_id != workspace_id: # print(f"[DEBUG] end_user.workspace_id={end_user.workspace_id}, api_key.workspace_id={workspace_id}") # logger.warning( # f"End user {end_user_id} belongs to workspace {end_user.workspace_id}, " # f"not authorized workspace {workspace_id}" # ) # raise BusinessException( # message=f"End user does not belong to authorized workspace. end_user.workspace_id={end_user.workspace_id}, api_key.workspace_id={workspace_id}", # code=BizCode.FORBIDDEN # ) logger.info(f"End user {end_user_id} validated successfully") return end_user def _update_end_user_config(self, end_user_id: str, config_id: str) -> None: """Update the end user's memory_config_id. Silently updates the config association. Logs warnings on failure but does not raise, so it won't block the main read/write operation. Args: end_user_id: End user identifier config_id: Memory configuration ID to assign """ try: config_uuid = uuid.UUID(config_id) from app.repositories.end_user_repository import EndUserRepository end_user_repo = EndUserRepository(self.db) end_user_repo.update_memory_config_id( end_user_id=uuid.UUID(end_user_id), memory_config_id=config_uuid, ) except Exception as e: logger.warning(f"Failed to update memory_config_id for end_user {end_user_id}: {e}") def write_memory( self, workspace_id: uuid.UUID, end_user_id: str, message: str, config_id: str, storage_type: str = "neo4j", user_rag_memory_id: Optional[str] = None, ) -> Dict[str, Any]: """Submit a memory write task via Celery. Validates end_user exists and belongs to workspace, updates the end user's memory_config_id, then dispatches write_message_task to Celery for async processing with per-user fair locking. Args: workspace_id: Workspace ID for resource validation end_user_id: End user identifier message: Message content to store config_id: Memory configuration ID (required) storage_type: Storage backend (neo4j or rag) user_rag_memory_id: Optional RAG memory ID Returns: Dict with task_id, status, and end_user_id Raises: ResourceNotFoundException: If end_user not found BusinessException: If validation fails """ logger.info(f"Submitting memory write for end_user: {end_user_id}, workspace: {workspace_id}") # Validate end_user exists and belongs to workspace self.validate_end_user(end_user_id, workspace_id) # Update end user's memory_config_id self._update_end_user_config(end_user_id, config_id) # Convert to message list format expected by write_message_task messages = message if isinstance(message, list) else [{"role": "user", "content": message}] # from app.tasks import write_message_task # task = write_message_task.delay( # end_user_id, # messages, # config_id, # storage_type, # user_rag_memory_id or "", # ) task_id = scheduler.push_task( "app.core.memory.agent.write_message", end_user_id, { "end_user_id": end_user_id, "message": messages, "config_id": config_id, "storage_type": storage_type, "user_rag_memory_id": user_rag_memory_id or "" } ) logger.info(f"Memory write task submitted, task_id={task_id} end_user_id={end_user_id}") return { "task_id": task_id, "status": "QUEUED", "end_user_id": end_user_id, } def read_memory( self, workspace_id: uuid.UUID, end_user_id: str, message: str, search_switch: str = "0", config_id: str = "", storage_type: str = "neo4j", user_rag_memory_id: Optional[str] = None, ) -> Dict[str, Any]: """Submit a memory read task via Celery. Validates end_user exists and belongs to workspace, updates the end user's memory_config_id, then dispatches read_message_task to Celery for async processing. Args: workspace_id: Workspace ID for resource validation end_user_id: End user identifier message: Query message search_switch: Search mode (0=deep search with verification, 1=deep search, 2=fast search) config_id: Memory configuration ID (required) storage_type: Storage backend (neo4j or rag) user_rag_memory_id: Optional RAG memory ID Returns: Dict with task_id, status, and end_user_id Raises: ResourceNotFoundException: If end_user not found BusinessException: If validation fails """ logger.info(f"Submitting memory read for end_user: {end_user_id}, workspace: {workspace_id}") # Validate end_user exists and belongs to workspace self.validate_end_user(end_user_id, workspace_id) # Update end user's memory_config_id self._update_end_user_config(end_user_id, config_id) from app.tasks import read_message_task task = read_message_task.delay( end_user_id, message, [], # history search_switch, config_id, storage_type, user_rag_memory_id or "", ) logger.info(f"Memory read task submitted: task_id={task.id}, end_user_id={end_user_id}") return { "task_id": task.id, "status": "PENDING", "end_user_id": end_user_id, } async def write_memory_sync( self, workspace_id: uuid.UUID, end_user_id: str, message: str, config_id: str, storage_type: str = "neo4j", user_rag_memory_id: Optional[str] = None, ) -> Dict[str, Any]: """Write memory synchronously (inline, no Celery). Validates end_user, then calls MemoryAgentService.write_memory directly. Blocks until the write completes. Use for cases where the caller needs immediate confirmation. Args: workspace_id: Workspace ID for resource validation end_user_id: End user identifier message: Message content to store config_id: Memory configuration ID (required) storage_type: Storage backend (neo4j or rag) user_rag_memory_id: Optional RAG memory ID Returns: Dict with status and end_user_id Raises: ResourceNotFoundException: If end_user not found BusinessException: If write fails """ logger.info(f"Writing memory (sync) for end_user: {end_user_id}, workspace: {workspace_id}") self.validate_end_user(end_user_id, workspace_id) self._update_end_user_config(end_user_id, config_id) try: messages = message if isinstance(message, list) else [{"role": "user", "content": message}] result = await MemoryAgentService().write_memory( WriteMemoryRequest( end_user_id=end_user_id, messages=messages, config_id=config_id, storage_type=storage_type, user_rag_memory_id=user_rag_memory_id or "", ), self.db, ) logger.info(f"Memory write (sync) successful for end_user: {end_user_id}") if isinstance(result, dict): return { **result, "status": result.get("status", "unknown"), "end_user_id": end_user_id, } return { "status": result if isinstance(result, str) else "success", "end_user_id": end_user_id, } except ConfigurationError as e: logger.error(f"Memory configuration error for end_user {end_user_id}: {e}") raise BusinessException(message=str(e), code=BizCode.MEMORY_CONFIG_NOT_FOUND) except BusinessException: raise except Exception as e: logger.error(f"Memory write (sync) failed for end_user {end_user_id}: {e}") raise BusinessException( message=f"Memory write failed: {str(e)}", code=BizCode.MEMORY_WRITE_FAILED ) async def read_memory_sync( self, workspace_id: uuid.UUID, end_user_id: str, message: str, search_switch: str = "0", config_id: str = "", storage_type: str = "neo4j", user_rag_memory_id: Optional[str] = None, ) -> Dict[str, Any]: """Read memory synchronously (inline, no Celery). Validates end_user, then calls MemoryAgentService.read_memory directly. Blocks until the read completes. Use for cases where the caller needs the answer immediately. Args: workspace_id: Workspace ID for resource validation end_user_id: End user identifier message: Query message search_switch: Search mode (0=deep search with verification, 1=deep search, 2=fast search) config_id: Memory configuration ID (required) storage_type: Storage backend (neo4j or rag) user_rag_memory_id: Optional RAG memory ID Returns: Dict with answer, intermediate_outputs, and end_user_id Raises: ResourceNotFoundException: If end_user not found BusinessException: If read fails """ logger.info(f"Reading memory (sync) for end_user: {end_user_id}, workspace: {workspace_id}") self.validate_end_user(end_user_id, workspace_id) self._update_end_user_config(end_user_id, config_id) try: result = await MemoryAgentService().read_memory( end_user_id=end_user_id, message=message, history=[], search_switch=search_switch, config_id=config_id, db=self.db, storage_type=storage_type, user_rag_memory_id=user_rag_memory_id or "" ) logger.info(f"Memory read (sync) successful for end_user: {end_user_id}") return { "answer": result.get("answer", ""), "intermediate_outputs": result.get("intermediate_outputs", []), "end_user_id": end_user_id } except ConfigurationError as e: logger.error(f"Memory configuration error for end_user {end_user_id}: {e}") raise BusinessException(message=str(e), code=BizCode.MEMORY_CONFIG_NOT_FOUND) except BusinessException: raise except Exception as e: logger.error(f"Memory read (sync) failed for end_user {end_user_id}: {e}") raise BusinessException( message=f"Memory read failed: {str(e)}", code=BizCode.MEMORY_READ_FAILED ) def create_end_user( self, workspace_id: uuid.UUID, other_id: str, ) -> Dict[str, Any]: """Create or retrieve an end user for the workspace. Uses get_or_create semantics: if an end user with the same other_id already exists in the workspace, returns the existing one. Args: workspace_id: Workspace ID from API key authorization other_id: External user identifier Returns: Dict with id, other_id, other_name, and workspace_id Raises: BusinessException: If creation fails """ logger.info(f"Creating end user - other_id: {other_id}, workspace_id: {workspace_id}") try: from app.repositories.end_user_repository import EndUserRepository end_user_repo = EndUserRepository(self.db) end_user = end_user_repo.get_or_create_end_user( app_id=None, workspace_id=workspace_id, other_id=other_id, ) logger.info(f"End user ready: {end_user.id}") return { "id": str(end_user.id), "other_id": end_user.other_id or "", "other_name": end_user.other_name or "", "workspace_id": str(end_user.workspace_id), } except Exception as e: logger.error(f"Failed to create end user for workspace {workspace_id}: {e}") raise BusinessException( message=f"Failed to create end user: {str(e)}", code=BizCode.INTERNAL_ERROR ) def list_memory_configs( self, workspace_id: uuid.UUID, ) -> Dict[str, Any]: """List all memory configs for a workspace. Args: workspace_id: Workspace ID from API key authorization Returns: Dict with configs list and total count Raises: BusinessException: If listing fails """ logger.info(f"Listing memory configs for workspace: {workspace_id}") try: from app.repositories.memory_config_repository import MemoryConfigRepository results = MemoryConfigRepository.get_all(self.db, workspace_id=workspace_id) configs = [] for config, scene_name in results: configs.append({ "config_id": str(config.config_id), "config_name": config.config_name, "config_desc": config.config_desc, "is_default": config.is_default or False, "scene_name": scene_name, "created_at": config.created_at.isoformat() if config.created_at else None, "updated_at": config.updated_at.isoformat() if config.updated_at else None, }) logger.info(f"Found {len(configs)} memory configs for workspace {workspace_id}") return { "configs": configs, "total": len(configs), } except Exception as e: logger.error(f"Failed to list memory configs for workspace {workspace_id}: {e}") raise BusinessException( message=f"Failed to list memory configs: {str(e)}", code=BizCode.MEMORY_READ_FAILED )