Merge pull request #63 from agentscope-ai/merge_fl

add flowllm v2

docs/deprecated.txt

@@ -1,7 +1,9 @@
 from loguru import logger
 
-用英文注释，完善module/class/function docstring，要一句话简洁，不要变更代码
-用英文注释，完善module/class/function docstring，要一句话简洁，代码要简洁，符合pep和pylint规范
+用英文注释，完善module/class/function docstring，要一句话简洁，不要变更代码逻辑，符合pep和pylint规范，使用list而不是typing.List/Dict，不使用typing.Union
+
+看看代码有什么问题
+用英文注释，完善module/class/function docstring，要一句话简洁，代码要简洁，符合pep和pylint规范，使用list而不是typing.List，不使用typing.Union
 C0114: Missing module docstring (missing-module-docstring)
 C0115: Missing class docstring (missing-class-docstring)
 C0116: Missing function or method docstring (missing-function-docstring)

reme_ai/core/application.py

@@ -0,0 +1,208 @@
+"""Main application module for managing ReMe AI service lifecycle and flow execution."""
+
+import asyncio
+import os
+
+from .context import C
+from .flow import BaseFlow
+from .schema import ServiceConfig, Response
+from .utils import PydanticConfigParser, init_logger, execute_stream_task, run_coro_safely
+
+
+class Application:
+    """
+    Main application class for managing the lifecycle of ReMe AI services.
+
+    Handles initialization, configuration, service management, and flow execution
+    for both synchronous and asynchronous contexts.
+    """
+
+    def __init__(
+        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,
+    ):
+        """
+        Initialize the Application with configuration settings.
+
+        Args:
+            llm_api_key: API key for LLM service
+            llm_api_base: Base URL for LLM service
+            embedding_api_key: API key for embedding service
+            embedding_api_base: Base URL for embedding service
+            service_config: Pre-built service configuration object
+            parser: Custom parser class for configuration (defaults to PydanticConfigParser)
+            config_path: Path to configuration file
+            enable_logo: Whether to display the ReMe logo on startup
+            llm: LLM configuration dictionary
+            embedding_model: Embedding model configuration dictionary
+            vector_store: Vector store configuration dictionary
+            token_counter: Token counter configuration dictionary
+            **kwargs: Additional configuration arguments
+        """
+        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)
+
+        init_logger()
+
+        # Use default parser if not provided
+        parser_class = parser if parser is not None else PydanticConfigParser
+        self.parser = parser_class(ServiceConfig)
+
+        if service_config is None:
+            input_args = []
+            if config_path:
+                input_args.append(f"config={config_path}")
+            if args:
+                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)
+
+        C.service_config = service_config
+
+        if llm:
+            C.update_section_config("llm", **llm)
+        if embedding_model:
+            C.update_section_config("embedding_model", **embedding_model)
+        if vector_store:
+            C.update_section_config("vector_store", **vector_store)
+        if token_counter:
+            C.update_section_config("token_counter", **token_counter)
+        C.service_config.enable_logo = enable_logo
+        C.print_logo()
+
+    @staticmethod
+    def _update_env(key: str, value: str | None) -> None:
+        """Update environment variable if value is provided."""
+        if value:
+            os.environ[key] = value
+
+    @staticmethod
+    async def start() -> None:
+        """Initialize the service context and prepare external MCP servers."""
+        C.initialize_service_context()
+        await C.prepare_mcp_servers()
+
+    @staticmethod
+    def start_sync() -> None:
+        """Synchronous version of start()."""
+        C.initialize_service_context()
+        run_coro_safely(C.prepare_mcp_servers())
+
+    @staticmethod
+    async def stop(wait_thread_pool: bool = True, wait_ray: bool = True) -> None:
+        """
+        Stop the application and cleanup resources.
+
+        Args:
+            wait_thread_pool: Whether to wait for thread pool shutdown
+            wait_ray: Whether to wait for Ray shutdown
+        """
+        await C.close()
+        C.shutdown_thread_pool(wait=wait_thread_pool)
+        C.shutdown_ray(wait=wait_ray)
+
+    @staticmethod
+    def stop_sync(wait_thread_pool: bool = True, wait_ray: bool = True) -> None:
+        """Synchronous version of stop()."""
+        C.close_sync()
+        C.shutdown_thread_pool(wait=wait_thread_pool)
+        C.shutdown_ray(wait=wait_ray)
+
+    async def __aenter__(self):
+        """Async context manager entry."""
+        await self.start()
+        return self
+
+    def __enter__(self):
+        """Context manager entry."""
+        self.start_sync()
+        return self
+
+    async def __aexit__(self, exc_type=None, exc_val=None, exc_tb=None):
+        """Async context manager exit."""
+        await self.stop()
+        return False
+
+    def __exit__(self, exc_type=None, exc_val=None, exc_tb=None):
+        """Context manager exit."""
+        self.stop_sync()
+        return False
+
+    @staticmethod
+    async def execute_flow(name: str, **kwargs) -> Response:
+        """
+        Execute a flow asynchronously.
+
+        Args:
+            name: Name of the flow to execute
+            **kwargs: Arguments to pass to the flow
+
+        Returns:
+            Response object from the flow execution
+        """
+        flow: BaseFlow = C.get_flow(name)
+        return await flow.call(**kwargs)
+
+    @staticmethod
+    def execute_flow_sync(name: str, **kwargs) -> Response:
+        """
+        Execute a flow synchronously.
+
+        Args:
+            name: Name of the flow to execute
+            **kwargs: Arguments to pass to the flow
+
+        Returns:
+            Response object from the flow execution
+        """
+        flow: BaseFlow = C.get_flow(name)
+        return flow.call_sync(**kwargs)
+
+    @staticmethod
+    async def execute_stream_flow(name: str, **kwargs):
+        """
+        Execute a streaming flow asynchronously.
+
+        Args:
+            name: Name of the streaming flow to execute
+            **kwargs: Arguments to pass to the flow
+
+        Yields:
+            Stream chunks from the flow execution
+
+        Raises:
+            AssertionError: If the flow is not configured for streaming
+        """
+        flow: BaseFlow = C.get_flow(name)
+        assert flow.stream is True, "non-stream flow is not supported in execute_stream_flow!"
+        stream_queue = asyncio.Queue()
+        task = asyncio.create_task(flow.call(stream_queue=stream_queue, **kwargs))
+
+        async for chunk in execute_stream_task(
+            queue=stream_queue,
+            task=task,
+            flow_name=name,
+            as_bytes=False,
+        ):
+            yield chunk
+
+    @staticmethod
+    def run_service():
+        """Run the configured service (HTTP, MCP, or CMD)."""
+        C.get_service().run()

reme_ai/core/config/__init__.py

@@ -0,0 +1,5 @@
+"""config"""
+
+from .reme_config_parser import ReMeConfigParser
+
+__all__ = ["ReMeConfigParser"]

reme_ai/core/config/default.yaml

@@ -0,0 +1,34 @@
+backend: http
+thread_pool_max_workers: 64
+
+mcp:
+  transport: sse
+  host: "0.0.0.0"
+  port: 8001
+
+http:
+  host: "0.0.0.0"
+  port: 8002
+  timeout_keep_alive: 600
+  limit_concurrency: 64
+
+llm:
+  default:
+    backend: openai
+    model_name: qwen3-30b-a3b-instruct-2507
+    temperature: 0.6
+
+embedding_model:
+  default:
+    backend: openai
+    model_name: text-embedding-v4
+    dimensions: 1024
+
+vector_store:
+  default:
+    backend: local
+    embedding_model: default
+
+token_counter:
+  default:
+    backend: base

reme_ai/core/config/reme_config_parser.py

@@ -0,0 +1,7 @@
+"""Configuration parser for ReMe framework."""
+
+from ..utils import PydanticConfigParser
+
+
+class ReMeConfigParser(PydanticConfigParser):
+    """Configuration parser for ReMe framework."""

reme_ai/core/context/runtime_context.py

@@ -1,15 +1,14 @@
-"""Module providing a runtime context for managing response states and asynchronous data streaming."""
+"""Runtime context for managing response states and asynchronous data streaming."""
 
 import asyncio
 
 from .base_context import BaseContext
 from ..enumeration import ChunkEnum
-from ..schema import Response
-from ..schema import StreamChunk
+from ..schema import Response, StreamChunk
 
 
 class RuntimeContext(BaseContext):
-    """A context class for handling execution state, including response metadata and stream queues."""
+    """Context for execution state, response metadata, and stream queues."""
 
     def __init__(
         self,
@@ -17,40 +16,64 @@ class RuntimeContext(BaseContext):
         stream_queue: asyncio.Queue | None = None,
         **kwargs,
     ):
-        """Initialize the runtime context with optional response objects and message queues."""
+        """Initialize the context with optional response and queue."""
         super().__init__(**kwargs)
+        self.response = response or Response()
+        self.stream_queue = stream_queue
 
-        self.response: Response | None = response if response is not None else Response()
-        self.stream_queue: asyncio.Queue | None = stream_queue
+    @classmethod
+    def from_context(cls, context: "RuntimeContext | None" = None, **kwargs) -> "RuntimeContext":
+        """Create a new context from an existing instance or keywords."""
+        if context is None:
+            return cls(**kwargs)
 
-    async def add_stream_string_and_type(self, chunk: str, chunk_type: ChunkEnum):
-        """Create and enqueue a stream chunk from a raw string and specific type."""
-        if self.stream_queue is None:
-            return self
+        context.update(kwargs)
+        return context
 
-        # Package raw data into a StreamChunk schema
-        stream_chunk = StreamChunk(chunk_type=chunk_type, chunk=chunk)
-        await self.stream_queue.put(stream_chunk)
+    async def _enqueue(self, chunk: StreamChunk) -> None:
+        """Internal helper to put a chunk into the queue if it exists."""
+        if self.stream_queue:
+            await self.stream_queue.put(chunk)
+
+    async def add_stream_string(self, chunk: str, chunk_type: ChunkEnum) -> "RuntimeContext":
+        """Enqueue a stream chunk from a raw string and type."""
+        await self._enqueue(StreamChunk(chunk_type=chunk_type, chunk=chunk))
         return self
 
-    async def add_stream_chunk(self, stream_chunk: StreamChunk):
-        """Directly enqueue an existing stream chunk into the stream queue."""
-        if self.stream_queue is None:
-            return self
-        await self.stream_queue.put(stream_chunk)
+    async def add_stream_chunk(self, stream_chunk: StreamChunk) -> "RuntimeContext":
+        """Enqueue an existing stream chunk."""
+        await self._enqueue(stream_chunk)
         return self
 
-    async def add_stream_done(self):
-        """Enqueue a termination chunk to signal the end of the data stream."""
-        if self.stream_queue is None:
-            return self
-
-        # Create a special chunk representing the completion state
-        done_chunk = StreamChunk(chunk_type=ChunkEnum.DONE, chunk="", done=True)
-        await self.stream_queue.put(done_chunk)
+    async def add_stream_done(self) -> "RuntimeContext":
+        """Enqueue a termination chunk to signal the end of the stream."""
+        await self._enqueue(StreamChunk(chunk_type=ChunkEnum.DONE, chunk="", done=True))
         return self
 
-    def add_response_error(self, e: Exception):
-        """Update the internal response object to reflect a failure state using exception details."""
+    def add_response_error(self, e: Exception) -> "RuntimeContext":
+        """Record an exception into the response object."""
         self.response.success = False
-        self.response.answer = str(e.args)
+        self.response.answer = str(e)
+        return self
+
+    def apply_mapping(self, mapping: dict[str, str]) -> "RuntimeContext":
+        """Copy internal values based on a source-to-target key map."""
+        if not mapping:
+            return self
+
+        for source, target in mapping.items():
+            if source in self:
+                self[target] = self[source]
+        return self
+
+    def validate_required_keys(self, required_keys: dict[str, bool], context_name: str = "context") -> "RuntimeContext":
+        """Ensure all required keys are present in the context.
+
+        Args:
+            required_keys: Dictionary mapping key names to boolean indicating if required
+            context_name: Name of the context for error messages (e.g., operator name)
+        """
+        for key, is_required in required_keys.items():
+            if is_required and key not in self:
+                raise ValueError(f"{context_name}: missing required input '{key}'")
+        return self

reme_ai/core/context/service_context.py

@@ -1,34 +1,80 @@
 """Module for managing global service configurations and component registries via a singleton context."""
 
 from concurrent.futures import ThreadPoolExecutor
-from typing import Dict
+from typing import TYPE_CHECKING
+
+from loguru import logger
 
 from .base_context import BaseContext
 from .registry import Registry
 from ..enumeration import RegistryEnum
 from ..schema import ServiceConfig
-from ..utils import singleton
+from ..utils import singleton, print_logo
+
+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
 
 
 @singleton
 class ServiceContext(BaseContext):
-    """A singleton container for global application state, thread pools, and component registries."""
+    """A singleton container for global application state, thread pools, and component registries.
+
+    This class serves as the central management hub for the entire ReMe application, providing:
+    - Service configuration management
+    - Component registration and instantiation (LLMs, embeddings, vector stores, etc.)
+    - Thread pool and Ray distributed computing management
+    - MCP (Model Context Protocol) server integration
+
+    The singleton pattern ensures only one instance exists throughout the application lifecycle,
+    accessible via the global `C` variable exported at the bottom of this module.
+    """
 
     def __init__(self, **kwargs):
-        """Initialize the global context with configuration objects and specialized registries."""
+        """Initialize the global context with configuration objects and specialized registries.
+
+        Sets up:
+        - Empty service configuration placeholder
+        - Thread pool for concurrent operations
+        - Registry dictionaries for class registration (templates)
+        - Instance dictionaries for instantiated objects (actual instances)
+        - MCP server mapping for external tool integration
+        """
         super().__init__(**kwargs)
 
+        # Service configuration and runtime settings
         self.service_config: ServiceConfig | None = None
         self.language: str = ""
         self.thread_pool: ThreadPoolExecutor | None = None
-        self.vector_store_dict: Dict[str, dict] = {}
-        self.external_mcp_tool_call_dict: dict = {}
-        # Initialize a registry for every category defined in RegistryEnum
-        self.registry_dict: Dict[RegistryEnum, Registry] = {v: Registry() for v in RegistryEnum.__members__.values()}
-        self.flow_dict: dict = {}
+
+        # Registry system: stores class definitions for different component types
+        self.registry_dict: dict[RegistryEnum, Registry] = {v: Registry() for v in RegistryEnum.__members__.values()}
+
+        # Instance system: stores instantiated objects created from registered classes
+        self.instance_dict: dict[RegistryEnum, dict] = {v: {} for v in RegistryEnum.__members__.values()}
+
+        # MCP server mapping: maps server_name -> {tool_name: ToolCall}
+        self.mcp_server_mapping: dict[str, dict] = {}
 
     def register(self, name: str, register_type: RegistryEnum):
-        """Return a decorator to register a component within a specific registry category."""
+        """Return a decorator to register a component within a specific registry category.
+
+        Args:
+            name: The registration name for the component (used for lookup)
+            register_type: The type of registry (LLM, EMBEDDING_MODEL, VECTOR_STORE, etc.)
+
+        Returns:
+            A decorator function that registers the decorated class
+
+        Example:
+            @C.register("my_llm", RegistryEnum.LLM)
+            class MyLLM(BaseLLM):
+                pass
+        """
         return self.registry_dict[register_type].register(name=name)
 
     def register_llm(self, name: str = ""):
@@ -60,18 +106,29 @@ class ServiceContext(BaseContext):
         return self.register(name=name, register_type=RegistryEnum.TOKEN_COUNTER)
 
     def get_model_class(self, name: str, register_type: RegistryEnum):
-        """Retrieve a registered class by name from a specific registry category."""
+        """Retrieve a registered class by name from a specific registry category.
+
+        Args:
+            name: The registration name of the class
+            register_type: The type of registry to search in
+
+        Returns:
+            The registered class (not an instance, but the class itself)
+
+        Raises:
+            AssertionError: If the class is not found in the registry
+        """
         assert name in self.registry_dict[register_type], f"{name} not in registry_dict[{register_type}]"
         return self.registry_dict[register_type][name]
 
-    def get_embedding_model_class(self, name: str):
-        """Get the embedding model class registered under the given name."""
-        return self.get_model_class(name, RegistryEnum.EMBEDDING_MODEL)
-
     def get_llm_class(self, name: str):
         """Get the LLM class registered under the given name."""
         return self.get_model_class(name, RegistryEnum.LLM)
 
+    def get_embedding_model_class(self, name: str):
+        """Get the embedding model class registered under the given name."""
+        return self.get_model_class(name, RegistryEnum.EMBEDDING_MODEL)
+
     def get_vector_store_class(self, name: str):
         """Get the vector store class registered under the given name."""
         return self.get_model_class(name, RegistryEnum.VECTOR_STORE)
@@ -92,14 +149,378 @@ class ServiceContext(BaseContext):
         """Get the token counter class registered under the given name."""
         return self.get_model_class(name, RegistryEnum.TOKEN_COUNTER)
 
-    def get_vector_store(self, name: str):
-        """Retrieve a specific vector store instance by name."""
-        return self.vector_store_dict[name]
+    def get_llm(self, name: str) -> "BaseLLM":
+        """Retrieve a specific LLM instance by name.
 
-    def get_flow(self, name: str):
-        """Retrieve a specific flow instance by name."""
-        return self.flow_dict[name]
+        Args:
+            name: The name of the LLM instance (typically 'default' or custom name)
+
+        Returns:
+            The instantiated LLM object
+
+        Raises:
+            KeyError: If no LLM with the given name exists
+        """
+        return self.instance_dict[RegistryEnum.LLM][name]
+
+    def get_embedding_model(self, name: str) -> "BaseEmbeddingModel":
+        """Retrieve a specific embedding model instance by name.
+
+        Args:
+            name: The name of the embedding model instance (typically 'default')
+
+        Returns:
+            The instantiated embedding model object
+
+        Raises:
+            KeyError: If no embedding model with the given name exists
+        """
+        return self.instance_dict[RegistryEnum.EMBEDDING_MODEL][name]
+
+    def get_vector_store(self, name: str) -> "BaseVectorStore":
+        """Retrieve a specific vector store instance by name.
+
+        Args:
+            name: The name of the vector store instance (typically 'default')
+
+        Returns:
+            The instantiated vector store object
+
+        Raises:
+            KeyError: If no vector store with the given name exists
+        """
+        return self.instance_dict[RegistryEnum.VECTOR_STORE][name]
+
+    def get_token_counter(self, name: str) -> "BaseTokenCounter":
+        """Retrieve a specific token counter instance by name.
+
+        Args:
+            name: The name of the token counter instance (typically 'default')
+
+        Returns:
+            The instantiated token counter object
+
+        Raises:
+            KeyError: If no token counter with the given name exists
+        """
+        return self.instance_dict[RegistryEnum.TOKEN_COUNTER][name]
+
+    def get_flow(self, name: str) -> "BaseFlow":
+        """Retrieve a specific flow instance by name.
+
+        Args:
+            name: The name of the flow instance
+
+        Returns:
+            The instantiated flow object
+
+        Raises:
+            KeyError: If no flow with the given name exists
+        """
+        return self.instance_dict[RegistryEnum.FLOW][name]
+
+    def get_service(self) -> "BaseService":
+        """Retrieve the default service instance.
+
+        Returns:
+            The instantiated service backend (HTTP, MCP, or CMD service)
+
+        Raises:
+            KeyError: If the default service was not initialized
+        """
+        return self.instance_dict[RegistryEnum.SERVICE]["default"]
+
+    def update_section_config(self, section_name: str, **kwargs):
+        """Update a specific section of the service config with new values.
+
+        Args:
+            section_name: Name of the config section (e.g., 'llm', 'embedding_model')
+            **kwargs: Key-value pairs to update in the default configuration
+
+        Raises:
+            KeyError: If the default config for the section doesn't exist
+
+        Example:
+            update_section_config('llm', temperature=0.8, max_tokens=1000)
+        """
+        if not hasattr(self.service_config, section_name) or not kwargs:
+            return
+
+        section_dict: dict = getattr(self.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 initialize_service_context(self):
+        """Initialize the service context with the configuration.
+
+        This is the main initialization method that sets up all system components in order:
+        1. Language settings
+        2. Thread pool for concurrent operations
+        3. Ray cluster (if configured for distributed computing)
+        4. LLM instances
+        5. Embedding model instances
+        6. Token counter instances
+        7. Vector store instances (with their embedding models)
+        8. Flow instances (both registered and configured)
+        9. Service backend instance
+
+        Note: This method should be called after service_config is set.
+        """
+        self.language = self.service_config.language
+        self.thread_pool = ThreadPoolExecutor(max_workers=self.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)
+
+        # Initialize components in dependency order
+        self._initialize_llm()
+        self._initialize_embedding_model()
+        self._initialize_token_counter()
+        self._initialize_vector_store()  # Depends on embedding models
+        self._initialize_flow()
+        self._initialize_service()
+
+    def _initialize_llm(self):
+        """Initialize all configured LLM instances.
+
+        For each LLM configuration:
+        - Retrieves the corresponding registered LLM class by backend name
+        - Instantiates it with model_name and additional configuration
+        - Stores the instance in instance_dict for later retrieval
+        """
+        for name, config in self.service_config.llm.items():
+            llm_cls = self.get_llm_class(config.backend)
+            self.instance_dict[RegistryEnum.LLM][name] = llm_cls(model_name=config.model_name, **config.model_extra)
+
+    def _initialize_embedding_model(self):
+        """Initialize all configured embedding model instances.
+
+        For each embedding model configuration:
+        - Retrieves the corresponding registered embedding model class by backend name
+        - Instantiates it with model_name and additional configuration
+        - Stores the instance in instance_dict for later retrieval
+        """
+        for name, config in self.service_config.embedding_model.items():
+            embedding_model_cls = self.get_embedding_model_class(config.backend)
+            self.instance_dict[RegistryEnum.EMBEDDING_MODEL][name] = embedding_model_cls(
+                model_name=config.model_name,
+                **config.model_extra,
+            )
+
+    def _initialize_token_counter(self):
+        """Initialize all configured token counter instances.
+
+        For each token counter configuration:
+        - Retrieves the corresponding registered token counter class by backend name
+        - Instantiates it with model_name and additional configuration
+        - Stores the instance in instance_dict for later retrieval
+        """
+        for name, config in self.service_config.token_counter.items():
+            token_counter_cls = self.get_token_counter_class(config.backend)
+            self.instance_dict[RegistryEnum.TOKEN_COUNTER][name] = token_counter_cls(
+                model_name=config.model_name,
+                **config.model_extra,
+            )
+
+    def _initialize_vector_store(self):
+        """Initialize all configured vector stores with their embedding models.
+
+        For each vector store configuration:
+        - Retrieves the corresponding registered vector store class by backend name
+        - Retrieves the associated embedding model instance by name
+        - Instantiates the vector store with collection name, embedding model, and extra config
+        - Stores the instance in instance_dict for later retrieval
+
+        Note: This must be called after _initialize_embedding_model() since vector stores
+        depend on embedding model instances.
+        """
+        for name, config in self.service_config.vector_store.items():
+            vector_store_cls = self.get_vector_store_class(config.backend)
+            self.instance_dict[RegistryEnum.VECTOR_STORE][name] = vector_store_cls(
+                collection_name=config.collection_name,
+                embedding_model=self.instance_dict[RegistryEnum.EMBEDDING_MODEL][config.embedding_model],
+                **config.model_extra,
+            )
+
+    def _filter_flows(self, name: str) -> bool:
+        """Filter flows based on enabled_flows and disabled_flows configuration.
+
+        The filtering logic follows this priority:
+        1. If enabled_flows is set: only flows in the list are loaded
+        2. Else if disabled_flows is set: all flows except those in the list are loaded
+        3. Otherwise: all flows are loaded
+
+        Args:
+            name: The flow name to check
+
+        Returns:
+            True if the flow should be loaded, False otherwise
+        """
+        if self.service_config.enabled_flows:
+            return name in self.service_config.enabled_flows
+        elif self.service_config.disabled_flows:
+            return name not in self.service_config.disabled_flows
+        else:
+            return True
+
+    def _initialize_flow(self):
+        """Initialize all flows from both registry and configuration.
+
+        Flows can be defined in two ways:
+        1. Registered flows: Python classes decorated with @register_flow
+        2. Configuration flows: Defined in config as ExpressionFlow instances
+
+        Process:
+        1. First, instantiate all registered flow classes (from decorators)
+           - Filter based on enabled_flows/disabled_flows
+           - Create instance with the flow name
+
+        2. Then, instantiate all configured flows (from config file)
+           - Filter based on enabled_flows/disabled_flows
+           - Create ExpressionFlow instances with flow configuration
+
+        Note: Configuration flows can override registered flows with the same name.
+        """
+
+        # Initialize flows from registry (decorator-based registration)
+        for name, flow_cls in self.registry_dict[RegistryEnum.FLOW].items():
+            if not self._filter_flows(name):
+                continue
+            flow: "BaseFlow" = flow_cls(name=name)
+            self.instance_dict[RegistryEnum.FLOW][flow.name] = flow
+
+        # Initialize flows from configuration (config-based definition)
+        from ..flow import ExpressionFlow
+
+        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)
+            self.instance_dict[RegistryEnum.FLOW][name] = flow
+
+    def _initialize_service(self):
+        """Initialize the service backend instance.
+
+        Creates an instance of the configured service backend (e.g., HTTP, MCP, or CMD service)
+        and stores it in the instance dictionary under the 'default' key.
+        """
+        service_cls = self.get_service_class(self.service_config.backend)
+        self.instance_dict[RegistryEnum.SERVICE]["default"] = service_cls()
+
+    async def prepare_mcp_servers(self):
+        """Prepare and initialize MCP (Model Context Protocol) server connections.
+
+        This method:
+        1. Checks if MCP servers are configured
+        2. Creates an MCP client instance
+        3. For each configured server:
+           - Lists available tool calls from the server
+           - Builds a mapping of tool_name -> ToolCall object
+           - Logs available tools for debugging
+
+        The mcp_server_mapping is structured as:
+        {
+            "server_name": {
+                "tool_name": ToolCall(...),
+                ...
+            },
+            ...
+        }
+
+        This allows the application to discover and use external tools provided by MCP servers.
+        """
+        if not self.service_config.mcp_servers:
+            return
+
+        from ..utils import MCPClient
+
+        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}")
+
+    def print_logo(self):
+        """Print the ReMe logo if enabled in configuration."""
+        if self.service_config.enable_logo:
+            print_logo(service_config=self.service_config)
+
+    async def close(self):
+        """Close all service components asynchronously.
+
+        Gracefully closes all instantiated components in order:
+        1. Vector stores (closes database connections)
+        2. LLMs (closes API clients and connections)
+        3. Embedding models (closes API clients and connections)
+
+        This method should be called when shutting down the application
+        to ensure all resources are properly released.
+        """
+        for _, vector_store in self.instance_dict[RegistryEnum.VECTOR_STORE].items():
+            await vector_store.close()
+
+        for _, llm in self.instance_dict[RegistryEnum.LLM].items():
+            await llm.close()
+
+        for _, embedding_model in self.instance_dict[RegistryEnum.EMBEDDING_MODEL].items():
+            await embedding_model.close()
+
+    def close_sync(self):
+        """Close all service components synchronously.
+
+        Synchronous version of close() for non-async contexts.
+        Closes LLMs and embedding models without using async/await.
+
+        Note: Vector stores are not closed here as they typically require async operations.
+        """
+        for _, llm in self.instance_dict[RegistryEnum.LLM].items():
+            llm.close_sync()
+
+        for _, embedding_model in self.instance_dict[RegistryEnum.EMBEDDING_MODEL].items():
+            embedding_model.close_sync()
+
+    def shutdown_thread_pool(self, wait: bool = True):
+        """Shutdown the thread pool executor.
+
+        Args:
+            wait: If True, blocks until all pending futures are executed.
+                  If False, returns immediately and pending futures may be cancelled.
+        """
+        if self.thread_pool:
+            self.thread_pool.shutdown(wait=wait)
+
+    def shutdown_ray(self, wait: bool = True):
+        """Shutdown Ray cluster if it was initialized.
+
+        Args:
+            wait: If True, waits for Ray to fully shutdown.
+                  If False, returns immediately without waiting.
+
+        Note: Only shuts down Ray if it was configured with ray_max_workers > 1.
+        """
+        if self.service_config and self.service_config.ray_max_workers > 1:
+            import ray
+
+            ray.shutdown(_exiting_interpreter=not wait)
 
 
-# Export a global instance for easy access across the application
+# Export a global singleton instance for easy access across the application
+# This is the primary way to access the service context throughout the codebase
 C = ServiceContext()

reme_ai/core/embedding/__init__.py

@@ -0,0 +1,11 @@
+"""embedding"""
+
+from .base_embedding_model import BaseEmbeddingModel
+from .openai_embedding_model import OpenAIEmbeddingModel
+from .openai_embedding_model_sync import OpenAIEmbeddingModelSync
+
+__all__ = [
+    "BaseEmbeddingModel",
+    "OpenAIEmbeddingModel",
+    "OpenAIEmbeddingModelSync",
+]

reme_ai/core/embedding/base_embedding_model.py

@@ -0,0 +1,156 @@
+"""Base embedding model interface for ReMe.
+
+Defines the abstract base class and standard API for all embedding model implementations.
+"""
+
+import asyncio
+import time
+from abc import ABC
+
+from loguru import logger
+
+from ..schema import VectorNode
+
+
+class BaseEmbeddingModel(ABC):
+    """Abstract base class for embedding model implementations.
+
+    Provides a standard interface for text-to-vector generation with
+    built-in batching, retry logic, and error handling.
+    """
+
+    def __init__(
+        self,
+        model_name: str = "",
+        dimensions: int = 1024,
+        max_batch_size: int = 10,
+        max_retries: int = 3,
+        raise_exception: bool = True,
+        **kwargs,
+    ):
+        """Initialize model configuration and parameters."""
+        self.model_name = model_name
+        self.dimensions = dimensions
+        self.max_batch_size = max_batch_size
+        self.max_retries = max_retries
+        self.raise_exception = raise_exception
+        self.kwargs = kwargs
+
+    async def _get_embeddings(self, input_text: list[str]) -> list[list[float]]:
+        """Internal async implementation for calling the embedding API with batch input."""
+
+    def _get_embeddings_sync(self, input_text: list[str]) -> list[list[float]]:
+        """Internal synchronous implementation for calling the embedding API with batch input."""
+
+    async def get_embedding(self, input_text: str) -> list[float]:
+        """Async get embedding for a single text with exponential backoff retries."""
+        for i in range(self.max_retries):
+            try:
+                result = await self._get_embeddings([input_text])
+                return result[0]
+            except Exception as e:
+                logger.error(f"Model {self.model_name} failed: {e}")
+                if i == self.max_retries - 1:
+                    if self.raise_exception:
+                        raise
+                    return []
+                await asyncio.sleep(i + 1)
+        return []
+
+    async def get_embeddings(self, input_text: list[str]) -> list[list[float]]:
+        """Async get embeddings with automatic batching and exponential backoff retries."""
+        # Split into batches and process sequentially to respect rate limits
+        results = []
+        for i in range(0, len(input_text), self.max_batch_size):
+            batch = input_text[i : i + self.max_batch_size]
+            # Process each batch with retry logic
+            for retry in range(self.max_retries):
+                try:
+                    batch_res = await self._get_embeddings(batch)
+                    if batch_res:
+                        results.extend(batch_res)
+                    break
+                except Exception as e:
+                    logger.error(f"Model {self.model_name} batch failed: {e}")
+                    if retry == self.max_retries - 1:
+                        if self.raise_exception:
+                            raise
+                    else:
+                        await asyncio.sleep(retry + 1)
+        return results
+
+    def get_embedding_sync(self, input_text: str) -> list[float]:
+        """Synchronous get embedding for a single text with retry logic."""
+        for i in range(self.max_retries):
+            try:
+                result = self._get_embeddings_sync([input_text])
+                return result[0]
+            except Exception as exc:
+                logger.error(f"Model {self.model_name} failed: {exc}")
+                if i == self.max_retries - 1:
+                    if self.raise_exception:
+                        raise
+                    return []
+                time.sleep(i + 1)
+        return []
+
+    def get_embeddings_sync(self, input_text: list[str]) -> list[list[float]]:
+        """Synchronous get embeddings with automatic batching and retry logic."""
+        results = []
+        for i in range(0, len(input_text), self.max_batch_size):
+            batch = input_text[i : i + self.max_batch_size]
+            # Process each batch with retry logic
+            for retry in range(self.max_retries):
+                try:
+                    batch_res = self._get_embeddings_sync(batch)
+                    if batch_res:
+                        results.extend(batch_res)
+                    break
+                except Exception as exc:
+                    logger.error(f"Model {self.model_name} batch failed: {exc}")
+                    if retry == self.max_retries - 1:
+                        if self.raise_exception:
+                            raise
+                    else:
+                        time.sleep(retry + 1)
+        return results
+
+    async def get_node_embedding(self, node: VectorNode) -> VectorNode:
+        """Async generate and populate vector field for a single VectorNode object."""
+        node.vector = await self.get_embedding(node.content)
+        return node
+
+    async def get_node_embeddings(self, nodes: list[VectorNode]) -> list[VectorNode]:
+        """Async generate and populate vector fields for a batch of VectorNode objects."""
+        contents = [node.content for node in nodes]
+        embeddings: list[list[float]] = await self.get_embeddings(contents)
+
+        if len(embeddings) == len(nodes):
+            for node, vec in zip(nodes, embeddings):
+                node.vector = vec
+        else:
+            logger.warning(f"Mismatch: got {len(embeddings)} vectors for {len(nodes)} nodes")
+        return nodes
+
+    def get_node_embedding_sync(self, node: VectorNode) -> VectorNode:
+        """Synchronously generate and populate vector field for a single VectorNode object."""
+        node.vector = self.get_embedding_sync(node.content)
+        return node
+
+    def get_node_embeddings_sync(self, nodes: list[VectorNode]) -> list[VectorNode]:
+        """Synchronously generate and populate vector fields for a batch of VectorNode objects."""
+        contents = [node.content for node in nodes]
+        embeddings: list[list[float]] = self.get_embeddings_sync(contents)
+
+        if len(embeddings) == len(nodes):
+            for node, vec in zip(nodes, embeddings):
+                node.vector = vec
+        else:
+            logger.warning(f"Mismatch: got {len(embeddings)} vectors for {len(nodes)} nodes")
+        return nodes
+
+    def close_sync(self):
+        """Synchronously release resources and close connections."""
+
+    async def close(self):
+        """Asynchronously release resources and close connections."""

