feat(core): add core modules for LLM integration and context management

Author: jinliyl

reme_ai/core/context/__init__.py

@@ -0,0 +1,14 @@
+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/flow_context.py

@@ -1,9 +1,3 @@
-"""Prompt handler for loading and managing prompts.
-
-This module provides a class for loading prompts from files, managing
-language-specific prompts, and formatting prompts with variables.
-"""
-
 from pathlib import Path
 
 import yaml
@@ -14,37 +8,11 @@
 
 
 class PromptHandler(BaseContext):
-    """Handler for loading, storing, and formatting prompts.
-
-    This class manages prompts loaded from YAML files, supports
-    language-specific prompts, and provides formatting with variables
-    and conditional flags.
-
-    Attributes:
-        language: Language code for language-specific prompts.
-    """
-
     def __init__(self, language: str = "", **kwargs):
-        """Initialize PromptHandler with language setting.
-
-        Args:
-            language: Language code for language-specific prompts.
-                If not provided, uses the language from service context.
-            **kwargs: Additional context data to store.
-        """
         super().__init__(**kwargs)
         self.language: str = language or C.language
 
     def load_prompt_by_file(self, prompt_file_path: Path | str = None):
-        """Load prompts from a YAML file.
-
-        Args:
-            prompt_file_path: Path to the YAML file containing prompts.
-                Can be a Path object or string. If None, does nothing.
-
-        Returns:
-            Self for method chaining.
-        """
         if prompt_file_path is None:
             return self
 
@@ -60,64 +28,27 @@ def load_prompt_by_file(self, prompt_file_path: Path | str = None):
         return self
 
     def load_prompt_dict(self, prompt_dict: dict = None):
-        """Load prompts from a dictionary.
-
-        Args:
-            prompt_dict: Dictionary of prompt names to prompt strings.
-                If None, does nothing.
-
-        Returns:
-            Self for method chaining.
-        """
         if not prompt_dict:
             return self
 
         for key, value in prompt_dict.items():
             if isinstance(value, str):
-                if key in self._data:
-                    self._data[key] = value
-                    logger.warning(f"prompt_dict key={key} overwrite!")
-
+                if key in self:
+                    logger.warning(f"Overwriting prompt key={key}, old_value={self[key]}, new_value={value}")
                 else:
-                    self._data[key] = value
-                    logger.debug(f"add prompt_dict key={key}")
+                    logger.debug(f"Adding new prompt key={key}, value={value}")
+                self[key] = value
         return self
 
     def get_prompt(self, prompt_name: str):
-        """Get a prompt by name, with language suffix if applicable.
-
-        Args:
-            prompt_name: Base name of the prompt to retrieve.
-
-        Returns:
-            The prompt string.
-
-        Raises:
-            AssertionError: If the prompt (with language suffix) is not found.
-        """
         key: str = prompt_name
         if self.language and not key.endswith(self.language.strip()):
             key += "_" + self.language.strip()
 
-        assert key in self._data, f"prompt_name={key} not found."
-        return self._data[key]
+        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 with variables and conditional flags.
-
-        This method supports two types of formatting:
-        1. Boolean flags: Lines starting with [flag_name] are included
-           only if the corresponding flag is True.
-        2. Variable substitution: Other kwargs are used for string
-           formatting with {variable_name}.
-
-        Args:
-            prompt_name: Name of the prompt to format.
-            **kwargs: Variables and flags for formatting.
-
-        Returns:
-            The formatted prompt string.
-        """
         prompt = self.get_prompt(prompt_name)
 
         flag_kwargs = {k: v for k, v in kwargs.items() if isinstance(v, bool)}
@@ -128,7 +59,7 @@ def prompt_format(self, prompt_name: str, **kwargs) -> str:
             for line in prompt.strip().split("\n"):
                 hit = False
                 hit_flag = True
-                for key, flag in kwargs.items():
+                for key, flag in flag_kwargs.items():
                     if not line.startswith(f"[{key}]"):
                         continue
 

reme_ai/core/context/prompt_handler.py

@@ -1,43 +1,12 @@
-"""Registry for class registration and lookup.
-
-This module provides a registry class for managing class registrations
-with dynamic lookup capabilities.
-"""
-
 from .base_context import BaseContext
 
 
 class Registry(BaseContext):
-    """Registry for storing and retrieving registered classes.
-
-    This class provides a decorator-based registration system for classes,
-    allowing dynamic class lookup by name.
-    """
-
     def register(self, name: str = "", add_cls: bool = True):
-        """Register a class in the registry.
-
-        Args:
-            name: Name to register the class under.
-            add_cls: Whether to actually add the class to the registry.
-                Defaults to True.
-
-        Returns:
-            Decorator function that registers the class when applied.
-        """
-
         def decorator(cls):
-            """Decorator function that registers the class.
-
-            Args:
-                cls: The class to register.
-
-            Returns:
-                The class unchanged (for use as decorator).
-            """
             if add_cls:
                 key = name or cls.__name__
-                self._data[key] = cls
+                self[key] = cls
             return cls
 
         return decorator

reme_ai/core/context/registry.py

@@ -0,0 +1,43 @@
+import asyncio
+
+from .base_context import BaseContext
+from ..enumeration import ChunkEnum
+from ..schema import Response
+from ..schema import StreamChunk
+
+
+class RuntimeContext(BaseContext):
+    def __init__(
+        self,
+        response: Response | None = None,
+        stream_queue: asyncio.Queue | None = None,
+        **kwargs,
+    ):
+        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):
+        if self.stream_queue is None:
+            return self
+        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):
+        if self.stream_queue is None:
+            return self
+        await self.stream_queue.put(stream_chunk)
+        return self
+
+    async def add_stream_done(self):
+        if self.stream_queue is None:
+            return self
+        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):
+        self.response.success = False
+        self.response.answer = str(e.args)