feat(context): add comprehensive context management system with prompt handling and registries

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",
+]

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__()

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

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

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)

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()

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
 

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"])

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"],
+            },
+        },
     }
 
     # 解析