reme_ai/core/embedding/openai_embedding_model.py

@@ -0,0 +1,52 @@
+"""Asynchronous OpenAI-compatible embedding model implementation for ReMe."""
+
+import os
+from typing import Literal
+
+from openai import AsyncOpenAI
+
+from .base_embedding_model import BaseEmbeddingModel
+from ..context import C
+
+
+@C.register_embedding_model("openai")
+class OpenAIEmbeddingModel(BaseEmbeddingModel):
+    """Asynchronous embedding model implementation compatible with OpenAI-style APIs."""
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        base_url: str | None = None,
+        encoding_format: Literal["float", "base64"] = "float",
+        **kwargs,
+    ):
+        """Initialize the OpenAI async embedding model with API credentials and configuration."""
+        super().__init__(**kwargs)
+        self.api_key: str = api_key or os.getenv("REME_EMBEDDING_API_KEY", "")
+        self.base_url: str = base_url or os.getenv("REME_EMBEDDING_BASE_URL", "")
+        self.encoding_format: Literal["float", "base64"] = encoding_format
+
+        # Create client using factory method
+        self._client = self._create_client()
+
+    def _create_client(self):
+        """Create and return an internal AsyncOpenAI client instance."""
+        return AsyncOpenAI(api_key=self.api_key, base_url=self.base_url)
+
+    async def _get_embeddings(self, input_text: list[str]) -> list[list[float]]:
+        """Fetch embeddings from the API for a batch of strings."""
+        completion = await self._client.embeddings.create(
+            model=self.model_name,
+            input=input_text,
+            dimensions=self.dimensions,
+            encoding_format=self.encoding_format,
+        )
+
+        result_emb = [[] for _ in range(len(input_text))]
+        for emb in completion.data:
+            result_emb[emb.index] = emb.embedding
+        return result_emb
+
+    async def close(self):
+        """Close the asynchronous OpenAI client and release network resources."""
+        await self._client.close()

reme_ai/core/embedding/openai_embedding_model_sync.py

@@ -0,0 +1,33 @@
+"""Synchronous OpenAI-compatible embedding model implementation for ReMe."""
+
+from openai import OpenAI
+
+from .openai_embedding_model import OpenAIEmbeddingModel
+from ..context import C
+
+
+@C.register_embedding_model("openai_sync")
+class OpenAIEmbeddingModelSync(OpenAIEmbeddingModel):
+    """Synchronous embedding model implementation that extends the asynchronous OpenAI model."""
+
+    def _create_client(self):
+        """Create and return an internal synchronous OpenAI client instance."""
+        return OpenAI(api_key=self.api_key, base_url=self.base_url)
+
+    def _get_embeddings_sync(self, input_text: list[str]) -> list[list[float]]:
+        """Fetch embeddings synchronously from the API for a batch of strings."""
+        completion = self._client.embeddings.create(
+            model=self.model_name,
+            input=input_text,
+            dimensions=self.dimensions,
+            encoding_format=self.encoding_format,
+        )
+
+        result_emb = [[] for _ in range(len(input_text))]
+        for emb in completion.data:
+            result_emb[emb.index] = emb.embedding
+        return result_emb
+
+    def close_sync(self):
+        """Close the synchronous OpenAI client and release network resources."""
+        self._client.close()

reme_ai/core/enumeration/json_schema_enum.py

@@ -3,17 +3,16 @@
 from enum import Enum
 
 
-class JsonSchemaEnum(str, Enum):
+class JsonSchemaEnum(Enum):
     """Enumeration of valid JSON Schema data types."""
 
-    STRING = "string"
-    NUMBER = "number"
-    INTEGER = "integer"
-    OBJECT = "object"
-    ARRAY = "array"
-    BOOLEAN = "boolean"
-    NULL = "null"
+    STRING = str
+    NUMBER = float
+    INTEGER = int
+    OBJECT = dict
+    ARRAY = list
+    BOOLEAN = bool
 
     def __str__(self) -> str:
         """Returns the string representation of the enum value."""
-        return self.value
+        return self.name.lower()

reme_ai/core/flow/__init__.py

@@ -0,0 +1,11 @@
+"""flow"""
+
+from .base_flow import BaseFlow
+from .cmd_flow import CmdFlow
+from .expression_flow import ExpressionFlow
+
+__all__ = [
+    "BaseFlow",
+    "CmdFlow",
+    "ExpressionFlow",
+]

reme_ai/core/flow/base_flow.py

@@ -0,0 +1,214 @@
+"""Base flow module providing abstract flow execution with caching and operation orchestration."""
+
+import asyncio
+import hashlib
+import json
+from abc import ABC, abstractmethod
+
+from loguru import logger
+
+from ..context import C, RuntimeContext
+from ..enumeration import ChunkEnum, RegistryEnum
+from ..op import BaseOp, SequentialOp, ParallelOp
+from ..schema import Response, ToolCall, ToolAttr
+from ..utils import camel_to_snake, CacheHandler
+
+
+class BaseFlow(ABC):
+    """Abstract base class for flow execution with caching, streaming, and operation tree management.
+
+    BaseFlow provides a framework for building complex workflows by composing operations
+    into executable trees. It supports both synchronous and asynchronous execution modes,
+    response caching, streaming outputs, and automatic tool call schema generation.
+    """
+
+    def __init__(
+        self,
+        name: str = "",
+        stream: bool = False,
+        raise_exception: bool = True,
+        enable_cache: bool = False,
+        cache_path: str = "cache/flow",
+        cache_expire_hours: float = 0.1,
+        **kwargs,
+    ):
+        """Initialize flow configuration and execution state."""
+        super().__init__()
+
+        self.name: str = name or camel_to_snake(self.__class__.__name__)
+        self.stream: bool = stream
+        self.raise_exception: bool = raise_exception
+        self.enable_cache: bool = enable_cache
+        self.cache_path: str = cache_path
+        self.cache_expire_hours: float = cache_expire_hours
+        self.flow_params: dict = kwargs
+
+        self._flow_op: BaseOp | None = None
+        self._cache: CacheHandler | None = None
+        self._flow_printed: bool = False
+        self._tool_call: ToolCall | None = None
+
+    def _build_tool_call(self) -> ToolCall | None:
+        """Generate the tool call schema definition for this flow."""
+
+    @abstractmethod
+    def _build_flow(self) -> BaseOp:
+        """Construct the root operation tree for flow execution."""
+
+    def _compute_cache_key(self, params: dict) -> str | None:
+        """Generate a SHA256 hash from input parameters for caching."""
+        try:
+            payload = json.dumps(params, sort_keys=True, ensure_ascii=False, default=str)
+            return hashlib.sha256(payload.encode("utf-8")).hexdigest()
+        except Exception as e:
+            logger.exception(f"{self.name} cache key serialization failed: {e}")
+            return None
+
+    def _maybe_load_cached(self, params: dict) -> Response | None:
+        """Retrieve a cached response if caching is enabled and available."""
+        if not self.enable_cache or self.stream:
+            return None
+
+        if key := self._compute_cache_key(params):
+            if cached := self.cache.load(key):
+                logger.info(f"Loaded {self.name} response from cache.")
+                return Response(**cached)
+        return None
+
+    def _maybe_save_cache(self, params: dict, response: Response):
+        """Persist the execution response to the cache."""
+        if not self.enable_cache or self.stream:
+            return
+
+        if key := self._compute_cache_key(params):
+            self.cache.save(
+                key,
+                response.model_dump(exclude_none=True),
+                expire_hours=self.cache_expire_hours,
+            )
+
+    def _print_operation_tree(self, name: str, op: BaseOp, indent: int):
+        """Recursively log the hierarchy of the flow's operation tree."""
+        prefix = "  " * indent
+        op_type = "sequential" if isinstance(op, SequentialOp) else "parallel" if isinstance(op, ParallelOp) else name
+        logger.info(f"{prefix}{op_type} execution")
+
+        for sub_op in op.sub_ops or []:
+            self._print_operation_tree(sub_op.name, sub_op, indent + 2)
+
+    @property
+    def tool_call(self) -> ToolCall | None:
+        """Lazily construct the ToolCall schema describing this flow."""
+        if self.flow_op.tool_call:
+            return self.flow_op.tool_call
+
+        if self._tool_call is None:
+            self._tool_call = self._build_tool_call()
+            if self._tool_call:
+                self._tool_call.name = self._tool_call.name or self.name
+                self._tool_call.output = self._tool_call.output or {
+                    f"{self.name}_result": ToolAttr(
+                        type="string",
+                        description=f"The execution result of the {self.name}",
+                    ),
+                }
+        return self._tool_call
+
+    @property
+    def cache(self) -> CacheHandler:
+        """Provide access to the internal CacheHandler instance."""
+        assert self.enable_cache, "Cache usage requested while disabled."
+        if self._cache is None:
+            self._cache = CacheHandler(f"{self.cache_path}/{self.name}")
+        return self._cache
+
+    @property
+    def flow_op(self) -> BaseOp:
+        """Lazily build and retrieve the root operation of the flow."""
+        if self._flow_op is None:
+            self._flow_op = self._build_flow()
+        return self._flow_op
+
+    @property
+    def async_mode(self) -> bool:
+        """Check if the current flow operation tree is asynchronous."""
+        return self.flow_op.async_mode
+
+    @staticmethod
+    def parse_expression(expression: str) -> BaseOp:
+        """Parse a string expression into an executable BaseOp instance."""
+        lines = [x.strip() for x in expression.strip().splitlines() if x.strip()]
+        if not lines:
+            raise ValueError("Expression is empty")
+
+        env: dict = C.registry_dict[RegistryEnum.OP]
+        if len(lines) > 1:
+            exec("\n".join(lines[:-1]), {"__builtins__": {}}, env)
+
+        result = eval(lines[-1], {"__builtins__": {}}, env)
+        if not isinstance(result, BaseOp):
+            raise TypeError(f"Expression evaluated to {type(result)}, expected BaseOp")
+        return result
+
+    def print_flow(self):
+        """Log the visual structure of the flow once."""
+        if not self._flow_printed:
+            logger.info(f"---------- [Flow Structure] {self.name} ----------")
+            self._print_operation_tree(self.name, self.flow_op, 0)
+            logger.info("-" * 50)
+            self._flow_printed = True
+
+    async def call(self, **kwargs) -> Response | asyncio.Queue:
+        """Execute the flow asynchronously with parameter caching."""
+        kwargs["stream"] = self.stream
+        logger.info(f"{self.name} incoming params: {kwargs}")
+        if cached := self._maybe_load_cached(kwargs):
+            return cached
+
+        context = RuntimeContext(**kwargs)
+        try:
+            self.print_flow()
+            flow_op: BaseOp = self._build_flow()
+            assert self.flow_op.async_mode, "Async call requires an async flow operation."
+
+            await flow_op.call(context=context)
+            result = context.stream_queue if self.stream else context.response
+
+            if self.stream:
+                await context.add_stream_done()
+
+            self._maybe_save_cache(kwargs, result)
+            return result
+        except Exception as e:
+            logger.exception(f"{self.name} async call failed: {e}")
+            if self.raise_exception:
+                raise e
+            if self.stream:
+                await context.add_stream_chunk_and_type(str(e), ChunkEnum.ERROR)
+                await context.add_stream_done()
+                return context.stream_queue
+            context.add_response_error(e)
+            return context.response
+
+    def call_sync(self, **kwargs) -> Response:
+        """Execute the flow synchronously with parameter caching."""
+        logger.info(f"{self.name} incoming sync params: {kwargs}")
+        assert not self.stream, "Synchronous call cannot be used in stream mode."
+        if cached := self._maybe_load_cached(kwargs):
+            return cached
+
+        context = RuntimeContext(**kwargs)
+        try:
+            self.print_flow()
+            flow_op: BaseOp = self._build_flow()
+            assert not self.flow_op.async_mode, "Sync call requires a sync flow operation."
+
+            flow_op.call_sync(context=context)
+            self._maybe_save_cache(kwargs, context.response)
+            return context.response
+        except Exception as e:
+            logger.exception(f"{self.name} sync call failed: {e}")
+            if self.raise_exception:
+                raise e
+            context.add_response_error(e)
+            return context.response

reme_ai/core/flow/cmd_flow.py

@@ -0,0 +1,18 @@
+"""Command-based flow implementation for parsing and executing operation sequences."""
+
+from .base_flow import BaseFlow
+from ..op import BaseOp
+
+
+class CmdFlow(BaseFlow):
+    """A flow class that builds an operation chain from a string expression."""
+
+    def __init__(self, flow: str = "", **kwargs):
+        """Initialize the command flow with a string-based operation definition."""
+        super().__init__(**kwargs)
+        self.flow = flow
+        assert flow, "add `flow=<op_flow>` in cmd!"
+
+    def _build_flow(self) -> BaseOp:
+        """Parse the stored flow expression into a functional operation object."""
+        return self.parse_expression(self.flow)

reme_ai/core/flow/expression_flow.py

@@ -0,0 +1,30 @@
+"""Expression-based flow implementation driven by configuration objects."""
+
+from .base_flow import BaseFlow
+from ..op import BaseOp
+from ..schema import FlowConfig, ToolCall
+
+
+class ExpressionFlow(BaseFlow):
+    """A flow implementation that constructs operations from a FlowConfig definition."""
+
+    def __init__(self, flow_config: FlowConfig):
+        """Initialize the flow using settings and metadata from a FlowConfig instance."""
+        self.flow_config: FlowConfig = flow_config
+        super().__init__(
+            name=flow_config.name,
+            stream=self.flow_config.stream,
+            raise_exception=self.flow_config.raise_exception,
+            enable_cache=self.flow_config.enable_cache,
+            cache_path=self.flow_config.cache_path,
+            cache_expire_hours=self.flow_config.cache_expire_hours,
+            **flow_config.model_extra,
+        )
+
+    def _build_flow(self) -> BaseOp:
+        """Generate the operation chain by parsing the flow content string."""
+        return self.parse_expression(self.flow_config.flow_content)
+
+    def _build_tool_call(self) -> ToolCall:
+        """Construct a tool call representation based on configuration parameters."""
+        return ToolCall(**{"description": self.flow_config.description, "parameters": self.flow_config.parameters})

reme_ai/core/llm/base_llm.py

@@ -4,7 +4,7 @@ import asyncio
 import json
 import time
 from abc import ABC
-from typing import List, Callable, Generator, AsyncGenerator, Any, Optional, Dict
+from typing import Callable, Generator, AsyncGenerator, Optional, Any
 
 from loguru import logger
 
@@ -74,7 +74,7 @@ class BaseLLM(ABC):
     @staticmethod
     def _accumulate_tool_call_chunk(
         tool_call,
-        ret_tools: List[ToolCall],
+        ret_tools: list[ToolCall],
     ) -> None:
         """Assemble incremental tool call fragments into complete ToolCall objects."""
         index = tool_call.index
@@ -95,15 +95,15 @@ class BaseLLM(ABC):
 
     @staticmethod
     def _validate_and_serialize_tools(
-        ret_tools: List[ToolCall],
-        tools: Optional[List[ToolCall]],
-    ) -> List[Dict]:
+        ret_tools: list[ToolCall],
+        tools: Optional[list[ToolCall]],
+    ) -> list[dict]:
         """Validate tool call integrity and return serialized tool dictionaries."""
         if not ret_tools:
             return []
 
         # Create lookup dict for tool validation
-        tool_dict: Dict[str, ToolCall] = {x.name: x for x in tools} if tools else {}
+        tool_dict: dict[str, ToolCall] = {x.name: x for x in tools} if tools else {}
         validated_tools = []
 
         for tool in ret_tools:
@@ -123,8 +123,8 @@ class BaseLLM(ABC):
 
     def _build_stream_kwargs(
         self,
-        messages: List[Message],
-        tools: Optional[List[ToolCall]] = None,
+        messages: list[Message],
+        tools: Optional[list[ToolCall]] = None,
         log_params: bool = True,
         **kwargs,
     ) -> dict:
@@ -133,8 +133,8 @@ class BaseLLM(ABC):
 
     async def _stream_chat(
         self,
-        messages: List[Message],
-        tools: Optional[List[ToolCall]] = None,
+        messages: list[Message],
+        tools: Optional[list[ToolCall]] = None,
         stream_kwargs: Optional[dict] = None,
     ) -> AsyncGenerator[StreamChunk, None]:
         """Internal async generator for streaming raw response chunks."""
@@ -142,8 +142,8 @@ class BaseLLM(ABC):
 
     def _stream_chat_sync(
         self,
-        messages: List[Message],
-        tools: Optional[List[ToolCall]] = None,
+        messages: list[Message],
+        tools: Optional[list[ToolCall]] = None,
         stream_kwargs: Optional[dict] = None,
     ) -> Generator[StreamChunk, None, None]:
         """Internal synchronous generator for streaming raw response chunks."""
@@ -152,8 +152,8 @@ class BaseLLM(ABC):
     async def _stream_with_retry(
         self,
         operation_name: str,
-        messages: List[Message],
-        tools: Optional[List[ToolCall]],
+        messages: list[Message],
+        tools: Optional[list[ToolCall]],
         stream_kwargs: dict,
     ) -> AsyncGenerator[StreamChunk, None]:
         """Execute the async streaming operation with retry logic and error recovery."""
@@ -178,8 +178,8 @@ class BaseLLM(ABC):
     def _stream_with_retry_sync(
         self,
         operation_name: str,
-        messages: List[Message],
-        tools: Optional[List[ToolCall]],
+        messages: list[Message],
+        tools: Optional[list[ToolCall]],
         stream_kwargs: dict,
     ) -> Generator[StreamChunk, None, None]:
         """Execute the synchronous streaming operation with retry logic and error recovery."""
@@ -202,8 +202,8 @@ class BaseLLM(ABC):
 
     async def stream_chat(
         self,
-        messages: List[Message],
-        tools: Optional[List[ToolCall]] = None,
+        messages: list[Message],
+        tools: Optional[list[ToolCall]] = None,
         **kwargs,
     ) -> AsyncGenerator[StreamChunk, None]:
         """Public async interface for streaming chat completions with retries."""
@@ -213,8 +213,8 @@ class BaseLLM(ABC):
 
     def stream_chat_sync(
         self,
-        messages: List[Message],
-        tools: Optional[List[ToolCall]] = None,
+        messages: list[Message],
+        tools: Optional[list[ToolCall]] = None,
         **kwargs,
     ) -> Generator[StreamChunk, None, None]:
         """Public synchronous interface for streaming chat completions with retries."""
@@ -223,8 +223,8 @@ class BaseLLM(ABC):
 
     async def _chat(
         self,
-        messages: List[Message],
-        tools: Optional[List[ToolCall]] = None,
+        messages: list[Message],
+        tools: Optional[list[ToolCall]] = None,
         enable_stream_print: bool = False,
         **kwargs,
     ) -> Message:
@@ -245,8 +245,8 @@ class BaseLLM(ABC):
 
     def _chat_sync(
         self,
-        messages: List[Message],
-        tools: Optional[List[ToolCall]] = None,
+        messages: list[Message],
+        tools: Optional[list[ToolCall]] = None,
         enable_stream_print: bool = False,
         **kwargs,
     ) -> Message:
