{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip available: \u001b[0m\u001b[31;49m22.3.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.0\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install --quiet -U langchain_openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import getpass\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = getpass.getpass(\"OpenAI API Key:\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langchain_core.messages import BaseMessage, HumanMessage\n",
    "from langgraph.graph import END, MessageGraph\n",
    "\n",
    "model = ChatOpenAI(temperature=0)\n",
    "\n",
    "graph = MessageGraph()\n",
    "\n",
    "def invoke_model(state: List[BaseMessage]):\n",
    "    return model.invoke(state)\n",
    "\n",
    "graph.add_node(\"oracle\", invoke_model)\n",
    "graph.add_edge(\"oracle\", END)\n",
    "\n",
    "graph.set_entry_point(\"oracle\")\n",
    "\n",
    "runnable = graph.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[HumanMessage(content='What is 1 + 1?'), AIMessage(content='1 + 1 equals 2.')]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "runnable.invoke(HumanMessage(\"What is 1 + 1?\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from langchain_core.messages import ToolMessage\n",
    "from langchain_core.tools import tool\n",
    "from langchain_core.utils.function_calling import convert_to_openai_tool\n",
    "\n",
    "@tool\n",
    "def multiply(first_number: int, second_number: int):\n",
    "    \"\"\"Multiplies two numbers together.\"\"\"\n",
    "    return first_number * second_number\n",
    "\n",
    "model = ChatOpenAI(temperature=0)\n",
    "model_with_tools = model.bind(tools=[convert_to_openai_tool(multiply)])\n",
    "\n",
    "graph = MessageGraph()\n",
    "\n",
    "def invoke_model(state: List[BaseMessage]):\n",
    "    return model_with_tools.invoke(state)\n",
    "\n",
    "graph.add_node(\"oracle\", invoke_model)\n",
    "\n",
    "def invoke_tool(state: List[BaseMessage]):\n",
    "    tool_calls = state[-1].additional_kwargs.get(\"tool_calls\", [])\n",
    "    multiply_call = None\n",
    "\n",
    "    for tool_call in tool_calls:\n",
    "        if tool_call.get(\"function\").get(\"name\") == \"multiply\":\n",
    "            multiply_call = tool_call\n",
    "\n",
    "    if multiply_call is None:\n",
    "        raise Exception(\"No adder input found.\")\n",
    "\n",
    "    res = multiply.invoke(\n",
    "        json.loads(multiply_call.get(\"function\").get(\"arguments\"))\n",
    "    )\n",
    "\n",
    "    return ToolMessage(\n",
    "        tool_call_id=multiply_call.get(\"id\"),\n",
    "        content=res\n",
    "    )\n",
    "\n",
    "graph.add_node(\"multiply\", invoke_tool)\n",
    "\n",
    "graph.add_edge(\"multiply\", END)\n",
    "\n",
    "graph.set_entry_point(\"oracle\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def router(state: List[BaseMessage]):\n",
    "    tool_calls = state[-1].additional_kwargs.get(\"tool_calls\", [])\n",
    "    if len(tool_calls):\n",
    "        return \"multiply\"\n",
    "    else:\n",
    "        return \"end\"\n",
    "\n",
    "graph.add_conditional_edges(\"oracle\", router, {\n",
    "    \"multiply\": \"multiply\",\n",
    "    \"end\": END,\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[HumanMessage(content='What is 123 * 456?'),\n",
       " AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_cKnzV4f7sEjnuetuyW8glC7u', 'function': {'arguments': '{\"first_number\":123,\"second_number\":456}', 'name': 'multiply'}, 'type': 'function'}]}),\n",
       " ToolMessage(content='56088', tool_call_id='call_cKnzV4f7sEjnuetuyW8glC7u')]"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "runnable = graph.compile()\n",
    "\n",
    "runnable.invoke(HumanMessage(\"What is 123 * 456?\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[HumanMessage(content='What is your name?'),\n",
       " AIMessage(content='My name is Assistant. How can I assist you today?')]"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "runnable.invoke(HumanMessage(\"What is your name?\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}