scriptoriumadmin/ReMe
1
0
Fork 0
mirror of https://github.com/agentscope-ai/ReMe.git synced 2026-08-28 05:25:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 11

refactor(core): simplify component registration and improve application lifecycle management

Browse source
This commit is contained in:
jinli.yl 2026-01-31 01:05:31 +08:00
parent 777e08ecc9
commit 3680571c94
27 changed files with 389 additions and 334 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
benchmark/halumem/eval_reme.py
View file

@ -43,7 +43,7 @@ class EvalConfig:
output_dir: str = "bench_results/reme"
reme_model_name: str = "qwen-flash"
eval_model_name: str = "qwen3-max"
algo_version: str = "v1"
algo_version: str = "halumem"
# ==================== Utilities ====================
@ -271,7 +271,7 @@ async def evaluation_for_question(
class MemoryProcessor:
"""Handles ReMe memory operations."""
def __init__(self, reme: ReMe, eval_model_name: str = "qwen3-max", algo_version: str = "v1"):
def __init__(self, reme: ReMe, eval_model_name: str = "qwen3-max", algo_version: str = "halumem"):
self.reme = reme
self.eval_model_name = eval_model_name
self.algo_version = algo_version
@ -567,21 +567,11 @@ class HaluMemEvaluator:
async def __aenter__(self):
"""Async context manager entry."""
return self
return await self.reme.start()
async def __aexit__(self, exc_type, exc_val, exc_tb):
"""Async context manager exit with cleanup."""
await self.reme.close()
return False
def __enter__(self):
"""Sync context manager entry."""
return self
def __exit__(self, exc_type, exc_val, exc_tb):
"""Sync context manager exit with cleanup."""
self.reme.close_sync()
return False
return await self.reme.close()
async def process_session(
self,
@ -875,7 +865,7 @@ async def main_async(
max_concurrency: int,
reme_model_name: str= "qwen-flash",
eval_model_name: str = "qwen3-max",
algo_version: str = "v1"
algo_version: str = "halumem"
):
"""Main async entry point for ReMe evaluation with proper resource cleanup."""
config = EvalConfig(
@ -900,7 +890,7 @@ def main(
max_concurrency: int,
reme_model_name: str= "qwen-flash",
eval_model_name: str = "qwen3-max",
algo_version: str = "v1"
algo_version: str = "halumem"
):
"""Main entry point for ReMe evaluation."""
asyncio.run(main_async(

4
reme/agent/chat/__init__.py
View file

@ -9,5 +9,5 @@ __all__ = [
"SimpleChat",
]
R.op.register("simple_chat")(SimpleChat)
R.op.register("stream_chat")(StreamChat)
R.op.register(SimpleChat)
R.op.register(StreamChat)

2
reme/agent/memory/__init__.py
View file

@ -38,4 +38,4 @@ for name in __all__:
and issubclass(agent_class, BaseMemoryAgent)
and agent_class is not BaseMemoryAgent
):
R.op.register()(agent_class)
R.op.register(agent_class)

5
reme/config/default.yaml
View file

@ -12,6 +12,11 @@ http:
timeout_keep_alive: 600
limit_concurrency: 64
flow:
test:
flow_content: TestOp()
description: "test"
llm:
default:
backend: openai

102
reme/core/application.py
View file

@ -30,7 +30,6 @@ class Application:
token_counter: dict | None = None,
**kwargs,
):
# ServiceContext
self.service_context = ServiceContext(
*args,
llm_api_key=llm_api_key,
@ -47,41 +46,68 @@ class Application:
token_counter=token_counter,
**kwargs,
)
# PromptHandler
self.prompt_handler = PromptHandler(language=self.service_context.language)
self._started: bool = False
# LLM & EmbeddingModel & VectorStore & TokenCounter
self.llm: BaseLLM | None = self.service_context.llms.get("default", None)
self.embedding_model: BaseEmbeddingModel | None = self.service_context.embedding_models.get("default", None)
self.vector_store: BaseVectorStore | None = self.service_context.vector_stores.get("default", None)
self.token_counter: BaseTokenCounter | None = self.service_context.token_counters.get("default", None)
@classmethod
async def create(
cls,
*args,
llm_api_key: str | None = None,
llm_api_base: str | None = None,
embedding_api_key: str | None = None,
embedding_api_base: str | None = None,
enable_logo: bool = True,
parser: type[PydanticConfigParser] | None = None,
llm: dict | None = None,
embedding_model: dict | None = None,
vector_store: dict | None = None,
token_counter: dict | None = None,
**kwargs,
) -> "Application":
"""Create and start an Application instance asynchronously."""
instance = cls(
*args,
llm_api_key=llm_api_key,
llm_api_base=llm_api_base,
embedding_api_key=embedding_api_key,
embedding_api_base=embedding_api_base,
enable_logo=enable_logo,
parser=parser,
llm=llm,
embedding_model=embedding_model,
vector_store=vector_store,
token_counter=token_counter,
**kwargs,
)
await instance.start()
return instance
async def start(self):
"""Start the application."""
if self._started:
return self
else:
await self.service_context.start()
self._started = True
return self
async def close(self):
"""Close the application."""
if self._started:
await self.service_context.close()
self._started = False
else:
raise RuntimeError("Application is not started")
return False
async def __aenter__(self):
"""Async context manager entry."""
return self
def __enter__(self):
"""Context manager entry."""
return self
async def close(self):
"""Close"""
return await self.service_context.close()
def close_sync(self):
"""Close synchronously"""
self.service_context.close_sync()
return await self.start()
async def __aexit__(self, exc_type=None, exc_val=None, exc_tb=None):
"""Async context manager exit."""
await self.close()
return False
def __exit__(self, exc_type=None, exc_val=None, exc_tb=None):
"""Context manager exit."""
self.close_sync()
return False
return await self.close()
async def execute_flow(self, name: str, **kwargs) -> Response:
"""Execute a flow with the given name and parameters."""
@ -104,6 +130,26 @@ class Application:
):
yield chunk
@property
def llm(self) -> BaseLLM:
"""Get the default LLM instance."""
return self.service_context.llms.get("default")
@property
def embedding_model(self) -> BaseEmbeddingModel:
"""Get the default embedding model instance."""
return self.service_context.embedding_models.get("default")
@property
def vector_store(self) -> BaseVectorStore:
"""Get the default vector store instance."""
return self.service_context.vector_stores.get("default")
@property
def token_counter(self) -> BaseTokenCounter:
"""Get the default token counter instance."""
return self.service_context.token_counters.get("default")
def run_service(self):
"""Run the configured service (HTTP, MCP, or CMD)."""
import warnings

253
reme/core/context/service_context.py
View file

@ -2,13 +2,22 @@
import os
from concurrent.futures import ThreadPoolExecutor
from typing import TYPE_CHECKING
from loguru import logger
from .base_context import BaseContext
from .registry_factory import R
from ..schema import ServiceConfig
from ..utils import MCPClient, print_logo, PydanticConfigParser, init_logger, load_env, run_coro_safely
from ..utils import load_env, MCPClient, print_logo, PydanticConfigParser, init_logger
if TYPE_CHECKING:
from ..llm import BaseLLM
from ..embedding import BaseEmbeddingModel
from ..vector_store import BaseVectorStore
from ..token_counter import BaseTokenCounter
from ..flow import BaseFlow
from ..service import BaseService
class ServiceContext(BaseContext):
@ -32,19 +41,75 @@ class ServiceContext(BaseContext):
**kwargs,
):
super().__init__()
# Set environment variables
self.service_config: ServiceConfig = self._build_service_config(
*args,
llm_api_key=llm_api_key,
llm_api_base=llm_api_base,
embedding_api_key=embedding_api_key,
embedding_api_base=embedding_api_base,
service_config=service_config,
parser=parser,
config_path=config_path,
enable_logo=enable_logo,
llm=llm,
embedding_model=embedding_model,
vector_store=vector_store,
token_counter=token_counter,
**kwargs,
)
if self.service_config.init_logger:
init_logger()
if self.service_config.enable_logo:
print_logo(service_config=self.service_config)
self.language: str = self.service_config.language
self.thread_pool: ThreadPoolExecutor = ThreadPoolExecutor(
max_workers=self.service_config.thread_pool_max_workers,
)
if self.service_config.ray_max_workers > 1:
import ray
ray.init(num_cpus=self.service_config.ray_max_workers)
self.llms: dict[str, "BaseLLM"] = {}
self.embedding_models: dict[str, "BaseEmbeddingModel"] = {}
self.token_counters: dict[str, "BaseTokenCounter"] = {}
self.vector_stores: dict[str, "BaseVectorStore"] = {}
self.flows: dict[str, "BaseFlow"] = {}
self.mcp_server_mapping: dict[str, dict] = {}
self.service: "BaseService" = R.service[self.service_config.backend](service_context=self)
self._build_flows()
def _build_service_config(
self,
*args,
llm_api_key: str | None = None,
llm_api_base: str | None = None,
embedding_api_key: str | None = None,
embedding_api_base: str | None = None,
service_config: ServiceConfig | None = None,
parser: type[PydanticConfigParser] | None = None,
config_path: str | None = None,
enable_logo: bool = True,
llm: dict | None = None,
embedding_model: dict | None = None,
vector_store: dict | None = None,
token_counter: dict | None = None,
**kwargs,
) -> ServiceConfig:
load_env()
self._update_env("REME_LLM_API_KEY", llm_api_key)
self._update_env("REME_LLM_BASE_URL", llm_api_base)
self._update_env("REME_EMBEDDING_API_KEY", embedding_api_key)
self._update_env("REME_EMBEDDING_BASE_URL", embedding_api_base)
# Use default parser if not provided
parser_class = parser if parser is not None else PydanticConfigParser
self.parser = parser_class(ServiceConfig)
# Service configuration
if service_config is None:
parser_class = parser if parser is not None else PydanticConfigParser
parser = parser_class(ServiceConfig)
input_args = []
if config_path:
input_args.append(f"config={config_path}")
@ -52,101 +117,18 @@ class ServiceContext(BaseContext):
input_args.extend(args)
if kwargs:
input_args.extend([f"{k}={v}" for k, v in kwargs.items()])
service_config = self.parser.parse_args(*input_args)
self.service_config: ServiceConfig = service_config
service_config = parser.parse_args(*input_args)
# Initialize logger
if self.service_config.init_logger:
init_logger()
# Update service config with provided arguments
service_config.enable_logo = enable_logo
if llm:
self.update_section_config("llm", **llm)
self._update_section_config(service_config, "llm", **llm)
if embedding_model:
self.update_section_config("embedding_model", **embedding_model)
self._update_section_config(service_config, "embedding_model", **embedding_model)
if token_counter:
self.update_section_config("token_counter", **token_counter)
self._update_section_config(service_config, "token_counter", **token_counter)
if vector_store:
self.update_section_config("vector_store", **vector_store)
# Print the ReMe logo if enabled in configuration.
self.service_config.enable_logo = enable_logo
if self.service_config.enable_logo:
print_logo(service_config=self.service_config)
# Service configuration and runtime settings
self.language: str = self.service_config.language
self.thread_pool: ThreadPoolExecutor = ThreadPoolExecutor(max_workers=service_config.thread_pool_max_workers)
# Initialize Ray for distributed computing if configured
if self.service_config.ray_max_workers > 1:
import ray
ray.init(num_cpus=self.service_config.ray_max_workers)
from ..llm import BaseLLM
from ..embedding import BaseEmbeddingModel
from ..vector_store import BaseVectorStore
from ..token_counter import BaseTokenCounter
from ..flow import BaseFlow, ExpressionFlow
from ..service import BaseService
# Initialize LLM instances
self.llms: dict[str, BaseLLM] = {}
for name, config in self.service_config.llm.items():
self.llms[name] = R.llm[config.backend](model_name=config.model_name, **config.model_extra)
# Initialize Embedding model instances
self.embedding_models: dict[str, BaseEmbeddingModel] = {}
for name, config in self.service_config.embedding_model.items():
self.embedding_models[name] = R.embedding_model[config.backend](
model_name=config.model_name,
**config.model_extra,
)
# Initialize Token counter instances
self.token_counters: dict[str, BaseTokenCounter] = {}
for name, config in self.service_config.token_counter.items():
self.token_counters[name] = R.token_counter[config.backend](
model_name=config.model_name,
**config.model_extra,
)
# Initialize Vector store instances
self.vector_stores: dict[str, BaseVectorStore] = {}
for name, config in self.service_config.vector_store.items():
self.vector_stores[name] = R.vector_store[config.backend](
collection_name=config.collection_name,
embedding_model=self.embedding_models[config.embedding_model],
thread_pool=self.thread_pool,
**config.model_extra,
)
run_coro_safely(self.vector_stores[name].create_collection(config.collection_name))
# Initialize flow instances
self.flows: dict[str, BaseFlow] = {}
for name, flow_cls in R.flow.items():
if not self._filter_flows(name):
continue
flow: "BaseFlow" = flow_cls(name=name, service_context=self)
self.flows[flow.name] = flow
# Initialize flow instances from service config
for name, flow_config in self.service_config.flow.items():
if not self._filter_flows(name):
continue
flow_config.name = name
flow: BaseFlow = ExpressionFlow(flow_config=flow_config, service_context=self)
self.flows[flow.name] = flow
# Initialize service instance
self.service: BaseService = R.service[self.service_config.backend](service_context=self)
# MCP server mapping: maps server_name -> {tool_name: ToolCall}
if self.service_config.mcp_servers:
self.mcp_server_mapping: dict[str, dict] = run_coro_safely(self.prepare_mcp_servers())
else:
self.mcp_server_mapping: dict[str, dict] = {}
self._update_section_config(service_config, "vector_store", **vector_store)
return service_config
@staticmethod
def _update_env(key: str, value: str | None):
@ -154,15 +136,66 @@ class ServiceContext(BaseContext):
if value:
os.environ[key] = value
def update_section_config(self, section_name: str, **kwargs):
@staticmethod
def _update_section_config(service_config: ServiceConfig, section_name: str, **kwargs):
"""Update a specific section of the service config with new values."""
section_dict: dict = getattr(self.service_config, section_name)
section_dict: dict = getattr(service_config, section_name)
if "default" not in section_dict:
raise KeyError(f"Default `{section_name}` config not found")
current_config = section_dict["default"]
section_dict["default"] = current_config.model_copy(update=kwargs, deep=True)
def _build_flows(self):
expression_flow_cls = None
for name, flow_cls in R.flow.items():
if not self._filter_flows(name):
continue
if name == "ExpressionFlow":
expression_flow_cls = flow_cls
else:
flow: "BaseFlow" = flow_cls(name=name, service_context=self)
self.flows[flow.name] = flow
if expression_flow_cls is not None:
for name, flow_config in self.service_config.flow.items():
if not self._filter_flows(name):
continue
flow_config.name = name
flow: BaseFlow = expression_flow_cls(flow_config=flow_config, service_context=self) # noqa
self.flows[flow.name] = flow
else:
logger.info("No expression flow found, please check your configuration.")
async def start(self):
"""Start the service context by initializing all configured components."""
for name, config in self.service_config.llm.items():
self.llms[name] = R.llm[config.backend](model_name=config.model_name, **config.model_extra)
for name, config in self.service_config.embedding_model.items():
self.embedding_models[name] = R.embedding_model[config.backend](
model_name=config.model_name,
**config.model_extra,
)
for name, config in self.service_config.token_counter.items():
self.token_counters[name] = R.token_counter[config.backend](
model_name=config.model_name,
**config.model_extra,
)
for name, config in self.service_config.vector_store.items():
self.vector_stores[name] = R.vector_store[config.backend](
collection_name=config.collection_name,
embedding_model=self.embedding_models[config.embedding_model],
thread_pool=self.thread_pool,
**config.model_extra,
)
await self.vector_stores[name].create_collection(config.collection_name)
if self.service_config.mcp_servers:
await self.prepare_mcp_servers()
def _filter_flows(self, name: str) -> bool:
"""Filter flows based on enabled_flows and disabled_flows configuration."""
if self.service_config.enabled_flows:
@ -177,16 +210,10 @@ class ServiceContext(BaseContext):
mcp_client = MCPClient(config={"mcpServers": self.service_config.mcp_servers})
for server_name in self.service_config.mcp_servers.keys():
try:
# Retrieve all available tool calls from this MCP server
tool_calls = await mcp_client.list_tool_calls(server_name=server_name, return_dict=False)
# Build mapping: tool_name -> ToolCall for quick lookup
self.mcp_server_mapping[server_name] = {tool_call.name: tool_call for tool_call in tool_calls}
# Log discovered tools for debugging
for tool_call in tool_calls:
logger.info(f"list_tool_calls: {server_name}@{tool_call.name} {tool_call.simple_input_dump()}")
except Exception as e:
logger.exception(f"list_tool_calls: {server_name} error: {e}")
@ -208,20 +235,6 @@ class ServiceContext(BaseContext):
self.shutdown_thread_pool()
self.shutdown_ray()
def close_sync(self):
"""Close all service components synchronously."""
for _, vector_store in self.vector_stores.items():
run_coro_safely(vector_store.close())
for _, llm in self.llms.items():
llm.close_sync()
for _, embedding_model in self.embedding_models.items():
embedding_model.close_sync()
self.shutdown_thread_pool()
self.shutdown_ray()
def shutdown_thread_pool(self, wait: bool = True):
"""Shutdown the thread pool executor."""
if self.thread_pool:

3
reme/core/flow/__init__.py
View file

@ -3,9 +3,12 @@
from .base_flow import BaseFlow
from .cmd_flow import CmdFlow
from .expression_flow import ExpressionFlow
from ..context import R
__all__ = [
"BaseFlow",
"CmdFlow",
"ExpressionFlow",
]
R.flow.register(ExpressionFlow)

4
reme/core/flow/base_flow.py
View file

@ -139,9 +139,9 @@ class BaseFlow(ABC):
def print_flow(self):
"""Log the visual structure of the flow once."""
if not self._flow_printed:
logger.info(f"[{self.__class__.__name__}] ---------- [Flow Structure] {self.name} ----------")
logger.info(f"[{self.__class__.__name__}] ---------- [Flow Structure] {self.name} [Start] ----------")
self._print_operation_tree(self.name, self.flow_op, 0)
logger.info(f"[{self.__class__.__name__}] " + "-" * 50)
logger.info(f"[{self.__class__.__name__}] ---------- [Flow Structure] {self.name} [End] ----------")
self._flow_printed = True
async def call(self, **kwargs) -> Response | asyncio.Queue:

2
reme/core/op/__init__.py
View file

@ -19,4 +19,4 @@ __all__ = [
"SequentialOp",
]
R.op.register("mcp_tool")(MCPTool)
R.op.register(MCPTool)

11
reme/core/service/http_service.py
View file

@ -2,6 +2,7 @@
import asyncio
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
import uvicorn
from fastapi import FastAPI
@ -20,7 +21,15 @@ class HttpService(BaseService):
def __init__(self, **kwargs):
"""Initialize FastAPI app with CORS and health checks."""
super().__init__(**kwargs)
self.app = FastAPI(title=self.service_config.app_name)
@asynccontextmanager
async def lifespan(_: FastAPI):
await self.service_context.start()
yield
await self.service_context.close()
self.app = FastAPI(title=self.service_config.app_name, lifespan=lifespan)
self.app.add_middleware(
CORSMiddleware,
allow_origins=["*"],

11
reme/core/service/mcp_service.py
View file

@ -1,5 +1,7 @@
"""Model Context Protocol (MCP) service implementation."""
from contextlib import asynccontextmanager
from fastmcp import FastMCP
from fastmcp.tools import FunctionTool
@ -13,7 +15,14 @@ class MCPService(BaseService):
def __init__(self, **kwargs):
"""Initialize FastMCP instance with service settings."""
super().__init__(**kwargs)
self.mcp = FastMCP(name=self.service_config.app_name)
@asynccontextmanager
async def lifespan(_: FastMCP):
await self.service_context.start()
yield {}
await self.service_context.close()
self.mcp = FastMCP(name=self.service_config.app_name, lifespan=lifespan)
def integrate_flow(self, flow: BaseFlow) -> str | None:
"""Register a non-streaming flow as an MCP tool."""

4
reme/core/utils/logger_utils.py
View file

@ -28,13 +28,13 @@ def init_logger(log_dir: str = "logs", level: str = "INFO") -> None:
retention="7 days",
compression="zip",
encoding="utf-8",
format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {file}:{line} | {message}",
format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {file}:{line} | {function} | {message}",
)
# Configure colorized standard output logging
logger.add(
sink=sys.stdout,
level=level,
format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {file}:{line} | {message}",
format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {file}:{line} | {function} | {message}",
colorize=True,
)

29
reme/core/utils/time.py
View file

@ -37,12 +37,23 @@ def timer(func: F) -> F:
file_path = "unknown"
line_no = 0
def patcher(record):
"""Modifies the log record to reflect the decorated function's location."""
record["function"] = func_name
record["file"].name = file_path.split("/")[-1]
record["file"].path = file_path
record["line"] = line_no
def create_patcher(instance):
"""Creates a patcher with runtime class information."""
# Get the actual class name at runtime if this is a method
if instance is not None and hasattr(instance, "__class__"):
class_name = instance.__class__.__name__
display_name = f"{class_name}.{func_name}"
else:
display_name = func_name
def patcher(record):
"""Modifies the log record to reflect the decorated function's location."""
record["function"] = display_name
record["file"].name = file_path.split("/")[-1]
record["file"].path = file_path
record["line"] = line_no
return patcher
@functools.wraps(func)
async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
@ -52,6 +63,9 @@ def timer(func: F) -> F:
return await func(*args, **kwargs)
finally:
duration = time.perf_counter() - start_time
# Get the instance (self) if this is a method call
instance = args[0] if args else None
patcher = create_patcher(instance)
# Use patch to inject metadata instead of relying on stack depth
logger.patch(patcher).info(
"========== cost={:.6f}s ==========",
@ -66,6 +80,9 @@ def timer(func: F) -> F:
return func(*args, **kwargs)
finally:
duration = time.perf_counter() - start_time
# Get the instance (self) if this is a method call
instance = args[0] if args else None
patcher = create_patcher(instance)
logger.patch(patcher).info(
"========== cost={:.6f}s ==========",
duration,

11
reme/core/vector_store/chroma_vector_store.py
View file

@ -231,7 +231,7 @@ class ChromaVectorStore(BaseVectorStore):
new_collection = await self._run_sync_in_executor(_create)
if collection_name == self.collection_name:
self.collection = new_collection
logger.info(f"Created collection {collection_name}")
logger.info(f"Created collection `{collection_name}`")
async def delete_collection(self, collection_name: str, **kwargs):
"""Delete a specified collection from the database."""
@ -469,15 +469,6 @@ class ChromaVectorStore(BaseVectorStore):
await self._run_sync_in_executor(_recreate)
logger.info(f"Collection {self.collection_name} has been reset")
def set_collection_name(self, collection_name: str):
"""Set the collection name and reinitialize the collection object."""
super().set_collection_name(collection_name)
self.collection = self.client.get_or_create_collection(
name=collection_name,
metadata={"hnsw:space": "cosine"},
)
logger.info(f"Collection name set to {collection_name}, collection object reinitialized")
async def close(self):
"""Close the vector store and log the shutdown process."""
logger.info(f"ChromaDB vector store for collection {self.collection_name} closed")

101
reme/core/vector_store/es_vector_store.py
View file

@ -69,37 +69,19 @@ class ESVectorStore(BaseVectorStore):
**kwargs,
)
# Store connection parameters for lazy initialization
self.hosts = hosts
self.cloud_id = cloud_id
self.api_key = api_key
self.basic_auth = basic_auth
self.verify_certs = verify_certs
self.headers = headers or {}
self._client: AsyncElasticsearch | None = None
async def _get_client(self) -> AsyncElasticsearch:
"""Create or return the existing AsyncElasticsearch client.
This lazy initialization ensures the client is created in the correct event loop.
"""
if self._client is None:
self._client = AsyncElasticsearch(
hosts=self.hosts,
cloud_id=self.cloud_id,
api_key=self.api_key,
basic_auth=self.basic_auth,
verify_certs=self.verify_certs,
headers=self.headers,
)
logger.info("AsyncElasticsearch client initialized")
return self._client
# Initialize AsyncElasticsearch client
self.client = AsyncElasticsearch(
hosts=hosts,
cloud_id=cloud_id,
api_key=api_key,
basic_auth=basic_auth,
verify_certs=verify_certs,
headers=headers or {},
)
async def list_collections(self) -> list[str]:
"""List all available index names in the Elasticsearch cluster."""
client = await self._get_client()
aliases = await client.indices.get_alias()
aliases = await self.client.indices.get_alias()
return list(aliases.keys())
async def create_collection(self, collection_name: str, **kwargs):
@ -110,9 +92,8 @@ class ESVectorStore(BaseVectorStore):
**kwargs: Settings like dimensions, similarity, shards, and replicas.
"""
collection_name = collection_name.lower()
client = await self._get_client()
if await client.indices.exists(index=collection_name):
if await self.client.indices.exists(index=collection_name):
return
dimensions = kwargs.get("dimensions", self.embedding_model.dimensions)
@ -144,8 +125,8 @@ class ESVectorStore(BaseVectorStore):
},
}
if not await client.indices.exists(index=collection_name):
await client.indices.create(index=collection_name, body=index_settings)
if not await self.client.indices.exists(index=collection_name):
await self.client.indices.create(index=collection_name, body=index_settings)
logger.info(f"Created index {collection_name} with dimensions={dimensions}")
else:
logger.info(f"Index {collection_name} already exists")
@ -158,10 +139,9 @@ class ESVectorStore(BaseVectorStore):
**kwargs: Additional parameters for the deletion request.
"""
collection_name = collection_name.lower()
client = await self._get_client()
if await client.indices.exists(index=collection_name):
await client.indices.delete(index=collection_name)
if await self.client.indices.exists(index=collection_name):
await self.client.indices.delete(index=collection_name)
logger.info(f"Deleted index {collection_name}")
else:
logger.warning(f"Index {collection_name} does not exist")
@ -174,9 +154,8 @@ class ESVectorStore(BaseVectorStore):
**kwargs: Additional parameters for the reindexing process.
"""
collection_name = collection_name.lower()
client = await self._get_client()
current_index = await client.indices.get(index=self.collection_name)
current_index = await self.client.indices.get(index=self.collection_name)
current_settings = current_index[self.collection_name]
settings_to_copy = current_settings.get("settings", {}).copy()
@ -195,7 +174,7 @@ class ESVectorStore(BaseVectorStore):
index_settings.pop(key, None)
settings_to_copy["index"] = index_settings
await client.indices.create(
await self.client.indices.create(
index=collection_name,
body={
"settings": settings_to_copy,
@ -203,7 +182,7 @@ class ESVectorStore(BaseVectorStore):
},
)
await client.reindex(
await self.client.reindex(
body={
"source": {"index": self.collection_name},
"dest": {"index": collection_name},
@ -245,8 +224,7 @@ class ESVectorStore(BaseVectorStore):
}
actions.append(action)
client = await self._get_client()
success, failed = await async_bulk(client, actions, raise_on_error=False)
success, failed = await async_bulk(self.client, actions, raise_on_error=False)
if failed:
logger.warning(f"Failed to insert {len(failed)} documents")
@ -254,7 +232,7 @@ class ESVectorStore(BaseVectorStore):
logger.info(f"Inserted {success} documents into {self.collection_name}")
if refresh:
await client.indices.refresh(index=self.collection_name)
await self.client.indices.refresh(index=self.collection_name)
async def search(
self,
@ -308,8 +286,7 @@ class ESVectorStore(BaseVectorStore):
filter_conditions.append({"term": {f"metadata.{key}": value}})
search_query["knn"]["filter"] = {"bool": {"must": filter_conditions}}
client = await self._get_client()
response = await client.search(index=self.collection_name, body=search_query)
response = await self.client.search(index=self.collection_name, body=search_query)
results = []
for hit in response["hits"]["hits"]:
@ -346,9 +323,8 @@ class ESVectorStore(BaseVectorStore):
},
)
client = await self._get_client()
success, failed = await async_bulk(
client,
self.client,
actions,
raise_on_error=False,
raise_on_exception=False,
@ -360,7 +336,7 @@ class ESVectorStore(BaseVectorStore):
logger.info(f"Deleted {success} documents from {self.collection_name}")
if refresh:
await client.indices.refresh(index=self.collection_name)
await self.client.indices.refresh(index=self.collection_name)
async def delete_all(self, **kwargs):
"""Remove all vectors from the collection.
@ -368,8 +344,7 @@ class ESVectorStore(BaseVectorStore):
Args:
**kwargs: Additional deletion parameters.
"""
client = await self._get_client()
response = await client.delete_by_query(
response = await self.client.delete_by_query(
index=self.collection_name,
body={"query": {"match_all": {}}},
)
@ -379,7 +354,7 @@ class ESVectorStore(BaseVectorStore):
refresh = kwargs.get("refresh", True)
if refresh:
await client.indices.refresh(index=self.collection_name)
await self.client.indices.refresh(index=self.collection_name)
async def update(self, nodes: VectorNode | list[VectorNode], refresh: bool = True, **kwargs):
"""Update existing documents with new content or metadata.
@ -419,9 +394,8 @@ class ESVectorStore(BaseVectorStore):
},
)
client = await self._get_client()
success, failed = await async_bulk(
client,
self.client,
actions,
raise_on_error=False,
raise_on_exception=False,
@ -433,7 +407,7 @@ class ESVectorStore(BaseVectorStore):
logger.info(f"Updated {success} documents in {self.collection_name}")
if refresh:
await client.indices.refresh(index=self.collection_name)
await self.client.indices.refresh(index=self.collection_name)
async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode]:
"""Fetch documents by their IDs from the current index.
@ -448,8 +422,7 @@ class ESVectorStore(BaseVectorStore):
if single_result:
vector_ids = [vector_ids]
client = await self._get_client()
response = await client.mget(
response = await self.client.mget(
index=self.collection_name,
body={"ids": vector_ids},
)
@ -526,8 +499,7 @@ class ESVectorStore(BaseVectorStore):
else:
query["size"] = 10000
client = await self._get_client()
response = await client.search(index=self.collection_name, body=query)
response = await self.client.search(index=self.collection_name, body=query)
results = []
for hit in response["hits"]["hits"]:
@ -542,15 +514,14 @@ class ESVectorStore(BaseVectorStore):
return results
def set_collection_name(self, collection_name: str):
"""Set the collection name and ensure it's lowercase for Elasticsearch compatibility."""
async def reset_collection(self, collection_name: str):
"""Reset collection with lowercase conversion for Elasticsearch compatibility."""
collection_name = collection_name.lower()
super().set_collection_name(collection_name)
logger.info(f"Collection name set to {collection_name} (converted to lowercase)")
self.collection_name = collection_name
await self.create_collection(collection_name)
logger.info(f"Collection reset to {collection_name}")
async def close(self):
"""Terminate the Elasticsearch client session and release resources."""
if self._client is not None:
await self._client.close()
self._client = None
logger.info("Elasticsearch client connection closed")
await self.client.close()
logger.info("Elasticsearch client connection closed")

