diff --git a/reme_ai/core/application.py b/reme_ai/core/application.py new file mode 100644 index 00000000..e15f4fcf --- /dev/null +++ b/reme_ai/core/application.py @@ -0,0 +1,208 @@ +"""Main application module for managing ReMe AI service lifecycle and flow execution.""" + +import asyncio +import os + +from .context import C +from .flow import BaseFlow +from .schema import ServiceConfig, Response +from .utils import PydanticConfigParser, init_logger, execute_stream_task, run_coro_safely + + +class Application: + """ + Main application class for managing the lifecycle of ReMe AI services. + + Handles initialization, configuration, service management, and flow execution + for both synchronous and asynchronous contexts. + """ + + def __init__( + self, + *args, + llm_api_key: str | None = None, + llm_api_base: str | None = None, + embedding_api_key: str | None = None, + embedding_api_base: str | None = None, + service_config: ServiceConfig | None = None, + parser: type[PydanticConfigParser] | None = None, + config_path: str | None = None, + enable_logo: bool = True, + llm: dict | None = None, + embedding_model: dict | None = None, + vector_store: dict | None = None, + token_counter: dict | None = None, + **kwargs, + ): + """ + Initialize the Application with configuration settings. + + Args: + llm_api_key: API key for LLM service + llm_api_base: Base URL for LLM service + embedding_api_key: API key for embedding service + embedding_api_base: Base URL for embedding service + service_config: Pre-built service configuration object + parser: Custom parser class for configuration (defaults to PydanticConfigParser) + config_path: Path to configuration file + enable_logo: Whether to display the ReMe logo on startup + llm: LLM configuration dictionary + embedding_model: Embedding model configuration dictionary + vector_store: Vector store configuration dictionary + token_counter: Token counter configuration dictionary + **kwargs: Additional configuration arguments + """ + self._update_env("REME_LLM_API_KEY", llm_api_key) + self._update_env("REME_LLM_BASE_URL", llm_api_base) + self._update_env("REME_EMBEDDING_API_KEY", embedding_api_key) + self._update_env("REME_EMBEDDING_BASE_URL", embedding_api_base) + + init_logger() + + # Use default parser if not provided + parser_class = parser if parser is not None else PydanticConfigParser + self.parser = parser_class(ServiceConfig) + + if service_config is None: + input_args = [] + if config_path: + input_args.append(f"config={config_path}") + if args: + input_args.extend(args) + if kwargs: + input_args.extend([f"{k}={v}" for k, v in kwargs.items()]) + service_config = self.parser.parse_args(*input_args) + + C.service_config = service_config + + if llm: + C.update_section_config("llm", **llm) + if embedding_model: + C.update_section_config("embedding_model", **embedding_model) + if vector_store: + C.update_section_config("vector_store", **vector_store) + if token_counter: + C.update_section_config("token_counter", **token_counter) + C.service_config.enable_logo = enable_logo + C.print_logo() + + @staticmethod + def _update_env(key: str, value: str | None) -> None: + """Update environment variable if value is provided.""" + if value: + os.environ[key] = value + + @staticmethod + async def start() -> None: + """Initialize the service context and prepare external MCP servers.""" + C.initialize_service_context() + await C.prepare_mcp_servers() + + @staticmethod + def start_sync() -> None: + """Synchronous version of start().""" + C.initialize_service_context() + run_coro_safely(C.prepare_mcp_servers()) + + @staticmethod + async def stop(wait_thread_pool: bool = True, wait_ray: bool = True) -> None: + """ + Stop the application and cleanup resources. + + Args: + wait_thread_pool: Whether to wait for thread pool shutdown + wait_ray: Whether to wait for Ray shutdown + """ + await C.close() + C.shutdown_thread_pool(wait=wait_thread_pool) + C.shutdown_ray(wait=wait_ray) + + @staticmethod + def stop_sync(wait_thread_pool: bool = True, wait_ray: bool = True) -> None: + """Synchronous version of stop().""" + C.close_sync() + C.shutdown_thread_pool(wait=wait_thread_pool) + C.shutdown_ray(wait=wait_ray) + + async def __aenter__(self): + """Async context manager entry.""" + await self.start() + return self + + def __enter__(self): + """Context manager entry.""" + self.start_sync() + return self + + async def __aexit__(self, exc_type=None, exc_val=None, exc_tb=None): + """Async context manager exit.""" + await self.stop() + return False + + def __exit__(self, exc_type=None, exc_val=None, exc_tb=None): + """Context manager exit.""" + self.stop_sync() + return False + + @staticmethod + async def execute_flow(name: str, **kwargs) -> Response: + """ + Execute a flow asynchronously. + + Args: + name: Name of the flow to execute + **kwargs: Arguments to pass to the flow + + Returns: + Response object from the flow execution + """ + flow: BaseFlow = C.get_flow(name) + return await flow.call(**kwargs) + + @staticmethod + def execute_flow_sync(name: str, **kwargs) -> Response: + """ + Execute a flow synchronously. + + Args: + name: Name of the flow to execute + **kwargs: Arguments to pass to the flow + + Returns: + Response object from the flow execution + """ + flow: BaseFlow = C.get_flow(name) + return flow.call_sync(**kwargs) + + @staticmethod + async def execute_stream_flow(name: str, **kwargs): + """ + Execute a streaming flow asynchronously. + + Args: + name: Name of the streaming flow to execute + **kwargs: Arguments to pass to the flow + + Yields: + Stream chunks from the flow execution + + Raises: + AssertionError: If the flow is not configured for streaming + """ + flow: BaseFlow = C.get_flow(name) + assert flow.stream is True, "non-stream flow is not supported in execute_stream_flow!" + stream_queue = asyncio.Queue() + task = asyncio.create_task(flow.call(stream_queue=stream_queue, **kwargs)) + + async for chunk in execute_stream_task( + queue=stream_queue, + task=task, + flow_name=name, + as_bytes=False, + ): + yield chunk + + @staticmethod + def run_service(): + """Run the configured service (HTTP, MCP, or CMD).""" + C.get_service().run() diff --git a/reme_ai/core/config/__init__.py b/reme_ai/core/config/__init__.py new file mode 100644 index 00000000..a8d954d2 --- /dev/null +++ b/reme_ai/core/config/__init__.py @@ -0,0 +1,5 @@ +"""config""" + +from .reme_config_parser import ReMeConfigParser + +__all__ = ["ReMeConfigParser"] diff --git a/reme_ai/core/config/default.yaml b/reme_ai/core/config/default.yaml new file mode 100644 index 00000000..cd04da1d --- /dev/null +++ b/reme_ai/core/config/default.yaml @@ -0,0 +1,34 @@ +backend: http +thread_pool_max_workers: 64 + +mcp: + transport: sse + host: "0.0.0.0" + port: 8001 + +http: + host: "0.0.0.0" + port: 8002 + timeout_keep_alive: 600 + limit_concurrency: 64 + +llm: + default: + backend: openai + model_name: qwen3-30b-a3b-instruct-2507 + temperature: 0.6 + +embedding_model: + default: + backend: openai + model_name: text-embedding-v4 + dimensions: 1024 + +vector_store: + default: + backend: local + embedding_model: default + +token_counter: + default: + backend: base diff --git a/reme_ai/core/config/reme_config_parser.py b/reme_ai/core/config/reme_config_parser.py new file mode 100644 index 00000000..798235b2 --- /dev/null +++ b/reme_ai/core/config/reme_config_parser.py @@ -0,0 +1,7 @@ +"""Configuration parser for ReMe framework.""" + +from ..utils import PydanticConfigParser + + +class ReMeConfigParser(PydanticConfigParser): + """Configuration parser for ReMe framework.""" diff --git a/reme_ai/core/context/service_context.py b/reme_ai/core/context/service_context.py index 6d67b695..c5bcde0c 100644 --- a/reme_ai/core/context/service_context.py +++ b/reme_ai/core/context/service_context.py @@ -1,44 +1,80 @@ """Module for managing global service configurations and component registries via a singleton context.""" from concurrent.futures import ThreadPoolExecutor +from typing import TYPE_CHECKING + +from loguru import logger from .base_context import BaseContext from .registry import Registry from ..enumeration import RegistryEnum from ..schema import ServiceConfig -from ..utils import singleton +from ..utils import singleton, print_logo + +if TYPE_CHECKING: + from ..llm import BaseLLM + from ..embedding import BaseEmbeddingModel + from ..vector_store import BaseVectorStore + from ..token_counter import BaseTokenCounter + from ..flow import BaseFlow + from ..service import BaseService @singleton class ServiceContext(BaseContext): - """A singleton container for global application state, thread pools, and component registries.""" + """A singleton container for global application state, thread pools, and component registries. + + This class serves as the central management hub for the entire ReMe application, providing: + - Service configuration management + - Component registration and instantiation (LLMs, embeddings, vector stores, etc.) + - Thread pool and Ray distributed computing management + - MCP (Model Context Protocol) server integration + + The singleton pattern ensures only one instance exists throughout the application lifecycle, + accessible via the global `C` variable exported at the bottom of this module. + """ def __init__(self, **kwargs): - """Initialize the global context with configuration objects and specialized registries.""" + """Initialize the global context with configuration objects and specialized registries. + + Sets up: + - Empty service configuration placeholder + - Thread pool for concurrent operations + - Registry dictionaries for class registration (templates) + - Instance dictionaries for instantiated objects (actual instances) + - MCP server mapping for external tool integration + """ super().__init__(**kwargs) + # Service configuration and runtime settings self.service_config: ServiceConfig | None = None self.language: str = "" self.thread_pool: ThreadPoolExecutor | None = None - self.vector_store_dict: dict[str, dict] = {} - self.mcp_server_tool_call_mapping: dict = {} + # Registry system: stores class definitions for different component types self.registry_dict: dict[RegistryEnum, Registry] = {v: Registry() for v in RegistryEnum.__members__.values()} - self.flow_dict: dict = {} - def _update_config_section(self, section_name: str, update_dict: dict | None): - if not update_dict: - return + # Instance system: stores instantiated objects created from registered classes + self.instance_dict: dict[RegistryEnum, dict] = {v: {} for v in RegistryEnum.__members__.values()} - target_registry = getattr(self.service_config, section_name) - if "default" not in target_registry: - raise KeyError(f"Default `{section_name}` config not found in service_config") - - current_config = target_registry["default"] - target_registry["default"] = current_config.model_copy(update=update_dict, deep=True) + # MCP server mapping: maps server_name -> {tool_name: ToolCall} + self.mcp_server_mapping: dict[str, dict] = {} def register(self, name: str, register_type: RegistryEnum): - """Return a decorator to register a component within a specific registry category.""" + """Return a decorator to register a component within a specific registry category. + + Args: + name: The registration name for the component (used for lookup) + register_type: The type of registry (LLM, EMBEDDING_MODEL, VECTOR_STORE, etc.) + + Returns: + A decorator function that registers the decorated class + + Example: + @C.register("my_llm", RegistryEnum.LLM) + class MyLLM(BaseLLM): + pass + """ return self.registry_dict[register_type].register(name=name) def register_llm(self, name: str = ""): @@ -70,18 +106,29 @@ class ServiceContext(BaseContext): 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.""" + """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_embedding_model_class(self, name: str): - """Get the embedding model class registered under the given name.""" - return self.get_model_class(name, RegistryEnum.EMBEDDING_MODEL) - def get_llm_class(self, name: str): """Get the LLM class registered under the given name.""" return self.get_model_class(name, RegistryEnum.LLM) + def get_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) @@ -102,14 +149,378 @@ class ServiceContext(BaseContext): """Get the token counter class registered under the given name.""" return self.get_model_class(name, RegistryEnum.TOKEN_COUNTER) - def get_vector_store(self, name: str): - """Retrieve a specific vector store instance by name.""" - return self.vector_store_dict[name] + def get_llm(self, name: str) -> "BaseLLM": + """Retrieve a specific LLM instance by name. - def get_flow(self, name: str): - """Retrieve a specific flow instance by name.""" - return self.flow_dict[name] + Args: + name: The name of the LLM instance (typically 'default' or custom name) + + Returns: + The instantiated LLM object + + Raises: + KeyError: If no LLM with the given name exists + """ + return self.instance_dict[RegistryEnum.LLM][name] + + def get_embedding_model(self, name: str) -> "BaseEmbeddingModel": + """Retrieve a specific embedding model instance by name. + + Args: + name: The name of the embedding model instance (typically 'default') + + Returns: + The instantiated embedding model object + + Raises: + KeyError: If no embedding model with the given name exists + """ + return self.instance_dict[RegistryEnum.EMBEDDING_MODEL][name] + + def get_vector_store(self, name: str) -> "BaseVectorStore": + """Retrieve a specific vector store instance by name. + + Args: + name: The name of the vector store instance (typically 'default') + + Returns: + The instantiated vector store object + + Raises: + KeyError: If no vector store with the given name exists + """ + return self.instance_dict[RegistryEnum.VECTOR_STORE][name] + + def get_token_counter(self, name: str) -> "BaseTokenCounter": + """Retrieve a specific token counter instance by name. + + Args: + name: The name of the token counter instance (typically 'default') + + Returns: + The instantiated token counter object + + Raises: + KeyError: If no token counter with the given name exists + """ + return self.instance_dict[RegistryEnum.TOKEN_COUNTER][name] + + def get_flow(self, name: str) -> "BaseFlow": + """Retrieve a specific flow instance by name. + + Args: + name: The name of the flow instance + + Returns: + The instantiated flow object + + Raises: + KeyError: If no flow with the given name exists + """ + return self.instance_dict[RegistryEnum.FLOW][name] + + def get_service(self) -> "BaseService": + """Retrieve the default service instance. + + Returns: + The instantiated service backend (HTTP, MCP, or CMD service) + + Raises: + KeyError: If the default service was not initialized + """ + return self.instance_dict[RegistryEnum.SERVICE]["default"] + + def update_section_config(self, section_name: str, **kwargs): + """Update a specific section of the service config with new values. + + Args: + section_name: Name of the config section (e.g., 'llm', 'embedding_model') + **kwargs: Key-value pairs to update in the default configuration + + Raises: + KeyError: If the default config for the section doesn't exist + + Example: + update_section_config('llm', temperature=0.8, max_tokens=1000) + """ + if not hasattr(self.service_config, section_name) or not kwargs: + return + + section_dict: dict = getattr(self.service_config, section_name) + if "default" not in section_dict: + raise KeyError(f"Default `{section_name}` config not found") + + current_config = section_dict["default"] + section_dict["default"] = current_config.model_copy(update=kwargs, deep=True) + + def initialize_service_context(self): + """Initialize the service context with the configuration. + + This is the main initialization method that sets up all system components in order: + 1. Language settings + 2. Thread pool for concurrent operations + 3. Ray cluster (if configured for distributed computing) + 4. LLM instances + 5. Embedding model instances + 6. Token counter instances + 7. Vector store instances (with their embedding models) + 8. Flow instances (both registered and configured) + 9. Service backend instance + + Note: This method should be called after service_config is set. + """ + self.language = self.service_config.language + self.thread_pool = ThreadPoolExecutor(max_workers=self.service_config.thread_pool_max_workers) + + # Initialize Ray for distributed computing if configured + if self.service_config.ray_max_workers > 1: + import ray + + ray.init(num_cpus=self.service_config.ray_max_workers) + + # Initialize components in dependency order + self._initialize_llm() + self._initialize_embedding_model() + self._initialize_token_counter() + self._initialize_vector_store() # Depends on embedding models + self._initialize_flow() + self._initialize_service() + + def _initialize_llm(self): + """Initialize all configured LLM instances. + + For each LLM configuration: + - Retrieves the corresponding registered LLM class by backend name + - Instantiates it with model_name and additional configuration + - Stores the instance in instance_dict for later retrieval + """ + for name, config in self.service_config.llm.items(): + llm_cls = self.get_llm_class(config.backend) + self.instance_dict[RegistryEnum.LLM][name] = llm_cls(model_name=config.model_name, **config.model_extra) + + def _initialize_embedding_model(self): + """Initialize all configured embedding model instances. + + For each embedding model configuration: + - Retrieves the corresponding registered embedding model class by backend name + - Instantiates it with model_name and additional configuration + - Stores the instance in instance_dict for later retrieval + """ + for name, config in self.service_config.embedding_model.items(): + embedding_model_cls = self.get_embedding_model_class(config.backend) + self.instance_dict[RegistryEnum.EMBEDDING_MODEL][name] = embedding_model_cls( + model_name=config.model_name, + **config.model_extra, + ) + + def _initialize_token_counter(self): + """Initialize all configured token counter instances. + + For each token counter configuration: + - Retrieves the corresponding registered token counter class by backend name + - Instantiates it with model_name and additional configuration + - Stores the instance in instance_dict for later retrieval + """ + for name, config in self.service_config.token_counter.items(): + token_counter_cls = self.get_token_counter_class(config.backend) + self.instance_dict[RegistryEnum.TOKEN_COUNTER][name] = token_counter_cls( + model_name=config.model_name, + **config.model_extra, + ) + + def _initialize_vector_store(self): + """Initialize all configured vector stores with their embedding models. + + For each vector store configuration: + - Retrieves the corresponding registered vector store class by backend name + - Retrieves the associated embedding model instance by name + - Instantiates the vector store with collection name, embedding model, and extra config + - Stores the instance in instance_dict for later retrieval + + Note: This must be called after _initialize_embedding_model() since vector stores + depend on embedding model instances. + """ + for name, config in self.service_config.vector_store.items(): + vector_store_cls = self.get_vector_store_class(config.backend) + self.instance_dict[RegistryEnum.VECTOR_STORE][name] = vector_store_cls( + collection_name=config.collection_name, + embedding_model=self.instance_dict[RegistryEnum.EMBEDDING_MODEL][config.embedding_model], + **config.model_extra, + ) + + def _filter_flows(self, name: str) -> bool: + """Filter flows based on enabled_flows and disabled_flows configuration. + + The filtering logic follows this priority: + 1. If enabled_flows is set: only flows in the list are loaded + 2. Else if disabled_flows is set: all flows except those in the list are loaded + 3. Otherwise: all flows are loaded + + Args: + name: The flow name to check + + Returns: + True if the flow should be loaded, False otherwise + """ + if self.service_config.enabled_flows: + return name in self.service_config.enabled_flows + elif self.service_config.disabled_flows: + return name not in self.service_config.disabled_flows + else: + return True + + def _initialize_flow(self): + """Initialize all flows from both registry and configuration. + + Flows can be defined in two ways: + 1. Registered flows: Python classes decorated with @register_flow + 2. Configuration flows: Defined in config as ExpressionFlow instances + + Process: + 1. First, instantiate all registered flow classes (from decorators) + - Filter based on enabled_flows/disabled_flows + - Create instance with the flow name + + 2. Then, instantiate all configured flows (from config file) + - Filter based on enabled_flows/disabled_flows + - Create ExpressionFlow instances with flow configuration + + Note: Configuration flows can override registered flows with the same name. + """ + + # Initialize flows from registry (decorator-based registration) + for name, flow_cls in self.registry_dict[RegistryEnum.FLOW].items(): + if not self._filter_flows(name): + continue + flow: "BaseFlow" = flow_cls(name=name) + self.instance_dict[RegistryEnum.FLOW][flow.name] = flow + + # Initialize flows from configuration (config-based definition) + from ..flow import ExpressionFlow + + for name, flow_config in self.service_config.flow.items(): + if not self._filter_flows(name): + continue + flow_config.name = name + flow: BaseFlow = ExpressionFlow(flow_config=flow_config) + self.instance_dict[RegistryEnum.FLOW][name] = flow + + def _initialize_service(self): + """Initialize the service backend instance. + + Creates an instance of the configured service backend (e.g., HTTP, MCP, or CMD service) + and stores it in the instance dictionary under the 'default' key. + """ + service_cls = self.get_service_class(self.service_config.backend) + self.instance_dict[RegistryEnum.SERVICE]["default"] = service_cls() + + async def prepare_mcp_servers(self): + """Prepare and initialize MCP (Model Context Protocol) server connections. + + This method: + 1. Checks if MCP servers are configured + 2. Creates an MCP client instance + 3. For each configured server: + - Lists available tool calls from the server + - Builds a mapping of tool_name -> ToolCall object + - Logs available tools for debugging + + The mcp_server_mapping is structured as: + { + "server_name": { + "tool_name": ToolCall(...), + ... + }, + ... + } + + This allows the application to discover and use external tools provided by MCP servers. + """ + if not self.service_config.mcp_servers: + return + + from ..utils import MCPClient + + mcp_client = MCPClient(config={"mcpServers": self.service_config.mcp_servers}) + for server_name in self.service_config.mcp_servers.keys(): + try: + # Retrieve all available tool calls from this MCP server + tool_calls = await mcp_client.list_tool_calls(server_name=server_name, return_dict=False) + + # Build mapping: tool_name -> ToolCall for quick lookup + self.mcp_server_mapping[server_name] = {tool_call.name: tool_call for tool_call in tool_calls} + + # Log discovered tools for debugging + for tool_call in tool_calls: + logger.info(f"list_tool_calls: {server_name}@{tool_call.name} {tool_call.simple_input_dump()}") + + except Exception as e: + logger.exception(f"list_tool_calls: {server_name} error: {e}") + + def print_logo(self): + """Print the ReMe logo if enabled in configuration.""" + if self.service_config.enable_logo: + print_logo(service_config=self.service_config) + + async def close(self): + """Close all service components asynchronously. + + Gracefully closes all instantiated components in order: + 1. Vector stores (closes database connections) + 2. LLMs (closes API clients and connections) + 3. Embedding models (closes API clients and connections) + + This method should be called when shutting down the application + to ensure all resources are properly released. + """ + for _, vector_store in self.instance_dict[RegistryEnum.VECTOR_STORE].items(): + await vector_store.close() + + for _, llm in self.instance_dict[RegistryEnum.LLM].items(): + await llm.close() + + for _, embedding_model in self.instance_dict[RegistryEnum.EMBEDDING_MODEL].items(): + await embedding_model.close() + + def close_sync(self): + """Close all service components synchronously. + + Synchronous version of close() for non-async contexts. + Closes LLMs and embedding models without using async/await. + + Note: Vector stores are not closed here as they typically require async operations. + """ + for _, llm in self.instance_dict[RegistryEnum.LLM].items(): + llm.close_sync() + + for _, embedding_model in self.instance_dict[RegistryEnum.EMBEDDING_MODEL].items(): + embedding_model.close_sync() + + def shutdown_thread_pool(self, wait: bool = True): + """Shutdown the thread pool executor. + + Args: + wait: If True, blocks until all pending futures are executed. + If False, returns immediately and pending futures may be cancelled. + """ + if self.thread_pool: + self.thread_pool.shutdown(wait=wait) + + def shutdown_ray(self, wait: bool = True): + """Shutdown Ray cluster if it was initialized. + + Args: + wait: If True, waits for Ray to fully shutdown. + If False, returns immediately without waiting. + + Note: Only shuts down Ray if it was configured with ray_max_workers > 1. + """ + if self.service_config and self.service_config.ray_max_workers > 1: + import ray + + ray.shutdown(_exiting_interpreter=not wait) -# Export a global instance for easy access across the application +# Export a global singleton instance for easy access across the application +# This is the primary way to access the service context throughout the codebase C = ServiceContext() diff --git a/reme_ai/core/op/base_op.py b/reme_ai/core/op/base_op.py index 26f1e7f5..7eeb5fbd 100644 --- a/reme_ai/core/op/base_op.py +++ b/reme_ai/core/op/base_op.py @@ -178,21 +178,16 @@ class BaseOp: @property def llm(self) -> BaseLLM: - """Lazily initialize and return the LLM instance.""" + """Get the LLM instance from ServiceContext.""" if isinstance(self._llm, str): - cfg = C.service_config.llm[self._llm] - self._llm = C.get_llm_class(cfg.backend)(model_name=cfg.model_name, **cfg.model_extra) + self._llm = C.get_llm(self._llm) return self._llm @property def embedding_model(self) -> BaseEmbeddingModel: - """Lazily initialize and return the embedding model instance.""" + """Get the embedding model instance from ServiceContext.""" if isinstance(self._embedding_model, str): - cfg = C.service_config.embedding_model[self._embedding_model] - self._embedding_model = C.get_embedding_model_class(cfg.backend)( - model_name=cfg.model_name, - **cfg.model_extra, - ) + self._embedding_model = C.get_embedding_model(self._embedding_model) return self._embedding_model @property @@ -204,10 +199,9 @@ class BaseOp: @property def token_counter(self) -> BaseTokenCounter: - """Lazily initialize and return the token counter instance.""" + """Get the token counter instance from ServiceContext.""" if isinstance(self._token_counter, str): - cfg = C.service_config.token_counter[self._token_counter] - self._token_counter = C.get_token_counter_class(cfg.backend)(model_name=cfg.model_name, **cfg.model_extra) + self._token_counter = C.get_token_counter(self._token_counter) return self._token_counter @property @@ -220,25 +214,6 @@ class BaseOp: """Get the response object.""" return self.context.response - async def before_execute(self): - """Prepare context and validate before async execution.""" - self.context.apply_mapping(self.input_mapping) - self._validate_inputs() - - async def execute(self): - """Define core async logic in subclasses.""" - - async def after_execute(self): - """Finalize context and mappings after async execution.""" - self.context.apply_mapping(self.output_mapping) - if self.tool_call is not None and self.save_response_result: - self.context.response.answer = self.output - - if not isinstance(self._llm, str) and hasattr(self._llm, "close"): - await self._llm.close() - if not isinstance(self._embedding_model, str) and hasattr(self._embedding_model, "close"): - await self._embedding_model.close() - def before_execute_sync(self): """Prepare context and validate before sync execution.""" self.context.apply_mapping(self.input_mapping) @@ -253,10 +228,16 @@ class BaseOp: if self.tool_call is not None and self.save_response_result: self.context.response.answer = self.output - if not isinstance(self._llm, str) and hasattr(self._llm, "close_sync"): - self._llm.close_sync() - if not isinstance(self._embedding_model, str) and hasattr(self._embedding_model, "close_sync"): - self._embedding_model.close_sync() + async def before_execute(self): + """Prepare context and validate before async execution.""" + self.before_execute_sync() + + async def execute(self): + """Define core async logic in subclasses.""" + + async def after_execute(self): + """Finalize context and mappings after async execution.""" + self.after_execute_sync() @timer def call_sync(self, context: RuntimeContext = None, **kwargs): diff --git a/reme_ai/core/reme.py b/reme_ai/core/reme.py new file mode 100644 index 00000000..d1240cb6 --- /dev/null +++ b/reme_ai/core/reme.py @@ -0,0 +1,92 @@ +"""ReMe application classes for simplified configuration and execution.""" + +import sys + +from .application import Application +from .config import ReMeConfigParser +from .context import C + + +class ReMe(Application): + """Simplified ReMe application that auto-initializes the service context.""" + + def __init__( + self, + *args, + llm_api_key: str | None = None, + llm_api_base: str | None = None, + embedding_api_key: str | None = None, + embedding_api_base: str | None = None, + enable_logo: bool = True, + llm: dict | None = None, + embedding_model: dict | None = None, + vector_store: dict | None = None, + token_counter: dict | None = None, + **kwargs, + ): + super().__init__( + *args, + llm_api_key=llm_api_key, + llm_api_base=llm_api_base, + embedding_api_key=embedding_api_key, + embedding_api_base=embedding_api_base, + service_config=None, + parser=ReMeConfigParser, + config_path=None, + enable_logo=enable_logo, + llm=llm, + embedding_model=embedding_model, + vector_store=vector_store, + token_counter=token_counter, + **kwargs, + ) + + C.initialize_service_context() + + async def summary(self): + """Execute summary operations.""" + + async def retrieve(self): + """Execute retrieve operations.""" + + +class ReMeApp(Application): + """ReMe application with config file support and flow execution methods.""" + + def __init__( + self, + *args, + llm_api_key: str | None = None, + llm_api_base: str | None = None, + embedding_api_key: str | None = None, + embedding_api_base: str | None = None, + config_path: str | None = None, + enable_logo: bool = True, + **kwargs, + ): + super().__init__( + *args, + llm_api_key=llm_api_key, + llm_api_base=llm_api_base, + embedding_api_key=embedding_api_key, + embedding_api_base=embedding_api_base, + service_config=None, + parser=ReMeConfigParser, + config_path=config_path, + enable_logo=enable_logo, + **kwargs, + ) + + async def async_execute(self, name: str, **kwargs) -> dict: + """Execute a flow asynchronously and return the result as a dictionary.""" + return (await self.execute_flow(name=name, **kwargs)).model_dump() + + +def main(): + """Main entry point for running ReMe application from command line.""" + with ReMeApp(*sys.argv[1:]) as app: + app.run_service() + + +if __name__ == "__main__": + main() diff --git a/reme_ai/core/service/base_service.py b/reme_ai/core/service/base_service.py index b81a710d..28c82fb3 100644 --- a/reme_ai/core/service/base_service.py +++ b/reme_ai/core/service/base_service.py @@ -29,7 +29,7 @@ class BaseService(ABC): model = create_pydantic_model(tool_call.name, tool_call.parameters) return tool_call, model - def run(self) -> None: + def run(self): """Initialize and integrate all flows registered in the global context.""" flow_names: list[str] = [] for _, flow in C.flow_dict.items(): diff --git a/reme_ai/core/service/cmd_service.py b/reme_ai/core/service/cmd_service.py index 0b10ac1e..bed8ec86 100644 --- a/reme_ai/core/service/cmd_service.py +++ b/reme_ai/core/service/cmd_service.py @@ -21,7 +21,7 @@ class CmdService(BaseService): """Integrate the workflow configuration into the command service.""" self._cmd_flow = CmdFlow(flow=C.service_config.flow) - def run(self) -> None: + def run(self): """Execute the command flow in either asynchronous or synchronous mode.""" super().run() diff --git a/reme_ai/core/service/http_service.py b/reme_ai/core/service/http_service.py index 9ed56e17..f63d7333 100644 --- a/reme_ai/core/service/http_service.py +++ b/reme_ai/core/service/http_service.py @@ -7,13 +7,12 @@ import uvicorn from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import StreamingResponse -from loguru import logger from .base_service import BaseService from ..context import C -from ..enumeration import ChunkEnum from ..flow import BaseFlow -from ..schema import Response, StreamChunk +from ..schema import Response +from ..utils.common_utils import execute_stream_task @C.register_service("http") @@ -57,34 +56,13 @@ class HttpService(BaseService): task = asyncio.create_task(flow.call(stream_queue=queue, **request.model_dump(exclude_none=True))) async def generate_stream() -> AsyncGenerator[bytes, None]: - done_bytes = b"data:[DONE]\n\n" - try: - while True: - # Wait for next chunk or check if task failed - get_chunk = asyncio.create_task(queue.get()) - done, _ = await asyncio.wait({get_chunk, task}, return_when=asyncio.FIRST_COMPLETED) - - if get_chunk in done: - chunk: StreamChunk = get_chunk.result() - if chunk.done: - yield done_bytes - break - yield f"data:{chunk.model_dump_json()}\n\n".encode() - else: - # Task finished unexpectedly or raised exception - await task - yield done_bytes - break - - except Exception as e: - logger.exception(f"Stream error in {tool_call.name}: {e}") - err = StreamChunk(chunk_type=ChunkEnum.ERROR, chunk=str(e), done=True) - yield f"data:{err.model_dump_json()}\n\n".encode() - yield done_bytes - - finally: - if not task.done(): - task.cancel() + async for chunk in execute_stream_task( + queue=queue, + task=task, + flow_name=tool_call.name, + as_bytes=True, + ): + yield chunk return StreamingResponse(generate_stream(), media_type="text/event-stream") @@ -95,7 +73,7 @@ class HttpService(BaseService): """Register a flow based on its streaming configuration.""" return self._integrate_stream_flow(flow) if flow.stream else self._integrate_flow(flow) - def run(self) -> None: + def run(self): """Start the Uvicorn server.""" super().run() cfg = C.service_config.http diff --git a/reme_ai/core/tool/mcp_tool.py b/reme_ai/core/tool/mcp_tool.py index 90f6ad7b..802e4138 100644 --- a/reme_ai/core/tool/mcp_tool.py +++ b/reme_ai/core/tool/mcp_tool.py @@ -48,7 +48,7 @@ class MCPTool(BaseOp): self._client = MCPClient(C.service_config.mcp_servers) def _build_tool_call(self) -> ToolCall: - tool_call_dict = C.mcp_server_tool_call_mapping[self.mcp_server] + tool_call_dict = C.mcp_server_mapping[self.mcp_server] tool_call: ToolCall = tool_call_dict[self.tool_name].model_copy(deep=True) # Initialize required list if not exists diff --git a/reme_ai/core/utils/__init__.py b/reme_ai/core/utils/__init__.py index 7767a13b..87e01299 100644 --- a/reme_ai/core/utils/__init__.py +++ b/reme_ai/core/utils/__init__.py @@ -2,7 +2,7 @@ from .cache_handler import CacheHandler from .case_converter import snake_to_camel, camel_to_snake -from .common_utils import run_coro_safely +from .common_utils import run_coro_safely, execute_stream_task from .env_utils import load_env from .http_client import HttpClient from .llm_utils import extract_content, format_messages @@ -19,6 +19,7 @@ __all__ = [ "snake_to_camel", "camel_to_snake", "run_coro_safely", + "execute_stream_task", "load_env", "HttpClient", "extract_content", diff --git a/reme_ai/core/utils/common_utils.py b/reme_ai/core/utils/common_utils.py index 5fe01a61..d370064d 100644 --- a/reme_ai/core/utils/common_utils.py +++ b/reme_ai/core/utils/common_utils.py @@ -1,9 +1,14 @@ """Common utility functions""" import asyncio -from collections.abc import Coroutine +from collections.abc import AsyncGenerator, Coroutine from typing import Any +from loguru import logger + +from ..enumeration import ChunkEnum +from ..schema import StreamChunk + def run_coro_safely(coro: Coroutine[Any, Any, Any]) -> Any | asyncio.Task[Any]: """Run a coroutine in the current event loop or a new one if none exists.""" @@ -18,3 +23,61 @@ def run_coro_safely(coro: Coroutine[Any, Any, Any]) -> Any | asyncio.Task[Any]: else: # Schedule the coroutine as a background task in the active loop return loop.create_task(coro) + + +async def execute_stream_task( + queue: asyncio.Queue, + task: asyncio.Task, + flow_name: str | None = None, + as_bytes: bool = False, +) -> AsyncGenerator[str | bytes, None]: + """ + Core stream flow execution logic. + + Handles streaming from a queue while monitoring the task completion. + Properly manages errors and resource cleanup. + + Args: + queue: Queue to receive StreamChunk objects from + task: Background task executing the flow + flow_name: Optional flow name for logging purposes + as_bytes: If True, yield bytes for HTTP responses; if False, yield strings + + Yields: + SSE-formatted data chunks (either str or bytes based on as_bytes) + """ + done_msg = b"data:[DONE]\n\n" if as_bytes else "data:[DONE]\n\n" + + try: + while True: + # Wait for next chunk or check if task failed + get_chunk = asyncio.create_task(queue.get()) + done, _ = await asyncio.wait({get_chunk, task}, return_when=asyncio.FIRST_COMPLETED) + + if get_chunk in done: + chunk: StreamChunk = get_chunk.result() + if chunk.done: + yield done_msg + break + + data = f"data:{chunk.model_dump_json()}\n\n" + yield data.encode() if as_bytes else data + else: + # Task finished unexpectedly or raised exception + await task + yield done_msg + break + + except Exception as e: + log_msg = f"Stream error in {flow_name}: {e}" if flow_name else f"Stream error: {e}" + logger.exception(log_msg) + + err = StreamChunk(chunk_type=ChunkEnum.ERROR, chunk=str(e), done=True) + err_data = f"data:{err.model_dump_json()}\n\n" + yield err_data.encode() if as_bytes else err_data + yield done_msg + + finally: + # Ensure task is cancelled if still running to avoid resource leaks + if not task.done(): + task.cancel()