diff --git a/reme_ai/core/context/__init__.py b/reme_ai/core/context/__init__.py index 7f26d600..27957fd2 100644 --- a/reme_ai/core/context/__init__.py +++ b/reme_ai/core/context/__init__.py @@ -2,15 +2,10 @@ 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 +from .registry_factory import R __all__ = [ "BaseContext", "PromptHandler", - "Registry", - "RuntimeContext", - "ServiceContext", - "C", + "R", ] diff --git a/reme_ai/core/context/prompt_handler.py b/reme_ai/core/context/prompt_handler.py index 9639cbf8..e6b6d737 100644 --- a/reme_ai/core/context/prompt_handler.py +++ b/reme_ai/core/context/prompt_handler.py @@ -28,25 +28,24 @@ class PromptNotFoundError(KeyError): super().__init__( f"Prompt '{prompt_name}' not found. " f"Available prompts: {', '.join(available_prompts[:10])}" - f"{'...' if len(available_prompts) > 10 else ''}" + f"{'...' if len(available_prompts) > 10 else ''}", ) class PromptFormattingError(ValueError): """Exception raised when prompt formatting fails.""" - pass class PromptHandler(BaseContext): """A context-aware handler for loading, retrieving, and formatting prompt templates. - + This handler supports: - Loading prompts from YAML/JSON files or dictionaries - Multi-language prompt support with automatic language suffix - Conditional line filtering using boolean flags (e.g., [debug], [verbose]) - Template variable substitution with validation - Method chaining for fluent API - + Examples: >>> handler = PromptHandler(language="en") >>> handler.load_prompt_dict({ @@ -61,7 +60,7 @@ class PromptHandler(BaseContext): def __init__(self, language: str = "", **kwargs): """Initialize the PromptHandler with optional language configuration. - + Args: language: Language code to append as suffix (e.g., "en", "zh", "ja"). If provided, get_prompt will automatically try to find @@ -72,24 +71,24 @@ class PromptHandler(BaseContext): self.language: str = language.strip() def load_prompt_by_file( - self, - prompt_file_path: Optional[Union[Path, str]] = None, - overwrite: bool = True + self, + prompt_file_path: Optional[Union[Path, str]] = None, + overwrite: bool = True, ) -> "PromptHandler": """Load prompt configurations from a YAML or JSON file into the context. - + Supports both YAML (.yaml, .yml) and JSON (.json) file formats. Non-existent files are silently skipped. - + Args: prompt_file_path: Path to the prompt configuration file. If None, returns self without changes. overwrite: If True, allows overwriting existing prompts with warnings. If False, skips existing prompts without overwriting. - + Returns: Self for method chaining. - + Raises: ValueError: If file format is not supported. yaml.YAMLError: If YAML parsing fails. @@ -115,8 +114,7 @@ class PromptHandler(BaseContext): prompt_dict = json.load(f) else: raise ValueError( - f"Unsupported file format: {suffix}. " - f"Supported formats: .yaml, .yml, .json" + f"Unsupported file format: {suffix}. " f"Supported formats: .yaml, .yml, .json", ) logger.info(f"Loaded {len(prompt_dict or {})} prompts from {prompt_file_path}") @@ -125,23 +123,23 @@ class PromptHandler(BaseContext): except (yaml.YAMLError, json.JSONDecodeError) as e: logger.error(f"Failed to parse prompt file {prompt_file_path}: {e}") raise - + return self def load_prompt_dict( - self, - prompt_dict: Optional[Dict[str, Any]] = None, - overwrite: bool = True + self, + prompt_dict: Optional[Dict[str, Any]] = None, + overwrite: bool = True, ) -> "PromptHandler": """Merge a dictionary of prompt strings into the current context. - + Only string values are stored as prompts. Non-string values are skipped. - + Args: prompt_dict: Dictionary mapping prompt names to prompt template strings. overwrite: If True, allows overwriting existing prompts with warnings. If False, skips existing prompts without overwriting. - + Returns: Self for method chaining. """ @@ -156,8 +154,7 @@ class PromptHandler(BaseContext): if key in self: if overwrite: logger.warning( - f"Overwriting prompt '{key}': " - f"old length={len(self[key])}, new length={len(value)}" + f"Overwriting prompt '{key}': " f"old length={len(self[key])}, new length={len(value)}", ) self[key] = value else: @@ -170,20 +167,20 @@ class PromptHandler(BaseContext): def get_prompt(self, prompt_name: str, fallback_to_base: bool = True) -> str: """Retrieve a prompt by name with automatic language suffix handling. - + If a language is configured, this method will: 1. First try to find the prompt with language suffix (e.g., "greeting_en") 2. If not found and fallback_to_base is True, try the base name (e.g., "greeting") 3. Otherwise, raise PromptNotFoundError - + Args: prompt_name: Name of the prompt to retrieve. fallback_to_base: If True and language-specific prompt not found, fallback to prompt without language suffix. - + Returns: The prompt template string, stripped of leading/trailing whitespace. - + Raises: PromptNotFoundError: If the prompt is not found. """ @@ -211,10 +208,10 @@ class PromptHandler(BaseContext): def has_prompt(self, prompt_name: str) -> bool: """Check if a prompt exists (with or without language suffix). - + Args: prompt_name: Name of the prompt to check. - + Returns: True if the prompt exists, False otherwise. """ @@ -226,11 +223,11 @@ class PromptHandler(BaseContext): def list_prompts(self, language_filter: Optional[str] = None) -> list[str]: """List all available prompt names. - + Args: language_filter: If provided, only return prompts for this language. If None, return all prompts. - + Returns: List of prompt names. """ @@ -243,31 +240,27 @@ class PromptHandler(BaseContext): @staticmethod def _extract_format_fields(template: str) -> set[str]: """Extract all format field names from a template string. - + Args: template: Template string with {variable} placeholders. - + Returns: Set of field names used in the template. """ - return { - field_name - for _, field_name, _, _ in Formatter().parse(template) - if field_name is not None - } + return {field_name for _, field_name, _, _ in Formatter().parse(template) if field_name is not None} @staticmethod def _filter_conditional_lines(prompt: str, flags: Dict[str, bool]) -> str: """Filter lines based on boolean flags. - + Lines starting with [flag_name] are conditionally included based on the value of flags[flag_name]. If True, the line is included (without the flag marker). If False, the line is excluded. - + Args: prompt: The prompt text with conditional markers. flags: Dictionary of flag names to boolean values. - + Returns: Filtered prompt text. """ @@ -288,38 +281,38 @@ class PromptHandler(BaseContext): elif flags[matched_flag]: # Flag is True, include without marker marker = f"[{matched_flag}]" - filtered_lines.append(line[len(marker):]) + filtered_lines.append(line[len(marker) :]) # else: Flag is False, skip this line return "\n".join(filtered_lines) def prompt_format( - self, - prompt_name: str, - validate: bool = True, - **kwargs + self, + prompt_name: str, + validate: bool = True, + **kwargs, ) -> str: """Format a prompt with conditional line filtering and variable substitution. - + This method performs two-stage formatting: 1. Conditional line filtering: Lines marked with [flag] are included only if the corresponding boolean kwarg is True. 2. Variable substitution: Template variables {var} are replaced with provided values. - + Args: prompt_name: Name of the prompt to format. validate: If True, check that all required template variables are provided. **kwargs: Keyword arguments for formatting. Boolean values are treated as conditional flags, other values are used for template substitution. - + Returns: Formatted prompt string. - + Raises: PromptNotFoundError: If the prompt is not found. PromptFormattingError: If validation fails or formatting errors occur. - + Examples: >>> handler = PromptHandler() >>> handler["test"] = "[debug]Debug: {info}\\nResult: {value}" @@ -347,7 +340,7 @@ class PromptHandler(BaseContext): if missing_fields: raise PromptFormattingError( f"Missing required format variables for prompt '{prompt_name}': " - f"{', '.join(sorted(missing_fields))}" + f"{', '.join(sorted(missing_fields))}", ) # Step 3: Format with variables @@ -356,18 +349,15 @@ class PromptHandler(BaseContext): prompt = prompt.format(**format_kwargs) except KeyError as e: raise PromptFormattingError( - f"Format error in prompt '{prompt_name}': missing variable {e}" + f"Format error in prompt '{prompt_name}': missing variable {e}", ) from e except (ValueError, IndexError) as e: raise PromptFormattingError( - f"Format error in prompt '{prompt_name}': {e}" + f"Format error in prompt '{prompt_name}': {e}", ) from e return prompt.strip() def __repr__(self) -> str: """Return a string representation of the PromptHandler.""" - return ( - f"PromptHandler(language='{self.language}', " - f"num_prompts={len(self)})" - ) + return f"PromptHandler(language='{self.language}', " f"num_prompts={len(self)})" diff --git a/reme_ai/core/context/registry.py b/reme_ai/core/context/registry.py deleted file mode 100644 index 9fa35d30..00000000 --- a/reme_ai/core/context/registry.py +++ /dev/null @@ -1,143 +0,0 @@ -"""Module providing a registry class for managing class-to-name mappings via decorators.""" - -import inspect -from typing import Callable, TypeVar - -from .base_context import BaseContext -from ..enumeration import RegistryEnum -from ...core_old.utils import singleton - -T = TypeVar("T") - - -@singleton -class Registry(BaseContext): - """A singleton registry manager that maintains separate registries for different component types. - - This class serves as the central registry hub for the entire ReMe application, providing: - - Component registration for different types (LLMs, embeddings, vector stores, etc.) - - Convenient access methods for retrieving registered classes - - Decorator-based registration API - - The singleton pattern ensures only one instance exists throughout the application lifecycle, - accessible via the global `R` variable exported at the bottom of this module. - """ - - def __init__(self, **kwargs): - """Initialize the registry manager with separate registries for each component type.""" - super().__init__(**kwargs) - - # Registry system: stores class definitions for different component types - self.registry_dict: dict[RegistryEnum, dict] = { - v: {} for v in RegistryEnum.__members__.values() - } - - def register(self, name: str | type = "", register_type: RegistryEnum = None) -> Callable[[type[T]], type[T]] | type[T]: - """Return a decorator to register a component within a specific registry category. - - Can be used in multiple ways: - - @R.register_op() # with empty parentheses, uses class name - - @R.register_op # without parentheses, uses class name - - @R.register_op("custom_name") # with custom name - - Args: - name: Either a string name for the class, or the class itself when used without parentheses - register_type: The type of registry (LLM, EMBEDDING_MODEL, VECTOR_STORE, etc.) - - Returns: - Either a decorator function or the registered class itself - - Example: - @R.register("my_llm", RegistryEnum.LLM) - class MyLLM(BaseLLM): - pass - """ - if inspect.isclass(name): - # Used without parentheses: @R.register_op - self.registry_dict[register_type][name.__name__] = name - return name - else: - # Used with parentheses: @R.register_op() or @R.register_op("name") - def decorator(cls): - key = name if isinstance(name, str) and name else cls.__name__ - self.registry_dict[register_type][key] = cls - return cls - - return decorator - - 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. - - Args: - name: The registration name of the class - register_type: The type of registry to search in - - Returns: - The registered class (not an instance, but the class itself) - - Raises: - AssertionError: If the class is not found in the registry - """ - 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_llm_class(self, name: str): - """Get the LLM class registered under the given name.""" - return self.get_model_class(name, RegistryEnum.LLM) - - 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_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) - - -# Export a global singleton instance for easy access across the application -# This is the primary way to access the registry throughout the codebase -R = Registry() diff --git a/reme_ai/core/context/registry_factory.py b/reme_ai/core/context/registry_factory.py new file mode 100644 index 00000000..28cb1c82 --- /dev/null +++ b/reme_ai/core/context/registry_factory.py @@ -0,0 +1,45 @@ +"""Module providing a registry class for managing class-to-name mappings via decorators.""" + +import inspect +from typing import Callable, TypeVar + +from .base_context import BaseContext +from ..utils import singleton + +T = TypeVar("T") + + +class Registry(BaseContext): + """A registry container that uses decorators to map and store class references.""" + + def register(self, name: str | type = "") -> Callable[[type[T]], type[T]] | type[T]: + """Return a decorator that registers a class under a specific name in the registry.""" + if inspect.isclass(name): + self[name.__name__] = name + return name + + else: + + def decorator(cls): + key: str = name if isinstance(name, str) and name else cls.__name__ + self[key] = cls + return cls + + return decorator + + +@singleton +class RegistryFactory: + """A factory class for creating registries.""" + + def __init__(self): + self.llm = Registry() + self.embedding_model = Registry() + self.vector_store = Registry() + self.op = Registry() + self.flow = Registry() + self.service = Registry() + self.token_counter = Registry() + + +R = RegistryFactory() diff --git a/reme_ai/core/context/runtime_context.py b/reme_ai/core/context/runtime_context.py deleted file mode 100644 index d7112e1c..00000000 --- a/reme_ai/core/context/runtime_context.py +++ /dev/null @@ -1,79 +0,0 @@ -"""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, StreamChunk - - -class RuntimeContext(BaseContext): - """Context for execution state, response metadata, and stream queues.""" - - def __init__( - self, - response: Response | None = None, - stream_queue: asyncio.Queue | None = None, - **kwargs, - ): - """Initialize the context with optional response and queue.""" - super().__init__(**kwargs) - self.response = response or Response() - self.stream_queue = stream_queue - - @classmethod - def from_context(cls, context: "RuntimeContext | None" = None, **kwargs) -> "RuntimeContext": - """Create a new context from an existing instance or keywords.""" - if context is None: - return cls(**kwargs) - - context.update(kwargs) - return context - - async def _enqueue(self, chunk: StreamChunk) -> None: - """Internal helper to put a chunk into the queue if it exists.""" - if self.stream_queue: - await self.stream_queue.put(chunk) - - async def add_stream_string(self, chunk: str, chunk_type: ChunkEnum) -> "RuntimeContext": - """Enqueue a stream chunk from a raw string and type.""" - await self._enqueue(StreamChunk(chunk_type=chunk_type, chunk=chunk)) - return self - - async def add_stream_chunk(self, stream_chunk: StreamChunk) -> "RuntimeContext": - """Enqueue an existing stream chunk.""" - await self._enqueue(stream_chunk) - return self - - async def add_stream_done(self) -> "RuntimeContext": - """Enqueue a termination chunk to signal the end of the stream.""" - await self._enqueue(StreamChunk(chunk_type=ChunkEnum.DONE, chunk="", done=True)) - return self - - def add_response_error(self, e: Exception) -> "RuntimeContext": - """Record an exception into the response object.""" - self.response.success = False - self.response.answer = str(e) - return self - - def apply_mapping(self, mapping: dict[str, str]) -> "RuntimeContext": - """Copy internal values based on a source-to-target key map.""" - if not mapping: - return self - - for source, target in mapping.items(): - if source in self: - self[target] = self[source] - return self - - def validate_required_keys(self, required_keys: dict[str, bool], context_name: str = "context") -> "RuntimeContext": - """Ensure all required keys are present in the context. - - Args: - required_keys: Dictionary mapping key names to boolean indicating if required - context_name: Name of the context for error messages (e.g., operator name) - """ - for key, is_required in required_keys.items(): - if is_required and key not in self: - raise ValueError(f"{context_name}: missing required input '{key}'") - return self diff --git a/reme_ai/core/schema/__init__.py b/reme_ai/core/schema/__init__.py new file mode 100644 index 00000000..b7b73719 --- /dev/null +++ b/reme_ai/core/schema/__init__.py @@ -0,0 +1,42 @@ +"""schema""" + +from .memory_node import MemoryNode +from .message import ContentBlock, Message, Trajectory +from .request import Request +from .response import Response +from .service_config import ( + CmdConfig, + EmbeddingModelConfig, + FlowConfig, + HttpConfig, + LLMConfig, + MCPConfig, + ServiceConfig, + TokenCounterConfig, + VectorStoreConfig, +) +from .stream_chunk import StreamChunk +from .tool_call import ToolAttr, ToolCall +from .vector_node import VectorNode + +__all__ = [ + "MemoryNode", + "ContentBlock", + "EmbeddingModelConfig", + "FlowConfig", + "HttpConfig", + "LLMConfig", + "MCPConfig", + "Message", + "Request", + "Response", + "ServiceConfig", + "StreamChunk", + "TokenCounterConfig", + "Trajectory", + "ToolAttr", + "ToolCall", + "VectorNode", + "VectorStoreConfig", + "CmdConfig", +] diff --git a/reme_ai/core/schema/memory_node.py b/reme_ai/core/schema/memory_node.py new file mode 100644 index 00000000..67ed7c43 --- /dev/null +++ b/reme_ai/core/schema/memory_node.py @@ -0,0 +1,198 @@ +"""Memory schema module for the ReMe AI system. + +This module defines the MemoryNode class for storing and retrieving +memories in the ReMe system. +""" + +import datetime +import hashlib +from typing import Any + +from pydantic import BaseModel, Field, model_validator + +from .vector_node import VectorNode +from ..enumeration import MemoryType + + +def get_now_time() -> str: + """Get current timestamp in YYYY-MM-DD HH:MM:SS format. + + Returns: + str: Current timestamp string in format 'YYYY-MM-DD HH:MM:SS'. + """ + return datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") + + +# Length of the memory ID (first N characters of SHA-256 hash) +MEMORY_ID_LENGTH: int = 16 + + +class MemoryNode(BaseModel): + """Memory node for storing memories in the ReMe system. + + Attributes: + memory_id: Unique identifier, auto-generated from content hash. + memory_type: Type of memory (e.g., SUMMARY, PERSONAL). + memory_target: Target or topic this memory relates to. + when_to_use: Condition description for vector retrieval. + content: Actual memory content. + ref_memory_id: Reference to related raw history memory. + time_created: Creation timestamp. + time_modified: Last modification timestamp. + author: Author or source of this memory. + score: Relevance or importance score. + metadata: Additional metadata for extensibility. + """ + + memory_id: str = Field(default="", description="Unique memory identifier") + memory_type: MemoryType = Field(default=..., description="Type of memory") + memory_target: str = Field(default="", description="Target or topic of the memory") + when_to_use: str = Field(default="", description="Condition description for vector retrieval") + content: str = Field(default="", description="Actual memory content") + ref_memory_id: str = Field(default="", description="Reference to related raw history memory ID") + + time_created: str = Field(default_factory=get_now_time, description="Creation timestamp") + time_modified: str = Field(default_factory=get_now_time, description="Last modification timestamp") + author: str = Field(default="", description="Author or source of the memory") + score: float = Field(default=0, description="Relevance or importance score") + + metadata: dict[str, Any] = Field(default_factory=dict, description="Additional metadata") + + def _update_modified_time(self) -> "MemoryNode": + """Update time_modified to current timestamp. + + Returns: + Self: Returns self for method chaining. + """ + self.time_modified = get_now_time() + return self + + def _update_memory_id(self) -> "MemoryNode": + """Generate memory_id from SHA-256 hash of content. + + Takes the first MEMORY_ID_LENGTH characters of the hash. + + Returns: + Self: Returns self for method chaining. + """ + if not self.content: + return self + + hash_obj = hashlib.sha256(self.content.encode("utf-8")) + hex_dig = hash_obj.hexdigest() + self.memory_id = hex_dig[:MEMORY_ID_LENGTH] + return self + + @model_validator(mode="after") + def _update_after_init(self) -> "MemoryNode": + """Post-initialization validator. + + Auto-generates memory_id from content if not provided. + + Returns: + Self: Returns self for method chaining. + """ + if not self.memory_id: + self._update_memory_id() + return self + + def __setattr__(self, name: str, value): + """Auto-update timestamps and memory_id when content or when_to_use changes. + + Args: + name: Attribute name being set. + value: New value for the attribute. + """ + should_update: bool = name in ("when_to_use", "content") and getattr(self, name, None) != value + super().__setattr__(name, value) + if should_update: + self._update_modified_time() + if name == "content": + self._update_memory_id() + + def to_vector_node(self) -> VectorNode: + """Convert to VectorNode for vector storage. + + When when_to_use is set, use it as vector content and store content in metadata. + When when_to_use is empty, use content as vector content directly. + + Returns: + VectorNode: Vector node representation of this memory. + """ + # Build base metadata (shared fields) + metadata: dict[str, Any] = { + "memory_type": self.memory_type.value, + "memory_target": self.memory_target, + "ref_memory_id": self.ref_memory_id, + "time_created": self.time_created, + "time_modified": self.time_modified, + "author": self.author, + "score": self.score, + **self.metadata, + } + + if self.when_to_use: + # Use when_to_use for vector embedding, store content in metadata + vector_content = self.when_to_use + metadata["content"] = self.content + else: + # Use content directly for vector embedding + vector_content = self.content + + return VectorNode( + vector_id=self.memory_id, + content=vector_content, + metadata=metadata, + ) + + @classmethod + def from_vector_node(cls, node: VectorNode) -> "MemoryNode": + """Reconstruct MemoryNode from VectorNode. + + Reverses the to_vector_node conversion: + - If metadata contains 'content': node.content -> when_to_use, metadata['content'] -> content + - Otherwise: node.content -> content, when_to_use remains empty + + Args: + node: VectorNode containing memory data. + + Returns: + Self: Reconstructed MemoryNode instance. + + Raises: + ValueError: If memory_type in metadata is invalid. + """ + metadata = node.metadata.copy() + memory_type_str = metadata.pop("memory_type", None) + + try: + memory_type: MemoryType = MemoryType(memory_type_str) + except ValueError as e: + raise ValueError( + f"Invalid memory_type '{memory_type_str}' in VectorNode metadata. " + f"Valid types are: {[t.value for t in MemoryType]}", + ) from e + + # Restore when_to_use and content based on metadata structure + if "content" in metadata: + # Original had when_to_use set + when_to_use = node.content + content = metadata.pop("content", "") + else: + # Original had empty when_to_use + when_to_use = "" + content = node.content + + return cls( + memory_id=node.vector_id, + memory_type=memory_type, + memory_target=metadata.pop("memory_target", ""), + when_to_use=when_to_use, + content=content, + ref_memory_id=metadata.pop("ref_memory_id", ""), + time_created=metadata.pop("time_created", ""), + time_modified=metadata.pop("time_modified", ""), + author=metadata.pop("author", ""), + score=metadata.pop("score", 0), + metadata=metadata, + ) diff --git a/reme_ai/core/schema/message.py b/reme_ai/core/schema/message.py new file mode 100644 index 00000000..321dd8c8 --- /dev/null +++ b/reme_ai/core/schema/message.py @@ -0,0 +1,165 @@ +"""Data models for multi-modal conversation history and LLM interaction trajectories.""" + +import datetime +import json +import re + +from pydantic import BaseModel, ConfigDict, Field, model_validator + +from .tool_call import ToolCall +from ..enumeration import Role + + +class ContentBlock(BaseModel): + """ + Individual unit of multi-modal content like text, images, or video. + examples: + { + "type": "image_url", + "image_url": { + "url": "https://img.alicdn.com/imgextra/i1/O1CN01gDEY8M1W114Hi3XcN_!!6000000002727-0-tps-1024-406.jpg" + }, + } + + { + "type": "video", + "video": [ + "https://help-static-aliyun-doc.aliyuncs.com/file-manage-files/zh-CN/20241108/xzsgiz/football1.jpg", + "https://help-static-aliyun-doc.aliyuncs.com/file-manage-files/zh-CN/20241108/tdescd/football2.jpg", + "https://help-static-aliyun-doc.aliyuncs.com/file-manage-files/zh-CN/20241108/zefdja/football3.jpg", + "https://help-static-aliyun-doc.aliyuncs.com/file-manage-files/zh-CN/20241108/aedbqh/football4.jpg", + ], + } + + { + "type": "text", + "text": "How do you solve this problem?" + } + """ + + model_config = ConfigDict(extra="allow") + + type: str = Field(default="") + content: str | dict | list = Field(default="") + + @model_validator(mode="before") + @classmethod + def init_block(cls, data: dict) -> dict: + """Dynamically maps the type-specific key to the content field.""" + content_type = data.get("type", "") + if content_type and content_type in data: + data["content"] = data[content_type] + return data + + def simple_dump(self) -> dict: + """Serializes the block into an API-compatible dictionary format.""" + return { + "type": self.type, + self.type: self.content, + **self.model_extra, + } + + +class Message(BaseModel): + """Data model for a single dialogue entry including roles and tool interactions.""" + + name: str | None = Field(default=None) + role: Role = Field(default=Role.USER) + content: str | list[ContentBlock] = Field(default="") + reasoning_content: str = Field(default="") + tool_calls: list[ToolCall] = Field(default_factory=list) + tool_call_id: str = Field(default="") + time_created: str = Field(default_factory=lambda: datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")) + metadata: dict = Field(default_factory=dict) + + def dump_content(self) -> str | list[dict]: + """Returns content as a raw string or a list of serialized blocks.""" + if isinstance(self.content, str): + return self.content + return [block.simple_dump() for block in self.content] + + def simple_dump( + self, + add_name: bool = False, + add_reasoning: bool = True, + add_time_created: bool = False, + add_metadata: bool = False, + enable_json_dump: bool = False, + ) -> dict | str: + """Transforms the message into a simplified dictionary for standard APIs.""" + result = {} + if add_name and self.name: + result["name"] = self.name + + result["role"] = self.role.value + result["content"] = self.dump_content() + + if add_reasoning and self.reasoning_content: + result["reasoning_content"] = self.reasoning_content + + if self.tool_calls: + result["tool_calls"] = [tc.simple_output_dump() for tc in self.tool_calls] + + if self.tool_call_id: + result["tool_call_id"] = self.tool_call_id + + if add_time_created: + result["time_created"] = self.time_created + + if add_metadata: + result["metadata"] = self.metadata + + if enable_json_dump: + return json.dumps(result, ensure_ascii=False) + else: + return result + + def format_message( + self, + index: int | None = None, + add_time: bool = False, + use_name: bool = False, + add_reasoning: bool = True, + add_tools: bool = True, + strip_markdown_headers: bool = False, + ) -> str: + """Generates a human-readable string representation of the message.""" + prefix = f"round{index} " if index is not None else "" + time_str = f"[{self.time_created}] " if add_time else "" + header = f"{self.name or self.role.value if use_name else self.role.value}:" + + lines = [f"{prefix}{time_str}{header}"] + + def strip_md_func(line): + if strip_markdown_headers: + line = re.sub(r"\n##+ +", "\n", line) + return line + + if add_reasoning and self.reasoning_content: + lines.append(self.reasoning_content) + + if isinstance(self.content, str): + lines.append(strip_md_func(self.content)) + + elif isinstance(self.content, list): + for block in self.content: + text = ( + block.content if isinstance(block.content, str) else json.dumps(block.content, ensure_ascii=False) + ) + text = str(text) + lines.append(strip_md_func(text)) + + if add_tools and self.tool_calls: + for tc in self.tool_calls: + lines.append(f" - tool_call={tc.name} params={tc.arguments}") + + return " ".join(lines).strip() + + +class Trajectory(BaseModel): + """Sequence of messages representing a full conversation session and its evaluation.""" + + task_id: str = Field(default="") + messages: list[Message] = Field(default_factory=list) + score: float = Field(default=0.0) + metadata: dict = Field(default_factory=dict) diff --git a/reme_ai/core/schema/request.py b/reme_ai/core/schema/request.py new file mode 100644 index 00000000..ece942b9 --- /dev/null +++ b/reme_ai/core/schema/request.py @@ -0,0 +1,11 @@ +"""Defines the data structure for processing incoming user requests and message history.""" + +from pydantic import Field, BaseModel, ConfigDict + + +class Request(BaseModel): + """Represents a structured request payload containing a query, message list, and metadata.""" + + model_config = ConfigDict(extra="allow") + + metadata: dict = Field(default_factory=dict) diff --git a/reme_ai/core/schema/response.py b/reme_ai/core/schema/response.py new file mode 100644 index 00000000..3104bc6e --- /dev/null +++ b/reme_ai/core/schema/response.py @@ -0,0 +1,11 @@ +"""Defines the standardized data structure for model output responses.""" + +from pydantic import Field, BaseModel + + +class Response(BaseModel): + """Represents a structured response containing the execution result, status, and metadata.""" + + answer: str | dict | list = Field(default="") + success: bool = Field(default=True) + metadata: dict = Field(default_factory=dict) diff --git a/reme_ai/core/schema/service_config.py b/reme_ai/core/schema/service_config.py new file mode 100644 index 00000000..e1ae9df7 --- /dev/null +++ b/reme_ai/core/schema/service_config.py @@ -0,0 +1,113 @@ +"""Configuration schemas for service components using Pydantic models.""" + +import os +from typing import Dict, List + +from pydantic import BaseModel, Field, ConfigDict + +from .tool_call import ToolCall + + +class MCPConfig(BaseModel): + """Configuration for Model Context Protocol transport and network settings.""" + + model_config = ConfigDict(extra="allow") + + transport: str = Field(default="stdio") + host: str = Field(default="0.0.0.0") + port: int = Field(default=8001) + + +class HttpConfig(BaseModel): + """Configuration for the HTTP server interface and connection lifecycle.""" + + model_config = ConfigDict(extra="allow") + + host: str = Field(default="0.0.0.0") + port: int = Field(default=8001) + timeout_keep_alive: int = Field(default=3600) + limit_concurrency: int = Field(default=1000) + + +class CmdConfig(BaseModel): + """Configuration for command-line flow execution parameters.""" + + model_config = ConfigDict(extra="allow") + + flow: str = Field(default="") + + +class FlowConfig(ToolCall): + """Configuration for workflow execution, caching, and error handling.""" + + model_config = ConfigDict(extra="allow") + + flow_content: str = Field(default="") + stream: bool = Field(default=False) + raise_exception: bool = Field(default=True) + enable_cache: bool = Field(default=False) + cache_path: str = Field(default="cache/flow") + cache_expire_hours: float = Field(default=0.1) + + +class LLMConfig(BaseModel): + """Configuration for Large Language Model backend and model identification.""" + + model_config = ConfigDict(extra="allow") + + backend: str = Field(default="") + model_name: str = Field(default="") + + +class EmbeddingModelConfig(BaseModel): + """Configuration for embedding model backends and identity.""" + + model_config = ConfigDict(extra="allow") + + backend: str = Field(default="") + model_name: str = Field(default="") + + +class VectorStoreConfig(BaseModel): + """Configuration for vector database storage and associated embeddings.""" + + model_config = ConfigDict(extra="allow") + + backend: str = Field(default="local") + collection_name: str = Field(default="reme") + embedding_model: str = Field(default="default") + + +class TokenCounterConfig(BaseModel): + """Configuration for token counting services and model mapping.""" + + model_config = ConfigDict(extra="allow") + + backend: str = Field(default="base") + model_name: str = Field(default="") + + +class ServiceConfig(BaseModel): + """Root configuration schema aggregating all service-level settings and components.""" + + model_config = ConfigDict(extra="allow") + + backend: str = Field(default="") + app_name: str = Field(default=os.getenv("APP_NAME", "ReMe")) + enable_logo: bool = Field(default=True) + language: str = Field(default="") + thread_pool_max_workers: int = Field(default=16) + ray_max_workers: int = Field(default=-1) + init_logger: bool = Field(default=True) + disabled_flows: List[str] = Field(default_factory=list) + enabled_flows: List[str] = Field(default_factory=list) + mcp_servers: Dict[str, dict] = Field(default_factory=dict) + + mcp: MCPConfig = Field(default_factory=MCPConfig) + http: HttpConfig = Field(default_factory=HttpConfig) + cmd: CmdConfig = Field(default_factory=CmdConfig) + flow: Dict[str, FlowConfig] = Field(default_factory=dict) + llm: Dict[str, LLMConfig] = Field(default_factory=dict) + embedding_model: Dict[str, EmbeddingModelConfig] = Field(default_factory=dict) + vector_store: Dict[str, VectorStoreConfig] = Field(default_factory=dict) + token_counter: Dict[str, TokenCounterConfig] = Field(default_factory=dict) diff --git a/reme_ai/core/schema/stream_chunk.py b/reme_ai/core/schema/stream_chunk.py new file mode 100644 index 00000000..764981fd --- /dev/null +++ b/reme_ai/core/schema/stream_chunk.py @@ -0,0 +1,14 @@ +"""Defines the data structure for individual data packets in a streaming response.""" + +from pydantic import Field, BaseModel + +from ..enumeration import ChunkEnum + + +class StreamChunk(BaseModel): + """Represents a single chunk of streamed data including its type, content, and completion status.""" + + chunk_type: ChunkEnum = Field(default=ChunkEnum.ANSWER) + chunk: str | dict | list = Field(default="") + done: bool = Field(default=False) + metadata: dict = Field(default_factory=dict) diff --git a/reme_ai/core/schema/tool_call.py b/reme_ai/core/schema/tool_call.py new file mode 100644 index 00000000..70355d60 --- /dev/null +++ b/reme_ai/core/schema/tool_call.py @@ -0,0 +1,226 @@ +"""MCP Tool Schema definitions for recursive JSON Schema representation.""" + +import json +from typing import Any, Dict, List, Optional, Union + +from mcp.types import Tool +from pydantic import BaseModel, ConfigDict, Field, model_validator, field_validator + +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: str = Field(default=str(JsonSchemaEnum.STRING), 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") + properties: Optional[Dict[str, "ToolAttr"]] = Field(default=None, description="Child properties for objects") + 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") + + @field_validator("type") + @classmethod + def validate_type_is_valid_enum(cls, v: str) -> str: + """Validates that the provided type string exists within JsonSchemaEnum values.""" + valid_types = [str(e) for e in JsonSchemaEnum] + + if v not in valid_types: + raise ValueError(f"Invalid type: '{v}'. Must be one of {valid_types}") + return v + + def simple_input_dump(self) -> dict: + """Serializes the attribute into a standard JSON Schema dictionary.""" + res: dict = {"type": self.type} + if self.description: + res["description"] = self.description + if self.enum: + res["enum"] = self.enum + + if self.type == "object" and self.properties is not None: + res["properties"] = { + k: v.simple_input_dump() if isinstance(v, ToolAttr) else v for k, v in self.properties.items() + } + if self.required is not None: + res["required"] = self.required + + if self.type == "array" and self.items is not None: + res["items"] = self.items.simple_input_dump() if isinstance(self.items, ToolAttr) else self.items + + return res + + +# Enable recursive type resolution +ToolAttr.model_rebuild() + + +class ToolCall(BaseModel): + """ + Model representing a tool definition and its call structure. + Supports parsing from standard JSON Schema formats and converting to MCP Tool objects. + input: + { + "type": "function", + "function": { + "name": "get_current_weather", + "description": "It is very useful when you want to check the weather of a specified city.", + "parameters": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "Cities or counties, such as Beijing, Hangzhou, Yuhang District, etc.", + } + }, + "required": ["location"] + } + } + } + output: + { + "index": 0, + "id": "call_6596dafa2a6a46f7a217da", + "function": { + "arguments": "{\"location\": \"Beijing\"}", + "name": "get_current_weather" + }, + "type": "function", + } + """ + + index: int = 0 + id: str = "" + type: str = "function" + name: str = "" + description: str = "" + + arguments: str = Field(default="", description="JSON string of tool execution arguments") + + parameters: ToolAttr = Field( + default_factory=lambda: ToolAttr(type="object", properties={}, required=[]), + description="Specification for input parameters", + ) + + output: ToolAttr = Field( + default_factory=lambda: ToolAttr(type="object", properties={}), + description="Specification for the execution result (Schema)", + ) + + @model_validator(mode="before") + @classmethod + def init_tool_call(cls, data: dict) -> dict: + """Initializes the model by parsing tool-specific body data.""" + data = data.copy() + t_type = data.get("type", "function") + body = data.get(t_type, {}) + + # Extract basic metadata + data["name"] = body.get("name", data.get("name", "")) + data["arguments"] = body.get("arguments", data.get("arguments", "")) + data["description"] = body.get("description", data.get("description", "")) + + # Handle parameters mapping + if "parameters" in body: + params = body["parameters"] + # If parameters is already a dict, ensure it matches ToolAttr structure + if isinstance(params, dict): + data["parameters"] = ToolAttr(**params) + + # Handle output mapping (if provided in source) + if "output" in body and isinstance(body["output"], dict): + data["output"] = ToolAttr(**body["output"]) + + return data + + def simple_input_dump(self) -> dict: + """Returns a standardized tool definition dictionary.""" + return { + "type": self.type, + self.type: { + "name": self.name, + "description": self.description, + "parameters": self.parameters.simple_input_dump(), + }, + } + + def simple_output_dump(self) -> dict: + """Convert ToolCall to output format dictionary for API responses.""" + return { + "index": self.index, + "id": self.id, + self.type: { + "arguments": self.arguments, + "name": self.name, + }, + "type": self.type, + } + + @property + def argument_dict(self) -> dict: + """Parse and return arguments as a dictionary.""" + return json.loads(self.arguments) + + def check_argument(self) -> bool: + """Check if arguments can be parsed as valid JSON.""" + try: + _ = self.argument_dict + return True + except Exception: + return False + + def sanitize_and_check_argument(self) -> bool: + """ + Attempt to sanitize and validate arguments JSON. + Common issues from LLM streaming: + - Extra closing brackets: }]}] -> }] + - Missing closing brackets + - Trailing commas + """ + if not self.arguments or not self.arguments.strip(): + return False + + try: + # First try parsing as-is + _ = json.loads(self.arguments) + return True + except json.JSONDecodeError: + pass + + # Try to fix common issues + sanitized = self.arguments.strip() + + # Remove trailing extra brackets/braces + # Pattern: if it ends with multiple closing chars, try removing extras + while len(sanitized) > 1: + try: + json.loads(sanitized) + self.arguments = sanitized # Update with sanitized version + return True + except json.JSONDecodeError: + # Try removing last character + if sanitized[-1] in "]}": + sanitized = sanitized[:-1].rstrip() + else: + break + + return False + + @classmethod + def from_mcp_tool(cls, tool: Tool) -> "ToolCall": + """Creates a ToolCall instance from an MCP Tool object.""" + # MCP Tool inputSchema maps directly to our parameters ToolAttr + return cls( + name=tool.name, + description=tool.description or "", + parameters=ToolAttr(**tool.inputSchema), + ) + + def to_mcp_tool(self) -> Tool: + """Converts the instance back into an MCP Tool object.""" + return Tool( + name=self.name, + description=self.description, + inputSchema=self.parameters.simple_input_dump(), + ) diff --git a/reme_ai/core/schema/vector_node.py b/reme_ai/core/schema/vector_node.py new file mode 100644 index 00000000..937ef4be --- /dev/null +++ b/reme_ai/core/schema/vector_node.py @@ -0,0 +1,15 @@ +"""Defines the data structure for individual vector embedding nodes within a retrieval system.""" + +from typing import List, Dict +from uuid import uuid4 + +from pydantic import BaseModel, Field + + +class VectorNode(BaseModel): + """Represents a discrete unit of text content paired with its corresponding vector embedding and metadata.""" + + vector_id: str = Field(default_factory=lambda: uuid4().hex) + content: str = Field(default="") + vector: List[float] | None = Field(default=None) + metadata: Dict[str, str | bool | int | float] = Field(default_factory=dict) diff --git a/reme_ai/core/utils/__init__.py b/reme_ai/core/utils/__init__.py new file mode 100644 index 00000000..ce9d9b2d --- /dev/null +++ b/reme_ai/core/utils/__init__.py @@ -0,0 +1,7 @@ +"""utils""" + +from .singleton import singleton + +__all__ = [ + "singleton", +] diff --git a/reme_ai/core_old/utils/singleton.py b/reme_ai/core/utils/singleton.py similarity index 100% rename from reme_ai/core_old/utils/singleton.py rename to reme_ai/core/utils/singleton.py