adai/MemoryBear
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(memory): support perception-aware memory writing in workflow and Neo4j nodes

Browse Source
This commit is contained in:
Eternity
2026-03-23 16:33:25 +08:00
parent 31085ed678
commit 2ff81ba101
22 changed files with 820 additions and 519 deletions
Download Patch File Download Diff File Expand all files Collapse all files

85
api/app/services/memory_api_service.py
View File

@@ -28,7 +28,7 @@ class MemoryAPIService:
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.
@@ -36,11 +36,11 @@ class MemoryAPIService:
db: SQLAlchemy database session
"""
self.db = db
def validate_end_user(
self,
end_user_id: str,
workspace_id: uuid.UUID
self,
end_user_id: str,
workspace_id: uuid.UUID
) -> EndUser:
"""Validate that end_user exists and belongs to the workspace.
@@ -56,7 +56,7 @@ class MemoryAPIService:
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)
@@ -66,7 +66,7 @@ class MemoryAPIService:
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:
@@ -75,13 +75,13 @@ class MemoryAPIService:
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(
@@ -99,7 +99,7 @@ class MemoryAPIService:
# 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
@@ -125,13 +125,14 @@ class MemoryAPIService:
logger.warning(f"Failed to update memory_config_id for end_user {end_user_id}: {e}")
async 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,
self,
workspace_id: uuid.UUID,
end_user_id: str,
message: str,
config_id: str,
storage_type: str = "neo4j",
files: Optional[list]=None,
user_rag_memory_id: Optional[str] = None,
) -> Dict[str, Any]:
"""Write memory with validation.
@@ -153,14 +154,16 @@ class MemoryAPIService:
ResourceNotFoundException: If end_user not found
BusinessException: If end_user not in authorized workspace or write fails
"""
if files is None:
files = list()
logger.info(f"Writing memory 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)
try:
# Delegate to MemoryAgentService
# Convert string message to list[dict] format expected by MemoryAgentService
@@ -171,11 +174,12 @@ class MemoryAPIService:
config_id=config_id,
db=self.db,
storage_type=storage_type,
user_rag_memory_id=user_rag_memory_id or ""
user_rag_memory_id=user_rag_memory_id or "",
files=files
)
logger.info(f"Memory write successful for end_user: {end_user_id}")
# result may be a string "success" or a dict with a "status" key
# Preserve the full dict so callers don't silently lose extra fields
# (e.g. error codes, metadata) returned by MemoryAgentService.
@@ -189,7 +193,7 @@ class MemoryAPIService:
"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(
@@ -204,16 +208,16 @@ class MemoryAPIService:
message=f"Memory write failed: {str(e)}",
code=BizCode.MEMORY_WRITE_FAILED
)
async 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,
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 with validation.
@@ -237,14 +241,13 @@ class MemoryAPIService:
BusinessException: If end_user not in authorized workspace or read fails
"""
logger.info(f"Reading memory 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)
try:
# Delegate to MemoryAgentService
result = await MemoryAgentService().read_memory(
@@ -257,15 +260,15 @@ class MemoryAPIService:
storage_type=storage_type,
user_rag_memory_id=user_rag_memory_id or ""
)
logger.info(f"Memory read 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(
@@ -282,8 +285,8 @@ class MemoryAPIService:
)
def list_memory_configs(
self,
workspace_id: uuid.UUID,
self,
workspace_id: uuid.UUID,
) -> Dict[str, Any]:
"""List all memory configs for a workspace.