autogen_core.tools._workbench 源代码

from abc import ABC, abstractmethod
from types import TracebackType
from typing import Any, List, Literal, Mapping, Optional, Type

from pydantic import BaseModel, Field
from typing_extensions import Annotated, Self

from .._cancellation_token import CancellationToken
from .._component_config import ComponentBase
from .._image import Image
from ._base import ToolSchema




class TextResultContent(BaseModel):
    """
    工具执行结果的文本内容。
    """

    type: Literal["TextResultContent"] = "TextResultContent"

    content: str
    """结果的文本内容。"""





class ImageResultContent(BaseModel):
    """
    工具执行结果的图像内容。
    """

    type: Literal["ImageResultContent"] = "ImageResultContent"

    content: Image
    """结果的图像内容。"""



ResultContent = Annotated[TextResultContent | ImageResultContent, Field(discriminator="type")]




class ToolResult(BaseModel):
    """
    工作台中工具执行的结果。
    """

    type: Literal["ToolResult"] = "ToolResult"

    name: str
    """已执行工具的名称。"""

    result: List[ResultContent]
    """工具执行的结果。"""

    is_error: bool = False
    """工具执行是否导致错误。"""



    def to_text(self, replace_image: str | None = None) -> str:
        """
        将结果转换为文本字符串。

        Args:
            replace_image (str | None): 用于替换图像内容的字符串。
                如果为None，图像内容将以base64字符串形式包含在文本中。

        Returns:
            str: 结果的文本表示形式。
        """
        parts: List[str] = []
        for content in self.result:
            if isinstance(content, TextResultContent):
                parts.append(content.content)
            elif isinstance(content, ImageResultContent):
                if replace_image is not None:
                    parts.append(replace_image)
                else:
                    parts.append(f"[Image: {content.content.to_base64()}]")
        return "\n".join(parts)






class Workbench(ABC, ComponentBase[BaseModel]):
    """
    工作台是一个提供一组可能共享资源和状态的工具的组件。

    工作台负责管理工具的生命周期，并提供调用它们的统一接口。工作台提供的工具可能是动态的，
    每次工具执行后它们的可用性可能会发生变化。

    可以通过调用 :meth:`~autogen_core.tools.Workbench.start` 方法启动工作台，
    通过调用 :meth:`~autogen_core.tools.Workbench.stop` 方法停止工作台。
    它也可以作为异步上下文管理器使用，在进入和退出上下文时会自动启动和停止工作台。
    """

    component_type = "workbench"



    @abstractmethod
    async def list_tools(self) -> List[ToolSchema]:
        """
        以 :class:`ToolSchema` 对象形式列出工作台中当前可用的工具。

        工具列表可能是动态的，其内容可能在工具执行后发生变化。
        """
        ...




    @abstractmethod
    async def call_tool(
        self, name: str, arguments: Mapping[str, Any] | None = None, cancellation_token: CancellationToken | None = None
    ) -> ToolResult:
        """
        调用工作台中的工具。

        Args:
            name (str): 要调用的工具名称。
            arguments (Mapping[str, Any] | None): 传递给工具的参数。
                如果为None，工具将在无参数情况下被调用。
            cancellation_token (CancellationToken | None): 可选的取消令牌，
                用于取消工具执行。
        Returns:
            ToolResult: 工具执行的结果。
        """
        ...




    @abstractmethod
    async def start(self) -> None:
        """
        启动工作台并初始化所有资源。

        在使用工作台前应调用此方法。
        """
        ...




    @abstractmethod
    async def stop(self) -> None:
        """
        停止工作台并释放所有资源。

        当不再需要工作台时应调用此方法。
        """
        ...




    @abstractmethod
    async def reset(self) -> None:
        """
        将工作台重置为初始化后的启动状态。
        """
        ...




    @abstractmethod
    async def save_state(self) -> Mapping[str, Any]:
        """
        保存工作台的状态。

        调用此方法以持久化工作台的状态。
        """
        ...




    @abstractmethod
    async def load_state(self, state: Mapping[str, Any]) -> None:
        """
        加载工作台的状态。

        Args:
            state (Mapping[str, Any]): 要加载到工作台中的状态。
        """
        ...


    async def __aenter__(self) -> Self:
        """
        进入工作台上下文管理器。

        当工作台在 `with` 语句中使用时调用此方法。
        它会调用 :meth:`~autogen_core.tools.WorkBench.start` 方法来启动工作台。
        """
        await self.start()
        return self

    async def __aexit__(
        self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
    ) -> None:
        """
        退出工作台上下文管理器。
        当工作台在 `with` 语句中使用时调用此方法。
        它会调用 :meth:`~autogen_core.tools.WorkBench.stop` 方法来停止工作台。
        """
        await self.stop()