@@ -315,8 +315,8 @@ class BaseLLM(ABC):
 
     async def chat(
         self,
-        messages: List[Message],
-        tools: Optional[List[ToolCall]] = None,
+        messages: list[Message],
+        tools: Optional[list[ToolCall]] = None,
         enable_stream_print: bool = False,
         callback_fn: Optional[Callable[[Message], Any]] = None,
         default_value: Any = None,
@@ -337,8 +337,8 @@ class BaseLLM(ABC):
 
     def chat_sync(
         self,
-        messages: List[Message],
-        tools: Optional[List[ToolCall]] = None,
+        messages: list[Message],
+        tools: Optional[list[ToolCall]] = None,
         enable_stream_print: bool = False,
         callback_fn: Optional[Callable[[Message], Any]] = None,
         default_value: Any = None,

reme_ai/core/llm/lite_llm_sync.py

@@ -1,11 +1,6 @@
-"""Synchronous LiteLLM-based LLM implementation for the ReMe framework.
+"""Synchronous LiteLLM-based LLM implementation for the ReMe framework."""
 
-This module provides a unified synchronous interface for 100+ LLM providers via LiteLLM,
-supporting streaming completions, tool calling, and reasoning content. For
-asynchronous operations, refer to the LiteLLM class in the lite_llm module.
-"""
-
-from typing import List, Generator, Optional
+from typing import Generator
 
 import litellm
 
@@ -19,54 +14,21 @@ from ..schema import ToolCall
 
 @C.register_llm("litellm_sync")
 class LiteLLMSync(LiteLLM):
-    """
-    Synchronous LiteLLM client for executing chat completions and streaming responses.
-
-    This class extends the base LiteLLM implementation to provide synchronous
-    execution of streaming methods, inheriting initialization and configuration
-    logic from the parent class.
-
-    Example:
-        >>> llm = LiteLLMSync(
-        ...     model_name="qwen3-max",
-        ...     api_key="sk-...",
-        ...     temperature=0.7
-        ... )
-        >>> messages = [Message(role=Role.USER, content="Hello!")]
-        >>> for chunk in llm.chat(messages):
-        ...     print(chunk)
-    """
+    """Synchronous LiteLLM client for executing chat completions and streaming responses."""
 
     def _stream_chat_sync(
         self,
-        messages: List[Message],
-        tools: Optional[List[ToolCall]] = None,
-        stream_kwargs: Optional[dict] = None,
+        messages: list[Message],
+        tools: list[ToolCall] | None = None,
+        stream_kwargs: dict | None = None,
     ) -> Generator[StreamChunk, None, None]:
-        """
-        Internal synchronous generator for processing streaming chat completion chunks.
-
-        This method orchestrates the LiteLLM completion lifecycle by categorizing
-        raw API chunks into usage data, reasoning content (thinking), regular
-        text responses, and aggregated tool calls.
-
-        Args:
-            messages: List of conversation messages to send to the model.
-            tools: Optional list of tool definitions available for the model to call.
-            stream_kwargs: Dictionary of pre-built parameters for the LiteLLM API.
-
-        Yields:
-            StreamChunk: Wrapped response fragments categorized by ChunkEnum.
-
-        Raises:
-            ValueError: If tool call arguments fail validation or serialization.
-        """
+        """Internal synchronous generator for processing streaming chat completion chunks."""
         # Create streaming completion request using LiteLLM
         stream_kwargs = stream_kwargs or {}
         completion = litellm.completion(**stream_kwargs)
 
         # Track accumulated tool calls across chunks
-        ret_tools: List[ToolCall] = []
+        ret_tools: list[ToolCall] = []
         # Flag to track if we've started receiving answer content
         is_answering: bool = False
 

reme_ai/core/llm/openai_llm.py

@@ -1,7 +1,7 @@
 """Asynchronous OpenAI-compatible LLM implementation supporting streaming, tool calls, and reasoning content."""
 
 import os
-from typing import List, AsyncGenerator, Optional
+from typing import AsyncGenerator, Optional
 
 from loguru import logger
 from openai import AsyncOpenAI
@@ -38,8 +38,8 @@ class OpenAILLM(BaseLLM):
 
     def _build_stream_kwargs(
         self,
-        messages: List[Message],
-        tools: Optional[List[ToolCall]] = None,
+        messages: list[Message],
+        tools: Optional[list[ToolCall]] = None,
         log_params: bool = True,
         **kwargs,
     ) -> dict:
@@ -68,8 +68,8 @@ class OpenAILLM(BaseLLM):
 
     async def _stream_chat(
         self,
-        messages: List[Message],
-        tools: Optional[List[ToolCall]] = None,
+        messages: list[Message],
+        tools: Optional[list[ToolCall]] = None,
         stream_kwargs: Optional[dict] = None,
     ) -> AsyncGenerator[StreamChunk, None]:
         """Generate a stream of chat completion chunks including text, reasoning content, and tool calls."""
@@ -78,7 +78,7 @@ class OpenAILLM(BaseLLM):
         completion = await self._client.chat.completions.create(**stream_kwargs)
 
         # Track accumulated tool calls across chunks
-        ret_tools: List[ToolCall] = []
+        ret_tools: list[ToolCall] = []
         # Flag to track if we've started receiving answer content
         is_answering: bool = False
 

reme_ai/core/llm/openai_llm_sync.py

@@ -1,6 +1,6 @@
 """Synchronous OpenAI-compatible LLM implementation supporting streaming, tool calls, and reasoning content."""
 
-from typing import List, Generator, Optional
+from typing import Generator, Optional
 
 from openai import OpenAI
 
@@ -22,8 +22,8 @@ class OpenAILLMSync(OpenAILLM):
 
     def _stream_chat_sync(
         self,
-        messages: List[Message],
-        tools: Optional[List[ToolCall]] = None,
+        messages: list[Message],
+        tools: Optional[list[ToolCall]] = None,
         stream_kwargs: Optional[dict] = None,
     ) -> Generator[StreamChunk, None, None]:
         """Synchronously generate a stream of chat completion chunks including text, reasoning, and tool calls."""
@@ -32,7 +32,7 @@ class OpenAILLMSync(OpenAILLM):
         completion = self._client.chat.completions.create(**stream_kwargs)
 
         # Track accumulated tool calls across chunks
-        ret_tools: List[ToolCall] = []
+        ret_tools: list[ToolCall] = []
         # Flag to track if we've started receiving answer content
         is_answering: bool = False
 

reme_ai/core/op/__init__.py

@@ -0,0 +1,13 @@
+"""op"""
+
+from .base_op import BaseOp
+from .base_ray_op import BaseRayOp
+from .parallel_op import ParallelOp
+from .sequential_op import SequentialOp
+
+__all__ = [
+    "BaseOp",
+    "BaseRayOp",
+    "ParallelOp",
+    "SequentialOp",
+]

reme_ai/core/op/base_op.py

@@ -0,0 +1,363 @@
+"""Base operator class for LLM workflow execution and composition."""
+
+import asyncio
+import copy
+import inspect
+from pathlib import Path
+from typing import Callable, Any, Optional
+
+from loguru import logger
+from tqdm import tqdm
+
+from ..context import RuntimeContext, PromptHandler, C
+from ..embedding import BaseEmbeddingModel
+from ..llm import BaseLLM
+from ..schema import ToolCall, ToolAttr, Response
+from ..token_counter import BaseTokenCounter
+from ..utils import camel_to_snake, CacheHandler, timer
+from ..vector_store import BaseVectorStore
+
+
+class BaseOp:
+    """Base operator class for LLM workflow execution and composition."""
+
+    def __new__(cls, *args, **kwargs):
+        """Capture initialization arguments for object cloning."""
+        instance = super().__new__(cls)
+        instance._init_args = copy.copy(args)
+        instance._init_kwargs = copy.copy(kwargs)
+        return instance
+
+    def __init__(
+        self,
+        name: str = "",
+        async_mode: bool = True,
+        language: str = "",
+        prompt_name: str = "",
+        llm: str | BaseLLM = "default",
+        embedding_model: str | BaseEmbeddingModel = "default",
+        vector_store: str | BaseVectorStore = "default",
+        token_counter: str | BaseTokenCounter = "default",
+        enable_cache: bool = False,
+        cache_path: str = "cache/op",
+        sub_ops: dict[str, "BaseOp"] | list["BaseOp"] | Optional["BaseOp"] = None,
+        input_mapping: dict[str, str] | None = None,
+        output_mapping: dict[str, str] | None = None,
+        save_response_result: bool = False,
+        enable_sync_thread_pool: bool = True,
+        max_retries: int = 1,
+        raise_exception: bool = False,
+        **kwargs,
+    ):
+        """Initialize operator configurations and internal state."""
+        self.name = name or camel_to_snake(self.__class__.__name__)
+        self.async_mode = async_mode
+        self.language = language or C.language
+        self.prompt = self._get_prompt_handler(prompt_name)
+
+        self._llm = llm
+        self._embedding_model = embedding_model
+        self._vector_store = vector_store
+        self._token_counter = token_counter
+
+        self.enable_cache = enable_cache
+        self.cache_path = cache_path
+        self.sub_ops: list[BaseOp] = []
+        self.add_sub_ops(sub_ops)
+
+        self.input_mapping = input_mapping
+        self.output_mapping = output_mapping
+        self.save_response_result = save_response_result
+        self.enable_sync_thread_pool = enable_sync_thread_pool
+        self.max_retries = max(1, max_retries)
+        self.raise_exception = raise_exception
+        self.op_params = kwargs
+
+        self._pending_tasks: list = []
+        self.context: RuntimeContext | None = None
+        self._cache: CacheHandler | None = None
+        self._tool_call: ToolCall | None = None
+
+    def _get_prompt_handler(self, prompt_name: str) -> PromptHandler:
+        """Load prompt configuration from the associated YAML file."""
+        path = Path(inspect.getfile(self.__class__))
+        path = path.with_stem(prompt_name) if prompt_name else path
+        return PromptHandler(language=self.language).load_prompt_by_file(path.with_suffix(".yaml"))
+
+    def _build_tool_call(self) -> ToolCall | None:
+        """Build and return the tool call schema; override in subclasses."""
+
+    def _validate_inputs(self):
+        """Ensure all required tool inputs are present in context."""
+        if self.tool_call is not None:
+            parameters = self.tool_call.parameters
+            if parameters.type == "object" and parameters.properties:
+                required_list = parameters.required or []
+                required_keys = {k: (k in required_list) for k in parameters.properties.keys()}
+                self.context.validate_required_keys(required_keys, self.name)
+
+    def _handle_failure(self, e: Exception, attempt: int):
+        """Log failures and handle final retry logic."""
+        message = f"{self.name} failed (attempt {attempt + 1}): {e}"
+        if attempt == self.max_retries - 1:
+            logger.exception(message)
+            if self.raise_exception:
+                raise e
+
+            if self.tool_call is not None:
+                self.output = f"{self.name} failed: {e}"
+        else:
+            logger.warning(message)
+
+    @property
+    def tool_call(self) -> ToolCall | None:
+        """Lazily construct and return the tool call metadata."""
+        if self._tool_call is None:
+            self._tool_call = self._build_tool_call()
+            if self._tool_call is None:
+                return None
+
+            self._tool_call.name = self._tool_call.name or self.name
+            if not self._tool_call.output.properties:
+                self._tool_call.output.properties = {
+                    f"{self.name}_result": ToolAttr(type="string", description=f"Execution result of {self.name}"),
+                }
+        return self._tool_call
+
+    def set_tool_call(self, tool_call: ToolCall | dict):
+        """Set the tool call."""
+        if isinstance(tool_call, dict):
+            self._tool_call = ToolCall(**tool_call)
+        elif isinstance(tool_call, ToolCall):
+            self._tool_call = tool_call
+        else:
+            raise ValueError(f"Invalid tool call: {tool_call}")
+
+        self._tool_call.name = self._tool_call.name or self.name
+        if not self._tool_call.output.properties:
+            self._tool_call.output.properties = {
+                f"{self.name}_result": ToolAttr(type="string", description=f"Execution result of {self.name}"),
+            }
+
+    @property
+    def input_dict(self) -> dict:
+        """Extract required and optional inputs from context based on schema."""
+        parameters = self.tool_call.parameters
+        if parameters.type != "object" or not parameters.properties:
+            return {}
+        required_keys = set(parameters.required or [])
+        return {k: self.context[k] for k in parameters.properties.keys() if (k in required_keys or k in self.context)}
+
+    @property
+    def output(self) -> Any:
+        """Get the single output value from context."""
+        output_properties = self.tool_call.output.properties
+        if not output_properties:
+            return None
+
+        keys = list(output_properties.keys())
+        return self.context[keys[0]]
+
+    @output.setter
+    def output(self, value: Any):
+        """Set the single output value into context."""
+        output_properties = self.tool_call.output.properties
+        if not output_properties:
+            return
+
+        keys = list(output_properties.keys())
+        self.context[keys[0]] = value
+
+    @property
+    def cache(self) -> CacheHandler:
+        """Access the operator-specific cache handler."""
+        assert self.enable_cache, "Cache is disabled!"
+        if not self._cache:
+            self._cache = CacheHandler(f"{self.cache_path}/{self.name}")
+        return self._cache
+
+    @property
+    def llm(self) -> BaseLLM:
+        """Get the LLM instance from ServiceContext."""
+        if isinstance(self._llm, str):
+            self._llm = C.get_llm(self._llm)
+        return self._llm
+
+    @property
+    def embedding_model(self) -> BaseEmbeddingModel:
+        """Get the embedding model instance from ServiceContext."""
+        if isinstance(self._embedding_model, str):
+            self._embedding_model = C.get_embedding_model(self._embedding_model)
+        return self._embedding_model
+
+    @property
+    def vector_store(self) -> BaseVectorStore:
+        """Lazily initialize and return the vector store instance."""
+        if isinstance(self._vector_store, str):
+            self._vector_store = C.get_vector_store(self._vector_store)
+        return self._vector_store
+
+    @property
+    def token_counter(self) -> BaseTokenCounter:
+        """Get the token counter instance from ServiceContext."""
+        if isinstance(self._token_counter, str):
+            self._token_counter = C.get_token_counter(self._token_counter)
+        return self._token_counter
+
+    @property
+    def service_metadata(self) -> dict:
+        """Get service configuration metadata."""
+        return C.service_config.model_extra
+
+    @property
+    def response(self) -> Response:
+        """Get the response object."""
+        return self.context.response
+
+    def before_execute_sync(self):
+        """Prepare context and validate before sync execution."""
+        self.context.apply_mapping(self.input_mapping)
+        self._validate_inputs()
+
+    def execute_sync(self):
+        """Define core sync logic in subclasses."""
+
+    def after_execute_sync(self):
+        """Finalize context and mappings after sync execution."""
+        self.context.apply_mapping(self.output_mapping)
+        if self.tool_call is not None and self.save_response_result:
+            self.context.response.answer = self.output
+
+    async def before_execute(self):
+        """Prepare context and validate before async execution."""
+        self.before_execute_sync()
+
+    async def execute(self):
+        """Define core async logic in subclasses."""
+
+    async def after_execute(self):
+        """Finalize context and mappings after async execution."""
+        self.after_execute_sync()
+
+    @timer
+    def call_sync(self, context: RuntimeContext = None, **kwargs):
+        """Execute the operator synchronously with retry logic."""
+        self.context = RuntimeContext.from_context(context, **kwargs)
+        for i in range(self.max_retries):
+            try:
+                self.before_execute_sync()
+                self.execute_sync()
+                self.after_execute_sync()
+                break
+            except Exception as e:
+                self._handle_failure(e, i)
+        return self.output if self.tool_call is not None else None
+
+    async def call(self, context: RuntimeContext = None, **kwargs):
+        """Execute the operator asynchronously with retry logic."""
+        self.context = RuntimeContext.from_context(context, **kwargs)
+        for i in range(self.max_retries):
+            try:
+                await self.before_execute()
+                await self.execute()
+                await self.after_execute()
+                break
+            except Exception as e:
+                self._handle_failure(e, i)
+        return self.output if self.tool_call is not None else None
+
+    def submit_sync_task(self, fn: Callable, *args, **kwargs) -> "BaseOp":
+        """Submit a task to the thread pool or local queue."""
+        task = C.thread_pool.submit(fn, *args, **kwargs) if self.enable_sync_thread_pool else (fn, args, kwargs)
+        self._pending_tasks.append(task)
+        return self
+
+    def submit_async_task(self, coro_fn: Callable, *args, **kwargs) -> "BaseOp":
+        """Submit an async task to the pending tasks queue."""
+        task = coro_fn(*args, **kwargs)
+        self._pending_tasks.append(task)
+        return self
+
+    def join_sync_tasks(self, task_desc: str = None) -> list:
+        """Wait for all pending sync tasks and return flattened results."""
+        results = []
+        for task in tqdm(self._pending_tasks, desc=task_desc or self.name):
+            res = task.result() if self.enable_sync_thread_pool else task[0](*task[1], **task[2])
+            if res:
+                results.extend(res if isinstance(res, list) else [res])
+        self._pending_tasks.clear()
+        return results
+
+    async def join_async_tasks(self, return_exceptions: bool = True) -> list:
+        """Wait for all pending async tasks and aggregate results."""
+        try:
+            raw_results = await asyncio.gather(*self._pending_tasks, return_exceptions=return_exceptions)
+            results = []
+            for res in raw_results:
+                if isinstance(res, Exception):
+                    logger.error(f"Async task failed: {res}")
+                    continue
+                if res:
+                    results.extend(res if isinstance(res, list) else [res])
+            return results
+        finally:
+            self._pending_tasks.clear()
+
+    def add_sub_ops(self, sub_ops: dict[str, "BaseOp"] | list["BaseOp"] | Optional["BaseOp"]):
+        """Add child operators to this operator's sub_ops."""
+        if not sub_ops:
+            return
+
+        if isinstance(sub_ops, dict):
+            for name, op in sub_ops.items():
+                assert self.async_mode == op.async_mode, "Async mode mismatch!"
+                op.name = name
+                self.sub_ops.append(op)
+        elif isinstance(sub_ops, list):
+            for op in sub_ops:
+                assert self.async_mode == op.async_mode, "Async mode mismatch!"
+                self.sub_ops.append(op)
+        else:
+            assert self.async_mode == sub_ops.async_mode, "Async mode mismatch!"
+            self.sub_ops.append(sub_ops)
+
+    def add_sub_op(self, sub_op: "BaseOp"):
+        """Add a single child operator to this operator's sub_ops."""
+        self.sub_ops.append(sub_op)
+
+    def __lshift__(self, ops):
+        """Operator overload for adding sub-operators."""
+        self.add_sub_ops(ops)
+        return self
+
+    def __rshift__(self, op: "BaseOp"):
+        """Operator overload for sequential execution composition."""
+        from .sequential_op import SequentialOp
+
+        seq = SequentialOp(sub_ops=[self], async_mode=self.async_mode)
+        seq.add_sub_ops(op.sub_ops if isinstance(op, SequentialOp) else op)
+        return seq
+
+    def __or__(self, op: "BaseOp"):
+        """Operator overload for parallel execution composition."""
+        from .parallel_op import ParallelOp
+
+        par = ParallelOp(sub_ops=[self], async_mode=self.async_mode)
+        par.add_sub_ops(op.sub_ops if isinstance(op, ParallelOp) else op)
+        return par
+
+    def prompt_format(self, prompt_name: str, **kwargs) -> str:
+        """Format a prompt template with provided keyword arguments."""
+        return self.prompt.prompt_format(prompt_name=prompt_name, **kwargs)
+
+    def get_prompt(self, prompt_name: str) -> str:
+        """Get a prompt template by name."""
+        return self.prompt.get_prompt(prompt_name=prompt_name)
+
+    def copy(self, **kwargs):
+        """Create a copy of this operator with optional parameter overrides."""
+        copy_op = self.__class__(*self._init_args, **{**self._init_kwargs, **kwargs})
+        if self.sub_ops:
+            copy_op.sub_ops.clear()
+            copy_op.add_sub_ops(self.sub_ops)
+        return copy_op

reme_ai/core/op/base_ray_op.py

@@ -0,0 +1,124 @@
+"""Base class for Ray-based parallel operations."""
+
+from abc import ABCMeta
+from typing import Callable
+
+import pandas as pd
+from loguru import logger
+from tqdm import tqdm
+
+from .base_op import BaseOp
+from ..context import BaseContext, C
+
+_RAY_IMPORT_ERROR = None
+
+try:
+    import ray
+except ImportError as e:
+    _RAY_IMPORT_ERROR = e
+    ray = None
+
+
+class BaseRayOp(BaseOp, metaclass=ABCMeta):
+    """Base class for Ray-based parallel operations."""
+
+    def __init__(self, **kwargs):
+        if _RAY_IMPORT_ERROR:
+            raise ImportError("Ray requires extra dependencies. Install with `pip install ray`")
+
+        super().__init__(**kwargs)
+        self._ray_task_list: list = []
+
+    def submit_and_join_parallel_op(self, op: BaseOp, **kwargs) -> list:
+        """Submit a BaseOp to be executed in parallel via Ray."""
+        return self.submit_and_join_ray_task(fn=op.call, task_desc=op.name, context=self.context, **kwargs)
+
+    def submit_and_join_ray_task(self, fn: Callable, parallel_key: str = "", task_desc: str = "", **kwargs) -> list:
+        """Divide data into chunks and execute them across Ray workers."""
+        max_workers = C.service_config.ray_max_workers
+        self._ray_task_list.clear()
+
+        # Automatically detect the key containing the list to parallelize
+        if not parallel_key:
+            for key, value in kwargs.items():
+                if isinstance(value, list):
+                    parallel_key = key
+                    break
+
+        if not parallel_key:
+            raise ValueError("No list found in kwargs to parallelize over.")
+
+        parallel_list = kwargs.pop(parallel_key)
+        logger.info(f"Parallelizing '{parallel_key}' across {max_workers} workers")
+
+        # Put large shared objects into the Ray Object Store once
+        optimized_kwargs = {
+            k: (ray.put(v) if isinstance(v, (pd.DataFrame, pd.Series, dict, list, BaseContext)) else v)
+            for k, v in kwargs.items()
+        }
+
+        # Submit sliced chunks to reduce inter-node data transfer
+        remote_task_loop = ray.remote(self._ray_task_loop)
+        for i in range(max_workers):
+            chunk = parallel_list[i::max_workers]
+            if not chunk:
+                continue
+
+            task = remote_task_loop.remote(
+                fn,
+                parallel_key,
+                chunk,
+                i,
+                **optimized_kwargs,
+            )
+            self._ray_task_list.append(task)
+            logger.info(f"Submitted task {i + 1}/{max_workers} for {task_desc}")
+
+        return self.join_ray_task(task_desc=task_desc)
+
+    @staticmethod
+    def _ray_task_loop(internal_fn: Callable, parallel_key: str, chunk: list, actor_index: int, **kwargs) -> list:
+        """Execute the function over a specific chunk of data on a worker."""
+        results = []
+        for value in chunk:
+            current_kwargs = {**kwargs, "actor_index": actor_index, parallel_key: value}
+            t_result = internal_fn(**current_kwargs)
+
+            if t_result is not None:
+                if isinstance(t_result, list):
+                    results.extend(t_result)
+                else:
+                    results.append(t_result)
+        return results
+
+    def submit_ray_task(self, fn, *args, **kwargs):
+        """Submit a single Ray task to the task list for later execution."""
+        if not ray.is_initialized():
+            ray.init(num_cpus=C.service_config.ray_max_workers, ignore_reinit_error=True)
+
+        remote_fn = ray.remote(fn)
+        task = remote_fn.remote(*args, **kwargs)
+        self._ray_task_list.append(task)
+        return self
+
+    def join_ray_task(self, task_desc: str | None = None) -> list:
+        """Collect results from Ray workers using a progress bar."""
+        results = []
+        unfinished = list(self._ray_task_list)
+
+        with tqdm(total=len(unfinished), desc=task_desc or f"{self.name}_ray") as pbar:
+            while unfinished:
+                ready, unfinished = ray.wait(unfinished, num_returns=1)
+                for obj_ref in ready:
+                    try:
+                        t_result = ray.get(obj_ref)
+                        if isinstance(t_result, list):
+                            results.extend(t_result)
+                        elif t_result is not None:
+                            results.append(t_result)
+                    except Exception as e:
+                        logger.error(f"Worker task failed: {e}")
+                    pbar.update(1)
+
+        self._ray_task_list.clear()
+        return results

reme_ai/core/op/parallel_op.py

@@ -0,0 +1,33 @@
+"""Module providing the ParallelOp class for concurrent operation execution."""
+
+from .base_op import BaseOp
+
+
+class ParallelOp(BaseOp):
+    """Operation class that executes multiple sub-operations in parallel."""
+
+    async def execute(self):
+        """Executes all sub-operations concurrently using asynchronous tasks."""
+        for op in self.sub_ops:
+            assert op.async_mode
+            self.submit_async_task(op.call, context=self.context)
+        await self.join_async_tasks()
+
+    def execute_sync(self):
+        """Executes all sub-operations concurrently using synchronous task management."""
+        for op in self.sub_ops:
+            assert not op.async_mode
+            self.submit_sync_task(op.call_sync, context=self.context)
+        self.join_sync_tasks()
+
+    def __lshift__(self, op: dict[str, BaseOp] | list[BaseOp] | BaseOp):
+        """Raises RuntimeError as the shift operator is not supported for parallel operations."""
+        raise RuntimeError(f"`<<` is not supported in `{self.name}`")
+
+    def __or__(self, op: BaseOp):
+        """Adds sub-operations to the current parallel group using the bitwise OR operator."""
+        if isinstance(op, ParallelOp) and op.sub_ops:
+            self.add_sub_ops(op.sub_ops)
+        else:
+            self.add_sub_op(op)
+        return self

reme_ai/core/op/sequential_op.py

@@ -0,0 +1,31 @@
+"""Module providing the SequentialOp class for serial operation execution."""
+
+from .base_op import BaseOp
+
+
+class SequentialOp(BaseOp):
+    """Operation class that executes sub-operations one after another in order."""
+
+    async def execute(self):
+        """Executes sub-operations sequentially using asynchronous awaits."""
+        for op in self.sub_ops:
+            assert op.async_mode
+            await op.call(context=self.context)
+
+    def execute_sync(self):
+        """Executes sub-operations sequentially in a synchronous blocking manner."""
+        for op in self.sub_ops:
+            assert not op.async_mode
+            op.call_sync(context=self.context)
+
+    def __lshift__(self, op: dict[str, BaseOp] | list[BaseOp] | BaseOp):
+        """Raises RuntimeError as the left shift operator is not supported."""
+        raise RuntimeError(f"`<<` is not supported in `{self.name}`")
+
+    def __rshift__(self, op: BaseOp):
+        """Appends operations to the sequence using the bitwise right shift operator."""
+        if isinstance(op, SequentialOp) and op.sub_ops:
+            self.add_sub_ops(op.sub_ops)
+        else:
+            self.add_sub_op(op)
+        return self

reme_ai/core/reme.py

@@ -0,0 +1,92 @@
+"""ReMe application classes for simplified configuration and execution."""
+
+import sys
+
+from .application import Application
+from .config import ReMeConfigParser
+from .context import C
+
+
+class ReMe(Application):
+    """Simplified ReMe application that auto-initializes the service context."""
+
+    def __init__(
+        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,
+        enable_logo: bool = True,
+        llm: dict | None = None,
+        embedding_model: dict | None = None,
+        vector_store: dict | None = None,
+        token_counter: dict | None = None,
+        **kwargs,
+    ):
+        super().__init__(
+            *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=None,
+            parser=ReMeConfigParser,
+            config_path=None,
+            enable_logo=enable_logo,
+            llm=llm,
+            embedding_model=embedding_model,
+            vector_store=vector_store,
+            token_counter=token_counter,
+            **kwargs,
+        )
+
+        C.initialize_service_context()
+
+    async def summary(self):
+        """Execute summary operations."""
+
+    async def retrieve(self):
+        """Execute retrieve operations."""
+
+
+class ReMeApp(Application):
+    """ReMe application with config file support and flow execution methods."""
+
+    def __init__(
+        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,
+        config_path: str | None = None,
+        enable_logo: bool = True,
+        **kwargs,
+    ):
+        super().__init__(
+            *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=None,
+            parser=ReMeConfigParser,
+            config_path=config_path,
+            enable_logo=enable_logo,
+            **kwargs,
+        )
+
+    async def async_execute(self, name: str, **kwargs) -> dict:
+        """Execute a flow asynchronously and return the result as a dictionary."""
+        return (await self.execute_flow(name=name, **kwargs)).model_dump()
+
+
+def main():
+    """Main entry point for running ReMe application from command line."""
+    with ReMeApp(*sys.argv[1:]) as app:
+        app.run_service()
+
+
+if __name__ == "__main__":
+    main()

reme_ai/core/schema/request.py

@@ -1,17 +1,11 @@
 """Defines the data structure for processing incoming user requests and message history."""
 
-from typing import List
-
 from pydantic import Field, BaseModel, ConfigDict
 
-from .message import Message
-
 
 class Request(BaseModel):
     """Represents a structured request payload containing a query, message list, and metadata."""
 
     model_config = ConfigDict(extra="allow")
 
-    query: str = Field(default="")
-    messages: List[Message] = Field(default_factory=list)
     metadata: dict = Field(default_factory=dict)

reme_ai/core/schema/service_config.py

@@ -1,5 +1,6 @@
 """Configuration schemas for service components using Pydantic models."""
 
+import os
 from typing import Dict, List
 
 from pydantic import BaseModel, Field, ConfigDict
@@ -12,7 +13,7 @@ class MCPConfig(BaseModel):
 
     model_config = ConfigDict(extra="allow")
 
-    transport: str = Field(default="")
+    transport: str = Field(default="stdio")
     host: str = Field(default="0.0.0.0")
     port: int = Field(default=8001)
 
@@ -92,16 +93,14 @@ class ServiceConfig(BaseModel):
     model_config = ConfigDict(extra="allow")
 
     backend: str = Field(default="")
+    app_name: str = Field(default=os.getenv("APP_NAME", "ReMe"))
     enable_logo: bool = Field(default=True)
     language: str = Field(default="")
     thread_pool_max_workers: int = Field(default=16)
     ray_max_workers: int = Field(default=-1)
     disabled_flows: List[str] = Field(default_factory=list)
     enabled_flows: List[str] = Field(default_factory=list)
-    external_mcp: Dict[str, dict] = Field(
-        default_factory=dict,
-        description="External MCP Server configuration",
-    )
+    mcp_servers: Dict[str, dict] = Field(default_factory=dict, description="External MCP Server configuration")
 
     mcp: MCPConfig = Field(default_factory=MCPConfig)
     http: HttpConfig = Field(default_factory=HttpConfig)

reme_ai/core/schema/tool_call.py

@@ -16,7 +16,7 @@ class ToolAttr(BaseModel):
 
     model_config = ConfigDict(extra="allow")
 
-    type: str = Field(default=JsonSchemaEnum.STRING.value, description="The data type of the attribute")
+    type: str = Field(default=str(JsonSchemaEnum.STRING), description="The data type of the attribute")
     description: Optional[str] = Field(default=None, description="Description of the attribute")
     required: Optional[List[str]] = Field(default=None, description="Required property names for object types")
     properties: Optional[Dict[str, "ToolAttr"]] = Field(default=None, description="Child properties for objects")
@@ -27,7 +27,7 @@ class ToolAttr(BaseModel):
     @classmethod
     def validate_type_is_valid_enum(cls, v: str) -> str:
         """Validates that the provided type string exists within JsonSchemaEnum values."""
-        valid_types = [e.value for e in JsonSchemaEnum]
+        valid_types = [str(e) for e in JsonSchemaEnum]
 
         if v not in valid_types:
             raise ValueError(f"Invalid type: '{v}'. Must be one of {valid_types}")

reme_ai/core/service/__init__.py

@@ -0,0 +1,13 @@
+"""service"""
+
+from .base_service import BaseService
+from .cmd_service import CmdService
+from .http_service import HttpService
+from .mcp_service import MCPService
+
+__all__ = [
+    "BaseService",
+    "CmdService",
+    "HttpService",
+    "MCPService",
+]

reme_ai/core/service/base_service.py

@@ -0,0 +1,41 @@
+"""Base service definitions for flow management."""
+
+from abc import ABC, abstractmethod
+
+from loguru import logger
+from pydantic import BaseModel
+
+from ..context import C
+from ..flow import BaseFlow
+from ..schema import ToolCall
+from ..utils import create_pydantic_model
+
+
+class BaseService(ABC):
+    """Abstract base class for services that integrate and execute flows."""
+
+    def __init__(self, **kwargs):
+        """Initialize the base service."""
+        self.kwargs = kwargs
+
+    @abstractmethod
+    def integrate_flow(self, flow: BaseFlow) -> str | None:
+        """Integrate a flow into the service and return its name if successful."""
+
+    @staticmethod
+    def _prepare_route(flow: BaseFlow) -> tuple[ToolCall, type[BaseModel]]:
+        """Generate the request model and route name for a flow."""
+        tool_call = flow.tool_call
+        model = create_pydantic_model(tool_call.name, tool_call.parameters)
+        return tool_call, model
+
+    def run(self):
+        """Initialize and integrate all flows registered in the global context."""
+        flow_names: list[str] = []
+        for _, flow in C.flow_dict.items():
+            flow_name = self.integrate_flow(flow)
+            if flow_name:
+                flow_names.append(flow_name)
+
+        if flow_names:
+            logger.info(f"integrate {','.join(flow_names)}")

reme_ai/core/service/cmd_service.py

@@ -0,0 +1,36 @@
+"""Command service module for managing and executing command-based workflows."""
+
+from loguru import logger
+
+from .base_service import BaseService
+from ..context import C
+from ..flow import CmdFlow, BaseFlow
+from ..utils.common_utils import run_coro_safely
+
+
+@C.register_service("cmd")
+class CmdService(BaseService):
+    """Service implementation for handling command flow execution logic."""
+
+    def __init__(self, **kwargs):
+        """Initialize the command service instance."""
+        super().__init__(**kwargs)
+        self._cmd_flow: CmdFlow | None = None
+
+    def integrate_flow(self, flow: BaseFlow) -> str | None:
+        """Integrate the workflow configuration into the command service."""
+        self._cmd_flow = CmdFlow(flow=C.service_config.flow)
+
+    def run(self):
+        """Execute the command flow in either asynchronous or synchronous mode."""
+        super().run()
+
+        if self._cmd_flow.async_mode:
+            response = run_coro_safely(
+                self._cmd_flow.call(**C.service_config.cmd.model_extra),
+            )
+        else:
+            response = self._cmd_flow.call_sync(**C.service_config.cmd.model_extra)
+
+        if response.answer:
+            logger.info(f"response.answer={response.answer}")

reme_ai/core/service/http_service.py

@@ -0,0 +1,87 @@
+"""HTTP service implementation using FastAPI."""
+
+import asyncio
+from collections.abc import AsyncGenerator
+
+import uvicorn
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse
+
+from .base_service import BaseService
+from ..context import C
+from ..flow import BaseFlow
+from ..schema import Response
+from ..utils.common_utils import execute_stream_task
+
+
+@C.register_service("http")
+class HttpService(BaseService):
+    """Expose flows via HTTP REST and SSE endpoints."""
+
+    def __init__(self, **kwargs):
+        """Initialize FastAPI app with CORS and health checks."""
+        super().__init__(**kwargs)
+        self.app = FastAPI(title=C.service_config.app_name)
+        self.app.add_middleware(
+            CORSMiddleware,
+            allow_origins=["*"],
+            allow_credentials=True,
+            allow_methods=["*"],
+            allow_headers=["*"],
+        )
+        self.app.get("/health")(lambda: {"status": "healthy"})
+
+    def _integrate_flow(self, flow: BaseFlow) -> str:
+        """Register a standard flow as a POST endpoint."""
+        tool_call, request_model = self._prepare_route(flow)
+
+        async def execute_endpoint(request: request_model) -> Response:
+            return await flow.call(**request.model_dump(exclude_none=True))
+
+        self.app.post(
+            path=f"/{tool_call.name}",
+            response_model=Response,
+            description=tool_call.description,
+        )(execute_endpoint)
+        return tool_call.name
+
+    def _integrate_stream_flow(self, flow: BaseFlow) -> str:
+        """Register a streaming flow as an SSE endpoint."""
+        tool_call, request_model = self._prepare_route(flow)
+
+        async def execute_stream_endpoint(request: request_model) -> StreamingResponse:
+            queue = asyncio.Queue()
+            # Start flow as a background task
+            task = asyncio.create_task(flow.call(stream_queue=queue, **request.model_dump(exclude_none=True)))
+
+            async def generate_stream() -> AsyncGenerator[bytes, None]:
+                async for chunk in execute_stream_task(
+                    queue=queue,
+                    task=task,
+                    flow_name=tool_call.name,
+                    as_bytes=True,
+                ):
+                    yield chunk
+
+            return StreamingResponse(generate_stream(), media_type="text/event-stream")
+
+        self.app.post(f"/{tool_call.name}")(execute_stream_endpoint)
+        return tool_call.name
+
+    def integrate_flow(self, flow: BaseFlow) -> str | None:
+        """Register a flow based on its streaming configuration."""
+        return self._integrate_stream_flow(flow) if flow.stream else self._integrate_flow(flow)
+
+    def run(self):
+        """Start the Uvicorn server."""
+        super().run()
+        cfg = C.service_config.http
+        uvicorn.run(
+            self.app,
+            host=cfg.host,
+            port=cfg.port,
+            timeout_keep_alive=cfg.timeout_keep_alive,
+            limit_concurrency=cfg.limit_concurrency,
+            **cfg.model_extra,
+        )

reme_ai/core/service/mcp_service.py

@@ -0,0 +1,56 @@
+"""Model Context Protocol (MCP) service implementation."""
+
+from typing import Any
+
+from fastmcp import FastMCP
+from fastmcp.tools import FunctionTool
+
+from .base_service import BaseService
+from ..context import C
+from ..flow import BaseFlow
+
+
+@C.register_service("mcp")
+class MCPService(BaseService):
+    """Expose flows as Model Context Protocol (MCP) tools."""
+
+    def __init__(self, **kwargs: Any):
+        """Initialize FastMCP instance with service settings."""
+        super().__init__(**kwargs)
+        self.mcp = FastMCP(name=C.service_config.app_name)
+
+    def integrate_flow(self, flow: BaseFlow) -> str | None:
+        """Register a non-streaming flow as an MCP tool."""
+        if flow.stream:
+            return None
+
+        tool_call, request_model = self._prepare_route(flow)
+
+        async def execute_tool(**kwargs):
+            """Execute flow logic and return the string answer."""
+            request_instance = request_model(**kwargs)
+            response = await flow.call(**request_instance.model_dump(exclude_none=True))
+            return response.answer
+
+        self.mcp.add_tool(
+            FunctionTool(
+                name=tool_call.name,
+                description=tool_call.description,
+                fn=execute_tool,
+                parameters=tool_call.parameters.simple_input_dump(),
+            ),
+        )
+        return tool_call.name
+
+    def run(self):
+        """Run the MCP server with specified transport protocol."""
+        super().run()
+        cfg = C.service_config.mcp
+
+        run_args: dict = {"transport": cfg.transport, "show_banner": False, **cfg.model_extra}
+
+        # Add network settings for non-stdio transports
+        if cfg.transport != "stdio":
+            run_args.update({"host": cfg.host, "port": cfg.port})
+
+        self.mcp.run(**run_args)

reme_ai/core/token_counter/__init__.py

@@ -0,0 +1,11 @@
+"""token counter"""
+
+from .base_token_counter import BaseTokenCounter
+from .hf_token_counter import HFTokenCounter
+from .openai_token_counter import OpenAITokenCounter
+
+__all__ = [
+    "BaseTokenCounter",
+    "HFTokenCounter",
+    "OpenAITokenCounter",
+]

reme_ai/core/token_counter/base_token_counter.py

@@ -0,0 +1,59 @@
+"""Token counting utility based on character-type rules."""
+
+import math
+import re
+from loguru import logger
+
+from ..context import C
+from ..schema import Message, ToolCall
+
+
+@C.register_token_counter("base")
+class BaseTokenCounter:
+    """A rule-based token counter for Chinese and non-Chinese text."""
+
+    def __init__(self, model_name: str, **kwargs):
+        """Initialize with model name and additional parameters."""
+        self.model_name = model_name
+        self.kwargs = kwargs
+        # Matches Chinese characters including extensions
+        self._cn_regex = re.compile(r"[\u4e00-\u9fff]")
+
+    def _count_chars(self, text: str) -> tuple[int, int]:
+        """Count Chinese and other characters in a string."""
+        if not text:
+            return 0, 0
+        cn_count = len(self._cn_regex.findall(text))
+        return cn_count, len(text) - cn_count
+
+    def count_token(
+        self,
+        messages: list[Message],
+        tools: list[ToolCall] | None = None,
+        **_kwargs,
+    ) -> int:
+        """Calculate total tokens using the 1:2 (CN) and 1:4 (Other) rule."""
+        cn_total = 0
+        ot_total = 0
+        logger.info("Calculating tokens using rule-based estimation.")
+
+        # Extract text from messages
+        segments = []
+        for msg in messages:
+            content = msg.content
+            if isinstance(content, bytes):
+                content = content.decode("utf-8", errors="ignore")
+            segments.extend([content, msg.reasoning_content])
+
+        # Extract text from tools
+        if tools:
+            for tool in tools:
+                segments.extend([tool.name, tool.description, tool.arguments])
+
+        # Process all segments
+        for text in filter(None, segments):
+            cn_chars, ot_chars = self._count_chars(text)
+            cn_total += cn_chars
+            ot_total += ot_chars
+
+        return math.ceil(cn_total / 2) + math.ceil(ot_total / 4)

reme_ai/core/token_counter/hf_token_counter.py

@@ -0,0 +1,82 @@
+"""HuggingFace token counting utilities."""
+
+import os
+
+from loguru import logger
+
+from .base_token_counter import BaseTokenCounter
+from ..context import C
+from ..schema import Message, ToolCall
+
+
+@C.register_token_counter("hf")
+class HFTokenCounter(BaseTokenCounter):
+    """Token counter using transformers.AutoTokenizer.apply_chat_template."""
+
+    def __init__(
+        self,
+        model_name: str,
+        use_fast: bool = False,
+        trust_remote_code: bool = False,
+        use_mirror: bool = True,
+        **kwargs,
+    ):
+        """Initialize the counter with model config and lazy tokenizer loading."""
+        super().__init__(model_name=model_name, **kwargs)
+        self.use_fast = use_fast
+        self.trust_remote_code = trust_remote_code
+        self.use_mirror = use_mirror
+        self._tokenizer = None
+
+    def _ensure_tokenizer(self):
+        """Initialize and cache the HuggingFace tokenizer safely."""
+        if self._tokenizer:
+            return self._tokenizer
+
+        if self.use_mirror:
+            os.environ.setdefault("HF_ENDPOINT", "https://hf-mirror.com")
+
+        try:
+            from transformers import AutoTokenizer
+
+            logger.info("Initializing HuggingFace tokenizer for {}", self.model_name)
+
+            tokenizer = AutoTokenizer.from_pretrained(
+                self.model_name,
+                use_fast=self.use_fast,
+                trust_remote_code=self.trust_remote_code,
+                **self.kwargs,
+            )
+
+            if not hasattr(tokenizer, "chat_template") or tokenizer.chat_template is None:
+                raise ValueError(f"Model {self.model_name} lacks a chat template.")
+
+            self._tokenizer = tokenizer
+            return tokenizer
+        except Exception as e:
+            logger.error("Failed to load tokenizer {}: {}", self.model_name, e)
+            raise
+
+    def count_token(
+        self,
+        messages: list[Message],
+        tools: list[ToolCall] | None = None,
+        **kwargs,
+    ) -> int:
+        """Calculate total tokens for messages and tools using the chat template."""
+        tokenizer = self._ensure_tokenizer()
+
+        # Serialize inputs for the template
+        formatted_msgs = [m.simple_dump() for m in messages]
+        formatted_tools = [t.simple_input_dump() for t in tools] if tools else None
+
+        # Setting tokenize=True and leaving return_tensors=None returns a List[int]
+        tokens = tokenizer.apply_chat_template(
+            formatted_msgs,
+            tools=formatted_tools,
+            add_generation_prompt=kwargs.pop("add_generation_prompt", False),
+            tokenize=True,
+            **kwargs,
+        )
+
+        return len(tokens)

reme_ai/core/token_counter/openai_token_counter.py

@@ -0,0 +1,58 @@
+"""Token counting implementation for OpenAI-compatible models."""
+
+import json
+from loguru import logger
+from .base_token_counter import BaseTokenCounter
+from ..context import C
+from ..schema import Message, ToolCall
+
+
+@C.register_token_counter("openai")
+class OpenAITokenCounter(BaseTokenCounter):
+    """Token counter for OpenAI models using tiktoken."""
+
+    def __init__(self, model_name: str, **kwargs):
+        super().__init__(model_name, **kwargs)
+        self._encoding = None
+
+    @property
+    def encoding(self):
+        """Get or initialize the tiktoken encoding for the specified model."""
+        if self._encoding is None:
+            import tiktoken
+
+            try:
+                self._encoding = tiktoken.encoding_for_model(self.model_name)
+            except KeyError:
+                logger.warning(f"Model {self.model_name} not found; falling back to o200k_base.")
+                self._encoding = tiktoken.get_encoding("o200k_base")
+        return self._encoding
+
+    def count_token(
+        self,
+        messages: list[Message],
+        tools: list[ToolCall] | None = None,
+        **_kwargs,
+    ) -> int:
+        """Calculate total tokens for a request including messages and tool definitions."""
+        enc = self.encoding
+        total_tokens = 0
+
+        for msg in messages:
+            # Every message has <|start|>{role/name}\n{content}<|end|>\n
+            total_tokens += 3  # Base overhead per message
+            if msg.content:
+                total_tokens += len(enc.encode(msg.content))
+
+            if msg.tool_calls:
+                for tc in msg.tool_calls:
+                    dump = json.dumps(tc.simple_output_dump(), ensure_ascii=False)
+                    total_tokens += len(enc.encode(dump))
+
+        if tools:
+            # Account for tool/function definitions if provided
+            tool_json = json.dumps([t.simple_input_dump() for t in tools], ensure_ascii=False)
+            total_tokens += len(enc.encode(tool_json))
+
+        total_tokens += 3  # Every reply is primed with <|start|>assistant<|message|>
+        return total_tokens

reme_ai/core/tool/__init__.py

@@ -0,0 +1,7 @@
+"""tool"""
+
+from .mcp_tool import MCPTool
+
+__all__ = [
+    "MCPTool",
+]

reme_ai/core/tool/mcp_tool.py

@@ -0,0 +1,82 @@
+"""MCP (Model Context Protocol) tool integration for remote tool execution."""
+
+from typing import List
+
+from ..context import C
+from ..op import BaseOp
+from ..schema import ToolCall
+from ..utils import MCPClient
+
+
+@C.register_op()
+class MCPTool(BaseOp):
+    """Operator for calling remote MCP (Model Context Protocol) tools.
+
+    This class enables integration with external MCP servers to execute tools
+    and retrieve their results. It supports parameter customization and retry logic.
+    """
+
+    def __init__(
+        self,
+        mcp_server: str = "",
+        tool_name: str = "",
+        save_response_result: bool = True,
+        parameter_required: List[str] | None = None,
+        parameter_optional: List[str] | None = None,
+        parameter_deleted: List[str] | None = None,
+        max_retries: int = 3,
+        timeout: float | None = None,
+        raise_exception: bool = False,
+        **kwargs,
+    ):
+
+        super().__init__(
+            save_response_result=save_response_result,
+            max_retries=max_retries,
+            raise_exception=raise_exception,
+            **kwargs,
+        )
+
+        self.mcp_server: str = mcp_server
+        self.tool_name: str = tool_name
+        self.parameter_required: List[str] | None = parameter_required
+        self.parameter_optional: List[str] | None = parameter_optional
+        self.parameter_deleted: List[str] | None = parameter_deleted
+        self.timeout: float | None = timeout
+        # Example MCP marketplace: https://bailian.console.aliyun.com/?tab=mcp#/mcp-market
+
+        self._client = MCPClient(C.service_config.mcp_servers)
+
+    def _build_tool_call(self) -> ToolCall:
+        tool_call_dict = C.mcp_server_mapping[self.mcp_server]
+        tool_call: ToolCall = tool_call_dict[self.tool_name].model_copy(deep=True)
+
+        # Initialize required list if not exists
+        if tool_call.parameters.required is None:
+            tool_call.parameters.required = []
+
+        if self.parameter_required:
+            for name in self.parameter_required:
+                if name not in tool_call.parameters.required:
+                    tool_call.parameters.required.append(name)
+
+        if self.parameter_optional:
+            for name in self.parameter_optional:
+                if name in tool_call.parameters.required:
+                    tool_call.parameters.required.remove(name)
+
+        if self.parameter_deleted:
+            for name in self.parameter_deleted:
+                tool_call.parameters.properties.pop(name, None)
+                if tool_call.parameters.required and name in tool_call.parameters.required:
+                    tool_call.parameters.required.remove(name)
+
+        return tool_call
+
+    async def execute(self):
+        self.output = await self._client.call_tool(
+            server_name=self.mcp_server,
+            tool_name=self.tool_name,
+            arguments=self.input_dict,
+            parse_text_result=True,
+        )

reme_ai/core/utils/__init__.py

@@ -1,14 +1,34 @@
 """utils"""
 
+from .cache_handler import CacheHandler
 from .case_converter import snake_to_camel, camel_to_snake
+from .common_utils import run_coro_safely, execute_stream_task
 from .env_utils import load_env
+from .http_client import HttpClient
+from .llm_utils import extract_content, format_messages
+from .logger_utils import init_logger
+from .logo_utils import print_logo
+from .mcp_client import MCPClient
+from .pydantic_config_parser import PydanticConfigParser
+from .pydantic_utils import create_pydantic_model
 from .singleton import singleton
 from .timer import timer
 
 __all__ = [
+    "CacheHandler",
     "snake_to_camel",
     "camel_to_snake",
+    "run_coro_safely",
+    "execute_stream_task",
     "load_env",
+    "HttpClient",
+    "extract_content",
+    "format_messages",
+    "init_logger",
+    "print_logo",
+    "MCPClient",
+    "PydanticConfigParser",
+    "create_pydantic_model",
     "singleton",
     "timer",
 ]

reme_ai/core/utils/cache_handler.py

@@ -0,0 +1,184 @@
+"""Local file-based cache utility for DataFrames, lists, dicts, and strings."""
+
+import json
+from datetime import datetime, timedelta
+from pathlib import Path
+from typing import Any
+
+import pandas as pd
+from loguru import logger
+
+
+class CacheHandler:
+    """Handles persistent data caching with expiration and type support."""
+
+    _EXTENSIONS = {
+        pd.DataFrame: ".csv",
+        dict: ".json",
+        list: ".json",
+        str: ".txt",
+    }
+
+    _TYPE_NAMES = {
+        "DataFrame": pd.DataFrame,
+        "dict": dict,
+        "list": list,
+        "str": str,
+    }
+
+    def __init__(self, cache_dir: str | Path = "cache"):
+        """Initialize cache directory and load existing metadata."""
+        self.cache_dir = Path(cache_dir)
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+        self.metadata_file = self.cache_dir / "metadata.json"
+        self.metadata: dict[str, Any] = self._load_metadata()
+
+    def set_cache_dir(self, cache_dir: str | Path) -> None:
+        """Change the cache directory and reload metadata."""
+        self.cache_dir = Path(cache_dir)
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+        self.metadata_file = self.cache_dir / "metadata.json"
+        self.metadata = self._load_metadata()
+        logger.info(f"Cache directory moved to: {self.cache_dir}")
+
+    def _load_metadata(self) -> dict[str, Any]:
+        """Load metadata from the JSON file."""
+        if self.metadata_file.exists():
+            try:
+                with open(self.metadata_file, "r", encoding="utf-8") as f:
+                    return json.load(f)
+            except (json.JSONDecodeError, OSError) as e:
+                logger.warning(f"Metadata load failed: {e}")
+        return {}
+
+    def _save_metadata(self) -> None:
+        """Persist metadata to the disk."""
+        try:
+            with open(self.metadata_file, "w", encoding="utf-8") as f:
+                json.dump(self.metadata, f, ensure_ascii=False, indent=2)
+        except OSError as e:
+            logger.error(f"Metadata save failed: {e}")
+
+    def _get_path(self, key: str, data_type: type | None = None) -> Path:
+        """Resolve the file path based on data type or metadata."""
+        ext = ".dat"
+        if data_type in self._EXTENSIONS:
+            ext = self._EXTENSIONS[data_type]
+        elif key in self.metadata:
+            stored_type = self.metadata[key].get("data_type")
+            ext = self._EXTENSIONS.get(self._TYPE_NAMES.get(stored_type, None), ".dat")
+        return self.cache_dir / f"{key}{ext}"
+
+    @staticmethod
+    def _execute_save(data: Any, path: Path, dtype: type, **kwargs) -> dict:
+        """Execute type-specific save operations."""
+        if dtype is pd.DataFrame:
+            data.to_csv(path, index=kwargs.get("index", False), encoding="utf-8")
+            return {"row_count": len(data), "file_size": path.stat().st_size}
+
+        if dtype in (dict, list):
+            with open(path, "w", encoding="utf-8") as f:
+                json.dump(data, f, ensure_ascii=False, indent=2)
+            return {"item_count": len(data), "file_size": path.stat().st_size}
+
+        if dtype is str:
+            path.write_text(data, encoding=kwargs.get("encoding", "utf-8"))
+            return {"char_count": len(data), "file_size": path.stat().st_size}
+
+        raise ValueError(f"Unsupported type: {dtype}")
+
+    @staticmethod
+    def _execute_load(path: Path, type_name: str, **kwargs) -> Any:
+        """Execute type-specific load operations."""
+        if type_name == "DataFrame":
+            return pd.read_csv(path, encoding=kwargs.get("encoding", "utf-8"))
+        if type_name in ("dict", "list"):
+            with open(path, "r", encoding="utf-8") as f:
+                return json.load(f)
+        if type_name == "str":
+            return path.read_text(encoding=kwargs.get("encoding", "utf-8"))
+        raise ValueError(f"Unknown data type in metadata: {type_name}")
+
+    def save(self, key: str, data: Any, expire_hours: float | None = None, **kwargs) -> bool:
+        """Save data to cache with optional expiration."""
+        try:
+            dtype = type(data)
+            path = self._get_path(key, dtype)
+            stats = self._execute_save(data, path, dtype, **kwargs)
+
+            now = datetime.now()
+            self.metadata[key] = {
+                "created_at": now.isoformat(),
+                "expire_at": (now + timedelta(hours=expire_hours)).isoformat() if expire_hours else None,
+                "data_type": dtype.__name__,
+                **stats,
+            }
+            self._save_metadata()
+            return True
+        except Exception as e:
+            logger.error(f"Save failed for {key}: {e}")
+            return False
+
+    def load(self, key: str, auto_clean: bool = True, **kwargs) -> Any | None:
+        """Load data from cache if not expired."""
+        if self._is_expired(key):
+            if auto_clean:
+                self.delete(key)
+            return None
+
+        path = self._get_path(key)
+        if not path.exists() or key not in self.metadata:
+            return None
+
+        try:
+            return self._execute_load(path, self.metadata[key]["data_type"], **kwargs)
+        except Exception as e:
+            logger.error(f"Load failed for {key}: {e}")
+            return None
+
+    def _is_expired(self, key: str) -> bool:
+        """Check if the cached entry has expired."""
+        entry = self.metadata.get(key)
+        if not entry or not entry.get("expire_at"):
+            return False
+        return datetime.now() > datetime.fromisoformat(entry["expire_at"])
+
+    def delete(self, key: str) -> bool:
+        """Remove a specific cache entry and its file."""
+        try:
+            path = self._get_path(key)
+            if path.exists():
+                path.unlink()
+            if key in self.metadata:
+                del self.metadata[key]
+                self._save_metadata()
+            return True
+        except OSError as e:
+            logger.error(f"Delete failed for {key}: {e}")
+            return False
+
+    def exists(self, key: str) -> bool:
+        """Check if a valid cache entry exists."""
+        return key in self.metadata and not self._is_expired(key)
+
+    def clear_all(self) -> bool:
+        """Purge all cache files and reset metadata."""
+        try:
+            for file in self.cache_dir.iterdir():
+                if file.is_file():
+                    file.unlink()
+            self.metadata = {}
+            self._save_metadata()
+            return True
+        except OSError as e:
+            logger.error(f"Clear all failed: {e}")
+            return False
+
+    def get_stats(self) -> dict[str, Any]:
+        """Return cache usage statistics."""
+        total_size = sum(f.stat().st_size for f in self.cache_dir.glob("*") if f.is_file())
+        return {
+            "count": len(self.metadata),
+            "size_mb": round(total_size / (1024 * 1024), 2),
+            "dir": str(self.cache_dir),
+        }

reme_ai/core/utils/common_utils.py

@@ -0,0 +1,83 @@
+"""Common utility functions"""
+
+import asyncio
+from collections.abc import AsyncGenerator, Coroutine
+from typing import Any
+
+from loguru import logger
+
+from ..enumeration import ChunkEnum
+from ..schema import StreamChunk
+
+
+def run_coro_safely(coro: Coroutine[Any, Any, Any]) -> Any | asyncio.Task[Any]:
+    """Run a coroutine in the current event loop or a new one if none exists."""
+    try:
+        # Attempt to retrieve the event loop associated with the current thread
+        loop = asyncio.get_running_loop()
+
+    except RuntimeError:
+        # Start a new event loop to run the coroutine to completion
+        return asyncio.run(coro)
+
+    else:
+        # Schedule the coroutine as a background task in the active loop
+        return loop.create_task(coro)
+
+
+async def execute_stream_task(
+    queue: asyncio.Queue,
+    task: asyncio.Task,
+    flow_name: str | None = None,
+    as_bytes: bool = False,
+) -> AsyncGenerator[str | bytes, None]:
+    """
+    Core stream flow execution logic.
+
+    Handles streaming from a queue while monitoring the task completion.
+    Properly manages errors and resource cleanup.
+
+    Args:
+        queue: Queue to receive StreamChunk objects from
+        task: Background task executing the flow
+        flow_name: Optional flow name for logging purposes
+        as_bytes: If True, yield bytes for HTTP responses; if False, yield strings
+
+    Yields:
+        SSE-formatted data chunks (either str or bytes based on as_bytes)
+    """
+    done_msg = b"data:[DONE]\n\n" if as_bytes else "data:[DONE]\n\n"
+
+    try:
+        while True:
+            # Wait for next chunk or check if task failed
+            get_chunk = asyncio.create_task(queue.get())
+            done, _ = await asyncio.wait({get_chunk, task}, return_when=asyncio.FIRST_COMPLETED)
+
+            if get_chunk in done:
+                chunk: StreamChunk = get_chunk.result()
+                if chunk.done:
+                    yield done_msg
+                    break
+
+                data = f"data:{chunk.model_dump_json()}\n\n"
+                yield data.encode() if as_bytes else data
+            else:
+                # Task finished unexpectedly or raised exception
+                await task
+                yield done_msg
+                break
+
+    except Exception as e:
+        log_msg = f"Stream error in {flow_name}: {e}" if flow_name else f"Stream error: {e}"
+        logger.exception(log_msg)
+
+        err = StreamChunk(chunk_type=ChunkEnum.ERROR, chunk=str(e), done=True)
+        err_data = f"data:{err.model_dump_json()}\n\n"
+        yield err_data.encode() if as_bytes else err_data
+        yield done_msg
+
+    finally:
+        # Ensure task is cancelled if still running to avoid resource leaks
+        if not task.done():
+            task.cancel()

reme_ai/core/utils/http_client.py

@@ -0,0 +1,91 @@
+"""Asynchronous HTTP client for executing flows with built-in retry logic."""
+
+import json
+from collections.abc import AsyncIterator
+from typing import Optional
+
+import httpx
+from loguru import logger
+
+from ..schema import Response
+
+
+class HttpClient:
+    """Async client for flow endpoints with automated retries and error handling."""
+
+    def __init__(
+        self,
+        base_url: str = "http://localhost:8001",
+        timeout: float = 3600.0,
+        max_retries: int = 3,
+        raise_exception: bool = True,
+    ):
+        """Initialize the client with base configuration."""
+        self.base_url = base_url.rstrip("/")
+        self.timeout = timeout
+        self.max_retries = max_retries
+        self.raise_exception = raise_exception
+        self.client = httpx.AsyncClient(timeout=timeout)
+
+    async def __aenter__(self):
+        """Enter async context manager."""
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Exit async context manager and close connection."""
+        await self.close()
+
+    async def close(self):
+        """Close the underlying HTTP client."""
+        await self.client.aclose()
+
+    async def health_check(self) -> dict[str, str]:
+        """Check the health status of the flow service."""
+        response = await self.client.get(f"{self.base_url}/health")
+        response.raise_for_status()
+        return response.json()
+
+    async def execute_flow(self, flow_name: str, **kwargs) -> Optional[Response]:
+        """Execute a flow with automated retry logic."""
+        endpoint = f"{self.base_url}/{flow_name}"
+
+        for attempt in range(self.max_retries):
+            try:
+                response = await self.client.post(endpoint, json=kwargs)
+                response.raise_for_status()
+                return Response(**response.json())
+
+            except (httpx.HTTPError, Exception) as e:
+                logger.error(f"Flow {flow_name} failed (attempt {attempt + 1}/{self.max_retries}): {e}")
+                if attempt == self.max_retries - 1 and self.raise_exception:
+                    raise e
+        return None
+
+    async def list_endpoints(self) -> dict:
+        """Retrieve available endpoints from OpenAPI specification."""
+        response = await self.client.get(f"{self.base_url}/openapi.json")
+        response.raise_for_status()
+        return response.json()
+
+    async def execute_stream_flow(self, flow_name: str, **kwargs) -> AsyncIterator[dict[str, str]]:
+        """Execute a flow and yield parsed SSE stream chunks."""
+        endpoint = f"{self.base_url}/{flow_name}"
+
+        async with self.client.stream("POST", endpoint, json=kwargs) as response:
+            response.raise_for_status()
+            async for line in response.aiter_lines():
+                if not line or not line.startswith("data:"):
+                    continue
+
+                content = line.removeprefix("data:").strip()
+                if content == "[DONE]":
+                    break
+
+                try:
+                    data = json.loads(content)
+                    yield {
+                        "type": data.get("chunk_type", "answer"),
+                        "content": data.get("chunk", ""),
+                    }
+                except json.JSONDecodeError:
+                    continue

reme_ai/core/utils/llm_utils.py

@@ -0,0 +1,40 @@
+"""Utility functions for processing and formatting LLM-related message data."""
+
+import json
+import re
+
+from ..enumeration import Role
+from ..schema import Message
+
+
+def format_messages(messages: list[Message | dict], enable_system: bool = False) -> str:
+    """Formats a list of messages into a single string, optionally filtering system roles."""
+    formatted_lines = []
+    for message in messages:
+        if isinstance(message, dict):
+            message = Message(**message)
+        if not enable_system and message.role is Role.SYSTEM:
+            continue
+
+        formatted_lines.append(message.format_message())
+    return "\n".join(formatted_lines)
+
+
+def extract_content(text: str, language_tag: str = "json", greedy: bool = False):
+    """Extracts content from Markdown code blocks and parses it if the tag is JSON."""
+    quantifier = ".*" if greedy else ".*?"
+    pattern = rf"```\s*{re.escape(language_tag)}\s*({quantifier})\s*```"
+    match = re.search(pattern, text, re.DOTALL)
+
+    if match:
+        result = match.group(1).strip()
+    else:
+        result = text
+
+    if language_tag == "json":
+        try:
+            result = json.loads(result)
+        except json.JSONDecodeError:
+            result = None
+
+    return result

reme_ai/core/utils/logger_utils.py

@@ -0,0 +1,40 @@
+"""Logging configuration module for application-wide tracing."""
+
+import os
+import sys
+from datetime import datetime
+
+
+def init_logger(log_dir: str = "logs", level: str = "INFO") -> None:
+    """Initialize the logger with both file and console handlers."""
+    from loguru import logger
+
+    # Remove default handler to avoid duplicate logs
+    logger.remove()
+
+    # Ensure the logging directory exists
+    os.makedirs(log_dir, exist_ok=True)
+
+    # Generate filename based on the current timestamp
+    current_ts = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+    log_filename = f"{current_ts}.log"
+    log_filepath = os.path.join(log_dir, log_filename)
+
+    # Configure file-based logging with rotation and compression
+    logger.add(
+        log_filepath,
+        level=level,
+        rotation="00:00",
+        retention="7 days",
+        compression="zip",
+        encoding="utf-8",
+        format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {file}:{line} | {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}",
+        colorize=True,
+    )

reme_ai/core/utils/logo_utils.py

@@ -0,0 +1,84 @@
+"""Terminal branding and configuration display utilities."""
+
+import importlib.metadata
+from typing import TYPE_CHECKING
+
+from rich.console import Console, Group
+from rich.panel import Panel
+from rich.table import Table
+from rich.text import Text
+
+if TYPE_CHECKING:
+    from ..schema import ServiceConfig
+
+
+def get_version(package_name: str) -> str:
+    """Return the installed version of a package or 'unknown'."""
+    try:
+        return importlib.metadata.version(package_name)
+    except importlib.metadata.PackageNotFoundError:
+        return ""
+
+
+def print_logo(service_config: "ServiceConfig"):
+    """Print a stylized ASCII logo and service metadata to the console."""
+    ascii_art = [
+        r" ██████╗  ███████╗ ███╗   ███╗ ███████╗ ",
+        r" ██╔══██╗ ██╔════╝ ████╗ ████║ ██╔════╝ ",
+        r" ██████╔╝ █████╗   ██╔████╔██║ █████╗   ",
+        r" ██╔══██╗ ██╔══╝   ██║╚██╔╝██║ ██╔══╝   ",
+        r" ██║  ██║ ███████╗ ██║ ╚═╝ ██║ ███████╗ ",
+        r" ╚═╝  ╚═╝ ╚══════╝ ╚═╝     ╚═╝ ╚══════╝ ",
+    ]
+
+    start_color = (85, 239, 196)
+    end_color = (162, 155, 254)
+
+    logo_text = Text()
+    for line in ascii_art:
+        line_len = max(1, len(line) - 1)
+        for i, char in enumerate(line):
+            # Calculate gradient shift per character
+            ratio = i / line_len
+            rgb = tuple(int(s + (e - s) * ratio) for s, e in zip(start_color, end_color))
+            logo_text.append(char, style=f"bold rgb({rgb[0]},{rgb[1]},{rgb[2]})")
+        logo_text.append("\n")
+
+    # Layout configuration info
+    info_table = Table.grid(padding=(0, 1))
+    info_table.add_column(style="bold", justify="center")
+    info_table.add_column(style="bold cyan", justify="left")
+    info_table.add_column(style="white", justify="left")
+
+    # Add core service info
+    info_table.add_row("📦", "Backend:", service_config.backend)
+
+    match service_config.backend:
+        case "http":
+            host, port = service_config.http.host, service_config.http.port
+            info_table.add_row("🔗", "URL:", f"http://{host}:{port}")
+            info_table.add_row("📚", "FastAPI:", Text(get_version("fastapi"), style="dim"))
+        case "mcp":
+            mcp = service_config.mcp
+            transport = mcp.transport if mcp.transport else "stdio"
+            info_table.add_row("🚌", "Transport:", transport)
+            if transport != "stdio":
+                url = f"http://{mcp.host}:{mcp.port}"
+                if transport == "sse":
+                    url += "/sse"
+                info_table.add_row("🔗", "URL:", url)
+            info_table.add_row("📚", "FastMCP:", Text(get_version("fastmcp"), style="dim"))
+
+    info_table.add_row("🚀", "ReMe:", Text(get_version("reme-ai"), style="dim"))
+
+    # Render layout within a panel
+    panel = Panel(
+        Group(logo_text, info_table),
+        title=service_config.app_name,
+        title_align="left",
+        border_style="dim",
+        padding=(1, 4),
+        expand=False,
+    )
+
+    Console().print(Group("\n", panel, "\n"))

reme_ai/core/utils/mcp_client.py

@@ -0,0 +1,121 @@
+"""Module for managing Model Context Protocol (MCP) server connections."""
+
+import os
+import re
+from contextlib import asynccontextmanager
+from typing import Any
+
+from mcp import ClientSession, StdioServerParameters, Tool
+from mcp.client.sse import sse_client
+from mcp.client.stdio import stdio_client
+from mcp.client.streamable_http import streamablehttp_client
+from mcp.types import CallToolResult, TextContent
+
+from ..schema import ToolCall
+
+
+class MCPClient:
+    """A client manager for handling multiple MCP transport protocols."""
+
+    def __init__(self, config: dict):
+        """Initialize the client with server configuration."""
+        self.config = config
+
+    @staticmethod
+    def _infer_transport_type(cfg: dict[str, Any]) -> str:
+        """Infer the transport type based on configuration fields."""
+        if "command" in cfg:
+            return "stdio"
+
+        if "url" in cfg:
+            url = cfg["url"].lower()
+            if url.endswith("/sse") or "sse" in url:
+                return "sse"
+            return "streamable-http"
+
+        raise ValueError(f"Could not infer transport type for: {cfg}")
+
+    def _replace_env_vars(self, data: str | dict | list) -> Any:
+        """Replace environment variable placeholders in configuration."""
+        if isinstance(data, str):
+            return re.sub(r"\$\{(\w+)\}", lambda m: os.getenv(m.group(1), m.group(0)), data)
+        if isinstance(data, dict):
+            return {k: self._replace_env_vars(v) for k, v in data.items()}
+        if isinstance(data, list):
+            return [self._replace_env_vars(i) for i in data]
+        return data
+
+    @asynccontextmanager
+    async def _get_transport(self, cfg: dict[str, Any]):
+        """Context manager to yield the appropriate MCP transport."""
+        # Pop 'type' if present, otherwise infer it
+        t_type = cfg.pop("type", None) or self._infer_transport_type(cfg)
+
+        try:
+            if t_type == "stdio":
+                params = StdioServerParameters(
+                    command=cfg["command"],
+                    args=cfg.get("args", []),
+                    env=cfg.get("env", None),
+                )
+                async with stdio_client(params) as transport:
+                    yield transport
+            elif t_type == "sse":
+                async with sse_client(**cfg) as transport:
+                    yield transport
+            elif t_type == "streamable-http":
+                async with streamablehttp_client(**cfg) as transport:
+                    yield transport
+            else:
+                raise NotImplementedError(f"Unsupported transport: {t_type}")
+        finally:
+            pass  # Ensure proper cleanup
+
+    @asynccontextmanager
+    async def connect_to_server(self, server_name: str):
+        """Establish a session with the specified MCP server."""
+        server_config = self.config.get("mcpServers", {}).get(server_name)
+        if not server_config:
+            raise ValueError(f"Config for '{server_name}' not found.")
+
+        # Process environment variables and transport selection
+        cfg = self._replace_env_vars(server_config)
+
+        async with self._get_transport(cfg) as (read, write):
+            async with ClientSession(read, write) as session:
+                await session.initialize()
+                yield session
+
+    async def list_tools(self, server_name: str) -> list[Tool]:
+        """Retrieve available tools from a specific server."""
+        async with self.connect_to_server(server_name) as session:
+            result = await session.list_tools()
+            return result.tools
+
+    async def list_tool_calls(self, server_name: str, return_dict: bool = True) -> list[dict | ToolCall]:
+        """Retrieve available tools from a specific server."""
+        tools = await self.list_tools(server_name)
+        tool_calls: list[ToolCall] = [ToolCall.from_mcp_tool(tool) for tool in tools]
+        if return_dict:
+            return [tool_call.simple_input_dump() for tool_call in tool_calls]
+
+        return tool_calls
+
+    async def call_tool(
+        self,
+        server_name: str,
+        tool_name: str,
+        arguments: dict[str, Any],
+        parse_text_result: bool = False,
+    ) -> CallToolResult | str:
+        """Execute a tool on a specific server."""
+        async with self.connect_to_server(server_name) as session:
+            tool_results: CallToolResult = await session.call_tool(tool_name, arguments)
+            if not parse_text_result:
+                return tool_results
+
+            text_result = []
+            for block in tool_results.content:
+                if isinstance(block, TextContent):
+                    text_result.append(block.text)
+            return "\n".join(text_result)

reme_ai/core/utils/pydantic_config_parser.py

@@ -0,0 +1,206 @@
+"""Parser for Pydantic config models with YAML and CLI argument support."""
+
+import inspect
+import json
+from pathlib import Path
+from typing import Any, TypeVar
+
+import yaml
+from loguru import logger
+from pydantic import BaseModel
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class PydanticConfigParser:
+    """Parser that loads and merges Pydantic configs from YAML files and CLI args."""
+
+    def __init__(self, config_class: type[T], default_config: str = "default"):
+        """Initialize parser with a Pydantic config class.
+
+        Args:
+            config_class: Pydantic BaseModel class to validate configs against.
+            default_config: Default config file name to use if not specified in args.
+        """
+        self.config_class = config_class
+        self.default_config = default_config
+        self.config_dict: dict = {}
+
+    def _deep_merge(self, base_dict: dict, update_dict: dict) -> dict:
+        """Recursively merge two dictionaries."""
+        result = base_dict.copy()
+        for key, value in update_dict.items():
+            if key in result and isinstance(result[key], dict) and isinstance(value, dict):
+                result[key] = self._deep_merge(result[key], value)
+            else:
+                result[key] = value
+        return result
+
+    @staticmethod
+    def _convert_value(value_str: str) -> Any:
+        """Convert string value to appropriate Python type."""
+        value_str = value_str.strip()
+        lower_str = value_str.lower()
+
+        # Boolean and None conversion
+        if lower_str in ("true", "false"):
+            return lower_str == "true"
+        if lower_str in ("none", "null"):
+            return None
+
+        # Numeric conversion
+        if "e" in lower_str or "." in value_str:
+            try:
+                return float(value_str)
+            except ValueError:
+                pass
+        else:
+            try:
+                return int(value_str)
+            except ValueError:
+                pass
+
+        # JSON conversion for complex types
+        try:
+            return json.loads(value_str)
+        except (json.JSONDecodeError, ValueError):
+            return value_str
+
+    @staticmethod
+    def load_from_yaml(yaml_path: str | Path) -> dict:
+        """Load configuration from YAML file.
+
+        Args:
+            yaml_path: Path to YAML configuration file.
+
+        Returns:
+            Dictionary containing configuration data.
+
+        Raises:
+            FileNotFoundError: If YAML file does not exist.
+        """
+        if isinstance(yaml_path, str):
+            yaml_path = Path(yaml_path)
+
+        if not yaml_path.exists():
+            raise FileNotFoundError(f"Configuration file does not exist: {yaml_path}")
+
+        with yaml_path.open(encoding="utf-8") as f:
+            return yaml.safe_load(f) or {}
+
+    def merge_configs(self, *config_dicts: dict) -> dict:
+        """Merge multiple config dictionaries in order.
+
+        Args:
+            *config_dicts: Variable number of config dictionaries to merge.
+
+        Returns:
+            Merged configuration dictionary.
+        """
+        result = {}
+        for config_dict in config_dicts:
+            result = self._deep_merge(result, config_dict)
+        return result
+
+    def parse_dot_notation(self, dot_list: list[str]) -> dict:
+        """Parse dot notation strings into nested dictionary.
+
+        Args:
+            dot_list: List of strings in format "key.subkey=value".
+
+        Returns:
+            Nested dictionary representation of dot notation.
+        """
+        config_dict = {}
+        for item in dot_list:
+            if "=" not in item:
+                continue
+
+            key_path, value_str = item.split("=", 1)
+            keys = key_path.split(".")
+
+            # Build nested dictionary
+            current = config_dict
+            for key in keys[:-1]:
+                current = current.setdefault(key, {})
+            current[keys[-1]] = self._convert_value(value_str)
+
+        return config_dict
+
+    def _find_config_path(self, config_name: str) -> Path:
+        """Find config file path, trying parser directory first then current directory."""
+        if not config_name.endswith(".yaml"):
+            config_name += ".yaml"
+
+        # Try parser class directory first
+        config_path = Path(inspect.getfile(self.__class__)).parent / config_name
+        if config_path.exists():
+            logger.info(f"load config={config_path}")
+            return config_path
+
+        # Try current directory
+        logger.warning(f"config={config_path} not found, try {config_name}")
+        config_path = Path(config_name)
+        if not config_path.exists():
+            raise FileNotFoundError(f"config={config_path} not found")
+        return config_path
+
+    def parse_args(self, *args: str) -> T:
+        """Parse CLI arguments and load configs from YAML files.
+
+        Args:
+            *args: CLI arguments in format "key=value" or "config=file.yaml".
+
+        Returns:
+            Validated Pydantic config instance.
+
+        Raises:
+            ValueError: If no config file is specified.
+            FileNotFoundError: If specified config file does not exist.
+        """
+        configs_to_merge = [self.config_class().model_dump()]
+
+        # Separate config file path from other arguments
+        config = ""
+        filter_args = []
+        for arg in args:
+            if "=" not in arg:
+                continue
+            arg = arg.lstrip("-")
+            if arg.startswith(("c=", "config=")):
+                config = arg.split("=", 1)[1]
+            else:
+                filter_args.append(arg)
+
+        # Use default config if not specified
+        config = config or self.default_config
+
+        # Load each config file
+        for single_config in (c.strip() for c in config.split(",") if c.strip()):
+            config_path = self._find_config_path(single_config)
+            configs_to_merge.append(self.load_from_yaml(config_path))
+
+        # Apply CLI overrides
+        if filter_args:
+            configs_to_merge.append(self.parse_dot_notation(filter_args))
+
+        # Merge all configs and validate
+        self.config_dict = self.merge_configs(*configs_to_merge)
+        return self.config_class.model_validate(self.config_dict)
+
+    def update_config(self, **kwargs) -> T:
+        """Update current config with new values using kwargs.
+
+        Args:
+            **kwargs: Key-value pairs where __ in keys represents nested levels.
+
+        Returns:
+            Updated and validated Pydantic config instance.
+        """
+        # Convert kwargs to dot notation and parse
+        dot_list = [f"{key.replace('__', '.')}={value}" for key, value in kwargs.items()]
+        override_config = self.parse_dot_notation(dot_list)
+
+        # Merge with existing config
+        final_config = self.merge_configs(self.config_dict, override_config)
+        return self.config_class.model_validate(final_config)

reme_ai/core/utils/pydantic_utils.py

@@ -0,0 +1,66 @@
+"""
+Utility module for dynamic Pydantic model generation based on schema definitions.
+"""
+
+from typing import Any, Literal
+
+from pydantic import create_model, Field
+
+from . import snake_to_camel
+from ..enumeration import JsonSchemaEnum
+from ..schema import ToolAttr, Request
+
+TYPE_MAPPING = {str(t): t.value for t in JsonSchemaEnum}
+
+
+def create_pydantic_model(name: str, parameters: ToolAttr | None = None) -> type[Request]:
+    """
+    Recursively generates a Pydantic model from a ToolAttr schema definition.
+    """
+    fields = {}
+
+    if not parameters or not parameters.properties:
+        return create_model(f"{snake_to_camel(name)}Model", __base__=Request)
+
+    for field_name, attr in parameters.properties.items():
+        # 1. Determine the base field type
+        if attr.type == "object" and attr.properties:
+            # Handle nested objects recursively
+            field_type = create_pydantic_model(field_name, attr)
+
+        elif attr.type == "array" and attr.items:
+            # Handle array/list types
+            if isinstance(attr.items, ToolAttr):
+                if attr.items.type == "object":
+                    inner_type = create_pydantic_model(f"{field_name}_item", attr.items)
+                else:
+                    inner_type = TYPE_MAPPING.get(attr.items.type, Any)
+                field_type = list[inner_type]
+            else:
+                # Fallback for simple dictionary item definitions
+                field_type = list[Any]
+
+        else:
+            # Handle primitive types
+            field_type = TYPE_MAPPING.get(attr.type, Any)
+
+        # 2. Handle enumeration constraints
+        if attr.enum:
+            # Dynamically create a Literal type from the enum list
+            field_type = Literal[tuple(attr.enum)]  # type: ignore
+
+        # 3. Determine requirement status and default values
+        is_required = False
+        if parameters.required and field_name in parameters.required:
+            is_required = True
+
+        # 4. Construct Field metadata
+        field_info = Field(default=... if is_required else None, description=attr.description)
+
+        if not is_required:
+            field_type = field_type | None
+
+        fields[field_name] = (field_type, field_info)
+
+    # Dynamically construct the final Pydantic model class
+    return create_model(f"{snake_to_camel(name)}Model", **fields, __base__=Request)

reme_ai/core/vector_store/__init__.py

@@ -0,0 +1,17 @@
+"""vector store"""
+
+from .base_vector_store import BaseVectorStore
+from .chroma_vector_store import ChromaVectorStore
+from .es_vector_store import ESVectorStore
+from .local_vector_store import LocalVectorStore
+from .pgvector_store import PGVectorStore
+from .qdrant_vector_store import QdrantVectorStore
+
+__all__ = [
+    "BaseVectorStore",
+    "ChromaVectorStore",
+    "ESVectorStore",
+    "LocalVectorStore",
+    "PGVectorStore",
+    "QdrantVectorStore",
+]

reme_ai/core/vector_store/base_vector_store.py

@@ -0,0 +1,92 @@
+"""Base vector store interface for managing vector embeddings and similarity search."""
+
+import asyncio
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from functools import partial
+
+from reme_ai.core.context import C
+from reme_ai.core.embedding import BaseEmbeddingModel
+from reme_ai.core.schema import VectorNode
+
+
+class BaseVectorStore(ABC):
+    """Abstract base class defining the interface for vector storage and retrieval."""
+
+    def __init__(
+        self,
+        collection_name: str,
+        embedding_model: BaseEmbeddingModel,
+        **kwargs,
+    ):
+        """Initialize the vector store with a collection name and an embedding model."""
+        if embedding_model is None:
+            raise ValueError("embedding_model is required")
+        self.collection_name: str = collection_name
+        self.embedding_model: BaseEmbeddingModel = embedding_model
+        self.kwargs: dict = kwargs
+
+    @staticmethod
+    async def _run_sync_in_executor(sync_func: Callable, *args, **kwargs):
+        """Run a synchronous function in the context-defined thread pool executor."""
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(C.thread_pool, partial(sync_func, *args, **kwargs))
+
+    async def get_node_embedding(self, node: VectorNode) -> VectorNode:
+        """Generate and assign embedding for a single vector node."""
+        return await self.embedding_model.get_node_embedding(node)
+
+    async def get_node_embeddings(self, nodes: list[VectorNode]) -> list[VectorNode]:
+        """Generate and assign embeddings for multiple vector nodes."""
+        return await self.embedding_model.get_node_embeddings(nodes)
+
+    async def get_embedding(self, query: str) -> list[float]:
+        """Convert a single text query into vector embedding using the configured model."""
+        return await self.embedding_model.get_embedding(query)
+
+    async def get_embeddings(self, queries: list[str]) -> list[list[float]]:
+        """Convert multiple text queries into vector embeddings using the configured model."""
+        return await self.embedding_model.get_embeddings(queries)
+
+    @abstractmethod
+    async def list_collections(self) -> list[str]:
+        """Retrieve a list of all existing collection names in the store."""
+
+    @abstractmethod
+    async def create_collection(self, collection_name: str, **kwargs) -> None:
+        """Create a new vector collection with the specified name and configuration."""
+
+    @abstractmethod
+    async def delete_collection(self, collection_name: str, **kwargs) -> None:
+        """Permanently remove a collection from the vector store."""
+
+    @abstractmethod
+    async def copy_collection(self, collection_name: str, **kwargs) -> None:
+        """Duplicate the current collection to a new one with the given name."""
+
+    @abstractmethod
+    async def insert(self, nodes: VectorNode | list[VectorNode], **kwargs) -> None:
+        """Add one or more vector nodes into the current collection."""
+
+    @abstractmethod
+    async def search(self, query: str, limit: int = 5, filters: dict | None = None, **kwargs) -> list[VectorNode]:
+        """Find the most similar vector nodes based on a text query."""
+
+    @abstractmethod
+    async def delete(self, vector_ids: str | list[str], **kwargs) -> None:
+        """Remove specific vectors from the collection using their identifiers."""
+
+    @abstractmethod
+    async def update(self, nodes: VectorNode | list[VectorNode], **kwargs) -> None:
+        """Update the data or metadata of existing vectors in the collection."""
+
+    @abstractmethod
+    async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode]:
+        """Fetch specific vector nodes from the collection by their IDs."""
+
+    @abstractmethod
+    async def list(self, filters: dict | None = None, limit: int | None = None) -> list[VectorNode]:
+        """Retrieve vectors from the collection that match the given filters."""
+
+    async def close(self) -> None:
+        """Release resources and close active connections to the vector store."""

reme_ai/core/vector_store/chroma_vector_store.py

@@ -0,0 +1,397 @@
+"""ChromaDB vector store implementation for the ReMe framework."""
+
+from typing import Any
+
+from loguru import logger
+
+from .base_vector_store import BaseVectorStore
+from ..context import C
+from ..embedding import BaseEmbeddingModel
+from ..schema import VectorNode
+
+_CHROMADB_IMPORT_ERROR = None
+
+try:
+    import chromadb
+    from chromadb.config import Settings
+except ImportError as e:
+    _CHROMADB_IMPORT_ERROR = e
+    chromadb = None
+    Settings = None
+
+
+@C.register_vector_store("chroma")
+class ChromaVectorStore(BaseVectorStore):
+    """ChromaDB-based vector store implementation for local or remote storage."""
+
+    def __init__(
+        self,
+        collection_name: str,
+        embedding_model: BaseEmbeddingModel,
+        client: chromadb.ClientAPI | None = None,
+        host: str | None = None,
+        port: int | None = None,
+        path: str | None = None,
+        api_key: str | None = None,
+        tenant: str | None = None,
+        database: str | None = None,
+        **kwargs,
+    ):
+        """Initialize the ChromaDB vector store with the provided configuration."""
+        if _CHROMADB_IMPORT_ERROR is not None:
+            raise ImportError(
+                "ChromaDB requires extra dependencies. Install with `pip install chromadb`",
+            ) from _CHROMADB_IMPORT_ERROR
+
+        super().__init__(
+            collection_name=collection_name,
+            embedding_model=embedding_model,
+            **kwargs,
+        )
+
+        self.client: chromadb.ClientAPI
+        self.collection: chromadb.Collection
+
+        if client:
+            self.client = client
+        elif api_key and tenant:
+            logger.info("Initializing ChromaDB Cloud client")
+            self.client = chromadb.CloudClient(
+                api_key=api_key,
+                tenant=tenant,
+                database=database or "default",
+            )
+        elif host and port:
+            logger.info(f"Initializing ChromaDB HTTP client at {host}:{port}")
+            self.client = chromadb.HttpClient(host=host, port=port)
+        else:
+            if path is None:
+                path = "./chroma_db"
+            logger.info(f"Initializing local ChromaDB at {path}")
+            self.client = chromadb.PersistentClient(
+                path=path,
+                settings=Settings(anonymized_telemetry=False),
+            )
+
+        self.collection = self.client.get_or_create_collection(
+            name=collection_name,
+            metadata={"hnsw:space": "cosine"},
+        )
+
+    @staticmethod
+    def _parse_results(
+        results: dict,
+        include_score: bool = False,
+    ) -> list[VectorNode]:
+        """Convert ChromaDB query results into a list of VectorNode objects."""
+        nodes = []
+
+        ids = results.get("ids", [])
+        documents = results.get("documents", [])
+        metadatas = results.get("metadatas", [])
+        embeddings = results.get("embeddings") if results.get("embeddings") is not None else []
+        distances = results.get("distances") if results.get("distances") is not None else []
+
+        if ids and isinstance(ids[0], list):
+            ids = ids[0] if ids else []
+            documents = documents[0] if documents else []
+            metadatas = metadatas[0] if metadatas else []
+            embeddings = embeddings[0] if embeddings and len(embeddings) > 0 else []
+            distances = distances[0] if distances and len(distances) > 0 else []
+
+        for i, vector_id in enumerate(ids):
+            metadata = metadatas[i] if i < len(metadatas) and metadatas[i] else {}
+
+            if include_score and distances and i < len(distances):
+                metadata["_score"] = 1.0 - distances[i]
+
+            node = VectorNode(
+                vector_id=vector_id,
+                content=documents[i] if i < len(documents) and documents[i] else "",
+                vector=embeddings[i] if len(embeddings) > i else None,
+                metadata=metadata,
+            )
+            nodes.append(node)
+
+        return nodes
+
+    @staticmethod
+    def _generate_where_clause(filters: dict | None) -> dict | None:
+        """Convert the universal filter format to a ChromaDB-compatible where clause."""
+        if not filters:
+            return None
+
+        def convert_condition(k: str, v: Any) -> dict | None:
+            """Convert a single filter condition to ChromaDB operator format."""
+            if v == "*":
+                return None
+            if isinstance(v, dict):
+                chroma_condition = {}
+                for op, val in v.items():
+                    mapping = {
+                        "eq": "$eq",
+                        "ne": "$ne",
+                        "gt": "$gt",
+                        "gte": "$gte",
+                        "lt": "$lt",
+                        "lte": "$lte",
+                        "in": "$in",
+                        "nin": "$nin",
+                    }
+                    chroma_op = mapping.get(op, "$eq")
+                    chroma_condition[k] = {chroma_op: val}
+                return chroma_condition
+            if isinstance(v, list):
+                return {k: {"$in": v}}
+            return {k: {"$eq": v}}
+
+        processed_filters = []
+
+        for key, value in filters.items():
+            if key == "$or":
+                or_conditions = []
+                for condition in value:
+                    or_condition = {}
+                    for sub_key, sub_value in condition.items():
+                        converted = convert_condition(sub_key, sub_value)
+                        if converted:
+                            or_condition.update(converted)
+                    if or_condition:
+                        or_conditions.append(or_condition)
+                if len(or_conditions) > 1:
+                    processed_filters.append({"$or": or_conditions})
+                elif len(or_conditions) == 1:
+                    processed_filters.append(or_conditions[0])
+
+            elif key == "$and":
+                for condition in value:
+                    for sub_key, sub_value in condition.items():
+                        converted = convert_condition(sub_key, sub_value)
+                        if converted:
+                            processed_filters.append(converted)
+            elif key == "$not":
+                continue
+            else:
+                converted = convert_condition(key, value)
+                if converted:
+                    processed_filters.append(converted)
+
+        if not processed_filters:
+            return None
+        return processed_filters[0] if len(processed_filters) == 1 else {"$and": processed_filters}
+
+    async def list_collections(self) -> list[str]:
+        """Retrieve a list of all existing collection names."""
+
+        def _list():
+            return [col.name for col in self.client.list_collections()]
+
+        return await self._run_sync_in_executor(_list)
+
+    async def create_collection(self, collection_name: str, **kwargs):
+        """Create a new collection with specified distance metrics and metadata."""
+
+        def _create():
+            distance_metric = kwargs.get("distance_metric", "cosine")
+            metadata = kwargs.get("metadata", {})
+            metadata["hnsw:space"] = distance_metric
+            return self.client.get_or_create_collection(name=collection_name, metadata=metadata)
+
+        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}")
+
+    async def delete_collection(self, collection_name: str, **kwargs):
+        """Delete a specified collection from the database."""
+
+        def _delete():
+            try:
+                self.client.delete_collection(name=collection_name)
+                return True
+            except Exception as e:
+                logger.warning(f"Failed to delete collection {collection_name}: {e}")
+                return False
+
+        deleted = await self._run_sync_in_executor(_delete)
+        if deleted and collection_name == self.collection_name:
+            self.collection = None
+        logger.info(f"Deleted collection {collection_name}")
+
+    async def copy_collection(self, collection_name: str, **kwargs):
+        """Copy all data from the current collection to a new collection."""
+
+        def _copy():
+            source_data = self.collection.get(include=["documents", "metadatas", "embeddings"])
+            if not source_data["ids"]:
+                logger.warning(f"Source collection {self.collection_name} is empty")
+                return
+
+            target_collection = self.client.get_or_create_collection(
+                name=collection_name,
+                metadata={"hnsw:space": "cosine"},
+            )
+            target_collection.add(
+                ids=source_data["ids"],
+                documents=source_data["documents"],
+                metadatas=source_data["metadatas"],
+                embeddings=source_data["embeddings"],
+            )
+
+        await self._run_sync_in_executor(_copy)
+        logger.info(f"Copied collection {self.collection_name} to {collection_name}")
+
+    async def insert(self, nodes: VectorNode | list[VectorNode], **kwargs):
+        """Insert vector nodes into the current collection in batches."""
+        if isinstance(nodes, VectorNode):
+            nodes = [nodes]
+        if not nodes:
+            return
+
+        # Batch generate embeddings for nodes that need them
+        nodes_without_vectors = [node for node in nodes if node.vector is None]
+        if nodes_without_vectors:
+            nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors)
+            # Create a mapping for quick lookup
+            vector_map = {n.vector_id: n for n in nodes_with_vectors}
+            nodes_to_insert = [vector_map.get(n.vector_id, n) if n.vector is None else n for n in nodes]
+        else:
+            nodes_to_insert = nodes
+
+        batch_size = kwargs.get("batch_size", 100)
+
+        def _insert_batch(batch_nodes: list[VectorNode]):
+            self.collection.add(
+                ids=[n.vector_id for n in batch_nodes],
+                documents=[n.content for n in batch_nodes],
+                embeddings=[n.vector for n in batch_nodes],
+                metadatas=[n.metadata for n in batch_nodes],
+            )
+
+        for i in range(0, len(nodes_to_insert), batch_size):
+            await self._run_sync_in_executor(_insert_batch, nodes_to_insert[i : i + batch_size])
+        logger.info(f"Inserted {len(nodes_to_insert)} nodes into {self.collection_name}")
+
+    async def search(
+        self,
+        query: str,
+        limit: int = 5,
+        filters: dict | None = None,
+        **kwargs,
+    ) -> list[VectorNode]:
+        """Search for the most similar vector nodes based on a text query."""
+        query_vector = await self.get_embedding(query)
+        where_clause = self._generate_where_clause(filters)
+        include_embeddings = kwargs.get("include_embeddings", False)
+
+        def _search():
+            include: list = ["documents", "metadatas", "distances"]
+            if include_embeddings:
+                include.append("embeddings")
+            return self.collection.query(
+                query_embeddings=[query_vector],
+                n_results=limit,
+                where=where_clause,
+                include=include,
+            )
+
+        results = await self._run_sync_in_executor(_search)
+        nodes = self._parse_results(results, include_score=True)
+
+        score_threshold = kwargs.get("score_threshold")
+        if score_threshold is not None:
+            nodes = [n for n in nodes if n.metadata.get("_score", 0) >= score_threshold]
+        return nodes
+
+    async def delete(self, vector_ids: str | list[str], **kwargs):
+        """Delete specific vector nodes by their IDs."""
+        if isinstance(vector_ids, str):
+            vector_ids = [vector_ids]
+        if not vector_ids:
+            return
+
+        def _delete():
+            self.collection.delete(ids=vector_ids)
+
+        await self._run_sync_in_executor(_delete)
+        logger.info(f"Deleted {len(vector_ids)} nodes from {self.collection_name}")
+
+    async def update(self, nodes: VectorNode | list[VectorNode], **kwargs):
+        """Update existing vector nodes with new content or metadata."""
+        if isinstance(nodes, VectorNode):
+            nodes = [nodes]
+        if not nodes:
+            return
+
+        # Batch generate embeddings for nodes that need them
+        nodes_without_vectors = [node for node in nodes if node.vector is None and node.content]
+        if nodes_without_vectors:
+            nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors)
+            # Create a mapping for quick lookup
+            vector_map = {n.vector_id: n for n in nodes_with_vectors}
+            nodes_to_update = [vector_map.get(n.vector_id, n) if n.vector is None and n.content else n for n in nodes]
+        else:
+            nodes_to_update = nodes
+
+        def _update():
+            self.collection.upsert(
+                ids=[n.vector_id for n in nodes_to_update],
+                documents=[n.content for n in nodes_to_update],
+                embeddings=[n.vector for n in nodes_to_update if n.vector] or None,
+                metadatas=[n.metadata for n in nodes_to_update],
+            )
+
+        await self._run_sync_in_executor(_update)
+        logger.info(f"Updated {len(nodes_to_update)} nodes in {self.collection_name}")
+
+    async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode] | None:
+        """Fetch vector nodes by their IDs from the collection."""
+        is_single = isinstance(vector_ids, str)
+        ids = [vector_ids] if is_single else vector_ids
+
+        def _get():
+            return self.collection.get(ids=ids, include=["documents", "metadatas", "embeddings"])
+
+        results = await self._run_sync_in_executor(_get)
+        nodes = self._parse_results(results)
+        return nodes[0] if is_single and nodes else (nodes if not is_single else None)
+
+    async def list(
+        self,
+        filters: dict | None = None,
+        limit: int | None = None,
+    ) -> list[VectorNode]:
+        """List vector nodes matching optional metadata filters."""
+        where_clause = self._generate_where_clause(filters)
+
+        def _list():
+            return self.collection.get(
+                where=where_clause,
+                limit=limit,
+                include=["documents", "metadatas", "embeddings"],
+            )
+
+        results = await self._run_sync_in_executor(_list)
+        return self._parse_results(results)
+
+    async def count(self) -> int:
+        """Return the total number of vectors in the current collection."""
+        return await self._run_sync_in_executor(self.collection.count)
+
+    async def reset(self):
+        """Reset the current collection by clearing all its data."""
+        logger.warning(f"Resetting collection {self.collection_name}...")
+        await self.delete_collection(self.collection_name)
+
+        def _recreate():
+            self.collection = self.client.get_or_create_collection(
+                name=self.collection_name,
+                metadata={"hnsw:space": "cosine"},
+            )
+
+        await self._run_sync_in_executor(_recreate)
+        logger.info(f"Collection {self.collection_name} has been reset")
+
+    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")

reme_ai/core/vector_store/es_vector_store.py

@@ -0,0 +1,458 @@
+"""Elasticsearch vector store implementation for ReMe.
+
+This module provides an Elasticsearch-based vector store that implements the BaseVectorStore
+interface for high-performance dense vector storage and retrieval.
+"""
+
+from typing import Any
+
+from loguru import logger
+
+from .base_vector_store import BaseVectorStore
+from ..context import C
+from ..embedding import BaseEmbeddingModel
+from ..schema import VectorNode
+
+_ELASTICSEARCH_IMPORT_ERROR = None
+
+try:
+    from elasticsearch import AsyncElasticsearch
+    from elasticsearch.helpers import async_bulk
+except ImportError as e:
+    _ELASTICSEARCH_IMPORT_ERROR = e
+    AsyncElasticsearch = None
+    async_bulk = None
+
+
+@C.register_vector_store("es")
+class ESVectorStore(BaseVectorStore):
+    """Elasticsearch-based vector store for dense vector storage and kNN search."""
+
+    def __init__(
+        self,
+        collection_name: str,
+        embedding_model: BaseEmbeddingModel,
+        hosts: str | list[str] | None = None,
+        basic_auth: tuple[str, str] | None = None,
+        cloud_id: str | None = None,
+        api_key: str | None = None,
+        verify_certs: bool = True,
+        headers: dict[str, str] | None = None,
+        **kwargs,
+    ):
+        """Initialize the Elasticsearch client and vector store configuration.
+
+        Args:
+            collection_name: Name of the Elasticsearch index (converted to lowercase).
+            embedding_model: Model instance used to generate vector embeddings.
+            hosts: Connection host(s) for the Elasticsearch cluster.
+            basic_auth: Credentials for basic authentication.
+            cloud_id: Deployment ID for Elastic Cloud.
+            api_key: API key for authentication.
+            verify_certs: Enable or disable SSL certificate verification.
+            headers: Custom HTTP headers for requests.
+            **kwargs: Additional configuration passed to the base class.
+        """
+        if _ELASTICSEARCH_IMPORT_ERROR is not None:
+            raise ImportError(
+                "Elasticsearch requires extra dependencies. Install with `pip install elasticsearch`",
+            ) from _ELASTICSEARCH_IMPORT_ERROR
+
+        # Elasticsearch requires lowercase index names
+        collection_name = collection_name.lower()
+
+        super().__init__(collection_name=collection_name, embedding_model=embedding_model, **kwargs)
+
+        # 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."""
+        aliases = await self.client.indices.get_alias()
+        return list(aliases.keys())
+
+    async def create_collection(self, collection_name: str, **kwargs):
+        """Create a new index with dense vector mappings for kNN search.
+
+        Args:
+            collection_name: Name of the index to create.
+            **kwargs: Settings like dimensions, similarity, shards, and replicas.
+        """
+        collection_name = collection_name.lower()
+
+        if await self.client.indices.exists(index=collection_name):
+            return
+
+        dimensions = kwargs.get("dimensions", self.embedding_model.dimensions)
+        similarity = kwargs.get("similarity", "cosine")
+        number_of_shards = kwargs.get("number_of_shards", 5)
+        number_of_replicas = kwargs.get("number_of_replicas", 1)
+        refresh_interval = kwargs.get("refresh_interval", "1s")
+
+        index_settings = {
+            "settings": {
+                "index": {
+                    "number_of_replicas": number_of_replicas,
+                    "number_of_shards": number_of_shards,
+                    "refresh_interval": refresh_interval,
+                },
+            },
+            "mappings": {
+                "properties": {
+                    "vector_id": {"type": "keyword"},
+                    "content": {"type": "text"},
+                    "vector": {
+                        "type": "dense_vector",
+                        "dims": dimensions,
+                        "index": True,
+                        "similarity": similarity,
+                    },
+                    "metadata": {"type": "object", "enabled": True},
+                },
+            },
+        }
+
+        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")
+
+    async def delete_collection(self, collection_name: str, **kwargs):
+        """Permanently delete an Elasticsearch index.
+
+        Args:
+            collection_name: Name of the index to delete.
+            **kwargs: Additional parameters for the deletion request.
+        """
+        collection_name = collection_name.lower()
+
+        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")
+
+    async def copy_collection(self, collection_name: str, **kwargs):
+        """Reindex the current collection into a new index with identical mappings.
+
+        Args:
+            collection_name: Name of the destination index.
+            **kwargs: Additional parameters for the reindexing process.
+        """
+        collection_name = collection_name.lower()
+
+        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()
+        if "index" in settings_to_copy:
+            index_settings = settings_to_copy["index"].copy()
+            internal_keys = [
+                "uuid",
+                "creation_date",
+                "provided_name",
+                "version",
+                "store",
+                "routing",
+                "replication",
+            ]
+            for key in internal_keys:
+                index_settings.pop(key, None)
+            settings_to_copy["index"] = index_settings
+
+        await self.client.indices.create(
+            index=collection_name,
+            body={
+                "settings": settings_to_copy,
+                "mappings": current_settings.get("mappings", {}),
+            },
+        )
+
+        await self.client.reindex(
+            body={
+                "source": {"index": self.collection_name},
+                "dest": {"index": collection_name},
+            },
+        )
+
+        logger.info(f"Copied collection {self.collection_name} to {collection_name}")
+
+    async def insert(self, nodes: VectorNode | list[VectorNode], refresh: bool = True, **kwargs):
+        """Insert nodes into the index, generating embeddings if missing.
+
+        Args:
+            nodes: Single or multiple VectorNode objects to index.
+            refresh: If True, makes the operation visible to search immediately.
+            **kwargs: Additional insertion options.
+        """
+        if isinstance(nodes, VectorNode):
+            nodes = [nodes]
+
+        nodes_without_vectors = [node for node in nodes if node.vector is None]
+        if nodes_without_vectors:
+            nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors)
+            vector_map = {n.vector_id: n for n in nodes_with_vectors}
+            nodes_to_insert = [vector_map.get(n.vector_id, n) if n.vector is None else n for n in nodes]
+        else:
+            nodes_to_insert = nodes
+
+        actions = []
+        for node in nodes_to_insert:
+            action = {
+                "_index": self.collection_name,
+                "_id": node.vector_id,
+                "_source": {
+                    "vector_id": node.vector_id,
+                    "content": node.content,
+                    "vector": node.vector,
+                    "metadata": node.metadata,
+                },
+            }
+            actions.append(action)
+
+        success, failed = await async_bulk(self.client, actions, raise_on_error=False)
+
+        if failed:
+            logger.warning(f"Failed to insert {len(failed)} documents")
+
+        logger.info(f"Inserted {success} documents into {self.collection_name}")
+
+        if refresh:
+            await self.client.indices.refresh(index=self.collection_name)
+
+    async def search(
+        self,
+        query: str,
+        limit: int = 5,
+        filters: dict | None = None,
+        **kwargs,
+    ) -> list[VectorNode]:
+        """Perform a kNN similarity search based on a text query.
+
+        Args:
+            query: The text to search for.
+            limit: Maximum number of nearest neighbors to return.
+            filters: Metadata filters for exact match or 'IN' operations.
+            **kwargs: Search parameters like num_candidates or score_threshold.
+
+        Returns:
+            List of VectorNode objects ordered by similarity.
+        """
+        query_vector = await self.get_embedding(query)
+        num_candidates = kwargs.get("num_candidates", limit * 2)
+
+        search_query: dict = {
+            "knn": {
+                "field": "vector",
+                "query_vector": query_vector,
+                "k": limit,
+                "num_candidates": num_candidates,
+            },
+            "size": limit,
+        }
+
+        if filters:
+            filter_conditions = []
+            for key, value in filters.items():
+                if isinstance(value, list):
+                    filter_conditions.append({"terms": {f"metadata.{key}": value}})
+                else:
+                    filter_conditions.append({"term": {f"metadata.{key}": value}})
+            search_query["knn"]["filter"] = {"bool": {"must": filter_conditions}}
+
+        response = await self.client.search(index=self.collection_name, body=search_query)
+
+        results = []
+        for hit in response["hits"]["hits"]:
+            source = hit["_source"]
+            node = VectorNode(
+                vector_id=source.get("vector_id", hit["_id"]),
+                content=source.get("content", ""),
+                vector=source.get("vector"),
+                metadata=source.get("metadata", {}),
+            )
+            node.metadata["_score"] = hit["_score"]
+            results.append(node)
+
+        return results
+
+    async def delete(self, vector_ids: str | list[str], refresh: bool = True, **kwargs):
+        """Delete specific vectors from the index by their IDs.
+
+        Args:
+            vector_ids: Single ID or list of IDs to remove.
+            refresh: If True, refreshes the index after deletion.
+            **kwargs: Additional deletion parameters.
+        """
+        if isinstance(vector_ids, str):
+            vector_ids = [vector_ids]
+
+        actions = []
+        for vector_id in vector_ids:
+            actions.append(
+                {
+                    "_op_type": "delete",
+                    "_index": self.collection_name,
+                    "_id": vector_id,
+                },
+            )
+
+        success, failed = await async_bulk(
+            self.client,
+            actions,
+            raise_on_error=False,
+            raise_on_exception=False,
+        )
+
+        if failed:
+            logger.warning(f"Failed to delete {len(failed)} documents")
+
+        logger.info(f"Deleted {success} documents from {self.collection_name}")
+
+        if refresh:
+            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.
+
+        Args:
+            nodes: Single or multiple VectorNode objects with updated data.
+            refresh: If True, refreshes the index after update.
+            **kwargs: Additional update parameters.
+        """
+        if isinstance(nodes, VectorNode):
+            nodes = [nodes]
+
+        nodes_without_vectors = [node for node in nodes if node.vector is None and node.content]
+        if nodes_without_vectors:
+            nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors)
+            vector_map = {n.vector_id: n for n in nodes_with_vectors}
+            nodes_to_update = [vector_map.get(n.vector_id, n) if n.vector is None and n.content else n for n in nodes]
+        else:
+            nodes_to_update = nodes
+
+        actions = []
+        for node in nodes_to_update:
+            doc = {
+                "vector_id": node.vector_id,
+                "content": node.content,
+                "metadata": node.metadata,
+            }
+            if node.vector is not None:
+                doc["vector"] = node.vector
+
+            actions.append(
+                {
+                    "_op_type": "update",
+                    "_index": self.collection_name,
+                    "_id": node.vector_id,
+                    "doc": doc,
+                },
+            )
+
+        success, failed = await async_bulk(
+            self.client,
+            actions,
+            raise_on_error=False,
+            raise_on_exception=False,
+        )
+
+        if failed:
+            logger.warning(f"Failed to update {len(failed)} documents")
+
+        logger.info(f"Updated {success} documents in {self.collection_name}")
+
+        if refresh:
+            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.
+
+        Args:
+            vector_ids: Single ID or list of IDs to retrieve.
+
+        Returns:
+            A single VectorNode or a list of VectorNodes.
+        """
+        single_result = isinstance(vector_ids, str)
+        if single_result:
+            vector_ids = [vector_ids]
+
+        response = await self.client.mget(
+            index=self.collection_name,
+            body={"ids": vector_ids},
+        )
+
+        results = []
+        for doc in response["docs"]:
+            if doc.get("found"):
+                source = doc["_source"]
+                node = VectorNode(
+                    vector_id=source.get("vector_id", doc["_id"]),
+                    content=source.get("content", ""),
+                    vector=source.get("vector"),
+                    metadata=source.get("metadata", {}),
+                )
+                results.append(node)
+            else:
+                logger.warning(f"Document with ID {doc['_id']} not found")
+
+        return results[0] if single_result and results else results
+
+    async def list(
+        self,
+        filters: dict | None = None,
+        limit: int | None = None,
+    ) -> list[VectorNode]:
+        """Retrieve a list of nodes filtered by metadata or limit.
+
+        Args:
+            filters: Optional metadata filtering criteria.
+            limit: Maximum number of nodes to return.
+
+        Returns:
+            A list of matching VectorNode objects.
+        """
+        query: dict[str, Any] = {"query": {"match_all": {}}}
+
+        if filters:
+            filter_conditions = []
+            for key, value in filters.items():
+                if isinstance(value, list):
+                    filter_conditions.append({"terms": {f"metadata.{key}": value}})
+                else:
+                    filter_conditions.append({"term": {f"metadata.{key}": value}})
+            query["query"] = {"bool": {"must": filter_conditions}}
+
+        if limit:
+            query["size"] = limit
+        else:
+            query["size"] = 10000
+
+        response = await self.client.search(index=self.collection_name, body=query)
+
+        results = []
+        for hit in response["hits"]["hits"]:
+            source = hit["_source"]
+            node = VectorNode(
+                vector_id=source.get("vector_id", hit["_id"]),
+                content=source.get("content", ""),
+                vector=source.get("vector"),
+                metadata=source.get("metadata", {}),
+            )
+            results.append(node)
+
+        return results
+
+    async def close(self):
+        """Terminate the Elasticsearch client session and release resources."""
+        await self.client.close()
+        logger.info("Elasticsearch client connection closed")

reme_ai/core/vector_store/local_vector_store.py

@@ -0,0 +1,281 @@
+"""Local file system vector store implementation for ReMe."""
+
+import json
+from pathlib import Path
+
+from loguru import logger
+
+from .base_vector_store import BaseVectorStore
+from ..context import C
+from ..embedding import BaseEmbeddingModel
+from ..schema import VectorNode
+
+
+@C.register_vector_store("local")
+class LocalVectorStore(BaseVectorStore):
+    """Local file system-based vector store using JSON files and manual cosine similarity."""
+
+    def __init__(
+        self,
+        collection_name: str,
+        embedding_model: BaseEmbeddingModel,
+        root_path: str = "./local_vector_store",
+        **kwargs,
+    ):
+        """Initialize the local vector store with a root path and collection name."""
+        super().__init__(collection_name=collection_name, embedding_model=embedding_model, **kwargs)
+        self.root_path = Path(root_path)
+        self.collection_path = self.root_path / collection_name
+        self.root_path.mkdir(parents=True, exist_ok=True)
+
+    def _get_collection_path(self, collection_name: str) -> Path:
+        """Get the file system path for a specific collection."""
+        return self.root_path / collection_name
+
+    def _get_node_file_path(self, vector_id: str, collection_name: str | None = None) -> Path:
+        """Get the JSON file path for a specific vector node."""
+        col_path = self._get_collection_path(collection_name or self.collection_name)
+        return col_path / f"{vector_id}.json"
+
+    def _save_node(self, node: VectorNode, collection_name: str | None = None):
+        """Save a vector node to a JSON file on disk."""
+        file_path = self._get_node_file_path(node.vector_id, collection_name)
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+
+        with open(file_path, "w", encoding="utf-8") as f:
+            json.dump(node.model_dump(), f, ensure_ascii=False, indent=2)
+
+    def _load_node(self, vector_id: str, collection_name: str | None = None) -> VectorNode | None:
+        """Load a vector node from a JSON file."""
+        file_path = self._get_node_file_path(vector_id, collection_name)
+
+        if not file_path.exists():
+            return None
+
+        with open(file_path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+            return VectorNode(**data)
+
+    def _load_all_nodes(self, collection_name: str | None = None) -> list[VectorNode]:
+        """Load all vector nodes existing in a collection."""
+        col_path = self._get_collection_path(collection_name or self.collection_name)
+
+        if not col_path.exists():
+            return []
+
+        nodes = []
+        for file_path in col_path.glob("*.json"):
+            try:
+                with open(file_path, "r", encoding="utf-8") as f:
+                    data = json.load(f)
+                    nodes.append(VectorNode(**data))
+            except Exception as e:
+                logger.warning(f"Failed to load node from {file_path}: {e}")
+
+        return nodes
+
+    @staticmethod
+    def _cosine_similarity(vec1: list[float], vec2: list[float]) -> float:
+        """Calculate the cosine similarity between two numeric vectors."""
+        if len(vec1) != len(vec2):
+            raise ValueError(f"Vectors must have same length: {len(vec1)} != {len(vec2)}")
+
+        dot_product = sum(a * b for a, b in zip(vec1, vec2))
+        magnitude1 = sum(a * a for a in vec1) ** 0.5
+        magnitude2 = sum(b * b for b in vec2) ** 0.5
+
+        if magnitude1 == 0 or magnitude2 == 0:
+            return 0.0
+
+        return dot_product / (magnitude1 * magnitude2)
+
+    @staticmethod
+    def _match_filters(node: VectorNode, filters: dict | None) -> bool:
+        """Check if a vector node matches the provided metadata filters."""
+        if not filters:
+            return True
+
+        for key, value in filters.items():
+            node_value = node.metadata.get(key)
+
+            if isinstance(value, list):
+                if node_value not in value:
+                    return False
+            else:
+                if node_value != value:
+                    return False
+
+        return True
+
+    async def list_collections(self) -> list[str]:
+        """List all collection directories in the root path."""
+        if not self.root_path.exists():
+            return []
+
+        return [d.name for d in self.root_path.iterdir() if d.is_dir() and not d.name.startswith(".")]
+
+    async def create_collection(self, collection_name: str, **kwargs):
+        """Create a new collection directory."""
+        col_path = self._get_collection_path(collection_name)
+        col_path.mkdir(parents=True, exist_ok=True)
+        logger.info(f"Created collection {collection_name} at {col_path}")
+
+    async def delete_collection(self, collection_name: str, **kwargs):
+        """Delete a collection directory and all its JSON files."""
+        col_path = self._get_collection_path(collection_name)
+
+        if not col_path.exists():
+            logger.warning(f"Collection {collection_name} does not exist")
+            return
+
+        for file_path in col_path.glob("*.json"):
+            file_path.unlink()
+
+        col_path.rmdir()
+        logger.info(f"Deleted collection {collection_name}")
+
+    async def copy_collection(self, collection_name: str, **kwargs):
+        """Copy all nodes from the current collection to a new one."""
+        source_path = self._get_collection_path(self.collection_name)
+        target_path = self._get_collection_path(collection_name)
+
+        if not source_path.exists():
+            logger.warning(f"Source collection {self.collection_name} does not exist")
+            return
+
+        target_path.mkdir(parents=True, exist_ok=True)
+
+        for file_path in source_path.glob("*.json"):
+            target_file = target_path / file_path.name
+            target_file.write_text(file_path.read_text(encoding="utf-8"), encoding="utf-8")
+
+        logger.info(f"Copied collection {self.collection_name} to {collection_name}")
+
+    async def insert(self, nodes: VectorNode | list[VectorNode], **kwargs):
+        """Insert vector nodes into the local store, generating embeddings if necessary."""
+        if isinstance(nodes, VectorNode):
+            nodes = [nodes]
+
+        nodes_without_vectors = [node for node in nodes if node.vector is None]
+        if nodes_without_vectors:
+            nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors)
+            vector_map = {n.vector_id: n for n in nodes_with_vectors}
+            nodes_to_insert = [vector_map.get(n.vector_id, n) if n.vector is None else n for n in nodes]
+        else:
+            nodes_to_insert = nodes
+
+        for node in nodes_to_insert:
+            self._save_node(node)
+
+        logger.info(f"Inserted {len(nodes_to_insert)} nodes into {self.collection_name}")
+
+    async def search(
+        self,
+        query: str,
+        limit: int = 5,
+        filters: dict | None = None,
+        **kwargs,
+    ) -> list[VectorNode]:
+        """Search for nodes similar to the query using brute-force cosine similarity."""
+        query_vector = await self.get_embedding(query)
+        all_nodes = self._load_all_nodes()
+        filtered_nodes = [node for node in all_nodes if self._match_filters(node, filters)]
+
+        scored_nodes = []
+        for node in filtered_nodes:
+            if node.vector is None:
+                logger.warning(f"Node {node.vector_id} has no vector, skipping")
+                continue
+
+            try:
+                score = self._cosine_similarity(query_vector, node.vector)
+                scored_nodes.append((node, score))
+            except ValueError as e:
+                logger.warning(f"Failed to calculate similarity for node {node.vector_id}: {e}")
+
+        scored_nodes.sort(key=lambda x: x[1], reverse=True)
+
+        score_threshold = kwargs.get("score_threshold")
+        if score_threshold is not None:
+            scored_nodes = [(node, score) for node, score in scored_nodes if score >= score_threshold]
+
+        scored_nodes = scored_nodes[:limit]
+        results = []
+        for node, score in scored_nodes:
+            node.metadata["_score"] = score
+            results.append(node)
+
+        return results
+
+    async def delete(self, vector_ids: str | list[str], **kwargs):
+        """Delete specific vector nodes by their IDs."""
+        if isinstance(vector_ids, str):
+            vector_ids = [vector_ids]
+
+        deleted_count = 0
+        for vector_id in vector_ids:
+            file_path = self._get_node_file_path(vector_id)
+            if file_path.exists():
+                file_path.unlink()
+                deleted_count += 1
+            else:
+                logger.warning(f"Node {vector_id} does not exist")
+
+        logger.info(f"Deleted {deleted_count} nodes from {self.collection_name}")
+
+    async def update(self, nodes: VectorNode | list[VectorNode], **kwargs):
+        """Update existing vector nodes with new data or embeddings."""
+        if isinstance(nodes, VectorNode):
+            nodes = [nodes]
+
+        nodes_without_vectors = [node for node in nodes if node.vector is None and node.content]
+        if nodes_without_vectors:
+            nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors)
+            vector_map = {n.vector_id: n for n in nodes_with_vectors}
+            nodes_to_update = [vector_map.get(n.vector_id, n) if n.vector is None and n.content else n for n in nodes]
+        else:
+            nodes_to_update = nodes
+
+        updated_count = 0
+        for node in nodes_to_update:
+            file_path = self._get_node_file_path(node.vector_id)
+            if file_path.exists():
+                self._save_node(node)
+                updated_count += 1
+            else:
+                logger.warning(f"Node {node.vector_id} does not exist, skipping update")
+
+        logger.info(f"Updated {updated_count} nodes in {self.collection_name}")
+
+    async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode]:
+        """Retrieve one or more vector nodes by their unique IDs."""
+        is_single = isinstance(vector_ids, str)
+        ids = [vector_ids] if is_single else vector_ids
+
+        results = []
+        for vector_id in ids:
+            node = self._load_node(vector_id)
+            if node:
+                results.append(node)
+            else:
+                logger.warning(f"Node {vector_id} not found")
+
+        return results[0] if is_single and results else results
+
+    async def list(
+        self,
+        filters: dict | None = None,
+        limit: int | None = None,
+    ) -> list[VectorNode]:
+        """List vector nodes in the collection with optional filtering and limits."""
+        all_nodes = self._load_all_nodes()
+        filtered_nodes = [node for node in all_nodes if self._match_filters(node, filters)]
+
+        if limit is not None:
+            filtered_nodes = filtered_nodes[:limit]
+
+        return filtered_nodes
+
+    async def close(self):
+        """Close the vector store (no-op for local file system)."""
+        logger.info("Local vector store closed")

