diff --git a/reme2/component/base_component.py b/reme2/component/base_component.py index 3cef1f92..9d8b0ee0 100644 --- a/reme2/component/base_component.py +++ b/reme2/component/base_component.py @@ -119,15 +119,6 @@ class BaseComponent(ABC): # ----- Lookup -------------------------------------------------------- - def get_component(self, component_type: ComponentEnum, name: str): - """Get a component by type and name from app_context.""" - if self.app_context is None: - raise ValueError("app_context is not set") - component_dict = self.app_context.components.get(component_type, {}) - if name not in component_dict: - raise ValueError(f"{component_type.value} '{name}' not found.") - return component_dict[name] - @property def working_path(self) -> Path: if self.app_context is None: diff --git a/reme2/component/base_step.py b/reme2/component/base_step.py index 3412cff6..01a557b1 100644 --- a/reme2/component/base_step.py +++ b/reme2/component/base_step.py @@ -8,8 +8,8 @@ from agentscope.model import ChatModelBase from agentscope.token import TokenCounterBase from .base_component import BaseComponent -from .file_store import BaseFileStore from .embedding import BaseEmbeddingModel +from .file_store import BaseFileStore from .prompt_handler import PromptHandler from .runtime_context import RuntimeContext from ..enumeration import ComponentEnum @@ -27,12 +27,12 @@ class BaseStep(BaseComponent): return instance def __init__( - self, - language: str = "", - prompt_dict: dict[str, str] | None = None, - input_mapping: dict[str, str] | None = None, - output_mapping: dict[str, str] | None = None, - **kwargs, + self, + language: str = "", + prompt_dict: dict[str, str] | None = None, + input_mapping: dict[str, str] | None = None, + output_mapping: dict[str, str] | None = None, + **kwargs, ): super().__init__(**kwargs) self.language = language @@ -49,15 +49,11 @@ class BaseStep(BaseComponent): async def __call__(self, context: RuntimeContext | None = None, **kwargs): self.context = RuntimeContext.from_context(context, **kwargs) assert self.context is not None - if self.input_mapping: self.context.apply_mapping(self.input_mapping) - result = await self.execute() - if self.output_mapping: self.context.apply_mapping(self.output_mapping) - return result def get_component(self, key: ComponentEnum, name: str, attr: str | None = None): @@ -65,47 +61,29 @@ class BaseStep(BaseComponent): comp = self.app_context.components[key][name] return getattr(comp, attr) if attr else comp - def _get_component_optional(self, key: ComponentEnum, name: str = "default", attr: str | None = None): - """Like `_get_component` but returns None instead of raising when - the component / attribute is missing. For features that should - gracefully degrade when an LLM (etc.) isn't configured.""" - if self.app_context is None: - return None - comp = self.app_context.components.get(key, {}).get(name) - if comp is None: - return None - return getattr(comp, attr, None) if attr else comp - @property def as_llm(self) -> ChatModelBase: name = self.kwargs.get("as_llm", "default") - return name if isinstance(name, ChatModelBase) else self.get_component(ComponentEnum.AS_LLM, name, "model") + if isinstance(name, ChatModelBase): + return name + else: + return self.get_component(ComponentEnum.AS_LLM, name, "model") @property def as_llm_formatter(self) -> FormatterBase: name = self.kwargs.get("as_llm_formatter", "default") - return ( - name - if isinstance(name, FormatterBase) - else self.get_component( - ComponentEnum.AS_LLM_FORMATTER, - name, - "formatter", - ) - ) + if isinstance(name, FormatterBase): + return name + else: + return self.get_component(ComponentEnum.AS_LLM_FORMATTER, name, "formatter") @property def as_token_counter(self) -> TokenCounterBase: name = self.kwargs.get("as_token_counter", "default") - return ( - name - if isinstance(name, TokenCounterBase) - else self.get_component( - ComponentEnum.AS_TOKEN_COUNTER, - name, - "token_counter", - ) - ) + if isinstance(name, TokenCounterBase): + return name + else: + return self.get_component(ComponentEnum.AS_TOKEN_COUNTER, name, "token_counter") @property def file_store(self) -> BaseFileStore: @@ -115,14 +93,10 @@ class BaseStep(BaseComponent): @property def embedding(self) -> BaseEmbeddingModel: name = self.kwargs.get("embedding", "default") - return ( - name - if isinstance(name, BaseEmbeddingModel) - else self.get_component( - ComponentEnum.EMBEDDING_MODEL, - name, - ) - ) + if isinstance(name, BaseEmbeddingModel): + return name + else: + return self.get_component(ComponentEnum.EMBEDDING_MODEL, name) def prompt_format(self, prompt_name: str, **kwargs) -> str: return self.prompt.prompt_format(prompt_name=prompt_name, **kwargs)