From 68bd95b494e46518d600e954e39701d6e398f423 Mon Sep 17 00:00:00 2001
From: jinliyl <6469360+jinliyl@users.noreply.github.com>
Date: Mon, 18 May 2026 15:00:02 +0800
Subject: [PATCH] =?UTF-8?q?refactor(steps):=20Add=20job=20management=20met?=
 =?UTF-8?q?hods=20and=20support=20registering=20t=E2=80=A6=20(#242)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* refactor(steps): Add job management methods and support registering them as tools

Added methods to the `BaseStep` class for retrieving, running, and registering jobs as tools, enhancing the functionality of the step class.

* fix doc
---
 reme4/steps/base_step.py | 34 ++++++++++++++++++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/reme4/steps/base_step.py b/reme4/steps/base_step.py
index a710f457..d32c4445 100644
--- a/reme4/steps/base_step.py
+++ b/reme4/steps/base_step.py
@@ -5,8 +5,10 @@ from abc import abstractmethod, ABC
 from typing import TypeVar, TYPE_CHECKING
 
 from agentscope.formatter import FormatterBase
+from agentscope.message import TextBlock
 from agentscope.model import ChatModelBase
 from agentscope.token import TokenCounterBase
+from agentscope.tool import Toolkit, ToolResponse
 
 from ..components.embedding import BaseEmbeddingModel
 from ..components.file_parser import BaseFileParser
@@ -15,10 +17,12 @@ from ..components.file_watcher import BaseFileWatcher
 from ..components.prompt_handler import PromptHandler
 from ..components.runtime_context import RuntimeContext
 from ..enumeration import ComponentEnum
+from ..schema import Response
 from ..utils import get_logger
 
 if TYPE_CHECKING:
     from ..components import ApplicationContext
+    from ..components.job import BaseJob
 
 T = TypeVar("T")
 
@@ -138,3 +142,33 @@ class BaseStep(ABC):
     def copy(self, **kwargs) -> "BaseStep":
         """Construct a new instance from the original init args, applying overrides."""
         return self.__class__(*self._init_args, **{**self._init_kwargs, **kwargs})
+
+    def get_job(self, name: str) -> "BaseJob | None":
+        """Return a job by name."""
+        if self.app_context is None:
+            raise RuntimeError("Cannot get job without an app context")
+        return self.app_context.jobs.get(name)
+
+    async def run_job(self, name: str, **kwargs) -> Response:
+        """Execute a job by name and kwargs, return the final response."""
+        job: "BaseJob | None" = self.get_job(name)
+        if job is None:
+            raise RuntimeError(f"Job {name} not found")
+        return await job(**kwargs)
+
+    def add_as_tool(self, toolkit: Toolkit, job_name: str) -> None:
+        """Add the step as a tool to the toolkit."""
+        job: "BaseJob | None" = self.get_job(job_name)
+        if job is None:
+            raise RuntimeError(f"Job {job_name} not found")
+
+        async def run_job(**kwargs) -> ToolResponse:
+            response = await job(**kwargs)
+            return ToolResponse(content=[TextBlock(type="text", text=response.answer)])
+
+        toolkit.register_tool_function(
+            tool_func=run_job,
+            func_name=job_name,
+            func_description=job.description,
+            json_schema=job.parameters,
+        )