diff --git a/reme_cli/component/__init__.py b/reme_cli/component/__init__.py index e2b6373d..3f1fea2d 100644 --- a/reme_cli/component/__init__.py +++ b/reme_cli/component/__init__.py @@ -1,8 +1,14 @@ +from .application_context import ApplicationConfig from .base_component import BaseComponent from .component_registry import ComponentRegistry, R +from .prompt_handler import PromptHandler +from .runtime_context import RuntimeContext __all__ = [ + "ApplicationConfig", "BaseComponent", "ComponentRegistry", "R", + "PromptHandler", + "RuntimeContext", ] diff --git a/reme_cli/component/as_llm/__init__.py b/reme_cli/component/as_llm/__init__.py index 1dda9bae..b501257f 100644 --- a/reme_cli/component/as_llm/__init__.py +++ b/reme_cli/component/as_llm/__init__.py @@ -1,14 +1,15 @@ """Module for registering AgentScope LLM models.""" -from agentscope.model import OpenAIChatModel +import asyncio + +from agentscope.model import OpenAIChatModel, ChatModelBase from ..base_component import BaseComponent from ..component_registry import R from ...enumeration import ComponentEnum -@R.register("openai") -class AsOpenAIChatModel(BaseComponent): +class BaseAsLLM(BaseComponent): """Simple wrapper for AgentScope LLM models.""" component_type = ComponentEnum.AS_LLM @@ -16,7 +17,18 @@ class AsOpenAIChatModel(BaseComponent): def __init__(self, **kwargs) -> None: """Initialize with model configuration.""" super().__init__(**kwargs) - self.model: OpenAIChatModel | None = None + self.model: ChatModelBase | None = None + + async def _start(self, app_context=None) -> None: + """Initialize the AgentScope model instance.""" + + async def _close(self) -> None: + """Close the AgentScope model and release resources.""" + self.model = None + + +@R.register("openai") +class OpenAIAsLLM(BaseAsLLM): async def _start(self, app_context=None) -> None: """Initialize the AgentScope model instance.""" @@ -25,10 +37,17 @@ class AsOpenAIChatModel(BaseComponent): async def _close(self) -> None: """Close the AgentScope model and release resources.""" if self.model is not None: - await self.model.client.close() + client = getattr(self.model, "client", None) + if client is not None and hasattr(client, "close"): + close_method = client.close + if asyncio.iscoroutinefunction(close_method): + await close_method() + else: + close_method() self.model = None __all__ = [ - "AsOpenAIChatModel", + "BaseAsLLM", + "OpenAIAsLLM", ] diff --git a/reme_cli/component/as_llm_formatter/__init__.py b/reme_cli/component/as_llm_formatter/__init__.py index 42d942e4..30022a30 100644 --- a/reme_cli/component/as_llm_formatter/__init__.py +++ b/reme_cli/component/as_llm_formatter/__init__.py @@ -8,9 +8,8 @@ from ..component_registry import R from ...enumeration import ComponentEnum -@R.register("openai") -class AsOpenAIChatFormatter(BaseComponent): - """Wrapper for ReMeOpenAIChatFormatter.""" +class BaseAsLLMFormatter(BaseComponent): + """Base wrapper for AgentScope LLM formatters.""" component_type = ComponentEnum.AS_LLM_FORMATTER @@ -21,13 +20,22 @@ class AsOpenAIChatFormatter(BaseComponent): async def _start(self, app_context=None) -> None: """Initialize the formatter instance.""" - self.formatter = ReMeOpenAIChatFormatter(**self.kwargs) async def _close(self) -> None: """Close the formatter (no-op for formatter).""" self.formatter = None +@R.register("openai") +class AsOpenAIChatFormatter(BaseAsLLMFormatter): + """Wrapper for ReMeOpenAIChatFormatter.""" + + async def _start(self, app_context=None) -> None: + """Initialize the formatter instance.""" + self.formatter = ReMeOpenAIChatFormatter(**self.kwargs) + + __all__ = [ + "BaseAsLLMFormatter", "AsOpenAIChatFormatter", ] diff --git a/reme_cli/component/op/base_op.py b/reme_cli/component/op/base_op.py deleted file mode 100644 index b016bb47..00000000 --- a/reme_cli/component/op/base_op.py +++ /dev/null @@ -1,427 +0,0 @@ -"""Base operator class for LLM workflow execution and composition.""" - -import asyncio -import copy -import inspect -from abc import ABCMeta -from pathlib import Path -from typing import Callable, Optional, Any - -from agentscope.formatter import FormatterBase -from agentscope.model import ChatModelBase -from agentscope.token import HuggingFaceTokenCounter -from loguru import logger -from tqdm import tqdm - -from ..embedding import BaseEmbeddingModel -from ..file_store import BaseFileStore -from ..llm import BaseLLM -from ..prompt_handler import PromptHandler -from ..runtime_context import RuntimeContext -from ..schema import Response, ServiceConfig -from ..schema.service_config import OpConfig -from ..service_context import ServiceContext -from ..token_counter import BaseTokenCounter -from ..utils import camel_to_snake, CacheHandler, timer -from ..vector_store import BaseVectorStore - - -class BaseOp(metaclass=ABCMeta): - """Base operator class for LLM workflow execution and composition.""" - - __alias_name__: str = "" - - def __new__(cls, *args, **kwargs): - """Capture initialization arguments for object cloning.""" - instance = super().__new__(cls) - instance._init_args = copy.copy(args) - instance._init_kwargs = copy.copy(kwargs) - return instance - - def __init__( - self, - name: str = "", - async_mode: bool = True, - language: str = "", - prompt_name: str = "", - prompt_path: str = "", - as_llm: str | ChatModelBase = "default", - as_llm_formatter: str | FormatterBase = "default", - as_token_counter: str | HuggingFaceTokenCounter = "default", - llm: str | BaseLLM = "default", - embedding_model: str | BaseEmbeddingModel = "default", - vector_store: str | BaseVectorStore = "default", - file_store: str | BaseFileStore = "default", - token_counter: str | BaseTokenCounter = "default", - enable_cache: bool = False, - cache_path: str = "cache/op", - cache_expire_hours: float | None = None, - sub_ops: dict[str, "BaseOp"] | list["BaseOp"] | Optional["BaseOp"] = None, - input_mapping: dict[str, str] | None = None, - output_mapping: dict[str, str] | None = None, - enable_parallel: bool = False, - max_retries: int = 1, - raise_exception: bool = False, - **kwargs, - ): - """Initialize operator configurations and internal state.""" - self.name = name or self.__alias_name__ or camel_to_snake(self.__class__.__name__) - self.async_mode = async_mode - self.language = language - self.prompt = self._get_prompt_handler(prompt_name, prompt_path) - - self._as_llm = as_llm - self._as_llm_formatter = as_llm_formatter - self._as_token_counter = as_token_counter - self._llm = llm - self._embedding_model = embedding_model - self._vector_store = vector_store - self._file_store = file_store - self._token_counter = token_counter - - self.enable_cache = enable_cache - self.cache_path = cache_path - self.cache_expire_hours = cache_expire_hours - - self.sub_ops: list["BaseOp"] = [] - self.add_sub_ops(sub_ops) - - self.input_mapping = input_mapping - self.output_mapping = output_mapping - self.enable_parallel = enable_parallel # Control whether to execute tasks in parallel - self.max_retries = max(1, max_retries) - self.raise_exception = raise_exception - self.op_params = kwargs - - self._pending_tasks: list = [] - self.context: RuntimeContext | None = None - self._cache: CacheHandler | None = None - - def _get_prompt_handler(self, prompt_name: str, prompt_path: str) -> PromptHandler: - """Load prompt configuration from the associated YAML file.""" - if prompt_path: - path = Path(prompt_path) - else: - path = Path(inspect.getfile(self.__class__)) - if prompt_name: - path = path.with_stem(prompt_name) - return PromptHandler(language=self.language).load_prompt_by_file(path.with_suffix(".yaml")) - - def _handle_failure(self, e: Exception, attempt: int) -> str | None: - """Log failures and handle final retry logic.""" - message = f"[{self.__class__.__name__}] failed (attempt {attempt + 1}): {e}" - if attempt == self.max_retries - 1: - logger.exception(message) - if self.raise_exception: - raise e - return f"[{self.__class__.__name__}] failed: {e}" - else: - logger.warning(message) - return None - - @property - def cache(self) -> CacheHandler: - """Access the operator-specific cache handler.""" - assert self.enable_cache, "Cache is disabled!" - if not self._cache: - self._cache = CacheHandler(f"{self.cache_path}/{self.name}") - return self._cache - - @property - def service_context(self) -> ServiceContext: - """Access the service context.""" - assert self.context, "Service context is not initialized!" - return self.context.service_context - - @property - def service_config(self) -> ServiceConfig: - """Access the service configuration.""" - return self.service_context.service_config - - @property - def as_llm(self) -> ChatModelBase: - """Get the AgentScope LLM instance from ServiceContext.""" - if isinstance(self._as_llm, str): - self._as_llm = self.service_context.as_llms[self._as_llm] - return self._as_llm - - @property - def as_llm_formatter(self) -> FormatterBase: - """Get the AgentScope LLM formatter instance from ServiceContext.""" - if isinstance(self._as_llm_formatter, str): - self._as_llm_formatter = self.service_context.as_llm_formatters[self._as_llm_formatter] - return self._as_llm_formatter - - @property - def as_token_counter(self) -> HuggingFaceTokenCounter: - """Get the token counter instance from ServiceContext.""" - if isinstance(self._as_token_counter, str): - self._as_token_counter = self.service_context.as_token_counters[self._as_token_counter] - return self._as_token_counter - - @property - def llm(self) -> BaseLLM: - """Get the LLM instance from ServiceContext.""" - if isinstance(self._llm, str): - self._llm = self.service_context.llms[self._llm] - return self._llm - - @property - def embedding_model(self) -> BaseEmbeddingModel: - """Get the embedding model instance from ServiceContext.""" - if isinstance(self._embedding_model, str): - self._embedding_model = self.service_context.embedding_models[self._embedding_model] - return self._embedding_model - - @property - def vector_store(self) -> BaseVectorStore: - """Lazily initialize and return the vector store instance.""" - if isinstance(self._vector_store, str): - self._vector_store = self.service_context.vector_stores[self._vector_store] - return self._vector_store - - @property - def file_store(self) -> BaseFileStore: - """Lazily initialize and return the file store instance.""" - if isinstance(self._file_store, str): - self._file_store = self.service_context.file_stores[self._file_store] - return self._file_store - - @property - def token_counter(self) -> BaseTokenCounter: - """Get the token counter instance from ServiceContext.""" - if isinstance(self._token_counter, str): - self._token_counter = self.service_context.token_counters[self._token_counter] - return self._token_counter - - @property - def service_metadata(self) -> dict: - """Get service configuration metadata.""" - return self.service_context.service_config.metadata - - @property - def response(self) -> Response: - """Access the response object.""" - return self.context.response - - def before_execute_sync(self): - """Prepare context and validate before sync execution. - - This method performs the following steps: - 1. Apply input mapping to transform context variables - 2. Load operator-specific configuration from service config if available - 3. Override operator parameters and prompts based on config - """ - self.context.apply_mapping(self.input_mapping) - - if self.context.service_context is None: - return - - service_config = self.service_context.service_config - if self.name not in service_config.ops: - return - - op_config: OpConfig = service_config.ops[self.name] - - # Override operator parameters from config - if op_config.params: - for k, v in op_config.params.items(): - if hasattr(self, k): - setattr(self, k, v) - logger.info(f"[{self.__class__.__name__}] Set attribute '{k}' = {v}") - else: - self.op_params[k] = v - logger.info(f"[{self.__class__.__name__}] Set op_param '{k}' = {v}") - - # Load custom prompt templates from config - if op_config.prompt_dict: - self.prompt.load_prompt_dict(op_config.prompt_dict) - logger.info(f"[{self.__class__.__name__}] Loaded prompt keys={list(op_config.prompt_dict.keys())}") - - async def before_execute(self): - """Prepare context and validate before async execution.""" - self.before_execute_sync() - - def execute_sync(self): - """Define core sync logic in subclasses.""" - - async def execute(self): - """Define core async logic in subclasses.""" - - def after_execute_sync(self, response: Any): - """Finalize context and mappings after sync execution.""" - self.context.apply_mapping(self.output_mapping) - if response is not None: - if isinstance(response, dict): - for k, v in response.items(): - if k == "answer": - self.response.answer = v - elif k == "success": - self.response.success = v if isinstance(v, bool) else v.lower() == "true" - else: - self.response.metadata[k] = v - else: - self.response.answer = response - return response - - async def after_execute(self, output: Any): - """Finalize context and mappings after async execution.""" - return self.after_execute_sync(output) - - @timer - def call_sync(self, context: RuntimeContext = None, **kwargs): - """Execute the operator synchronously with retry logic.""" - self.context = RuntimeContext.from_context(context, **kwargs) - response = None - for i in range(self.max_retries): - try: - self.before_execute_sync() - response = self.execute_sync() - response = self.after_execute_sync(response) - break - except Exception as e: - response = self._handle_failure(e, i) - - return response - - @timer - async def call(self, context: RuntimeContext = None, **kwargs): - """Execute the operator asynchronously with retry logic.""" - self.context = RuntimeContext.from_context(context, **kwargs) - response = None - for i in range(self.max_retries): - try: - await self.before_execute() - response = await self.execute() - response = await self.after_execute(response) - break - except Exception as e: - response = self._handle_failure(e, i) - return response - - def submit_sync_task(self, fn: Callable, *args, **kwargs) -> "BaseOp": - """Submit a task to the thread pool or local queue.""" - if self.enable_parallel and self.service_context.thread_pool is not None: - task = self.service_context.thread_pool.submit(fn, *args, **kwargs) - else: - task = (fn, args, kwargs) - self._pending_tasks.append(task) - return self - - def submit_async_task(self, coro_fn: Callable, *args, **kwargs) -> "BaseOp": - """Submit an async task to the pending tasks queue.""" - task = coro_fn(*args, **kwargs) - self._pending_tasks.append(task) - return self - - def join_sync_tasks(self, task_desc: str = None) -> list: - """Wait for all pending sync tasks and return flattened results.""" - results = [] - for task in tqdm(self._pending_tasks, desc=task_desc or self.name): - if self.enable_parallel: - result = task.result() - else: - result = task[0](*task[1], **task[2]) - if result: - if isinstance(result, list): - results.extend(result) - else: - results.append(result) - self._pending_tasks.clear() - return results - - async def join_async_tasks(self, return_exceptions: bool = True) -> list: - """Wait for all pending async tasks and aggregate results.""" - if self.enable_parallel: - raw_results = await asyncio.gather(*self._pending_tasks, return_exceptions=return_exceptions) - else: - raw_results = [] - for task in self._pending_tasks: - try: - result = await task - raw_results.append(result) - except Exception as e: - if return_exceptions: - raw_results.append(e) - else: - raise - - results = [] - for result in raw_results: - if isinstance(result, Exception): - logger.error(f"[{self.__class__.__name__}] Async task failed: {result}") - elif result: - if isinstance(result, list): - results.extend(result) - else: - results.append(result) - self._pending_tasks.clear() - return results - - def add_sub_ops(self, sub_ops: dict[str, "BaseOp"] | list["BaseOp"] | Optional["BaseOp"]): - """Add child operators to this operator's sub_ops.""" - if not sub_ops: - return - - if isinstance(sub_ops, dict): - for name, op in sub_ops.items(): - assert self.async_mode == op.async_mode, "Async mode mismatch!" - op.name = name - if self.language: - op.language = self.language - self.sub_ops.append(op) - - elif isinstance(sub_ops, list): - for op in sub_ops: - assert self.async_mode == op.async_mode, "Async mode mismatch!" - if self.language: - op.language = self.language - self.sub_ops.append(op) - - else: - assert self.async_mode == sub_ops.async_mode, "Async mode mismatch!" - if self.language: - sub_ops.language = self.language - self.sub_ops.append(sub_ops) - - def add_sub_op(self, sub_op: "BaseOp"): - """Add a single child operator to this operator's sub_ops.""" - self.sub_ops.append(sub_op) - - def __lshift__(self, ops): - """Operator overload for adding sub-operators.""" - self.add_sub_ops(ops) - return self - - def __rshift__(self, op: "BaseOp"): - """Operator overload for sequential execution composition.""" - from .sequential_op import SequentialOp - - seq = SequentialOp(sub_ops=[self], async_mode=self.async_mode) - seq.add_sub_ops(op.sub_ops if isinstance(op, SequentialOp) else op) - return seq - - def __or__(self, op: "BaseOp"): - """Operator overload for parallel execution composition.""" - from .parallel_op import ParallelOp - - par = ParallelOp(sub_ops=[self], async_mode=self.async_mode) - par.add_sub_ops(op.sub_ops if isinstance(op, ParallelOp) else op) - return par - - def prompt_format(self, prompt_name: str, **kwargs) -> str: - """Format a prompt template with provided keyword arguments.""" - return self.prompt.prompt_format(prompt_name=prompt_name, **kwargs) - - def get_prompt(self, prompt_name: str) -> str: - """Get a prompt template by name.""" - return self.prompt.get_prompt(prompt_name=prompt_name) - - def copy(self, **kwargs): - """Create a copy of this operator with optional parameter overrides.""" - copy_op = self.__class__(*self._init_args, **{**self._init_kwargs, **kwargs}) - if self.sub_ops: - copy_op.sub_ops.clear() - for op in self.sub_ops: - copy_op.add_sub_op(op.copy()) - return copy_op diff --git a/reme_cli/component/prompt_handler.py b/reme_cli/component/prompt_handler.py new file mode 100644 index 00000000..3832c11f --- /dev/null +++ b/reme_cli/component/prompt_handler.py @@ -0,0 +1,129 @@ +"""Module for managing and formatting prompt templates.""" + +import json +from pathlib import Path +from string import Formatter + +import yaml + +from ..utils import get_logger + + +class PromptHandler: + """A handler for loading, retrieving, and formatting prompt templates.""" + + _SUPPORTED_EXTENSIONS = {".yaml", ".yml", ".json"} + + def __init__(self, language: str = "", **kwargs): + self.data: dict[str, str] = {k: v for k, v in kwargs.items() if isinstance(v, str)} + self.language: str = language.strip() + self.logger = get_logger() + + def load_prompt_by_file(self, prompt_file_path: str | Path | None = None, + overwrite: bool = True) -> "PromptHandler": + """Load prompts from a YAML or JSON file.""" + if prompt_file_path is None: + return self + + path = Path(prompt_file_path) + if not path.exists(): + self.logger.warning(f"Prompt file not found: {path}") + return self + + suffix = path.suffix.lower() + if suffix not in self._SUPPORTED_EXTENSIONS: + self.logger.warning(f"Unsupported file extension '{suffix}', expected one of {self._SUPPORTED_EXTENSIONS}") + return self + + try: + with path.open(encoding="utf-8") as f: + prompt_dict = yaml.safe_load(f) if suffix in (".yaml", ".yml") else json.load(f) + except (json.JSONDecodeError, yaml.YAMLError) as e: + self.logger.error(f"Failed to parse prompt file {path}: {e}") + return self + except OSError as e: + self.logger.error(f"Failed to read prompt file {path}: {e}") + return self + + return self.load_prompt_dict(prompt_dict, overwrite) + + def load_prompt_dict(self, prompt_dict: dict | None = None, overwrite: bool = True) -> "PromptHandler": + """Merge prompts from a dictionary.""" + if not prompt_dict: + return self + + for key, value in prompt_dict.items(): + if not isinstance(value, str): + continue + if key in self.data and not overwrite: + continue + if key in self.data: + self.logger.warning(f"Overwriting prompt '{key}'") + self.data[key] = value + + return self + + def get_prompt(self, prompt_name: str) -> str: + """Retrieve a prompt by name with language suffix fallback.""" + if self.language: + key = f"{prompt_name}_{self.language}" + if key in self.data: + return self.data[key].strip() + + if prompt_name in self.data: + return self.data[prompt_name].strip() + + raise KeyError(f"Prompt '{prompt_name}' not found. Available: {list(self.data.keys())[:10]}") + + def has_prompt(self, prompt_name: str) -> bool: + """Check if a prompt exists.""" + return prompt_name in self.data or f"{prompt_name}_{self.language}" in self.data + + def list_prompts(self, language_filter: str | None = None) -> list[str]: + """List all available prompt names.""" + if not language_filter: + return list(self.data.keys()) + suffix = f"_{language_filter.strip()}" + return [k for k in self.data if k.endswith(suffix)] + + def prompt_format(self, prompt_name: str, validate: bool = True, **kwargs) -> str: + """Format a prompt with conditional line filtering and variable substitution.""" + prompt = self.get_prompt(prompt_name) + + flags = {k: v for k, v in kwargs.items() if isinstance(v, bool)} + formats = {k: v for k, v in kwargs.items() if not isinstance(v, bool)} + + if flags: + lines = [] + for line in prompt.split("\n"): + remaining = line + has_flag = False + should_include = False + while True: + matched = False + for flag, enabled in flags.items(): + prefix = f"[{flag}]" + if remaining.startswith(prefix): + remaining = remaining[len(prefix):] + has_flag = True + if enabled: + should_include = True + matched = True + break + if not matched: + break + # Include line if: no flag prefix, or at least one flag is enabled + if not has_flag or should_include: + lines.append(remaining) + prompt = "\n".join(lines) + + if validate: + required = {f for _, f, _, _ in Formatter().parse(prompt) if f} + missing = required - set(formats.keys()) + if missing: + raise ValueError(f"Missing format variables for '{prompt_name}': {sorted(missing)}") + + return prompt.format(**formats).strip() if formats else prompt.strip() + + def __repr__(self) -> str: + return f"PromptHandler(language='{self.language}', num_prompts={len(self.data)})" diff --git a/reme_cli/component/runtime_context.py b/reme_cli/component/runtime_context.py new file mode 100644 index 00000000..700ff3cd --- /dev/null +++ b/reme_cli/component/runtime_context.py @@ -0,0 +1,105 @@ +"""Runtime context for managing response states and asynchronous data streaming.""" + +import asyncio + +from .application_context import ApplicationContext +from ..enumeration import ChunkEnum +from ..schema import Response, StreamChunk + + +class RuntimeContext: + """Context for execution state, response metadata, and stream queues.""" + + def __init__(self, **kwargs): + """Initialize the context with all keyword arguments stored in data.""" + self.data: dict = kwargs + + @property + def response(self) -> Response: + """Get or create the response object.""" + return self.data.setdefault("response", Response()) + + @response.setter + def response(self, value: Response) -> None: + """Set the response object.""" + self.data["response"] = value + + @property + def stream_queue(self) -> asyncio.Queue | None: + """Get the stream queue.""" + return self.data.get("stream_queue") + + @stream_queue.setter + def stream_queue(self, value: asyncio.Queue | None) -> None: + """Set the stream queue.""" + self.data["stream_queue"] = value + + @property + def application_context(self) -> ApplicationContext | None: + """Get the application context.""" + return self.data.get("application_context") + + @application_context.setter + def application_context(self, value: ApplicationContext | None) -> None: + """Set the application context.""" + self.data["application_context"] = value + + @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.data.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.data: + self.data[target] = self.data[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.data: + raise ValueError(f"{context_name}: missing required input '{key}'") + return self diff --git a/reme_cli/enumeration/__init__.py b/reme_cli/enumeration/__init__.py index a394e324..f5126f45 100644 --- a/reme_cli/enumeration/__init__.py +++ b/reme_cli/enumeration/__init__.py @@ -1,9 +1,11 @@ """enumeration""" +from .chunk_enum import ChunkEnum from .component_enum import ComponentEnum from .json_schema_enum import JsonSchemaEnum __all__ = [ + "ChunkEnum", "ComponentEnum", "JsonSchemaEnum", ] diff --git a/reme_cli/enumeration/chunk_enum.py b/reme_cli/enumeration/chunk_enum.py new file mode 100644 index 00000000..736d9c67 --- /dev/null +++ b/reme_cli/enumeration/chunk_enum.py @@ -0,0 +1,31 @@ +"""Defines the types of data chunks used in streaming responses.""" + +from enum import Enum + + +class ChunkEnum(str, Enum): + """Enumeration of possible chunk categories for stream processing.""" + + # Internal reasoning or chain-of-thought process + THINK = "think" + + # The final generated response content + ANSWER = "answer" + + # Metadata or calls related to external tools + TOOL = "tool" + + # Resource consumption and token usage statistics + USAGE = "usage" + + # Error messages or exception details + ERROR = "error" + + # Signal indicating the start of a new ReAct step + STEP_START = "step_start" + + # Tool execution result + TOOL_RESULT = "tool_result" + + # Final signal indicating the completion of the stream + DONE = "done" diff --git a/reme_cli/component/flow/__init__.py b/reme_cli/flow/__init__.py similarity index 100% rename from reme_cli/component/flow/__init__.py rename to reme_cli/flow/__init__.py diff --git a/reme_cli/component/flow/base_flow.py b/reme_cli/flow/base_flow.py similarity index 100% rename from reme_cli/component/flow/base_flow.py rename to reme_cli/flow/base_flow.py diff --git a/reme_cli/component/flow/cmd_flow.py b/reme_cli/flow/cmd_flow.py similarity index 100% rename from reme_cli/component/flow/cmd_flow.py rename to reme_cli/flow/cmd_flow.py diff --git a/reme_cli/component/flow/expression_flow.py b/reme_cli/flow/expression_flow.py similarity index 100% rename from reme_cli/component/flow/expression_flow.py rename to reme_cli/flow/expression_flow.py diff --git a/reme_cli/component/op/__init__.py b/reme_cli/op/__init__.py similarity index 100% rename from reme_cli/component/op/__init__.py rename to reme_cli/op/__init__.py diff --git a/reme_cli/op/base_op.py b/reme_cli/op/base_op.py new file mode 100644 index 00000000..ce4a520c --- /dev/null +++ b/reme_cli/op/base_op.py @@ -0,0 +1,193 @@ +"""Base operator class for LLM workflow execution and composition.""" + +import asyncio +import copy +import inspect +from abc import ABCMeta +from pathlib import Path +from typing import Callable, TYPE_CHECKING + +from ..component import PromptHandler +from ..component import RuntimeContext +from ..component.application_context import ApplicationContext +from ..component.embedding import BaseEmbeddingModel +from ..component.file_store import BaseFileStore +from ..schema import ApplicationConfig +from ..enumeration import ComponentEnum +if TYPE_CHECKING: + from ..component.as_llm import AsOpenAIChatModel + + +class BaseOp(metaclass=ABCMeta): + """Base operator class for LLM workflow execution and composition.""" + + def __new__(cls, *args, **kwargs): + """Capture initialization arguments for object cloning.""" + instance = super().__new__(cls) + instance._init_args = copy.copy(args) + instance._init_kwargs = copy.copy(kwargs) + return instance + + def __init__( + self, + name: str = "", + language: str = "", + prompt_dict: dict[str, str] | None = None, + input_mapping: dict[str, str] | None = None, + output_mapping: dict[str, str] | None = None, + **kwargs, + ): + """Initialize operator configurations and internal state.""" + self.name = name or self.__class__.__name__ + self.language = language + self.prompt = PromptHandler(language=self.language) + self.prompt.load_prompt_by_file(Path(inspect.getfile(self.__class__)).with_suffix(".yaml")) \ + .load_prompt_dict(prompt_dict) + + self.input_mapping = input_mapping + self.output_mapping = output_mapping + self.enable_parallel = enable_parallel + self.max_retries = max(1, max_retries) + self.raise_exception = raise_exception + self.op_params = kwargs + + self._pending_tasks: list = [] + self.context: RuntimeContext | None = None + + async def before_execute(self): + self.context.apply_mapping(self.input_mapping) + + async def execute(self): + """""" + + async def after_execute(self, output): + self.context.apply_mapping(self.output_mapping) + return output + + async def call(self, context: RuntimeContext = None, **kwargs): + self.context = RuntimeContext.from_context(context, **kwargs) + await self.before_execute() + response = await self.execute() + response = await self.after_execute(response) + return response + + @property + def application_context(self) -> ApplicationContext: + return self.context.application_context + + @property + def app_config(self) -> ApplicationConfig: + return self.application_context.app_config + + @property + def as_llm(self) -> ChatModelBase: + """Get the AgentScope LLM instance from ServiceContext.""" + as_llm_name = self.op_params.get("as_llm", "default") + as_llm_dict = self.application_context.components[ComponentEnum.AS_LLM] + + return as_llm_dict[as_llm_name] + + @property + def as_llm_formatter(self) -> FormatterBase: + """Get the AgentScope LLM formatter instance from ServiceContext.""" + if isinstance(self._as_llm_formatter, str): + self._as_llm_formatter = self.service_context.as_llm_formatters[self._as_llm_formatter] + return self._as_llm_formatter + + @property + def as_token_counter(self) -> HuggingFaceTokenCounter: + """Get the token counter instance from ServiceContext.""" + if isinstance(self._as_token_counter, str): + self._as_token_counter = self.service_context.as_token_counters[self._as_token_counter] + return self._as_token_counter + + @property + def llm(self) -> BaseLLM: + """Get the LLM instance from ServiceContext.""" + if isinstance(self._llm, str): + self._llm = self.service_context.llms[self._llm] + return self._llm + + @property + def embedding_model(self) -> BaseEmbeddingModel: + """Get the embedding model instance from ServiceContext.""" + if isinstance(self._embedding_model, str): + self._embedding_model = self.service_context.embedding_models[self._embedding_model] + return self._embedding_model + + @property + def vector_store(self) -> BaseVectorStore: + """Lazily initialize and return the vector store instance.""" + if isinstance(self._vector_store, str): + self._vector_store = self.service_context.vector_stores[self._vector_store] + return self._vector_store + + @property + def file_store(self) -> BaseFileStore: + """Lazily initialize and return the file store instance.""" + if isinstance(self._file_store, str): + self._file_store = self.service_context.file_stores[self._file_store] + return self._file_store + + @property + def token_counter(self) -> BaseTokenCounter: + """Get the token counter instance from ServiceContext.""" + if isinstance(self._token_counter, str): + self._token_counter = self.service_context.token_counters[self._token_counter] + return self._token_counter + + @property + def service_metadata(self) -> dict: + """Get service configuration metadata.""" + return self.service_context.service_config.metadata + + @property + def response(self) -> Response: + """Access the response object.""" + return self.context.response + + def submit_async_task(self, coro_fn: Callable, *args, **kwargs) -> "BaseOp": + """Submit an async task to the pending tasks queue.""" + task = coro_fn(*args, **kwargs) + self._pending_tasks.append(task) + return self + + async def join_async_tasks(self, return_exceptions: bool = True) -> list: + """Wait for all pending async tasks and aggregate results.""" + if self.enable_parallel: + raw_results = await asyncio.gather(*self._pending_tasks, return_exceptions=return_exceptions) + else: + raw_results = [] + for task in self._pending_tasks: + try: + result = await task + raw_results.append(result) + except Exception as e: + if return_exceptions: + raw_results.append(e) + else: + raise + + results = [] + for result in raw_results: + if isinstance(result, Exception): + logger.error(f"[{self.__class__.__name__}] Async task failed: {result}") + elif result: + if isinstance(result, list): + results.extend(result) + else: + results.append(result) + self._pending_tasks.clear() + return results + + def prompt_format(self, prompt_name: str, **kwargs) -> str: + """Format a prompt template with provided keyword arguments.""" + return self.prompt.prompt_format(prompt_name=prompt_name, **kwargs) + + def get_prompt(self, prompt_name: str) -> str: + """Get a prompt template by name.""" + return self.prompt.get_prompt(prompt_name=prompt_name) + + def copy(self, **kwargs): + """Create a copy of this operator with optional parameter overrides.""" + return self.__class__(*self._init_args, **{**self._init_kwargs, **kwargs}) diff --git a/reme_cli/schema/__init__.py b/reme_cli/schema/__init__.py index 3bab477e..7b984fb7 100644 --- a/reme_cli/schema/__init__.py +++ b/reme_cli/schema/__init__.py @@ -2,10 +2,17 @@ from .application_config import ApplicationConfig from .base_node import BaseNode from .file_chunk import FileChunk from .file_metadata import FileMetadata +from .response import Response +from .stream_chunk import StreamChunk +from .tool_call import ToolAttr, ToolCall __all__ = [ "ApplicationConfig", "BaseNode", "FileChunk", "FileMetadata", + "Response", + "StreamChunk", + "ToolAttr", + "ToolCall", ] diff --git a/reme_cli/schema/application_config.py b/reme_cli/schema/application_config.py index 0d439f1b..f8f142d8 100644 --- a/reme_cli/schema/application_config.py +++ b/reme_cli/schema/application_config.py @@ -13,6 +13,7 @@ class ApplicationConfig(BaseModel): enable_logo: bool = Field(default=False) language: str = Field(default="") log_to_console: bool = Field(default=True) + log_to_file: bool = Field(default=True) mcp_servers: dict[str, dict] = Field(default_factory=dict) service: dict = Field(default_factory=dict) ops: dict[str, dict] = Field(default_factory=dict) diff --git a/reme_cli/schema/response.py b/reme_cli/schema/response.py new file mode 100644 index 00000000..fe753232 --- /dev/null +++ b/reme_cli/schema/response.py @@ -0,0 +1,13 @@ +"""Defines the standardized data structure for model output responses.""" + +from typing import Any + +from pydantic import Field, BaseModel + + +class Response(BaseModel): + """Represents a structured response containing the execution result, status, and metadata.""" + + answer: str | Any = Field(default="") + success: bool = Field(default=True) + metadata: dict = Field(default_factory=dict) diff --git a/reme_cli/schema/stream_chunk.py b/reme_cli/schema/stream_chunk.py new file mode 100644 index 00000000..764981fd --- /dev/null +++ b/reme_cli/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_cli/utils/logger_utils.py b/reme_cli/utils/logger_utils.py index e7636f64..c33a880c 100644 --- a/reme_cli/utils/logger_utils.py +++ b/reme_cli/utils/logger_utils.py @@ -13,6 +13,7 @@ def get_logger( log_dir: str = "logs", level: str = "INFO", log_to_console: bool = True, + log_to_file: bool = True, force_init: bool = False, ): """Get a configured logger instance. @@ -24,6 +25,7 @@ def get_logger( log_dir: Directory path for log files. level: Logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL). log_to_console: Whether to print logs to console/screen. + log_to_file: Whether to write logs to file. force_init: Force re-initialization even if already initialized. Returns: @@ -46,24 +48,25 @@ def get_logger( colorize=True, ) - # Configure file-based logging (skip if permission denied) - try: - os.makedirs(log_dir, exist_ok=True) + # Configure file-based logging if enabled + if log_to_file: + try: + os.makedirs(log_dir, exist_ok=True) - current_ts = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") - log_filepath = os.path.join(log_dir, f"{current_ts}.log") + current_ts = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + log_filepath = os.path.join(log_dir, f"{current_ts}.log") - logger.add( - log_filepath, - level=level, - rotation="00:00", - retention="7 days", - compression="zip", - encoding="utf-8", - format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {file}:{line} | {function} | {message}", - ) - except Exception as e: - logger.error(f"Error configuring file logging: {e}") + logger.add( + log_filepath, + level=level, + rotation="00:00", + retention="7 days", + compression="zip", + encoding="utf-8", + format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {file}:{line} | {function} | {message}", + ) + except Exception as e: + logger.error(f"Error configuring file logging: {e}") _initialized = True return logger \ No newline at end of file