reme_ai/core/vector_store/pgvector_store.py

@@ -0,0 +1,533 @@
+"""PostgreSQL pgvector implementation for vector storage and retrieval."""
+
+import json
+from typing import Any
+
+from loguru import logger
+
+from .base_vector_store import BaseVectorStore
+from ..context import C
+from ..embedding import BaseEmbeddingModel
+from ..schema import VectorNode
+
+_ASYNCPG_IMPORT_ERROR = None
+
+try:
+    import asyncpg
+    from asyncpg import Pool
+except ImportError as e:
+    _ASYNCPG_IMPORT_ERROR = e
+    asyncpg = None
+    Pool = None
+
+
+@C.register_vector_store("pgvector")
+class PGVectorStore(BaseVectorStore):
+    """Vector store implementation using PostgreSQL and pgvector for efficient similarity search."""
+
+    def __init__(
+        self,
+        collection_name: str,
+        embedding_model: BaseEmbeddingModel,
+        host: str = "localhost",
+        port: int = 5432,
+        database: str = "postgres",
+        user: str = "postgres",
+        password: str = "",
+        min_size: int = 1,
+        max_size: int = 10,
+        dsn: str | None = None,
+        use_hnsw: bool = True,
+        use_diskann: bool = False,
+        **kwargs,
+    ):
+        """Initialize the PGVector store with connection parameters and index settings."""
+        if _ASYNCPG_IMPORT_ERROR is not None:
+            raise ImportError(
+                "PGVector requires extra dependencies. Install with `pip install asyncpg pgvector`",
+            ) from _ASYNCPG_IMPORT_ERROR
+
+        super().__init__(collection_name=collection_name, embedding_model=embedding_model, **kwargs)
+
+        self.dsn = dsn
+        self.host = host
+        self.port = port
+        self.database = database
+        self.user = user
+        self.password = password
+        self.min_size = min_size
+        self.max_size = max_size
+        self.use_hnsw = use_hnsw
+        self.use_diskann = use_diskann
+        self._pool: Pool | None = None
+        self.embedding_model_dims = embedding_model.dimensions
+
+    async def _get_pool(self) -> Pool:
+        """Create or return the existing asyncpg connection pool."""
+        if self._pool is None:
+            if self.dsn:
+                self._pool = await asyncpg.create_pool(
+                    dsn=self.dsn,
+                    min_size=self.min_size,
+                    max_size=self.max_size,
+                )
+            else:
+                self._pool = await asyncpg.create_pool(
+                    host=self.host,
+                    port=self.port,
+                    database=self.database,
+                    user=self.user,
+                    password=self.password,
+                    min_size=self.min_size,
+                    max_size=self.max_size,
+                )
+
+            async with self._pool.acquire() as conn:
+                await conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
+
+            logger.info(f"PGVector connection pool created for database {self.database}")
+
+        return self._pool
+
+    async def _ensure_collection_exists(self):
+        """Check if the collection table exists and create it if missing."""
+        collections = await self.list_collections()
+        if self.collection_name not in collections:
+            await self.create_collection(self.collection_name)
+
+    async def list_collections(self) -> list[str]:
+        """List all available table names in the current database."""
+        pool = await self._get_pool()
+        async with pool.acquire() as conn:
+            rows = await conn.fetch(
+                "SELECT table_name FROM information_schema.tables WHERE table_schema = 'public'",
+            )
+            return [row["table_name"] for row in rows]
+
+    async def create_collection(self, collection_name: str, **kwargs):
+        """Create a new PostgreSQL table with vector support and appropriate indexing."""
+        pool = await self._get_pool()
+        dimensions = kwargs.get("dimensions", self.embedding_model_dims)
+
+        async with pool.acquire() as conn:
+            await conn.execute(
+                f"""
+                CREATE TABLE IF NOT EXISTS {collection_name} (
+                    id TEXT PRIMARY KEY,
+                    content TEXT,
+                    vector vector({dimensions}),
+                    metadata JSONB
+                )
+            """,
+            )
+
+            if self.use_diskann and dimensions < 2000:
+                result = await conn.fetchval(
+                    "SELECT 1 FROM pg_extension WHERE extname = 'vectorscale'",
+                )
+                if result:
+                    await conn.execute(
+                        f"""
+                        CREATE INDEX IF NOT EXISTS {collection_name}_diskann_idx
+                        ON {collection_name}
+                        USING diskann (vector)
+                    """,
+                    )
+                    logger.info(f"Created DiskANN index for collection {collection_name}")
+                else:
+                    logger.warning("vectorscale extension not available, skipping DiskANN index")
+            elif self.use_hnsw:
+                await conn.execute(
+                    f"""
+                    CREATE INDEX IF NOT EXISTS {collection_name}_hnsw_idx
+                    ON {collection_name}
+                    USING hnsw (vector vector_cosine_ops)
+                """,
+                )
+                logger.info(f"Created HNSW index for collection {collection_name}")
+
+        logger.info(f"Created collection {collection_name} with dimensions={dimensions}")
+
+    async def delete_collection(self, collection_name: str, **kwargs):
+        """Remove the specified collection table from the database."""
+        pool = await self._get_pool()
+        async with pool.acquire() as conn:
+            await conn.execute(f"DROP TABLE IF EXISTS {collection_name}")
+        logger.info(f"Deleted collection {collection_name}")
+
+    async def copy_collection(self, collection_name: str, **kwargs):
+        """Duplicate the structure and content of the current collection to a new table."""
+        pool = await self._get_pool()
+
+        async with pool.acquire() as conn:
+            columns = await conn.fetch(
+                """
+                SELECT column_name, data_type, udt_name
+                FROM information_schema.columns
+                WHERE table_name = $1 AND table_schema = 'public'
+            """,
+                self.collection_name,
+            )
+
+            if not columns:
+                raise ValueError(f"Source collection {self.collection_name} does not exist")
+
+            await conn.execute(f"CREATE TABLE {collection_name} AS TABLE {self.collection_name}")
+            await conn.execute(f"ALTER TABLE {collection_name} ADD PRIMARY KEY (id)")
+
+            if self.use_hnsw:
+                await conn.execute(
+                    f"""
+                    CREATE INDEX IF NOT EXISTS {collection_name}_hnsw_idx
+                    ON {collection_name}
+                    USING hnsw (vector vector_cosine_ops)
+                """,
+                )
+
+        logger.info(f"Copied collection {self.collection_name} to {collection_name}")
+
+    async def insert(self, nodes: VectorNode | list[VectorNode], **kwargs):
+        """Insert or upsert vector nodes into the PostgreSQL collection."""
+        await self._ensure_collection_exists()
+
+        if isinstance(nodes, VectorNode):
+            nodes = [nodes]
+
+        if not nodes:
+            return
+
+        nodes_without_vectors = [node for node in nodes if node.vector is None]
+        if nodes_without_vectors:
+            nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors)
+            vector_map = {n.vector_id: n for n in nodes_with_vectors}
+            nodes_to_insert = [vector_map.get(n.vector_id, n) if n.vector is None else n for n in nodes]
+        else:
+            nodes_to_insert = nodes
+
+        pool = await self._get_pool()
+        data = [
+            (
+                node.vector_id,
+                node.content,
+                f"[{','.join(map(str, node.vector))}]",
+                json.dumps(node.metadata),
+            )
+            for node in nodes_to_insert
+        ]
+
+        async with pool.acquire() as conn:
+            on_conflict = kwargs.get("on_conflict", "update")
+
+            if on_conflict == "update":
+                await conn.executemany(
+                    f"""
+                    INSERT INTO {self.collection_name} (id, content, vector, metadata)
+                    VALUES ($1, $2, $3::vector, $4::jsonb)
+                    ON CONFLICT (id) DO UPDATE SET
+                        content = EXCLUDED.content,
+                        vector = EXCLUDED.vector,
+                        metadata = EXCLUDED.metadata
+                """,
+                    data,
+                )
+            elif on_conflict == "ignore":
+                await conn.executemany(
+                    f"""
+                    INSERT INTO {self.collection_name} (id, content, vector, metadata)
+                    VALUES ($1, $2, $3::vector, $4::jsonb)
+                    ON CONFLICT (id) DO NOTHING
+                """,
+                    data,
+                )
+            else:
+                await conn.executemany(
+                    f"""
+                    INSERT INTO {self.collection_name} (id, content, vector, metadata)
+                    VALUES ($1, $2, $3::vector, $4::jsonb)
+                """,
+                    data,
+                )
+
+        logger.info(f"Inserted {len(nodes_to_insert)} documents into {self.collection_name}")
+
+    @staticmethod
+    def _build_filter_clause(filters: dict | None) -> tuple[str, list]:
+        """Generate an SQL WHERE clause and parameter list from a filter dictionary."""
+        if not filters:
+            return "", []
+
+        conditions = []
+        params = []
+        param_idx = 1
+
+        for key, value in filters.items():
+            if isinstance(value, list):
+                placeholders = ", ".join([f"${param_idx + i}" for i in range(len(value))])
+                conditions.append(f"metadata->>'{key}' IN ({placeholders})")
+                params.extend([str(v) for v in value])
+                param_idx += len(value)
+            else:
+                conditions.append(f"metadata->>'{key}' = ${param_idx}")
+                params.append(str(value))
+                param_idx += 1
+
+        filter_clause = "WHERE " + " AND ".join(conditions) if conditions else ""
+        return filter_clause, params
+
+    async def search(
+        self,
+        query: str,
+        limit: int = 5,
+        filters: dict | None = None,
+        **kwargs,
+    ) -> list[VectorNode]:
+        """Perform vector similarity search with optional metadata filtering."""
+        await self._ensure_collection_exists()
+
+        query_vector = await self.get_embedding(query)
+        vector_str = f"[{','.join(map(str, query_vector))}]"
+        pool = await self._get_pool()
+
+        filter_clause, filter_params = self._build_filter_clause(filters)
+
+        if filter_clause:
+            for i in range(len(filter_params)):
+                old_idx = i + 1
+                new_idx = i + 2
+                filter_clause = filter_clause.replace(f"${old_idx}", f"${new_idx}", 1)
+
+        async with pool.acquire() as conn:
+            sql = f"""
+                SELECT id, content, vector, metadata, vector <=> $1::vector AS distance
+                FROM {self.collection_name}
+                {filter_clause}
+                ORDER BY distance
+                LIMIT ${len(filter_params) + 2}
+            """
+            rows = await conn.fetch(sql, vector_str, *filter_params, limit)
+
+        results = []
+        score_threshold = kwargs.get("score_threshold")
+
+        for row in rows:
+            distance = row["distance"]
+            if score_threshold is not None and distance > score_threshold:
+                continue
+
+            vector_data = None
+            if row["vector"]:
+                vector_str_raw = str(row["vector"])
+                if vector_str_raw.startswith("[") and vector_str_raw.endswith("]"):
+                    vector_data = [float(x) for x in vector_str_raw[1:-1].split(",")]
+
+            metadata = row["metadata"] if row["metadata"] else {}
+            if isinstance(metadata, str):
+                metadata = json.loads(metadata)
+
+            metadata["_score"] = 1 - distance
+            metadata["_distance"] = distance
+
+            node = VectorNode(
+                vector_id=row["id"],
+                content=row["content"] or "",
+                vector=vector_data,
+                metadata=metadata,
+            )
+            results.append(node)
+
+        return results
+
+    async def delete(self, vector_ids: str | list[str], **kwargs):
+        """Remove specific vector records from the collection by their IDs."""
+        await self._ensure_collection_exists()
+
+        if isinstance(vector_ids, str):
+            vector_ids = [vector_ids]
+
+        if not vector_ids:
+            return
+
+        pool = await self._get_pool()
+        async with pool.acquire() as conn:
+            placeholders = ", ".join([f"${i + 1}" for i in range(len(vector_ids))])
+            await conn.execute(
+                f"DELETE FROM {self.collection_name} WHERE id IN ({placeholders})",
+                *vector_ids,
+            )
+
+        logger.info(f"Deleted {len(vector_ids)} documents from {self.collection_name}")
+
+    async def update(self, nodes: VectorNode | list[VectorNode], **kwargs):
+        """Update existing vector nodes with new content, embeddings, or metadata."""
+        await self._ensure_collection_exists()
+
+        if isinstance(nodes, VectorNode):
+            nodes = [nodes]
+
+        if not nodes:
+            return
+
+        nodes_without_vectors = [node for node in nodes if node.vector is None and node.content]
+        if nodes_without_vectors:
+            nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors)
+            vector_map = {n.vector_id: n for n in nodes_with_vectors}
+            nodes_to_update = [vector_map.get(n.vector_id, n) if n.vector is None and n.content else n for n in nodes]
+        else:
+            nodes_to_update = nodes
+
+        pool = await self._get_pool()
+        async with pool.acquire() as conn:
+            for node in nodes_to_update:
+                update_fields = []
+                params = []
+                idx = 1
+
+                if node.content:
+                    update_fields.append(f"content = ${idx}")
+                    params.append(node.content)
+                    idx += 1
+
+                if node.vector:
+                    vector_str = f"[{','.join(map(str, node.vector))}]"
+                    update_fields.append(f"vector = ${idx}::vector")
+                    params.append(vector_str)
+                    idx += 1
+
+                if node.metadata:
+                    update_fields.append(f"metadata = ${idx}::jsonb")
+                    params.append(json.dumps(node.metadata))
+                    idx += 1
+
+                if update_fields:
+                    params.append(node.vector_id)
+                    await conn.execute(
+                        f"UPDATE {self.collection_name} SET {', '.join(update_fields)} WHERE id = ${idx}",
+                        *params,
+                    )
+
+        logger.info(f"Updated {len(nodes_to_update)} documents in {self.collection_name}")
+
+    async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode] | None:
+        """Retrieve vector nodes by their unique identifiers."""
+        await self._ensure_collection_exists()
+
+        single_result = isinstance(vector_ids, str)
+        if single_result:
+            vector_ids = [vector_ids]
+
+        if not vector_ids:
+            return [] if not single_result else None
+
+        pool = await self._get_pool()
+        async with pool.acquire() as conn:
+            placeholders = ", ".join([f"${i + 1}" for i in range(len(vector_ids))])
+            rows = await conn.fetch(
+                f"SELECT id, content, vector, metadata FROM {self.collection_name} WHERE id IN ({placeholders})",
+                *vector_ids,
+            )
+
+        results = []
+        for row in rows:
+            vector_data = None
+            if row["vector"]:
+                vector_str_raw = str(row["vector"])
+                if vector_str_raw.startswith("[") and vector_str_raw.endswith("]"):
+                    vector_data = [float(x) for x in vector_str_raw[1:-1].split(",")]
+
+            metadata = row["metadata"] if row["metadata"] else {}
+            if isinstance(metadata, str):
+                metadata = json.loads(metadata)
+
+            results.append(
+                VectorNode(
+                    vector_id=row["id"],
+                    content=row["content"] or "",
+                    vector=vector_data,
+                    metadata=metadata,
+                ),
+            )
+
+        if single_result:
+            return results[0] if results else None
+        return results
+
+    async def list(
+        self,
+        filters: dict | None = None,
+        limit: int | None = None,
+    ) -> list[VectorNode]:
+        """Return a list of vector nodes matching the provided filters and limit."""
+        await self._ensure_collection_exists()
+
+        pool = await self._get_pool()
+        filter_clause, filter_params = self._build_filter_clause(filters)
+
+        limit_clause = ""
+        if limit:
+            limit_clause = f"LIMIT ${len(filter_params) + 1}"
+            filter_params.append(limit)
+
+        async with pool.acquire() as conn:
+            sql = f"""
+                SELECT id, content, vector, metadata
+                FROM {self.collection_name}
+                {filter_clause}
+                {limit_clause}
+            """
+            rows = await conn.fetch(sql, *filter_params)
+
+        results = []
+        for row in rows:
+            vector_data = None
+            if row["vector"]:
+                vector_str_raw = str(row["vector"])
+                if vector_str_raw.startswith("[") and vector_str_raw.endswith("]"):
+                    vector_data = [float(x) for x in vector_str_raw[1:-1].split(",")]
+
+            metadata = row["metadata"] if row["metadata"] else {}
+            if isinstance(metadata, str):
+                metadata = json.loads(metadata)
+
+            results.append(
+                VectorNode(
+                    vector_id=row["id"],
+                    content=row["content"] or "",
+                    vector=vector_data,
+                    metadata=metadata,
+                ),
+            )
+
+        return results
+
+    async def collection_info(self) -> dict[str, Any]:
+        """Fetch metadata including record count and disk usage for the collection."""
+        pool = await self._get_pool()
+
+        async with pool.acquire() as conn:
+            row = await conn.fetchrow(
+                f"""
+                SELECT
+                    '{self.collection_name}' as name,
+                    (SELECT COUNT(*) FROM {self.collection_name}) as row_count,
+                    pg_size_pretty(pg_total_relation_size('{self.collection_name}')) as total_size
+            """,
+            )
+
+        return {
+            "name": row["name"],
+            "count": row["row_count"],
+            "size": row["total_size"],
+        }
+
+    async def reset(self):
+        """Purge all data by dropping and recreating the collection table."""
+        logger.warning(f"Resetting collection {self.collection_name}...")
+        await self.delete_collection(self.collection_name)
+        await self.create_collection(self.collection_name)
+
+    async def close(self):
+        """Terminate the database connection pool and release associated resources."""
+        if self._pool is not None:
+            await self._pool.close()
+            self._pool = None
+            logger.info("PGVector connection pool closed")

