ReMe/reme_ai/core/application.py

"""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, load_env


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:
            *args: Additional arguments passed to parser. Examples:
                - "llm.default.model_name=qwen3-30b-a3b-thinking-2507"
                - "llm.default.backend=openai_compatible"
                - "llm.default.temperature=0.6"
                - "embedding_model.default.model_name=text-embedding-v4"
                - "embedding_model.default.backend=openai_compatible"
                - "embedding_model.default.dimensions=1024"
                - "vector_store.default.backend=memory"
                - "vector_store.default.embedding_model=default"
            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 keyword arguments passed to parser. Same format as args but as kwargs. Examples:
                - **{"llm.default.model_name": "qwen3-30b-a3b-thinking-2507"}
        """

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

        # 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 C.service_config.init_logger:
            init_logger()

        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):
        """Update environment variable if value is provided."""
        if value:
            os.environ[key] = value

    @staticmethod
    async def start():
        """Initialize the service context and prepare external MCP servers."""
        C.initialize_service_context()
        await C.prepare_mcp_servers()

    @staticmethod
    def start_sync():
        """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):
        """
        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):
        """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(
            stream_queue=stream_queue,
            task=task,
            task_name=name,
            as_bytes=False,
        ):
            yield chunk

    @staticmethod
    def run_service():
        """Run the configured service (HTTP, MCP, or CMD)."""
        C.get_service().run()