Lesson 14 of 23

Multi-Agent System Design

Tool Registry Design

4 min read

A tool registry is a central system for managing, discovering, and providing tools to agents. Good registry design enables agents to dynamically select the right tools for any task.

Why a Tool Registry?

Problem How Registry Solves It
Too many tools for context Semantic search for relevant tools
Tools change frequently Dynamic registration, no agent restart
Tool versioning Multiple versions, gradual migration
Access control Per-agent tool permissions

Core Registry Architecture

from dataclasses import dataclass
from typing import Callable, Optional
import json

@dataclass
class Tool:
    name: str
    description: str
    parameters: dict  # JSON Schema
    handler: Callable
    version: str = "1.0.0"
    category: str = "general"
    embedding: Optional[list] = None  # For semantic search

class ToolRegistry:
    def __init__(self, embedding_model):
        self.tools = {}
        self.embedder = embedding_model
        self.embeddings = {}  # name -> embedding

    async def register(self, tool: Tool):
        """Register a tool with semantic embedding."""
        # Generate embedding for semantic discovery
        description_text = f"{tool.name}: {tool.description}"
        tool.embedding = await self.embedder.embed(description_text)

        self.tools[tool.name] = tool
        self.embeddings[tool.name] = tool.embedding

    def get(self, name: str) -> Optional[Tool]:
        """Get tool by exact name."""
        return self.tools.get(name)

    async def search(self, query: str, top_k: int = 5) -> list:
        """Semantic search for relevant tools."""
        query_embedding = await self.embedder.embed(query)

        scores = []
        for name, embedding in self.embeddings.items():
            score = self._cosine_similarity(query_embedding, embedding)
            scores.append((name, score))

        # Sort by similarity
        scores.sort(key=lambda x: x[1], reverse=True)

        return [self.tools[name] for name, _ in scores[:top_k]]

    def list_by_category(self, category: str) -> list:
        """List tools by category."""
        return [t for t in self.tools.values() if t.category == category]

    def to_openai_format(self, tools: list) -> list:
        """Convert tools to OpenAI function calling format."""
        return [
            {
                "type": "function",
                "function": {
                    "name": tool.name,
                    "description": tool.description,
                    "parameters": tool.parameters
                }
            }
            for tool in tools
        ]

Dynamic Tool Selection

Select relevant tools based on the user's query:

class SmartAgent:
    def __init__(self, llm, registry: ToolRegistry, max_tools: int = 10):
        self.llm = llm
        self.registry = registry
        self.max_tools = max_tools

    async def run(self, task: str) -> str:
        # Step 1: Find relevant tools for this task
        relevant_tools = await self.registry.search(task, top_k=self.max_tools)

        # Step 2: Run agent with selected tools
        tools_formatted = self.registry.to_openai_format(relevant_tools)

        messages = [{"role": "user", "content": task}]

        while True:
            response = await self.llm.complete(
                messages=messages,
                tools=tools_formatted
            )

            if response.tool_calls:
                for call in response.tool_calls:
                    tool = self.registry.get(call.function.name)
                    result = await tool.handler(**json.loads(call.function.arguments))
                    messages.append({
                        "role": "tool",
                        "content": str(result),
                        "tool_call_id": call.id
                    })
            else:
                return response.content

Tool Versioning

Support multiple versions for gradual migration:

class VersionedToolRegistry(ToolRegistry):
    def __init__(self, embedding_model):
        super().__init__(embedding_model)
        self.versions = {}  # name -> {version -> Tool}

    async def register(self, tool: Tool):
        """Register with version tracking."""
        if tool.name not in self.versions:
            self.versions[tool.name] = {}

        self.versions[tool.name][tool.version] = tool

        # Latest version is the default
        await super().register(tool)

    def get(self, name: str, version: str = "latest") -> Optional[Tool]:
        """Get specific version of a tool."""
        if name not in self.versions:
            return None

        if version == "latest":
            return self.tools.get(name)

        return self.versions[name].get(version)

    def deprecate(self, name: str, version: str, replacement_version: str):
        """Mark a version as deprecated."""
        if name in self.versions and version in self.versions[name]:
            tool = self.versions[name][version]
            tool.deprecated = True
            tool.replacement = replacement_version

Tool Categories and Permissions

@dataclass
class ToolPermission:
    tool_name: str
    allowed_agents: list  # Agent IDs that can use this tool
    rate_limit: Optional[int] = None  # Calls per minute
    requires_approval: bool = False

class SecureToolRegistry(ToolRegistry):
    def __init__(self, embedding_model):
        super().__init__(embedding_model)
        self.permissions = {}

    def set_permission(self, permission: ToolPermission):
        self.permissions[permission.tool_name] = permission

    async def get_for_agent(self, agent_id: str, query: str) -> list:
        """Get tools the agent is allowed to use."""
        # Get relevant tools
        relevant = await self.search(query)

        # Filter by permissions
        allowed = []
        for tool in relevant:
            perm = self.permissions.get(tool.name)
            if perm is None or agent_id in perm.allowed_agents:
                allowed.append(tool)

        return allowed

Tool Composition

Build complex tools from simpler ones:

class ComposedTool:
    """A tool that combines multiple tools."""

    def __init__(self, name: str, steps: list):
        self.name = name
        self.steps = steps  # List of (tool_name, input_transform)

    async def execute(self, registry: ToolRegistry, initial_input: dict) -> str:
        current_result = initial_input

        for tool_name, transform in self.steps:
            tool = registry.get(tool_name)
            # Transform previous result to tool input
            tool_input = transform(current_result)
            current_result = await tool.handler(**tool_input)

        return current_result

# Example: Composed "research_and_summarize" tool
research_summarize = ComposedTool(
    name="research_and_summarize",
    steps=[
        ("web_search", lambda x: {"query": x["topic"]}),
        ("summarize", lambda x: {"text": x["results"]})
    ]
)

Interview Tip

When discussing tool registries, highlight:

  1. Semantic discovery - Agents don't need to know all tools
  2. Scalability - Handle hundreds of tools efficiently
  3. Security - Not all agents should access all tools
  4. Evolution - Tools change, agents shouldn't break

Next, we'll explore state and memory management for agents. :::

Quiz

Module 4: Multi-Agent System Design

Take Quiz