reme_ai/core/vector_store/qdrant_vector_store.py

@@ -0,0 +1,444 @@
+"""Qdrant vector store implementation for the ReMe project."""
+
+from typing import Any
+
+from loguru import logger
+
+from .base_vector_store import BaseVectorStore
+from ..context import C
+from ..embedding import BaseEmbeddingModel
+from ..schema import VectorNode
+
+_QDRANT_IMPORT_ERROR = None
+
+try:
+    from qdrant_client import AsyncQdrantClient
+    from qdrant_client.models import (
+        Distance,
+        FieldCondition,
+        Filter,
+        MatchValue,
+        PointIdsList,
+        PointStruct,
+        Range,
+        VectorParams,
+    )
+except ImportError as e:
+    _QDRANT_IMPORT_ERROR = e
+    AsyncQdrantClient = None
+    Distance = None
+    FieldCondition = None
+    Filter = None
+    MatchValue = None
+    PointIdsList = None
+    PointStruct = None
+    Range = None
+    VectorParams = None
+
+
+@C.register_vector_store("qdrant")
+class QdrantVectorStore(BaseVectorStore):
+    """Vector store implementation using Qdrant for dense vector search."""
+
+    def __init__(
+        self,
+        collection_name: str,
+        embedding_model: BaseEmbeddingModel,
+        host: str | None = None,
+        port: int = 6333,
+        path: str | None = None,
+        url: str | None = None,
+        api_key: str | None = None,
+        https: bool | None = None,
+        grpc_port: int = 6334,
+        prefer_grpc: bool = False,
+        distance: str = "cosine",
+        on_disk: bool = False,
+        **kwargs: Any,
+    ):
+        """Initialize the Qdrant client and collection configuration.
+
+        Args:
+            collection_name: Name of the collection.
+            embedding_model: Model used for generating vector embeddings.
+            host: Server host address.
+            port: HTTP port for the server.
+            path: Local storage path for on-disk/in-memory mode.
+            url: Full connection URL.
+            api_key: Authentication key for Qdrant Cloud.
+            https: Use secure connection if True.
+            grpc_port: gRPC interface port.
+            prefer_grpc: Use gRPC instead of HTTP if True.
+            distance: Metric for similarity (cosine, euclid, dot).
+            on_disk: Enable persistent storage for vectors.
+            **kwargs: Additional client configuration.
+        """
+        if _QDRANT_IMPORT_ERROR is not None:
+            raise ImportError(
+                "Qdrant requires extra dependencies. Install with `pip install qdrant-client`",
+            ) from _QDRANT_IMPORT_ERROR
+
+        super().__init__(collection_name=collection_name, embedding_model=embedding_model, **kwargs)
+
+        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,
+            **kwargs,
+        )
+
+        self.is_local = path is not None
+        distance_map = {
+            "cosine": Distance.COSINE,
+            "euclid": Distance.EUCLID,
+            "dot": Distance.DOT,
+        }
+        self.distance = distance_map.get(distance.lower(), Distance.COSINE)
+        self.on_disk = on_disk
+
+    async def list_collections(self) -> list[str]:
+        """Retrieve names of all existing collections in the Qdrant instance."""
+        collections = await self.client.get_collections()
+        return [collection.name for collection in collections.collections]
+
+    async def create_collection(self, collection_name: str, **kwargs: Any):
+        """Create a new collection with the specified vector configuration.
+
+        Args:
+            collection_name: Name of the collection to create.
+            **kwargs: Overrides for dimensions, distance, or on_disk settings.
+        """
+        collections = await self.list_collections()
+        if collection_name in collections:
+            logger.info(f"Collection {collection_name} already exists")
+            return
+
+        dimensions = kwargs.get("dimensions", self.embedding_model.dimensions)
+        distance = kwargs.get("distance", self.distance)
+        on_disk = kwargs.get("on_disk", self.on_disk)
+
+        await self.client.create_collection(
+            collection_name=collection_name,
+            vectors_config=VectorParams(
+                size=dimensions,
+                distance=distance,
+                on_disk=on_disk,
+            ),
+        )
+
+        logger.info(f"Created collection {collection_name} with dimensions={dimensions}")
+
+        if not self.is_local:
+            await self._create_payload_indexes(collection_name)
+
+    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"]
+
+        for field in common_fields:
+            try:
+                await self.client.create_payload_index(
+                    collection_name=collection_name,
+                    field_name=field,
+                    field_schema="keyword",
+                )
+                logger.debug(f"Created index for {field} in collection {collection_name}")
+            except Exception as e:
+                logger.debug(f"Index for {field} might already exist: {e}")
+
+    async def delete_collection(self, collection_name: str, **kwargs: Any):
+        """Permanently remove a collection from the Qdrant instance."""
+        collections = await self.list_collections()
+        if collection_name in collections:
+            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."""
+        collection_info = await self.client.get_collection(collection_name=self.collection_name)
+
+        await self.client.create_collection(
+            collection_name=collection_name,
+            vectors_config=collection_info.config.params.vectors,
+        )
+
+        offset = None
+        batch_size = 100
+
+        while True:
+            records, next_offset = await self.client.scroll(
+                collection_name=self.collection_name,
+                limit=batch_size,
+                offset=offset,
+                with_payload=True,
+                with_vectors=True,
+            )
+
+            if not records:
+                break
+
+            points = [
+                PointStruct(
+                    id=record.id,
+                    vector=record.vector,
+                    payload=record.payload,
+                )
+                for record in records
+            ]
+
+            await self.client.upsert(
+                collection_name=collection_name,
+                points=points,
+            )
+
+            offset = next_offset
+            if offset is None:
+                break
+
+        logger.info(f"Copied collection {self.collection_name} to {collection_name}")
+
+    async def insert(self, nodes: VectorNode | list[VectorNode], **kwargs: Any):
+        """Insert vector nodes into the collection, generating embeddings as needed."""
+        if isinstance(nodes, VectorNode):
+            nodes = [nodes]
+
+        nodes_without_vectors = [node for node in nodes if node.vector is None]
+        if nodes_without_vectors:
+            nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors)
+            vector_map = {n.vector_id: n for n in nodes_with_vectors}
+            nodes_to_insert = [vector_map.get(n.vector_id, n) if n.vector is None else n for n in nodes]
+        else:
+            nodes_to_insert = nodes
+
+        points = []
+        for node in nodes_to_insert:
+            try:
+                point_id = int(node.vector_id)
+            except ValueError:
+                point_id = abs(hash(node.vector_id)) % (10**18)
+
+            point = PointStruct(
+                id=point_id,
+                vector=node.vector,
+                payload={
+                    "vector_id": node.vector_id,
+                    "content": node.content,
+                    "metadata": node.metadata,
+                },
+            )
+            points.append(point)
+
+        wait = kwargs.get("wait", True)
+        await self.client.upsert(
+            collection_name=self.collection_name,
+            points=points,
+            wait=wait,
+        )
+
+        logger.info(f"Inserted {len(points)} documents into {self.collection_name}")
+
+    @staticmethod
+    def _create_filter(filters: dict) -> Filter | None:
+        """Convert a dictionary of filter conditions into a Qdrant Filter object."""
+        if not filters:
+            return None
+
+        conditions = []
+        for key, value in filters.items():
+            if isinstance(value, dict) and ("gte" in value or "lte" in value):
+                range_params = {}
+                if "gte" in value:
+                    range_params["gte"] = value["gte"]
+                if "lte" in value:
+                    range_params["lte"] = value["lte"]
+                conditions.append(
+                    FieldCondition(
+                        key=f"metadata.{key}",
+                        range=Range(**range_params),
+                    ),
+                )
+            elif isinstance(value, list):
+                conditions.append(
+                    FieldCondition(key=f"metadata.{key}", match=MatchValue(value=value[0])),
+                )
+            else:
+                conditions.append(
+                    FieldCondition(key=f"metadata.{key}", match=MatchValue(value=value)),
+                )
+
+        return Filter(must=conditions) if conditions else None
+
+    async def search(
+        self,
+        query: str,
+        limit: int = 5,
+        filters: dict | None = None,
+        **kwargs: Any,
+    ) -> list[VectorNode]:
+        """Search for the most similar vectors based on a text query."""
+        query_vector = await self.get_embedding(query)
+        query_filter = self._create_filter(filters) if filters else None
+        score_threshold = kwargs.get("score_threshold", None)
+
+        results = await self.client.query_points(
+            collection_name=self.collection_name,
+            query=query_vector,
+            query_filter=query_filter,
+            limit=limit,
+            score_threshold=score_threshold,
+        )
+
+        nodes = []
+        for point in results.points:
+            payload = point.payload or {}
+            node = VectorNode(
+                vector_id=payload.get("vector_id", str(point.id)),
+                content=payload.get("content", ""),
+                vector=point.vector if hasattr(point, "vector") else None,
+                metadata=payload.get("metadata", {}),
+            )
+            node.metadata["_score"] = point.score
+            nodes.append(node)
+
+        return nodes
+
+    async def delete(self, vector_ids: str | list[str], **kwargs: Any):
+        """Delete specific vectors from the collection using their IDs."""
+        if isinstance(vector_ids, str):
+            vector_ids = [vector_ids]
+
+        point_ids = []
+        for vector_id in vector_ids:
+            try:
+                point_id = int(vector_id)
+            except ValueError:
+                point_id = abs(hash(vector_id)) % (10**18)
+            point_ids.append(point_id)
+
+        wait = kwargs.get("wait", True)
+        await self.client.delete(
+            collection_name=self.collection_name,
+            points_selector=PointIdsList(points=point_ids),
+            wait=wait,
+        )
+
+        logger.info(f"Deleted {len(point_ids)} documents from {self.collection_name}")
+
+    async def update(self, nodes: VectorNode | list[VectorNode], **kwargs: Any):
+        """Update existing vector nodes with new content or metadata."""
+        if isinstance(nodes, VectorNode):
+            nodes = [nodes]
+
+        nodes_without_vectors = [node for node in nodes if node.vector is None and node.content]
+        if nodes_without_vectors:
+            nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors)
+            vector_map = {n.vector_id: n for n in nodes_with_vectors}
+            nodes_to_update = [vector_map.get(n.vector_id, n) if n.vector is None and n.content else n for n in nodes]
+        else:
+            nodes_to_update = nodes
+
+        points = []
+        for node in nodes_to_update:
+            try:
+                point_id = int(node.vector_id)
+            except ValueError:
+                point_id = abs(hash(node.vector_id)) % (10**18)
+
+            point = PointStruct(
+                id=point_id,
+                vector=node.vector,
+                payload={
+                    "vector_id": node.vector_id,
+                    "content": node.content,
+                    "metadata": node.metadata,
+                },
+            )
+            points.append(point)
+
+        wait = kwargs.get("wait", True)
+        await self.client.upsert(
+            collection_name=self.collection_name,
+            points=points,
+            wait=wait,
+        )
+
+        logger.info(f"Updated {len(points)} documents in {self.collection_name}")
+
+    async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode]:
+        """Retrieve vector nodes by their IDs from the collection."""
+        single_result = isinstance(vector_ids, str)
+        if single_result:
+            vector_ids = [vector_ids]
+
+        point_ids = []
+        for vector_id in vector_ids:
+            try:
+                point_id = int(vector_id)
+            except ValueError:
+                point_id = abs(hash(vector_id)) % (10**18)
+            point_ids.append(point_id)
+
+        points = await self.client.retrieve(
+            collection_name=self.collection_name,
+            ids=point_ids,
+            with_payload=True,
+            with_vectors=True,
+        )
+
+        results = []
+        for point in points:
+            if point:
+                payload = point.payload or {}
+                node = VectorNode(
+                    vector_id=payload.get("vector_id", str(point.id)),
+                    content=payload.get("content", ""),
+                    vector=point.vector,
+                    metadata=payload.get("metadata", {}),
+                )
+                results.append(node)
+            else:
+                logger.warning("Point not found")
+
+        return results[0] if single_result and results else results
+
+    async def list(
+        self,
+        filters: dict | None = None,
+        limit: int | None = None,
+    ) -> list[VectorNode]:
+        """List all vector nodes in the collection matching the filter criteria."""
+        scroll_filter = self._create_filter(filters) if filters else None
+
+        limit = limit or 10000
+        records, _ = await self.client.scroll(
+            collection_name=self.collection_name,
+            scroll_filter=scroll_filter,
+            limit=limit,
+            with_payload=True,
+            with_vectors=True,
+        )
+
+        results = []
+        for record in records:
+            payload = record.payload or {}
+            node = VectorNode(
+                vector_id=payload.get("vector_id", str(record.id)),
+                content=payload.get("content", ""),
+                vector=record.vector,
+                metadata=payload.get("metadata", {}),
+            )
+            results.append(node)
+
+        return results
+
+    async def close(self):
+        """Close the AsyncQdrantClient connection and release resources."""
+        await self.client.close()
+        logger.info("Qdrant client connection closed")

tests/mcp_servers_demo.json

@@ -0,0 +1,37 @@
+{
+  "mcpServers": {
+    "sqlite-explorer": {
+      "type": "stdio",
+      "command": "uv",
+      "args": [
+        "run",
+        "--with",
+        "mcp-server-sqlite",
+        "mcp-server-sqlite",
+        "--db-path",
+        "/path/to/your/database.db"
+      ],
+      "env": {
+        "CUSTOM_VAR": "optional_value"
+      }
+    },
+    "remote-fetcher": {
+      "type": "sse",
+      "url": "https://mcp-server.example.com/sse",
+      "headers": {
+        "Authorization": "Bearer {BAILIAN_MCP_API_KEY}",
+        "Content-Type": "application/json"
+      },
+      "timeout": 5,
+      "sse_read_timeout": 300
+    },
+    "my-modern-remote": {
+      "type": "streamable-http",
+      "url": "https://api.example.com/mcp",
+      "headers": {
+        "Authorization": "Bearer {BAILIAN_MCP_API_KEY}"
+      },
+      "timeout": 5
+    }
+  }
+}

tests/test_base_context.py

@@ -0,0 +1,72 @@
+"""
+Unit tests for the BaseContext class in reme_ai.core.context.
+Ensures attribute-style and dict-style access work interchangeably.
+"""
+
+import pickle
+from reme_ai.core.context import BaseContext
+
+
+def test_attribute_access():
+    """Test setting values via attributes and retrieving via items."""
+    context = BaseContext()
+    context.xxx = 123
+    assert context.xxx == 123
+    assert context["xxx"] == 123
+
+
+def test_dict_access():
+    """Test setting values via items and retrieving via attributes."""
+    context = BaseContext()
+    context["yyy"] = 456
+    assert context.yyy == 456
+    assert context["yyy"] == 456
+
+
+def test_delete_attribute():
+    """Test that deleting an attribute removes it from the internal state."""
+    context = BaseContext()
+    context.zzz = 789
+    del context.zzz
+    assert "zzz" not in context
+
+
+def test_attribute_error():
+    """Test that accessing non-existent attributes raises the correct error."""
+    context = BaseContext()
+    try:
+        _ = context.nonexistent
+        assert False, "Should raise AttributeError"
+    except AttributeError as error:
+        assert "nonexistent" in str(error)
+
+
+def test_pickling():
+    """Test that BaseContext instances can be serialized and deserialized."""
+    context = BaseContext()
+    context.test_value = "bar"
+    context.num = 42
+
+    pickled = pickle.dumps(context)
+    restored = pickle.loads(pickled)
+
+    assert restored.test_value == "bar"
+    assert restored.num == 42
+    assert isinstance(restored, BaseContext)
+
+
+def test_init_with_data():
+    """Test that the constructor correctly handles initial dictionary data."""
+    context = BaseContext({"a": 1, "b": 2})
+    assert context.a == 1
+    assert context.b == 2
+
+
+if __name__ == "__main__":
+    test_attribute_access()
+    test_dict_access()
+    test_delete_attribute()
+    test_attribute_error()
+    test_pickling()
+    test_init_with_data()
+    print("All tests passed!")

tests/test_cache_handler.py

@@ -0,0 +1,94 @@
+"""
+Self-contained script for CacheHandler's comprehensive test suite.
+"""
+
+import shutil
+from datetime import datetime, timedelta
+from pathlib import Path
+
+import pandas as pd
+from loguru import logger
+
+from reme_ai.core.utils.cache_handler import CacheHandler
+
+
+def run_tests():
+    """Execute comprehensive tests for CacheHandler."""
+    test_dir = Path("test_cache_system")
+    if test_dir.exists():
+        shutil.rmtree(test_dir)
+
+    handler = CacheHandler(cache_dir=test_dir)
+    logger.info("Starting CacheHandler tests...")
+
+    # 1. Test Data Types
+    logger.info("Testing data types support...")
+
+    # DataFrame
+    df = pd.DataFrame({"a": [1, 2], "b": [3, 4]})
+    assert handler.save("df_test", df)
+    assert isinstance(handler.load("df_test"), pd.DataFrame)
+    assert handler.load("df_test").shape == (2, 2)
+
+    # Dict & List
+    d = {"key": "value", "nested": [1, 2]}
+    l_value = [1, "string", {"a": 1}]
+    assert handler.save("dict_test", d)
+    assert handler.save("list_test", l_value)
+    assert handler.load("dict_test")["key"] == "value"
+    assert handler.load("list_test")[1] == "string"
+
+    # String
+    s = "Hello World"
+    assert handler.save("str_test", s)
+    assert handler.load("str_test") == "Hello World"
+
+    # 2. Test Expiration
+    logger.info("Testing expiration logic...")
+    # Save with 1 second expiry (approx 0.00027 hours)
+    handler.save("exp_test", {"data": 1}, expire_hours=0.00001)
+    assert handler.exists("exp_test") is True
+
+    # Manually modify metadata to force expiration for instant test
+    handler.metadata["exp_test"]["expire_at"] = (datetime.now() - timedelta(seconds=1)).isoformat()
+    assert handler.exists("exp_test") is False
+    assert handler.load("exp_test") is None
+    assert "exp_test" not in handler.metadata  # Auto-cleaned
+
+    # 3. Test Existence and Deletion
+    logger.info("Testing delete and exists...")
+    handler.save("del_test", "delete me")
+    assert handler.exists("del_test") is True
+    handler.delete("del_test")
+    assert handler.exists("del_test") is False
+    assert not (test_dir / "del_test.txt").exists()
+
+    # 4. Test Persistence (Reload handler)
+    logger.info("Testing persistence...")
+    handler.save("persist_test", [1, 2, 3])
+    new_handler = CacheHandler(cache_dir=test_dir)
+    assert new_handler.exists("persist_test") is True
+    assert new_handler.load("persist_test") == [1, 2, 3]
+
+    # 5. Test Statistics and Clear
+    logger.info("Testing stats and clear...")
+    stats = handler.get_stats()
+    assert stats["count"] > 0
+    handler.clear_all()
+    assert handler.get_stats()["count"] == 0
+    assert len(list(test_dir.glob("*"))) == 1  # Only metadata.json remains
+
+    # 6. Test Error Handling
+    logger.info("Testing error handling...")
+    assert handler.load("non_existent_key") is None
+    # Test unsupported type
+    assert handler.save("invalid", {1, 2}) is False
+
+    logger.success("All tests passed successfully!")
+
+    # Cleanup after tests
+    shutil.rmtree(test_dir)
+
+
+if __name__ == "__main__":
+    run_tests()

tests/test_embedding.py

@@ -0,0 +1,349 @@
+"""
+Async unit tests for Embedding classes (OpenAIEmbeddingModel) covering:
+- Async single text embedding
+- Async batch text embeddings
+- Async large batch with automatic batching
+- Async VectorNode embedding (single and batch)
+- Error handling and retries
+
+Usage:
+    python test_embedding.py --openai      # Test OpenAIEmbeddingModel only
+    python test_embedding.py --all         # Test all embedding models
+"""
+
+# flake8: noqa: E402
+# pylint: disable=C0413
+
+import asyncio
+import argparse
+from typing import Type, List
+
+from reme_ai.core.utils import load_env
+
+load_env()
+
+from reme_ai.core.embedding import OpenAIEmbeddingModel, BaseEmbeddingModel
+from reme_ai.core.schema import VectorNode
+
+
+def get_embedding_model(model_class: Type[BaseEmbeddingModel]) -> BaseEmbeddingModel:
+    """Create and return an embedding model instance."""
+    return model_class(
+        model_name="text-embedding-v4",
+        dimensions=1024,
+        max_retries=2,
+        raise_exception=True,
+    )
+
+
+def get_test_texts() -> List[str]:
+    """Create test texts for embedding."""
+    return [
+        "The quick brown fox jumps over the lazy dog.",
+        "Machine learning is a subset of artificial intelligence.",
+        "Python is a popular programming language for data science.",
+        "Solar energy is a renewable source of power.",
+        "The capital of France is Paris.",
+    ]
+
+
+def get_large_batch_texts() -> List[str]:
+    """Create a large batch of test texts to test automatic batching."""
+    texts = []
+    topics = [
+        "Climate change and global warming",
+        "Artificial intelligence and machine learning",
+        "Renewable energy sources",
+        "Space exploration and astronomy",
+        "Medical research and healthcare",
+        "Financial markets and economics",
+        "Education and learning systems",
+        "Transportation and urban planning",
+    ]
+
+    for i, topic in enumerate(topics):
+        for j in range(3):
+            texts.append(f"Text {i*3+j+1}: This is a sample text about {topic}.")
+
+    return texts  # 24 texts total
+
+
+def get_test_nodes() -> List[VectorNode]:
+    """Create test VectorNodes for embedding."""
+    texts = get_test_texts()
+    return [
+        VectorNode(
+            content=text,
+            metadata={"index": str(i), "category": "test"},
+        )
+        for i, text in enumerate(texts)
+    ]
+
+
+async def test_async_single_embedding(model_class: Type[BaseEmbeddingModel], model_name: str):
+    """Test asynchronous single text embedding."""
+    print(f"\n{'='*60}")
+    print(f"Testing {model_name}: Async Single Text Embedding")
+    print(f"{'='*60}")
+
+    model = get_embedding_model(model_class)
+    test_text = "Hello, this is a test sentence for embedding."
+
+    print(f"Input text: {test_text}")
+
+    embedding = await model.get_embedding(test_text)
+
+    assert embedding is not None, f"{model_name}: Embedding is None"
+    assert isinstance(embedding, list), f"{model_name}: Embedding is not a list"
+    assert len(embedding) > 0, f"{model_name}: Empty embedding"
+    assert len(embedding) == model.dimensions, f"{model_name}: Embedding dimension mismatch"
+    assert all(isinstance(x, float) for x in embedding), f"{model_name}: Not all elements are floats"
+
+    print("\n✓ Embedding generated successfully")
+    print(f"  - Dimension: {len(embedding)}")
+    print(f"  - First 5 values: {embedding[:5]}")
+    print(f"  - Value range: [{min(embedding):.4f}, {max(embedding):.4f}]")
+
+    await model.close()
+    print(f"✓ PASSED: {model_name} async single embedding")
+
+
+async def test_async_batch_embeddings(model_class: Type[BaseEmbeddingModel], model_name: str):
+    """Test asynchronous batch text embeddings."""
+    print(f"\n{'='*60}")
+    print(f"Testing {model_name}: Async Batch Text Embeddings")
+    print(f"{'='*60}")
+
+    model = get_embedding_model(model_class)
+    test_texts = get_test_texts()
+
+    print(f"Input: {len(test_texts)} texts")
+    for i, text in enumerate(test_texts[:3], 1):
+        print(f"  {i}. {text[:50]}...")
+
+    embeddings = await model.get_embeddings(test_texts)
+
+    assert embeddings is not None, f"{model_name}: Embeddings is None"
+    assert isinstance(embeddings, list), f"{model_name}: Embeddings is not a list"
+    assert len(embeddings) == len(test_texts), f"{model_name}: Embeddings count mismatch"
+
+    for i, emb in enumerate(embeddings):
+        assert isinstance(emb, list), f"{model_name}: Embedding {i} is not a list"
+        assert len(emb) == model.dimensions, f"{model_name}: Embedding {i} dimension mismatch"
+        assert all(isinstance(x, float) for x in emb), f"{model_name}: Embedding {i} has non-float values"
+
+    print("\n✓ Batch embeddings generated successfully")
+    print(f"  - Count: {len(embeddings)}")
+    print(f"  - Dimension: {len(embeddings[0])}")
+    print(f"  - First embedding preview: {embeddings[0][:3]}...")
+
+    await model.close()
+    print(f"✓ PASSED: {model_name} async batch embeddings")
+
+
+async def test_async_large_batch_embeddings(model_class: Type[BaseEmbeddingModel], model_name: str):
+    """Test asynchronous large batch embeddings with automatic batching."""
+    print(f"\n{'='*60}")
+    print(f"Testing {model_name}: Async Large Batch with Auto-Batching")
+    print(f"{'='*60}")
+
+    model = get_embedding_model(model_class)
+    test_texts = get_large_batch_texts()
+
+    print(f"Input: {len(test_texts)} texts")
+    print(f"Max batch size: {model.max_batch_size}")
+    print(f"Expected batches: {(len(test_texts) + model.max_batch_size - 1) // model.max_batch_size}")
+
+    embeddings = await model.get_embeddings(test_texts)
+
+    assert embeddings is not None, f"{model_name}: Embeddings is None"
+    assert isinstance(embeddings, list), f"{model_name}: Embeddings is not a list"
+    assert len(embeddings) == len(test_texts), f"{model_name}: Embeddings count mismatch"
+
+    # Check all embeddings are valid
+    for i, emb in enumerate(embeddings):
+        assert isinstance(emb, list), f"{model_name}: Embedding {i} is not a list"
+        assert len(emb) == model.dimensions, f"{model_name}: Embedding {i} dimension mismatch"
+
+    print("\n✓ Large batch embeddings generated successfully")
+    print(f"  - Total texts: {len(test_texts)}")
+    print(f"  - Total embeddings: {len(embeddings)}")
+    print(f"  - Dimension: {len(embeddings[0])}")
+
+    await model.close()
+    print(f"✓ PASSED: {model_name} async large batch embeddings")
+
+
+async def test_async_single_node_embedding(model_class: Type[BaseEmbeddingModel], model_name: str):
+    """Test asynchronous single VectorNode embedding."""
+    print(f"\n{'='*60}")
+    print(f"Testing {model_name}: Async Single VectorNode Embedding")
+    print(f"{'='*60}")
+
+    model = get_embedding_model(model_class)
+    node = VectorNode(
+        content="This is a test node for embedding.",
+        metadata={"test": "true"},
+    )
+
+    print(f"Input node content: {node.content}")
+    print(f"Initial vector: {node.vector}")
+
+    result_node = await model.get_node_embedding(node)
+
+    assert result_node is not None, f"{model_name}: Result node is None"
+    assert result_node.vector is not None, f"{model_name}: Node vector is None"
+    assert isinstance(result_node.vector, list), f"{model_name}: Vector is not a list"
+    assert len(result_node.vector) == model.dimensions, f"{model_name}: Vector dimension mismatch"
+
+    print("\n✓ Node embedding generated successfully")
+    print(f"  - Vector dimension: {len(result_node.vector)}")
+    print(f"  - First 5 values: {result_node.vector[:5]}")
+    print(f"  - Metadata preserved: {result_node.metadata}")
+
+    await model.close()
+    print(f"✓ PASSED: {model_name} async single node embedding")
+
+
+async def test_async_batch_node_embeddings(model_class: Type[BaseEmbeddingModel], model_name: str):
+    """Test asynchronous batch VectorNode embeddings."""
+    print(f"\n{'='*60}")
+    print(f"Testing {model_name}: Async Batch VectorNode Embeddings")
+    print(f"{'='*60}")
+
+    model = get_embedding_model(model_class)
+    nodes = get_test_nodes()
+
+    print(f"Input: {len(nodes)} nodes")
+    for i, node in enumerate(nodes[:3], 1):
+        print(f"  {i}. {node.content[:50]}...")
+
+    result_nodes = await model.get_node_embeddings(nodes)
+
+    assert result_nodes is not None, f"{model_name}: Result nodes is None"
+    assert isinstance(result_nodes, list), f"{model_name}: Result is not a list"
+    assert len(result_nodes) == len(nodes), f"{model_name}: Nodes count mismatch"
+
+    for i, node in enumerate(result_nodes):
+        assert node.vector is not None, f"{model_name}: Node {i} vector is None"
+        assert isinstance(node.vector, list), f"{model_name}: Node {i} vector is not a list"
+        assert len(node.vector) == model.dimensions, f"{model_name}: Node {i} dimension mismatch"
+        assert node.metadata is not None, f"{model_name}: Node {i} metadata is None"
+
+    print("\n✓ Batch node embeddings generated successfully")
+    print(f"  - Count: {len(result_nodes)}")
+    print(f"  - All vectors populated: {all(n.vector is not None for n in result_nodes)}")
+    print(f"  - All metadata preserved: {all(n.metadata is not None for n in result_nodes)}")
+
+    await model.close()
+    print(f"✓ PASSED: {model_name} async batch node embeddings")
+
+
+async def test_async_large_batch_node_embeddings(model_class: Type[BaseEmbeddingModel], model_name: str):
+    """Test asynchronous large batch VectorNode embeddings with automatic batching."""
+    print(f"\n{'='*60}")
+    print(f"Testing {model_name}: Async Large Batch Node Embeddings")
+    print(f"{'='*60}")
+
+    model = get_embedding_model(model_class)
+    texts = get_large_batch_texts()
+    nodes = [VectorNode(content=text, metadata={"index": str(i)}) for i, text in enumerate(texts)]
+
+    print(f"Input: {len(nodes)} nodes")
+    print(f"Max batch size: {model.max_batch_size}")
+    print(f"Expected batches: {(len(nodes) + model.max_batch_size - 1) // model.max_batch_size}")
+
+    result_nodes = await model.get_node_embeddings(nodes)
+
+    assert result_nodes is not None, f"{model_name}: Result nodes is None"
+    assert len(result_nodes) == len(nodes), f"{model_name}: Nodes count mismatch"
+
+    # Check all nodes have embeddings
+    nodes_with_vectors = sum(1 for n in result_nodes if n.vector is not None)
+    print("\n✓ Large batch node embeddings generated successfully")
+    print(f"  - Total nodes: {len(result_nodes)}")
+    print(f"  - Nodes with vectors: {nodes_with_vectors}")
+    print(f"  - Success rate: {nodes_with_vectors/len(result_nodes)*100:.1f}%")
+
+    assert nodes_with_vectors == len(nodes), f"{model_name}: Not all nodes have vectors"
+
+    await model.close()
+    print(f"✓ PASSED: {model_name} async large batch node embeddings")
+
+
+async def run_all_tests_for_model(model_class: Type[BaseEmbeddingModel], model_name: str):
+    """Run all tests for a specific embedding model class."""
+    print(f"\n\n{'#'*60}")
+    print(f"# Running all tests for: {model_name}")
+    print(f"{'#'*60}")
+
+    await test_async_single_embedding(model_class, model_name)
+    await test_async_batch_embeddings(model_class, model_name)
+    await test_async_large_batch_embeddings(model_class, model_name)
+    await test_async_single_node_embedding(model_class, model_name)
+    await test_async_batch_node_embeddings(model_class, model_name)
+    await test_async_large_batch_node_embeddings(model_class, model_name)
+
+    print(f"\n{'='*60}")
+    print(f"✓ All tests passed for {model_name}!")
+    print(f"{'='*60}")
+
+
+async def main():
+    """Main entry point for running tests."""
+    parser = argparse.ArgumentParser(
+        description="Run async embedding model tests",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  python test_embedding.py --openai      # Test OpenAIEmbeddingModel only
+  python test_embedding.py --all         # Test all embedding models
+        """,
+    )
+    parser.add_argument(
+        "--openai",
+        action="store_true",
+        help="Test OpenAIEmbeddingModel",
+    )
+    parser.add_argument(
+        "--all",
+        action="store_true",
+        help="Run tests for all available embedding models",
+    )
+
+    args = parser.parse_args()
+
+    # Determine which models to test
+    models_to_test = []
+
+    if args.openai:
+        models_to_test.append((OpenAIEmbeddingModel, "OpenAIEmbeddingModel"))
+    elif args.all:
+        models_to_test.append((OpenAIEmbeddingModel, "OpenAIEmbeddingModel"))
+    else:
+        # Default to all models if no argument provided
+        models_to_test = [(OpenAIEmbeddingModel, "OpenAIEmbeddingModel")]
+        print("No model specified, defaulting to all models")
+        print("Use --openai to test OpenAI specifically\n")
+
+    # Run tests for each model
+    for model_class, model_name in models_to_test:
+        try:
+            await run_all_tests_for_model(model_class, model_name)
+        except Exception as e:
+            print(f"\n✗ FAILED: {model_name} tests failed with error:")
+            print(f"  {type(e).__name__}: {e}")
+            raise
+
+    # Final summary
+    print(f"\n\n{'#'*60}")
+    print("# TEST SUMMARY")
+    print(f"{'#'*60}")
+    print(f"✓ All tests passed for {len(models_to_test)} embedding model(s):")
+    for _, model_name in models_to_test:
+        print(f"  - {model_name}")
+    print(f"{'#'*60}\n")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