6
reme/core/vector_store/local_vector_store.py
View file

@ -335,12 +335,6 @@ class LocalVectorStore(BaseVectorStore):
return filtered_nodes
def set_collection_name(self, collection_name: str):
"""Set the collection name and reinitialize the collection path."""
super().set_collection_name(collection_name)
self.collection_path = self.root_path / collection_name
logger.info(f"Collection name set to {collection_name}, path updated to {self.collection_path}")
async def close(self):
"""Close the vector store (no-op for local file system)."""
logger.info("Local vector store closed")

7
reme/core/vector_store/pgvector_store.py
View file

@ -606,6 +606,13 @@ class PGVectorStore(BaseVectorStore):
await self.delete_collection(self.collection_name)
await self.create_collection(self.collection_name)
async def reset_collection(self, collection_name: str):
"""Reset collection with table name validation for SQL injection prevention."""
self._validate_table_name(collection_name)
self.collection_name = collection_name
await self.create_collection(collection_name)
logger.info(f"Collection reset to {collection_name}")
async def close(self):
"""Terminate the database connection pool and release associated resources."""
if self._pool is not None:

93
reme/core/vector_store/qdrant_vector_store.py
View file

@ -86,17 +86,19 @@ class QdrantVectorStore(BaseVectorStore):
**kwargs,
)
# Store connection parameters for lazy initialization
self.host = host
self.port = port
self.path = path
self.url = url
self.api_key = api_key
self.https = https
self.grpc_port = grpc_port
self.prefer_grpc = prefer_grpc
self.client_kwargs = {k: v for k, v in kwargs.items() if k != "thread_pool"}
self._client: AsyncQdrantClient | None = None
client_kwargs = {k: v for k, v in kwargs.items() if k != "thread_pool"}
self.client = AsyncQdrantClient(
host=host,
port=port,
path=path,
url=url,
api_key=api_key,
https=https,
grpc_port=grpc_port,
prefer_grpc=prefer_grpc,
**client_kwargs,
)
self.is_local = path is not None
distance_map = {
@ -107,31 +109,9 @@ class QdrantVectorStore(BaseVectorStore):
self.distance = distance_map.get(distance.lower(), Distance.COSINE)
self.on_disk = on_disk
async def _get_client(self) -> AsyncQdrantClient:
"""Create or return the existing AsyncQdrantClient.
This lazy initialization ensures the client is created in the correct event loop.
"""
if self._client is None:
self._client = AsyncQdrantClient(
host=self.host,
port=self.port,
path=self.path,
url=self.url,
api_key=self.api_key,
https=self.https,
grpc_port=self.grpc_port,
prefer_grpc=self.prefer_grpc,
**self.client_kwargs,
)
logger.info("AsyncQdrantClient initialized")
return self._client
async def list_collections(self) -> list[str]:
"""Retrieve names of all existing collections in the Qdrant instance."""
client = await self._get_client()
collections = await client.get_collections()
collections = await self.client.get_collections()
return [collection.name for collection in collections.collections]
async def create_collection(self, collection_name: str, **kwargs: Any):
@ -150,8 +130,7 @@ class QdrantVectorStore(BaseVectorStore):
distance = kwargs.get("distance", self.distance)
on_disk = kwargs.get("on_disk", self.on_disk)
client = await self._get_client()
await client.create_collection(
await self.client.create_collection(
collection_name=collection_name,
vectors_config=VectorParams(
size=dimensions,
@ -168,11 +147,10 @@ class QdrantVectorStore(BaseVectorStore):
async def _create_payload_indexes(self, collection_name: str):
"""Create keyword indexes for common metadata fields to optimize filtering."""
common_fields = ["user_id", "agent_id", "run_id", "actor_id", "source"]
client = await self._get_client()
for field in common_fields:
try:
await client.create_payload_index(
await self.client.create_payload_index(
collection_name=collection_name,
field_name=field,
field_schema="keyword",
@ -185,18 +163,16 @@ class QdrantVectorStore(BaseVectorStore):
"""Permanently remove a collection from the Qdrant instance."""
collections = await self.list_collections()
if collection_name in collections:
client = await self._get_client()
await client.delete_collection(collection_name=collection_name)
await self.client.delete_collection(collection_name=collection_name)
logger.info(f"Deleted collection {collection_name}")
else:
logger.warning(f"Collection {collection_name} does not exist")
async def copy_collection(self, collection_name: str, **kwargs: Any):
"""Duplicate an existing collection to a new one including all data."""
client = await self._get_client()
collection_info = await client.get_collection(collection_name=self.collection_name)
collection_info = await self.client.get_collection(collection_name=self.collection_name)
await client.create_collection(
await self.client.create_collection(
collection_name=collection_name,
vectors_config=collection_info.config.params.vectors,
)
@ -205,7 +181,7 @@ class QdrantVectorStore(BaseVectorStore):
batch_size = 100
while True:
records, next_offset = await client.scroll(
records, next_offset = await self.client.scroll(
collection_name=self.collection_name,
limit=batch_size,
offset=offset,
@ -225,7 +201,7 @@ class QdrantVectorStore(BaseVectorStore):
for record in records
]
await client.upsert(
await self.client.upsert(
collection_name=collection_name,
points=points,
)
@ -268,8 +244,7 @@ class QdrantVectorStore(BaseVectorStore):
points.append(point)
wait = kwargs.get("wait", True)
client = await self._get_client()
await client.upsert(
await self.client.upsert(
collection_name=self.collection_name,
points=points,
wait=wait,
@ -358,8 +333,7 @@ class QdrantVectorStore(BaseVectorStore):
query_filter = self._create_filter(filters) if filters else None
score_threshold = kwargs.get("score_threshold", None)
client = await self._get_client()
results = await client.query_points(
results = await self.client.query_points(
collection_name=self.collection_name,
query=query_vector,
query_filter=query_filter,
@ -395,8 +369,7 @@ class QdrantVectorStore(BaseVectorStore):
point_ids.append(point_id)
wait = kwargs.get("wait", True)
client = await self._get_client()
await client.delete(
await self.client.delete(
collection_name=self.collection_name,
points_selector=PointIdsList(points=point_ids),
wait=wait,
@ -411,8 +384,7 @@ class QdrantVectorStore(BaseVectorStore):
# Delete all points by using an empty filter (matches all)
from qdrant_client.models import FilterSelector
client = await self._get_client()
await client.delete(
await self.client.delete(
collection_name=self.collection_name,
points_selector=FilterSelector(filter=Filter(must=[])),
wait=wait,
@ -452,8 +424,7 @@ class QdrantVectorStore(BaseVectorStore):
points.append(point)
wait = kwargs.get("wait", True)
client = await self._get_client()
await client.upsert(
await self.client.upsert(
collection_name=self.collection_name,
points=points,
wait=wait,
@ -475,8 +446,7 @@ class QdrantVectorStore(BaseVectorStore):
point_id = abs(hash(vector_id)) % (10**18)
point_ids.append(point_id)
client = await self._get_client()
points = await client.retrieve(
points = await self.client.retrieve(
collection_name=self.collection_name,
ids=point_ids,
with_payload=True,
@ -519,8 +489,7 @@ class QdrantVectorStore(BaseVectorStore):
# If sorting is needed, fetch more records than the limit to ensure correct sorting
fetch_limit = 10000 if sort_key else (limit or 10000)
client = await self._get_client()
records, _ = await client.scroll(
records, _ = await self.client.scroll(
collection_name=self.collection_name,
scroll_filter=scroll_filter,
limit=fetch_limit,
@ -559,7 +528,5 @@ class QdrantVectorStore(BaseVectorStore):
async def close(self):
"""Close the AsyncQdrantClient connection and release resources."""
if self._client is not None:
await self._client.close()
self._client = None
logger.info("Qdrant client connection closed")
await self.client.close()
logger.info("Qdrant client connection closed")

5
reme/reme.py
View file

@ -32,8 +32,6 @@ from .tool.memory import (
UpdateMemoryV2,
AddDraftAndReadAllProfiles,
UpdateProfile,
# DeleteProfile,
# AddProfile,
AddHistory,
ReadAllProfiles,
AddMemory,
@ -430,8 +428,7 @@ class ReMe(Application):
def main():
"""Main entry point for running ReMe from command line."""
with ReMe(*sys.argv[1:]) as app:
app.run_service()
ReMe(*sys.argv[1:]).run_service()
if __name__ == "__main__":

2
reme/tool/gallery/__init__.py
View file

@ -13,4 +13,4 @@ __all__ = [
for name in __all__:
tool_class = globals()[name]
R.op.register()(tool_class)
R.op.register(tool_class)

2
reme/tool/memory/__init__.py
View file

@ -48,4 +48,4 @@ __all__ = [
for name in __all__:
tool_class = globals()[name]
if isinstance(tool_class, type) and issubclass(tool_class, BaseMemoryTool) and tool_class is not BaseMemoryTool:
R.op.register()(tool_class)
R.op.register(tool_class)

2
reme/tool/search/__init__.py
View file

@ -13,4 +13,4 @@ __all__ = [
for name in __all__:
tool_class = globals()[name]
R.op.register()(tool_class)
R.op.register(tool_class)

9
reme/workflow/__init__.py
View file

@ -0,0 +1,9 @@
"""workflow"""
from . import test
from . import procedural_memory
__all__ = [
"test",
"procedural_memory",
]

2
reme/workflow/procedural_memory/summarizer/__init__.py
View file

@ -12,4 +12,4 @@ __all__ = ["TrajectoryPreprocess", "SuccessExtraction"]
for name in __all__:
tool_class = globals()[name]
R.op.register()(tool_class)
R.op.register(tool_class)

10
reme/workflow/test/__init__.py Normal file
View file

@ -0,0 +1,10 @@
"""test"""
from .test_op import TestOp
from ...core import R
__all__ = [
"TestOp",
]
R.op.register(TestOp)

15
reme/workflow/test/test_op.py Normal file
View file

@ -0,0 +1,15 @@
"""Test workflow operations."""
from loguru import logger
from ...core.op import BaseOp
class TestOp(BaseOp):
"""Test operation for workflow testing."""
async def execute(self):
logger.info("delete start")
# await self.vector_store.delete_all()
await self.vector_store.delete("123")
logger.info("delete end")

6
tests/test_reme.py
View file

@ -5,12 +5,14 @@ import asyncio
from reme import ReMe
from reme.core.schema import VectorNode, MemoryNode
reme = ReMe(vector_store={"collection_name": "reme"})
async def test_reme():
"""Tests ReMe memory system with personal information storage and retrieval."""
# 构建一段包含个人信息的对话
reme = ReMe(vector_store={"collection_name": "reme"})
await reme.start()
# reme = await ReMe.create(vector_store={"collection_name": "reme"})
await reme.vector_store.delete_all()
messages = [