mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-15 23:31:05 +00:00
96 lines
3.5 KiB
Python
96 lines
3.5 KiB
Python
"""Base class for memory tool"""
|
|
|
|
from abc import ABCMeta
|
|
from pathlib import Path
|
|
|
|
from ...core.enumeration import MemoryType
|
|
from ...core.op import BaseTool
|
|
from ...core.schema import ToolCall, MemoryNode, ToolAttr
|
|
from ...core.utils import CacheHandler
|
|
|
|
|
|
class BaseMemoryTool(BaseTool, metaclass=ABCMeta):
|
|
"""Base class for memory tool"""
|
|
|
|
def __init__(
|
|
self,
|
|
enable_multiple: bool = True,
|
|
enable_thinking_params: bool = False,
|
|
local_memory_path: str = "./reme_local_memory",
|
|
**kwargs,
|
|
):
|
|
super().__init__(**kwargs)
|
|
self.enable_multiple: bool = enable_multiple
|
|
self.enable_thinking_params: bool = enable_thinking_params
|
|
self.local_memory_path: str = local_memory_path
|
|
self.memory_nodes: list[MemoryNode | str] = []
|
|
|
|
def _build_tool_call(self) -> ToolCall:
|
|
"""Build and return the tool call schema"""
|
|
|
|
def _build_multiple_tool_call(self) -> ToolCall:
|
|
"""Build and return the multiple tool call schema"""
|
|
|
|
@property
|
|
def tool_call(self) -> ToolCall | None:
|
|
"""Get the tool call schema."""
|
|
if self._tool_call is None:
|
|
if self.enable_multiple:
|
|
self._tool_call = self._build_multiple_tool_call()
|
|
else:
|
|
self._tool_call = self._build_tool_call()
|
|
self._tool_call.name = self._tool_call.name or self.name
|
|
|
|
# Add thinking parameter if enabled
|
|
if self.enable_thinking_params:
|
|
parameters = self._tool_call.parameters
|
|
if parameters and parameters.properties is not None:
|
|
if "thinking" not in parameters.properties:
|
|
parameters.properties = {
|
|
"thinking": ToolAttr(
|
|
type="string",
|
|
description="Your complete and detailed thinking process "
|
|
"about how to fill in each parameter",
|
|
),
|
|
**parameters.properties,
|
|
}
|
|
if parameters.required is not None:
|
|
parameters.required = ["thinking", *parameters.required]
|
|
else:
|
|
parameters.required = ["thinking"]
|
|
return self._tool_call
|
|
|
|
@property
|
|
def local_memory(self) -> CacheHandler:
|
|
"""Create the meta memory cache handler."""
|
|
return CacheHandler(Path(self.local_memory_path) / self.vector_store.collection_name)
|
|
|
|
@property
|
|
def memory_type(self) -> MemoryType:
|
|
"""Get the memory type from context."""
|
|
return MemoryType(self.context.get("memory_type"))
|
|
|
|
@property
|
|
def memory_target(self) -> str:
|
|
"""Get the memory target from context."""
|
|
return self.context.get("memory_target", "")
|
|
|
|
@property
|
|
def memory_cache_key(self) -> str:
|
|
"""Get the memory cache key from context."""
|
|
return f"{self.memory_type.value}_{self.memory_target}".replace(" ", "_").lower()
|
|
|
|
@property
|
|
def history_node(self) -> MemoryNode:
|
|
"""Get the history node from context."""
|
|
return self.context.get("history_node")
|
|
|
|
@property
|
|
def retrieved_nodes(self) -> list[MemoryNode]:
|
|
"""Get the retrieved nodes from context."""
|
|
return self.context.get("retrieved_nodes")
|
|
|
|
@property
|
|
def author(self) -> str:
|
|
"""Get the author from context."""
|
|
return self.context.get("author", "")
|