diff --git a/docs/work_memory/message_offload_ops.md b/docs/work_memory/message_offload_ops.md index d3965918..c05e22f7 100644 --- a/docs/work_memory/message_offload_ops.md +++ b/docs/work_memory/message_offload_ops.md @@ -96,7 +96,7 @@ When context grows too large, model performance degrades significantly—a pheno ### Usage Pattern For complete working examples of how to use MessageOffloadOp in practice, please refer to: -[test_message_offload_op.py](../../test_op/test_message_offload_op.py) +[test_message_offload_op.py](../../test/test_message_offload_op.py) This test file demonstrates: - **Compact mode**: How to configure and use compaction-only strategy diff --git a/docs/work_memory/message_reload_ops.md b/docs/work_memory/message_reload_ops.md index 6ee80136..c03f478d 100644 --- a/docs/work_memory/message_reload_ops.md +++ b/docs/work_memory/message_reload_ops.md @@ -114,7 +114,7 @@ Example: Reading `/workspace/context_store/tool_call_123.txt` with `offset=0` an ## Usage Pattern: Combining Grep and ReadFile For a complete working example of how to use these operations in practice, please refer to: -[test_agentic_retrieve_op.py](../../test_op/test_agentic_retrieve_op.py) +[test_agentic_retrieve_op.py](../../test/test_agentic_retrieve_op.py) This test file demonstrates: - How to configure the system prompt to guide AI in using Grep and ReadFile operations diff --git a/pyproject.toml b/pyproject.toml index ba7eb72d..47aa478d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,7 +57,7 @@ full = [ [tool.setuptools.packages.find] where = ["."] -include = ["reme_ai*"] +include = ["reme_ai*", "reme*"] exclude = ["test*", "cookbook*", "doc*", "library*", "dist*"] [tool.setuptools.package-data] diff --git a/reme/__init__.py b/reme/__init__.py index e69de29b..32f34911 100644 --- a/reme/__init__.py +++ b/reme/__init__.py @@ -0,0 +1,19 @@ +"""ReMe""" + +from . import agent +from . import config +from . import core +from . import tool +from . import workflow +from .reme_app import ReMeApp + +__all__ = [ + "agent", + "config", + "core", + "tool", + "workflow", + "ReMeApp", +] + +__version__ = "0.3.0.0a1" diff --git a/reme/agent/__init__.py b/reme/agent/__init__.py new file mode 100644 index 00000000..45fed6e8 --- /dev/null +++ b/reme/agent/__init__.py @@ -0,0 +1,7 @@ +"""A simple chatbot.""" + +from . import chat + +__all__ = [ + "chat", +] diff --git a/reme_ai/mem_agent/chat/__init__.py b/reme/agent/chat/__init__.py similarity index 57% rename from reme_ai/mem_agent/chat/__init__.py rename to reme/agent/chat/__init__.py index be2fc055..ed3049ef 100644 --- a/reme_ai/mem_agent/chat/__init__.py +++ b/reme/agent/chat/__init__.py @@ -1,11 +1,13 @@ """chat agent""" -from .remy_agent import ReMyAgent from .simple_chat import SimpleChat from .stream_chat import StreamChat +from ...core import R __all__ = [ - "ReMyAgent", "StreamChat", "SimpleChat", ] + +R.op.register("simple_chat")(SimpleChat) +R.op.register("stream_chat")(StreamChat) diff --git a/reme_ai/mem_agent/chat/simple_chat.py b/reme/agent/chat/simple_chat.py similarity index 93% rename from reme_ai/mem_agent/chat/simple_chat.py rename to reme/agent/chat/simple_chat.py index 8a71c7a8..36181346 100644 --- a/reme_ai/mem_agent/chat/simple_chat.py +++ b/reme/agent/chat/simple_chat.py @@ -2,14 +2,12 @@ from loguru import logger -from ...core.context import C from ...core.enumeration import Role -from ...core.op import BaseOp +from ...core.op import BaseTool from ...core.schema import Message, ToolCall -@C.register_op() -class SimpleChat(BaseOp): +class SimpleChat(BaseTool): """Simple chat agent that handles non-streaming conversations.""" def _build_tool_call(self) -> ToolCall: @@ -59,4 +57,4 @@ class SimpleChat(BaseOp): logger.info(f"messages={messages}") assistant_message = await self.llm.chat(messages=messages) logger.info(f"assistant_message={assistant_message.simple_dump()}") - self.output = assistant_message.content + return assistant_message.content diff --git a/reme_ai/mem_agent/chat/stream_chat.py b/reme/agent/chat/stream_chat.py similarity index 95% rename from reme_ai/mem_agent/chat/stream_chat.py rename to reme/agent/chat/stream_chat.py index 470e4647..f5cd7a5a 100644 --- a/reme_ai/mem_agent/chat/stream_chat.py +++ b/reme/agent/chat/stream_chat.py @@ -2,14 +2,12 @@ from loguru import logger -from ...core.context import C from ...core.enumeration import Role, ChunkEnum -from ...core.op import BaseOp +from ...core.op import BaseTool from ...core.schema import Message, ToolCall -@C.register_op() -class StreamChat(BaseOp): +class StreamChat(BaseTool): """Streaming chat agent that handles real-time conversation streaming.""" def _build_tool_call(self) -> ToolCall: diff --git a/reme/core/__init__.py b/reme/core/__init__.py index e69de29b..251c2433 100644 --- a/reme/core/__init__.py +++ b/reme/core/__init__.py @@ -0,0 +1,29 @@ +"""Core""" + +from . import context +from . import embedding +from . import enumeration +from . import flow +from . import llm +from . import op +from . import schema +from . import service +from . import token_counter +from . import utils +from . import vector_store +from .context import R + +__all__ = [ + "context", + "embedding", + "enumeration", + "flow", + "llm", + "op", + "schema", + "service", + "token_counter", + "utils", + "vector_store", + "R", +] diff --git a/reme/core/context/runtime_context.py b/reme/core/context/runtime_context.py index 44167ca0..d4c02050 100644 --- a/reme/core/context/runtime_context.py +++ b/reme/core/context/runtime_context.py @@ -30,7 +30,8 @@ class RuntimeContext(BaseContext): if context is None: return cls(**kwargs) else: - context.update(kwargs) + if kwargs: + context.update(kwargs) return context async def _enqueue(self, chunk: StreamChunk) -> None: diff --git a/reme/core/vector_store/chroma_vector_store.py b/reme/core/vector_store/chroma_vector_store.py index eaddd0e7..710ce73e 100644 --- a/reme/core/vector_store/chroma_vector_store.py +++ b/reme/core/vector_store/chroma_vector_store.py @@ -1,5 +1,6 @@ """ChromaDB vector store implementation for the ReMe framework.""" +from concurrent.futures import ThreadPoolExecutor from typing import Any from loguru import logger @@ -26,6 +27,7 @@ class ChromaVectorStore(BaseVectorStore): self, collection_name: str, embedding_model: BaseEmbeddingModel, + thread_pool: ThreadPoolExecutor, client: chromadb.ClientAPI | None = None, host: str | None = None, port: int | None = None, @@ -44,6 +46,7 @@ class ChromaVectorStore(BaseVectorStore): super().__init__( collection_name=collection_name, embedding_model=embedding_model, + thread_pool=thread_pool, **kwargs, ) diff --git a/reme/core/vector_store/es_vector_store.py b/reme/core/vector_store/es_vector_store.py index 49054afc..9a1019b8 100644 --- a/reme/core/vector_store/es_vector_store.py +++ b/reme/core/vector_store/es_vector_store.py @@ -4,6 +4,7 @@ This module provides an Elasticsearch-based vector store that implements the Bas interface for high-performance dense vector storage and retrieval. """ +from concurrent.futures import ThreadPoolExecutor from typing import Any from loguru import logger @@ -30,6 +31,7 @@ class ESVectorStore(BaseVectorStore): self, collection_name: str, embedding_model: BaseEmbeddingModel, + thread_pool: ThreadPoolExecutor, hosts: str | list[str] | None = None, basic_auth: tuple[str, str] | None = None, cloud_id: str | None = None, @@ -43,6 +45,7 @@ class ESVectorStore(BaseVectorStore): Args: collection_name: Name of the Elasticsearch index (converted to lowercase). embedding_model: Model instance used to generate vector embeddings. + thread_pool: ThreadPoolExecutor for running synchronous operations. hosts: Connection host(s) for the Elasticsearch cluster. basic_auth: Credentials for basic authentication. cloud_id: Deployment ID for Elastic Cloud. @@ -59,7 +62,12 @@ class ESVectorStore(BaseVectorStore): # Elasticsearch requires lowercase index names collection_name = collection_name.lower() - super().__init__(collection_name=collection_name, embedding_model=embedding_model, **kwargs) + super().__init__( + collection_name=collection_name, + embedding_model=embedding_model, + thread_pool=thread_pool, + **kwargs, + ) # Initialize AsyncElasticsearch client self.client = AsyncElasticsearch( diff --git a/reme/core/vector_store/local_vector_store.py b/reme/core/vector_store/local_vector_store.py index 24e1977c..a86e226a 100644 --- a/reme/core/vector_store/local_vector_store.py +++ b/reme/core/vector_store/local_vector_store.py @@ -1,6 +1,7 @@ """Local file system vector store implementation for ReMe.""" import json +from concurrent.futures import ThreadPoolExecutor from pathlib import Path from loguru import logger @@ -17,11 +18,17 @@ class LocalVectorStore(BaseVectorStore): self, collection_name: str, embedding_model: BaseEmbeddingModel, + thread_pool: ThreadPoolExecutor, 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) + super().__init__( + collection_name=collection_name, + embedding_model=embedding_model, + thread_pool=thread_pool, + **kwargs, + ) self.root_path = Path(root_path) self.collection_path = self.root_path / collection_name self.root_path.mkdir(parents=True, exist_ok=True) diff --git a/reme/core/vector_store/pgvector_store.py b/reme/core/vector_store/pgvector_store.py index 7d1b3401..279975d0 100644 --- a/reme/core/vector_store/pgvector_store.py +++ b/reme/core/vector_store/pgvector_store.py @@ -2,6 +2,7 @@ import json import re +from concurrent.futures import ThreadPoolExecutor from typing import Any from loguru import logger @@ -47,6 +48,7 @@ class PGVectorStore(BaseVectorStore): self, collection_name: str, embedding_model: BaseEmbeddingModel, + thread_pool: ThreadPoolExecutor, host: str = "localhost", port: int = 5432, database: str = "postgres", @@ -68,7 +70,12 @@ class PGVectorStore(BaseVectorStore): # Validate collection name to prevent SQL injection self._validate_table_name(collection_name) - super().__init__(collection_name=collection_name, embedding_model=embedding_model, **kwargs) + super().__init__( + collection_name=collection_name, + embedding_model=embedding_model, + thread_pool=thread_pool, + **kwargs, + ) self.dsn = dsn self.host = host diff --git a/reme/core/vector_store/qdrant_vector_store.py b/reme/core/vector_store/qdrant_vector_store.py index 5ca74cd0..93ccee70 100644 --- a/reme/core/vector_store/qdrant_vector_store.py +++ b/reme/core/vector_store/qdrant_vector_store.py @@ -1,5 +1,6 @@ """Qdrant vector store implementation for the ReMe project.""" +from concurrent.futures import ThreadPoolExecutor from typing import Any from loguru import logger @@ -42,6 +43,7 @@ class QdrantVectorStore(BaseVectorStore): self, collection_name: str, embedding_model: BaseEmbeddingModel, + thread_pool: ThreadPoolExecutor, host: str | None = None, port: int = 6333, path: str | None = None, @@ -59,6 +61,7 @@ class QdrantVectorStore(BaseVectorStore): Args: collection_name: Name of the collection. embedding_model: Model used for generating vector embeddings. + thread_pool: ThreadPoolExecutor for running synchronous operations. host: Server host address. port: HTTP port for the server. path: Local storage path for on-disk/in-memory mode. @@ -76,7 +79,14 @@ class QdrantVectorStore(BaseVectorStore): "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) + super().__init__( + collection_name=collection_name, + embedding_model=embedding_model, + thread_pool=thread_pool, + **kwargs, + ) + + client_kwargs = {k: v for k, v in kwargs.items() if k != "thread_pool"} self.client = AsyncQdrantClient( host=host, @@ -87,7 +97,7 @@ class QdrantVectorStore(BaseVectorStore): https=https, grpc_port=grpc_port, prefer_grpc=prefer_grpc, - **kwargs, + **client_kwargs, ) self.is_local = path is not None diff --git a/reme/tool/__init__.py b/reme/tool/__init__.py new file mode 100644 index 00000000..60082176 --- /dev/null +++ b/reme/tool/__init__.py @@ -0,0 +1,9 @@ +"""Tool""" + +from . import execute +from . import search + +__all__ = [ + "execute", + "search", +] diff --git a/reme_ai/tool/execute/__init__.py b/reme/tool/execute/__init__.py similarity index 64% rename from reme_ai/tool/execute/__init__.py rename to reme/tool/execute/__init__.py index f78a73ac..28e3ecf3 100644 --- a/reme_ai/tool/execute/__init__.py +++ b/reme/tool/execute/__init__.py @@ -2,8 +2,12 @@ from .execute_code import ExecuteCode from .execute_shell import ExecuteShell +from ...core import R __all__ = [ "ExecuteCode", "ExecuteShell", ] + +R.op.register()(ExecuteCode) +R.op.register()(ExecuteShell) diff --git a/reme_ai/tool/execute/execute_code.py b/reme/tool/execute/execute_code.py similarity index 87% rename from reme_ai/tool/execute/execute_code.py rename to reme/tool/execute/execute_code.py index ea259487..6cb77c85 100644 --- a/reme_ai/tool/execute/execute_code.py +++ b/reme/tool/execute/execute_code.py @@ -4,15 +4,13 @@ This module provides an operation that can execute Python code strings and return the output or error messages. """ -from ...core.context import C -from ...core.op import BaseOp +from ...core.op import BaseTool from ...core.schema import ToolCall from ...core.utils import exec_code -@C.register_op() -class ExecuteCode(BaseOp): +class ExecuteCode(BaseTool): """Operation for executing Python code dynamically. This operation takes Python code as input, executes it in a safe context, @@ -40,4 +38,4 @@ class ExecuteCode(BaseOp): self.execute_sync() def execute_sync(self): - self.output = exec_code(self.context.code) + return exec_code(self.context.code) diff --git a/reme_ai/tool/execute/execute_code.yaml b/reme/tool/execute/execute_code.yaml similarity index 100% rename from reme_ai/tool/execute/execute_code.yaml rename to reme/tool/execute/execute_code.yaml diff --git a/reme_ai/tool/execute/execute_shell.py b/reme/tool/execute/execute_shell.py similarity index 90% rename from reme_ai/tool/execute/execute_shell.py rename to reme/tool/execute/execute_shell.py index 6e244ddb..de862602 100644 --- a/reme_ai/tool/execute/execute_shell.py +++ b/reme/tool/execute/execute_shell.py @@ -4,15 +4,13 @@ This module provides an operation that can execute shell commands asynchronously and return the output, error, and exit code. """ -from ...core.context import C -from ...core.op import BaseOp +from ...core.op import BaseTool from ...core.schema import ToolCall from ...core.utils import run_shell_command -@C.register_op() -class ExecuteShell(BaseOp): +class ExecuteShell(BaseTool): """Operation for executing shell commands asynchronously. This operation takes a shell command as input, executes it asynchronously, @@ -46,4 +44,4 @@ class ExecuteShell(BaseOp): f"Exit Code: {return_code if return_code is not None else '(none)'}", ] - self.output = "\n".join(result_parts) + return "\n".join(result_parts) diff --git a/reme_ai/tool/execute/execute_shell.yaml b/reme/tool/execute/execute_shell.yaml similarity index 100% rename from reme_ai/tool/execute/execute_shell.yaml rename to reme/tool/execute/execute_shell.yaml diff --git a/reme_ai/tool/search/__init__.py b/reme/tool/search/__init__.py similarity index 65% rename from reme_ai/tool/search/__init__.py rename to reme/tool/search/__init__.py index 5a73dc2e..a9e87ef5 100644 --- a/reme_ai/tool/search/__init__.py +++ b/reme/tool/search/__init__.py @@ -3,9 +3,14 @@ from .dashscope_search import DashscopeSearch from .mock_search import MockSearch from .tavily_search import TavilySearch +from ...core import R __all__ = [ "DashscopeSearch", "MockSearch", "TavilySearch", ] + +R.op.register()(DashscopeSearch) +R.op.register()(MockSearch) +R.op.register()(TavilySearch) diff --git a/reme_ai/tool/search/dashscope_search.py b/reme/tool/search/dashscope_search.py similarity index 92% rename from reme_ai/tool/search/dashscope_search.py rename to reme/tool/search/dashscope_search.py index 19bd8104..f0e0ae9d 100644 --- a/reme_ai/tool/search/dashscope_search.py +++ b/reme/tool/search/dashscope_search.py @@ -9,13 +9,11 @@ from typing import Literal from loguru import logger -from ...core.context import C -from ...core.op import BaseOp +from ...core.op import BaseTool from ...core.schema import ToolCall -@C.register_op() -class DashscopeSearch(BaseOp): +class DashscopeSearch(BaseTool): """Operation for performing web searches using Dashscope API. This operation uses Alibaba Cloud's Dashscope service to search the web @@ -61,8 +59,7 @@ class DashscopeSearch(BaseOp): if self.enable_cache: cached_result = self.cache.load(query) if cached_result: - self.output = cached_result["response_content"] - return + return cached_result["response_content"] if self.enable_role_prompt: user_query = self.prompt_format("role_prompt", query=query) @@ -108,4 +105,4 @@ class DashscopeSearch(BaseOp): if self.enable_cache: self.cache.save(query, final_result, expire_hours=self.cache_expire_hours) - self.output = final_result["response_content"] + return final_result["response_content"] diff --git a/reme_ai/tool/search/dashscope_search.yaml b/reme/tool/search/dashscope_search.yaml similarity index 100% rename from reme_ai/tool/search/dashscope_search.yaml rename to reme/tool/search/dashscope_search.yaml diff --git a/reme_ai/tool/search/mock_search.py b/reme/tool/search/mock_search.py similarity index 91% rename from reme_ai/tool/search/mock_search.py rename to reme/tool/search/mock_search.py index 187463dc..ed695b0c 100644 --- a/reme_ai/tool/search/mock_search.py +++ b/reme/tool/search/mock_search.py @@ -9,15 +9,13 @@ import random from loguru import logger -from ...core.context import C from ...core.enumeration import Role -from ...core.op import BaseOp +from ...core.op import BaseTool from ...core.schema import ToolCall, Message from ...core.utils import extract_content -@C.register_op() -class MockSearch(BaseOp): +class MockSearch(BaseTool): """Operation for generating mock search results. This operation generates simulated search results using an LLM, @@ -61,4 +59,4 @@ class MockSearch(BaseOp): return extract_content(message.content, "json") search_results: str = await self.llm.chat(messages=messages, callback_fn=callback_fn) - self.output = json.dumps(search_results, ensure_ascii=False, indent=2) + return json.dumps(search_results, ensure_ascii=False, indent=2) diff --git a/reme_ai/tool/search/mock_search.yaml b/reme/tool/search/mock_search.yaml similarity index 100% rename from reme_ai/tool/search/mock_search.yaml rename to reme/tool/search/mock_search.yaml diff --git a/reme_ai/tool/search/tavily_search.py b/reme/tool/search/tavily_search.py similarity index 90% rename from reme_ai/tool/search/tavily_search.py rename to reme/tool/search/tavily_search.py index 5c194bdc..bfe29879 100644 --- a/reme_ai/tool/search/tavily_search.py +++ b/reme/tool/search/tavily_search.py @@ -9,13 +9,11 @@ import os from loguru import logger -from ...core.context import C -from ...core.op import BaseOp +from ...core.op import BaseTool from ...core.schema import ToolCall -@C.register_op() -class TavilySearch(BaseOp): +class TavilySearch(BaseTool): """Operation for performing web searches using Tavily API. This operation uses the Tavily search service to find web content @@ -73,8 +71,7 @@ class TavilySearch(BaseOp): if self.enable_cache: cached_result = self.cache.load(query) if cached_result: - self.output = json.dumps(cached_result, ensure_ascii=False, indent=2) - return + return json.dumps(cached_result, ensure_ascii=False, indent=2) response = await self.client.search(query=query) logger.info(f"tavily_search response={response}") @@ -88,8 +85,7 @@ class TavilySearch(BaseOp): if self.enable_cache and final_result: self.cache.save(query, final_result, expire_hours=self.cache_expire_hours) - self.output = json.dumps(final_result, ensure_ascii=False, indent=2) - return + return json.dumps(final_result, ensure_ascii=False, indent=2) url_info_dict = {item["url"]: item for item in response["results"]} response_extract = await self.client.extract(urls=[item["url"] for item in response["results"]]) @@ -116,4 +112,4 @@ class TavilySearch(BaseOp): if self.enable_cache and final_result: self.cache.save(query, final_result, expire_hours=self.cache_expire_hours) - self.output = json.dumps(final_result, ensure_ascii=False, indent=2) + return json.dumps(final_result, ensure_ascii=False, indent=2) diff --git a/reme_ai/tool/search/tavily_search.yaml b/reme/tool/search/tavily_search.yaml similarity index 100% rename from reme_ai/tool/search/tavily_search.yaml rename to reme/tool/search/tavily_search.yaml diff --git a/reme/workflow/__init__.py b/reme/workflow/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/reme_ai/mem_agent/chat/remy_agent.py b/reme_ai/mem_agent/chat/remy_agent.py deleted file mode 100644 index 1eaa9ac7..00000000 --- a/reme_ai/mem_agent/chat/remy_agent.py +++ /dev/null @@ -1,51 +0,0 @@ -"""ReMy agent with identity and meta memory capabilities.""" - -from typing import List - -from ..base_memory_agent import BaseMemoryAgent -from ...core.context import C -from ...core.enumeration import Role -from ...core.schema import Message -from ...core.utils import get_now_time - - -@C.register_op() -class ReMyAgent(BaseMemoryAgent): - """Memory agent with identity awareness and meta memory retrieval.""" - - def __init__(self, enable_tool_memory: bool = True, enable_identity_memory: bool = True, **kwargs): - """Initialize ReMy agent with memory options.""" - super().__init__(**kwargs) - self.enable_tool_memory = enable_tool_memory - self.enable_identity_memory = enable_identity_memory - - @staticmethod - async def _read_identity_memory() -> str: - """Read and return identity memory as string.""" - from ...mem_tool import ReadIdentityMemory - - op = ReadIdentityMemory() - await op.call() - return str(op.output) - - async def _read_meta_memories(self) -> str: - """Read and return meta memories as string.""" - from ...mem_tool import ReadMetaMemory - - op = ReadMetaMemory( - enable_tool_memory=self.enable_tool_memory, - enable_identity_memory=self.enable_identity_memory, - ) - await op.call() - return str(op.output) - - async def build_messages(self) -> List[Message]: - """Build messages with system prompt and user messages.""" - system_prompt = self.prompt_format( - prompt_name="system_prompt", - now_time=get_now_time(), - identity_memory=await self._read_identity_memory(), - meta_memory_info=await self._read_meta_memories(), - ) - - return [Message(role=Role.SYSTEM, content=system_prompt)] + self.get_messages() diff --git a/reme_ai/mem_agent/chat/remy_agent.yaml b/reme_ai/mem_agent/chat/remy_agent.yaml deleted file mode 100644 index 5498b3cc..00000000 --- a/reme_ai/mem_agent/chat/remy_agent.yaml +++ /dev/null @@ -1,36 +0,0 @@ -tool: | - Conversational AI assistant with integrated memory capabilities. - Use this tool to engage in natural conversations with users while leveraging - stored identity and memory context. The agent can access historical information, - user preferences, and procedural knowledge through its memory system, and can - use various tools to accomplish tasks and answer questions. - -system_prompt: | - You are ReMy, an intelligent AI assistant with memory capabilities. - - ## Current Time - {now_time} - - ## Self-Awareness - {identity_memory} - - ## Available Meta Memories - Format: "- (): " - {meta_memory_info} - - ## Guiding Principles - 1. **Be Helpful and Accurate**: Provide clear and correct information. - 2. **Use Memory Wisely**: Retrieve relevant memories when they can improve your response. - 3. **Use Tools Appropriately**: Select the right tool for each task. - 4. **Stay Conversational**: Maintain a natural and friendly tone. - 5. **Seek Clarification**: Ask questions if the user’s intent is unclear. - 6. **Acknowledge Limitations**: Be honest about what you can and cannot do. - - ## How to Use the Memory Retrieval Tool - When using `vector_retrieve_memory` to search memories: - - Choose an appropriate `memory_type` and `memory_target` from the "Available Meta Memories" list above. - - Formulate a clear and specific query based on the information you need. - - **Important**: When retrieving tool-related memories (`memory_type` is "tool"), the query must use the tool’s exact name (not a description or a question). - - If retrieval results include a `ref_memory_id` and you need more details, use `read_history_memory` with the `ref_memory_id` as the `memory_id` parameter. - - If the initial retrieval yields no results, try rephrasing your query or using a different memory type. - - You may generate multiple queries with different phrasings or perspectives for the same memory type/target. diff --git a/test_op/test_agentic_retrieve_op.py b/test/test_agentic_retrieve_op.py similarity index 100% rename from test_op/test_agentic_retrieve_op.py rename to test/test_agentic_retrieve_op.py diff --git a/test_op/test_message_compact_op.py b/test/test_message_compact_op.py similarity index 100% rename from test_op/test_message_compact_op.py rename to test/test_message_compact_op.py diff --git a/test_op/test_message_compress_op.py b/test/test_message_compress_op.py similarity index 100% rename from test_op/test_message_compress_op.py rename to test/test_message_compress_op.py diff --git a/test_op/test_message_offload_op.py b/test/test_message_offload_op.py similarity index 100% rename from test_op/test_message_offload_op.py rename to test/test_message_offload_op.py diff --git a/test/test_op_composition.py b/test/test_op_composition.py deleted file mode 100644 index 8d32b51a..00000000 --- a/test/test_op_composition.py +++ /dev/null @@ -1,325 +0,0 @@ -""" -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) diff --git a/test/test_reme.py b/test/test_reme.py index c9e6843e..3279da2a 100644 --- a/test/test_reme.py +++ b/test/test_reme.py @@ -2,8 +2,9 @@ import asyncio -from reme_ai.core.schema import VectorNode, MemoryNode -from reme_ai.reme import ReMe +from reme.reme import ReMe + +from reme.core.schema import VectorNode, MemoryNode reme = ReMe( vector_store={"collection_name": "reme"}, diff --git a/test/mcp_servers_demo.json b/tests/mcp_servers_demo.json similarity index 100% rename from test/mcp_servers_demo.json rename to tests/mcp_servers_demo.json diff --git a/test/test_base_context.py b/tests/test_base_context.py similarity index 97% rename from test/test_base_context.py rename to tests/test_base_context.py index 316a6796..61a355c0 100644 --- a/test/test_base_context.py +++ b/tests/test_base_context.py @@ -4,7 +4,8 @@ Ensures attribute-style and dict-style access work interchangeably. """ import pickle -from reme_ai.core.context import BaseContext + +from reme.core.context import BaseContext def test_attribute_access(): diff --git a/test/test_cache_handler.py b/tests/test_cache_handler.py similarity index 98% rename from test/test_cache_handler.py rename to tests/test_cache_handler.py index ddcac86f..f1facc0f 100644 --- a/test/test_cache_handler.py +++ b/tests/test_cache_handler.py @@ -9,7 +9,7 @@ from pathlib import Path import pandas as pd from loguru import logger -from reme_ai.core.utils.cache_handler import CacheHandler +from reme.core.utils.cache_handler import CacheHandler def run_tests(): diff --git a/test/test_embedding.py b/tests/test_embedding.py similarity index 98% rename from test/test_embedding.py rename to tests/test_embedding.py index d1c404f5..867037f4 100644 --- a/test/test_embedding.py +++ b/tests/test_embedding.py @@ -14,16 +14,16 @@ Usage: # flake8: noqa: E402 # pylint: disable=C0413 -import asyncio import argparse +import asyncio from typing import Type, List -from reme_ai.core.utils import load_env +from reme.core.utils import load_env load_env() -from reme_ai.core.embedding import OpenAIEmbeddingModel, BaseEmbeddingModel -from reme_ai.core.schema import VectorNode +from reme.core.embedding import OpenAIEmbeddingModel, BaseEmbeddingModel +from reme.core.schema import VectorNode def get_embedding_model(model_class: Type[BaseEmbeddingModel]) -> BaseEmbeddingModel: diff --git a/test/test_embedding_sync.py b/tests/test_embedding_sync.py similarity index 98% rename from test/test_embedding_sync.py rename to tests/test_embedding_sync.py index 361a42b3..34ddd350 100644 --- a/test/test_embedding_sync.py +++ b/tests/test_embedding_sync.py @@ -17,12 +17,12 @@ Usage: import argparse from typing import Type, List -from reme_ai.core.utils import load_env +from reme.core.utils import load_env load_env() -from reme_ai.core.embedding import OpenAIEmbeddingModelSync, BaseEmbeddingModel -from reme_ai.core.schema import VectorNode +from reme.core.embedding import OpenAIEmbeddingModelSync, BaseEmbeddingModel +from reme.core.schema import VectorNode def get_embedding_model(model_class: Type[BaseEmbeddingModel]) -> BaseEmbeddingModel: diff --git a/test/test_llm.py b/tests/test_llm.py similarity index 98% rename from test/test_llm.py rename to tests/test_llm.py index 12c6eca3..5e1755fe 100644 --- a/test/test_llm.py +++ b/tests/test_llm.py @@ -14,17 +14,17 @@ Usage: # flake8: noqa: E402 # pylint: disable=C0413 -import asyncio import argparse +import asyncio from typing import Type -from reme_ai.core.utils import load_env +from reme.core.utils import load_env load_env() -from reme_ai.core.llm import OpenAILLM, LiteLLM, BaseLLM -from reme_ai.core.schema import Message, ToolCall -from reme_ai.core.enumeration import Role, ChunkEnum +from reme.core.llm import OpenAILLM, LiteLLM, BaseLLM +from reme.core.schema import Message, ToolCall +from reme.core.enumeration import Role, ChunkEnum def get_llm(llm_class: Type[BaseLLM]) -> BaseLLM: diff --git a/test/test_llm_sync.py b/tests/test_llm_sync.py similarity index 98% rename from test/test_llm_sync.py rename to tests/test_llm_sync.py index 98751f87..0b3c8cce 100644 --- a/test/test_llm_sync.py +++ b/tests/test_llm_sync.py @@ -17,13 +17,13 @@ Usage: import argparse from typing import Type -from reme_ai.core.utils import load_env +from reme.core.utils import load_env load_env() -from reme_ai.core.llm import OpenAILLMSync, LiteLLMSync, BaseLLM -from reme_ai.core.schema import Message, ToolCall -from reme_ai.core.enumeration import Role, ChunkEnum +from reme.core.llm import OpenAILLMSync, LiteLLMSync, BaseLLM +from reme.core.schema import Message, ToolCall +from reme.core.enumeration import Role, ChunkEnum def get_llm(llm_class: Type[BaseLLM]) -> BaseLLM: diff --git a/test/test_logo.py b/tests/test_logo.py similarity index 61% rename from test/test_logo.py rename to tests/test_logo.py index eeede81e..12b757c4 100644 --- a/test/test_logo.py +++ b/tests/test_logo.py @@ -1,9 +1,9 @@ """test logo""" -from reme_ai.core.schema import ServiceConfig, MCPConfig +from reme.core.schema import ServiceConfig, MCPConfig if __name__ == "__main__": - from reme_ai.core.utils import print_logo + from reme.core.utils import print_logo c = ServiceConfig(app_name="reme", backend="mcp", mcp=MCPConfig(transport="sse")) print_logo(service_config=c) diff --git a/test/test_mcp_client.py b/tests/test_mcp_client.py similarity index 99% rename from test/test_mcp_client.py rename to tests/test_mcp_client.py index d2fef40e..a369d7cf 100644 --- a/test/test_mcp_client.py +++ b/tests/test_mcp_client.py @@ -5,7 +5,7 @@ import asyncio import json -from reme_ai.core.utils import MCPClient +from reme.core.utils import MCPClient async def main(): diff --git a/test/test_mcp_server.py b/tests/test_mcp_server.py similarity index 97% rename from test/test_mcp_server.py rename to tests/test_mcp_server.py index 67f9c542..f5744f99 100644 --- a/test/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -5,8 +5,8 @@ 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 +from reme.core.schema import ToolCall +from reme.core.utils import create_pydantic_model mcp = FastMCP("DynamicSchemaServer", port=8010) diff --git a/test/test_memory_vector_conversion.py b/tests/test_memory_vector_conversion.py similarity index 100% rename from test/test_memory_vector_conversion.py rename to tests/test_memory_vector_conversion.py diff --git a/test/test_message.py b/tests/test_message.py similarity index 98% rename from test/test_message.py rename to tests/test_message.py index 141174c5..7b7e44bc 100644 --- a/test/test_message.py +++ b/tests/test_message.py @@ -4,8 +4,8 @@ import unittest from mcp.types import Tool -from reme_ai.core.enumeration import Role -from reme_ai.core.schema import ToolAttr, ToolCall, ContentBlock, Message +from reme.core.enumeration import Role +from reme.core.schema import ToolAttr, ToolCall, ContentBlock, Message class TestModelDefinitions(unittest.TestCase): diff --git a/test/test_timer.py b/tests/test_timer.py similarity index 97% rename from test/test_timer.py rename to tests/test_timer.py index c9714e38..98b1d816 100644 --- a/test/test_timer.py +++ b/tests/test_timer.py @@ -7,7 +7,7 @@ import time from loguru import logger -from reme_ai.core.utils import timer +from reme.core.utils import timer @timer diff --git a/test/test_token_counter.py b/tests/test_token_counter.py similarity index 99% rename from test/test_token_counter.py rename to tests/test_token_counter.py index 3c44a298..2fd8bdfd 100644 --- a/test/test_token_counter.py +++ b/tests/test_token_counter.py @@ -14,9 +14,9 @@ Usage: 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 +from reme.core.enumeration import Role +from reme.core.schema import Message, ToolCall +from reme.core.token_counter import BaseTokenCounter, OpenAITokenCounter, HFTokenCounter def get_token_counter(counter_class: Type[BaseTokenCounter], **kwargs) -> BaseTokenCounter: diff --git a/test/test_tool.py b/tests/test_tool.py similarity index 75% rename from test/test_tool.py rename to tests/test_tool.py index 9db5a3ed..02346d0a 100644 --- a/test/test_tool.py +++ b/tests/test_tool.py @@ -8,9 +8,9 @@ search tools (Dashscope, Mock, Tavily) and execution tools (Code, Shell). import asyncio -from reme_ai.reme import ReMe +from reme.reme_app import ReMeApp -ReMe() +app = ReMeApp() def test_search(): @@ -19,7 +19,7 @@ def test_search(): Tests DashscopeSearch, MockSearch, and TavilySearch operations with a sample query to verify they work correctly. """ - from reme_ai.tool.search import DashscopeSearch, MockSearch, TavilySearch + from reme.tool.search import DashscopeSearch, MockSearch, TavilySearch query = "今天杭州的天气如何？" @@ -32,8 +32,8 @@ def test_search(): print(f"Testing {op.__class__.__name__}") print("=" * 60) print(f"Query: {query}") - asyncio.run(op.call(query=query)) - print(f"Output:\n{op.output}") + output = asyncio.run(op.call(query=query, service_context=app.service_context)) + print(f"Output:\n{output}") def test_execute(): @@ -43,7 +43,7 @@ def test_execute(): including successful execution, syntax errors, runtime errors, and invalid commands to verify error handling. """ - from reme_ai.tool.execute import ExecuteCode, ExecuteShell + from reme.tool.execute import ExecuteCode, ExecuteShell # Test ExecuteCode print("\n" + "=" * 60) @@ -53,8 +53,8 @@ def test_execute(): op = ExecuteCode() code_to_execute = "print('hello world')" print(f"Executing Python code: {code_to_execute}") - asyncio.run(op.call(code=code_to_execute)) - print(f"Output:\n{op.output}") + output = asyncio.run(op.call(code=code_to_execute)) + print(f"Output:\n{output}") # Test ExecuteCode with more complex code print("\n" + "=" * 60) @@ -64,8 +64,8 @@ def test_execute(): op = ExecuteCode() code_to_execute = "result = sum(range(1, 11))\nprint(f'Sum of 1-10: {result}')" print(f"Executing Python code:\n{code_to_execute}") - asyncio.run(op.call(code=code_to_execute)) - print(f"Output:\n{op.output}") + output = asyncio.run(op.call(code=code_to_execute)) + print(f"Output:\n{output}") # Test ExecuteShell print("\n" + "=" * 60) @@ -75,8 +75,8 @@ def test_execute(): op = ExecuteShell() command = "ls" print(f"Executing shell command: {command}") - asyncio.run(op.call(command=command)) - print(f"Output:\n{op.output}") + output = asyncio.run(op.call(command=command)) + print(f"Output:\n{output}") # Test ExecuteShell with echo print("\n" + "=" * 60) @@ -86,8 +86,8 @@ def test_execute(): op = ExecuteShell() command = "echo 'Hello from shell!'" print(f"Executing shell command: {command}") - asyncio.run(op.call(command=command)) - print(f"Output:\n{op.output}") + output = asyncio.run(op.call(command=command)) + print(f"Output:\n{output}") # Test ExecuteCode with error (syntax error) print("\n" + "=" * 60) @@ -97,8 +97,8 @@ def test_execute(): op = ExecuteCode() code_to_execute = "print('missing closing quote)" print(f"Executing Python code with syntax error:\n{code_to_execute}") - asyncio.run(op.call(code=code_to_execute)) - print(f"Output:\n{op.output}") + output = asyncio.run(op.call(code=code_to_execute)) + print(f"Output:\n{output}") # Test ExecuteCode with runtime error print("\n" + "=" * 60) @@ -108,8 +108,8 @@ def test_execute(): op = ExecuteCode() code_to_execute = "x = 1 / 0" print(f"Executing Python code with runtime error:\n{code_to_execute}") - asyncio.run(op.call(code=code_to_execute)) - print(f"Output:\n{op.output}") + output = asyncio.run(op.call(code=code_to_execute)) + print(f"Output:\n{output}") # Test ExecuteCode with undefined variable print("\n" + "=" * 60) @@ -119,8 +119,8 @@ def test_execute(): op = ExecuteCode() code_to_execute = "print(undefined_variable)" print(f"Executing Python code with undefined variable:\n{code_to_execute}") - asyncio.run(op.call(code=code_to_execute)) - print(f"Output:\n{op.output}") + output = asyncio.run(op.call(code=code_to_execute)) + print(f"Output:\n{output}") # Test ExecuteShell with invalid command print("\n" + "=" * 60) @@ -130,8 +130,8 @@ def test_execute(): op = ExecuteShell() command = "this_command_does_not_exist" print(f"Executing invalid shell command: {command}") - asyncio.run(op.call(command=command)) - print(f"Output:\n{op.output}") + output = asyncio.run(op.call(command=command)) + print(f"Output:\n{output}") # Test ExecuteShell with command that returns non-zero exit code print("\n" + "=" * 60) @@ -141,8 +141,8 @@ def test_execute(): op = ExecuteShell() command = "ls /nonexistent_directory_12345" print(f"Executing shell command that should fail: {command}") - asyncio.run(op.call(command=command)) - print(f"Output:\n{op.output}") + output = asyncio.run(op.call(command=command)) + print(f"Output:\n{output}") print("\n" + "=" * 60) print("All tests completed!") @@ -155,11 +155,11 @@ def test_simple_chat(): Tests the SimpleChat agent with a basic query to verify it can process and respond to user input. """ - from reme_ai.mem_agent.chat import SimpleChat + from reme.agent.chat import SimpleChat op = SimpleChat() - asyncio.run(op.call(query="你好")) - print(op.output) + output = asyncio.run(op.call(query="你好", service_context=app.service_context)) + print(output) async def test_stream_chat(): @@ -168,13 +168,13 @@ async def test_stream_chat(): Tests the StreamChat agent with a query to verify it can process and stream responses in real-time using async operations. """ - from reme_ai.mem_agent.chat import StreamChat - from reme_ai.core.utils import execute_stream_task - from reme_ai.core.context import RuntimeContext + from reme.agent.chat import StreamChat + from reme.core.utils import execute_stream_task + from reme.core.context import RuntimeContext from asyncio import Queue op = StreamChat() - context = RuntimeContext(query="你好，详细介绍一下自己", stream_queue=Queue()) + context = RuntimeContext(query="你好，详细介绍一下自己", stream_queue=Queue(), service_context=app.service_context) async def task(): await op.call(context) @@ -192,5 +192,5 @@ async def test_stream_chat(): if __name__ == "__main__": # test_search() # test_execute() - # test_simple_chat() + test_simple_chat() asyncio.run(test_stream_chat()) diff --git a/test/test_tool_call.py b/tests/test_tool_call.py similarity index 99% rename from test/test_tool_call.py rename to tests/test_tool_call.py index 30c0a37e..7be3079e 100644 --- a/test/test_tool_call.py +++ b/tests/test_tool_call.py @@ -2,7 +2,7 @@ import json -from reme_ai.core.schema.tool_call import ToolCall +from reme.core.schema.tool_call import ToolCall def test_simple_schema(): diff --git a/test/test_vector_store.py b/tests/test_vector_store.py similarity index 98% rename from test/test_vector_store.py rename to tests/test_vector_store.py index 00927d0c..14ed1c1d 100644 --- a/test/test_vector_store.py +++ b/tests/test_vector_store.py @@ -18,14 +18,16 @@ Usage: import argparse import asyncio import shutil +from concurrent.futures import ThreadPoolExecutor 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 ( +from reme.core.embedding import OpenAIEmbeddingModel +from reme.core.schema import VectorNode +from reme.core.utils import load_env +from reme.core.vector_store import ( BaseVectorStore, ChromaVectorStore, LocalVectorStore, @@ -34,6 +36,7 @@ from reme_ai.core.vector_store import ( QdrantVectorStore, ) +load_env() # ==================== Configuration ==================== @@ -199,7 +202,7 @@ def create_vector_store(store_type: str, collection_name: str) -> BaseVectorStor """Create a vector store instance based on type. Args: - store_type: Type of vector store ("local", "es", or "qdrant") + store_type: Type of vector store ("local", "es", "pgvector", "qdrant", or "chroma") collection_name: Name of the collection Returns: @@ -213,16 +216,21 @@ def create_vector_store(store_type: str, collection_name: str) -> BaseVectorStor dimensions=config.EMBEDDING_DIMENSIONS, ) + # Create thread pool executor for vector stores + thread_pool = ThreadPoolExecutor(max_workers=4) + if store_type == "local": return LocalVectorStore( collection_name=collection_name, embedding_model=embedding_model, + thread_pool=thread_pool, root_path=config.LOCAL_ROOT_PATH, ) elif store_type == "es": return ESVectorStore( collection_name=collection_name, embedding_model=embedding_model, + thread_pool=thread_pool, hosts=config.ES_HOSTS, basic_auth=config.ES_BASIC_AUTH, ) @@ -230,6 +238,7 @@ def create_vector_store(store_type: str, collection_name: str) -> BaseVectorStor return QdrantVectorStore( collection_name=collection_name, embedding_model=embedding_model, + thread_pool=thread_pool, path=config.QDRANT_PATH, host=config.QDRANT_HOST, port=config.QDRANT_PORT, @@ -242,6 +251,7 @@ def create_vector_store(store_type: str, collection_name: str) -> BaseVectorStor return PGVectorStore( collection_name=collection_name, embedding_model=embedding_model, + thread_pool=thread_pool, dsn=config.PG_DSN, min_size=config.PG_MIN_SIZE, max_size=config.PG_MAX_SIZE, @@ -252,6 +262,7 @@ def create_vector_store(store_type: str, collection_name: str) -> BaseVectorStor return ChromaVectorStore( collection_name=collection_name, embedding_model=embedding_model, + thread_pool=thread_pool, path=config.CHROMA_PATH, host=config.CHROMA_HOST, port=config.CHROMA_PORT, @@ -1267,7 +1278,7 @@ async def test_range_query_filters(store: BaseVectorStore, _store_name: str): for r in results_3: ts = r.metadata.get("timestamp") category = r.metadata.get("category") - assert ts >= start_time and ts <= end_time, "Timestamp should be in range" + assert start_time <= ts <= end_time, "Timestamp should be in range" assert category == "tech", f"Category should be 'tech', got '{category}'" logger.debug(f" Node {r.vector_id}: timestamp={ts}, category={category}") @@ -1290,8 +1301,8 @@ async def test_range_query_filters(store: BaseVectorStore, _store_name: str): for r in results_4: ts = r.metadata.get("timestamp") rating = r.metadata.get("rating") - assert ts >= base_timestamp + 8000 and ts <= base_timestamp + 12000, "Timestamp out of range" - assert rating >= 65 and rating <= 75, f"Rating {rating} out of range [65, 75]" + assert base_timestamp + 8000 <= ts <= base_timestamp + 12000, "Timestamp out of range" + assert 65 <= rating <= 75, f"Rating {rating} out of range [65, 75]" logger.debug(f" Node {r.vector_id}: timestamp={ts}, rating={rating}") # Expected: nodes 8-12 (5 nodes) with overlapping ranges @@ -1311,7 +1322,7 @@ async def test_range_query_filters(store: BaseVectorStore, _store_name: str): # Verify rating range in list results for r in results_5: rating = r.metadata.get("rating") - assert rating >= 60 and rating <= 70, f"Rating {rating} should be in range [60, 70]" + assert 60 <= rating <= 70, f"Rating {rating} should be in range [60, 70]" logger.info("✓ Range query in list operation validated") @@ -1349,7 +1360,7 @@ async def test_range_query_filters(store: BaseVectorStore, _store_name: str): rating1 = results_7[i].metadata.get("rating") rating2 = results_7[i + 1].metadata.get("rating") assert rating1 >= rating2, f"Results not sorted: {rating1} < {rating2}" - assert rating1 >= 60 and rating1 <= 80, "Rating out of range" + assert 60 <= rating1 <= 80, "Rating out of range" logger.info("✓ Range query with sorting validated") @@ -1419,7 +1430,6 @@ async def test_string_range_queries(store: BaseVectorStore, store_name: str): logger.info(f"Test 1 - String date range ['2024-02-01', '2024-03-15']: {len(results)} results") # Verify all results are within range - expected_dates = ["2024-02-01", "2024-02-15", "2024-03-01", "2024-03-15"] for r in results: date = r.metadata.get("date") assert date >= "2024-02-01", f"Date {date} should be >= '2024-02-01'" @@ -1453,16 +1463,15 @@ async def test_sql_injection_protection(store: BaseVectorStore, store_name: str) # Test 1: Invalid collection name (SQL injection attempt) try: - from reme_ai.core.vector_store import PGVectorStore - from reme_ai.core.embedding import OpenAIEmbeddingModel - embedding_model = OpenAIEmbeddingModel() + thread_pool = ThreadPoolExecutor(max_workers=4) # This should raise ValueError due to invalid table name try: - invalid_store = PGVectorStore( + _ = PGVectorStore( collection_name="test'; DROP TABLE users; --", embedding_model=embedding_model, + thread_pool=thread_pool, ) logger.error("❌ FAILED: Invalid collection name was accepted (SQL injection risk!)") assert False, "Should have raised ValueError for invalid collection name" @@ -1471,7 +1480,7 @@ async def test_sql_injection_protection(store: BaseVectorStore, store_name: str) # Test 2: Invalid metadata key in filters try: - results = await store.search( + _ = await store.search( query="test", filters={ "normal_key": "value",