ChatDev/runtime/node/agent/memory/builtin_stores.py

56 lines
1.6 KiB
Python
Executable File

"""Register built-in memory stores."""
from entity.configs.node.memory import (
BlackboardMemoryConfig,
FileMemoryConfig,
Mem0MemoryConfig,
SimpleMemoryConfig,
MemoryStoreConfig,
)
from runtime.node.agent.memory.blackboard_memory import BlackboardMemory
from runtime.node.agent.memory.file_memory import FileMemory
from runtime.node.agent.memory.memory_base import MemoryBase
from runtime.node.agent.memory.simple_memory import SimpleMemory
from runtime.node.agent.memory.registry import register_memory_store, get_memory_store_registration
register_memory_store(
"simple",
config_cls=SimpleMemoryConfig,
factory=lambda store: SimpleMemory(store),
summary="In-memory store that resets between runs; best for testing",
)
register_memory_store(
"file",
config_cls=FileMemoryConfig,
factory=lambda store: FileMemory(store),
summary="Persists documents on disk and supports embedding search",
)
register_memory_store(
"blackboard",
config_cls=BlackboardMemoryConfig,
factory=lambda store: BlackboardMemory(store),
summary="Shared blackboard memory allowing multiple nodes to read/write",
)
def _create_mem0_memory(store):
from runtime.node.agent.memory.mem0_memory import Mem0Memory
return Mem0Memory(store)
register_memory_store(
"mem0",
config_cls=Mem0MemoryConfig,
factory=_create_mem0_memory,
summary="Mem0 managed memory with semantic search and graph relationships",
)
class MemoryFactory:
@staticmethod
def create_memory(store: MemoryStoreConfig) -> MemoryBase:
registration = get_memory_store_registration(store.type)
return registration.factory(store)