Workbench (aka "the tool pool") design proposal

[JMLX42]

What feature would you like to be added?

Hello,

following up on #2666 (comment), I would like to propose a design for an abstraction on List[Tool]. That abstraction would be a "tool pool" that makes it possible to add, remove and find tools. It would make it easier to make agents that:

• have access to virtually limitless set of tools
• dynamically handle the actual set of available tools for a task
• use various complex (other agents, LLMs, vector distance search...) to select/elect tools

Introducing: the workbench.

class Workbench(ABC):
    @property
    @abstractmethod
    def tools(self) -> List[Tool]: ...

    @abstractmethod
    def get_tools_for_context(self, context: Sequence[LLMMessage]) -> List[Tool]: ...

    @abstractmethod
    def get_tool_names_for_query(self, query: str) -> List[str]: ...

    @abstractmethod
    def get_tool_documentation(self, tool_name: str) -> str: ...


class DynamicWorkbench(ABC):
    @abstractmethod
    def add_tool(
        self, tool: Tool | Callable[..., Any] | Callable[..., Awaitable[Any]], enabled: bool = True
    ) -> None: ...

    @abstractmethod
    def remove_tool(self, tool_name: str) -> None: ...

    @abstractmethod
    def remove_all_tools(self) -> None: ...

    @abstractmethod
    def enable_tool(self, tool_name: str) -> None: ...

    @abstractmethod
    def disable_tool(self, tool_name: str) -> None: ...

    @abstractmethod
    def tool_is_enabled(self, tool_name: str) -> bool: ...

    @property
    @abstractmethod
    def enabled_tools(self) -> List[Tool]: ...

An example implementation: VectorWorkbench is a workbench that uses vector distance matching on tools description/documentation to find/elect tools.

class VectorWorkbench(Workbench, DynamicWorkbench):
   def get_tools_for_context(self, context: Sequence[LLMMessage]) -> List[Tool]:
        query = "\n".join([msg.content for msg in context if isinstance(msg, TextMessage)]).strip()
        return self._get_tools_for_query(query)
   
   def _get_tools_for_query(self, query: str) -> List[str]:
        if len(query) == 0:
            return []

        logger.info(f"Query: {query}")

        # Generate the query embedding.
        query_embedding = self._embedding_func(query).data[0]["embedding"]
        query_embedding = np.array(query_embedding).reshape(1, -1)

        # Search for the closest matching tool embeddings.
        k = min(self._top_k, len(self._tools))
        scores, indexes = self._index.search(query_embedding, k=k)  # Retrieve top matches
        results = [
            {"tool": self._tools[i], "score": scores[0][index]}
            for index, i in enumerate(indexes[0])
            if scores[0][index] <= self._score_threshold  # Filter out results below the threshold
            and self.tool_is_enabled(self._tools[i].name)  # Filter out disabled tools
        ]

        return [result["tool"] for result in results]

   def get_tool_names_for_query(self, query: str) -> List[str]:
        """Returns a list of tool names that are relevant to the given query."""
        logger.info(f"getting tool names for query: {query}")
        tools = self._get_tools_for_query(query)
        return [tool.name for tool in tools]

    def get_tool_documentation(self, tool_name: str) -> str:
        """Returns the documentation for the given tool name."""
        logger.info(f"getting documentation for tool: {tool_name}")
        # We assume that all tools have a documentation.
        if tool_name not in self._tool_docs:
            raise ValueError(f"tool {tool_name} does not exist")
        return self._tool_docs[tool_name]

   # the rest is boilerplate...

Then, the following agents can be provided OOB:

class WorkbenchSearchAgent(AssistantAgent):
    """ Searches for tools, but does *not* call them. Good for planning for example."""
    def __init__(
        self,
        name: str,
        model_client: ChatCompletionClient,
        workbench: Workbench,
        handoffs: List[Handoff | str] | None = None,
        description: str = "An agent that provides assistance to find relevant tools.",
        system_message: str | None = "You are a helpful AI assistant to find relevant tools to solve a problem.",
    ):
        super().__init__(
            name=name,
            model_client=model_client,
            handoffs=handoffs,
            description=description,
            system_message=system_message,
            tools=[
                workbench.get_tool_names_for_query,
                workbench.get_tool_documentation,
            ],
        )
        self._workbench = workbench

    @property
    def produced_message_types(self) -> List[type[ChatMessage]]:
        return [TextMessage]

    @property
    def workbench(self) -> Workbench:
        return self._workbench


