Tool Selection for LLM Agents: Routing Strategies and Implementation › Code Implementation: Building a Tool Router

from dataclasses import dataclass, field
from typing import Dict, List, Tuple

import numpy as np
from sentence_transformers import SentenceTransformer


@dataclass
class Tool:
    name: str
    description: str
    parameters: Dict
    examples: List[str] = field(default_factory=list)


class HybridToolRouter:
    def __init__(self, embedding_model: str = "all-MiniLM-L6-v2"):
        self.tools: List[Tool] = []
        self.embedder = SentenceTransformer(
            embedding_model
        )  # Run once locally, then comment out
        self.tool_embeddings = None
        self.tool_texts = []

    def register_tool(self, tool: Tool):
        """Add a tool to the registry."""
        self.tools.append(tool)
        # Create rich text representation for embedding
        text = f"{tool.name}: {tool.description}"
        if tool.examples:
            text += f" Examples: {'; '.join(tool.examples)}"
        self.tool_texts.append(text)

    def build_index(self):
        """Compute embeddings for all tools."""
        if not self.tool_texts:
            raise ValueError("No tools registered")
        self.tool_embeddings = self.embedder.encode(self.tool_texts)

    def retrieve_candidates(
        self, query: str, top_k: int = 3
    ) -> List[Tuple[Tool, float]]:
        """Retrieve top-k candidate tools using embeddings."""
        if self.tool_embeddings is None:
            self.build_index()

        query_embedding = self.embedder.encode([query])[0]

        # Compute cosine similarities
        similarities = np.dot(self.tool_embeddings, query_embedding) / (
            np.linalg.norm(self.tool_embeddings, axis=1)
            * np.linalg.norm(query_embedding)
        )

        # Get top-k indices
        top_indices = np.argsort(similarities)[-top_k:][::-1]

        candidates = []
        for idx in top_indices:
            candidates.append((self.tools[idx], float(similarities[idx])))

        return candidates

# Initialize router and register tools
router = HybridToolRouter()

# Register our tool ecosystem
router.register_tool(
    Tool(
        name="calculator",
        description="Evaluate mathematical expressions and perform calculations",
        parameters={"expression": "string"},
        examples=["What is 15 times 23?", "Calculate the square root of 144"],
    )
)

router.register_tool(
    Tool(
        name="weather",
        description="Get current weather conditions and forecasts for locations",
        parameters={"location": "string", "days": "integer"},
        examples=[
            "What's the weather in Paris?",
            "Will it rain tomorrow in London?",
        ],
    )
)

router.register_tool(
    Tool(
        name="search",
        description="Search for current information on the internet",
        parameters={"query": "string"},
        examples=["Who is the current president?", "Latest news about AI"],
    )
)

router.register_tool(
    Tool(
        name="calendar",
        description="Check schedule and create calendar events",
        parameters={"date": "string", "event": "string"},
        examples=["Do I have meetings tomorrow?", "Schedule a call for 3pm"],
    )
)

router.register_tool(
    Tool(
        name="translate",
        description="Translate text between languages",
        parameters={"text": "string", "target_language": "string"},
        examples=[
            "Translate 'hello' to French",
            "How do you say 'thank you' in Japanese?",
        ],
    )
)

router.build_index()

# Test retrieval on ambiguous and clear queries
test_cases = [
    "What's 25 divided by 5?",  # Clear: calculator
    "Will it rain in Seattle tomorrow?",  # Clear: weather
    "How do you say 'goodbye' in Spanish?",  # Clear: translate
    "Find me the latest stock prices",  # Ambiguous: search vs calculator
    "Schedule a meeting for 2pm",  # Clear: calendar
    "Calculate the weather forecast for next week",  # Nonsensical: should fail
]

from typing import Optional


def simulate_llm_selection(
    query: str, candidates: List[Tuple[Tool, float]]
) -> Optional[Tool]:
    """
    Simulate LLM-based selection from candidates.
    In production, this would be an actual LLM API call.
    """
    if not candidates:
        return None

    best_tool, best_score = candidates[0]

    # Decline if best match is below threshold
    if best_score < 0.3:
        return None

    # Check for parameter compatibility (simulated)
    if best_tool.name == "calculator" and not any(c.isdigit() for c in query):
        # No numbers in query, likely not a calculation
        if len(candidates) > 1 and candidates[1][1] > 0.25:
            return candidates[1][0]  # Select second best

    return best_tool


def generate_selection_prompt(
    query: str, candidates: List[Tuple[Tool, float]]
) -> str:
    """Generate the prompt that would be sent to an LLM for final selection."""
    prompt = f'User query: "{query}"\n\nAvailable tools:\n'
    for i, (tool, score) in enumerate(candidates, 1):
        prompt += f"\n{i}. {tool.name}: {tool.description}\n"
        prompt += f"   Parameters: {json.dumps(tool.parameters)}\n"
        prompt += f"   Examples: {tool.examples}\n"
    prompt += "\nSelect the most appropriate tool by name, or respond 'none' if no tool is suitable."
    return prompt