feat(core): add core modules for AI framework

Author: jinliyl

reme_ai/core/context/__init__.py

@@ -0,0 +1,26 @@
+class BaseContext(dict):
+    """A dict subclass that supports attribute-style access and pickling."""
+
+    def __getattr__(self, name):
+        try:
+            return self[name]
+        except KeyError:
+            raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
+
+    def __setattr__(self, name, value):
+        self[name] = value
+
+    def __delattr__(self, name):
+        try:
+            del self[name]
+        except KeyError:
+            raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
+
+    def __getstate__(self):
+        return dict(self)
+
+    def __setstate__(self, state):
+        self.update(state)
+
+    def __reduce__(self):
+        return self.__class__, (), self.__getstate__()

reme_ai/core/context/base_context.py

@@ -0,0 +1,98 @@
+"""Flow context for managing flow execution state.
+
+This module provides a context class for managing the state of flow execution,
+including flow identification, response handling, and streaming capabilities.
+"""
+
+import asyncio
+import uuid
+from typing import Optional
+
+from .base_context import BaseContext
+from ..enumeration import ChunkEnum
+from ..schema import FlowResponse
+from ..schema import FlowStreamChunk
+
+
+class FlowContext(BaseContext):
+    """Context for managing flow execution state and streaming.
+
+    This class manages the state of a single flow execution, including:
+    - Flow identification
+    - Response handling
+    - Stream queue for asynchronous streaming
+
+    Attributes:
+        flow_id: Unique identifier for the flow instance.
+        response: FlowResponse object for storing flow results.
+        stream_queue: Asynchronous queue for streaming chunks.
+    """
+
+    def __init__(
+        self,
+        flow_id: str = uuid.uuid4().hex,
+        response: Optional[FlowResponse] = None,
+        stream_queue: Optional[asyncio.Queue] = None,
+        **kwargs,
+    ):
+        """Initialize FlowContext with flow ID and optional components.
+
+        Args:
+            flow_id: Unique identifier for the flow instance.
+                Defaults to a random UUID hex string.
+            response: FlowResponse object. If None, a new FlowResponse
+                will be created.
+            stream_queue: Asynchronous queue for streaming chunks.
+            **kwargs: Additional context data to store.
+        """
+        super().__init__(**kwargs)
+
+        self.flow_id: str = flow_id
+        self.response: Optional[FlowResponse] = response if response is not None else FlowResponse()
+        self.stream_queue: Optional[asyncio.Queue] = stream_queue
+
+    async def add_stream_string_and_type(self, chunk: str, chunk_type: ChunkEnum):
+        """Add a stream chunk with string content and type to the stream queue.
+
+        Args:
+            chunk: The string content to stream.
+            chunk_type: The type of the chunk.
+
+        Returns:
+            Self for method chaining.
+        """
+        stream_chunk = FlowStreamChunk(flow_id=self.flow_id, chunk_type=chunk_type, chunk=chunk)
+        await self.stream_queue.put(stream_chunk)
+        return self
+
+    async def add_stream_chunk(self, stream_chunk: FlowStreamChunk):
+        """Add a stream chunk to the stream queue.
+
+        Args:
+            stream_chunk: The stream chunk to add.
+
+        Returns:
+            Self for method chaining.
+        """
+        stream_chunk.flow_id = self.flow_id
+        await self.stream_queue.put(stream_chunk)
+        return self
+
+    async def add_stream_done(self):
+        """Add a done signal to the stream queue.
+
+        Returns:
+            Self for method chaining.
+        """
+        done_chunk = FlowStreamChunk(flow_id=self.flow_id, chunk_type=ChunkEnum.DONE, chunk="", done=True)
+        await self.stream_queue.put(done_chunk)
+        return self
+
+    def add_response_error(self, e: Exception):
+        """Add an error to the flow response.
+
+        Args:
+            e: The exception to record as an error.
+        """
+        self.response.success = False
+        self.response.answer = str(e.args)

reme_ai/core/context/flow_context.py

@@ -0,0 +1,150 @@
+"""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
+from loguru import logger
+
+from .base_context import BaseContext
+from .service_context import C
+
+
+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
+
+        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:
+            prompt_dict = yaml.load(f, yaml.FullLoader)
+            self.load_prompt_dict(prompt_dict)
+        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!")
+
+                else:
+                    self._data[key] = value
+                    logger.debug(f"add prompt_dict key={key}")
+        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]
+
+    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)}
+        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 kwargs.items():
+                    if not line.startswith(f"[{key}]"):
+                        continue
+
+                    hit = True
+                    hit_flag = flag
+                    line = line.strip(f"[{key}]")
+                    break
+
+                if not hit:
+                    split_prompt.append(line)
+                elif hit_flag:
+                    split_prompt.append(line)
+
+            prompt = "\n".join(split_prompt)
+
+        if other_kwargs:
+            prompt = prompt.format(**other_kwargs)
+
+        return prompt

reme_ai/core/context/prompt_handler.py

@@ -0,0 +1,43 @@
+"""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
+            return cls
+
+        return decorator

reme_ai/core/context/registry.py

@@ -0,0 +1,11 @@
+from .chunk_enum import ChunkEnum
+from .http_enum import HttpEnum
+from .registry_enum import RegistryEnum
+from .role import Role
+
+__all__ = [
+    "ChunkEnum",
+    "HttpEnum",
+    "RegistryEnum",
+    "Role",
+]

reme_ai/core/enumeration/__init__.py

@@ -0,0 +1,15 @@
+from enum import Enum
+
+
+class ChunkEnum(str, Enum):
+    THINK = "think"
+
+    ANSWER = "answer"
+
+    TOOL = "tool"
+
+    USAGE = "usage"
+
+    ERROR = "error"
+
+    DONE = "done"

reme_ai/core/enumeration/chunk_enum.py

@@ -0,0 +1,13 @@
+from enum import Enum
+
+
+class HttpEnum(str, Enum):
+    GET = "get"
+
+    POST = "post"
+
+    HEAD = "head"
+
+    PUT = "put"
+
+    DELETE = "delete"

reme_ai/core/enumeration/http_enum.py

@@ -0,0 +1,17 @@
+from enum import Enum
+
+
+class RegistryEnum(str, Enum):
+    LLM = "llm"
+
+    EMBEDDING_MODEL = "embedding_model"
+
+    VECTOR_STORE = "vector_store"
+
+    OP = "op"
+
+    FLOW = "flow"
+
+    SERVICE = "service"
+
+    TOKEN_COUNTER = "token_counter"

reme_ai/core/enumeration/registry_enum.py

@@ -0,0 +1,11 @@
+from enum import Enum
+
+
+class Role(str, Enum):
+    SYSTEM = "system"
+
+    USER = "user"
+
+    ASSISTANT = "assistant"
+
+    TOOL = "tool"