Introduce workbench

python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/llamaindex-agent.ipynb

@@ -38,7 +38,6 @@
    "outputs": [],
    "source": [
     "import os\n",
-    "from pydantic import BaseModel\n",
     "from typing import List, Optional\n",
     "\n",
     "from autogen_core import AgentId, MessageContext, RoutedAgent, SingleThreadedAgentRuntime, message_handler\n",
@@ -57,7 +56,8 @@
     "from llama_index.embeddings.openai import OpenAIEmbedding\n",
     "from llama_index.llms.azure_openai import AzureOpenAI\n",
     "from llama_index.llms.openai import OpenAI\n",
-    "from llama_index.tools.wikipedia import WikipediaToolSpec"
+    "from llama_index.tools.wikipedia import WikipediaToolSpec\n",
+    "from pydantic import BaseModel"
    ]
   },
   {

python/packages/autogen-core/src/autogen_core/_component_config.py

@@ -7,7 +7,7 @@
 from pydantic import BaseModel
 from typing_extensions import Self, TypeVar
 
-ComponentType = Literal["model", "agent", "tool", "termination", "token_provider"] | str
+ComponentType = Literal["model", "agent", "tool", "termination", "token_provider", "workbench"] | str
 ConfigT = TypeVar("ConfigT", bound=BaseModel)
 FromConfigT = TypeVar("FromConfigT", bound=BaseModel, contravariant=True)
 ToConfigT = TypeVar("ToConfigT", bound=BaseModel, covariant=True)

python/packages/autogen-core/src/autogen_core/tools/__init__.py

@@ -1,5 +1,7 @@
 from ._base import BaseTool, BaseToolWithState, ParametersSchema, Tool, ToolSchema
 from ._function_tool import FunctionTool
+from ._static_workbench import StaticWorkbench
+from ._workbench import ImageResultContent, TextResultContent, ToolResult, Workbench
 
 __all__ = [
     "Tool",
@@ -8,4 +10,9 @@
     "BaseTool",
     "BaseToolWithState",
     "FunctionTool",
+    "Workbench",
+    "ToolResult",
+    "TextResultContent",
+    "ImageResultContent",
+    "StaticWorkbench",
 ]

python/packages/autogen-core/src/autogen_core/tools/_function_tool.py

@@ -178,4 +178,4 @@ def _from_config(cls, config: FunctionToolConfig) -> Self:
         if not callable(func):
             raise TypeError(f"Expected function but got {type(func)}")
 
-        return cls(func, "", None)
+        return cls(func, name=config.name, description=config.description, global_imports=config.global_imports)

python/packages/autogen-core/src/autogen_core/tools/_static_workbench.py

@@ -0,0 +1,88 @@
+import asyncio
+from typing import Any, Dict, List, Literal, Mapping
+
+from pydantic import BaseModel
+from typing_extensions import Self
+
+from .._cancellation_token import CancellationToken
+from .._component_config import Component, ComponentModel
+from ._base import BaseTool, ToolSchema
+from ._workbench import TextResultContent, ToolResult, Workbench
+
+
+class StaticWorkbenchConfig(BaseModel):
+    tools: List[ComponentModel] = []
+
+
+class StateicWorkbenchState(BaseModel):
+    type: Literal["StaticWorkbenchState"] = "StaticWorkbenchState"
+    tools: Dict[str, Mapping[str, Any]] = {}
+
+
+class StaticWorkbench(Workbench, Component[StaticWorkbenchConfig]):
+    """
+    A workbench that provides a static set of tools that do not change after
+    each tool execution.
+
+    Args:
+        tools (List[BaseTool[Any, Any]]): A list of tools to be included in the workbench.
+            The tools should be subclasses of :class:`~autogen_core.tools.BaseTool`.
+    """
+
+    component_provider_override = "autogen_core.tools.StaticWorkbench"
+    component_config_schema = StaticWorkbenchConfig
+
+    def __init__(self, tools: List[BaseTool[Any, Any]]) -> None:
+        self._tools = tools
+
+    async def list_tools(self) -> List[ToolSchema]:
+        return [tool.schema for tool in self._tools]
+
+    async def call_tool(
+        self, name: str, arguments: Mapping[str, Any] | None = None, cancellation_token: CancellationToken | None = None
+    ) -> ToolResult:
+        tool = next((tool for tool in self._tools if tool.name == name), None)
+        if tool is None:
+            raise ValueError(f"Tool {name} not found in workbench.")
+        if not cancellation_token:
+            cancellation_token = CancellationToken()
+        if not arguments:
+            arguments = {}
+        try:
+            result_future = asyncio.ensure_future(tool.run_json(arguments, cancellation_token))
+            cancellation_token.link_future(result_future)
+            result = await result_future
+            is_error = False
+        except Exception as e:
+            result = str(e)
+            is_error = True
+        result_str = tool.return_value_as_string(result)
+        return ToolResult(name=tool.name, result=[TextResultContent(content=result_str)], is_error=is_error)
+
+    async def start(self) -> None:
+        return None
+
+    async def stop(self) -> None:
+        return None
+
+    async def reset(self) -> None:
+        return None
+
+    async def save_state(self) -> Mapping[str, Any]:
+        tool_states = StateicWorkbenchState()
+        for tool in self._tools:
+            tool_states.tools[tool.name] = await tool.save_state_json()
+        return tool_states.model_dump()
+
+    async def load_state(self, state: Mapping[str, Any]) -> None:
+        parsed_state = StateicWorkbenchState.model_validate(state)
+        for tool in self._tools:
+            if tool.name in parsed_state.tools:
+                await tool.load_state_json(parsed_state.tools[tool.name])
+
+    def _to_config(self) -> StaticWorkbenchConfig:
+        return StaticWorkbenchConfig(tools=[tool.dump_component() for tool in self._tools])
+
+    @classmethod
+    def _from_config(cls, config: StaticWorkbenchConfig) -> Self:
+        return cls(tools=[BaseTool.load_component(tool) for tool in config.tools])

python/packages/autogen-core/src/autogen_core/tools/_workbench.py

@@ -0,0 +1,164 @@
+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):
+    """
+    Text result content of a tool execution.
+    """
+
+    type: Literal["TextResultContent"] = "TextResultContent"
+
+    content: str
+    """The text content of the result."""
+
+
+class ImageResultContent(BaseModel):
+    """
+    Image result content of a tool execution.
+    """
+
+    type: Literal["ImageResultContent"] = "ImageResultContent"
+
+    content: Image
+    """The image content of the result."""
+
+
+ResultContent = Annotated[TextResultContent | ImageResultContent, Field(discriminator="type")]
+
+
+class ToolResult(BaseModel):
+    """
+    A result of a tool execution by a workbench.
+    """
+
+    type: Literal["ToolResult"] = "ToolResult"
+
+    name: str
+    """The name of the tool that was executed."""
+
+    result: List[ResultContent]
+    """The result of the tool execution."""
+
+    is_error: bool = False
+    """Whether the tool execution resulted in an error."""
+
+
+class Workbench(ABC, ComponentBase[BaseModel]):
+    """
+    A workbench is a component that provides a set of tools that may share
+    resources and state.
+
+    A workbench is responsible for managing the lifecycle of the tools and
+    providing a single interface to call them. The tools provided by the workbench
+    may be dynamic and their availabilities may change after each tool execution.
+
+    A workbench can be started by calling the :meth:`~autogen_core.tools.Workbench.start` method
+    and stopped by calling the :meth:`~autogen_core.tools.Workbench.stop` method.
+    It can also be used as an asynchronous context manager, which will automatically
+    start and stop the workbench when entering and exiting the context.
+    """
+
+    component_type = "workbench"
+
+    @abstractmethod
+    async def list_tools(self) -> List[ToolSchema]:
+        """
+        List the currently available tools in the workbench as :class:`ToolSchema`
+        objects.
+
+        The list of tools may be dynamic, and their content may change after
+        tool execution.
+        """
+        ...
+
+    @abstractmethod
+    async def call_tool(
+        self, name: str, arguments: Mapping[str, Any] | None = None, cancellation_token: CancellationToken | None = None
+    ) -> ToolResult:
+        """
+        Call a tool in the workbench.
+
+        Args:
+            name (str): The name of the tool to call.
+            arguments (Mapping[str, Any] | None): The arguments to pass to the tool.
+                If None, the tool will be called with no arguments.
+            cancellation_token (CancellationToken | None): An optional cancellation token
+                to cancel the tool execution.
+        Returns:
+            ToolResult: The result of the tool execution.
+        """
+        ...
+
+    @abstractmethod
+    async def start(self) -> None:
+        """
+        Start the workbench and initialize any resources.
+
+        This method should be called before using the workbench.
+        """
+        ...
+
+    @abstractmethod
+    async def stop(self) -> None:
+        """
+        Stop the workbench and release any resources.
+
+        This method should be called when the workbench is no longer needed.
+        """
+        ...
+
+    @abstractmethod
+    async def reset(self) -> None:
+        """
+        Reset the workbench to its initialized, started state.
+        """
+        ...
+
+    @abstractmethod
+    async def save_state(self) -> Mapping[str, Any]:
+        """
+        Save the state of the workbench.
+
+        This method should be called to persist the state of the workbench.
+        """
+        ...
+
+    @abstractmethod
+    async def load_state(self, state: Mapping[str, Any]) -> None:
+        """
+        Load the state of the workbench.
+
+        Args:
+            state (Mapping[str, Any]): The state to load into the workbench.
+        """
+        ...
+
+    async def __aenter__(self) -> Self:
+        """
+        Enter the workbench context manager.
+
+        This method is called when the workbench is used in a `with` statement.
+        It calls the :meth:`~autogen_core.tools.WorkBench.start` method to start the workbench.
+        """
+        await self.start()
+        return self
+
+    async def __aexit__(
+        self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
+    ) -> None:
+        """
+        Exit the workbench context manager.
+        This method is called when the workbench is used in a `with` statement.
+        It calls the :meth:`~autogen_core.tools.WorkBench.stop` method to stop the workbench.
+        """
+        await self.stop()