2. LangChain.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"width: 100%; overflow: hidden;\">\n",
    "    <div style=\"width: 150px; float: left;\"> <img src=\"data/D4Sci_logo_ball.png\" alt=\"Data For Science, Inc\" align=\"left\" border=\"0\"> </div>\n",
    "    <div style=\"float: left; margin-left: 10px;\"> <h1>LangChain for Generative AI</h1>\n",
    "<h1>LangChain</h1>\n",
    "        <p>Bruno Gonçalves<br/>\n",
    "        <a href=\"http://www.data4sci.com/\">www.data4sci.com</a><br/>\n",
    "            @bgoncalves, @data4sci</p></div>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "from pprint import pprint\n",
    "from operator import itemgetter\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt \n",
    "\n",
    "import torch\n",
    "\n",
    "import openai\n",
    "from openai import OpenAI\n",
    "\n",
    "import transformers\n",
    "from transformers import pipeline\n",
    "from transformers import set_seed\n",
    "set_seed(42) # Set the seed to get reproducible results\n",
    "\n",
    "\n",
    "import langchain\n",
    "from langchain.chains import create_sql_query_chain\n",
    "from langchain.tools import DuckDuckGoSearchRun\n",
    "\n",
    "import langchain_openai\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "import langchain_anthropic\n",
    "from langchain_anthropic import ChatAnthropic\n",
    "\n",
    "import langchain_core\n",
    "from langchain_core.messages import HumanMessage, SystemMessage\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder, PromptTemplate\n",
    "from langchain_core.chat_history import BaseChatMessageHistory\n",
    "from langchain_core.runnables.history import RunnableWithMessageHistory\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "\n",
    "import langchain_community\n",
    "from langchain_community.chat_message_histories import ChatMessageHistory\n",
    "from langchain_community.utilities import SQLDatabase\n",
    "from langchain_community.tools.sql_database.tool import QuerySQLDataBaseTool\n",
    "\n",
    "import watermark\n",
    "\n",
    "%load_ext watermark\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We start by print out the versions of the libraries we're using for future reference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Python implementation: CPython\n",
      "Python version       : 3.11.7\n",
      "IPython version      : 8.12.3\n",
      "\n",
      "Compiler    : Clang 14.0.6 \n",
      "OS          : Darwin\n",
      "Release     : 23.6.0\n",
      "Machine     : arm64\n",
      "Processor   : arm\n",
      "CPU cores   : 16\n",
      "Architecture: 64bit\n",
      "\n",
      "Git hash: 0b932bf32c4e9b1e4ff9126fc7e8e7ac7b4205ff\n",
      "\n",
      "langchain_core     : 0.2.3\n",
      "langchain_openai   : 0.1.8\n",
      "transformers       : 4.41.1\n",
      "langchain          : 0.2.2\n",
      "langchain_anthropic: 0.1.15\n",
      "pandas             : 2.2.3\n",
      "langchain_community: 0.2.1\n",
      "torch              : 2.3.0\n",
      "watermark          : 2.4.3\n",
      "matplotlib         : 3.8.0\n",
      "openai             : 1.30.5\n",
      "numpy              : 1.26.4\n",
      "\n"
     ]
    }
   ],
   "source": [
    "%watermark -n -v -m -g -iv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load default figure style"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.style.use('./d4sci.mplstyle')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# OpenAI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first step is generate API key on the OpenAI website and store it as the \"OPENAI_API_KEY\" variable in your local environment. Without it we won't be able to do anything. You can find your API key in your using settings: https://help.openai.com/en/articles/4936850-where-do-i-find-my-secret-api-key. Then we are ready to instantiate the client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = OpenAI()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = client.chat.completions.create(\n",
    "  model=\"gpt-4o\",\n",
    "  messages=[\n",
    "        {\n",
    "            \"role\": \"user\", \n",
    "            \"content\": \"What was Superman's weakness?\"\n",
    "        },\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ChatCompletion(id='chatcmpl-AVlBKZ3AK1ebP607ObfRLejPGembN', choices=[Choice(finish_reason='stop', index=0, logprobs=None, message=ChatCompletionMessage(content=\"Superman's most well-known weakness is kryptonite, a mineral from his home planet of Krypton. Exposure to kryptonite weakens Superman and can potentially be fatal to him. Different forms of kryptonite have varied effects; for example, green kryptonite weakens him, while red kryptonite can cause unpredictable, often bizarre changes in his behavior and powers. Besides kryptonite, Superman is also vulnerable to magic, which can affect him just as it would any other being. Additionally, he's not immune to the effects of a red sun, which can drain him of his powers.\", role='assistant', function_call=None, tool_calls=None, refusal=None))], created=1732132898, model='gpt-4o-2024-08-06', object='chat.completion', system_fingerprint='fp_7f6be3efb0', usage=CompletionUsage(completion_tokens=115, prompt_tokens=13, total_tokens=128, prompt_tokens_details={'cached_tokens': 0, 'audio_tokens': 0}, completion_tokens_details={'reasoning_tokens': 0, 'audio_tokens': 0, 'accepted_prediction_tokens': 0, 'rejected_prediction_tokens': 0}))\n"
     ]
    }
   ],
   "source": [
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Superman's most well-known weakness is kryptonite, a mineral from his home planet of Krypton. Exposure to kryptonite weakens Superman and can potentially be fatal to him. Different forms of kryptonite have varied effects; for example, green kryptonite weakens him, while red kryptonite can cause unpredictable, often bizarre changes in his behavior and powers. Besides kryptonite, Superman is also vulnerable to magic, which can affect him just as it would any other being. Additionally, he's not immune to the effects of a red sun, which can drain him of his powers.\n"
     ]
    }
   ],
   "source": [
    "print(response.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LangChain"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We instantiate the LangChain interface for OpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ChatOpenAI(model=\"gpt-4o\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "content=\"Superman's primary weakness is Kryptonite, a mineral from his home planet of Krypton. When exposed to Kryptonite, Superman's powers are significantly weakened, and prolonged exposure can be fatal to him. There are different types of Kryptonite, with green Kryptonite being the most common and harmful to Superman. Other versions, like red, blue, and gold Kryptonite, have varying effects, often altering his powers or behavior in unpredictable ways.\\n\\nIn addition to Kryptonite, Superman is also vulnerable to magic. His powers, which are based on Earth's yellow sun, do not protect him from magical attacks or spells. Furthermore, Superman's moral code and compassion for others are often considered weaknesses, as his enemies can exploit these traits to manipulate or distract him.\" response_metadata={'token_usage': {'completion_tokens': 150, 'prompt_tokens': 13, 'total_tokens': 163, 'prompt_tokens_details': {'cached_tokens': 0, 'audio_tokens': 0}, 'completion_tokens_details': {'reasoning_tokens': 0, 'audio_tokens': 0, 'accepted_prediction_tokens': 0, 'rejected_prediction_tokens': 0}}, 'model_name': 'gpt-4o', 'system_fingerprint': 'fp_831e067d82', 'finish_reason': 'stop', 'logprobs': None} id='run-3ef83ee5-4b56-4bb1-abae-2cbfdb3d5539-0' usage_metadata={'input_tokens': 13, 'output_tokens': 150, 'total_tokens': 163}\n"
     ]
    }
   ],
   "source": [
    "messages = [\n",
    "    SystemMessage(content=\"What was Superman's weakness?\"),\n",
    "]\n",
    "\n",
    "output = model.invoke(messages)\n",
    "print(output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'completion_tokens': 150,\n",
       " 'prompt_tokens': 13,\n",
       " 'total_tokens': 163,\n",
       " 'prompt_tokens_details': {'cached_tokens': 0, 'audio_tokens': 0},\n",
       " 'completion_tokens_details': {'reasoning_tokens': 0,\n",
       "  'audio_tokens': 0,\n",
       "  'accepted_prediction_tokens': 0,\n",
       "  'rejected_prediction_tokens': 0}}"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output.response_metadata[\"token_usage\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "parser = StrOutputParser()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = model.invoke(messages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Superman\\'s most well-known weakness is Kryptonite, a mineral from his home planet of Krypton. When exposed to it, Superman loses his powers and becomes vulnerable, and prolonged exposure can be fatal. Different forms of Kryptonite have different effects, but green Kryptonite is the most common and harmful to him.\\n\\nAdditionally, Superman is vulnerable to magic and can be harmed by magical spells and artifacts. While not a \"weakness\" in the traditional sense, his powers don\\'t protect him from magical attacks as they do against conventional threats.\\n\\nLastly, Superman\\'s compassion and strong moral code can sometimes be seen as a weakness, as his enemies may exploit his desire to protect others to manipulate or trap him.'"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "parser.invoke(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us create our first chain. Stages of the chain are conencted with the pipe '|' character"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "chain = model | parser"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now whenver we call __invoke()__ on the chain, it automatically runs all the steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"Superman's most well-known weakness is Kryptonite, a radioactive substance from his home planet, Krypton. Exposure to Kryptonite weakens Superman, drains his powers, and can even be lethal with prolonged exposure. There are various forms of Kryptonite, each having different effects; for example, green Kryptonite is the most common and harmful, while red Kryptonite can cause unpredictable changes in his behavior or abilities. \\n\\nIn addition to Kryptonite, Superman has other vulnerabilities, such as magic, which can affect him just as it would any other being, bypassing his usual invulnerability. Also, while he is incredibly strong and durable, Superman can be overpowered by beings of equal or greater strength, such as other Kryptonians or powerful cosmic entities.\""
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.invoke(messages)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also create templates for our prompts, following conventions similar to the Jinja templating system"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_template = \"Translate the following into {language}:\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we can combine multiple messages into a single template"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_template = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "     (\"system\", system_template), \n",
    "     (\"user\", \"{text}\")]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To instantiate the prompt, we must provide the correct fields"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ChatPromptValue(messages=[SystemMessage(content='Translate the following into italian:'), HumanMessage(content='Be the change that you wish to see in the world.')])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result = prompt_template.invoke({\n",
    "    \"language\": \"italian\", \n",
    "    \"text\": \"Be the change that you wish to see in the world.\"\n",
    "})\n",
    "\n",
    "result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The full interaction is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[SystemMessage(content='Translate the following into italian:'),\n",
       " HumanMessage(content='Be the change that you wish to see in the world.')]"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result.to_messages()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "chain = prompt_template | model | parser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Sei addestrato sui dati fino a ottobre 2023.'"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.invoke({\n",
    "    \"language\": \"italian\", \n",
    "    \"text\": \"Be the change that you wish to see in the world.\"\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Anthropic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_a = ChatAnthropic(model=\"claude-3-opus-20240229\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "chain_a = prompt_template | model_a | parser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ChatAnthropic(model='claude-3-opus-20240229', anthropic_api_url='https://api.anthropic.com', anthropic_api_key=SecretStr('**********'), _client=<anthropic.Anthropic object at 0x176cbabd0>, _async_client=<anthropic.AsyncAnthropic object at 0x176cb8d10>)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Sii il cambiamento che desideri vedere nel mondo.'"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain_a.invoke({\n",
    "    \"language\": \"italian\", \n",
    "    \"text\": \"Be the change that you wish to see in the world.\"\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Message History"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "store = {}\n",
    "\n",
    "def get_session_history(session_id: str) -> BaseChatMessageHistory:\n",
    "    if session_id not in store:\n",
    "        store[session_id] = ChatMessageHistory()\n",
    "    return store[session_id]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "with_message_history = RunnableWithMessageHistory(model, get_session_history)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = {\"configurable\": {\"session_id\": \"abc2\"}}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Hi Bob! How can I assist you today?'"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response = with_message_history.invoke(\n",
    "    [HumanMessage(content=\"Hi! I'm Bob\")],\n",
    "    config=config,\n",
    ")\n",
    "\n",
    "response.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'You mentioned that your name is Bob.'"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response = with_message_history.invoke(\n",
    "    [HumanMessage(content=\"What's my name?\")],\n",
    "    config=config,\n",
    ")\n",
    "\n",
    "response.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"I'm sorry, but I can't determine your name based on the current information.\""
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "config = {\"configurable\": {\"session_id\": \"abc3\"}}\n",
    "\n",
    "response = with_message_history.invoke(\n",
    "    [HumanMessage(content=\"What's my name?\")],\n",
    "    config=config,\n",
    ")\n",
    "\n",
    "response.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'You said your name is Bob.'"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "config = {\"configurable\": {\"session_id\": \"abc2\"}}\n",
    "\n",
    "response = with_message_history.invoke(\n",
    "    [HumanMessage(content=\"What's my name?\")],\n",
    "    config=config,\n",
    ")\n",
    "\n",
    "response.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            \"You are a helpful assistant. Answer all questions to the best of your ability.\",\n",
    "        ),\n",
    "        MessagesPlaceholder(variable_name=\"messages\"),\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "chain = prompt | model | parser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Hello, Bob! How can I assist you today?'"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response = chain.invoke({\"messages\": [HumanMessage(content=\"hi! I'm bob\")]})\n",
    "\n",
    "response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "with_message_history = RunnableWithMessageHistory(chain, get_session_history)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = {\"configurable\": {\"session_id\": \"abc5\"}}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Hello Jim! How can I assist you today?'"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response = with_message_history.invoke(\n",
    "    [HumanMessage(content=\"Hi! I'm Jim\")],\n",
    "    config=config,\n",
    ")\n",
    "\n",
    "response"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Database Integration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "db = SQLDatabase.from_uri(\"sqlite:///data/Northwind_small.sqlite\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sqlite\n"
     ]
    }
   ],
   "source": [
    "print(db.dialect)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Category', 'Customer', 'CustomerCustomerDemo', 'CustomerDemographic', 'Employee', 'EmployeeTerritory', 'Order', 'OrderDetail', 'Product', 'Region', 'Shipper', 'Supplier', 'Territory']\n"
     ]
    }
   ],
   "source": [
    "print(db.get_usable_table_names())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = ChatOpenAI(model=\"gpt-3.5-turbo\", temperature=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "write_query = create_sql_query_chain(llm, db)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'SELECT COUNT(\"Id\") AS TotalCustomers\\nFROM Customer;'"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response = write_query.invoke({\"question\": \"How many customers are there\"}) \n",
    "response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'[(91,)]'"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "db.run(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "execute_query = QuerySQLDataBaseTool(db=db)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "sql_chain = write_query | execute_query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'[(9,)]'"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sql_chain.invoke({\"question\": \"How many employees are there\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'There are 9 employees.'"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "answer_prompt = PromptTemplate.from_template(\n",
    "    \"\"\"Given the following user question, corresponding SQL query, and SQL result, answer the user question.\n",
    "\n",
    "Question: {question}\n",
    "SQL Query: {query}\n",
    "SQL Result: {result}\n",
    "Answer: \"\"\"\n",
    ")\n",
    "\n",
    "answer = answer_prompt | llm | StrOutputParser()\n",
    "chain = (\n",
    "    RunnablePassthrough.assign(query=write_query).assign(\n",
    "        result=itemgetter(\"query\") | execute_query\n",
    "    )\n",
    "    | answer\n",
    ")\n",
    "\n",
    "chain.invoke({\"question\": \"How many employees are there\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"Get Ready for These Upcoming Eclipses! More Eclipses Solar Eclipses Date Solar Eclipse Type Geographic Region of Visibility Oct. 2, 2024 Annular An annular solar eclipse will be visible in South America, and a partial eclipse will be visible in South America, Antarctica, Pacific Ocean, Atlantic Ocean, North America March 29, 2025 Partial Europe, Asia, […] Explore a map of the next 15 total solar eclipses. In case you miss this year's solar eclipse, there are 14 more in the next 20 years. This map of eclipse paths from 2024 to 2044 reveals that ... It will be 20 years before there's a chance to witness a total solar eclipse in the United States again. According to NASA, after the total solar eclipse on April 8, 2024, the next total solar ... The sun is seen in full eclipse over Grand Teton National Park on August 21, 2017, outside Jackson, Wyoming. After 2024, the next total solar eclipse visible in the U.S. will be in 2044. The total length of the 2024 eclipse path is 9,190 miles (14,790 km). The magnitude of this eclipse is 1.0565, which means the Moon's diameter is 5.65 percent larger than the Sun's. Only when ...\""
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "search = DuckDuckGoSearchRun()\n",
    "search.run(\"When will the next solar eclipse be?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<center>\n",
    "     <img src=\"data/D4Sci_logo_full.png\" alt=\"Data For Science, Inc\" align=\"center\" border=\"0\" width=300px> \n",
    "</center>"
   ]
  }
 ],
 "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.7"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": true,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}