{
 "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": [
      "143.83831971965762\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",
    "In AutoGen, every tool is a subclass of {py:class}`~autogen_core.tools.BaseTool`,\n",
    "which automatically generates the JSON schema for the tool.\n",
    "For example, to get the JSON schema for the `stock_price_tool`, we can use the\n",
    "{py:attr}`~autogen_core.tools.BaseTool.schema` property."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'name': 'get_stock_price',\n",
       " 'description': 'Get the stock price.',\n",
       " 'parameters': {'type': 'object',\n",
       "  'properties': {'ticker': {'description': 'ticker',\n",
       "    'title': 'Ticker',\n",
       "    'type': 'string'},\n",
       "   'date': {'description': 'Date in YYYY/MM/DD',\n",
       "    'title': 'Date',\n",
       "    'type': 'string'}},\n",
       "  'required': ['ticker', 'date'],\n",
       "  'additionalProperties': False},\n",
       " 'strict': False}"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "stock_price_tool.schema"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Model clients use the JSON schema of the tools to generate tool calls.\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",
    "model_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 model_client.create(\n",
    "    messages=[user_message], tools=[stock_price_tool], cancellation_token=cancellation_token\n",
    ")\n",
    "create_result.content"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What is actually going on under the hood of the call to the\n",
    "{py:class}`~autogen_ext.models.openai.BaseOpenAIChatCompletionClient.create`\n",
    "method? The model client takes the list of tools and generates a JSON schema\n",
    "for the parameters of each tool. Then, it generates a request to the model\n",
    "API with the tool's JSON schema and the other messages to obtain a result.\n",
    "\n",
    "Many models, such as OpenAI's GPT-4o and Llama-3.2, are trained to produce\n",
    "tool calls in the form of structured JSON strings that conform to the\n",
    "JSON schema of the tool. AutoGen's model clients then parse the model's response\n",
    "and extract the tool call from the JSON string.\n",
    "\n",
    "The result is a list of {py:class}`~autogen_core.FunctionCall` objects, which can be\n",
    "used to run the corresponding tools.\n",
    "\n",
    "We use `json.loads` to parse the JSON string in the {py:class}`~autogen_core.FunctionCall.arguments`\n",
    "field into a Python dictionary. The {py:meth}`~autogen_core.tools.BaseTool.run_json`\n",
    "method takes the dictionary and runs the tool with the provided arguments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'32.381250753393104'"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "assert isinstance(create_result.content, list)\n",
    "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 model_client.create(messages=messages, cancellation_token=cancellation_token)  # type: ignore\n",
    "print(create_result.content)\n",
    "await model_client.close()"
   ]
  },
  {
   "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.\n",
    "\n",
    "```{note}\n",
    "The Core API is designed to be minimal and you need to build your own agent logic around model clients and tools.\n",
    "For \"pre-built\" agents that can use tools, please refer to the [AgentChat API](../../agentchat-user-guide/index.md).\n",
    "```"
   ]
  },
  {
   "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 the model client.\n",
    "model_client = OpenAIChatCompletionClient(model=\"gpt-4o-mini\")\n",
    "# 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",
    "        model_client=model_client,\n",
    "        tool_schema=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()\n",
    "await model_client.close()"
   ]
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}