489 lines
18 KiB
Python
489 lines
18 KiB
Python
"""
|
|
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.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(
|
|
end_user_id=end_user_id,
|
|
messages=messages,
|
|
config_id=config_id,
|
|
db=self.db,
|
|
storage_type=storage_type,
|
|
user_rag_memory_id=user_rag_memory_id or "",
|
|
)
|
|
|
|
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
|
|
)
|