Closed
Description
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.
Metadata
Metadata
Assignees
Labels
No labels