diff --git a/reme_ai/core/context/__init__.py b/reme_ai/core/context/__init__.py new file mode 100644 index 00000000..7f26d600 --- /dev/null +++ b/reme_ai/core/context/__init__.py @@ -0,0 +1,16 @@ +"""context""" + +from .base_context import BaseContext +from .prompt_handler import PromptHandler +from .registry import Registry +from .runtime_context import RuntimeContext +from .service_context import ServiceContext, C + +__all__ = [ + "BaseContext", + "PromptHandler", + "Registry", + "RuntimeContext", + "ServiceContext", + "C", +] diff --git a/reme_ai/core/context/base_context.py b/reme_ai/core/context/base_context.py new file mode 100644 index 00000000..dabd8cdb --- /dev/null +++ b/reme_ai/core/context/base_context.py @@ -0,0 +1,41 @@ +"""Module providing a dictionary subclass with attribute-style access and pickling support.""" + +from typing import Generic, TypeVar + +_KT = TypeVar("_KT") +_VT = TypeVar("_VT") + + +class BaseContext(dict, Generic[_KT, _VT]): + """A dictionary subclass that enables accessing and modifying keys as attributes.""" + + def __getattr__(self, name: str) -> _VT: + """Retrieve a dictionary item as an attribute.""" + try: + return self[name] + except KeyError as e: + raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'") from e + + def __setattr__(self, name: str, value: _VT) -> None: + """Assign a value to a dictionary item using attribute syntax.""" + self[name] = value + + def __delattr__(self, name: str) -> None: + """Remove a dictionary item using attribute syntax.""" + try: + # Delete item from dict via key + del self[name] + except KeyError as e: + raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'") from e + + def __getstate__(self) -> dict: + """Return the dictionary representation for pickling.""" + return dict(self) + + def __setstate__(self, state: dict) -> None: + """Restore the dictionary state from a pickled object.""" + self.update(state) + + def __reduce__(self): + """Define the reconstruction logic for pickling processes.""" + return self.__class__, (), self.__getstate__() diff --git a/reme_ai/core/context/prompt_handler.py b/reme_ai/core/context/prompt_handler.py new file mode 100644 index 00000000..e48f98ac --- /dev/null +++ b/reme_ai/core/context/prompt_handler.py @@ -0,0 +1,95 @@ +"""Module for managing and formatting prompt templates from files or dictionaries.""" + +from pathlib import Path + +import yaml +from loguru import logger + +from .base_context import BaseContext +from .service_context import C + + +class PromptHandler(BaseContext): + """A context-aware handler for loading, retrieving, and formatting prompt templates.""" + + def __init__(self, language: str = "", **kwargs): + """Initialize the handler with a specific language and optional context data.""" + super().__init__(**kwargs) + self.language: str = language or C.language + + def load_prompt_by_file(self, prompt_file_path: Path | str = None): + """Load prompt configurations from a YAML file into the context.""" + if prompt_file_path is None: + return self + + if isinstance(prompt_file_path, str): + prompt_file_path = Path(prompt_file_path) + + if not prompt_file_path.exists(): + return self + + with prompt_file_path.open(encoding="utf-8") as f: + # Load YAML content using the full loader + prompt_dict = yaml.load(f, yaml.FullLoader) + self.load_prompt_dict(prompt_dict) + return self + + def load_prompt_dict(self, prompt_dict: dict = None): + """Merge a dictionary of prompt strings into the current context.""" + if not prompt_dict: + return self + + for key, value in prompt_dict.items(): + if isinstance(value, str): + if key in self: + logger.warning(f"Overwriting prompt key={key}, old_value={self[key]}, new_value={value}") + else: + logger.debug(f"Adding new prompt key={key}, value={value}") + self[key] = value + return self + + def get_prompt(self, prompt_name: str): + """Retrieve a prompt by name, automatically appending the language suffix if needed.""" + key: str = prompt_name + if self.language and not key.endswith(self.language.strip()): + key += "_" + self.language.strip() + + assert key in self, f"prompt_name={key} not found." + return self[key] + + def prompt_format(self, prompt_name: str, **kwargs) -> str: + """Format a prompt by filtering flagged lines and filling template variables.""" + prompt = self.get_prompt(prompt_name) + + # Separate boolean flags from string formatting arguments + flag_kwargs = {k: v for k, v in kwargs.items() if isinstance(v, bool)} + other_kwargs = {k: v for k, v in kwargs.items() if not isinstance(v, bool)} + + if flag_kwargs: + split_prompt = [] + for line in prompt.strip().split("\n"): + hit = False + hit_flag = True + for key, flag in flag_kwargs.items(): + if not line.startswith(f"[{key}]"): + continue + + hit = True + hit_flag = flag + # Remove the flag prefix from the line + line = line.strip(f"[{key}]") + break + + # Include line if no flag is present or if the flag evaluates to True + if not hit: + split_prompt.append(line) + elif hit_flag: + split_prompt.append(line) + + prompt = "\n".join(split_prompt) + + if other_kwargs: + # Apply standard Python string formatting + prompt = prompt.format(**other_kwargs) + + return prompt diff --git a/reme_ai/core/context/registry.py b/reme_ai/core/context/registry.py new file mode 100644 index 00000000..f403037d --- /dev/null +++ b/reme_ai/core/context/registry.py @@ -0,0 +1,19 @@ +"""Module providing a registry class for managing class-to-name mappings via decorators.""" + +from .base_context import BaseContext + + +class Registry(BaseContext): + """A registry container that uses decorators to map and store class references.""" + + def register(self, name: str = "", add_cls: bool = True): + """Return a decorator that registers a class under a specific name in the registry.""" + + def decorator(cls): + if add_cls: + # Use provided name or default to the class name as the key + key = name or cls.__name__ + self[key] = cls + return cls + + return decorator diff --git a/reme_ai/core/context/runtime_context.py b/reme_ai/core/context/runtime_context.py new file mode 100644 index 00000000..ded35ded --- /dev/null +++ b/reme_ai/core/context/runtime_context.py @@ -0,0 +1,56 @@ +"""Module providing a 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 + + +class RuntimeContext(BaseContext): + """A context class for handling execution state, including response metadata and stream queues.""" + + def __init__( + self, + response: Response | None = None, + stream_queue: asyncio.Queue | None = None, + **kwargs, + ): + """Initialize the runtime context with optional response objects and message queues.""" + super().__init__(**kwargs) + + self.response: Response | None = response if response is not None else Response() + self.stream_queue: asyncio.Queue | None = stream_queue + + 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 + + # Package raw data into a StreamChunk schema + stream_chunk = StreamChunk(chunk_type=chunk_type, chunk=chunk) + await self.stream_queue.put(stream_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) + 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) + return self + + def add_response_error(self, e: Exception): + """Update the internal response object to reflect a failure state using exception details.""" + self.response.success = False + self.response.answer = str(e.args) diff --git a/reme_ai/core/context/service_context.py b/reme_ai/core/context/service_context.py new file mode 100644 index 00000000..af99a870 --- /dev/null +++ b/reme_ai/core/context/service_context.py @@ -0,0 +1,105 @@ +"""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 +from ..enumeration import RegistryEnum +from ..schema import ServiceConfig +from ..utils import singleton + + +@singleton +class ServiceContext(BaseContext): + """A singleton container for global application state, thread pools, and component registries.""" + + def __init__(self, **kwargs): + """Initialize the global context with configuration objects and specialized registries.""" + super().__init__(**kwargs) + + self.service_config: ServiceConfig | None = None + self.language: str = "" + self.thread_pool: ThreadPoolExecutor | None = None + self.vector_store_dict: Dict[str, dict] = {} + self.external_mcp_tool_call_dict: dict = {} + # Initialize a registry for every category defined in RegistryEnum + self.registry_dict: Dict[RegistryEnum, Registry] = {v: Registry() for v in RegistryEnum.__members__.values()} + self.flow_dict: dict = {} + + 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) + + def register_llm(self, name: str = ""): + """Register a Large Language Model class.""" + return self.register(name=name, register_type=RegistryEnum.LLM) + + def register_embedding_model(self, name: str = ""): + """Register an embedding model class.""" + return self.register(name=name, register_type=RegistryEnum.EMBEDDING_MODEL) + + def register_vector_store(self, name: str = ""): + """Register a vector store implementation class.""" + return self.register(name=name, register_type=RegistryEnum.VECTOR_STORE) + + def register_op(self, name: str = ""): + """Register an operation (Op) class.""" + return self.register(name=name, register_type=RegistryEnum.OP) + + def register_flow(self, name: str = ""): + """Register a workflow or logic flow class.""" + return self.register(name=name, register_type=RegistryEnum.FLOW) + + def register_service(self, name: str = ""): + """Register a backend service class.""" + return self.register(name=name, register_type=RegistryEnum.SERVICE) + + def register_token_counter(self, name: str = ""): + """Register a token counting utility class.""" + 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.""" + 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_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) + + def get_op_class(self, name: str): + """Get the operation class registered under the given name.""" + return self.get_model_class(name, RegistryEnum.OP) + + def get_flow_class(self, name: str): + """Get the flow class registered under the given name.""" + return self.get_model_class(name, RegistryEnum.FLOW) + + def get_service_class(self, name: str): + """Get the service class registered under the given name.""" + return self.get_model_class(name, RegistryEnum.SERVICE) + + def get_token_counter_class(self, name: str): + """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_flow(self, name: str): + """Retrieve a specific flow instance by name.""" + return self.flow_dict[name] + + +# Export a global instance for easy access across the application +C = ServiceContext() diff --git a/reme_ai/core/schema/tool_call.py b/reme_ai/core/schema/tool_call.py index 3688da54..d190974e 100644 --- a/reme_ai/core/schema/tool_call.py +++ b/reme_ai/core/schema/tool_call.py @@ -1,6 +1,7 @@ """ MCP Tool Schema definitions for recursive JSON Schema representation. """ + import json from typing import Any, Dict, List, Literal, Optional, Union @@ -12,6 +13,7 @@ from ..enumeration.json_schema_enum import JsonSchemaEnum class ToolAttr(BaseModel): """Recursive model representing JSON Schema attributes for tool parameters.""" + model_config = ConfigDict(extra="allow") type: Literal[ @@ -23,8 +25,8 @@ class ToolAttr(BaseModel): JsonSchemaEnum.BOOLEAN.value, JsonSchemaEnum.NULL.value, ] = Field( - default=JsonSchemaEnum.STRING.value, - description="The data type of the attribute" + default=JsonSchemaEnum.STRING.value, + 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") @@ -32,7 +34,6 @@ class ToolAttr(BaseModel): items: Optional[Union[Dict[str, Any], "ToolAttr"]] = Field(default=None, description="Schema for array items") enum: Optional[List[str]] = Field(default=None, description="Allowed values for the attribute") - def simple_input_dump(self) -> dict: """Serializes the attribute into a standard JSON Schema dictionary.""" res: dict = {"type": self.type} @@ -42,8 +43,9 @@ class ToolAttr(BaseModel): res["enum"] = self.enum if self.type == "object" and self.properties: - res["properties"] = {k: v.simple_input_dump() if isinstance(v, ToolAttr) else v - for k, v in self.properties.items()} + res["properties"] = { + k: v.simple_input_dump() if isinstance(v, ToolAttr) else v for k, v in self.properties.items() + } if self.required: res["required"] = self.required diff --git a/tests/test_message.py b/tests/test_message.py index 35d8598e..b45563dd 100644 --- a/tests/test_message.py +++ b/tests/test_message.py @@ -27,7 +27,7 @@ class TestModelDefinitions(unittest.TestCase): self.assertEqual(dump["type"], "string") self.assertEqual(dump["enum"], ["Beijing", "London"]) self.assertIn("description", dump) - + # Test object attribute with required child properties obj_attr = ToolAttr( type="object", @@ -39,10 +39,10 @@ class TestModelDefinitions(unittest.TestCase): required=["name"], # 'name' is required, 'age' is optional ) obj_dump = obj_attr.simple_input_dump() - + print("\n=== ToolAttr.simple_input_dump() (object with required) ===") print(obj_dump) - + self.assertEqual(obj_dump["type"], "object") self.assertIn("properties", obj_dump) self.assertEqual(obj_dump["required"], ["name"]) diff --git a/tests/test_tool_call.py b/tests/test_tool_call.py index 242f07ee..f914dca7 100644 --- a/tests/test_tool_call.py +++ b/tests/test_tool_call.py @@ -1,3 +1,5 @@ +"""simple tool call test""" + import json from reme_ai.core.schema.tool_call import ToolCall @@ -16,11 +18,11 @@ def test_simple_schema(): "type": "object", "properties": { "city": {"type": "string", "description": "城市名称"}, - "unit": {"type": "string", "description": "温度单位", "enum": ["celsius", "fahrenheit"]} + "unit": {"type": "string", "description": "温度单位", "enum": ["celsius", "fahrenheit"]}, }, - "required": ["city"] - } - } + "required": ["city"], + }, + }, } # 解析 @@ -58,14 +60,14 @@ def test_medium_nested_schema(): "properties": { "name": {"type": "string", "description": "客户姓名"}, "email": {"type": "string", "description": "客户邮箱"}, - "phone": {"type": "string", "description": "联系电话"} + "phone": {"type": "string", "description": "联系电话"}, }, - "required": ["name", "email"] - } + "required": ["name", "email"], + }, }, - "required": ["order_id", "customer"] - } - } + "required": ["order_id", "customer"], + }, + }, } # 解析 @@ -106,9 +108,9 @@ def test_nested_schema(): "description": "用户元数据", "properties": { "age": {"type": "integer"}, - "location": {"type": "string"} + "location": {"type": "string"}, }, - "required": ["age"] + "required": ["age"], }, "tags": { "type": "array", @@ -117,15 +119,15 @@ def test_nested_schema(): "type": "object", "properties": { "tag_id": {"type": "string"}, - "level": {"type": "number"} + "level": {"type": "number"}, }, - "required": ["tag_id"] - } - } + "required": ["tag_id"], + }, + }, }, - "required": ["username", "metadata"] - } - } + "required": ["username", "metadata"], + }, + }, } # 2. 解析:将原始字典转化为 ToolCall 实例 @@ -171,17 +173,17 @@ def test_array_of_primitives(): "file_paths": { "type": "array", "description": "文件路径列表", - "items": {"type": "string"} + "items": {"type": "string"}, }, "priorities": { "type": "array", "description": "优先级列表", - "items": {"type": "integer"} - } + "items": {"type": "integer"}, + }, }, - "required": ["file_paths"] - } - } + "required": ["file_paths"], + }, + }, } # 解析 @@ -191,8 +193,8 @@ def test_array_of_primitives(): file_paths_attr = tool_call.input_schema["file_paths"] print(f"file_paths 类型: {file_paths_attr.type}") - print( - f"file_paths items 类型: {file_paths_attr.items.type if hasattr(file_paths_attr.items, 'type') else file_paths_attr.items}") + t_items_type = file_paths_attr.items.type if hasattr(file_paths_attr.items, "type") else file_paths_attr.items + print(f"file_paths items 类型: {t_items_type}") # 导出并验证相等性 dumped_data = tool_call.simple_input_dump() @@ -230,12 +232,12 @@ def test_deep_nested_schema(): "type": "object", "properties": { "email": {"type": "string"}, - "phone": {"type": "string"} + "phone": {"type": "string"}, }, - "required": ["email"] - } + "required": ["email"], + }, }, - "required": ["name", "contact"] + "required": ["name", "contact"], }, "members": { "type": "array", @@ -247,19 +249,19 @@ def test_deep_nested_schema(): "role": {"type": "string"}, "skills": { "type": "array", - "items": {"type": "string"} - } + "items": {"type": "string"}, + }, }, - "required": ["name", "role"] - } - } + "required": ["name", "role"], + }, + }, }, - "required": ["leader"] - } + "required": ["leader"], + }, }, - "required": ["project_name", "team"] - } - } + "required": ["project_name", "team"], + }, + }, } # 解析 @@ -302,12 +304,12 @@ def test_mixed_types_schema(): "mode": { "type": "string", "description": "运行模式", - "enum": ["development", "production", "testing"] + "enum": ["development", "production", "testing"], }, "allowed_ips": { "type": "array", "description": "允许的IP地址列表", - "items": {"type": "string"} + "items": {"type": "string"}, }, "database": { "type": "object", @@ -315,14 +317,14 @@ def test_mixed_types_schema(): "properties": { "host": {"type": "string"}, "port": {"type": "integer"}, - "ssl_enabled": {"type": "boolean"} + "ssl_enabled": {"type": "boolean"}, }, - "required": ["host", "port"] - } + "required": ["host", "port"], + }, }, - "required": ["enabled", "mode"] - } - } + "required": ["enabled", "mode"], + }, + }, } # 解析