{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tools\n",
    "\n",
    "Tools are code that can be executed by an agent to perform actions. A tool\n",
    "can be a simple function such as a calculator, or an API call to a third-party service\n",
    "such as stock price lookup or weather forecast.\n",
    "In the context of AI agents, tools are designed to be executed by agents in\n",
    "response to model-generated function calls.\n",
    "\n",
    "AutoGen provides the {py:mod}`autogen_core.tools` module with a suite of built-in\n",
    "tools and utilities for creating and running custom tools."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Built-in Tools\n",
    "\n",
    "One of the built-in tools is the {py:class}`~autogen_ext.tools.code_execution.PythonCodeExecutionTool`,\n",
    "which allows agents to execute Python code snippets.\n",
    "\n",
    "Here is how you create the tool and use it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hello, world!\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from autogen_core import CancellationToken\n",
    "from autogen_ext.code_executors.docker import DockerCommandLineCodeExecutor\n",
    "from autogen_ext.tools.code_execution import PythonCodeExecutionTool\n",
    "\n",
    "# Create the tool.\n",
    "code_executor = DockerCommandLineCodeExecutor()\n",
    "await code_executor.start()\n",
    "code_execution_tool = PythonCodeExecutionTool(code_executor)\n",
    "cancellation_token = CancellationToken()\n",
    "\n",
    "# Use the tool directly without an agent.\n",
    "code = \"print('Hello, world!')\"\n",
    "result = await code_execution_tool.run_json({\"code\": code}, cancellation_token)\n",
    "print(code_execution_tool.return_value_as_string(result))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The {py:class}`~autogen_ext.code_executors.docker.DockerCommandLineCodeExecutor`\n",
    "class is a built-in code executor that runs Python code snippets in a subprocess\n",
    "in the command line environment of a docker container.\n",
    "The {py:class}`~autogen_ext.tools.code_execution.PythonCodeExecutionTool` class wraps the code executor\n",
    "and provides a simple interface to execute Python code snippets.\n",
    "\n",
    "Examples of other built-in tools\n",
    "- {py:class}`~autogen_ext.tools.graphrag.LocalSearchTool` and {py:class}`~autogen_ext.tools.graphrag.GlobalSearchTool` for using [GraphRAG](https://github.com/microsoft/graphrag).\n",
    "- {py:class}`~autogen_ext.tools.mcp.mcp_server_tools` for using [Model Context Protocol (MCP)](https://modelcontextprotocol.io/introduction) servers as tools.\n",
    "- {py:class}`~autogen_ext.tools.http.HttpTool` for making HTTP requests to REST APIs.\n",
    "- {py:class}`~autogen_ext.tools.langchain.LangChainToolAdapter` for using LangChain tools."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Custom Function Tools\n",
    "\n",
    "A tool can also be a simple Python function that performs a specific action.\n",
    "To create a custom function tool, you just need to create a Python function\n",
    "and use the {py:class}`~autogen_core.tools.FunctionTool` class to wrap it.\n",
    "\n",
    "The {py:class}`~autogen_core.tools.FunctionTool` class uses descriptions and type annotations\n",
    "to inform the LLM when and how to use a given function. The description provides context\n",
    "about the function’s purpose and intended use cases, while type annotations inform the LLM about\n",
    "the expected parameters and return type.\n",
    "\n",
    "For example, a simple tool to obtain the stock price of a company might look like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "36.63801673457121\n"
     ]
    }
   ],
   "source": [
    "import random\n",
    "\n",
    "from autogen_core import CancellationToken\n",
    "from autogen_core.tools import FunctionTool\n",
    "from typing_extensions import Annotated\n",
    "\n",
    "\n",
    "async def get_stock_price(ticker: str, date: Annotated[str, \"Date in YYYY/MM/DD\"]) -> float:\n",
    "    # Returns a random stock price for demonstration purposes.\n",
    "    return random.uniform(10, 200)\n",
    "\n",
    "\n",
    "# Create a function tool.\n",
    "stock_price_tool = FunctionTool(get_stock_price, description=\"Get the stock price.\")\n",
    "\n",
    "# Run the tool.\n",
    "cancellation_token = CancellationToken()\n",
    "result = await stock_price_tool.run_json({\"ticker\": \"AAPL\", \"date\": \"2021/01/01\"}, cancellation_token)\n",
    "\n",
    "# Print the result.\n",
    "print(stock_price_tool.return_value_as_string(result))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calling Tools with Model Clients\n",
    "\n",
    "Model clients can generate tool calls when they are provided with a list of tools.\n",
    "\n",
    "Here is an example of how to use the {py:class}`~autogen_core.tools.FunctionTool` class\n",
    "with a {py:class}`~autogen_ext.models.openai.OpenAIChatCompletionClient`.\n",
    "Other model client classes can be used in a similar way. See [Model Clients](./model-clients.ipynb)\n",
    "for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[FunctionCall(id='call_tpJ5J1Xoxi84Sw4v0scH0qBM', arguments='{\"ticker\":\"AAPL\",\"date\":\"2021/01/01\"}', name='get_stock_price')]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import json\n",
    "\n",
    "from autogen_core.models import AssistantMessage, FunctionExecutionResult, FunctionExecutionResultMessage, UserMessage\n",
    "from autogen_ext.models.openai import OpenAIChatCompletionClient\n",
    "\n",
    "# Create the OpenAI chat completion client. Using OPENAI_API_KEY from environment variable.\n",
    "client = OpenAIChatCompletionClient(model=\"gpt-4o-mini\")\n",
    "\n",
    "# Create a user message.\n",
    "user_message = UserMessage(content=\"What is the stock price of AAPL on 2021/01/01?\", source=\"user\")\n",
    "\n",
    "# Run the chat completion with the stock_price_tool defined above.\n",
    "cancellation_token = CancellationToken()\n",
    "create_result = await client.create(\n",
    "    messages=[user_message], tools=[stock_price_tool], cancellation_token=cancellation_token\n",
    ")\n",
    "create_result.content"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The result is a list of {py:class}`~autogen_core.FunctionCall` objects, which can be\n",
    "used to run the corresponding tools."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'32.381250753393104'"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "arguments = json.loads(create_result.content[0].arguments)  # type: ignore\n",
    "tool_result = await stock_price_tool.run_json(arguments, cancellation_token)\n",
    "tool_result_str = stock_price_tool.return_value_as_string(tool_result)\n",
    "tool_result_str"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now you can make another model client call to have the model generate a reflection\n",
    "on the result of the tool execution.\n",
    "\n",
    "The result of the tool call is wrapped in a {py:class}`~autogen_core.models.FunctionExecutionResult`\n",
    "object, which contains the result of the tool execution and the ID of the tool that was called.\n",
    "The model client can use this information to generate a reflection on the result of the tool execution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The stock price of AAPL (Apple Inc.) on January 1, 2021, was approximately $32.38.\n"
     ]
    }
   ],
   "source": [
    "# Create a function execution result\n",
    "exec_result = FunctionExecutionResult(\n",
    "    call_id=create_result.content[0].id,  # type: ignore\n",
    "    content=tool_result_str,\n",
    "    is_error=False,\n",
    "    name=stock_price_tool.name,\n",
    ")\n",
    "\n",
    "# Make another chat completion with the history and function execution result message.\n",
    "messages = [\n",
    "    user_message,\n",
    "    AssistantMessage(content=create_result.content, source=\"assistant\"),  # assistant message with tool call\n",
    "    FunctionExecutionResultMessage(content=[exec_result]),  # function execution result message\n",
    "]\n",
    "create_result = await client.create(messages=messages, cancellation_token=cancellation_token)  # type: ignore\n",
    "print(create_result.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tool-Equipped Agent\n",
    "\n",
    "Putting the model client and the tools together, you can create a tool-equipped agent\n",
    "that can use tools to perform actions, and reflect on the results of those actions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import asyncio\n",
    "import json\n",
    "from dataclasses import dataclass\n",
    "from typing import List\n",
    "\n",
    "from autogen_core import (\n",
    "    AgentId,\n",
    "    FunctionCall,\n",
    "    MessageContext,\n",
    "    RoutedAgent,\n",
    "    SingleThreadedAgentRuntime,\n",
    "    message_handler,\n",
    ")\n",
    "from autogen_core.models import (\n",
    "    ChatCompletionClient,\n",
    "    LLMMessage,\n",
    "    SystemMessage,\n",
    "    UserMessage,\n",
    ")\n",
    "from autogen_core.tools import FunctionTool, Tool\n",
    "from autogen_ext.models.openai import OpenAIChatCompletionClient\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class Message:\n",
    "    content: str\n",
    "\n",
    "\n",
    "class ToolUseAgent(RoutedAgent):\n",
    "    def __init__(self, model_client: ChatCompletionClient, tool_schema: List[Tool]) -> None:\n",
    "        super().__init__(\"An agent with tools\")\n",
    "        self._system_messages: List[LLMMessage] = [SystemMessage(content=\"You are a helpful AI assistant.\")]\n",
    "        self._model_client = model_client\n",
    "        self._tools = tool_schema\n",
    "\n",
    "    @message_handler\n",
    "    async def handle_user_message(self, message: Message, ctx: MessageContext) -> Message:\n",
    "        # Create a session of messages.\n",
    "        session: List[LLMMessage] = self._system_messages + [UserMessage(content=message.content, source=\"user\")]\n",
    "\n",
    "        # Run the chat completion with the tools.\n",
    "        create_result = await self._model_client.create(\n",
    "            messages=session,\n",
    "            tools=self._tools,\n",
    "            cancellation_token=ctx.cancellation_token,\n",
    "        )\n",
    "\n",
    "        # If there are no tool calls, return the result.\n",
    "        if isinstance(create_result.content, str):\n",
    "            return Message(content=create_result.content)\n",
    "        assert isinstance(create_result.content, list) and all(\n",
    "            isinstance(call, FunctionCall) for call in create_result.content\n",
    "        )\n",
    "\n",
    "        # Add the first model create result to the session.\n",
    "        session.append(AssistantMessage(content=create_result.content, source=\"assistant\"))\n",
    "\n",
    "        # Execute the tool calls.\n",
    "        results = await asyncio.gather(\n",
    "            *[self._execute_tool_call(call, ctx.cancellation_token) for call in create_result.content]\n",
    "        )\n",
    "\n",
    "        # Add the function execution results to the session.\n",
    "        session.append(FunctionExecutionResultMessage(content=results))\n",
    "\n",
    "        # Run the chat completion again to reflect on the history and function execution results.\n",
    "        create_result = await self._model_client.create(\n",
    "            messages=session,\n",
    "            cancellation_token=ctx.cancellation_token,\n",
    "        )\n",
    "        assert isinstance(create_result.content, str)\n",
    "\n",
    "        # Return the result as a message.\n",
    "        return Message(content=create_result.content)\n",
    "\n",
    "    async def _execute_tool_call(\n",
    "        self, call: FunctionCall, cancellation_token: CancellationToken\n",
    "    ) -> FunctionExecutionResult:\n",
    "        # Find the tool by name.\n",
    "        tool = next((tool for tool in self._tools if tool.name == call.name), None)\n",
    "        assert tool is not None\n",
    "\n",
    "        # Run the tool and capture the result.\n",
    "        try:\n",
    "            arguments = json.loads(call.arguments)\n",
    "            result = await tool.run_json(arguments, cancellation_token)\n",
    "            return FunctionExecutionResult(\n",
    "                call_id=call.id, content=tool.return_value_as_string(result), is_error=False, name=tool.name\n",
    "            )\n",
    "        except Exception as e:\n",
    "            return FunctionExecutionResult(call_id=call.id, content=str(e), is_error=True, name=tool.name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When handling a user message, the `ToolUseAgent` class first use the model client\n",
    "to generate a list of function calls to the tools, and then run the tools\n",
    "and generate a reflection on the results of the tool execution.\n",
    "The reflection is then returned to the user as the agent's response.\n",
    "\n",
    "To run the agent, let's create a runtime and register the agent with the runtime."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AgentType(type='tool_use_agent')"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create a runtime.\n",
    "runtime = SingleThreadedAgentRuntime()\n",
    "# Create the tools.\n",
    "tools: List[Tool] = [FunctionTool(get_stock_price, description=\"Get the stock price.\")]\n",
    "# Register the agents.\n",
    "await ToolUseAgent.register(\n",
    "    runtime,\n",
    "    \"tool_use_agent\",\n",
    "    lambda: ToolUseAgent(\n",
    "        OpenAIChatCompletionClient(model=\"gpt-4o-mini\"),\n",
    "        tools,\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This example uses the {py:class}`~autogen_ext.models.openai.OpenAIChatCompletionClient`,\n",
    "for Azure OpenAI and other clients, see [Model Clients](./model-clients.ipynb).\n",
    "Let's test the agent with a question about stock price."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The stock price of NVIDIA (NVDA) on June 1, 2024, was approximately $140.05.\n"
     ]
    }
   ],
   "source": [
    "# Start processing messages.\n",
    "runtime.start()\n",
    "# Send a direct message to the tool agent.\n",
    "tool_use_agent = AgentId(\"tool_use_agent\", \"default\")\n",
    "response = await runtime.send_message(Message(\"What is the stock price of NVDA on 2024/06/01?\"), tool_use_agent)\n",
    "print(response.content)\n",
    "# Stop processing messages.\n",
    "await runtime.stop()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}