{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e9751170-8a93-45f3-b0a0-c9edc7b2eafb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv, find_dotenv\n",
    "_ = load_dotenv(find_dotenv())\n",
    "openai_api_key = os.environ[\"OPENAI_API_KEY\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69d4b2ec-2155-4126-a18e-774c79cafeba",
   "metadata": {},
   "source": [
    "## Callbacks\n",
    "Calling functions in the middle of larger processes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "448908ed-5d1a-4e0e-aa39-88e19fa4190b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.callbacks import StdOutCallbackHandler\n",
    "from langchain.chains import LLMChain\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langchain.prompts import PromptTemplate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e7c79bb8-3cf4-4112-9b62-2455b92fb903",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = ChatOpenAI()\n",
    "\n",
    "prompt_template = PromptTemplate(\n",
    "    input_variables=[\"input\"], \n",
    "    template=\"Tell me a joke about {input}\")\n",
    "\n",
    "chain = LLMChain(llm=llm, prompt=prompt_template)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19981a0b-7b53-487e-ae19-2501dfa185bf",
   "metadata": {},
   "source": [
    "**Without callback**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0f62b733-c636-49f1-a843-01739e93ee72",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/juliocolomer/.pyenv/versions/3.11.4/envs/venv020124/lib/python3.11/site-packages/langchain_core/_api/deprecation.py:117: LangChainDeprecationWarning: The function `run` was deprecated in LangChain 0.1.0 and will be removed in 0.2.0. Use invoke instead.\n",
      "  warn_deprecated(\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Why did the bear bring a flashlight into the woods?\\n\\nBecause it wanted to be a \"dare\" bear!'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.run(input=\"bear\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e9f9903-6a19-4933-92b2-c237f26a2852",
   "metadata": {},
   "source": [
    "**With callback**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "915dad73-860e-479b-83b9-87b53a2dc34f",
   "metadata": {},
   "outputs": [],
   "source": [
    "handler = StdOutCallbackHandler()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f16886ef-d2e2-4712-808c-c69a658c70ee",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new LLMChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mTell me a joke about bear\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\"Why don't bears wear shoes?\\n\\nBecause they have bear feet!\""
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.run(input=\"bear\", callbacks=[handler])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68fde176-676a-4fe6-b59c-940ba28b021f",
   "metadata": {},
   "source": [
    "## Customized Callback"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1c7de8f9-b52b-4b80-9996-350181033382",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.callbacks.base import BaseCallbackHandler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "66c8aa3b-07db-4316-a38b-2b8cd5cf686f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyCustomHandler(BaseCallbackHandler):\n",
    "    def on_llm_end(self, response, **kwargs) -> None:\n",
    "        print(f\"REPONSE: \", response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ba22baac-3296-4998-a096-65ea7d000d4e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "REPONSE:  generations=[[ChatGeneration(text=\"Why don't whales ever use computers?\\n\\nBecause they're afraid of the net!\", generation_info={'finish_reason': 'stop', 'logprobs': None}, message=AIMessage(content=\"Why don't whales ever use computers?\\n\\nBecause they're afraid of the net!\"))]] llm_output={'token_usage': {'completion_tokens': 16, 'prompt_tokens': 13, 'total_tokens': 29}, 'model_name': 'gpt-3.5-turbo', 'system_fingerprint': None} run=None\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\"Why don't whales ever use computers?\\n\\nBecause they're afraid of the net!\""
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.run(input=\"whale\", callbacks=[MyCustomHandler()])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04bf49ee-cdca-43b6-9b63-35d5d8cf16c9",
   "metadata": {},
   "source": [
    "## Callback to check OpenAI costs, token usage, etc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "7f67fbae-9285-4eb4-ab42-2d70f2d5275d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.callbacks import get_openai_callback"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "30897e7a-beb0-48ad-865a-33d82635c999",
   "metadata": {},
   "outputs": [],
   "source": [
    "with get_openai_callback() as cb:\n",
    "    chain.run(input=\"elephant\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e5de3850-42a5-4645-81b8-28c0b07afe75",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tokens Used: 28\n",
      "\tPrompt Tokens: 13\n",
      "\tCompletion Tokens: 15\n",
      "Successful Requests: 1\n",
      "Total Cost (USD): $4.95e-05\n"
     ]
    }
   ],
   "source": [
    "print(cb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db832973-cbdb-46fb-b3fd-afe41ba99eed",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}