{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Optional, Type\n",
    "from langchain_core.pydantic_v1 import BaseModel, Field\n",
    "from langchain_core.tools import BaseTool\n",
    "\n",
    "class MultiplierInput(BaseModel):\n",
    "    a: int = Field(description=\"first number\")\n",
    "    b: int = Field(description=\"second number\")\n",
    "\n",
    "class CustomMultiplierTool(BaseTool):\n",
    "    name = \"Calculator\"\n",
    "    description = \"useful for when you need to answer questions about math\"\n",
    "    args_schema: Type[BaseModel] = MultiplierInput\n",
    "\n",
    "    def _run(\n",
    "        self, a: int, b: int\n",
    "    ) -> str:\n",
    "        \"\"\"Use the tool.\"\"\"\n",
    "        return a * b\n",
    "\n",
    "    async def _arun(\n",
    "        self,\n",
    "        a: int,\n",
    "        b: int,\n",
    "    ) -> str:\n",
    "        \"\"\"Use the tool asynchronously.\"\"\"\n",
    "        return self._run(a,b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import AgentExecutor, create_openai_tools_agent,Tool\n",
    "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langchain.memory import ConversationBufferMemory\n",
    "from langchain.prompts import MessagesPlaceholder\n",
    "from langchain.globals import set_verbose\n",
    "import dotenv\n",
    "\n",
    "dotenv.load_dotenv()\n",
    "\n",
    "def create_agent(model_name):\n",
    "\n",
    "    memory = ConversationBufferMemory(memory_key=\"chat_history\", return_messages=True)\n",
    "    \n",
    "    tools = [CustomMultiplierTool()]\n",
    "    \n",
    "    llm = ChatOpenAI(model=model_name, temperature=0)\n",
    "\n",
    "    system_message = \"\"\"You are a general AI assistant.\n",
    "    You have access to several tool. Don't answer the question if you are not getting the answer from a tool.\n",
    "    Don't put safety and cultural warnings. Only warn about security. \n",
    "    \"\"\"\n",
    "\n",
    "    prompt = ChatPromptTemplate.from_messages(\n",
    "        [\n",
    "            (\"system\", system_message),\n",
    "            MessagesPlaceholder(\"chat_history\", optional=True),\n",
    "            (\"human\", \"{input}\"),\n",
    "            MessagesPlaceholder(\"agent_scratchpad\"),\n",
    "        ]\n",
    "    )\n",
    "    agent = create_openai_tools_agent(llm, tools, prompt)\n",
    "    agent_exe = AgentExecutor(agent=agent, tools=tools,memory=memory)\n",
    "    return agent_exe\n",
    "\n",
    "async def run_agent(agent,user_query):\n",
    "    #print(agent.memory.chat_memory.messages[-2:] if len(agent.memory.chat_memory.messages) > 1 else \"\")\n",
    "    #set_verbose(True)\n",
    "    # print(agent.memory.chat_memory)\n",
    "    print('********************')\n",
    "    # print()\n",
    "    return await agent.ainvoke(input={\"input\":user_query},verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "agent_exe = create_agent(\"gpt-3.5-turbo\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "********************\n",
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `Calculator` with `{'a': 2, 'b': 3}`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3m6\u001b[0m\u001b[32;1m\u001b[1;3m2 multiplied by 3 is equal to 6.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'input': '2 multiply by 3',\n",
       " 'chat_history': [HumanMessage(content='2 multiply by 3'),\n",
       "  AIMessage(content='2 multiplied by 3 is equal to 6.')],\n",
       " 'output': '2 multiplied by 3 is equal to 6.'}"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "await run_agent(agent_exe,\"2 multiply by 3\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "stevens-chat-py311-v1",
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}