tests/test_embedding_sync.py

@@ -0,0 +1,348 @@
+"""
+Sync unit tests for Embedding classes (OpenAIEmbeddingModelSync) covering:
+- Sync single text embedding
+- Sync batch text embeddings
+- Sync large batch with automatic batching
+- Sync VectorNode embedding (single and batch)
+- Error handling and retries
+
+Usage:
+    python test_embedding_sync.py --openai      # Test OpenAIEmbeddingModelSync only
+    python test_embedding_sync.py --all         # Test all embedding models
+"""
+
+# flake8: noqa: E402
+# pylint: disable=C0413
+
+import argparse
+from typing import Type, List
+
+from reme_ai.core.utils import load_env
+
+load_env()
+
+from reme_ai.core.embedding import OpenAIEmbeddingModelSync, BaseEmbeddingModel
+from reme_ai.core.schema import VectorNode
+
+
+def get_embedding_model(model_class: Type[BaseEmbeddingModel]) -> BaseEmbeddingModel:
+    """Create and return an embedding model instance."""
+    return model_class(
+        model_name="text-embedding-v4",
+        dimensions=1024,
+        max_retries=2,
+        raise_exception=True,
+    )
+
+
+def get_test_texts() -> List[str]:
+    """Create test texts for embedding."""
+    return [
+        "The quick brown fox jumps over the lazy dog.",
+        "Machine learning is a subset of artificial intelligence.",
+        "Python is a popular programming language for data science.",
+        "Solar energy is a renewable source of power.",
+        "The capital of France is Paris.",
+    ]
+
+
+def get_large_batch_texts() -> List[str]:
+    """Create a large batch of test texts to test automatic batching."""
+    texts = []
+    topics = [
+        "Climate change and global warming",
+        "Artificial intelligence and machine learning",
+        "Renewable energy sources",
+        "Space exploration and astronomy",
+        "Medical research and healthcare",
+        "Financial markets and economics",
+        "Education and learning systems",
+        "Transportation and urban planning",
+    ]
+
+    for i, topic in enumerate(topics):
+        for j in range(3):
+            texts.append(f"Text {i*3+j+1}: This is a sample text about {topic}.")
+
+    return texts  # 24 texts total
+
+
+def get_test_nodes() -> List[VectorNode]:
+    """Create test VectorNodes for embedding."""
+    texts = get_test_texts()
+    return [
+        VectorNode(
+            content=text,
+            metadata={"index": str(i), "category": "test"},
+        )
+        for i, text in enumerate(texts)
+    ]
+
+
+def test_sync_single_embedding(model_class: Type[BaseEmbeddingModel], model_name: str):
+    """Test synchronous single text embedding."""
+    print(f"\n{'='*60}")
+    print(f"Testing {model_name}: Sync Single Text Embedding")
+    print(f"{'='*60}")
+
+    model = get_embedding_model(model_class)
+    test_text = "Hello, this is a test sentence for embedding."
+
+    print(f"Input text: {test_text}")
+
+    embedding = model.get_embedding_sync(test_text)
+
+    assert embedding is not None, f"{model_name}: Embedding is None"
+    assert isinstance(embedding, list), f"{model_name}: Embedding is not a list"
+    assert len(embedding) > 0, f"{model_name}: Empty embedding"
+    assert len(embedding) == model.dimensions, f"{model_name}: Embedding dimension mismatch"
+    assert all(isinstance(x, float) for x in embedding), f"{model_name}: Not all elements are floats"
+
+    print("\n✓ Embedding generated successfully")
+    print(f"  - Dimension: {len(embedding)}")
+    print(f"  - First 5 values: {embedding[:5]}")
+    print(f"  - Value range: [{min(embedding):.4f}, {max(embedding):.4f}]")
+
+    model.close_sync()
+    print(f"✓ PASSED: {model_name} sync single embedding")
+
+
+def test_sync_batch_embeddings(model_class: Type[BaseEmbeddingModel], model_name: str):
+    """Test synchronous batch text embeddings."""
+    print(f"\n{'='*60}")
+    print(f"Testing {model_name}: Sync Batch Text Embeddings")
+    print(f"{'='*60}")
+
+    model = get_embedding_model(model_class)
+    test_texts = get_test_texts()
+
+    print(f"Input: {len(test_texts)} texts")
+    for i, text in enumerate(test_texts[:3], 1):
+        print(f"  {i}. {text[:50]}...")
+
+    embeddings = model.get_embeddings_sync(test_texts)
+
+    assert embeddings is not None, f"{model_name}: Embeddings is None"
+    assert isinstance(embeddings, list), f"{model_name}: Embeddings is not a list"
+    assert len(embeddings) == len(test_texts), f"{model_name}: Embeddings count mismatch"
+
+    for i, emb in enumerate(embeddings):
+        assert isinstance(emb, list), f"{model_name}: Embedding {i} is not a list"
+        assert len(emb) == model.dimensions, f"{model_name}: Embedding {i} dimension mismatch"
+        assert all(isinstance(x, float) for x in emb), f"{model_name}: Embedding {i} has non-float values"
+
+    print("\n✓ Batch embeddings generated successfully")
+    print(f"  - Count: {len(embeddings)}")
+    print(f"  - Dimension: {len(embeddings[0])}")
+    print(f"  - First embedding preview: {embeddings[0][:3]}...")
+
+    model.close_sync()
+    print(f"✓ PASSED: {model_name} sync batch embeddings")
+
+
+def test_sync_large_batch_embeddings(model_class: Type[BaseEmbeddingModel], model_name: str):
+    """Test synchronous large batch embeddings with automatic batching."""
+    print(f"\n{'='*60}")
+    print(f"Testing {model_name}: Sync Large Batch with Auto-Batching")
+    print(f"{'='*60}")
+
+    model = get_embedding_model(model_class)
+    test_texts = get_large_batch_texts()
+
+    print(f"Input: {len(test_texts)} texts")
+    print(f"Max batch size: {model.max_batch_size}")
+    print(f"Expected batches: {(len(test_texts) + model.max_batch_size - 1) // model.max_batch_size}")
+
+    embeddings = model.get_embeddings_sync(test_texts)
+
+    assert embeddings is not None, f"{model_name}: Embeddings is None"
+    assert isinstance(embeddings, list), f"{model_name}: Embeddings is not a list"
+    assert len(embeddings) == len(test_texts), f"{model_name}: Embeddings count mismatch"
+
+    # Check all embeddings are valid
+    for i, emb in enumerate(embeddings):
+        assert isinstance(emb, list), f"{model_name}: Embedding {i} is not a list"
+        assert len(emb) == model.dimensions, f"{model_name}: Embedding {i} dimension mismatch"
+
+    print("\n✓ Large batch embeddings generated successfully")
+    print(f"  - Total texts: {len(test_texts)}")
+    print(f"  - Total embeddings: {len(embeddings)}")
+    print(f"  - Dimension: {len(embeddings[0])}")
+
+    model.close_sync()
+    print(f"✓ PASSED: {model_name} sync large batch embeddings")
+
+
+def test_sync_single_node_embedding(model_class: Type[BaseEmbeddingModel], model_name: str):
+    """Test synchronous single VectorNode embedding."""
+    print(f"\n{'='*60}")
+    print(f"Testing {model_name}: Sync Single VectorNode Embedding")
+    print(f"{'='*60}")
+
+    model = get_embedding_model(model_class)
+    node = VectorNode(
+        content="This is a test node for embedding.",
+        metadata={"test": "true"},
+    )
+
+    print(f"Input node content: {node.content}")
+    print(f"Initial vector: {node.vector}")
+
+    result_node = model.get_node_embedding_sync(node)
+
+    assert result_node is not None, f"{model_name}: Result node is None"
+    assert result_node.vector is not None, f"{model_name}: Node vector is None"
+    assert isinstance(result_node.vector, list), f"{model_name}: Vector is not a list"
+    assert len(result_node.vector) == model.dimensions, f"{model_name}: Vector dimension mismatch"
+
+    print("\n✓ Node embedding generated successfully")
+    print(f"  - Vector dimension: {len(result_node.vector)}")
+    print(f"  - First 5 values: {result_node.vector[:5]}")
+    print(f"  - Metadata preserved: {result_node.metadata}")
+
+    model.close_sync()
+    print(f"✓ PASSED: {model_name} sync single node embedding")
+
+
+def test_sync_batch_node_embeddings(model_class: Type[BaseEmbeddingModel], model_name: str):
+    """Test synchronous batch VectorNode embeddings."""
+    print(f"\n{'='*60}")
+    print(f"Testing {model_name}: Sync Batch VectorNode Embeddings")
+    print(f"{'='*60}")
+
+    model = get_embedding_model(model_class)
+    nodes = get_test_nodes()
+
+    print(f"Input: {len(nodes)} nodes")
+    for i, node in enumerate(nodes[:3], 1):
+        print(f"  {i}. {node.content[:50]}...")
+
+    result_nodes = model.get_node_embeddings_sync(nodes)
+
+    assert result_nodes is not None, f"{model_name}: Result nodes is None"
+    assert isinstance(result_nodes, list), f"{model_name}: Result is not a list"
+    assert len(result_nodes) == len(nodes), f"{model_name}: Nodes count mismatch"
+
+    for i, node in enumerate(result_nodes):
+        assert node.vector is not None, f"{model_name}: Node {i} vector is None"
+        assert isinstance(node.vector, list), f"{model_name}: Node {i} vector is not a list"
+        assert len(node.vector) == model.dimensions, f"{model_name}: Node {i} dimension mismatch"
+        assert node.metadata is not None, f"{model_name}: Node {i} metadata is None"
+
+    print("\n✓ Batch node embeddings generated successfully")
+    print(f"  - Count: {len(result_nodes)}")
+    print(f"  - All vectors populated: {all(n.vector is not None for n in result_nodes)}")
+    print(f"  - All metadata preserved: {all(n.metadata is not None for n in result_nodes)}")
+
+    model.close_sync()
+    print(f"✓ PASSED: {model_name} sync batch node embeddings")
+
+
+def test_sync_large_batch_node_embeddings(model_class: Type[BaseEmbeddingModel], model_name: str):
+    """Test synchronous large batch VectorNode embeddings with automatic batching."""
+    print(f"\n{'='*60}")
+    print(f"Testing {model_name}: Sync Large Batch Node Embeddings")
+    print(f"{'='*60}")
+
+    model = get_embedding_model(model_class)
+    texts = get_large_batch_texts()
+    nodes = [VectorNode(content=text, metadata={"index": str(i)}) for i, text in enumerate(texts)]
+
+    print(f"Input: {len(nodes)} nodes")
+    print(f"Max batch size: {model.max_batch_size}")
+    print(f"Expected batches: {(len(nodes) + model.max_batch_size - 1) // model.max_batch_size}")
+
+    result_nodes = model.get_node_embeddings_sync(nodes)
+
+    assert result_nodes is not None, f"{model_name}: Result nodes is None"
+    assert len(result_nodes) == len(nodes), f"{model_name}: Nodes count mismatch"
+
+    # Check all nodes have embeddings
+    nodes_with_vectors = sum(1 for n in result_nodes if n.vector is not None)
+    print("\n✓ Large batch node embeddings generated successfully")
+    print(f"  - Total nodes: {len(result_nodes)}")
+    print(f"  - Nodes with vectors: {nodes_with_vectors}")
+    print(f"  - Success rate: {nodes_with_vectors/len(result_nodes)*100:.1f}%")
+
+    assert nodes_with_vectors == len(nodes), f"{model_name}: Not all nodes have vectors"
+
+    model.close_sync()
+    print(f"✓ PASSED: {model_name} sync large batch node embeddings")
+
+
+def run_all_tests_for_model(model_class: Type[BaseEmbeddingModel], model_name: str):
+    """Run all tests for a specific embedding model class."""
+    print(f"\n\n{'#'*60}")
+    print(f"# Running all tests for: {model_name}")
+    print(f"{'#'*60}")
+
+    test_sync_single_embedding(model_class, model_name)
+    test_sync_batch_embeddings(model_class, model_name)
+    test_sync_large_batch_embeddings(model_class, model_name)
+    test_sync_single_node_embedding(model_class, model_name)
+    test_sync_batch_node_embeddings(model_class, model_name)
+    test_sync_large_batch_node_embeddings(model_class, model_name)
+
+    print(f"\n{'='*60}")
+    print(f"✓ All tests passed for {model_name}!")
+    print(f"{'='*60}")
+
+
+def main():
+    """Main entry point for running tests."""
+    parser = argparse.ArgumentParser(
+        description="Run sync embedding model tests",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  python test_embedding_sync.py --openai      # Test OpenAIEmbeddingModelSync only
+  python test_embedding_sync.py --all         # Test all embedding models
+        """,
+    )
+    parser.add_argument(
+        "--openai",
+        action="store_true",
+        help="Test OpenAIEmbeddingModelSync",
+    )
+    parser.add_argument(
+        "--all",
+        action="store_true",
+        help="Run tests for all available embedding models",
+    )
+
+    args = parser.parse_args()
+
+    # Determine which models to test
+    models_to_test = []
+
+    if args.openai:
+        models_to_test.append((OpenAIEmbeddingModelSync, "OpenAIEmbeddingModelSync"))
+    elif args.all:
+        models_to_test.append((OpenAIEmbeddingModelSync, "OpenAIEmbeddingModelSync"))
+    else:
+        # Default to all models if no argument provided
+        models_to_test = [(OpenAIEmbeddingModelSync, "OpenAIEmbeddingModelSync")]
+        print("No model specified, defaulting to all models")
+        print("Use --openai to test OpenAI specifically\n")
+
+    # Run tests for each model
+    for model_class, model_name in models_to_test:
+        try:
+            run_all_tests_for_model(model_class, model_name)
+        except Exception as e:
+            print(f"\n✗ FAILED: {model_name} tests failed with error:")
+            print(f"  {type(e).__name__}: {e}")
+            raise
+
+    # Final summary
+    print(f"\n\n{'#'*60}")
+    print("# TEST SUMMARY")
+    print(f"{'#'*60}")
+    print(f"✓ All tests passed for {len(models_to_test)} embedding model(s):")
+    for _, model_name in models_to_test:
+        print(f"  - {model_name}")
+    print(f"{'#'*60}\n")
+
+
+if __name__ == "__main__":
+    main()

tests/test_logo.py

@@ -0,0 +1,9 @@
+"""test logo"""
+
+from reme_ai.core.schema import ServiceConfig, MCPConfig
+
+if __name__ == "__main__":
+    from reme_ai.core.utils import print_logo
+
+    c = ServiceConfig(app_name="reme", backend="mcp", mcp=MCPConfig(transport="sse"))
+    print_logo(service_config=c)

tests/test_mcp_client.py

@@ -0,0 +1,127 @@
+"""Test module for demonstrating MCPClient functionality."""
+
+# pylint: disable=too-many-return-statements,too-many-statements
+
+import asyncio
+import json
+
+from reme_ai.core.utils import MCPClient
+
+
+async def main():
+    """Execute demonstration of the MCPClient."""
+    test_mcp = "test_mcp"
+    config_data = {
+        "mcpServers": {
+            test_mcp: {
+                "url": "http://127.0.0.1:8010/sse",
+            },
+        },
+    }
+
+    client = MCPClient(config_data)
+
+    try:
+        # List all available tools
+        print("=" * 50)
+        print("Listing available tools:")
+        print("=" * 50)
+        t_list = await client.list_tool_calls(test_mcp)
+        for t in t_list:
+            print(json.dumps(t, ensure_ascii=False, indent=2))
+
+        # Helper function to build default values
+        def build_default_value(param_info: dict) -> any:
+            """Build a default value for a parameter based on its schema."""
+            param_type = param_info.get("type", "string")
+
+            # Handle enum types - use the first enum value
+            if "enum" in param_info and param_info["enum"]:
+                return param_info["enum"][0]
+
+            # Handle different types
+            if param_type == "string":
+                return "example_string"
+            elif param_type == "number":
+                return 0.0
+            elif param_type == "integer":
+                return 0
+            elif param_type == "boolean":
+                return False
+            elif param_type == "array":
+                return []
+            elif param_type == "object":
+                # Recursively build nested objects
+                obj = {}
+                nested_properties = param_info.get("properties", {})
+                nested_required = param_info.get("required", [])
+
+                for nested_param_name in nested_required:
+                    if nested_param_name in nested_properties:
+                        nested_param_info = nested_properties[nested_param_name]
+                        obj[nested_param_name] = build_default_value(nested_param_info)
+
+                return obj
+            else:
+                return None
+
+        # Call tools if available
+        if t_list:
+            # Execute the first two tools (or fewer if not enough tools available)
+            tools_to_execute = min(2, len(t_list))
+
+            for idx in range(tools_to_execute):
+                print("\n" + "=" * 50)
+                print(f"Calling tool #{idx + 1}:")
+                print("=" * 50)
+
+                # Get the tool's information
+                current_tool = t_list[idx]
+                tool_type = current_tool.get("type", "function")
+                tool_body = current_tool.get(tool_type, {})
+
+                tool_name = tool_body.get("name")
+                tool_description = tool_body.get("description", "")
+
+                # Prepare arguments based on the tool's input schema
+                tool_arguments = {}
+                parameters = tool_body.get("parameters", {})
+                properties = parameters.get("properties", {})
+                required = parameters.get("required", [])
+
+                # Build minimal arguments for required parameters
+                for param_name in required:
+                    if param_name in properties:
+                        param_info = properties[param_name]
+                        tool_arguments[param_name] = build_default_value(param_info)
+
+                # Validate tool name exists
+                if not tool_name:
+                    print("Error: Tool name not found in the tool definition")
+                    print(f"Tool structure: {json.dumps(current_tool, ensure_ascii=False, indent=2)}")
+                else:
+                    print(f"Tool name: {tool_name}")
+                    print(f"Tool description: {tool_description}")
+                    print(f"Arguments: {json.dumps(tool_arguments, ensure_ascii=False, indent=2)}")
+
+                    # Call the tool
+                    result = await client.call_tool(test_mcp, tool_name, tool_arguments, parse_text_result=True)
+
+                    print("\n" + "-" * 50)
+                    print("Tool call result:")
+                    print("-" * 50)
+                    print(f"Content: {result}")
+                    if hasattr(result, "isError"):
+                        print(f"Is Error: {result.isError}")
+        else:
+            print("\nNo tools available to call.")
+
+    except Exception as e:
+        print(f"Error occurred: {e}")
+        import traceback
+
+        traceback.print_exc()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

tests/test_mcp_server.py

@@ -0,0 +1,125 @@
+"""Dynamic MCP server implementation with JSON-schema based tool registration."""
+
+from typing import Any
+
+from fastmcp import FastMCP
+from fastmcp.tools import FunctionTool
+
+from reme_ai.core.schema import ToolCall
+from reme_ai.core.utils import create_pydantic_model
+
+mcp = FastMCP("DynamicSchemaServer", port=8010)
+
+# Configuration including enum examples
+MODES_CONFIG = {
+    "register_user": ToolCall(
+        **{
+            "name": "register_user",
+            "description": "Register a new user with metadata, tags, and roles.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "username": {"type": "string", "description": "Unique username"},
+                    "role": {
+                        "type": "string",
+                        "enum": ["admin", "editor", "viewer"],
+                        "description": "User access level",
+                    },
+                    "metadata": {
+                        "type": "object",
+                        "description": "User metadata",
+                        "properties": {
+                            "age": {"type": "integer"},
+                            "location": {"type": "string"},
+                        },
+                        "required": ["age"],
+                    },
+                    "tags": {
+                        "type": "array",
+                        "description": "User tags",
+                        "items": {
+                            "type": "object",
+                            "properties": {
+                                "tag_id": {"type": "string"},
+                                "level": {"type": "number"},
+                            },
+                            "required": ["tag_id"],
+                        },
+                    },
+                },
+                "required": ["username", "metadata", "role"],
+            },
+        },
+    ),
+    "create_order": ToolCall(
+        **{
+            "name": "create_order",
+            "description": "创建订单",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "order_id": {"type": "string", "description": "订单ID"},
+                    "amount": {"type": "number", "description": "订单金额"},
+                    "customer": {
+                        "type": "object",
+                        "description": "客户信息",
+                        "properties": {
+                            "name": {"type": "string", "description": "客户姓名"},
+                            "email": {"type": "string", "description": "客户邮箱"},
+                            "phone": {"type": "string", "description": "联系电话"},
+                        },
+                        "required": ["name", "email"],
+                    },
+                },
+                "required": ["order_id", "customer"],
+            },
+        },
+    ),
+}
+
+
+async def core_handler(mode: str, **kwargs: Any) -> dict[str, Any]:
+    """Process dynamic tool requests and return execution results."""
+    print(f"Executing Mode: {mode}, Parameters: {kwargs}")
+    return {
+        "status": "success",
+        "mode": mode,
+        "received_data": kwargs,
+    }
+
+
+def register_dynamic_tools() -> None:
+    """Iterate over tool configurations and register them to the MCP instance."""
+    for mode_name, tool_call in MODES_CONFIG.items():
+        # Create Pydantic model from tool parameters
+        request_model = create_pydantic_model(tool_call.name, tool_call.parameters)
+
+        # Create execution function with closure to capture current mode and model
+        def create_tool_func(current_mode: str, model: type):
+            async def execute_tool(**kwargs: Any) -> dict[str, Any]:
+                # Validate and normalize input using Pydantic model
+                validated_data = model(**kwargs).model_dump(exclude_none=True)
+                return await core_handler(current_mode, **validated_data)
+
+            return execute_tool
+
+        tool_fn = create_tool_func(mode_name, request_model)
+
+        # Extract parameters schema
+        tool_call_schema = tool_call.simple_input_dump()
+        parameters = tool_call_schema[tool_call_schema["type"]]["parameters"]
+
+        # Create FunctionTool and register
+        tool = FunctionTool(
+            name=tool_call.name,
+            description=tool_call.description,
+            fn=tool_fn,
+            parameters=parameters,
+        )
+
+        mcp.add_tool(tool)
+
+
+if __name__ == "__main__":
+    register_dynamic_tools()
+    mcp.run(transport="sse")

tests/test_op_composition.py

@@ -0,0 +1,325 @@
+"""
+Unit tests for BaseOp and operator composition (>>, <<, |).
+Tests asynchronous execution mode.
+"""
+
+import asyncio
+
+from reme_ai.core.op import BaseOp
+from reme_ai.core.schema import ToolCall, ToolAttr
+
+
+class AddOp(BaseOp):
+    """Simple operator that adds a value to a number in context."""
+
+    def __init__(self, value: int = 1, **kwargs):
+        super().__init__(**kwargs)
+        self.value = value
+
+    def _build_tool_call(self) -> ToolCall:
+        return ToolCall(
+            **{
+                "name": self.name,
+                "description": f"Add {self.value} to input",
+                "parameters": ToolAttr(
+                    **{
+                        "type": "object",
+                        "properties": {
+                            "number": {"type": "integer", "description": "Input number"},
+                        },
+                        "required": ["number"],
+                    },
+                ),
+            },
+        )
+
+    async def execute(self):
+        """Async execution: add value to input number."""
+        self.context["number"] += self.value
+        self.output = self.context["number"]
+
+
+class MultiplyOp(BaseOp):
+    """Simple operator that multiplies a number in context."""
+
+    def __init__(self, factor: int = 2, **kwargs):
+        super().__init__(**kwargs)
+        self.factor = factor
+
+    def _build_tool_call(self) -> ToolCall:
+        return ToolCall(
+            **{
+                "name": self.name,
+                "description": f"Multiply by {self.factor}",
+                "parameters": ToolAttr(
+                    **{
+                        "type": "object",
+                        "properties": {
+                            "number": {"type": "integer", "description": "Input number"},
+                        },
+                        "required": ["number"],
+                    },
+                ),
+            },
+        )
+
+    async def execute(self):
+        """Async execution: multiply input number."""
+        self.context["number"] *= self.factor
+        self.output = self.context["number"]
+
+
+class AppendOp(BaseOp):
+    """Operator that appends a value to a list in context."""
+
+    def __init__(self, value: str = "", **kwargs):
+        super().__init__(**kwargs)
+        self.value = value
+
+    def _build_tool_call(self) -> ToolCall:
+        return ToolCall(
+            **{
+                "name": self.name,
+                "description": f"Append {self.value} to list",
+                "parameters": ToolAttr(
+                    **{
+                        "type": "object",
+                        "properties": {
+                            "items": {"type": "array", "description": "List of items"},
+                        },
+                        "required": ["items"],
+                    },
+                ),
+            },
+        )
+
+    async def execute(self):
+        """Async execution: append value to list."""
+        self.context["items"].append(self.value)
+        self.output = self.context["items"]
+
+
+async def test_basic_async_call():
+    """Test basic asynchronous operator execution."""
+    op = AddOp(value=5, name="add_5")
+    await op.call(number=10)
+    number = op.context["number"]
+    assert number == 15, f"Expected context result 15, got {number}"
+    print("✓ test_basic_async_call passed")
+
+
+async def test_sequential_composition_async():
+    """Test >> operator for sequential composition in async mode."""
+    add_op = AddOp(value=5, name="add_5")
+    multiply_op = MultiplyOp(factor=2, name="multiply_2")
+    composed = add_op >> multiply_op
+    await composed.call(number=10)
+
+    # (10 + 5) * 2 = 30
+    assert composed.context["number"] == 30, f"Expected 30, got {composed.context['number']}"
+    print("✓ test_sequential_composition_async passed")
+
+
+async def test_parallel_composition_async():
+    """Test | operator for parallel composition in async mode."""
+    append_a = AppendOp(value="A", name="append_a")
+    append_b = AppendOp(value="B", name="append_b")
+    append_c = AppendOp(value="C", name="append_c")
+
+    composed = append_a | append_b | append_c
+
+    await composed.call(items=[])
+
+    # All should append to the list
+    items = composed.context["items"]
+    assert len(items) == 3, f"Expected 3 items, got {len(items)}"
+    assert set(items) == {"A", "B", "C"}, f"Expected A,B,C, got {items}"
+    print("✓ test_parallel_composition_async passed")
+
+
+async def test_add_sub_ops_async():
+    """Test << operator for adding sub-operations in async mode."""
+    parent_op = BaseOp(name="parent")
+    child1 = AddOp(value=5, name="child1")
+    child2 = MultiplyOp(factor=2, name="child2")
+
+    _ = parent_op << child1
+    _ = parent_op << child2
+
+    assert len(parent_op.sub_ops) == 2, f"Expected 2 sub_ops, got {len(parent_op.sub_ops)}"
+    sub_op_names = [op.name for op in parent_op.sub_ops]
+    assert "child1" in sub_op_names, "child1 not in sub_ops"
+    assert "child2" in sub_op_names, "child2 not in sub_ops"
+    print("✓ test_add_sub_ops_async passed")
+
+
+async def test_add_sub_ops_dict():
+    """Test << operator with dictionary of operations."""
+    parent_op = BaseOp(name="parent")
+    ops_dict = {
+        "add": AddOp(value=5, name="add"),
+        "multiply": MultiplyOp(factor=2, name="multiply"),
+    }
+
+    _ = parent_op << ops_dict
+
+    assert len(parent_op.sub_ops) == 2, f"Expected 2 ops_dict, got {len(parent_op.sub_ops)}"
+    sub_op_names = [op.name for op in parent_op.sub_ops]
+    assert "add" in sub_op_names, "add not in ops_dict"
+    assert "multiply" in sub_op_names, "multiply not in ops_dict"
+    print("✓ test_add_sub_ops_dict passed")
+
+
+async def test_add_sub_ops_list():
+    """Test << operator with list of operations."""
+    parent_op = BaseOp(name="parent")
+    sub_ops = [
+        AddOp(value=5, name="add"),
+        MultiplyOp(factor=2, name="multiply"),
+    ]
+
+    _ = parent_op << sub_ops
+
+    assert len(parent_op.sub_ops) == 2, f"Expected 2 sub_ops, got {len(parent_op.sub_ops)}"
+    sub_op_names = [op.name for op in parent_op.sub_ops]
+    assert "add" in sub_op_names, "add not in sub_ops"
+    assert "multiply" in sub_op_names, "multiply not in sub_ops"
+    print("✓ test_add_sub_ops_list passed")
+
+
+async def test_mixed_composition_async():
+    """Test mixing >> and | operators in async mode."""
+    # (add_5 >> multiply_2) | (add_10 >> multiply_3)
+    seq1 = AddOp(value=5, name="add_5") >> MultiplyOp(factor=2, name="multiply_2")
+    seq2 = AddOp(value=10, name="add_10") >> MultiplyOp(factor=3, name="multiply_3")
+
+    composed = seq1 | seq2
+
+    await composed.call(number=10)
+
+    # Both sequences execute in parallel with shared context
+    # seq1: (10 + 5) * 2 = 30
+    # seq2: (30 + 10) * 3 = 120 (builds on seq1's result due to shared context)
+    # The exact result depends on execution order and timing
+    # With current implementation, result is 120
+    assert composed.context["number"] == 120, f"Expected 120, got {composed.context['number']}"
+    print("✓ test_mixed_composition_async passed")
+
+
+async def test_op_copy():
+    """Test operator copy functionality."""
+    original = AddOp(value=5, name="original")
+    copy_op = original.copy(name="copy")
+
+    assert copy_op.name == "copy", f"Expected name 'copy', got {copy_op.name}"
+    assert copy_op.value == 5, f"Expected value 5, got {copy_op.value}"
+    assert copy_op is not original, "Copy should be a different object"
+    print("✓ test_op_copy passed")
+
+
+async def test_input_mapping():
+    """Test input_mapping parameter."""
+    op = AddOp(
+        value=5,
+        name="add_5",
+        input_mapping={"x": "number"},  # Map x to number
+    )
+
+    await op.call(x=10)  # Input is 'x' not 'number'
+
+    assert op.context["number"] == 15, f"Expected number=15, got {op.context['number']}"
+    print("✓ test_input_mapping passed")
+
+
+async def test_output_mapping():
+    """Test output_mapping parameter."""
+    op = AddOp(
+        value=5,
+        name="add_5",
+        output_mapping={"number": "final_result"},  # Map number to final_result
+    )
+
+    await op.call(number=10)
+
+    assert op.context["number"] == 15, f"Expected number=15, got {op.context['number']}"
+    assert op.context["final_result"] == 15, f"Expected final_result=15, got {op.context['final_result']}"
+    print("✓ test_output_mapping passed")
+
+
+async def test_validation_missing_required():
+    """Test that missing required inputs raise an error."""
+    op = AddOp(value=5, name="add_5", raise_exception=True)
+
+    try:
+        await op.call()  # Missing 'number' field
+        assert False, "Should have raised ValueError for missing required input"
+    except ValueError as e:
+        assert "number" in str(e), f"Expected error about 'number', got: {e}"
+        print("✓ test_validation_missing_required passed")
+
+
+async def test_max_retries():
+    """Test max_retries parameter with failing operation."""
+
+    class FailingOp(BaseOp):
+        """An operation that always fails."""
+
+        def __init__(self, **kwargs):
+            super().__init__(**kwargs)
+            self.attempt_count = 0
+
+        def _build_tool_call(self) -> ToolCall:
+            return ToolCall(
+                **{
+                    "name": self.name,
+                    "description": "Always fails",
+                    "parameters": ToolAttr(**{"type": "object", "properties": {}}),
+                    "output": ToolAttr(
+                        **{
+                            "type": "object",
+                            "properties": {
+                                "result": ToolAttr(**{"type": "string", "description": "Result"}),
+                            },
+                        },
+                    ),
+                },
+            )
+
+        async def execute(self):
+            self.attempt_count += 1
+            raise RuntimeError(f"Attempt {self.attempt_count} failed")
+
+    op = FailingOp(max_retries=3, name="failing")
+
+    await op.call()
+
+    assert op.attempt_count == 3, f"Expected 3 attempts, got {op.attempt_count}"
+    print("✓ test_max_retries passed")
+
+
+async def async_main():
+    """Run all async tests."""
+    await test_basic_async_call()
+    await test_sequential_composition_async()
+    await test_parallel_composition_async()
+    await test_add_sub_ops_async()
+    await test_add_sub_ops_dict()
+    await test_add_sub_ops_list()
+    await test_mixed_composition_async()
+    await test_op_copy()
+    await test_input_mapping()
+    await test_output_mapping()
+    await test_validation_missing_required()
+    await test_max_retries()
+
+
+if __name__ == "__main__":
+    print("Running BaseOp composition tests...\n")
+
+    # Async tests
+    print("=== Asynchronous Tests ===")
+    asyncio.run(async_main())
+
+    print("\n" + "=" * 50)
+    print("All tests passed! ✓")
+    print("=" * 50)

tests/test_token_counter.py

@@ -0,0 +1,511 @@
+"""
+Unit tests for TokenCounter classes covering:
+- BaseTokenCounter (rule-based estimation)
+- OpenAITokenCounter (tiktoken-based)
+- HFTokenCounter (HuggingFace tokenizer-based)
+
+Usage:
+    python test_token_counter.py --base        # Test BaseTokenCounter only
+    python test_token_counter.py --openai      # Test OpenAITokenCounter only
+    python test_token_counter.py --hf          # Test HFTokenCounter only
+    python test_token_counter.py --all         # Test all token counters
+"""
+
+import argparse
+from typing import Type, List
+
+from reme_ai.core.enumeration import Role
+from reme_ai.core.schema import Message, ToolCall
+from reme_ai.core.token_counter import BaseTokenCounter, OpenAITokenCounter, HFTokenCounter
+
+
+def get_token_counter(counter_class: Type[BaseTokenCounter], **kwargs) -> BaseTokenCounter:
+    """Create and return a token counter instance."""
+    default_kwargs = {
+        "model_name": "gpt-4o",
+    }
+    default_kwargs.update(kwargs)
+    return counter_class(**default_kwargs)
+
+
+def get_test_messages() -> List[Message]:
+    """Create test messages for token counting."""
+    return [
+        Message(role=Role.SYSTEM, content="You are a helpful assistant."),
+        Message(role=Role.USER, content="Hello, how are you today?"),
+        Message(role=Role.ASSISTANT, content="I'm doing well, thank you for asking! How can I help you?"),
+        Message(role=Role.USER, content="Can you explain what machine learning is?"),
+        Message(
+            role=Role.ASSISTANT,
+            content="Machine learning is a subset of artificial intelligence that enables computers to "
+            "learn from data without being explicitly programmed.",
+        ),
+    ]
+
+
+def get_chinese_messages() -> List[Message]:
+    """Create test messages with Chinese content."""
+    return [
+        Message(role=Role.SYSTEM, content="你是一个有帮助的助手。"),
+        Message(role=Role.USER, content="你好，今天天气怎么样？"),
+        Message(role=Role.ASSISTANT, content="今天天气很好，阳光明媚，适合外出活动。"),
+        Message(role=Role.USER, content="能给我推荐一些好看的电影吗？"),
+        Message(role=Role.ASSISTANT, content="当然可以！我推荐《肖申克的救赎》、《阿甘正传》和《泰坦尼克号》。"),
+    ]
+
+
+def get_mixed_messages() -> List[Message]:
+    """Create test messages with mixed English and Chinese content."""
+    return [
+        Message(role=Role.SYSTEM, content="You are a bilingual assistant. 你是一个双语助手。"),
+        Message(role=Role.USER, content="What is AI? 什么是人工智能？"),
+        Message(
+            role=Role.ASSISTANT,
+            content="AI (Artificial Intelligence) 是人工智能的英文缩写，它是计算机科学的一个分支。",
+        ),
+    ]
+
+
+def get_messages_with_reasoning() -> List[Message]:
+    """Create test messages with reasoning content."""
+    return [
+        Message(role=Role.USER, content="What is 2 + 2?"),
+        Message(
+            role=Role.ASSISTANT,
+            content="The answer is 4.",
+            reasoning_content="Let me think about this step by step. 2 + 2 "
+            "equals 4 because addition combines two quantities.",
+        ),
+    ]
+
+
+def get_test_tools() -> List[ToolCall]:
+    """Create test tool calls for token counting."""
+    return [
+        ToolCall(
+            **{
+                "type": "function",
+                "function": {
+                    "name": "get_weather",
+                    "description": "Get the current weather for a specified location.",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "location": {
+                                "type": "string",
+                                "description": "The city and country, e.g., 'Beijing, China'",
+                            },
+                            "unit": {
+                                "type": "string",
+                                "description": "Temperature unit: 'celsius' or 'fahrenheit'",
+                                "enum": ["celsius", "fahrenheit"],
+                            },
+                        },
+                        "required": ["location"],
+                    },
+                },
+            },
+        ),
+        ToolCall(
+            **{
+                "type": "function",
+                "function": {
+                    "name": "search_web",
+                    "description": "Search the web for information.",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "query": {
+                                "type": "string",
+                                "description": "The search query",
+                            },
+                            "num_results": {
+                                "type": "integer",
+                                "description": "Number of results to return",
+                            },
+                        },
+                        "required": ["query"],
+                    },
+                },
+            },
+        ),
+    ]
+
+
+def get_tool_call_messages() -> List[Message]:
+    """Create messages with tool call responses."""
+    return [
+        Message(role=Role.USER, content="What's the weather in Beijing?"),
+        Message(
+            role=Role.ASSISTANT,
+            content="",
+            tool_calls=[
+                ToolCall(
+                    id="call_123",
+                    name="get_weather",
+                    arguments='{"location": "Beijing, China", "unit": "celsius"}',
+                ),
+            ],
+        ),
+        Message(
+            role=Role.TOOL,
+            content='{"temperature": 25, "condition": "sunny", "humidity": 60}',
+            tool_call_id="call_123",
+        ),
+        Message(
+            role=Role.ASSISTANT,
+            content="The weather in Beijing is sunny with a temperature of 25°C and 60% humidity.",
+        ),
+    ]
+
+
+def test_basic_token_count(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs):
+    """Test basic token counting with simple messages."""
+    print(f"\n{'=' * 60}")
+    print(f"Testing {counter_name}: Basic Token Count")
+    print(f"{'=' * 60}")
+
+    counter = get_token_counter(counter_class, **kwargs)
+    messages = get_test_messages()
+
+    print(f"Input: {len(messages)} messages")
+    for i, msg in enumerate(messages, 1):
+        content_preview = msg.content[:50] + "..." if len(msg.content) > 50 else msg.content
+        print(f"  {i}. [{msg.role.value}] {content_preview}")
+
+    token_count = counter.count_token(messages)
+
+    assert token_count is not None, f"{counter_name}: Token count is None"
+    assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer"
+    assert token_count > 0, f"{counter_name}: Token count should be positive"
+
+    print(f"\n✓ Token count: {token_count}")
+    print(f"✓ PASSED: {counter_name} basic token count")
+
+
+def test_chinese_token_count(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs):
+    """Test token counting with Chinese content."""
+    print(f"\n{'=' * 60}")
+    print(f"Testing {counter_name}: Chinese Token Count")
+    print(f"{'=' * 60}")
+
+    counter = get_token_counter(counter_class, **kwargs)
+    messages = get_chinese_messages()
+
+    print(f"Input: {len(messages)} Chinese messages")
+    for i, msg in enumerate(messages, 1):
+        content_preview = msg.content[:30] + "..." if len(msg.content) > 30 else msg.content
+        print(f"  {i}. [{msg.role.value}] {content_preview}")
+
+    token_count = counter.count_token(messages)
+
+    assert token_count is not None, f"{counter_name}: Token count is None"
+    assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer"
+    assert token_count > 0, f"{counter_name}: Token count should be positive"
+
+    print(f"\n✓ Token count: {token_count}")
+    print(f"✓ PASSED: {counter_name} Chinese token count")
+
+
+def test_mixed_language_token_count(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs):
+    """Test token counting with mixed English and Chinese content."""
+    print(f"\n{'=' * 60}")
+    print(f"Testing {counter_name}: Mixed Language Token Count")
+    print(f"{'=' * 60}")
+
+    counter = get_token_counter(counter_class, **kwargs)
+    messages = get_mixed_messages()
+
+    print(f"Input: {len(messages)} mixed language messages")
+    for i, msg in enumerate(messages, 1):
+        content_preview = msg.content[:40] + "..." if len(msg.content) > 40 else msg.content
+        print(f"  {i}. [{msg.role.value}] {content_preview}")
+
+    token_count = counter.count_token(messages)
+
+    assert token_count is not None, f"{counter_name}: Token count is None"
+    assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer"
+    assert token_count > 0, f"{counter_name}: Token count should be positive"
+
+    print(f"\n✓ Token count: {token_count}")
+    print(f"✓ PASSED: {counter_name} mixed language token count")
+
+
+def test_reasoning_content_token_count(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs):
+    """Test token counting with reasoning content."""
+    print(f"\n{'=' * 60}")
+    print(f"Testing {counter_name}: Reasoning Content Token Count")
+    print(f"{'=' * 60}")
+
+    counter = get_token_counter(counter_class, **kwargs)
+    messages = get_messages_with_reasoning()
+
+    print(f"Input: {len(messages)} messages with reasoning content")
+    for i, msg in enumerate(messages, 1):
+        print(f"  {i}. [{msg.role.value}] content: {msg.content[:30]}...")
+        if msg.reasoning_content:
+            print(f"     reasoning: {msg.reasoning_content[:30]}...")
+
+    token_count = counter.count_token(messages)
+
+    assert token_count is not None, f"{counter_name}: Token count is None"
+    assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer"
+    assert token_count > 0, f"{counter_name}: Token count should be positive"
+
+    print(f"\n✓ Token count: {token_count}")
+    print(f"✓ PASSED: {counter_name} reasoning content token count")
+
+
+def test_token_count_with_tools(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs):
+    """Test token counting with tool definitions."""
+    print(f"\n{'=' * 60}")
+    print(f"Testing {counter_name}: Token Count with Tools")
+    print(f"{'=' * 60}")
+
+    counter = get_token_counter(counter_class, **kwargs)
+    messages = get_test_messages()[:2]
+    tools = get_test_tools()
+
+    print(f"Input: {len(messages)} messages, {len(tools)} tools")
+    for tool in tools:
+        print(f"  Tool: {tool.name} - {tool.description[:40]}...")
+
+    token_count = counter.count_token(messages, tools=tools)
+
+    assert token_count is not None, f"{counter_name}: Token count is None"
+    assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer"
+    assert token_count > 0, f"{counter_name}: Token count should be positive"
+
+    # Token count with tools should be higher than without
+    token_count_no_tools = counter.count_token(messages)
+    assert token_count > token_count_no_tools, f"{counter_name}: Token count with tools should be higher"
+
+    print(f"\n✓ Token count without tools: {token_count_no_tools}")
+    print(f"✓ Token count with tools: {token_count}")
+    print(f"✓ Tools added {token_count - token_count_no_tools} tokens")
+    print(f"✓ PASSED: {counter_name} token count with tools")
+
+
+def test_tool_call_messages_token_count(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs):
+    """Test token counting with messages containing tool calls."""
+    print(f"\n{'=' * 60}")
+    print(f"Testing {counter_name}: Tool Call Messages Token Count")
+    print(f"{'=' * 60}")
+
+    counter = get_token_counter(counter_class, **kwargs)
+    messages = get_tool_call_messages()
+
+    print(f"Input: {len(messages)} messages with tool calls")
+    for i, msg in enumerate(messages, 1):
+        if msg.tool_calls:
+            print(f"  {i}. [{msg.role.value}] tool_calls: {[tc.name for tc in msg.tool_calls]}")
+        else:
+            content_preview = msg.content[:40] + "..." if len(msg.content) > 40 else msg.content
+            print(f"  {i}. [{msg.role.value}] {content_preview}")
+
+    token_count = counter.count_token(messages)
+
+    assert token_count is not None, f"{counter_name}: Token count is None"
+    assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer"
+    assert token_count > 0, f"{counter_name}: Token count should be positive"
+
+    print(f"\n✓ Token count: {token_count}")
+    print(f"✓ PASSED: {counter_name} tool call messages token count")
+
+
+def test_empty_messages(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs):
+    """Test token counting with empty message list."""
+    print(f"\n{'=' * 60}")
+    print(f"Testing {counter_name}: Empty Messages")
+    print(f"{'=' * 60}")
+
+    # HFTokenCounter does not support empty message list (apply_chat_template requires at least one message)
+    if counter_class == HFTokenCounter:
+        print("⊘ SKIPPED: HFTokenCounter does not support empty message list")
+        return
+
+    counter = get_token_counter(counter_class, **kwargs)
+    messages: List[Message] = []
+
+    print("Input: 0 messages")
+
+    token_count = counter.count_token(messages)
+
+    assert token_count is not None, f"{counter_name}: Token count is None"
+    assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer"
+
+    # OpenAITokenCounter adds 3 tokens for reply priming even with empty messages
+    if counter_class == OpenAITokenCounter:
+        assert token_count == 3, f"{counter_name}: Empty messages should have 3 tokens (reply priming)"
+        print(f"\n✓ Token count: {token_count} (includes 3 tokens for reply priming)")
+    else:
+        assert token_count == 0, f"{counter_name}: Empty messages should have 0 tokens"
+        print(f"\n✓ Token count: {token_count}")
+
+    print(f"✓ PASSED: {counter_name} empty messages")
+
+
+def test_single_message(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs):
+    """Test token counting with a single message."""
+    print(f"\n{'=' * 60}")
+    print(f"Testing {counter_name}: Single Message")
+    print(f"{'=' * 60}")
+
+    counter = get_token_counter(counter_class, **kwargs)
+    messages = [Message(role=Role.USER, content="Hello!")]
+
+    print(f"Input: 1 message - '{messages[0].content}'")
+
+    token_count = counter.count_token(messages)
+
+    assert token_count is not None, f"{counter_name}: Token count is None"
+    assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer"
+    assert token_count > 0, f"{counter_name}: Token count should be positive"
+
+    print(f"\n✓ Token count: {token_count}")
+    print(f"✓ PASSED: {counter_name} single message")
+
+
+def test_long_content(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs):
+    """Test token counting with long content."""
+    print(f"\n{'=' * 60}")
+    print(f"Testing {counter_name}: Long Content")
+    print(f"{'=' * 60}")
+
+    counter = get_token_counter(counter_class, **kwargs)
+
+    # Create a long message
+    long_text = "This is a test sentence. " * 100
+    messages = [Message(role=Role.USER, content=long_text)]
+
+    print(f"Input: 1 message with {len(long_text)} characters")
+
+    token_count = counter.count_token(messages)
+
+    assert token_count is not None, f"{counter_name}: Token count is None"
+    assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer"
+    assert token_count > 0, f"{counter_name}: Token count should be positive"
+
+    # Long content should have more tokens
+    short_messages = [Message(role=Role.USER, content="This is a test sentence.")]
+    short_token_count = counter.count_token(short_messages)
+    assert token_count > short_token_count, f"{counter_name}: Long content should have more tokens"
+
+    print(f"\n✓ Short content token count: {short_token_count}")
+    print(f"✓ Long content token count: {token_count}")
+    print(f"✓ PASSED: {counter_name} long content")
+
+
+def run_all_tests_for_counter(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs):
+    """Run all tests for a specific token counter class."""
+    print(f"\n\n{'#' * 60}")
+    print(f"# Running all tests for: {counter_name}")
+    print(f"{'#' * 60}")
+
+    test_basic_token_count(counter_class, counter_name, **kwargs)
+    test_chinese_token_count(counter_class, counter_name, **kwargs)
+    test_mixed_language_token_count(counter_class, counter_name, **kwargs)
+    test_reasoning_content_token_count(counter_class, counter_name, **kwargs)
+    test_token_count_with_tools(counter_class, counter_name, **kwargs)
+    test_tool_call_messages_token_count(counter_class, counter_name, **kwargs)
+    test_empty_messages(counter_class, counter_name, **kwargs)
+    test_single_message(counter_class, counter_name, **kwargs)
+    test_long_content(counter_class, counter_name, **kwargs)
+
+    print(f"\n{'=' * 60}")
+    print(f"✓ All tests passed for {counter_name}!")
+    print(f"{'=' * 60}")
+
+
+def main():
+    """Main entry point for running tests."""
+    parser = argparse.ArgumentParser(
+        description="Run token counter tests",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  python test_token_counter.py --base        # Test BaseTokenCounter only
+  python test_token_counter.py --openai      # Test OpenAITokenCounter only
+  python test_token_counter.py --hf          # Test HFTokenCounter only
+  python test_token_counter.py --all         # Test all token counters
+        """,
+    )
+    parser.add_argument(
+        "--base",
+        action="store_true",
+        help="Test BaseTokenCounter (rule-based)",
+    )
+    parser.add_argument(
+        "--openai",
+        action="store_true",
+        help="Test OpenAITokenCounter (tiktoken-based)",
+    )
+    parser.add_argument(
+        "--hf",
+        action="store_true",
+        help="Test HFTokenCounter (HuggingFace tokenizer-based)",
+    )
+    parser.add_argument(
+        "--hf-model",
+        type=str,
+        default="Qwen/Qwen2.5-0.5B-Instruct",
+        help="HuggingFace model name for HFTokenCounter (default: Qwen/Qwen2.5-0.5B-Instruct)",
+    )
+    parser.add_argument(
+        "--all",
+        action="store_true",
+        help="Run tests for all available token counters",
+    )
+
+    args = parser.parse_args()
+
+    # Determine which counters to test
+    counters_to_test = []
+
+    if args.all:
+        counters_to_test.append((BaseTokenCounter, "BaseTokenCounter", {}))
+        counters_to_test.append((OpenAITokenCounter, "OpenAITokenCounter", {}))
+        counters_to_test.append(
+            (HFTokenCounter, "HFTokenCounter", {"model_name": args.hf_model, "trust_remote_code": True}),
+        )
+    else:
+        if args.base:
+            counters_to_test.append((BaseTokenCounter, "BaseTokenCounter", {}))
+        if args.openai:
+            counters_to_test.append((OpenAITokenCounter, "OpenAITokenCounter", {}))
+        if args.hf:
+            counters_to_test.append(
+                (HFTokenCounter, "HFTokenCounter", {"model_name": args.hf_model, "trust_remote_code": True}),
+            )
+
+    if not counters_to_test:
+        # Default to all counters if no argument provided
+        counters_to_test = [
+            (BaseTokenCounter, "BaseTokenCounter", {}),
+            (OpenAITokenCounter, "OpenAITokenCounter", {}),
+            (HFTokenCounter, "HFTokenCounter", {"model_name": args.hf_model, "trust_remote_code": True}),
+        ]
+        print("No counter specified, defaulting to test all counters")
+        print("Use --base/--openai/--hf to test specific ones\n")
+
+    # Run tests for each counter
+    for counter_class, counter_name, kwargs in counters_to_test:
+        try:
+            run_all_tests_for_counter(counter_class, counter_name, **kwargs)
+        except Exception as e:
+            print(f"\n✗ FAILED: {counter_name} tests failed with error:")
+            print(f"  {type(e).__name__}: {e}")
+            raise
+
+    # Final summary
+    print(f"\n\n{'#' * 60}")
+    print("# TEST SUMMARY")
+    print(f"{'#' * 60}")
+    print(f"✓ All tests passed for {len(counters_to_test)} token counter(s):")
+    for _, counter_name, _ in counters_to_test:
+        print(f"  - {counter_name}")
+    print(f"{'#' * 60}\n")
+
+
+if __name__ == "__main__":
+    main()