scriptoriumadmin/ReMe
1
0
Fork 0
mirror of https://github.com/agentscope-ai/ReMe.git synced 2026-08-28 05:25:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 11

feat(core): refactor context management and add schema definitions

Browse source
This commit is contained in:
jinli.yl 2026-01-21 17:05:56 +08:00
parent c7fc8255b1
commit 32ff65ca0d
16 changed files with 897 additions and 287 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
reme_ai/core/context/__init__.py
View file

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

106
reme_ai/core/context/prompt_handler.py
View file

@ -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)})"

143
reme_ai/core/context/registry.py
View file

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

45
reme_ai/core/context/registry_factory.py Normal file
View file

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

79
reme_ai/core/context/runtime_context.py
View file

@ -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

42
reme_ai/core/schema/__init__.py Normal file
View file

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

198
reme_ai/core/schema/memory_node.py Normal file
View file

@ -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,
)

165
reme_ai/core/schema/message.py Normal file
View file

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

11
reme_ai/core/schema/request.py Normal file
View file

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

11
reme_ai/core/schema/response.py Normal file
View file

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

113
reme_ai/core/schema/service_config.py Normal file
View file

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

14
reme_ai/core/schema/stream_chunk.py Normal file
View file

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

226
reme_ai/core/schema/tool_call.py Normal file
View file

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

15
reme_ai/core/schema/vector_node.py Normal file
View file

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

7
reme_ai/core/utils/__init__.py Normal file
View file

@ -0,0 +1,7 @@
"""utils"""
from .singleton import singleton
__all__ = [
"singleton",
]

0
reme_ai/core_old/utils/singleton.py → reme_ai/core/utils/singleton.py
View file