Weaviate Cookbooks

import dspy
from typing import Any, Dict

class AgentResponse(dspy.Signature):

    # Input Fields
    history: dspy.History = dspy.InputField()
    user_prompt: str = dspy.InputField()
    available_tools: str = dspy.InputField()

    # Output Fields
    response: str = dspy.OutputField(
        description="The response to the user's prompt whilst the tool is running. Update the user on the progress of their request (if a tool is picked), or the final response to the user (if no tool is picked)."
    )
    tool: str | None = dspy.OutputField(
        description="The tool that needs to be used. Return None if no tool is needed."
    )
    tool_inputs: Dict[str, Any] | None = dspy.OutputField(
        description=(
            "The inputs for the tool. Return an empty dictionary (still include the field) if no inputs are needed. "
            "The key is the name of the input, the value is the value of the input."
        )
    )

from typing import Callable, List, Tuple

class RouterAgent:
    def __init__(self, model: str, tools: List[Callable] = []):
        self.tools: list[Callable] = tools
        self.model = dspy.LM(model)
        self.agent = dspy.ChainOfThought(AgentResponse)
        self.conversation_history = dspy.History(messages=[])

    def add_conversation_history(self, message: str, response: dspy.Prediction):
        self.conversation_history.messages.append({"user_prompt": message, **response})

    def get_tools_and_descriptions(self) -> str:
        return "\n".join(
            [
                f"{tool.__name__}:\nDescription: {tool.__doc__ or ''}\nInputs: { {k: v for k, v in tool.__annotations__.items() if k != 'return'} }"
                for tool in self.tools
            ]
        )

    def get_response(self, user_prompt: str) -> Tuple[str, str | None]:
        result = self.agent(
            history=self.conversation_history,
            user_prompt=user_prompt,
            available_tools=self.get_tools_and_descriptions(),
            lm=self.model,
        )
        self.add_conversation_history(message=user_prompt, response=result)
        if result.tool and result.tool.lower() not in ["null", "none"]:
            tool_function = next(
                (tool for tool in self.tools if tool.__name__ == result.tool), None
            )
            if tool_function is None:
                raise ValueError(f"Tool {result.tool} not found")
            tool_inputs = {k: v for k, v in result.tool_inputs.items() if k != "return"}
            tool_result = tool_function(**tool_inputs)
        else:
            tool_result = None
        return result.response, tool_result

router = RouterAgent(
    model="<model_name>",  # e.g. claude-sonnet-4-5, gpt-5.2, gemini-2.5-pro
    tools=[your_tool_function]
)
response, tool_result = router.get_response("user query here")

def your_tool(param1: str, param2: int) -> str:
    """Clear description of what this tool does and when to use it."""
    # tool logic here
    return "result as string"

class AgentFollowup(dspy.Signature):

    # Input Fields
    history: dspy.History = dspy.InputField()
    user_prompt: str = dspy.InputField()
    tool_output: str = dspy.InputField(description="The output of the previous tool.")
    available_tools: str = dspy.InputField(
        description="The available tools and their descriptions."
    )

    # Output Fields
    response: str = dspy.OutputField(
        description="The response to the user's prompt whilst the tool is running. Update the user on the progress of their request (if a tool is picked), or the final response to the user (if no tool is picked)."
    )
    tool: str | None = dspy.OutputField(
        description="The tool that needs to be used. Return None if no tool is needed."
    )
    tool_inputs: Dict[str, Any] | None = dspy.OutputField(
        description="The inputs for the tool. Return an empty dictionary (still include the field) if no inputs are needed. The key is the name of the input, the value is the value of the input.",
    )

class RouterAgent:
    def __init__(self, model: str, tools: List[Callable] = []):
        self.tools: list[Callable] = tools
        self.model = dspy.LM(model)
        self.agent = dspy.ChainOfThought(AgentResponse)
        self.followup_agent = dspy.ChainOfThought(AgentFollowup)
        self.conversation_history = dspy.History(messages=[])

    def add_conversation_history(self, message: str, response: dspy.Prediction):
        self.conversation_history.messages.append({"user_prompt": message, **response})

    def get_tools_and_descriptions(self) -> str:
        return "\n".join(
            [
                f"{tool.__name__}:\nDescription: {tool.__doc__ or ''}\nInputs: { {k: v for k, v in tool.__annotations__.items() if k != 'return'} }"
                for tool in self.tools
            ]
        )

    def get_response(self, user_prompt: str) -> str:
        result = self.agent(
            history=self.conversation_history,
            user_prompt=user_prompt,
            available_tools=self.get_tools_and_descriptions(),
            lm=self.model,
        )
        self.add_conversation_history(message=user_prompt, response=result)

        max_iter = 10
        iter = 0

        while result.tool is not None and result.tool.lower() not in ["null", "none"]:
            iter += 1
            if iter > max_iter:
                break

            tool_function = next(
                (tool for tool in self.tools if tool.__name__ == result.tool), None
            )
            if tool_function is None:
                raise ValueError(f"Tool {result.tool} not found")

            tool_inputs = {k: v for k, v in result.tool_inputs.items() if k != "return"}
            tool_result = tool_function(**tool_inputs)

            result = self.followup_agent(
                history=self.conversation_history,
                tool_output=tool_result,
                available_tools=self.get_tools_and_descriptions(),
                lm=self.model,
            )
            self.add_conversation_history(message=tool_result, response=result)

        return result.response