class WorkbenchAgent(AssistantAgent):
    """Find and actually call tools. Good for the plan execution for example."""
    def __init__(
        self,
        name: str,
        model_client: ChatCompletionClient,
        workbench: Workbench,
        handoffs: List[Handoff | str] | None = None,
        description: str = "An agent that provides assistance to find relevant tools.",
        system_message: str | None = "You are a helpful AI assistant to find relevant tools to solve a problem.",
    ):
        super().__init__(
            name=name,
            description=description,
            model_client=model_client,
            system_message=system_message,
            # The tools are dynamically set based on the message list using
            # embedding search.
            tools=[],
            handoffs=handoffs,
        )

        self._workbench = workbench

    @property
    def workbench(self) -> Workbench:
        return self._workbench

    @property
    def produced_message_types(self) -> List[type[ChatMessage]]:
        return [ToolCallMessage, ToolCallResultMessage]

    async def on_messages_stream(
        self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken
    ) -> AsyncGenerator[AgentMessage | Response, None]:
        self.remove_all_tools()
        tools = self._workbench.get_tools_for_context(messages)
        logger.info(f"selected tools: {[tool.name for tool in tools]}")
        for tool in tools:
            self.add_tool(tool)

        async for msg in super().on_messages_stream(messages, cancellation_token):
            yield msg

        self.remove_all_tools()

Used like this:

workbench_search_agent = WorkbenchSearchAgent(
        name="workbench_search_agent",
        model_client=model_client,
        workbench=openapi_agent.workbench,
        system_message="You are a helpful AI assistant to find tools. When you are done, handoff to planning_agent.",
        handoffs=[
            Handoff(
                target="planning_agent",
                message="Here are the tools your are looking for.",
            ),
        ],
    )

    planning_agent = AssistantAgent(
        name="planning_agent",
        model_client=model_client,
        system_message="You are a planning agent. Your task is to create a step-by-step plan for the given task."
        "In order to plan, you need to know the tools you can use."
        f"You start by asking {workbench_search_agent.name} to find the relevant tools."
        "Finally, you create a step-by-step plan for the given task."
        "When you are done, handoff to coordinator_agent.",
        handoffs=[
            "coordinator_agent",
            Handoff(
                target=workbench_search_agent.name,
                message="Find the tools to use for the given task.",
            ),
        ],
    )

Why is this needed?

1. Challenges in Current Tool Management

• Static Tool Sets: Current systems rely on pre-defined, fixed tools, limiting flexibility and adaptability to different tasks.
• Dynamic Task Requirements: Tasks often need tools tailored to specific contexts, but static systems can’t dynamically adapt.
• Scalability Issues: Managing large pools of tools becomes inefficient without dynamic filtering or selection mechanisms.

2. The Role of the Workbench

• Centralized Tool Management: The Workbench provides a unified system to add, remove, enable, and filter tools dynamically.
• Context-Aware Tool Selection: Uses techniques like vector search to find tools relevant to the current task or query.
• Agent Collaboration: Supports modular agents (e.g., Search Agent for finding tools and Execution Agent for using them) to streamline workflows.

3. Advantages for Autogen

• Scalability: Handles large, virtually unlimited toolsets dynamically and efficiently.
• Flexibility: Adapts tools in real-time to match task-specific needs.
• Extensibility: Compatible with various retrieval methods, enabling advanced filtering mechanisms like LLM embeddings or vector matching.
• Better Agent Collaboration: Facilitates seamless interactions between agents using a shared tool resource.

[ekzhu]

This is an interesting idea. Would you be available to discuss this in our community office hour meeting? Every wednesday 10 AM PST.

See: #4059

[JMLX42]

@ekzhu sure thing! See you tomorrow.