{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Agents with Structured Outputs\n",
    "\n",
    "When you run your agent or multi-agent framework, you might want it to output the result in a specific format. In this notebook, we will see a simple example of how to apply this to a FunctionAgent!🦙🚀"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's first install the needed dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install llama-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from getpass import getpass\n",
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = getpass()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now define our structured output format\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseModel, Field\n",
    "\n",
    "\n",
    "class MathResult(BaseModel):\n",
    "    operation: str = Field(description=\"The operation that has been performed\")\n",
    "    result: int = Field(description=\"Result of the operation\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And a very simple calculator agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.core.agent.workflow import FunctionAgent\n",
    "\n",
    "llm = OpenAI(model=\"gpt-4.1\")\n",
    "\n",
    "\n",
    "def add(x: int, y: int):\n",
    "    \"\"\"Add two numbers\"\"\"\n",
    "    return x + y\n",
    "\n",
    "\n",
    "def multiply(x: int, y: int):\n",
    "    \"\"\"Multiply two numbers\"\"\"\n",
    "    return x * y\n",
    "\n",
    "\n",
    "agent = FunctionAgent(\n",
    "    llm=llm,\n",
    "    output_cls=MathResult,\n",
    "    tools=[add, multiply],\n",
    "    system_prompt=\"You are a calculator agent that can add or multiply two numbers by calling tools\",\n",
    "    name=\"calculator\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now run the agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = await agent.run(\"What is the result of 10 multiplied by 4?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can get the structured output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print the structured output as a plain dictionary\n",
    "print(response.structured_response)\n",
    "# print the structured output as a Pydantic model\n",
    "print(response.get_pydantic_model(MathResult))"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}