From af52c0cc04f3074c9995968a23e04270fbd5493b Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Wed, 31 Dec 2025 11:00:57 +0800 Subject: [PATCH 01/11] feat(embedding): add embedding model framework with OpenAI implementation - Create BaseEmbeddingModel abstract class with async/sync interfaces - Implement OpenAIEmbeddingModel for asynchronous operations - Implement OpenAIEmbeddingModelSync for synchronous operations - Add automatic batching and retry logic for embedding operations - Add VectorNode embedding functionality for content indexing - Update type hints in LLM modules to use generic list instead of List - Refactor LiteLLMSync class documentation and type annotations - Add comprehensive async and sync unit tests for embedding models - Add unit tests for BaseContext attribute access patterns --- reme_ai/core/context/service_context.py | 5 +- reme_ai/core/embedding/__init__.py | 11 + .../core/embedding/base_embedding_model.py | 156 ++++++++ .../core/embedding/openai_embedding_model.py | 52 +++ .../embedding/openai_embedding_model_sync.py | 33 ++ reme_ai/core/llm/base_llm.py | 56 +-- reme_ai/core/llm/lite_llm_sync.py | 54 +-- reme_ai/core/llm/openai_llm.py | 12 +- reme_ai/core/llm/openai_llm_sync.py | 8 +- tests/test_base_context.py | 72 ++++ tests/test_embedding.py | 349 ++++++++++++++++++ tests/test_embedding_sync.py | 348 +++++++++++++++++ 12 files changed, 1069 insertions(+), 87 deletions(-) create mode 100644 reme_ai/core/embedding/__init__.py create mode 100644 reme_ai/core/embedding/base_embedding_model.py create mode 100644 reme_ai/core/embedding/openai_embedding_model.py create mode 100644 reme_ai/core/embedding/openai_embedding_model_sync.py create mode 100644 tests/test_base_context.py create mode 100644 tests/test_embedding.py create mode 100644 tests/test_embedding_sync.py diff --git a/reme_ai/core/context/service_context.py b/reme_ai/core/context/service_context.py index af99a870..65702be1 100644 --- a/reme_ai/core/context/service_context.py +++ b/reme_ai/core/context/service_context.py @@ -1,7 +1,6 @@ """Module for managing global service configurations and component registries via a singleton context.""" from concurrent.futures import ThreadPoolExecutor -from typing import Dict from .base_context import BaseContext from .registry import Registry @@ -21,10 +20,10 @@ class ServiceContext(BaseContext): self.service_config: ServiceConfig | None = None self.language: str = "" self.thread_pool: ThreadPoolExecutor | None = None - self.vector_store_dict: Dict[str, dict] = {} + 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.registry_dict: dict[RegistryEnum, Registry] = {v: Registry() for v in RegistryEnum.__members__.values()} self.flow_dict: dict = {} def register(self, name: str, register_type: RegistryEnum): diff --git a/reme_ai/core/embedding/__init__.py b/reme_ai/core/embedding/__init__.py new file mode 100644 index 00000000..c1d92375 --- /dev/null +++ b/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", +] diff --git a/reme_ai/core/embedding/base_embedding_model.py b/reme_ai/core/embedding/base_embedding_model.py new file mode 100644 index 00000000..3725f96d --- /dev/null +++ b/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.""" diff --git a/reme_ai/core/embedding/openai_embedding_model.py b/reme_ai/core/embedding/openai_embedding_model.py new file mode 100644 index 00000000..e10c7e66 --- /dev/null +++ b/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() diff --git a/reme_ai/core/embedding/openai_embedding_model_sync.py b/reme_ai/core/embedding/openai_embedding_model_sync.py new file mode 100644 index 00000000..b777e6d6 --- /dev/null +++ b/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() diff --git a/reme_ai/core/llm/base_llm.py b/reme_ai/core/llm/base_llm.py index 6f1da121..a7ba0f76 100644 --- a/reme_ai/core/llm/base_llm.py +++ b/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, diff --git a/reme_ai/core/llm/lite_llm_sync.py b/reme_ai/core/llm/lite_llm_sync.py index 5f814a2e..8015a275 100644 --- a/reme_ai/core/llm/lite_llm_sync.py +++ b/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 diff --git a/reme_ai/core/llm/openai_llm.py b/reme_ai/core/llm/openai_llm.py index 2c593d6f..df368021 100644 --- a/reme_ai/core/llm/openai_llm.py +++ b/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 diff --git a/reme_ai/core/llm/openai_llm_sync.py b/reme_ai/core/llm/openai_llm_sync.py index 3aadab1a..408ce272 100644 --- a/reme_ai/core/llm/openai_llm_sync.py +++ b/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 diff --git a/tests/test_base_context.py b/tests/test_base_context.py new file mode 100644 index 00000000..316a6796 --- /dev/null +++ b/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!") diff --git a/tests/test_embedding.py b/tests/test_embedding.py new file mode 100644 index 00000000..d1c404f5 --- /dev/null +++ b/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()) diff --git a/tests/test_embedding_sync.py b/tests/test_embedding_sync.py new file mode 100644 index 00000000..361a42b3 --- /dev/null +++ b/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() From fbdfdfee578037604b3d779e96263ef0ef4a9208 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Wed, 31 Dec 2025 11:17:12 +0800 Subject: [PATCH 02/11] feat(token-counter): add token counting system with multiple implementations --- docs/deprecated.txt | 4 +- reme_ai/core/token_counter/__init__.py | 11 + .../core/token_counter/base_token_counter.py | 59 ++ .../core/token_counter/hf_token_counter.py | 82 +++ .../token_counter/openai_token_counter.py | 58 ++ tests/test_token_counter.py | 511 ++++++++++++++++++ 6 files changed, 724 insertions(+), 1 deletion(-) create mode 100644 reme_ai/core/token_counter/__init__.py create mode 100644 reme_ai/core/token_counter/base_token_counter.py create mode 100644 reme_ai/core/token_counter/hf_token_counter.py create mode 100644 reme_ai/core/token_counter/openai_token_counter.py create mode 100644 tests/test_token_counter.py diff --git a/docs/deprecated.txt b/docs/deprecated.txt index 6d7cb14f..18618372 100644 --- a/docs/deprecated.txt +++ b/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，不使用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) diff --git a/reme_ai/core/token_counter/__init__.py b/reme_ai/core/token_counter/__init__.py new file mode 100644 index 00000000..a9b50826 --- /dev/null +++ b/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", +] diff --git a/reme_ai/core/token_counter/base_token_counter.py b/reme_ai/core/token_counter/base_token_counter.py new file mode 100644 index 00000000..76ccdb78 --- /dev/null +++ b/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) diff --git a/reme_ai/core/token_counter/hf_token_counter.py b/reme_ai/core/token_counter/hf_token_counter.py new file mode 100644 index 00000000..0bad4c78 --- /dev/null +++ b/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) diff --git a/reme_ai/core/token_counter/openai_token_counter.py b/reme_ai/core/token_counter/openai_token_counter.py new file mode 100644 index 00000000..5793e88a --- /dev/null +++ b/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 diff --git a/tests/test_token_counter.py b/tests/test_token_counter.py new file mode 100644 index 00000000..3c44a298 --- /dev/null +++ b/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() From 91a07e41860514f60db9112a5543312207614331 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Wed, 31 Dec 2025 13:18:33 +0800 Subject: [PATCH 03/11] feat(vector-store): add base vector store interface and multiple implementations --- docs/deprecated.txt | 2 +- reme_ai/core/vector_store/__init__.py | 17 + .../core/vector_store/base_vector_store.py | 92 ++ .../core/vector_store/chroma_vector_store.py | 397 +++++ reme_ai/core/vector_store/es_vector_store.py | 458 ++++++ .../core/vector_store/local_vector_store.py | 281 ++++ reme_ai/core/vector_store/pgvector_store.py | 533 +++++++ .../core/vector_store/qdrant_vector_store.py | 444 ++++++ tests/test_vector_store.py | 1417 +++++++++++++++++ 9 files changed, 3640 insertions(+), 1 deletion(-) create mode 100644 reme_ai/core/vector_store/__init__.py create mode 100644 reme_ai/core/vector_store/base_vector_store.py create mode 100644 reme_ai/core/vector_store/chroma_vector_store.py create mode 100644 reme_ai/core/vector_store/es_vector_store.py create mode 100644 reme_ai/core/vector_store/local_vector_store.py create mode 100644 reme_ai/core/vector_store/pgvector_store.py create mode 100644 reme_ai/core/vector_store/qdrant_vector_store.py create mode 100644 tests/test_vector_store.py diff --git a/docs/deprecated.txt b/docs/deprecated.txt index 18618372..06f9648f 100644 --- a/docs/deprecated.txt +++ b/docs/deprecated.txt @@ -1,6 +1,6 @@ from loguru import logger -用英文注释，完善module/class/function docstring，要一句话简洁，不要变更代码 +用英文注释，完善module/class/function docstring，要一句话简洁，不要变更代码逻辑，符合pep和pylint规范，使用list而不是typing.List/Dict，不使用typing.Union 看看代码有什么问题 用英文注释，完善module/class/function docstring，要一句话简洁，代码要简洁，符合pep和pylint规范，使用list而不是typing.List，不使用typing.Union diff --git a/reme_ai/core/vector_store/__init__.py b/reme_ai/core/vector_store/__init__.py new file mode 100644 index 00000000..79500294 --- /dev/null +++ b/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", +] diff --git a/reme_ai/core/vector_store/base_vector_store.py b/reme_ai/core/vector_store/base_vector_store.py new file mode 100644 index 00000000..e64b6dbf --- /dev/null +++ b/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.""" diff --git a/reme_ai/core/vector_store/chroma_vector_store.py b/reme_ai/core/vector_store/chroma_vector_store.py new file mode 100644 index 00000000..24b88d36 --- /dev/null +++ b/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") diff --git a/reme_ai/core/vector_store/es_vector_store.py b/reme_ai/core/vector_store/es_vector_store.py new file mode 100644 index 00000000..06dcafa5 --- /dev/null +++ b/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") diff --git a/reme_ai/core/vector_store/local_vector_store.py b/reme_ai/core/vector_store/local_vector_store.py new file mode 100644 index 00000000..25beef8d --- /dev/null +++ b/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") diff --git a/reme_ai/core/vector_store/pgvector_store.py b/reme_ai/core/vector_store/pgvector_store.py new file mode 100644 index 00000000..38c5667a --- /dev/null +++ b/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") diff --git a/reme_ai/core/vector_store/qdrant_vector_store.py b/reme_ai/core/vector_store/qdrant_vector_store.py new file mode 100644 index 00000000..97eb5b61 --- /dev/null +++ b/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") diff --git a/tests/test_vector_store.py b/tests/test_vector_store.py new file mode 100644 index 00000000..76de21e1 --- /dev/null +++ b/tests/test_vector_store.py @@ -0,0 +1,1417 @@ +# pylint: disable=too-many-lines +"""Unified test suite for vector store implementations. + +This module provides comprehensive test coverage for LocalVectorStore, ESVectorStore, +PGVectorStore, QdrantVectorStore, and ChromaVectorStore implementations. Tests can be +run for specific vector stores or all implementations. + +Usage: + python test_vector_store.py --local # Test LocalVectorStore only + python test_vector_store.py --es # Test ESVectorStore only + python test_vector_store.py --pgvector # Test PGVectorStore only + python test_vector_store.py --qdrant # Test QdrantVectorStore only + python test_vector_store.py --chroma # Test ChromaVectorStore only + python test_vector_store.py --all # Test all vector stores + +""" + +import argparse +import asyncio +import shutil +from pathlib import Path +from typing import List + +from loguru import logger + +from reme_ai.core.embedding import OpenAIEmbeddingModel +from reme_ai.core.schema import VectorNode +from reme_ai.core.vector_store import ( + BaseVectorStore, + ChromaVectorStore, + LocalVectorStore, + ESVectorStore, + PGVectorStore, + QdrantVectorStore, +) + + +# ==================== Configuration ==================== + + +class TestConfig: + """Configuration for test execution.""" + + # LocalVectorStore settings + LOCAL_ROOT_PATH = "./test_vector_store_local" + + # ESVectorStore settings + ES_HOSTS = "http://11.160.132.46:8200" + ES_BASIC_AUTH = None # Set to ("username", "password") if authentication is required + + # QdrantVectorStore settings + QDRANT_PATH = None # "./test_vector_store_qdrant" # For local mode + QDRANT_HOST = None # Set to host address for remote mode (e.g., "localhost") + QDRANT_PORT = None # Set to port for remote mode (e.g., 6333) + QDRANT_URL = "http://11.160.132.46:6333" # Alternative to host/port (e.g., http://localhost:6333) + QDRANT_API_KEY = None # Set for Qdrant Cloud authentication + + # PGVectorStore settings + PG_DSN = "postgresql://localhost/postgres" # PostgreSQL connection string + PG_MIN_SIZE = 1 # Minimum connections in pool + PG_MAX_SIZE = 5 # Maximum connections in pool + PG_USE_HNSW = True # Use HNSW index for faster search + PG_USE_DISKANN = False # Use DiskANN index (requires vectorscale extension) + + # ChromaVectorStore settings + CHROMA_PATH = "./test_vector_store_chroma" # For local persistent mode + CHROMA_HOST = None # Set to host address for remote mode (e.g., "localhost") + CHROMA_PORT = None # Set to port for remote mode (e.g., 8000) + CHROMA_API_KEY = None # Set for ChromaDB Cloud authentication + CHROMA_TENANT = None # Set for ChromaDB Cloud tenant + CHROMA_DATABASE = None # Set for ChromaDB Cloud database + + # Embedding model settings + EMBEDDING_MODEL_NAME = "text-embedding-v4" + EMBEDDING_DIMENSIONS = 64 + + # Test collection naming + TEST_COLLECTION_PREFIX = "test_vector_store" + + +# ==================== Sample Data Generator ==================== + + +class SampleDataGenerator: + """Generator for sample test data.""" + + @staticmethod + def create_sample_nodes(prefix: str = "") -> List[VectorNode]: + """Create sample VectorNode instances for testing. + + Args: + prefix: Optional prefix for vector_id to avoid conflicts + + Returns: + List[VectorNode]: List of sample nodes with diverse metadata + """ + id_prefix = f"{prefix}_" if prefix else "" + return [ + VectorNode( + vector_id=f"{id_prefix}node1", + content="Artificial intelligence is a technology that simulates human intelligence.", + metadata={ + "node_type": "tech", + "category": "AI", + "source": "research", + "priority": "high", + "year": "2023", + "department": "engineering", + "language": "english", + "status": "published", + }, + ), + VectorNode( + vector_id=f"{id_prefix}node2", + content="Machine learning is a subset of artificial intelligence.", + metadata={ + "node_type": "tech", + "category": "ML", + "source": "research", + "priority": "high", + "year": "2022", + "department": "engineering", + "language": "english", + "status": "published", + }, + ), + VectorNode( + vector_id=f"{id_prefix}node3", + content="Deep learning uses neural networks with multiple layers.", + metadata={ + "node_type": "tech_new", + "category": "DL", + "source": "blog", + "priority": "medium", + "year": "2024", + "department": "marketing", + "language": "chinese", + "status": "draft", + }, + ), + VectorNode( + vector_id=f"{id_prefix}node4", + content="I love eating delicious seafood, especially fresh fish.", + metadata={ + "node_type": "food", + "category": "preference", + "source": "personal", + "priority": "low", + "year": "2023", + "department": "lifestyle", + "language": "english", + "status": "published", + }, + ), + VectorNode( + vector_id=f"{id_prefix}node5", + content="Natural language processing enables computers to understand human language.", + metadata={ + "node_type": "tech", + "category": "NLP", + "source": "research", + "priority": "high", + "year": "2024", + "department": "engineering", + "language": "english", + "status": "review", + }, + ), + ] + + +# ==================== Vector Store Factory ==================== + + +def get_store_type(store: BaseVectorStore) -> str: + """Get the type identifier of a vector store instance. + + Args: + store: Vector store instance + + Returns: + str: Type identifier ("local", "es", "pgvector", "qdrant", or "chroma") + """ + if isinstance(store, LocalVectorStore): + return "local" + elif isinstance(store, QdrantVectorStore): + return "qdrant" + elif isinstance(store, ESVectorStore): + return "es" + elif isinstance(store, PGVectorStore): + return "pgvector" + elif isinstance(store, ChromaVectorStore): + return "chroma" + else: + raise ValueError(f"Unknown vector store type: {type(store)}") + + +def create_vector_store(store_type: str, collection_name: str) -> BaseVectorStore: + """Create a vector store instance based on type. + + Args: + store_type: Type of vector store ("local", "es", or "qdrant") + collection_name: Name of the collection + + Returns: + BaseVectorStore: Initialized vector store instance + """ + config = TestConfig() + + # Initialize embedding model + embedding_model = OpenAIEmbeddingModel( + model_name=config.EMBEDDING_MODEL_NAME, + dimensions=config.EMBEDDING_DIMENSIONS, + ) + + if store_type == "local": + return LocalVectorStore( + collection_name=collection_name, + embedding_model=embedding_model, + root_path=config.LOCAL_ROOT_PATH, + ) + elif store_type == "es": + return ESVectorStore( + collection_name=collection_name, + embedding_model=embedding_model, + hosts=config.ES_HOSTS, + basic_auth=config.ES_BASIC_AUTH, + ) + elif store_type == "qdrant": + return QdrantVectorStore( + collection_name=collection_name, + embedding_model=embedding_model, + path=config.QDRANT_PATH, + host=config.QDRANT_HOST, + port=config.QDRANT_PORT, + url=config.QDRANT_URL, + api_key=config.QDRANT_API_KEY, + distance="cosine", + on_disk=False, + ) + elif store_type == "pgvector": + return PGVectorStore( + collection_name=collection_name, + embedding_model=embedding_model, + dsn=config.PG_DSN, + min_size=config.PG_MIN_SIZE, + max_size=config.PG_MAX_SIZE, + use_hnsw=config.PG_USE_HNSW, + use_diskann=config.PG_USE_DISKANN, + ) + elif store_type == "chroma": + return ChromaVectorStore( + collection_name=collection_name, + embedding_model=embedding_model, + path=config.CHROMA_PATH, + host=config.CHROMA_HOST, + port=config.CHROMA_PORT, + api_key=config.CHROMA_API_KEY, + tenant=config.CHROMA_TENANT, + database=config.CHROMA_DATABASE, + ) + else: + raise ValueError(f"Unknown store type: {store_type}") + + +# ==================== Test Functions ==================== + + +async def test_create_collection(store: BaseVectorStore, _store_name: str): + """Test collection creation.""" + logger.info("=" * 20 + " CREATE COLLECTION TEST " + "=" * 20) + + # Clean up if exists + collections = await store.list_collections() + if store.collection_name in collections: + await store.delete_collection(store.collection_name) + logger.info(f"Cleaned up existing collection: {store.collection_name}") + + # Create collection + await store.create_collection(store.collection_name) + + # Verify creation + collections = await store.list_collections() + assert store.collection_name in collections, "Collection should exist after creation" + logger.info(f"✓ Created collection: {store.collection_name}") + + +async def test_insert(store: BaseVectorStore, _store_name: str) -> List[VectorNode]: + """Test node insertion (single and batch).""" + logger.info("=" * 20 + " INSERT TEST " + "=" * 20) + + # Test single node insertion + single_node = VectorNode( + vector_id="test_single_insert", + content="This is a single node insertion test", + metadata={"test_type": "single_insert"}, + ) + await store.insert(single_node) + logger.info("✓ Inserted single node") + + # Test batch insertion + sample_nodes = SampleDataGenerator.create_sample_nodes("test") + await store.insert(sample_nodes) + logger.info(f"✓ Batch inserted {len(sample_nodes)} nodes") + + # Verify total insertions + all_nodes = await store.list(limit=20) + assert len(all_nodes) >= len(sample_nodes) + 1, "Should have at least sample nodes + single node" + logger.info(f"✓ Total nodes in collection: {len(all_nodes)}") + + return sample_nodes + + +async def test_search(store: BaseVectorStore, _store_name: str): + """Test basic vector search.""" + logger.info("=" * 20 + " SEARCH TEST " + "=" * 20) + + results = await store.search( + query="What is artificial intelligence?", + limit=3, + ) + + logger.info(f"Search returned {len(results)} results") + for i, r in enumerate(results, 1): + score = r.metadata.get("_score", "N/A") + logger.info(f" Result {i}: {r.content[:60]}... (score: {score})") + + assert len(results) > 0, "Search should return results" + logger.info("✓ Basic search test passed") + + +async def test_search_with_single_filter(store: BaseVectorStore, _store_name: str): + """Test vector search with single metadata filter.""" + logger.info("=" * 20 + " SINGLE FILTER SEARCH TEST " + "=" * 20) + + # Test single value filter + filters = {"node_type": "tech"} + results = await store.search( + query="What is artificial intelligence?", + limit=5, + filters=filters, + ) + + logger.info(f"Filtered search (node_type=tech) returned {len(results)} results") + for i, r in enumerate(results, 1): + node_type = r.metadata.get("node_type") + logger.info(f" Result {i}: type={node_type}, content={r.content[:50]}...") + assert node_type == "tech", "Result should have node_type='tech'" + + logger.info("✓ Single filter search test passed") + + +async def test_search_with_list_filter(store: BaseVectorStore, _store_name: str): + """Test vector search with list filter (IN operation).""" + logger.info("=" * 20 + " LIST FILTER SEARCH TEST " + "=" * 20) + + # Test list filter (IN operation) + filters = {"node_type": ["tech", "tech_new"]} + results = await store.search( + query="What is artificial intelligence?", + limit=5, + filters=filters, + ) + + logger.info(f"Filtered search (node_type IN [tech, tech_new]) returned {len(results)} results") + for i, r in enumerate(results, 1): + node_type = r.metadata.get("node_type") + logger.info(f" Result {i}: type={node_type}, content={r.content[:50]}...") + assert node_type in ["tech", "tech_new"], "Result should have node_type in [tech, tech_new]" + + logger.info("✓ List filter search test passed") + + +async def test_search_with_multiple_filters(store: BaseVectorStore, _store_name: str): + """Test vector search with multiple metadata filters (AND operation).""" + logger.info("=" * 20 + " MULTIPLE FILTERS SEARCH TEST " + "=" * 20) + + # Test multiple filters (AND operation) + filters = { + "node_type": ["tech", "tech_new"], + "source": "research", + } + results = await store.search( + query="What is artificial intelligence?", + limit=5, + filters=filters, + ) + + logger.info( + f"Multi-filter search (node_type IN [tech, tech_new] AND source=research) " f"returned {len(results)} results", + ) + for i, r in enumerate(results, 1): + node_type = r.metadata.get("node_type") + source = r.metadata.get("source") + logger.info(f" Result {i}: type={node_type}, source={source}, content={r.content[:40]}...") + assert node_type in ["tech", "tech_new"], "Result should have node_type in [tech, tech_new]" + assert source == "research", "Result should have source='research'" + + logger.info("✓ Multiple filters search test passed") + + +async def test_get_by_id(store: BaseVectorStore, _store_name: str): + """Test retrieving nodes by vector_id (single and batch).""" + logger.info("=" * 20 + " GET BY ID TEST " + "=" * 20) + + # Test single ID retrieval + target_id = "test_node1" + result = await store.get(target_id) + + assert isinstance(result, VectorNode), "Should return a VectorNode for single ID" + assert result.vector_id == target_id, f"Result should have vector_id={target_id}" + logger.info(f"✓ Retrieved single node: {result.vector_id}") + + # Test batch retrieval (small batch) + target_ids = ["test_node1", "test_node2"] + results = await store.get(target_ids) + + assert isinstance(results, list), "Should return a list for multiple IDs" + assert len(results) == 2, f"Should return 2 results, got {len(results)}" + result_ids = {r.vector_id for r in results} + assert result_ids == set(target_ids), f"Result IDs should match {target_ids}" + logger.info(f"✓ Batch retrieved {len(results)} nodes") + + # Test larger batch retrieval + large_batch_ids = ["test_node1", "test_node2", "test_node3", "test_node5"] + large_results = await store.get(large_batch_ids) + assert isinstance(large_results, list), "Should return a list for batch IDs" + assert len(large_results) >= 3, "Should return at least 3 results" + logger.info(f"✓ Large batch retrieved {len(large_results)} nodes") + + +async def test_list_all(store: BaseVectorStore, _store_name: str): + """Test listing all nodes in collection.""" + logger.info("=" * 20 + " LIST ALL TEST " + "=" * 20) + + results = await store.list(limit=10) + + logger.info(f"Collection contains {len(results)} nodes") + for i, node in enumerate(results, 1): + logger.info(f" Node {i}: id={node.vector_id}, content={node.content[:50]}...") + + assert len(results) > 0, "Collection should contain nodes" + logger.info("✓ List all nodes test passed") + + +async def test_list_with_filters(store: BaseVectorStore, _store_name: str): + """Test listing nodes with metadata filters.""" + logger.info("=" * 20 + " LIST WITH FILTERS TEST " + "=" * 20) + + filters = {"category": "AI"} + results = await store.list(filters=filters, limit=10) + + logger.info(f"Filtered list (category=AI) returned {len(results)} nodes") + for i, node in enumerate(results, 1): + category = node.metadata.get("category") + logger.info(f" Node {i}: category={category}, id={node.vector_id}") + assert category == "AI", "All nodes should have category=AI" + + logger.info("✓ List with filters test passed") + + +async def test_update(store: BaseVectorStore, _store_name: str): + """Test updating existing nodes (single and batch).""" + logger.info("=" * 20 + " UPDATE TEST " + "=" * 20) + + # Test single node update + updated_node = VectorNode( + vector_id="test_node2", + content="Machine learning is a powerful subset of AI that learns from data.", + metadata={ + "node_type": "tech", + "category": "ML", + "updated": "true", + "update_timestamp": "2024-12-26", + }, + ) + + await store.update(updated_node) + + # Verify single update + result = await store.get("test_node2") + assert "updated" in result.metadata, "Updated metadata should be present" + logger.info(f"✓ Updated single node: {result.vector_id}") + logger.info(f" New content: {result.content[:60]}...") + + # Test batch update (update multiple nodes at once) + batch_update_nodes = [ + VectorNode( + vector_id="test_node1", + content="Artificial intelligence is evolving rapidly with new breakthroughs.", + metadata={ + "node_type": "tech", + "category": "AI", + "batch_updated": "true", + "update_timestamp": "2024-12-31", + }, + ), + VectorNode( + vector_id="test_node3", + content="Deep learning revolutionizes neural network architectures.", + metadata={ + "node_type": "tech_new", + "category": "DL", + "batch_updated": "true", + "update_timestamp": "2024-12-31", + }, + ), + ] + + await store.update(batch_update_nodes) + logger.info(f"✓ Batch updated {len(batch_update_nodes)} nodes") + + # Verify batch updates + results = await store.get(["test_node1", "test_node3"]) + for r in results: + assert r.metadata.get("batch_updated") == "true", f"Node {r.vector_id} should have batch_updated metadata" + logger.info(f"✓ Verified batch update for {len(results)} nodes") + + +async def test_delete(store: BaseVectorStore, _store_name: str): + """Test deleting nodes (single and batch).""" + logger.info("=" * 20 + " DELETE TEST " + "=" * 20) + + # Test single node deletion + node_to_delete = "test_node4" + await store.delete(node_to_delete) + + # Verify single deletion - try to get the deleted node + try: + result = await store.get(node_to_delete) + # If result is empty list or None, deletion was successful + if isinstance(result, list): + assert len(result) == 0, "Deleted node should not be retrievable" + else: + assert result is None, "Deleted node should not be retrievable" + except Exception: + pass # Expected if node doesn't exist + + logger.info(f"✓ Deleted single node: {node_to_delete}") + + # Test batch deletion - first insert some nodes to delete + batch_delete_nodes = [ + VectorNode( + vector_id=f"delete_test_{i}", + content=f"Node to be deleted {i}", + metadata={"test_type": "delete_batch"}, + ) + for i in range(5) + ] + await store.insert(batch_delete_nodes) + logger.info(f"✓ Inserted {len(batch_delete_nodes)} nodes for batch delete test") + + # Batch delete + delete_ids = [f"delete_test_{i}" for i in range(5)] + await store.delete(delete_ids) + logger.info(f"✓ Batch deleted {len(delete_ids)} nodes") + + # Verify batch deletion + try: + results = await store.get(delete_ids) + if isinstance(results, list): + assert len(results) == 0, "All deleted nodes should not be retrievable" + except Exception: + pass # Expected if nodes don't exist + logger.info("✓ Verified batch deletion") + + +async def test_copy_collection(store: BaseVectorStore, store_name: str): + """Test copying a collection.""" + logger.info("=" * 20 + " COPY COLLECTION TEST " + "=" * 20) + + config = TestConfig() + copy_collection_name = f"{config.TEST_COLLECTION_PREFIX}_{store_name}_copy" + + # Elasticsearch and PostgreSQL require lowercase table/index names + store_type = get_store_type(store) + if store_type in ("es", "pgvector"): + copy_collection_name = copy_collection_name.lower() + + # Clean up if exists + collections = await store.list_collections() + if copy_collection_name in collections: + await store.delete_collection(copy_collection_name) + + # Copy collection + await store.copy_collection(copy_collection_name) + + # Verify copy + collections = await store.list_collections() + assert copy_collection_name in collections, "Copied collection should exist" + logger.info(f"✓ Copied collection to: {copy_collection_name}") + + # Verify content in copied collection + copied_store = create_vector_store(store_type, copy_collection_name) + copied_nodes = await copied_store.list() + logger.info(f"✓ Copied collection has {len(copied_nodes)} nodes") + await copied_store.close() + + # Clean up copied collection + await store.delete_collection(copy_collection_name) + logger.info("✓ Cleaned up copied collection") + + +async def test_list_collections(store: BaseVectorStore, _store_name: str): + """Test listing all collections.""" + logger.info("=" * 20 + " LIST COLLECTIONS TEST " + "=" * 20) + + collections = await store.list_collections() + + logger.info(f"Found {len(collections)} collections") + config = TestConfig() + test_collections = [c for c in collections if c.startswith(config.TEST_COLLECTION_PREFIX)] + logger.info(f" Test collections: {test_collections}") + + assert store.collection_name in collections, "Main test collection should be listed" + logger.info("✓ List collections test passed") + + +async def test_delete_collection(store: BaseVectorStore, _store_name: str): + """Test deleting a collection.""" + logger.info("=" * 20 + " DELETE COLLECTION TEST " + "=" * 20) + + await store.delete_collection(store.collection_name) + + # Verify deletion + collections = await store.list_collections() + assert store.collection_name not in collections, "Collection should not exist after deletion" + logger.info(f"✓ Deleted collection: {store.collection_name}") + + +async def test_cosine_similarity(store_name: str): + """Test manual cosine similarity calculation (LocalVectorStore only).""" + if store_name != "LocalVectorStore": + logger.info("=" * 20 + " COSINE SIMILARITY TEST (SKIPPED) " + "=" * 20) + logger.info("⊘ Skipped: Only applicable to LocalVectorStore") + return + + logger.info("=" * 20 + " COSINE SIMILARITY TEST " + "=" * 20) + + vec1 = [1.0, 0.0, 0.0] + vec2 = [0.0, 1.0, 0.0] + vec3 = [1.0, 0.0, 0.0] + + # Test perpendicular vectors (similarity = 0) + sim1 = LocalVectorStore._cosine_similarity(vec1, vec2) # pylint: disable=protected-access + logger.info(f"Similarity between perpendicular vectors: {sim1:.4f}") + assert abs(sim1) < 0.0001, "Perpendicular vectors should have similarity close to 0" + + # Test identical vectors (similarity = 1) + sim2 = LocalVectorStore._cosine_similarity(vec1, vec3) # pylint: disable=protected-access + logger.info(f"Similarity between identical vectors: {sim2:.4f}") + assert abs(sim2 - 1.0) < 0.0001, "Identical vectors should have similarity close to 1" + + # Test with real-world like vectors + vec4 = [0.5, 0.5, 0.5] + vec5 = [0.6, 0.4, 0.5] + sim3 = LocalVectorStore._cosine_similarity(vec4, vec5) # pylint: disable=protected-access + logger.info(f"Similarity between similar vectors: {sim3:.4f}") + assert sim3 > 0.9, "Similar vectors should have high similarity" + + logger.info("✓ Cosine similarity tests passed") + + +async def test_batch_operations(store: BaseVectorStore, _store_name: str): + """Test large-scale batch insert, update, and delete operations. + + This test validates the efficiency and correctness of batch operations + by processing 100 nodes at once, which is more realistic for production use. + """ + logger.info("=" * 20 + " BATCH OPERATIONS TEST " + "=" * 20) + + # Create a large batch of nodes (100 nodes) + batch_nodes = [] + for i in range(100): + batch_nodes.append( + VectorNode( + vector_id=f"batch_node_{i}", + content=f"This is batch test content number {i} about various topics in technology and science.", + metadata={ + "batch_id": str(i // 10), # Group into batches of 10 + "index": str(i), + "category": ["tech", "science", "business"][i % 3], + "priority": ["high", "medium", "low"][i % 3], + }, + ), + ) + + # Batch insert + await store.insert(batch_nodes) + logger.info(f"✓ Inserted {len(batch_nodes)} nodes in batch") + + # Verify batch insert + results = await store.list(limit=150) + assert len(results) >= 100, f"Should have at least 100 nodes, got {len(results)}" + logger.info(f"✓ Verified batch insert: {len(results)} total nodes") + + # Batch update (update first 20 nodes) + update_nodes = [] + for i in range(20): + update_nodes.append( + VectorNode( + vector_id=f"batch_node_{i}", + content=f"UPDATED: This is updated batch content {i}", + metadata={ + "batch_id": str(i // 10), + "index": str(i), + "updated": "true", + "update_timestamp": "2024-12-31", + }, + ), + ) + + await store.update(update_nodes) + logger.info(f"✓ Updated {len(update_nodes)} nodes in batch") + + # Verify updates + updated_results = await store.list(filters={"updated": "true"}, limit=50) + assert len(updated_results) >= 20, "Should have at least 20 updated nodes" + logger.info(f"✓ Verified batch update: {len(updated_results)} updated nodes") + + # Batch delete (delete nodes with batch_id >= 5) + delete_ids = [f"batch_node_{i}" for i in range(50, 100)] + await store.delete(delete_ids) + logger.info(f"✓ Deleted {len(delete_ids)} nodes in batch") + + # Verify deletions + remaining = await store.list(limit=150) + batch_nodes_remaining = [n for n in remaining if n.vector_id.startswith("batch_node_")] + assert len(batch_nodes_remaining) <= 50, "Should have at most 50 batch nodes remaining" + logger.info(f"✓ Verified batch delete: {len(batch_nodes_remaining)} nodes remaining") + + +async def test_complex_metadata_queries(store: BaseVectorStore, _store_name: str): + """Test complex metadata filtering with nested conditions.""" + logger.info("=" * 20 + " COMPLEX METADATA QUERIES TEST " + "=" * 20) + + # Insert nodes with rich metadata + complex_nodes = [ + VectorNode( + vector_id="complex_1", + content="Advanced neural networks for computer vision applications", + metadata={ + "domain": "AI", + "subdomain": "computer_vision", + "year": "2024", + "citations": "150", + "impact_factor": "high", + "tags": "neural_networks,vision,deep_learning", + }, + ), + VectorNode( + vector_id="complex_2", + content="Natural language processing with transformer models", + metadata={ + "domain": "AI", + "subdomain": "nlp", + "year": "2023", + "citations": "200", + "impact_factor": "high", + "tags": "transformers,nlp,language_models", + }, + ), + VectorNode( + vector_id="complex_3", + content="Reinforcement learning for robotics control", + metadata={ + "domain": "AI", + "subdomain": "robotics", + "year": "2024", + "citations": "80", + "impact_factor": "medium", + "tags": "reinforcement_learning,robotics,control", + }, + ), + VectorNode( + vector_id="complex_4", + content="Quantum computing algorithms and applications", + metadata={ + "domain": "quantum", + "subdomain": "algorithms", + "year": "2024", + "citations": "50", + "impact_factor": "medium", + "tags": "quantum,algorithms,computing", + }, + ), + ] + + await store.insert(complex_nodes) + logger.info(f"✓ Inserted {len(complex_nodes)} nodes with complex metadata") + + # Test 1: Multiple field filters with list values + filters_1 = { + "domain": "AI", + "year": ["2023", "2024"], + "impact_factor": "high", + } + results_1 = await store.search( + query="artificial intelligence research", + limit=10, + filters=filters_1, + ) + logger.info(f"Test 1 - AI + high impact + recent years: {len(results_1)} results") + for r in results_1: + assert r.metadata.get("domain") == "AI" + assert r.metadata.get("impact_factor") == "high" + assert r.metadata.get("year") in ["2023", "2024"] + + # Test 2: List filter with multiple subdomains + filters_2 = { + "subdomain": ["nlp", "computer_vision"], + } + results_2 = await store.search( + query="deep learning applications", + limit=10, + filters=filters_2, + ) + logger.info(f"Test 2 - NLP or Computer Vision: {len(results_2)} results") + for r in results_2: + assert r.metadata.get("subdomain") in ["nlp", "computer_vision"] + + # Test 3: Year-based filtering + filters_3 = { + "year": "2024", + } + results_3 = await store.list(filters=filters_3, limit=10) + logger.info(f"Test 3 - Year 2024 only: {len(results_3)} results") + for r in results_3: + assert r.metadata.get("year") == "2024" + + logger.info("✓ Complex metadata queries test passed") + + +async def test_edge_cases(store: BaseVectorStore, _store_name: str): + """Test edge cases and boundary conditions.""" + logger.info("=" * 20 + " EDGE CASES TEST " + "=" * 20) + + # Test 1: Empty content + edge_node_1 = VectorNode( + vector_id="edge_empty_content", + content="", + metadata={"type": "empty"}, + ) + try: + await store.insert([edge_node_1]) + logger.info("✓ Handled empty content node") + except Exception as e: + logger.info(f"⊘ Empty content not supported: {e}") + + # Test 2: Very long content + edge_node_2 = VectorNode( + vector_id="edge_long_content", + content="A" * 10000, # 10k characters + metadata={"type": "long_content"}, + ) + await store.insert([edge_node_2]) + result = await store.get("edge_long_content") + assert len(result.content) == 10000 + logger.info("✓ Handled very long content (10k chars)") + + # Test 3: Special characters in content + edge_node_3 = VectorNode( + vector_id="edge_special_chars", + content="Special chars: @#$%^&*()[]{}|\\;:'\",.<>?/~`+=−×÷", + metadata={"type": "special_chars"}, + ) + await store.insert([edge_node_3]) + result = await store.get("edge_special_chars") + assert "@#$%^&*()" in result.content + logger.info("✓ Handled special characters in content") + + # Test 4: Unicode and emoji content + edge_node_4 = VectorNode( + vector_id="edge_unicode", + content="Unicode test: 你好世界 🌍 مرحبا العالم Привет мир", + metadata={"type": "unicode", "language": "multi"}, + ) + await store.insert([edge_node_4]) + result = await store.get("edge_unicode") + assert "你好世界" in result.content + assert "🌍" in result.content + logger.info("✓ Handled unicode and emoji content") + + # Test 5: Search with empty query + try: + results = await store.search(query="", limit=5) + logger.info(f"✓ Empty query returned {len(results)} results") + except Exception as e: + logger.info(f"⊘ Empty query not supported: {e}") + + # Test 6: Search with very high limit + results = await store.search(query="test", limit=1000) + logger.info(f"✓ High limit search returned {len(results)} results") + + # Test 7: Get non-existent ID + result = await store.get("non_existent_id_12345") + if isinstance(result, list): + assert len(result) == 0, "Non-existent ID should return empty list" + else: + assert result is None, "Non-existent ID should return None" + logger.info("✓ Handled non-existent ID gracefully") + + # Test 8: Metadata with empty string values + edge_node_5 = VectorNode( + vector_id="edge_empty_metadata", + content="Testing empty string values in metadata", + metadata={"field1": "value1", "field2": "", "field3": "value3"}, + ) + await store.insert([edge_node_5]) + logger.info("✓ Handled empty string values in metadata") + + logger.info("✓ Edge cases test passed") + + +async def test_concurrent_operations(store: BaseVectorStore, _store_name: str): + """Test concurrent read/write operations.""" + logger.info("=" * 20 + " CONCURRENT OPERATIONS TEST " + "=" * 20) + + # Prepare concurrent insert nodes + concurrent_nodes = [ + VectorNode( + vector_id=f"concurrent_{i}", + content=f"Concurrent test content {i}", + metadata={"thread_id": str(i % 5), "index": str(i)}, + ) + for i in range(50) + ] + + # Test concurrent inserts + insert_tasks = [] + for i in range(0, 50, 10): + batch = concurrent_nodes[i : i + 10] + insert_tasks.append(store.insert(batch)) + + await asyncio.gather(*insert_tasks) + logger.info("✓ Completed concurrent inserts") + + # Test concurrent searches + search_tasks = [store.search(query=f"concurrent test {i}", limit=5) for i in range(10)] + search_results = await asyncio.gather(*search_tasks) + logger.info(f"✓ Completed {len(search_results)} concurrent searches") + + # Test concurrent reads + get_tasks = [store.get(f"concurrent_{i}") for i in range(0, 50, 5)] + get_results = await asyncio.gather(*get_tasks) + logger.info(f"✓ Completed {len(get_results)} concurrent reads") + + # Test batch updates (using batch update instead of concurrent individual updates) + update_nodes = [ + VectorNode( + vector_id=f"concurrent_{i}", + content=f"UPDATED concurrent content {i}", + metadata={"thread_id": str(i % 5), "updated": "true"}, + ) + for i in range(0, 20, 2) + ] + await store.update(update_nodes) + logger.info(f"✓ Completed batch update of {len(update_nodes)} nodes") + + logger.info("✓ Concurrent operations test passed") + + +async def test_search_relevance_ranking(store: BaseVectorStore, _store_name: str): + """Test search result relevance and ranking.""" + logger.info("=" * 20 + " SEARCH RELEVANCE RANKING TEST " + "=" * 20) + + # Insert nodes with varying relevance + relevance_nodes = [ + VectorNode( + vector_id="relevance_exact", + content="Machine learning is a subset of artificial intelligence focused on learning from data.", + metadata={"relevance": "exact"}, + ), + VectorNode( + vector_id="relevance_high", + content="Artificial intelligence and machine learning are transforming technology.", + metadata={"relevance": "high"}, + ), + VectorNode( + vector_id="relevance_medium", + content="Deep learning uses neural networks for pattern recognition.", + metadata={"relevance": "medium"}, + ), + VectorNode( + vector_id="relevance_low", + content="Software engineering best practices for code quality.", + metadata={"relevance": "low"}, + ), + VectorNode( + vector_id="relevance_none", + content="Cooking recipes for delicious Italian pasta dishes.", + metadata={"relevance": "none"}, + ), + ] + + await store.insert(relevance_nodes) + logger.info(f"✓ Inserted {len(relevance_nodes)} nodes with varying relevance") + + # Search with specific query + query = "What is machine learning and artificial intelligence?" + results = await store.search(query=query, limit=5) + + logger.info(f"Search results for: '{query}'") + for i, result in enumerate(results, 1): + score = result.metadata.get("_score", "N/A") + relevance = result.metadata.get("relevance", "unknown") + logger.info(f" {i}. [{relevance}] score={score}: {result.content[:60]}...") + + # Verify that more relevant results appear first + if len(results) >= 2: + # The exact match should be in top results + top_ids = [r.vector_id for r in results[:3]] + assert ( + "relevance_exact" in top_ids or "relevance_high" in top_ids + ), "Most relevant results should appear in top 3" + logger.info("✓ Relevance ranking verified") + + # Test with different query + query2 = "neural networks deep learning" + results2 = await store.search(query=query2, limit=5) + logger.info(f"\nSearch results for: '{query2}'") + for i, result in enumerate(results2, 1): + score = result.metadata.get("_score", "N/A") + logger.info(f" {i}. score={score}: {result.content[:60]}...") + + logger.info("✓ Search relevance ranking test passed") + + +async def test_metadata_statistics(store: BaseVectorStore, _store_name: str): + """Test metadata aggregation and statistics.""" + logger.info("=" * 20 + " METADATA STATISTICS TEST " + "=" * 20) + + # Get all nodes and analyze metadata + all_nodes = await store.list(limit=500) + logger.info(f"Total nodes in collection: {len(all_nodes)}") + + # Count by category + category_counts = {} + for node in all_nodes: + category = node.metadata.get("category", "unknown") + category_counts[category] = category_counts.get(category, 0) + 1 + + logger.info("Category distribution:") + for category, count in sorted(category_counts.items()): + logger.info(f" {category}: {count}") + + # Count by node_type + type_counts = {} + for node in all_nodes: + node_type = node.metadata.get("node_type", "unknown") + type_counts[node_type] = type_counts.get(node_type, 0) + 1 + + logger.info("Node type distribution:") + for node_type, count in sorted(type_counts.items()): + logger.info(f" {node_type}: {count}") + + # Verify we can filter by each category + for category in category_counts: + if category != "unknown": + filtered = await store.list(filters={"category": category}, limit=100) + logger.info(f"✓ Filter by category '{category}': {len(filtered)} results") + + logger.info("✓ Metadata statistics test passed") + + +async def test_update_metadata_only(store: BaseVectorStore, _store_name: str): + """Test updating only metadata without changing content.""" + logger.info("=" * 20 + " UPDATE METADATA ONLY TEST " + "=" * 20) + + # Get an existing node + original = await store.get("test_node1") + original_content = original.content + + # Update with same content but different metadata + updated_node = VectorNode( + vector_id="test_node1", + content=original_content, # Keep same content + metadata={ + **original.metadata, + "metadata_updated": "true", + "update_count": "1", + "last_modified": "2024-12-31", + }, + ) + + await store.update(updated_node) + logger.info("✓ Updated metadata without changing content") + + # Verify update + result = await store.get("test_node1") + assert result.content == original_content, "Content should remain unchanged" + assert result.metadata.get("metadata_updated") == "true", "Metadata should be updated" + logger.info("✓ Verified metadata-only update") + + # Update metadata again + updated_node_2 = VectorNode( + vector_id="test_node1", + content=original_content, + metadata={ + **result.metadata, + "update_count": "2", + "last_modified": "2024-12-31T12:00:00", + }, + ) + await store.update(updated_node_2) + + result_2 = await store.get("test_node1") + assert result_2.metadata.get("update_count") == "2", "Metadata should be updated again" + logger.info("✓ Multiple metadata updates successful") + + logger.info("✓ Update metadata only test passed") + + +async def test_filter_combinations(store: BaseVectorStore, _store_name: str): + """Test various filter combinations and edge cases.""" + logger.info("=" * 20 + " FILTER COMBINATIONS TEST " + "=" * 20) + + # Test 1: Empty filter (should return all results) + results_1 = await store.search(query="technology", filters={}, limit=10) + logger.info(f"Test 1 - Empty filter: {len(results_1)} results") + + # Test 2: Single value filter + results_2 = await store.search( + query="technology", + filters={"node_type": "tech"}, + limit=10, + ) + logger.info(f"Test 2 - Single value filter: {len(results_2)} results") + for r in results_2: + assert r.metadata.get("node_type") == "tech" + + # Test 3: List filter with single item + results_3 = await store.search( + query="technology", + filters={"node_type": ["tech"]}, + limit=10, + ) + logger.info(f"Test 3 - List filter (single item): {len(results_3)} results") + + # Test 4: List filter with multiple items + results_4 = await store.search( + query="technology", + filters={"category": ["AI", "ML", "DL"]}, + limit=10, + ) + logger.info(f"Test 4 - List filter (multiple items): {len(results_4)} results") + for r in results_4: + assert r.metadata.get("category") in ["AI", "ML", "DL"] + + # Test 5: Multiple filters (AND operation) + results_5 = await store.search( + query="technology", + filters={ + "node_type": ["tech", "tech_new"], + "source": "research", + "priority": "high", + }, + limit=10, + ) + logger.info(f"Test 5 - Multiple filters (AND): {len(results_5)} results") + for r in results_5: + assert r.metadata.get("node_type") in ["tech", "tech_new"] + assert r.metadata.get("source") == "research" + assert r.metadata.get("priority") == "high" + + # Test 6: Filter with non-existent value + results_6 = await store.search( + query="technology", + filters={"category": "NON_EXISTENT_CATEGORY"}, + limit=10, + ) + logger.info(f"Test 6 - Non-existent filter value: {len(results_6)} results") + assert len(results_6) == 0, "Should return no results for non-existent filter value" + + # Test 7: List operation with filters + list_results = await store.list( + filters={"node_type": "tech", "priority": "high"}, + limit=20, + ) + logger.info(f"Test 7 - List with filters: {len(list_results)} results") + for r in list_results: + assert r.metadata.get("node_type") == "tech" + assert r.metadata.get("priority") == "high" + + logger.info("✓ Filter combinations test passed") + + +# ==================== Test Runner ==================== + + +async def run_all_tests_for_store(store_type: str, store_name: str): + """Run all tests for a specific vector store type. + + Args: + store_type: Type of vector store ("local" or "es") + store_name: Display name for the vector store + """ + logger.info(f"\n\n{'#' * 60}") + logger.info(f"# Running all tests for: {store_name}") + logger.info(f"{'#' * 60}") + + config = TestConfig() + collection_name = f"{config.TEST_COLLECTION_PREFIX}_{store_type}_main" + + # Create vector store instance + store = create_vector_store(store_type, collection_name) + + try: + # Run cosine similarity test first (only for LocalVectorStore) + await test_cosine_similarity(store_name) + + # ========== Basic Tests ========== + logger.info(f"\n{'#' * 60}") + logger.info("# BASIC FUNCTIONALITY TESTS") + logger.info(f"{'#' * 60}") + + await test_create_collection(store, store_name) + await test_insert(store, store_name) + await test_search(store, store_name) + await test_search_with_single_filter(store, store_name) + await test_search_with_list_filter(store, store_name) + await test_search_with_multiple_filters(store, store_name) + await test_get_by_id(store, store_name) + await test_list_all(store, store_name) + await test_list_with_filters(store, store_name) + await test_update(store, store_name) + await test_delete(store, store_name) + + # ========== Advanced Tests ========== + logger.info(f"\n{'#' * 60}") + logger.info("# ADVANCED FUNCTIONALITY TESTS") + logger.info(f"{'#' * 60}") + + await test_batch_operations(store, store_name) + await test_complex_metadata_queries(store, store_name) + await test_edge_cases(store, store_name) + await test_concurrent_operations(store, store_name) + await test_search_relevance_ranking(store, store_name) + await test_metadata_statistics(store, store_name) + await test_update_metadata_only(store, store_name) + await test_filter_combinations(store, store_name) + + # ========== Collection Management Tests ========== + logger.info(f"\n{'#' * 60}") + logger.info("# COLLECTION MANAGEMENT TESTS") + logger.info(f"{'#' * 60}") + + await test_list_collections(store, store_name) + await test_copy_collection(store, store_name) + await test_delete_collection(store, store_name) + + logger.info(f"\n{'=' * 60}") + logger.info(f"✓ All tests passed for {store_name}!") + logger.info(f"{'=' * 60}") + + except Exception as e: + logger.error(f"Test failed: {e}") + raise + finally: + # Cleanup + await cleanup_store(store, store_type) + + +async def cleanup_store(store: BaseVectorStore, store_type: str): + """Clean up test resources for a vector store. + + Args: + store: Vector store instance + store_type: Type of vector store ("local" or "es") + """ + logger.info("=" * 20 + " CLEANUP " + "=" * 20) + + try: + # Clean up test collections + config = TestConfig() + collections = await store.list_collections() + test_collections = [c for c in collections if c.startswith(config.TEST_COLLECTION_PREFIX)] + + for collection in test_collections: + try: + await store.delete_collection(collection) + logger.info(f"Deleted test collection: {collection}") + except Exception as e: + logger.warning(f"Failed to delete collection {collection}: {e}") + + # Close connections + await store.close() + + # Clean up local directory if LocalVectorStore + if store_type == "local": + test_dir = Path(config.LOCAL_ROOT_PATH) + if test_dir.exists(): + shutil.rmtree(test_dir) + logger.info(f"Cleaned up local directory: {config.LOCAL_ROOT_PATH}") + + # Clean up local directory if ChromaVectorStore + if store_type == "chroma" and config.CHROMA_PATH: + test_dir = Path(config.CHROMA_PATH) + if test_dir.exists(): + shutil.rmtree(test_dir) + logger.info(f"Cleaned up chroma directory: {config.CHROMA_PATH}") + + logger.info("✓ Cleanup completed") + except Exception as e: + logger.error(f"Cleanup error: {e}") + + +# ==================== Main Entry Point ==================== + + +async def main(): + """Main entry point for running tests.""" + parser = argparse.ArgumentParser( + description="Run vector store tests", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + python test_vector_store.py --local # Test LocalVectorStore only + python test_vector_store.py --es # Test ESVectorStore only + python test_vector_store.py --pgvector # Test PGVectorStore only + python test_vector_store.py --qdrant # Test QdrantVectorStore only + python test_vector_store.py --chroma # Test ChromaVectorStore only + python test_vector_store.py --all # Test all vector stores + """, + ) + parser.add_argument( + "--local", + action="store_true", + help="Test LocalVectorStore", + ) + parser.add_argument( + "--es", + action="store_true", + help="Test ESVectorStore", + ) + parser.add_argument( + "--qdrant", + action="store_true", + help="Test QdrantVectorStore", + ) + parser.add_argument( + "--pgvector", + action="store_true", + help="Test PGVectorStore", + ) + parser.add_argument( + "--chroma", + action="store_true", + help="Test ChromaVectorStore", + ) + parser.add_argument( + "--all", + action="store_true", + help="Run tests for all available vector stores", + ) + + args = parser.parse_args() + + # Determine which vector stores to test + stores_to_test = [] + + if args.all: + stores_to_test = [ + ("local", "LocalVectorStore"), + ("es", "ESVectorStore"), + ("pgvector", "PGVectorStore"), + ("qdrant", "QdrantVectorStore"), + ("chroma", "ChromaVectorStore"), + ] + else: + # Build list based on individual flags + if args.local: + stores_to_test.append(("local", "LocalVectorStore")) + if args.es: + stores_to_test.append(("es", "ESVectorStore")) + if args.pgvector: + stores_to_test.append(("pgvector", "PGVectorStore")) + if args.qdrant: + stores_to_test.append(("qdrant", "QdrantVectorStore")) + if args.chroma: + stores_to_test.append(("chroma", "ChromaVectorStore")) + + if not stores_to_test: + # Default to all vector stores if no argument provided + stores_to_test = [ + ("local", "LocalVectorStore"), + ("es", "ESVectorStore"), + ("pgvector", "PGVectorStore"), + ("qdrant", "QdrantVectorStore"), + ("chroma", "ChromaVectorStore"), + ] + print("No vector store specified, defaulting to test all vector stores") + print( + "Use --local/--es/--pgvector/--qdrant/--chroma to test specific ones\n", + ) + + # Run tests for each vector store + for store_type, store_name in stores_to_test: + try: + await run_all_tests_for_store(store_type, store_name) + except Exception as e: + logger.error(f"\n✗ FAILED: {store_name} tests failed with error:") + logger.error(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(stores_to_test)} vector store(s):") + for _, store_name in stores_to_test: + print(f" - {store_name}") + print(f"{'#' * 60}\n") + + +if __name__ == "__main__": + asyncio.run(main()) From 90a53737c85784766b61d8d501db88662a2b8a0a Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Wed, 31 Dec 2025 17:16:53 +0800 Subject: [PATCH 04/11] feat(core): add operator framework and utility modules --- reme_ai/core/context/runtime_context.py | 86 +++-- reme_ai/core/enumeration/json_schema_enum.py | 17 +- reme_ai/core/op/__init__.py | 11 + reme_ai/core/op/base_op.py | 336 +++++++++++++++++++ reme_ai/core/op/parallel_op.py | 33 ++ reme_ai/core/op/sequential_op.py | 31 ++ reme_ai/core/schema/request.py | 6 - reme_ai/core/schema/tool_call.py | 4 +- reme_ai/core/utils/__init__.py | 8 + reme_ai/core/utils/cache_handler.py | 184 ++++++++++ reme_ai/core/utils/http_client.py | 91 +++++ reme_ai/core/utils/mcp_client.py | 107 ++++++ reme_ai/core/utils/pydantic_utils.py | 66 ++++ tests/mcp_servers_demo.json | 37 ++ tests/test_cache_handler.py | 94 ++++++ tests/test_mcp_client.py | 31 ++ tests/test_mcp_server.py | 125 +++++++ 17 files changed, 1220 insertions(+), 47 deletions(-) create mode 100644 reme_ai/core/op/__init__.py create mode 100644 reme_ai/core/op/base_op.py create mode 100644 reme_ai/core/op/parallel_op.py create mode 100644 reme_ai/core/op/sequential_op.py create mode 100644 reme_ai/core/utils/cache_handler.py create mode 100644 reme_ai/core/utils/http_client.py create mode 100644 reme_ai/core/utils/mcp_client.py create mode 100644 reme_ai/core/utils/pydantic_utils.py create mode 100644 tests/mcp_servers_demo.json create mode 100644 tests/test_cache_handler.py create mode 100644 tests/test_mcp_client.py create mode 100644 tests/test_mcp_server.py diff --git a/reme_ai/core/context/runtime_context.py b/reme_ai/core/context/runtime_context.py index ded35ded..b2362161 100644 --- a/reme_ai/core/context/runtime_context.py +++ b/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,67 @@ 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 + new_instance = cls(response=context.response, stream_queue=context.stream_queue) + new_instance.update(context) + if kwargs: + new_instance.update(kwargs) + return new_instance - # 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 diff --git a/reme_ai/core/enumeration/json_schema_enum.py b/reme_ai/core/enumeration/json_schema_enum.py index 17b59380..507645f4 100644 --- a/reme_ai/core/enumeration/json_schema_enum.py +++ b/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() diff --git a/reme_ai/core/op/__init__.py b/reme_ai/core/op/__init__.py new file mode 100644 index 00000000..a8ba9a12 --- /dev/null +++ b/reme_ai/core/op/__init__.py @@ -0,0 +1,11 @@ +"""op""" + +from .base_op import BaseOp +from .parallel_op import ParallelOp +from .sequential_op import SequentialOp + +__all__ = [ + "BaseOp", + "ParallelOp", + "SequentialOp", +] diff --git a/reme_ai/core/op/base_op.py b/reme_ai/core/op/base_op.py new file mode 100644 index 00000000..7a131a28 --- /dev/null +++ b/reme_ai/core/op/base_op.py @@ -0,0 +1,336 @@ +"""Base operator class for LLM workflow execution and composition.""" + +import asyncio +import copy +import inspect +from pathlib import Path +from typing import Callable, Any, Union + +from loguru import logger +from tqdm import tqdm + +from ..context import RuntimeContext, PromptHandler, C, BaseContext +from ..embedding import BaseEmbeddingModel +from ..llm import BaseLLM +from ..schema import ToolCall, ToolAttr +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: Union[list["BaseOp"], dict[str, "BaseOp"], "BaseOp", None] = None, + input_mapping: dict[str, str] | None = None, + output_mapping: dict[str, str] | None = None, + enable_tool_response: 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 = BaseContext[str, BaseOp]() + self.add_sub_ops(sub_ops) + + self.input_mapping = input_mapping + self.output_mapping = output_mapping + self.enable_tool_response = enable_tool_response + 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: + 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.""" + logger.exception(f"{self.name} failed (attempt {attempt + 1}): {e}") + if attempt == self.max_retries - 1: + if self.raise_exception: + raise e + + if self.tool_call: + self.output = f"{self.name} failed: {e}" + + @property + def tool_call(self) -> ToolCall: + """Lazily construct and return the tool call metadata.""" + if self._tool_call is None: + self._tool_call = self._build_tool_call() + assert self._tool_call, "tool_call is not defined!" + self._tool_call.name = self._tool_call.name or self.name + if not self._tool_call.output.properties: + self._tool_call.output = ToolAttr( + type="object", + properties={ + f"{self.name}_result": ToolAttr(type="string", description=f"Execution result of {self.name}"), + }, + ) + return self._tool_call + + @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: + """Lazily initialize and return the LLM instance.""" + if isinstance(self._llm, str): + cfg = C.service_config.llm[self._llm] + self._llm = C.get_llm_class(cfg.backend)(model_name=cfg.model_name, **cfg.model_extra) + return self._llm + + @property + def embedding_model(self) -> BaseEmbeddingModel: + """Lazily initialize and return the embedding model instance.""" + if isinstance(self._embedding_model, str): + cfg = C.service_config.embedding_model[self._embedding_model] + self._embedding_model = C.get_embedding_model_class(cfg.backend)( + model_name=cfg.model_name, + **cfg.model_extra, + ) + 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: + """Lazily initialize and return the token counter instance.""" + if isinstance(self._token_counter, str): + cfg = C.service_config.token_counter[self._token_counter] + self._token_counter = C.get_token_counter_class(cfg.backend)( + model_name=cfg.model_name, + **cfg.model_extra, + ) + return self._token_counter + + async def before_execute(self): + """Prepare context and validate before async execution.""" + self.context.apply_mapping(self.input_mapping) + self._validate_inputs() + + async def execute(self): + """Define core async logic in subclasses.""" + + async def after_execute(self): + """Finalize context and mappings after async execution.""" + self.context.apply_mapping(self.output_mapping) + if self.tool_call and self.enable_tool_response: + self.context.response.answer = self.output + + if not isinstance(self._llm, str) and hasattr(self._llm, "close"): + await self._llm.close() + if not isinstance(self._embedding_model, str) and hasattr(self._embedding_model, "close"): + await self._embedding_model.close() + + 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 and self.enable_tool_response: + self.context.response.answer = self.output + + if not isinstance(self._llm, str) and hasattr(self._llm, "close_sync"): + self._llm.close_sync() + if not isinstance(self._embedding_model, str) and hasattr(self._embedding_model, "close_sync"): + self._embedding_model.close_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 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 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: Union[list["BaseOp"], dict[str, "BaseOp"], "BaseOp", None]): + """Add child operators to this operator's sub_ops context.""" + if not sub_ops: + return + + if isinstance(sub_ops, dict): + ops_dict = sub_ops + else: + ops_dict = {op.name: op for op in (sub_ops if isinstance(sub_ops, list) else [sub_ops])} + + for name, op in ops_dict.items(): + assert self.async_mode == op.async_mode, "Async mode mismatch!" + self.sub_ops[name] = op + + def add_sub_op(self, sub_op: "BaseOp"): + """Add a single child operator to this operator's sub_ops context.""" + self.add_sub_ops(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 diff --git a/reme_ai/core/op/parallel_op.py b/reme_ai/core/op/parallel_op.py new file mode 100644 index 00000000..746485ca --- /dev/null +++ b/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.values(): + 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.values(): + 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 diff --git a/reme_ai/core/op/sequential_op.py b/reme_ai/core/op/sequential_op.py new file mode 100644 index 00000000..79243cff --- /dev/null +++ b/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.values(): + 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.values(): + 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 diff --git a/reme_ai/core/schema/request.py b/reme_ai/core/schema/request.py index aa054219..ece942b9 100644 --- a/reme_ai/core/schema/request.py +++ b/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) diff --git a/reme_ai/core/schema/tool_call.py b/reme_ai/core/schema/tool_call.py index a94dfb1d..0a96f8c4 100644 --- a/reme_ai/core/schema/tool_call.py +++ b/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}") diff --git a/reme_ai/core/utils/__init__.py b/reme_ai/core/utils/__init__.py index ed1916b2..6b217d21 100644 --- a/reme_ai/core/utils/__init__.py +++ b/reme_ai/core/utils/__init__.py @@ -1,14 +1,22 @@ """utils""" +from .cache_handler import CacheHandler from .case_converter import snake_to_camel, camel_to_snake from .env_utils import load_env +from .http_client import HttpClient +from .mcp_client import MCPClient +from .pydantic_utils import create_pydantic_model from .singleton import singleton from .timer import timer __all__ = [ + "CacheHandler", "snake_to_camel", "camel_to_snake", "load_env", + "HttpClient", + "MCPClient", + "create_pydantic_model", "singleton", "timer", ] diff --git a/reme_ai/core/utils/cache_handler.py b/reme_ai/core/utils/cache_handler.py new file mode 100644 index 00000000..70c8585a --- /dev/null +++ b/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), + } diff --git a/reme_ai/core/utils/http_client.py b/reme_ai/core/utils/http_client.py new file mode 100644 index 00000000..8c8e92c3 --- /dev/null +++ b/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 diff --git a/reme_ai/core/utils/mcp_client.py b/reme_ai/core/utils/mcp_client.py new file mode 100644 index 00000000..3a1c70e7 --- /dev/null +++ b/reme_ai/core/utils/mcp_client.py @@ -0,0 +1,107 @@ +"""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 + +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]) -> CallToolResult: + """Execute a tool on a specific server.""" + async with self.connect_to_server(server_name) as session: + return await session.call_tool(tool_name, arguments) diff --git a/reme_ai/core/utils/pydantic_utils.py b/reme_ai/core/utils/pydantic_utils.py new file mode 100644 index 00000000..06b1f4fe --- /dev/null +++ b/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) diff --git a/tests/mcp_servers_demo.json b/tests/mcp_servers_demo.json new file mode 100644 index 00000000..e8f9a858 --- /dev/null +++ b/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 + } + } +} \ No newline at end of file diff --git a/tests/test_cache_handler.py b/tests/test_cache_handler.py new file mode 100644 index 00000000..ddcac86f --- /dev/null +++ b/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() diff --git a/tests/test_mcp_client.py b/tests/test_mcp_client.py new file mode 100644 index 00000000..7db2ad61 --- /dev/null +++ b/tests/test_mcp_client.py @@ -0,0 +1,31 @@ +"""Test module for demonstrating MCPClient functionality.""" + +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: + t_list = await client.list_tool_calls(test_mcp) + for t in t_list: + print(json.dumps(t, ensure_ascii=False, indent=2)) + except Exception as e: + print(f"Error occurred: {e}") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py new file mode 100644 index 00000000..67f9c542 --- /dev/null +++ b/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") From 1e7b8fbdadc7b35b7e5e093d9d36c83f368bc7a4 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Wed, 31 Dec 2025 17:36:43 +0800 Subject: [PATCH 05/11] feat(core): add service metadata and prompt formatting capabilities to base operator --- reme_ai/core/op/base_op.py | 21 +++++++ reme_ai/core/utils/mcp_client.py | 20 ++++++- tests/test_mcp_client.py | 96 ++++++++++++++++++++++++++++++++ 3 files changed, 134 insertions(+), 3 deletions(-) diff --git a/reme_ai/core/op/base_op.py b/reme_ai/core/op/base_op.py index 7a131a28..e5f8450e 100644 --- a/reme_ai/core/op/base_op.py +++ b/reme_ai/core/op/base_op.py @@ -194,6 +194,11 @@ class BaseOp: ) return self._token_counter + @property + def service_metadata(self) -> dict: + """Get service configuration metadata.""" + return C.service_config.model_extra + async def before_execute(self): """Prepare context and validate before async execution.""" self.context.apply_mapping(self.input_mapping) @@ -334,3 +339,19 @@ class BaseOp: 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 diff --git a/reme_ai/core/utils/mcp_client.py b/reme_ai/core/utils/mcp_client.py index 3a1c70e7..4c13be9f 100644 --- a/reme_ai/core/utils/mcp_client.py +++ b/reme_ai/core/utils/mcp_client.py @@ -9,7 +9,7 @@ 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 +from mcp.types import CallToolResult, TextContent from ..schema import ToolCall @@ -101,7 +101,21 @@ class MCPClient: return tool_calls - async def call_tool(self, server_name: str, tool_name: str, arguments: dict[str, Any]) -> CallToolResult: + 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: - return await session.call_tool(tool_name, arguments) + 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) diff --git a/tests/test_mcp_client.py b/tests/test_mcp_client.py index 7db2ad61..d2fef40e 100644 --- a/tests/test_mcp_client.py +++ b/tests/test_mcp_client.py @@ -1,5 +1,7 @@ """Test module for demonstrating MCPClient functionality.""" +# pylint: disable=too-many-return-statements,too-many-statements + import asyncio import json @@ -20,11 +22,105 @@ async def main(): 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__": From 472e069bc51428b452eae6aabd18820396b319ce Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Wed, 31 Dec 2025 23:43:56 +0800 Subject: [PATCH 06/11] feat(core): add MCP tool integration and Ray-based parallel operations --- reme_ai/core/context/service_context.py | 2 +- reme_ai/core/op/base_ray_op.py | 124 ++++++++++++++++++++++++ reme_ai/core/schema/service_config.py | 5 +- reme_ai/core/tool/__init__.py | 7 ++ reme_ai/core/tool/mcp_tool.py | 82 ++++++++++++++++ 5 files changed, 215 insertions(+), 5 deletions(-) create mode 100644 reme_ai/core/op/base_ray_op.py create mode 100644 reme_ai/core/tool/__init__.py create mode 100644 reme_ai/core/tool/mcp_tool.py diff --git a/reme_ai/core/context/service_context.py b/reme_ai/core/context/service_context.py index 65702be1..97426594 100644 --- a/reme_ai/core/context/service_context.py +++ b/reme_ai/core/context/service_context.py @@ -21,7 +21,7 @@ class ServiceContext(BaseContext): self.language: str = "" self.thread_pool: ThreadPoolExecutor | None = None self.vector_store_dict: dict[str, dict] = {} - self.external_mcp_tool_call_dict: dict = {} + self.mcp_server_tool_call_mapping: 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 = {} diff --git a/reme_ai/core/op/base_ray_op.py b/reme_ai/core/op/base_ray_op.py new file mode 100644 index 00000000..88cc2985 --- /dev/null +++ b/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 diff --git a/reme_ai/core/schema/service_config.py b/reme_ai/core/schema/service_config.py index 67cac68f..023eb3a7 100644 --- a/reme_ai/core/schema/service_config.py +++ b/reme_ai/core/schema/service_config.py @@ -98,10 +98,7 @@ class ServiceConfig(BaseModel): 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) diff --git a/reme_ai/core/tool/__init__.py b/reme_ai/core/tool/__init__.py new file mode 100644 index 00000000..02dbd468 --- /dev/null +++ b/reme_ai/core/tool/__init__.py @@ -0,0 +1,7 @@ +"""tool""" + +from .mcp_tool import MCPTool + +__all__ = [ + "MCPTool", +] diff --git a/reme_ai/core/tool/mcp_tool.py b/reme_ai/core/tool/mcp_tool.py new file mode 100644 index 00000000..f9afa029 --- /dev/null +++ b/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 = "", + enable_tool_response: 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__( + enable_tool_response=enable_tool_response, + 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_tool_call_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, + ) From 566a773591b2b20f386a252723e90336dd04fb58 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Thu, 1 Jan 2026 13:10:21 +0800 Subject: [PATCH 07/11] refactor(core): update BaseOp to improve sub-ops handling and tool call management --- reme_ai/core/context/runtime_context.py | 7 +- reme_ai/core/op/base_op.py | 99 +++++--- reme_ai/core/op/parallel_op.py | 4 +- reme_ai/core/op/sequential_op.py | 4 +- reme_ai/core/tool/mcp_tool.py | 4 +- tests/test_op_composition.py | 325 ++++++++++++++++++++++++ 6 files changed, 395 insertions(+), 48 deletions(-) create mode 100644 tests/test_op_composition.py diff --git a/reme_ai/core/context/runtime_context.py b/reme_ai/core/context/runtime_context.py index b2362161..d7112e1c 100644 --- a/reme_ai/core/context/runtime_context.py +++ b/reme_ai/core/context/runtime_context.py @@ -27,11 +27,8 @@ class RuntimeContext(BaseContext): if context is None: return cls(**kwargs) - new_instance = cls(response=context.response, stream_queue=context.stream_queue) - new_instance.update(context) - if kwargs: - new_instance.update(kwargs) - return new_instance + context.update(kwargs) + return context async def _enqueue(self, chunk: StreamChunk) -> None: """Internal helper to put a chunk into the queue if it exists.""" diff --git a/reme_ai/core/op/base_op.py b/reme_ai/core/op/base_op.py index e5f8450e..26f1e7f5 100644 --- a/reme_ai/core/op/base_op.py +++ b/reme_ai/core/op/base_op.py @@ -4,15 +4,15 @@ import asyncio import copy import inspect from pathlib import Path -from typing import Callable, Any, Union +from typing import Callable, Any, Optional from loguru import logger from tqdm import tqdm -from ..context import RuntimeContext, PromptHandler, C, BaseContext +from ..context import RuntimeContext, PromptHandler, C from ..embedding import BaseEmbeddingModel from ..llm import BaseLLM -from ..schema import ToolCall, ToolAttr +from ..schema import ToolCall, ToolAttr, Response from ..token_counter import BaseTokenCounter from ..utils import camel_to_snake, CacheHandler, timer from ..vector_store import BaseVectorStore @@ -40,10 +40,10 @@ class BaseOp: token_counter: str | BaseTokenCounter = "default", enable_cache: bool = False, cache_path: str = "cache/op", - sub_ops: Union[list["BaseOp"], dict[str, "BaseOp"], "BaseOp", None] = None, + sub_ops: dict[str, "BaseOp"] | list["BaseOp"] | Optional["BaseOp"] = None, input_mapping: dict[str, str] | None = None, output_mapping: dict[str, str] | None = None, - enable_tool_response: bool = False, + save_response_result: bool = False, enable_sync_thread_pool: bool = True, max_retries: int = 1, raise_exception: bool = False, @@ -62,12 +62,12 @@ class BaseOp: self.enable_cache = enable_cache self.cache_path = cache_path - self.sub_ops = BaseContext[str, BaseOp]() + self.sub_ops: list[BaseOp] = [] self.add_sub_ops(sub_ops) self.input_mapping = input_mapping self.output_mapping = output_mapping - self.enable_tool_response = enable_tool_response + 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 @@ -89,7 +89,7 @@ class BaseOp: def _validate_inputs(self): """Ensure all required tool inputs are present in context.""" - if self.tool_call: + if self.tool_call is not None: parameters = self.tool_call.parameters if parameters.type == "object" and parameters.properties: required_list = parameters.required or [] @@ -98,30 +98,47 @@ class BaseOp: def _handle_failure(self, e: Exception, attempt: int): """Log failures and handle final retry logic.""" - logger.exception(f"{self.name} failed (attempt {attempt + 1}): {e}") + 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: + if self.tool_call is not None: self.output = f"{self.name} failed: {e}" + else: + logger.warning(message) @property - def tool_call(self) -> ToolCall: + 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() - assert self._tool_call, "tool_call is not defined!" + 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 = ToolAttr( - type="object", - properties={ - f"{self.name}_result": ToolAttr(type="string", description=f"Execution result of {self.name}"), - }, - ) + 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.""" @@ -137,6 +154,7 @@ class BaseOp: output_properties = self.tool_call.output.properties if not output_properties: return None + keys = list(output_properties.keys()) return self.context[keys[0]] @@ -146,6 +164,7 @@ class BaseOp: output_properties = self.tool_call.output.properties if not output_properties: return + keys = list(output_properties.keys()) self.context[keys[0]] = value @@ -188,10 +207,7 @@ class BaseOp: """Lazily initialize and return the token counter instance.""" if isinstance(self._token_counter, str): cfg = C.service_config.token_counter[self._token_counter] - self._token_counter = C.get_token_counter_class(cfg.backend)( - model_name=cfg.model_name, - **cfg.model_extra, - ) + self._token_counter = C.get_token_counter_class(cfg.backend)(model_name=cfg.model_name, **cfg.model_extra) return self._token_counter @property @@ -199,6 +215,11 @@ class BaseOp: """Get service configuration metadata.""" return C.service_config.model_extra + @property + def response(self) -> Response: + """Get the response object.""" + return self.context.response + async def before_execute(self): """Prepare context and validate before async execution.""" self.context.apply_mapping(self.input_mapping) @@ -210,7 +231,7 @@ class BaseOp: async def after_execute(self): """Finalize context and mappings after async execution.""" self.context.apply_mapping(self.output_mapping) - if self.tool_call and self.enable_tool_response: + if self.tool_call is not None and self.save_response_result: self.context.response.answer = self.output if not isinstance(self._llm, str) and hasattr(self._llm, "close"): @@ -229,7 +250,7 @@ class BaseOp: def after_execute_sync(self): """Finalize context and mappings after sync execution.""" self.context.apply_mapping(self.output_mapping) - if self.tool_call and self.enable_tool_response: + if self.tool_call is not None and self.save_response_result: self.context.response.answer = self.output if not isinstance(self._llm, str) and hasattr(self._llm, "close_sync"): @@ -249,7 +270,7 @@ class BaseOp: break except Exception as e: self._handle_failure(e, i) - return self.output if self.tool_call else None + 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.""" @@ -262,7 +283,7 @@ class BaseOp: break except Exception as e: self._handle_failure(e, i) - return self.output if self.tool_call else None + 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.""" @@ -301,23 +322,27 @@ class BaseOp: finally: self._pending_tasks.clear() - def add_sub_ops(self, sub_ops: Union[list["BaseOp"], dict[str, "BaseOp"], "BaseOp", None]): - """Add child operators to this operator's sub_ops context.""" + 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): - ops_dict = sub_ops + 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: - ops_dict = {op.name: op for op in (sub_ops if isinstance(sub_ops, list) else [sub_ops])} - - for name, op in ops_dict.items(): - assert self.async_mode == op.async_mode, "Async mode mismatch!" - self.sub_ops[name] = op + 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 context.""" - self.add_sub_ops(sub_op) + """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.""" @@ -350,7 +375,7 @@ class BaseOp: 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) + 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) diff --git a/reme_ai/core/op/parallel_op.py b/reme_ai/core/op/parallel_op.py index 746485ca..18b84d0c 100644 --- a/reme_ai/core/op/parallel_op.py +++ b/reme_ai/core/op/parallel_op.py @@ -8,14 +8,14 @@ class ParallelOp(BaseOp): async def execute(self): """Executes all sub-operations concurrently using asynchronous tasks.""" - for op in self.sub_ops.values(): + 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.values(): + 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() diff --git a/reme_ai/core/op/sequential_op.py b/reme_ai/core/op/sequential_op.py index 79243cff..3dabb0c9 100644 --- a/reme_ai/core/op/sequential_op.py +++ b/reme_ai/core/op/sequential_op.py @@ -8,13 +8,13 @@ class SequentialOp(BaseOp): async def execute(self): """Executes sub-operations sequentially using asynchronous awaits.""" - for op in self.sub_ops.values(): + 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.values(): + for op in self.sub_ops: assert not op.async_mode op.call_sync(context=self.context) diff --git a/reme_ai/core/tool/mcp_tool.py b/reme_ai/core/tool/mcp_tool.py index f9afa029..90f6ad7b 100644 --- a/reme_ai/core/tool/mcp_tool.py +++ b/reme_ai/core/tool/mcp_tool.py @@ -20,7 +20,7 @@ class MCPTool(BaseOp): self, mcp_server: str = "", tool_name: str = "", - enable_tool_response: bool = True, + save_response_result: bool = True, parameter_required: List[str] | None = None, parameter_optional: List[str] | None = None, parameter_deleted: List[str] | None = None, @@ -31,7 +31,7 @@ class MCPTool(BaseOp): ): super().__init__( - enable_tool_response=enable_tool_response, + save_response_result=save_response_result, max_retries=max_retries, raise_exception=raise_exception, **kwargs, diff --git a/tests/test_op_composition.py b/tests/test_op_composition.py new file mode 100644 index 00000000..8d32b51a --- /dev/null +++ b/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) From 1ea8f07aac3076c6bf2743199ff3164ed83bf75e Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Fri, 2 Jan 2026 00:16:24 +0800 Subject: [PATCH 08/11] feat(core): add flow module and enhance documentation --- docs/reme_v2_design.md | 929 ++++++++++++++++++--------- reme_ai/core/flow/__init__.py | 11 + reme_ai/core/flow/base_flow.py | 214 ++++++ reme_ai/core/flow/cmd_flow.py | 18 + reme_ai/core/flow/expression_flow.py | 30 + reme_ai/core/op/__init__.py | 2 + 6 files changed, 918 insertions(+), 286 deletions(-) create mode 100644 reme_ai/core/flow/__init__.py create mode 100644 reme_ai/core/flow/base_flow.py create mode 100644 reme_ai/core/flow/cmd_flow.py create mode 100644 reme_ai/core/flow/expression_flow.py diff --git a/docs/reme_v2_design.md b/docs/reme_v2_design.md index 18d14657..6fa78140 100644 --- a/docs/reme_v2_design.md +++ b/docs/reme_v2_design.md @@ -1,378 +1,735 @@ -# ReMeV2 - agent memory kit design +# ReMeV2 深度设计文档：渐进式 Agentic Memory 方案 -ReMeV2 是一个面向智能体（Agent）的分层记忆系统，用于从连续的对话与交互中，逐步沉淀出稳定、可检索、可更新的长期记忆。 -系统强调从**原始轨迹 → 事件索引 → 高维抽象记忆**的分层存储，并通过**检索器（Retriever）**与**总结器（Summarizer）**实现动态更新。 +## 一、 背景与现状分析 + +### 1.1 当前面临的挑战 + +* **外功修炼（接口易用性）**：现有的 `server-client` 模式对新手开发者不够友好，集成成本高，需要更直观、纯 Pythonic 的调用方式。 +* **内功修炼（架构深度）**：受 `skills` 和 `agentic memory` 启发，现有的存储检索较为机械。我们需要一种基于**渐进式检索（Progressive Retrieval）**与**渐进式总结（Progressive Summarization）**的智能体记忆方案。 + +### 1.2 核心目标 + +1. **极简开发体验**：开发者友好，全异步接口，支持本地直接运行与 CLI 体验。 +2. **认知架构升级**：引入 渐进式检索 & 渐进式总结 的 Agentic 模式，融合多种记忆，让记忆的存取具备“思考”过程。 +3. **生态融合**：原生支持 AgentScope、LangChain 等主流框架。 --- -## 设计目标 +## 二、 竞品调研与启示 -本设计旨在构建一个可嵌入任意智能体框架（如对话 Agent、任务 Agent 等）的记忆子系统，使其具备以下能力： +### 2.1 主流竞品深度对比 + +| 产品 | 设计哲学 | 核心优势 | 局限性 | +|-------------|----------|---------------------------------------------|-------------------| +| **mem0** | 智能便签本 | 原子事实提取，极高 Token 效率。 | 缺乏对复杂逻辑链条的支持。 | +| **Letta** | 带硬盘的 CPU | 模拟计算机三级存储（Core/Recall/Archival），Agent 自主控存。 | 状态机管理相对复杂。 | +| **MIRIX** | 认知架构图谱 | 实体-关系双引擎，支持记忆“进化”与“固化”。 | 侧重研究，落地集成门槛较高。 | +| **LangMem** | 用户档案系统 | 异步 Compaction（压缩），Schema 驱动，强一致性。 | 偏向 SaaS 应用，灵活性略逊。 | + +### 2.2 mem0 +- https://github.com/mem0ai/mem0 +- https://docs.mem0.ai/core-concepts/memory-operations/add +- https://docs.mem0.ai/core-concepts/memory-operations/search +- https://docs.mem0.ai/core-concepts/memory-operations/update +- https://docs.mem0.ai/core-concepts/memory-operations/delete + +#### 2.2.1 API Reference +| 接口名称 | 核心输入参数 (Inputs) | 核心输出 (Outputs) | 背后逻辑 (Internal Logic) | +| --- | --- | --- | --- | +| **Add** | `messages` (文本/对话), `user_id`, `metadata` | `id`, `event` (ADD/UPDATE), `data` | **提取与合并**：LLM 提取事实，自动去重并更新已有记忆，而非简单堆叠。 | +| **Search** | `query` (自然语言), `filters`, `limit` | `id`, `memory` (事实文本), `score`, `metadata` | **语义检索**：基于向量相似度查找最相关的“原子事实”，支持多维过滤。 | +| **Update** | `memory_id` (必填), `data` (新内容) | 操作状态 (Success/Fail) | **手动干预**：允许开发者对特定的事实进行精确修正。 | +| **Delete** | `memory_id` 或 `user_id` (清空) | 操作状态 (Success/Fail) | **遗忘机制**：物理删除或逻辑移除不再需要的信息。 | + +#### 2.2.2 Tech Strategy & Benefits +| 维度 | 技术方案 (Technical Solution) | 核心优势 (Key Advantages) | +| --- | --- | --- | +| **存储架构** | **混合存储**：向量数据库 (Vector) + 图数据库 (Graph) + 关系型元数据。 | **多维关联**：不仅能搜到相似内容，还能理解实体间的逻辑关系（如“父子”、“因果”）。 | +| **数据处理** | **原子化事实提取**：利用 LLM 将长篇对话压缩为简短的 Fact。 | **极高 Token 效率**：注入 Prompt 的内容更精炼，减少 90% 以上的冗余信息，大幅降本。 | +| **管理层级** | **多级联动**：User (长期) Agent (专业) Session (短期)。 | **个性化定制**：实现跨会话的“长效记忆”，AI 能记住用户一个月前说过的偏好。 | +| **冲突处理** | **自适应更新算法**：新信息进入时自动比对旧记忆。 | **数据一致性**：自动处理矛盾信息（如用户更换了住址），确保记忆库始终是“最新真理”。 | +| **兼容性** | **解耦设计**：支持多种 Embedding 模型与向量数据库后端。 | **快速集成**：几行代码即可为现有 LLM 应用增加记忆层，适配各种生产环境。 | -- **支持多天、多会话的长期记忆积累**，不丢失原始细节。 -- **在推理时按需检索**与当前问题高度相关的记忆片段。 -- **在后台持续做分层总结与压缩**，提炼出抽象的“世界观 / 自我认知 / 用户画像 / 任务知识 / 工具经验”。 -- **支持记忆的增删改**，能够对过时或冲突信息进行更新，保持长期记忆的时效性与一致性。 -- **支持渐进式检索**，从高层抽象记忆逐步回溯到中层事件索引与底层原始轨迹，按需展开细节、控制上下文长度与推理成本。 -- **支持渐进式总结**，从最新对话开始按块增量总结并更新索引与主题记忆，而非每次对全量历史做重复总结。 --- -## 架构概览 +### 2.3 Letta +- https://github.com/letta-ai/letta +- https://docs.letta.com/guides/agents/archival-memory/ +- https://docs.letta.com/guides/agents/archival-search/ -系统整体可以分为两个维度来看： +#### 2.3.1 存储架构层级 (Memory Tiering) -- **静态维度**：记忆如何被分层存储（底层消息块 → 中层索引 → 顶层主题记忆）。 -- **动态维度**：记忆如何被检索与更新（上行 Summarizer、下行 Retriever）。 +Letta 将记忆分为三个物理/逻辑层，模拟计算机的存储架构： -核心组件包括： +| 记忆层级 | 存储介质 | 访问方式 | 核心作用 | +| --- | --- | --- | --- | +| **Core Memory** | **上下文窗口 (Prompt)** | 直接读写 | **即时意识**：包含 `Persona`（AI 设定）和 `Human`（用户信息）。Agent 随时可见，响应最快。 | +| **Recall Memory** | **关系型数据库 (SQL)** | 分页检索 | **短期/历史回顾**：存储完整的对话流（Messages）。用于回答“你刚才说了什么”。 | +| **Archival Memory** | **向量数据库 (Vector)** | 语义搜索 | **长期知识库**：存储海量事实或文档。Agent 通过工具自主检索或存入。 | -- **底层记忆层（Raw Dialogue Store）**：多天对话轨迹切分成 `messages block`。 -- **中层索引层（Index & Summary Store）**：为每个 `messages block` 生成 `block_id: time + summary`。 -- **顶层主题记忆层（High-level Memories）**：自我认知、任务记忆、工具记忆、用户画像等高维抽象。 -- **Meta Agent（记忆中枢）**：持有高层记忆视图与索引，负责路由检索与更新请求。 -- **Sub Agents（记忆子智能体）**：面向具体主题（如某个用户、某类任务、某组工具）的记忆管理单元。 -- **Retriever（检索器）**：从高层到低层，渐进式查找相关记忆。 -- **Summarizer（总结器）**：从低层到高层，渐进式抽象和更新长期记忆。 +#### 2.3.2 核心操作接口 (API & Tool Reference) + +在 Letta 中，记忆的操作通常封装为 **Tools**，由 Agent 根据推理需求主动调用。 + +| 接口/工具名称 | 输入参数 (Inputs) | 核心输出 (Outputs) | 背后逻辑 (Internal Logic) | +| --- | --- | --- | --- | +| **`core_memory_update`** | `section`, `new_content` | 更新后的段落内容 | **原子替换**：直接修改 System Prompt 中的特定块（如：更新用户的职业或 AI 的性格偏好）。 | +| **`archival_memory_insert`** | `content` (字符串) | 写入状态/ID | **知识沉淀**：将当前对话中的重要信息或外部文件片段“持久化”到向量数据库。 | +| **`archival_memory_search`** | `query`, `page` | 匹配的文本块列表 | **主动 RAG**：Agent 意识到知识不足时，自主发起向量检索，并将结果拉入临时上下文。 | +| **`conversation_search`** | `query`, `start_date` | 历史消息记录 | **全文检索**：在 Recall Memory 中根据关键词或时间戳查找历史对话详情。 | +| **`send_message`** | `message`, `agent_id` | 响应流/状态更新 | **状态循环**：这是主入口，触发 Agent 的“思考-行动-观察”循环，自动处理内存同步。 | + +#### 2.3.3 技术策略与核心优势 (Tech Strategy & Benefits) + +| 维度 | 技术方案 (Technical Solution) | 核心优势 (Key Advantages) | +| --- | --- | --- | +| **状态持久化** | **Agent State Snapshot**：将 Agent 的所有内存、工具定义和历史记录打包存入数据库。 | **无限存续**：Agent 不再是无状态的 API 调用。重启服务器后，Agent 依然记得所有细节。 | +| **自主演进** | **Self-Editing Loop**：Agent 拥有修改自己 Core Memory 的权限（通过函数调用）。 | **认知闭环**：AI 能在交流中发现矛盾并自我更正，例如发现用户搬家后自动更新 `Human` 模块。 | +| **算力调度** | **OOC (Out-of-Context) 管理**：当对话过长，系统自动将旧消息从 Core 移入 Recall。 | **突破 Context 限制**：在 8k 窗口的模型上也能处理相当于 1M 窗口的逻辑量，且成本更低。 | +| **多代理协同** | **Letta Server 中控**：统一管理多个 Agent 的状态机与资源访问权限。 | **企业级扩展**：支持创建 Agent 团队，每个 Agent 拥有独立的记忆空间但可共享 Archival 库。 | +| **解耦灵活性** | **Provider Agnostic**：后端支持 Postgres/Chroma，前端支持 OpenAI/Anthropic/Local LLMs。 | **无缝迁移**：不绑定特定模型，开发者可以根据成本或能力随时更换底座。 | + +#### 2.3.4 与 mem0 的深度对比 + +* **设计哲学**： +* **mem0** 像是一个**“智能记事本”**，它在后台默默地帮你总结事实。 +* **Letta** 像是一个**“带硬盘的 CPU”**，它把记忆管理完全交给了 Agent 自己的逻辑推理。 + + +* **交互模式**： +* **mem0** 通常是外部干预（Add/Search）。 +* **Letta** 强调 **Agentic Control**（Agent 意识到需要搜索时才去搜索），这种模式更接近人类的思维过程。 --- -## 静态记忆三层架构 +### 2.4 MIRIX +- https://github.com/Mirix-AI/MIRIX +- https://docs.mirix.io/ -### 垂直层级划分 +#### 2.4.1 API Reference -系统将记忆分为自下而上的三个层级，抽象程度逐级升高： +| 接口名称 | 核心输入参数 (Inputs) | 核心输出 (Outputs) | 背后逻辑 (Internal Logic) | +| --- | --- | --- | --- | +| **Add** | `content` (观察/对话), `agent_id`, `context_type` (如任务/闲聊) | `memory_id`, `graph_nodes`, `status` | **实体建模**：不只是提取事实，而是将信息拆解为实体（Entities）与关系（Relations），并挂载到智能体的知识图谱中。 | +| **Query** | `query` (意图), `scope` (全局/局部), `top_k` | `retrieved_memories`, `relation_paths`, `score` | **混合检索**：结合向量（Vector）的语义相关性和图（Graph）的拓扑连接性，寻找具有逻辑深度背景的记忆。 | +| **Evolve** | `target_memories` (可选), `agent_id` | `optimized_structure`, `merged_nodes` | **记忆固化/压缩**：模仿人类大脑的“睡眠”机制，自动合并碎片化记忆，将短期经验转化为长期的结构化知识。 | +| **Observe** | `interaction_stream`, `feedback` | `insights`, `priority_update` | **实时学习**：根据用户反馈或环境变化，动态调整记忆的权重（Importance）和置信度。 | -- **底层：记忆碎片（历史对话 / 轨迹）** - - 存放原始的、未经压缩的对话与交互记录。 - - **规则**：不去重、不解冲突，只做追加与切片。 +#### 2.4.2 Tech Strategy & Benefits -- **中层：记忆索引（事件总结）** - - 对原始消息块做首轮总结，形成“时间 + 简要摘要”的事件索引。 - - 连接底层细节与顶层抽象。 - - **规则**：允许存在部分冗余与矛盾，主要目标是加速检索。 +| 维度 | 技术方案 (Technical Solution) | 核心优势 (Key Advantages) | +| --- | --- | --- | +| **存储架构** | **语义-关系双引擎**：向量索引（Vector Index）+ 属性图（Property Graph）。 | **深度上下文**：不仅知道“是什么”，还能通过图路径推理出“为什么”，有效解决 LLM 幻觉问题。 | +| **记忆层级** | **三层架构**：感知记忆 (Perception) -> 语义记忆 (Semantic) -> 经验记忆 (Episodic)。 | **任务适应性**：不同任务自动匹配不同的记忆深度，短期任务关注细节，长期任务关注模式。 | +| **演化机制** | **自主固化 (Self-Consolidation)**：通过 LLM 定期对冗余、矛盾信息进行清洗和逻辑抽象。 | **永久生命力**：解决随时间推移记忆库膨胀导致的检索噪声，确保记忆库“越用越聪明”。 | +| **推理增强** | **基于记忆的 RAG+**：在检索到的事实基础上，额外提供关联的逻辑链条（Logic Chains）。 | **辅助决策**：为 Agent 提供决策支撑，使其在处理复杂流程时具备类似“长期经验值”的直觉。 | +| **多代理协同** | **内存共享协议**：支持 Agent 之间的记忆交换与知识同步。 | **群体智能**：多个 Agent 可以共享同一套底层知识体系，同时保留各自的私有工作记忆。 | -- **顶层：主题记忆（高维抽象知识）** - - 聚合并抽象出关于“自我、任务、工具、用户”等长期知识。 - - **规则**：需要去重（Deduplication）与解冲突（Resolve Conflict），追求紧凑、一致、可解释。 +#### 2.4.3 与 mem0 的主要区别 -### 底层：原始对话与轨迹（Raw Dialogue） - -- 按时间顺序堆叠的对话与行为轨迹： - - `Day1 Messages → Day2 Messages → ... → DayN Messages` -- 通过时间窗口、业务逻辑、固定长度（可重叠）等策略切分为多个 `messages block`： - - `Day1 Messages Block1` - - `Day1 Messages Block2` - - `DayN Messages BlockN` -- 特性： - - 不丢失任何历史细节。 - - 支持后续回溯、审计与再训练。 - -### 中层：记忆索引与事件总结（Index & Summary） - -- 对每个 `messages block` 生成一个结构化索引： - - 格式：`block_id: time + summary` -- 示例： - - `block_id1: 2025-11-11 + 用户提到他非常喜欢吃西瓜。` - - `block_id2: 2025-11-12 + 用户表达自己很孤独。` -- 作用： - - 作为底层原始消息与顶层抽象记忆之间的“桥梁”与“目录”。 - - 为时间感知检索、语义检索提供最小粒度的事件单元。 - -### 顶层：主题记忆（High-Dimensional Abstraction） - -顶层记忆由多个“记忆气泡”（Memory Bubble）构成，每个气泡代表一类主题知识，并通过关联的 `block_id` 与中层索引相连。主要包括： - -- **自我认知记忆（Self-Cognition Memory）** - - 示例： - - “我叫 Remy” - - “我每天需要上班” - - “我是一名面向长期陪伴的智能助手” - - 位置：处于记忆拓扑中心，与其他记忆类别广泛连接。 - -- **金融任务记忆（Financial Task Memory）** - - 示例： - - “做财报分析前需要调研企业的经营数据” - - “估值时可以优先考虑 forward PE” - - 作用：支撑在特定领域（如金融）的任务推理能力。 - -- **工具记忆（Tool Memory）** - - 示例： - - “`bing_search` 很不稳定，经常失败” - - “`tongyi_search` 的结果质量较好，常需要配合 `web_extract` 使用” - - 作用：让智能体在工具选择、调用顺序和错误预期上具有“经验”。 - -- **用户画像记忆（Per-User Memory，如用户 A / 用户 B）** - - 对用户 A 的记忆示例： - - “A 很喜欢吃西瓜” - - “A 非常有礼貌” - - “A 计划下周出行” - - 与中层事件关联： - - 顶层的“喜欢吃西瓜”记忆与 `block_id1` 建立逻辑连接。 - - 每个用户独立维护一套长期画像，避免混淆。 +* **Mem0** 侧重于**个性化偏好存储**（Personalization），核心是记住“用户喜欢什么”。 +* **MIRIX** 侧重于**智能体认知架构**（Agent Cognition），核心是让 Agent 具备类似人类的“知识归纳”和“逻辑推理”记忆能力。 --- -## 动态流程一：ReMeV2 Retriever（检索流程） +### 2.5 LangMem +- https://github.com/langchain-ai/langmem +- https://langchain-ai.github.io/langmem/ -检索流程是一个**自上而下**的渐进式过程：从高维抽象出发，逐步定位到具体事件甚至原始对话块，为当前问答或决策提供证据。 +#### 2.5.1 API Reference -```mermaid -flowchart TD - op0[op: 当前查询 / 上下文输入] - op1[op: Meta Agent 读取自我认知与主题索引] - op2[op: Meta Agent 基于 topic/intent 选择记忆子空间] - op3[op: 调度对应 Sub Agent 实例] - op4[op: 在顶层/中层记忆空间检索相关 Memories / 索引条目] - op5[op: 根据需要回溯到底层 messages block（retrieve_block）] - op6[op: 汇总得到的 messages block 列表，构造检索上下文] - op7[op: 将上下文交给上层 Agent 做回答 / 决策] +| 接口名称 | 核心输入参数 (Inputs) | 核心输出 (Outputs) | 背后逻辑 (Internal Logic) | +| --- | --- | --- | --- | +| **Add Messages** | `thread_id`, `messages` (List), `user_id` | 操作确认 / 任务 ID | **流式注入**：将原始对话追加到指定的 Thread。LangMem 会自动关联用户上下文，准备进行后续的异步处理。 | +| **Query Memory** | `user_id`, `query` (语义描述), `namespace` | 结构化记忆对象 (JSON / Text) | **多维检索**：不仅支持向量相似度搜索，还能根据定义的 Schema 返回结构化的用户画像或知识状态。 | +| **Trigger Logic** | `thread_id`, `memory_type` | 更新后的 Memory State | **异步固化**：后台启动 LLM 任务，将长篇对话“压缩”并“提取”到长期存储中。支持自定义提取逻辑（如更新用户信息）。 | +| **Manage State** | `user_id`, `patch_data` (增量更新) | 成功/失败 状态 | **精确受控**：开发者可以直接修改持久化的状态（State），支持类似于 Git 的状态管理。 | - op0 --> op1 --> op2 --> op3 --> op4 --> op5 --> op6 --> op7 +#### 2.5.2 Tech Strategy & Benefits + +| 维度 | 技术方案 (Technical Solution) | 核心优势 (Key Advantages) | +| --- | --- | --- | +| **存储架构** | **Stateful Persistence**：基于关系型数据库 (Postgres) + 向量索引。 | **强一致性**：利用数据库事务确保记忆更新的可靠性，支持复杂的结构化查询与过滤。 | +| **数据处理** | **异步化 Compaction (压缩)**：在对话间隙通过后台 Worker 提取知识。 | **无感延迟**：核心对话流程不被记忆提取阻塞，通过定时或事件驱动完成“记忆固化”，优化用户体验。 | +| **管理层级** | **Thread -> User -> Organization**：三层级联记忆。 | **上下文隔离**：完美适配 SaaS 应用场景，既能记住单次对话（Thread），也能沉淀用户习惯（User）。 | +| **逻辑引擎** | **Schema-Driven (模式驱动)**：允许定义 JSON Schema 来规范记忆内容。 | **高度可预测**：输出不再是散乱的句子，而是结构化的字段，方便下游程序直接调用逻辑（如自动填充表单）。 | +| **集成生态** | **LangGraph 原生集成**：作为 Checkpointer 或存储节点直接接入。 | **生态协同**：如果你已经在用 LangChain，LangMem 可以无缝接管状态流转，无需重写底层存储逻辑。 | + +#### 2.5.3 与 mem0 的核心差异 + +* **mem0** 像是一个**“便签本”**：它擅长从每一句话里抠出零散的事实（如“我喜欢吃苹果”），然后把它们存成一条条语义片段。 +* **LangMem** 像是一个**“用户档案系统”**：它更擅长分析一整段对话，然后更新一个复杂的 JSON 档案（如更新用户的偏好模型、性格标签、历史任务状态）。 + +--- + +## 三、 ReMeV2 API 接口设计 + +### 3.1 Long-Term Memory (长期记忆) + +#### 3.1.1 Basic Usage (基础用法) + +The most straightforward way to use ReMe for long-term memory management. Supports basic summary and retrieval operations. + +```python +import os +from reme_ai import ReMe + +os.environ["REME_LLM_API_KEY"] = "sk-..." +os.environ["REME_LLM_BASE_URL"] = "https://dashscope.aliyuncs.com/compatible-mode/v1" +os.environ["REME_EMBEDDING_API_KEY"] = "sk-..." +os.environ["REME_EMBEDDING_BASE_URL"] = "https://dashscope.aliyuncs.com/compatible-mode/v1" + +memory = ReMe( + memory_space="remy", # workspace identifier + llm={"backend": "openai", "model": "qwen-plus", "temperature": 0.6}, + embedding={"backend": "openai", "model": "text-embedding-v4", "dimension": 1024}, + vector_store={"backend": "local_file"}, # supported: local_file, chromadb, qdrant, etc. +) + +# Summarize conversation into memory +result = await memory.summary( + messages=[ + {"role": "user", "content": "I'm travelling to SF"}, + {"role": "assistant", "content": "That's great to hear!"} + ], + user_id="Alice", + # memory_type="auto" # default: auto (auto, personal, procedural, tool) +) + +# Retrieve relevant memories +memories = await memory.retrieve( + query="what is your travel plan?", + limit=3, + user_id="Alice", + # memory_type="auto" # default: auto +) +memories_str = "\n".join(f"- {m['memory']}" for m in memories["results"]) +print(memories_str) ``` -### Meta Agent：记忆入口与路由 +#### 3.1.2 CLI Chat Application (命令行聊天应用) -- 持有： - - **常驻的自我认知记忆**（总是加载在工作内存中）。 - - 各类主题记忆的**名称 + 描述索引**，例如： - - `1. 金融任务记忆` - - `2. 工具记忆` - - `3. 用户A记忆` - - `4. 用户B记忆` -- 职责： - - 根据当前查询的 **topic / intent**，选择合适的记忆子空间。 - - 按需创建或调度具体的 `Sub_Agent` 来完成深入检索。 +A complete example demonstrating how to build a memory-enhanced chatbot with CLI interface. -### Sub Agent：按主题检索 +```python +import os +from reme_ai import ReMe +from openai import OpenAI -- **预定义子智能体**（如：工具记忆 Agent、某个用户画像 Agent 等）： - - 支持接口： - - `retrieve_by_query(query)`：基于语义相似度的检索。 - - `time_aware_retrieve(time)`：基于时间的窗口检索（适合用户时间线类记忆）。 -- **自定义子智能体**： - - 可以实现更复杂的检索逻辑（例如结合时间衰减、多源融合等）。 +os.environ["REME_LLM_API_KEY"] = "sk-..." +os.environ["REME_LLM_BASE_URL"] = "https://dashscope.aliyuncs.com/compatible-mode/v1" +os.environ["REME_EMBEDDING_API_KEY"] = "sk-..." +os.environ["REME_EMBEDDING_BASE_URL"] = "https://dashscope.aliyuncs.com/compatible-mode/v1" -### 渐进式检索与回溯 +memory = ReMe( + memory_space="remy", + llm={"backend": "openai", "model": "qwen-plus", "temperature": 0.6}, + embedding={"backend": "openai", "model": "text-embedding-v4", "dimension": 1024}, + vector_store={"backend": "local_file"}, +) -- 子智能体首先在顶层 / 中层记忆空间中找到语义相关的 **Memories / 索引条目**： - - 每条记忆通常携带 `block_id` 与简要 summary。 -- 若需要更细节的证据，系统可以进一步调用： - - `retrieve_block(block_id)`：回溯到底层 `messages block`。 - - `time_match_memories(time_range)`：根据时间窗口获取原始对话块。 -- 最终输出： - - 一组与当前问题高度相关的 `messages block List`，作为上下文，供 Agent 生成回答或做决策。 +os.environ["OPENAI_API_KEY"] = "sk-..." +os.environ["OPENAI_BASE_URL"] = "https://dashscope.aliyuncs.com/compatible-mode/v1" +openai_client = OpenAI() ---- +def chat_with_memories( + query: str, + history_messages: list[dict], + user_name: str = "", + start_summary_size: int = 2, + keep_size: int = 0 +) -> str: + # Retrieve relevant memories for the query + memories = memory.retrieve(query=query, user_id=user_name, limit=3) -## 动态流程二：ReMeV2 Summarizer（总结流程） + # Build system prompt with memories + system_prompt = ( + "You are a helpful AI named `Remy`. Use the user memories to answer the question. " + "If you don't know the answer, just say you don't know. Don't try to make up an answer.\n" + ) + if memories: + memories_str = "\n".join(f"- {m['memory']}" for m in memories["results"]) + system_prompt += f"User Memories:\n{memories_str}\n" -总结流程是一个**自下而上**的学习与更新过程：将新发生的对话转化为结构化事件，再进一步纳入长期记忆。 + # Generate response + system_message = {"role": "system", "content": system_prompt} + history_messages.append({"role": "user", "content": query}) + response = openai_client.chat.completions.create( + model="qwen-plus", + messages=[system_message] + history_messages + ) + history_messages.append({"role": "assistant", "content": response.choices[0].message.content}) -```mermaid -flowchart TD - op0[op: 收集 DayN 的原始对话与行为轨迹] - op1[op: 按时间/业务逻辑/固定长度切片为 messages block] - op2[op: 对每个 block 做摘要，生成 block_id: time + summary] - op3[op: 将 block summary 流交给 Meta Agent] - op4[op: Meta Agent 判断主题并路由/创建对应 Sub Agent] - op5[op: Sub Agent 加载当前主题的长期记忆与新 block] - op6[op: 通过 recall_similar_memory 检索历史相似记忆] - op7[op: 对比新旧信息，决定 delete/new/update 记忆操作] - op8[op: 写回更新后的高层主题记忆，维护 block_id 链接] - op9[op: 将最新记忆视图回流给 Meta Agent，用于后续检索] + # Summarize history when it gets too long + if len(history_messages) >= start_summary_size: + memory.summary(history_messages[:-keep_size], user_id=user_name) + print("Current memories: " + memory.list_memories(user_id=user_name)) + history_messages = history_messages[-keep_size:] - op0 --> op1 --> op2 --> op3 --> op4 --> op5 --> op6 --> op7 --> op8 --> op9 + return history_messages[-1]["content"] + +def main(): + user_name = input("Enter your name: ").strip() + print("Chat with Remy (type 'exit' to quit)") + + messages = [] + while True: + user_input = input(f"{user_name}: ").strip() + if user_input.lower() == 'exit': + print("Goodbye!") + break + + print(f"Remy: {chat_with_memories(user_input, messages, user_name)}") + + # Cleanup + memory.delete_all_memories(user_id=user_name) + print("All memories deleted") + +if __name__ == "__main__": + main() ``` -### 输入与切片 +#### 3.1.3 Advanced Usage (高级用法) -- 输入源： - - `DayN Messages / Trajectory`（第 N 天的完整对话与行为轨迹）。 -- 切片策略： - - 按时间窗口切片（例如每 5 分钟一个 block）。 - - 按业务逻辑边界切片（如一个任务完成即成一个 block）。 - - 按固定长度（可带重叠）切片，兼顾局部上下文。 -- 结果： - - `messages block1, block2, ... blockN` +For advanced users who want to customize retriever and summarizer behavior with Agentic mode. -### 初步总结：生成索引条目 +```python +import os +from reme_ai import ReMe +from reme_ai.retriever import AgenticRetriever +from reme_ai.summarizer import AgenticSummarizer +from reme_ai.tools import ATool, BTool, CTool -- 对每个 `messages block` 进行摘要，形成： - - `block_id: time + summary` -- 示例： - - `id: 20251111_01 - 用户询问“什么是 AI”并表达对技术的好奇。` -- 这一阶段的目标： - - 为中层索引层提供结构化事件单元，尚不过度抽象。 +os.environ["REME_LLM_API_KEY"] = "sk-..." +os.environ["REME_LLM_BASE_URL"] = "https://dashscope.aliyuncs.com/compatible-mode/v1" +os.environ["REME_EMBEDDING_API_KEY"] = "sk-..." +os.environ["REME_EMBEDDING_BASE_URL"] = "https://dashscope.aliyuncs.com/compatible-mode/v1" -### Meta Agent 路由与子智能体更新 +memory = ReMe( + memory_space="remy", + llm={"backend": "openai", "model": "qwen-plus", "temperature": 0.6}, + embedding={"backend": "openai", "model": "text-embedding-v4", "dimension": 1024}, + vector_store={"backend": "local_file"}, + use_agentic_mode=True, +) -- Meta Agent 接收新的 `block summary` 流： - - 根据内容判断其归属主题（例如用户相关、任务相关、工具体验相关等）。 - - 调用 `create_sub_agent(topic)` 或路由到已有的 `Sub_Agent`。 -- 子智能体内的处理步骤： - - **Step 1：加载上下文记忆** - - 加载与当前主题相关的长期记忆（如“对用户 A 的记忆”）与当前 `messages block`。 - - **Step 2：回忆（Recall）** - - 通过 `recall_similar_memory(query)` 拉取历史相似记忆，用于对比新旧信息。 - - **Step 3：更新（Update）** - - 基于新旧信息的对比，决定对顶层记忆的操作： - - `delete_memory(mem_id)`：删除已过时或被否定的记忆。 - - `new_memory(mem_content)`：新增一条新的长期记忆。 - - `update_memory(mem_id, mem_content)`：对既有记忆做内容更新或细化。 +# Customize retriever and summarizer with custom tools and prompts +memory.set_retriever( + AgenticRetriever(tools=[ATool(), BTool(), CTool()]), + system_prompt="Custom retrieval instructions..." +) +memory.set_summarizer( + AgenticSummarizer(tools=[ATool(), BTool(), CTool()]) +) + +# Use the customized memory system +result = memory.summary( + messages=[ + {"role": "user", "content": "I'm travelling to SF"}, + {"role": "assistant", "content": "That's great to hear!"} + ], + user_id="Alice", + memory_type="auto", # auto, personal, procedural, tool +) + +memories = memory.retrieve( + query="what is your travel plan?", + limit=3, + user_id="Alice", + memory_type="auto", +) +memories_str = "\n".join(f"- {m['memory']}" for m in memories["results"]) +print(memories_str) +``` + +### 3.2 Short-Term Memory (短期记忆) + +#### 3.2.1 Basic Usage (基础用法) + +Context offload/reload API for managing short-term conversational memory within a session. + +```python +import os +from reme_ai import ReMe + +os.environ["REME_LLM_API_KEY"] = "sk-..." +os.environ["REME_LLM_BASE_URL"] = "https://dashscope.aliyuncs.com/compatible-mode/v1" +os.environ["REME_EMBEDDING_API_KEY"] = "sk-..." +os.environ["REME_EMBEDDING_BASE_URL"] = "https://dashscope.aliyuncs.com/compatible-mode/v1" + +memory = ReMe( + memory_space="remy", + llm={"backend": "openai", "model": "qwen-plus", "temperature": 0.6}, + embedding={"backend": "openai", "model": "text-embedding-v4", "dimension": 1024}, + vector_store={"backend": "local_file"}, +) + +# Offload context when conversation gets too long +result = memory.offload_context( + messages=[ + {"role": "user", "content": "I'm travelling to SF"}, + {"role": "assistant", "content": "That's great to hear!"} + ], +) + +# Reload relevant context when needed +memories = memory.reload_context( + query="what is your travel plan?", + limit=3, +) +memories_str = "\n".join(f"- {m['memory']}" for m in memories["results"]) +print(memories_str) +``` + +### 3.3 Framework Integration (框架集成) + +#### 3.3.1 Integration with AgentScope + +Integration example for AgentScope ReActAgent with long-term memory support. + +```python +# TODO: Provide AgentScope integration example +``` + +#### 3.3.2 Integration with LangChain + +Integration example for LangChain agents with ReMe memory layer. + +```python +# TODO: Provide LangChain integration example +``` + +### 3.4 OpenAI Compatible Interface + +OpenAI-compatible API interface for seamless integration with existing OpenAI-based applications. + +```python +# TODO: Research and implement OpenAI-compatible interface +# - Support for threads and assistants API +# - Compatible with OpenAI SDK +# - Support for streaming responses +``` -### 输出与闭环 -- Summarizer 输出一批新的或更新后的 **高层记忆项（New Memories）**，并附带 `block_id` 链接。 -- 这些记忆将： - - 被 Meta Agent 纳入其主题索引中。 - - 为后续 Retriever 流程提供可检索的抽象知识。 -- 至此，形成从“原始对话 → 索引总结 → 高维记忆 → 检索使用 → 再次总结更新”的完整闭环。 --- -## 总结 +## 四、核心方案设计 -- **分层存储**：底层保留原始细节，中层作为事件索引，顶层沉淀抽象认知；中层通过 `block_id` 将两端连通。 -- **按需调用**：检索时从高层主题记忆出发，由 Meta Agent 路由到 Sub Agent，再逐级回溯到底层对话。 -- **动态增删改**：新对话先被切片与总结，再与旧记忆对比，通过新增、更新、删除操作保持长期记忆“准确 + 精简 + 有时序感”。 -- **渐进式检索与总结**：在检索方向自上而下逐级展开，在学习方向自下而上逐级抽象，实现“用多少展开多少、学多少沉淀多少”的流式记忆管理。 +### 4.1 设计概述 +ReMeV2 采用简洁的架构设计，核心理念为：**ReMeV2 = Tool(s) + Agent(s)** -``` python code -base_memory中增加 -# block_id: str = Field(default=..., description="block_id") -block_ids: List[str] = Field(default=..., description="block_id") +- **Tool层**：提供原子化的记忆操作能力，包括增删改查、检索、元数据管理等基础操作 +- **Agent层**：基于Tool层构建的智能代理，负责复杂的记忆管理逻辑，如分类总结、渐进式检索等 +- **Runtime层**：内部调度机制，协调Tool和Agent的交互流程 +### 4.2 Tool层设计 -from typing import List, Tuple +Tool层提供装饰器形式的记忆操作工具，每个工具类通过 `@tool` 装饰器注册，明确定义初始化参数和调用参数。 -from flowllm.core.schema import Message -from pydantic import BaseModel, Field +#### 4.2.1 基类：BaseMemoryToolOp -class MessageModel(Message): +**初始化参数：** +- `enable_multiple` (bool): Enable multi-item operation mode. Default: `True` +- `enable_thinking_params` (bool): Include thinking parameter in tool schema for model reasoning. Default: `False` +- `memory_metadata_dir` (str): Directory path for storing memory metadata. Default: `"./memory_metadata"` + +#### 4.2.2 Tool操作列表 + +以下是所有Tool操作的完整定义，包括继承关系、初始化参数和调用参数： + +| Tool类 | 继承自 | 初始化参数（除基类外） | Tool Call参数（单项模式） | Tool Call参数（多项模式） | +|----------------------------|------------------|------------------------------------------------------------------------------------------------------------------|-----------------------------------------------------------------------------------------------------------------|-----------------------------------------------------------------------------------------------------------------------| +| **AddMemoryOp** | BaseMemoryToolOp | `add_when_to_use` (bool, 默认: False)
`add_metadata` (bool, 默认: True) | `when_to_use` (str, 可选)
`memory_content` (str, 必需)
`metadata` (dict, 可选) | `memories` (array, 必需):
- `when_to_use` (str, 可选)
- `memory_content` (str, 必需)
- `metadata` (dict, 可选) | +| **UpdateMemoryOp** | BaseMemoryToolOp | 无 | `memory_id` (str, 必需)
`memory_content` (str, 必需)
`metadata` (dict, 可选) | `memories` (array, 必需):
- `memory_id` (str, 必需)
- `memory_content` (str, 必需)
- `metadata` (dict, 可选) | +| **DeleteMemoryOp** | BaseMemoryToolOp | 无 | `memory_id` (str, 必需) | `memory_ids` (array[str], 必需) | +| **VectorRetrieveMemoryOp** | BaseMemoryToolOp | `enable_summary_memory` (bool, 默认: False)
`add_memory_type_target` (bool, 默认: False)
`top_k` (int, 默认: 20) | `query` (str, 必需)
`memory_type` (str, 可选, 枚举: [identity, personal, procedural])
`memory_target` (str, 可选) | `query_items` (array, 必需):
- `query` (str, 必需)
- `memory_type` (str, 可选)
- `memory_target` (str, 可选) | +| **AddMetaMemoryOp** | BaseMemoryToolOp | 无 | `memory_type` (str, 必需, 枚举: [personal, procedural])
`memory_target` (str, 必需) | `meta_memories` (array, 必需):
- `memory_type` (str, 必需)
- `memory_target` (str, 必需) | +| **ReadMetaMemoryOp** | BaseMemoryToolOp | `enable_tool_memory` (bool, 默认: False)
`enable_identity_memory` (bool, 默认: False) | 无（无输入schema） | N/A (enable_multiple=False) | +| **AddHistoryMemoryOp** | BaseMemoryToolOp | 无 | `messages` (array[object], 必需) | N/A (enable_multiple=False) | +| **ReadHistoryMemoryOp** | BaseMemoryToolOp | 无 | `memory_id` (str, 必需) | `memory_ids` (array[str], 必需) | +| **AddSummaryMemoryOp** | AddMemoryOp | 无（继承自AddMemoryOp） | `summary_memory` (str, 必需)
`metadata` (dict, 可选) | N/A (enable_multiple=False) | +| **ReadIdentityMemoryOp** | BaseMemoryToolOp | 无 | 无（无输入schema） | N/A (enable_multiple=False) | +| **UpdateIdentityMemoryOp** | BaseMemoryToolOp | 无 | `identity_memory` (str, 必需) | N/A (enable_multiple=False) | +| **ThinkToolOp** | BaseAsyncToolOp | `add_output_reflection` (bool, 默认: False) | `reflection` (str, 必需) | N/A | +| **HandsOffOp** | BaseMemoryToolOp | 无 | `memory_type` (str, 必需, 枚举: [identity, personal, procedural, tool])
`memory_target` (str, 必需) | `memory_tasks` (array, 必需):
- `memory_type` (str, 必需)
- `memory_target` (str, 必需) | + +### 4.3 Agent层设计 + +#### 4.3.1 基类：BaseMemoryAgentOp + +Agent层构建在Tool层之上，封装复杂的记忆管理逻辑。每个Agent通过组合多个Tool实现特定的记忆管理任务。 + +#### 4.3.2 Agent操作列表 + +以下是所有Agent操作的完整定义，包括初始化参数、调用参数和可用工具： + +| Agent类 | 继承自 | 初始化参数（基类外） | Tool Call参数 | 可用工具 | +|--------------------------------|-------------------|------------------------------------------------------------------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------| +| **PersonalSummaryAgentV1Op** | BaseMemoryAgentOp | None | `workspace_id` (str, required)
`memory_target` (str, required)
`query` (str, optional)
`messages` (array, optional)
`ref_memory_id` (str, required) | add_memory
update_memory
delete_memory
vector_retrieve_memory | +| **ProceduralSummaryAgentV1Op** | BaseMemoryAgentOp | None | `workspace_id` (str, required)
`memory_target` (str, required)
`query` (str, optional)
`messages` (array, optional)
`ref_memory_id` (str, required) | add_memory
update_memory
delete_memory
vector_retrieve_memory | +| **ToolSummaryAgentV1Op** | BaseMemoryAgentOp | None | `workspace_id` (str, required)
`memory_target` (str, required)
`query` (str, optional)
`messages` (array, optional)
`ref_memory_id` (str, required) | add_memory
update_memory
vector_retrieve_memory | +| **IdentitySummaryAgentV1Op** | BaseMemoryAgentOp | None | `workspace_id` (str, required)
`query` (str, optional)
`messages` (array, optional) | read_identity_memory
update_identity_memory | +| **ReMeSummaryAgentV1Op** | BaseMemoryAgentOp | `enable_tool_memory` (bool, 默认: True)
`enable_identity_memory` (bool, 默认: True) | `workspace_id` (str, required)
`query` (str, optional)
`messages` (array, optional) | add_meta_memory
add_summary_memory
hands_off
(内部调用: add_history_memory, read_identity_memory, read_meta_memory) | +| **ReMeRetrieveAgentV1Op** | BaseMemoryAgentOp | `enable_tool_memory` (bool, 默认: True) | `workspace_id` (str, required)
`query` (str, optional)
`messages` (array, optional) | vector_retrieve_memory
read_history_memory
(内部调用: read_meta_memory) | +| **ReMyAgentV1Op** | BaseMemoryAgentOp | `enable_tool_memory` (bool, 默认: True)
`enable_identity_memory` (bool, 默认: True) | `workspace_id` (str, required)
`query` (str, optional)
`messages` (array, optional) | vector_retrieve_memory
read_history_memory
(内部调用: read_identity_memory, read_meta_memory) | + +### 4.4 Runtime层设计（内部实现） + +Runtime层负责协调Tool和Agent的调用流程，实现记忆的渐进式处理。 + +#### 4.4.1 渐进式总结流程（Summary） + +总结流程采用分层处理策略，首先保存历史对话，读取元信息，然后由主Agent协调多个专用Agent完成分类总结。 + +**流程结构：** + +```python +# Step 1: Save conversation history +AddHistoryMemoryOp() + +# Step 2: Load meta information (memory types and targets) +ReadMetaMemoryOp() + +# Step 3: Progressive summarization with delegation +ReMeSummaryAgentV1Op(tools=[ + # Add meta memory entries for new memory types/targets + AddMetaMemoryOp(list(memory_type, memory_target)), + + # Add general summary memory as fallback + AddSummaryMemoryOp(summary_memory), + + # Delegate to specialized summary agents + HandsOffOp(list(memory_type, memory_target), agents=[ + PersonalSummaryAgentV1Op, # Summarize personal memories + ProceduralSummaryAgentV1Op, # Summarize procedural memories + ToolSummaryAgentV1Op, # Summarize tool-related memories + IdentitySummaryAgentV1Op # Update identity memory + ]), +]) + +# Specialized agents and their available tools +PersonalSummaryAgentV1Op(tools=[AddMemoryOp, UpdateMemoryOp, DeleteMemoryOp, VectorRetrieveMemoryOp]) +ProceduralSummaryAgentV1Op(tools=[AddMemoryOp, UpdateMemoryOp, DeleteMemoryOp, VectorRetrieveMemoryOp]) +ToolSummaryAgentV1Op(tools=[AddMemoryOp, UpdateMemoryOp, VectorRetrieveMemoryOp]) +IdentitySummaryAgentV1Op(tools=[ReadIdentityMemoryOp, UpdateIdentityMemoryOp]) +``` + +#### 4.4.2 渐进式检索流程（Retrieve） + +检索流程采用三层检索策略，类似于技能系统的加载机制，逐层加载和过滤记忆。 + +**流程结构：** + +```python +# Progressive retrieval with three layers +ReMeRetrieveAgentV1Op(tools=[ + # Layer 0: Load meta memory (all available memory types and targets) + ReadMetaMemoryOp(), + # Output format example: + # - personal(jinli): Information about Jinli's personal life and preferences + # - personal(jiaji): Information about Jiaji's background and interests + # - personal(jinli&jiaji): Shared memories between Jinli and Jiaji + # - procedural(appworld): Procedural knowledge for AppWorld tasks + # - procedural(bfcl-v3): Procedural knowledge for BFCL-v3 benchmark + # - tool(tool_guidelines): Guidelines for tool usage + # - identity(self): Agent's self-identity information + + # Layer 1+2: Vector-based retrieval on structured memories + VectorRetrieveMemoryOp(list(memory_type, memory_target, query)), + + # Layer 3: Load full conversation history for specific memory + ReadHistoryMemoryOp(ref_memory_id), +]) +``` + +**与技能系统的类比：** + +```python +# Skill system hierarchy (for reference) +load_meta_skills # Load skill metadata +load_skills # Load skill implementations +load_reference_skills # Load detailed skill documentation +execute_shell # Execute actual commands +``` + +## 五、扩展设计与实验方向 + +### 5.1 Summary Memory机制 + +Summary Memory作为通用维度的记忆类型，提供兜底的原始对话索引能力。 + +**工作流程示例：** + +```txt +Step 1: Progressive summarization across sessions + session1: List[Message] -> session2: List[Message] -> session3: List[Message] -> ... +summary ✓ (always) ✓ (always) ✓ (always) +personal ✗ ✗ ✓ (when applicable) +procedural ✗ ✓ (when applicable) ✗ + +Step 2: Retrieval with fallback strategy +vector_retrieve_memory(query, memory_type="personal", memory_target="jinli") + -> Search in memory_type: ["personal", "summary"] # Fallback to summary if personal not found +``` + +**设计优势：** +1. Provides a universal dimension for memory extraction across all memory types +2. Ensures fallback indexing of original conversations when specific meta memory is not available +3. Maintains conversation context even when specialized memory extraction fails + +### 5.2 Thinking参数实验 + +探索不同的模型推理能力增强方案，受AgentScope和Claude启发。 + +#### 5.2.1 Thinking参数设计 + +```python +async def record_to_memory( + self, + thinking: str, + content: list[str], + **kwargs: Any, +) -> ToolResponse: + """Use this function to record important information that you may + need later. The target content should be specific and concise, e.g. + who, when, where, do what, why, how, etc. + + Args: + thinking (`str`): + Your thinking and reasoning about what to record + content (`list[str]`): + The content to remember, which is a list of strings. """ - 被存储到传统的db中的数据结构，一行一条message - def retrieve_message(key_word=None, time_start=None, limit: int=50): - 支持关键字检索和时间检索 - content like %{key_word}% - 暂时不支持向量检索，向量只做抽取的memory的 - """ - name: str = Field(default="", description="The name of the character who actually outputs this message") - message_id: str = Field(default=..., description="message_id, uuid，每一条message不重样") - block_id: str = Field(default=..., description="block_id，message分块的id，用户链接summary") +``` +#### 5.2.2 实验对比方案 -from reme_ai.schema.memory import BaseMemory +| 方案类型 | 说明 | 灵感来源 | +|-------------------------------|--------------------------------------------|----------------| +| Thinking Model | Native reasoning-capable models (e.g., o1) | OpenAI | +| Instruct Model | Standard instruction-following models | Baseline | +| Instruct Model + Thinking Params | Add thinking parameter to tool schema | AgentScope | +| Instruct Model + Thinking Tool | Dedicated thinking tool for explicit reasoning | Claude | -""" -基本和原来一致，增加block compress的memory +### 5.3 多项操作模式实验 -工具记忆放到传统db中 -BlockCompress放到向量库中 -Personal + Task + Self-Cognition 放到向量库中 -""" +对比单次调用和批量调用的性能与准确性差异。 +**两种模式对比：** +| 模式 | Tool调用方式 | Model调用次数 | 优势 | 劣势 | +|--------------|----------------------------|---------------|------------------------------|--------------------------| +| 单项模式 | Single-item per call | Multiple | Fine-grained control | Higher latency, more tokens | +| 多项模式 | Batch multiple items | Single | Lower latency, fewer tokens | Potential batch errors | +**实验目标：** +- Evaluate accuracy: single vs. batch operations +- Measure latency and token efficiency +- Identify optimal use cases for each mode -class BlockCompressMemory(BaseMemory): - start_time: str = Field(default=..., description="block的开始时间") - end_time: str = Field(default=..., description="block的结束时间") +### 5.4 多版本与扩展性 -class SelfCognitionMemory(BaseMemory): - target_name: str = Field(default="自我认知") - memory_type: str = Field(default="self") +支持从基类继承实现自定义Agent，便于团队协作和功能迭代。 -class ToolMemory(BaseMemory): - target_name: str = Field(default="", description="工具名称") - memory_type: str = Field(default="tool") +**扩展示例：** -class PersonalMemory(BaseMemory): - target_name: str = Field(default="", description="这是记录谁的记忆？比如TOM、Jerry？") - memory_type: str = Field(default="personal") +```python +# Version 2 implementations by different team members +PersonalSummaryAgentV2Op / PersonalRetrieveAgentV2Op # @weikang +ProceduralSummaryAgentV2Op / ProceduralRetrieveAgentV2Op # @zouyin -class TaskMemory(BaseMemory): - target_name: str = Field(default="", description="这是什么任务的记忆、比如bfcl、appworld、金融?") - memory_type: str = Field(default="task") +# Inherit from BaseMemoryAgentOp +class PersonalSummaryAgentV2Op(BaseMemoryAgentOp): + """Enhanced personal memory summarization with improved algorithms""" + pass +``` +### 5.5 文件系统集成（未来方向） +探索将文件操作能力集成到记忆系统中，支持基于文件的记忆管理。 -# Handoffs allow meta agent to delegate tasks to another agent. -""" -summarizer -""" +**挑战与考虑：** -def compress_message_block(message_block: List[Message]) -> BlockCompressMemory: - """ - 压缩message_block信息到BlockCompressMemory - """ +1. **操作适配性**：Current operations (retrieve/add/update/delete) need adaptation for file-based storage +2. **工具选择**：Consider file operation tools: `grep`, `glob`, `ls`, `read_file`, `write_file`, `edit_file` +3. **模型能力**：Base models have limited file operation capabilities; `qwen3-code` shows better performance -def handoff_to_summary_agent(agent_name: str) -> str: - """ - 如果agent_name是存在的，那直接使用现成的，否则可以创建新的agent(默认是不开启的) - """ +**潜在架构：** -def retrieve_memory_by_query(query: str, limit_size: int = 50, threshold = 0.5) -> List[BaseMemory]: - """ - 如果之前使用过handoff_to_summary_agent - 在自我认知 - """ +```python +# File-based memory operations +FileMemoryOp(tools=[ + grep, # Search within files + glob, # File pattern matching + ls, # List directory contents + read_file, # Read file contents + write_file, # Write new memory files + edit_file, # Update existing memory files +]) +``` -def datetime_tool(**kwargs) - """ - 时间工具 - Returns: +### 5.6 自我修改上下文 - """ +支持Agent动态修改自身的上下文状态，实现自适应记忆管理。 -def delete_memory(memory_id: str): - ... +**实现方式：** -# 做实验 -def update_memory(memory_id: str, content: str, time_stamp: str): - ... +1. **Summary Agent 主动修改**： + - `add_meta_memory` directly modifies agent context + - Updates available memory types and targets during execution +2. **ReMy Agent 被动修改**： + - Retrieves `identity_memory` at each interaction + - Dynamically updates self-state based on retrieved identity + - Enables adaptive behavior based on accumulated identity knowledge -def new_memory(content: str, time_stamp: str) -> BaseMemory: - ... +## ReMe V2 开发路线图与实施计划 +### 技术改造阶段 +1. **代码整合与兼容**：合并flowllm中reme必要的代码，保留现在server-client的依赖，兼容现在各个仓库的依赖代码 +2. **核心接口重构**：新的ReMe接口设计，支持summary，retrieve，context_offload, context_reload 4个核心接口 +3. **Agentic算法升级**：新的agentic算法方案开发 -""" -summary reviser / Reward Model -""" +### 评估验证阶段 +4. **Benchmark测试** + - halumem + - locomo + - longmemevel + - personal-v2 ? + - appworld/bfcl-v3 -def revise_memory(messages: List[Message], - origin_memory: List[BaseMemory], - output_memory: List[BaseMemory]) -> Tuple["", bool]: - """ - reasoning and result - 是否所有信息都加入了？ - 当前结果是不是有冲突和重复。 - 如果有问题，就重新react或者重试 - """ +### 发布推广阶段 +5. **技术报告**撰写与发布 +6. **生态更新**：更新各个仓库的依赖代码 + - agentscope + - agentscope-runtime + - evotraders + - alias(tool-memory) + - agentscope-java + - AgentEvolver + - cookbook: reme procedural memory paper + - tool-memory-upgrade（将要合并） -""" -retriever -""" +**里程碑目标**：春节前完成小版本发布 -def retrieve_memory_by_query(query: str, limit_size: int = 50, threshold = 0.5) -> List[BaseMemory]: - """ - 和上面不一致的在于，这里可以搜索到compress的记忆 - """ +--- -def retrieve_message(key_word=None, time_offset=None, time_limit=None): - """ - 支持关键字检索和时间检索 - content like %{key_word}% - 暂时不支持向量检索，向量只做抽取的memory的 - """ +## ReMe V2 核心竞争优势 -def retrieve_block_id(block_id): - """ - 支持关键字检索和时间检索 - content like %{key_word}% - 暂时不支持向量检索，向量只做抽取的memory的 - """ +### 1. 渐进式 Agentic Memory 架构【核心创新】 +融合了多种记忆的渐进式agentic方案，实现从短期到长期记忆的智能化演进 -``` \ No newline at end of file +### 2. 全生命周期记忆管理 +同时支持长期记忆（Long-term Memory）和短期记忆（Working Memory），完整覆盖Agent认知周期 + +### 3. 模型 +提供开源小模型 + +### 4. 开发者友好生态 + 1. **简洁接口**：提供简洁的接口设计，全异步接口 + 2. **即开即用**：提供CLI工具，开箱即用的体验 + 3. **生态融合**：提供和AgentScope、LangChain无缝集成的方案 + 4. **高度可扩展**：支持Agentic算法的二次开发与定制 \ No newline at end of file diff --git a/reme_ai/core/flow/__init__.py b/reme_ai/core/flow/__init__.py new file mode 100644 index 00000000..6d5a053b --- /dev/null +++ b/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", +] diff --git a/reme_ai/core/flow/base_flow.py b/reme_ai/core/flow/base_flow.py new file mode 100644 index 00000000..1802af78 --- /dev/null +++ b/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 diff --git a/reme_ai/core/flow/cmd_flow.py b/reme_ai/core/flow/cmd_flow.py new file mode 100644 index 00000000..2fe3b97a --- /dev/null +++ b/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=` in cmd!" + + def _build_flow(self) -> BaseOp: + """Parse the stored flow expression into a functional operation object.""" + return self.parse_expression(self.flow) diff --git a/reme_ai/core/flow/expression_flow.py b/reme_ai/core/flow/expression_flow.py new file mode 100644 index 00000000..5ac8f0d8 --- /dev/null +++ b/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}) diff --git a/reme_ai/core/op/__init__.py b/reme_ai/core/op/__init__.py index a8ba9a12..c009e89f 100644 --- a/reme_ai/core/op/__init__.py +++ b/reme_ai/core/op/__init__.py @@ -1,11 +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", ] From afa319698837a8c73b28399ff1d86c49e1a6f3d3 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Mon, 5 Jan 2026 14:27:07 +0800 Subject: [PATCH 09/11] feat(core): add service infrastructure and utility modules --- reme_ai/core/schema/service_config.py | 5 +- reme_ai/core/service/__init__.py | 12 ++ reme_ai/core/service/base_service.py | 41 ++++ reme_ai/core/service/cmd_service.py | 39 ++++ reme_ai/core/service/http_service.py | 109 ++++++++++ reme_ai/core/service/mcp_service.py | 54 +++++ reme_ai/core/utils/__init__.py | 11 + reme_ai/core/utils/common_utils.py | 20 ++ reme_ai/core/utils/llm_utils.py | 40 ++++ reme_ai/core/utils/logger_utils.py | 40 ++++ reme_ai/core/utils/logo_utils.py | 84 ++++++++ reme_ai/core/utils/pydantic_config_parser.py | 208 +++++++++++++++++++ tests/test_logo.py | 7 + 13 files changed, 668 insertions(+), 2 deletions(-) create mode 100644 reme_ai/core/service/__init__.py create mode 100644 reme_ai/core/service/base_service.py create mode 100644 reme_ai/core/service/cmd_service.py create mode 100644 reme_ai/core/service/http_service.py create mode 100644 reme_ai/core/service/mcp_service.py create mode 100644 reme_ai/core/utils/common_utils.py create mode 100644 reme_ai/core/utils/llm_utils.py create mode 100644 reme_ai/core/utils/logger_utils.py create mode 100644 reme_ai/core/utils/logo_utils.py create mode 100644 reme_ai/core/utils/pydantic_config_parser.py create mode 100644 tests/test_logo.py diff --git a/reme_ai/core/schema/service_config.py b/reme_ai/core/schema/service_config.py index 023eb3a7..cec1c034 100644 --- a/reme_ai/core/schema/service_config.py +++ b/reme_ai/core/schema/service_config.py @@ -1,5 +1,5 @@ """Configuration schemas for service components using Pydantic models.""" - +import os from typing import Dict, List from pydantic import BaseModel, Field, ConfigDict @@ -12,7 +12,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,6 +92,7 @@ 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) diff --git a/reme_ai/core/service/__init__.py b/reme_ai/core/service/__init__.py new file mode 100644 index 00000000..5a752beb --- /dev/null +++ b/reme_ai/core/service/__init__.py @@ -0,0 +1,12 @@ +"""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", +] diff --git a/reme_ai/core/service/base_service.py b/reme_ai/core/service/base_service.py new file mode 100644 index 00000000..b81a710d --- /dev/null +++ b/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) -> None: + """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)}") diff --git a/reme_ai/core/service/cmd_service.py b/reme_ai/core/service/cmd_service.py new file mode 100644 index 00000000..5e6f21bf --- /dev/null +++ b/reme_ai/core/service/cmd_service.py @@ -0,0 +1,39 @@ +"""Command service module for managing and executing command-based workflows.""" + +from typing import Any + +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) + return None + + def run(self) -> None: + """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}") diff --git a/reme_ai/core/service/http_service.py b/reme_ai/core/service/http_service.py new file mode 100644 index 00000000..48a28598 --- /dev/null +++ b/reme_ai/core/service/http_service.py @@ -0,0 +1,109 @@ +"""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 loguru import logger + +from .base_service import BaseService +from ..context import C +from ..enumeration import ChunkEnum +from ..flow import BaseFlow +from ..schema import Response, StreamChunk + + +@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]: + done_bytes = b"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_bytes + break + yield f"data:{chunk.model_dump_json()}\n\n".encode() + else: + # Task finished unexpectedly or raised exception + await task + yield done_bytes + break + + except Exception as e: + logger.exception(f"Stream error in {tool_call.name}: {e}") + err = StreamChunk(chunk_type=ChunkEnum.ERROR, chunk=str(e), done=True) + yield f"data:{err.model_dump_json()}\n\n".encode() + yield done_bytes + + finally: + if not task.done(): + task.cancel() + + 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) -> None: + """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, + ) diff --git a/reme_ai/core/service/mcp_service.py b/reme_ai/core/service/mcp_service.py new file mode 100644 index 00000000..4ad2ba0a --- /dev/null +++ b/reme_ai/core/service/mcp_service.py @@ -0,0 +1,54 @@ +"""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) diff --git a/reme_ai/core/utils/__init__.py b/reme_ai/core/utils/__init__.py index 6b217d21..7767a13b 100644 --- a/reme_ai/core/utils/__init__.py +++ b/reme_ai/core/utils/__init__.py @@ -2,9 +2,14 @@ from .cache_handler import CacheHandler from .case_converter import snake_to_camel, camel_to_snake +from .common_utils import run_coro_safely 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 @@ -13,9 +18,15 @@ __all__ = [ "CacheHandler", "snake_to_camel", "camel_to_snake", + "run_coro_safely", "load_env", "HttpClient", + "extract_content", + "format_messages", + "init_logger", + "print_logo", "MCPClient", + "PydanticConfigParser", "create_pydantic_model", "singleton", "timer", diff --git a/reme_ai/core/utils/common_utils.py b/reme_ai/core/utils/common_utils.py new file mode 100644 index 00000000..942b1a21 --- /dev/null +++ b/reme_ai/core/utils/common_utils.py @@ -0,0 +1,20 @@ +"""Common utility functions""" + +import asyncio +from collections.abc import Coroutine +from typing import Any + + +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) \ No newline at end of file diff --git a/reme_ai/core/utils/llm_utils.py b/reme_ai/core/utils/llm_utils.py new file mode 100644 index 00000000..f51dc9db --- /dev/null +++ b/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 \ No newline at end of file diff --git a/reme_ai/core/utils/logger_utils.py b/reme_ai/core/utils/logger_utils.py new file mode 100644 index 00000000..56478683 --- /dev/null +++ b/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, + ) diff --git a/reme_ai/core/utils/logo_utils.py b/reme_ai/core/utils/logo_utils.py new file mode 100644 index 00000000..08b70ebf --- /dev/null +++ b/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")) \ No newline at end of file diff --git a/reme_ai/core/utils/pydantic_config_parser.py b/reme_ai/core/utils/pydantic_config_parser.py new file mode 100644 index 00000000..3856758d --- /dev/null +++ b/reme_ai/core/utils/pydantic_config_parser.py @@ -0,0 +1,208 @@ +"""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 + if not config: + raise ValueError("add `config=` in cmd!") + + # 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) diff --git a/tests/test_logo.py b/tests/test_logo.py new file mode 100644 index 00000000..e8bbd0ba --- /dev/null +++ b/tests/test_logo.py @@ -0,0 +1,7 @@ +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) From 36e88b26ddb526f6ae47348c549c2ff38b769edd Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Mon, 5 Jan 2026 14:41:22 +0800 Subject: [PATCH 10/11] style(core): format code according to team style guide --- reme_ai/core/context/service_context.py | 13 ++++++- reme_ai/core/schema/service_config.py | 1 + reme_ai/core/service/__init__.py | 1 + reme_ai/core/service/cmd_service.py | 5 +-- reme_ai/core/service/http_service.py | 2 +- reme_ai/core/service/mcp_service.py | 14 ++++---- reme_ai/core/utils/common_utils.py | 2 +- reme_ai/core/utils/llm_utils.py | 2 +- reme_ai/core/utils/logo_utils.py | 2 +- reme_ai/core/utils/pydantic_config_parser.py | 36 +++++++++----------- tests/test_logo.py | 2 ++ 11 files changed, 46 insertions(+), 34 deletions(-) diff --git a/reme_ai/core/context/service_context.py b/reme_ai/core/context/service_context.py index 97426594..6d67b695 100644 --- a/reme_ai/core/context/service_context.py +++ b/reme_ai/core/context/service_context.py @@ -22,10 +22,21 @@ class ServiceContext(BaseContext): self.thread_pool: ThreadPoolExecutor | None = None self.vector_store_dict: dict[str, dict] = {} self.mcp_server_tool_call_mapping: 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 = {} + def _update_config_section(self, section_name: str, update_dict: dict | None): + if not update_dict: + return + + target_registry = getattr(self.service_config, section_name) + if "default" not in target_registry: + raise KeyError(f"Default `{section_name}` config not found in service_config") + + current_config = target_registry["default"] + target_registry["default"] = current_config.model_copy(update=update_dict, deep=True) + def register(self, name: str, register_type: RegistryEnum): """Return a decorator to register a component within a specific registry category.""" return self.registry_dict[register_type].register(name=name) diff --git a/reme_ai/core/schema/service_config.py b/reme_ai/core/schema/service_config.py index cec1c034..6471e982 100644 --- a/reme_ai/core/schema/service_config.py +++ b/reme_ai/core/schema/service_config.py @@ -1,4 +1,5 @@ """Configuration schemas for service components using Pydantic models.""" + import os from typing import Dict, List diff --git a/reme_ai/core/service/__init__.py b/reme_ai/core/service/__init__.py index 5a752beb..e9f00a65 100644 --- a/reme_ai/core/service/__init__.py +++ b/reme_ai/core/service/__init__.py @@ -1,4 +1,5 @@ """service""" + from .base_service import BaseService from .cmd_service import CmdService from .http_service import HttpService diff --git a/reme_ai/core/service/cmd_service.py b/reme_ai/core/service/cmd_service.py index 5e6f21bf..0b10ac1e 100644 --- a/reme_ai/core/service/cmd_service.py +++ b/reme_ai/core/service/cmd_service.py @@ -1,7 +1,5 @@ """Command service module for managing and executing command-based workflows.""" -from typing import Any - from loguru import logger from .base_service import BaseService @@ -22,7 +20,6 @@ class CmdService(BaseService): 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) - return None def run(self) -> None: """Execute the command flow in either asynchronous or synchronous mode.""" @@ -30,7 +27,7 @@ class CmdService(BaseService): if self._cmd_flow.async_mode: response = run_coro_safely( - self._cmd_flow.call(**C.service_config.cmd.model_extra) + self._cmd_flow.call(**C.service_config.cmd.model_extra), ) else: response = self._cmd_flow.call_sync(**C.service_config.cmd.model_extra) diff --git a/reme_ai/core/service/http_service.py b/reme_ai/core/service/http_service.py index 48a28598..9ed56e17 100644 --- a/reme_ai/core/service/http_service.py +++ b/reme_ai/core/service/http_service.py @@ -43,7 +43,7 @@ class HttpService(BaseService): self.app.post( path=f"/{tool_call.name}", response_model=Response, - description=tool_call.description + description=tool_call.description, )(execute_endpoint) return tool_call.name diff --git a/reme_ai/core/service/mcp_service.py b/reme_ai/core/service/mcp_service.py index 4ad2ba0a..e60aec85 100644 --- a/reme_ai/core/service/mcp_service.py +++ b/reme_ai/core/service/mcp_service.py @@ -32,12 +32,14 @@ class MCPService(BaseService): 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(), - )) + 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): diff --git a/reme_ai/core/utils/common_utils.py b/reme_ai/core/utils/common_utils.py index 942b1a21..5fe01a61 100644 --- a/reme_ai/core/utils/common_utils.py +++ b/reme_ai/core/utils/common_utils.py @@ -17,4 +17,4 @@ def run_coro_safely(coro: Coroutine[Any, Any, Any]) -> Any | asyncio.Task[Any]: else: # Schedule the coroutine as a background task in the active loop - return loop.create_task(coro) \ No newline at end of file + return loop.create_task(coro) diff --git a/reme_ai/core/utils/llm_utils.py b/reme_ai/core/utils/llm_utils.py index f51dc9db..ae0e763f 100644 --- a/reme_ai/core/utils/llm_utils.py +++ b/reme_ai/core/utils/llm_utils.py @@ -37,4 +37,4 @@ def extract_content(text: str, language_tag: str = "json", greedy: bool = False) except json.JSONDecodeError: result = None - return result \ No newline at end of file + return result diff --git a/reme_ai/core/utils/logo_utils.py b/reme_ai/core/utils/logo_utils.py index 08b70ebf..91c46a02 100644 --- a/reme_ai/core/utils/logo_utils.py +++ b/reme_ai/core/utils/logo_utils.py @@ -81,4 +81,4 @@ def print_logo(service_config: "ServiceConfig"): expand=False, ) - Console().print(Group("\n", panel, "\n")) \ No newline at end of file + Console().print(Group("\n", panel, "\n")) diff --git a/reme_ai/core/utils/pydantic_config_parser.py b/reme_ai/core/utils/pydantic_config_parser.py index 3856758d..2f4f2ae0 100644 --- a/reme_ai/core/utils/pydantic_config_parser.py +++ b/reme_ai/core/utils/pydantic_config_parser.py @@ -17,7 +17,7 @@ class PydanticConfigParser: 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. @@ -69,13 +69,13 @@ class PydanticConfigParser: @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. """ @@ -90,10 +90,10 @@ class PydanticConfigParser: 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. """ @@ -104,10 +104,10 @@ class PydanticConfigParser: 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. """ @@ -118,7 +118,7 @@ class PydanticConfigParser: key_path, value_str = item.split("=", 1) keys = key_path.split(".") - + # Build nested dictionary current = config_dict for key in keys[:-1]: @@ -131,13 +131,13 @@ class PydanticConfigParser: """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) @@ -147,13 +147,13 @@ class PydanticConfigParser: 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. @@ -174,8 +174,6 @@ class PydanticConfigParser: # Use default config if not specified config = config or self.default_config - if not config: - raise ValueError("add `config=` in cmd!") # Load each config file for single_config in (c.strip() for c in config.split(",") if c.strip()): @@ -192,17 +190,17 @@ class PydanticConfigParser: 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) diff --git a/tests/test_logo.py b/tests/test_logo.py index e8bbd0ba..eeede81e 100644 --- a/tests/test_logo.py +++ b/tests/test_logo.py @@ -1,3 +1,5 @@ +"""test logo""" + from reme_ai.core.schema import ServiceConfig, MCPConfig if __name__ == "__main__": From 3c8eca8a3bb995851e95e878a85004758a07effb Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Mon, 5 Jan 2026 17:15:13 +0800 Subject: [PATCH 11/11] feat(core): add application lifecycle management and streaming flow execution --- reme_ai/core/application.py | 208 ++++++++++ reme_ai/core/config/__init__.py | 5 + reme_ai/core/config/default.yaml | 34 ++ reme_ai/core/config/reme_config_parser.py | 7 + reme_ai/core/context/service_context.py | 467 ++++++++++++++++++++-- reme_ai/core/op/base_op.py | 51 +-- reme_ai/core/reme.py | 92 +++++ reme_ai/core/service/base_service.py | 2 +- reme_ai/core/service/cmd_service.py | 2 +- reme_ai/core/service/http_service.py | 42 +- reme_ai/core/tool/mcp_tool.py | 2 +- reme_ai/core/utils/__init__.py | 3 +- reme_ai/core/utils/common_utils.py | 65 ++- 13 files changed, 880 insertions(+), 100 deletions(-) create mode 100644 reme_ai/core/application.py create mode 100644 reme_ai/core/config/__init__.py create mode 100644 reme_ai/core/config/default.yaml create mode 100644 reme_ai/core/config/reme_config_parser.py create mode 100644 reme_ai/core/reme.py diff --git a/reme_ai/core/application.py b/reme_ai/core/application.py new file mode 100644 index 00000000..e15f4fcf --- /dev/null +++ b/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() diff --git a/reme_ai/core/config/__init__.py b/reme_ai/core/config/__init__.py new file mode 100644 index 00000000..a8d954d2 --- /dev/null +++ b/reme_ai/core/config/__init__.py @@ -0,0 +1,5 @@ +"""config""" + +from .reme_config_parser import ReMeConfigParser + +__all__ = ["ReMeConfigParser"] diff --git a/reme_ai/core/config/default.yaml b/reme_ai/core/config/default.yaml new file mode 100644 index 00000000..cd04da1d --- /dev/null +++ b/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 diff --git a/reme_ai/core/config/reme_config_parser.py b/reme_ai/core/config/reme_config_parser.py new file mode 100644 index 00000000..798235b2 --- /dev/null +++ b/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.""" diff --git a/reme_ai/core/context/service_context.py b/reme_ai/core/context/service_context.py index 6d67b695..c5bcde0c 100644 --- a/reme_ai/core/context/service_context.py +++ b/reme_ai/core/context/service_context.py @@ -1,44 +1,80 @@ """Module for managing global service configurations and component registries via a singleton context.""" from concurrent.futures import ThreadPoolExecutor +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.mcp_server_tool_call_mapping: dict = {} + # Registry system: stores class definitions for different component types self.registry_dict: dict[RegistryEnum, Registry] = {v: Registry() for v in RegistryEnum.__members__.values()} - self.flow_dict: dict = {} - def _update_config_section(self, section_name: str, update_dict: dict | None): - if not update_dict: - return + # Instance system: stores instantiated objects created from registered classes + self.instance_dict: dict[RegistryEnum, dict] = {v: {} for v in RegistryEnum.__members__.values()} - target_registry = getattr(self.service_config, section_name) - if "default" not in target_registry: - raise KeyError(f"Default `{section_name}` config not found in service_config") - - current_config = target_registry["default"] - target_registry["default"] = current_config.model_copy(update=update_dict, deep=True) + # 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 = ""): @@ -70,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) @@ -102,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() diff --git a/reme_ai/core/op/base_op.py b/reme_ai/core/op/base_op.py index 26f1e7f5..7eeb5fbd 100644 --- a/reme_ai/core/op/base_op.py +++ b/reme_ai/core/op/base_op.py @@ -178,21 +178,16 @@ class BaseOp: @property def llm(self) -> BaseLLM: - """Lazily initialize and return the LLM instance.""" + """Get the LLM instance from ServiceContext.""" if isinstance(self._llm, str): - cfg = C.service_config.llm[self._llm] - self._llm = C.get_llm_class(cfg.backend)(model_name=cfg.model_name, **cfg.model_extra) + self._llm = C.get_llm(self._llm) return self._llm @property def embedding_model(self) -> BaseEmbeddingModel: - """Lazily initialize and return the embedding model instance.""" + """Get the embedding model instance from ServiceContext.""" if isinstance(self._embedding_model, str): - cfg = C.service_config.embedding_model[self._embedding_model] - self._embedding_model = C.get_embedding_model_class(cfg.backend)( - model_name=cfg.model_name, - **cfg.model_extra, - ) + self._embedding_model = C.get_embedding_model(self._embedding_model) return self._embedding_model @property @@ -204,10 +199,9 @@ class BaseOp: @property def token_counter(self) -> BaseTokenCounter: - """Lazily initialize and return the token counter instance.""" + """Get the token counter instance from ServiceContext.""" if isinstance(self._token_counter, str): - cfg = C.service_config.token_counter[self._token_counter] - self._token_counter = C.get_token_counter_class(cfg.backend)(model_name=cfg.model_name, **cfg.model_extra) + self._token_counter = C.get_token_counter(self._token_counter) return self._token_counter @property @@ -220,25 +214,6 @@ class BaseOp: """Get the response object.""" return self.context.response - async def before_execute(self): - """Prepare context and validate before async execution.""" - self.context.apply_mapping(self.input_mapping) - self._validate_inputs() - - async def execute(self): - """Define core async logic in subclasses.""" - - async def after_execute(self): - """Finalize context and mappings after async 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 - - if not isinstance(self._llm, str) and hasattr(self._llm, "close"): - await self._llm.close() - if not isinstance(self._embedding_model, str) and hasattr(self._embedding_model, "close"): - await self._embedding_model.close() - def before_execute_sync(self): """Prepare context and validate before sync execution.""" self.context.apply_mapping(self.input_mapping) @@ -253,10 +228,16 @@ class BaseOp: if self.tool_call is not None and self.save_response_result: self.context.response.answer = self.output - if not isinstance(self._llm, str) and hasattr(self._llm, "close_sync"): - self._llm.close_sync() - if not isinstance(self._embedding_model, str) and hasattr(self._embedding_model, "close_sync"): - self._embedding_model.close_sync() + 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): diff --git a/reme_ai/core/reme.py b/reme_ai/core/reme.py new file mode 100644 index 00000000..d1240cb6 --- /dev/null +++ b/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() diff --git a/reme_ai/core/service/base_service.py b/reme_ai/core/service/base_service.py index b81a710d..28c82fb3 100644 --- a/reme_ai/core/service/base_service.py +++ b/reme_ai/core/service/base_service.py @@ -29,7 +29,7 @@ class BaseService(ABC): model = create_pydantic_model(tool_call.name, tool_call.parameters) return tool_call, model - def run(self) -> None: + def run(self): """Initialize and integrate all flows registered in the global context.""" flow_names: list[str] = [] for _, flow in C.flow_dict.items(): diff --git a/reme_ai/core/service/cmd_service.py b/reme_ai/core/service/cmd_service.py index 0b10ac1e..bed8ec86 100644 --- a/reme_ai/core/service/cmd_service.py +++ b/reme_ai/core/service/cmd_service.py @@ -21,7 +21,7 @@ class CmdService(BaseService): """Integrate the workflow configuration into the command service.""" self._cmd_flow = CmdFlow(flow=C.service_config.flow) - def run(self) -> None: + def run(self): """Execute the command flow in either asynchronous or synchronous mode.""" super().run() diff --git a/reme_ai/core/service/http_service.py b/reme_ai/core/service/http_service.py index 9ed56e17..f63d7333 100644 --- a/reme_ai/core/service/http_service.py +++ b/reme_ai/core/service/http_service.py @@ -7,13 +7,12 @@ import uvicorn from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import StreamingResponse -from loguru import logger from .base_service import BaseService from ..context import C -from ..enumeration import ChunkEnum from ..flow import BaseFlow -from ..schema import Response, StreamChunk +from ..schema import Response +from ..utils.common_utils import execute_stream_task @C.register_service("http") @@ -57,34 +56,13 @@ class HttpService(BaseService): task = asyncio.create_task(flow.call(stream_queue=queue, **request.model_dump(exclude_none=True))) async def generate_stream() -> AsyncGenerator[bytes, None]: - done_bytes = b"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_bytes - break - yield f"data:{chunk.model_dump_json()}\n\n".encode() - else: - # Task finished unexpectedly or raised exception - await task - yield done_bytes - break - - except Exception as e: - logger.exception(f"Stream error in {tool_call.name}: {e}") - err = StreamChunk(chunk_type=ChunkEnum.ERROR, chunk=str(e), done=True) - yield f"data:{err.model_dump_json()}\n\n".encode() - yield done_bytes - - finally: - if not task.done(): - task.cancel() + 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") @@ -95,7 +73,7 @@ class HttpService(BaseService): """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) -> None: + def run(self): """Start the Uvicorn server.""" super().run() cfg = C.service_config.http diff --git a/reme_ai/core/tool/mcp_tool.py b/reme_ai/core/tool/mcp_tool.py index 90f6ad7b..802e4138 100644 --- a/reme_ai/core/tool/mcp_tool.py +++ b/reme_ai/core/tool/mcp_tool.py @@ -48,7 +48,7 @@ class MCPTool(BaseOp): self._client = MCPClient(C.service_config.mcp_servers) def _build_tool_call(self) -> ToolCall: - tool_call_dict = C.mcp_server_tool_call_mapping[self.mcp_server] + 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 diff --git a/reme_ai/core/utils/__init__.py b/reme_ai/core/utils/__init__.py index 7767a13b..87e01299 100644 --- a/reme_ai/core/utils/__init__.py +++ b/reme_ai/core/utils/__init__.py @@ -2,7 +2,7 @@ from .cache_handler import CacheHandler from .case_converter import snake_to_camel, camel_to_snake -from .common_utils import run_coro_safely +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 @@ -19,6 +19,7 @@ __all__ = [ "snake_to_camel", "camel_to_snake", "run_coro_safely", + "execute_stream_task", "load_env", "HttpClient", "extract_content", diff --git a/reme_ai/core/utils/common_utils.py b/reme_ai/core/utils/common_utils.py index 5fe01a61..d370064d 100644 --- a/reme_ai/core/utils/common_utils.py +++ b/reme_ai/core/utils/common_utils.py @@ -1,9 +1,14 @@ """Common utility functions""" import asyncio -from collections.abc import Coroutine +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.""" @@ -18,3 +23,61 @@ def run_coro_safely(coro: Coroutine[Any, Any, Any]) -> Any | asyncio.Task[Any]: 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()