mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-16 23:41:10 +00:00
212 lines
7.5 KiB
Python
212 lines
7.5 KiB
Python
"""Task memory service for managing task-oriented memories.
|
|
|
|
This module provides the TaskMemoryService class which extends the base
|
|
AgentscopeRuntimeMemoryService to handle task-related memory operations.
|
|
It supports creating, retrieving, listing, and deleting task memories
|
|
using flow-based execution.
|
|
"""
|
|
|
|
import asyncio
|
|
from typing import Optional, Dict, Any, List
|
|
|
|
from flowllm.core.schema import FlowResponse
|
|
from loguru import logger
|
|
from pydantic import Field, BaseModel
|
|
|
|
from reme_ai.schema.memory import TaskMemory
|
|
from reme_ai.service.agentscope_runtime_memory_service import AgentscopeRuntimeMemoryService
|
|
|
|
|
|
class TaskMemoryService(AgentscopeRuntimeMemoryService):
|
|
"""Service for managing task-oriented memories.
|
|
|
|
TaskMemoryService helps efficiently manage and schedule task-related memories,
|
|
enhancing both the accuracy and efficiency of task execution. Powered by LLM
|
|
capabilities, it supports flexible creation, retrieval, update, and deletion
|
|
of memories across diverse task scenarios.
|
|
"""
|
|
|
|
async def start(self):
|
|
"""Start the task memory service.
|
|
|
|
Returns:
|
|
The result of starting the underlying application.
|
|
"""
|
|
return await self.app.async_start()
|
|
|
|
async def stop(self) -> None:
|
|
"""Stop the task memory service.
|
|
|
|
Releases resources and stops the underlying application.
|
|
"""
|
|
return await self.app.async_stop()
|
|
|
|
async def health(self) -> bool:
|
|
"""Check the health status of the service.
|
|
|
|
Returns:
|
|
True if the service is healthy, False otherwise.
|
|
"""
|
|
return True
|
|
|
|
async def add_memory(self, user_id: str, messages: list, session_id: Optional[str] = None) -> None:
|
|
"""Add task memory from messages.
|
|
|
|
Processes the provided messages and creates task memories using
|
|
the summary_task_memory flow. The created memories are associated
|
|
with the given session_id.
|
|
|
|
Args:
|
|
user_id: The user identifier.
|
|
messages: List of messages (dict or BaseModel instances) to process.
|
|
session_id: Optional session identifier to associate with the memory.
|
|
"""
|
|
new_messages: List[dict] = []
|
|
for message in messages:
|
|
if isinstance(message, dict):
|
|
new_messages.append(message)
|
|
elif isinstance(message, BaseModel):
|
|
new_messages.append(message.model_dump())
|
|
else:
|
|
raise ValueError(f"Invalid message type={type(message)}")
|
|
|
|
kwargs = {
|
|
"workspace_id": user_id,
|
|
"trajectories": [
|
|
{"messages": new_messages, "score": 1.0},
|
|
],
|
|
}
|
|
|
|
result: FlowResponse = await self.app.async_execute_flow(name="summary_task_memory", **kwargs)
|
|
memory_list: List[TaskMemory] = result.metadata.get("memory_list", [])
|
|
for memory in memory_list:
|
|
memory_id = memory.memory_id
|
|
self.add_session_memory_id(session_id, memory_id)
|
|
logger.info(f"[task_memory_service] user_id={user_id} session_id={session_id} add memory: {memory}")
|
|
|
|
async def search_memory(
|
|
self,
|
|
user_id: str,
|
|
messages: list,
|
|
filters: Optional[Dict[str, Any]] = Field(
|
|
description="Associated filters for the messages, " "such as top_k, score etc.",
|
|
default=None,
|
|
),
|
|
) -> list:
|
|
"""Search for task memories matching the given messages.
|
|
|
|
Searches the memory store for task memories relevant to the provided
|
|
messages using the retrieve_task_memory flow.
|
|
|
|
Args:
|
|
user_id: The user identifier.
|
|
messages: List of messages (dict or BaseModel instances) to search with.
|
|
filters: Optional filters including top_k for controlling search results.
|
|
|
|
Returns:
|
|
List containing the search result answer.
|
|
"""
|
|
new_messages: List[dict] = []
|
|
for message in messages:
|
|
if isinstance(message, dict):
|
|
new_messages.append(message)
|
|
elif isinstance(message, BaseModel):
|
|
new_messages.append(message.model_dump())
|
|
else:
|
|
raise ValueError(f"Invalid message type={type(message)}")
|
|
|
|
kwargs = {
|
|
"workspace_id": user_id,
|
|
"messages": new_messages,
|
|
"top_k": filters.get("top_k", 1) if filters else 1,
|
|
}
|
|
|
|
result: FlowResponse = await self.app.async_execute_flow(name="retrieve_task_memory", **kwargs)
|
|
logger.info(f"[task_memory_service] user_id={user_id} add result: {result.model_dump_json()}")
|
|
|
|
return [result.answer]
|
|
|
|
async def list_memory(
|
|
self,
|
|
user_id: str,
|
|
filters: Optional[Dict[str, Any]] = Field(
|
|
description="Associated filters for the messages, " "such as top_k, score etc.",
|
|
default=None,
|
|
),
|
|
) -> list:
|
|
"""List all task memories for a user.
|
|
|
|
Retrieves all task memories associated with the given user_id
|
|
from the vector store.
|
|
|
|
Args:
|
|
user_id: The user identifier.
|
|
filters: Optional filters (currently not used but kept for API consistency).
|
|
|
|
Returns:
|
|
List of memory items for the user.
|
|
"""
|
|
result = await self.app.async_execute_flow(name="vector_store", workspace_id=user_id, action="list")
|
|
print("list_memory result:", result)
|
|
|
|
result = result.metadata["action_result"]
|
|
for i, line in enumerate(result):
|
|
logger.info(f"[task_memory_service] list memory.{i}={line}")
|
|
return result
|
|
|
|
async def delete_memory(self, user_id: str, session_id: Optional[str] = None) -> None:
|
|
"""Delete task memories for a user session.
|
|
|
|
Deletes all memories associated with the given session_id for the user.
|
|
If no session_id is provided or no memories exist for the session,
|
|
no deletion is performed.
|
|
|
|
Args:
|
|
user_id: The user identifier.
|
|
session_id: Optional session identifier. If provided, only memories
|
|
associated with this session will be deleted.
|
|
"""
|
|
delete_ids = self.session_id_dict.get(session_id, [])
|
|
if not delete_ids:
|
|
return
|
|
|
|
result = await self.app.async_execute_flow(
|
|
name="vector_store",
|
|
workspace_id=user_id,
|
|
action="delete_ids",
|
|
memory_ids=delete_ids,
|
|
)
|
|
result = result.metadata["action_result"]
|
|
logger.info(f"[task_memory_service] delete memory result={result}")
|
|
|
|
|
|
async def main():
|
|
"""Main function for testing the TaskMemoryService.
|
|
|
|
Demonstrates the usage of TaskMemoryService by adding, searching,
|
|
listing, and deleting task memories.
|
|
"""
|
|
async with TaskMemoryService() as service:
|
|
logger.info("========== start task memory service ==========")
|
|
|
|
await service.add_memory(
|
|
user_id="u_123456",
|
|
messages=[{"content": "please use web search tool to search financial news:"}],
|
|
session_id="s_123456",
|
|
)
|
|
|
|
await service.search_memory(
|
|
user_id="u_123456",
|
|
messages=[{"content": "please use web search tool to search financial news"}],
|
|
filters={"top_k": 1},
|
|
)
|
|
|
|
await service.list_memory(user_id="u_123456")
|
|
await service.delete_memory(user_id="u_123456", session_id="s_123456")
|
|
await service.list_memory(user_id="u_123456")
|
|
|
|
logger.info("========== end task memory service ==========")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|