mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-16 23:41:10 +00:00
5b801c0d3e
* refactor(file_io): update file I/O operations and truncation logic * refactor(memory): update file-based memory compaction logic
427 lines
16 KiB
Python
427 lines
16 KiB
Python
"""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
|