playground.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from decouple import config\n",
    "from langchain_community.utilities import SQLDatabase\n",
    "from langchain_experimental.sql import SQLDatabaseChain\n",
    "from sqlalchemy import create_engine, MetaData\n",
    "from sqlalchemy.orm import sessionmaker\n",
    "from operator import itemgetter\n",
    "\n",
    "from langchain.chains import create_sql_query_chain, LLMChain\n",
    "from langchain_openai import ChatOpenAI, OpenAI\n",
    "from langchain_mistralai import ChatMistralAI\n",
    "from langchain_core.prompts import PromptTemplate, FewShotPromptTemplate, ChatPromptTemplate, SystemMessagePromptTemplate\n",
    "from langchain_community.tools.sql_database.tool import QuerySQLDataBaseTool\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "from langchain_core.messages import AIMessage, HumanMessage\n",
    "from langchain_core.output_parsers import PydanticOutputParser, StrOutputParser, PydanticToolsParser\n",
    "from langchain_core.utils.function_calling import convert_to_openai_tool\n",
    "\n",
    "\n",
    "import psycopg2\n",
    "import json\n",
    "import boto3\n",
    "from typing import Literal, Union, Optional, List\n",
    "from pydantic.v1 import BaseModel, Field"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup S3 bucket connection\n",
    "AWS_ACCESS_KEY_ID = config(\"AWS_ACCESS_KEY_ID\")\n",
    "AWS_SECRET_ACCESS_KEY = config(\"AWS_SECRET_ACCESS_KEY\")\n",
    "\n",
    "\n",
    "SESSION = boto3.Session(\n",
    "    aws_access_key_id=AWS_ACCESS_KEY_ID ,\n",
    "    aws_secret_access_key=AWS_SECRET_ACCESS_KEY)\n",
    "S3 = SESSION.client('s3')\n",
    "BUCKET_NAME = \"chichewa-ai\"\n",
    "USERS_FILE_KEY = 'chichewa-ai/phoso-ai-files/welcomed_users.json'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "AWS_SECRET_ACCESS_KEY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_welcomed_users():\n",
    "    \"\"\"\n",
    "    Loads the list of welcomed users from an S3 bucket.\n",
    "\n",
    "    This function attempts to download and load a JSON file from an S3 bucket \n",
    "    that contains the list of phone numbers for users who have already received \n",
    "    a welcome message. If the file does not exist, the function returns an empty \n",
    "    dictionary.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    dict\n",
    "        A dictionary where the keys are user phone numbers and the values indicate \n",
    "        whether the user has been welcomed. If the S3 file does not exist, an \n",
    "        empty dictionary is returned.\n",
    "\n",
    "    Raises\n",
    "    ------\n",
    "    botocore.exceptions.BotoCoreError\n",
    "        If there is an error in accessing the S3 bucket, such as network issues \n",
    "        or incorrect credentials.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        # Download the file from S3\n",
    "        s3_response = S3.get_object(Bucket=BUCKET_NAME, Key=USERS_FILE_KEY)\n",
    "        users_data = s3_response['Body'].read().decode('utf-8')\n",
    "        return json.loads(users_data)\n",
    "    except S3.exceptions.NoSuchKey:\n",
    "        # If the file doesn't exist, return an empty dictionary\n",
    "        return {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_welcomed_users(welcomed_users):\n",
    "    \"\"\"\n",
    "    Saves the list of welcomed users to an S3 bucket.\n",
    "\n",
    "    This function takes a dictionary of welcomed users, converts it into a JSON string,\n",
    "    and uploads it to a specified S3 bucket. The JSON file stores the phone numbers of users \n",
    "    who have already received a welcome message.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    welcomed_users : dict\n",
    "        A dictionary where the keys are user phone numbers and the values indicate \n",
    "        whether the user has been welcomed.\n",
    "\n",
    "    Raises\n",
    "    ------\n",
    "    botocore.exceptions.BotoCoreError\n",
    "        If there is an error in uploading the file to the S3 bucket, such as network issues \n",
    "        or incorrect credentials.\n",
    "    \"\"\"\n",
    "    # Convert the dictionary to a JSON string\n",
    "    users_data = json.dumps(welcomed_users)\n",
    "    # Upload the JSON string to S3\n",
    "    S3.put_object(Bucket=BUCKET_NAME, Key=USERS_FILE_KEY, Body=users_data)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain import LLMChain, PromptTemplate\n",
    "from langchain.llms import OpenAI\n",
    "\n",
    "def generate_welcome_message(language):\n",
    "    # Initialize the LLM (e.g., OpenAI GPT-3.5)\n",
    "    llm = OpenAI(temperature=0.7)\n",
    "\n",
    "    if language == 'Chichewa':\n",
    "        prompt_text = \"\"\"\n",
    "        You are a helpful assistant who speaks Chichewa. Generate a warm welcome message in Chichewa.\n",
    "        \n",
    "        Message:\n",
    "        \"\"\"\n",
    "    else:  # Default to English\n",
    "        prompt_text = \"\"\"\n",
    "        You are a helpful assistant who speaks English. Generate a warm welcome message in English.\n",
    "        \n",
    "        Message:\n",
    "        \"\"\"\n",
    "\n",
    "    prompt = PromptTemplate.from_template(prompt_text)\n",
    "    chain = LLMChain(llm=llm, prompt=prompt)\n",
    "    \n",
    "    # Generate the welcome message\n",
    "    return chain.run({})\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain import LLMChain, PromptTemplate\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain.schema import HumanMessage\n",
    "\n",
    "# Define the prompt template for the welcome message\n",
    "welcome_message_prompt = PromptTemplate.from_template(\n",
    "    \"\"\"\n",
    "    You are a helpful assistant who interacts with users in their preferred language. \n",
    "    Your task is to greet the user in their language and briefly inform them about the type of information you can provide.\n",
    "\n",
    "    Specifically, you provide the following types of information:\n",
    "    - Prices for food and other agricultural commodities (e.g., Maize, rice, soy beans).\n",
    "    - Agricultural production details (e.g., Maize, Tobacco).\n",
    "    - The situation of food security (e.g., how many people are lacking food).\n",
    "    - All information is specific to Malawi.\n",
    "\n",
    "    Greet the user and let them know that they can ask questions about these topics. Please generate the welcome message in the user's language.\n",
    "\n",
    "    Here are some examples:\n",
    "    \n",
    "    Example 1:\n",
    "    Text: \"Hello, how are you?\"\n",
    "    Language: English\n",
    "    Welcome Message: \"Welcome! I can help you with information about food prices, agricultural production, and food security in Malawi. How can I assist you today?\"\n",
    "\n",
    "    Example 2:\n",
    "    Text: \"Moni, muli bwanji?\"\n",
    "    Language: Chichewa\n",
    "    Welcome Message: \"Takulandirani! Ndikhoza kukuthandizani ndi zambiri zokhudza mitengo ya chakudya, ulimi, ndi chitetezo cha chakudya ku Malawi. Kodi ndingakuthandizeni bwanji lero?\"\n",
    "\n",
    "    Now, based on the user's input, detect the language and generate a suitable welcome message.\n",
    "    \n",
    "    Text: \"{text_to_detect}\"\n",
    "    Language:\n",
    "    Welcome Message:\n",
    "    \"\"\"\n",
    ")\n",
    "\n",
    "# Initialize the chat-based model (e.g., GPT-3.5-turbo)\n",
    "llm = ChatOpenAI(temperature=0.7)\n",
    "\n",
    "# Create the LLMChain with the prompt\n",
    "welcome_message_chain = LLMChain(llm=llm, prompt=welcome_message_prompt)\n",
    "\n",
    "# Generate the welcome message\n",
    "def generate_welcome_message(user_message):\n",
    "    return welcome_message_chain.run({\"text_to_detect\": user_message})\n",
    "\n",
    "# Example usage\n",
    "user_message = \"Moni, muli bwanji?\"\n",
    "welcome_message = generate_welcome_message(user_message)\n",
    "print(welcome_message)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "users = load_welcomed_users()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "users['+12023726721'] = True\n",
    "users['+265999392358'] = True\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_welcomed_users(welcomed_users=users)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "welcomed_users = load_welcomed_users()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "welcomed_users"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def detect_language(text, llm=None):\n",
    "    \n",
    "    language_detection_prompt = PromptTemplate.from_template(\n",
    "    \"\"\"\n",
    "    You are a language detection expert. Your task is to identify the language of the given text accurately.\n",
    "    Respond with only the name of the language (e.g., \"English\", \"Chichewa\", \"Spanish\", etc.).\n",
    "\n",
    "    Here are some examples:\n",
    "    \n",
    "    Example 1:\n",
    "    Text: \"Hello, how are you?\"\n",
    "    Language: English\n",
    "    \n",
    "    Example 2:\n",
    "    Text: \"Moni, muli bwanji?\"\n",
    "    Language: Chichewa\n",
    "    \n",
    "    Example 3:\n",
    "    Text: \"ndikuti kukupezeka nyemba zambiri\"\n",
    "    Language: Chichewa\n",
    "    \n",
    "    Now, identify the language for the following text:\n",
    "    \n",
    "    Text: \"{text_to_detect}\"\n",
    "    Language:\n",
    "    \"\"\")\n",
    "\n",
    "    if not llm:\n",
    "        llm = ChatOpenAI(model=\"gpt-3.5-turbo\", temperature=0, openai_api_key=OPENAI_API_KEY)\n",
    "\n",
    "    # Create the LLMChain with the prompt\n",
    "    language_detection_chain = language_detection_prompt | llm | StrOutputParser()\n",
    "\n",
    "    # Detect the language\n",
    "    detected_language = language_detection_chain.invoke({\"text_to_detect\": text})\n",
    "\n",
    "    return detected_language"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_welcome_message(user_message, llm=None):\n",
    "    \n",
    "    # Define the prompt template for the welcome message\n",
    "    welcome_message_prompt = PromptTemplate.from_template(\n",
    "    \"\"\"\n",
    "    You are a helpful assistant who interacts with users in their preferred language. \n",
    "    Your task is to greet the user in their language and briefly inform them about the type of information you can provide.\n",
    "\n",
    "    Specifically, you provide the following types of information:\n",
    "    - Prices for food and other agricultural commodities (e.g., Maize, rice, soy beans).\n",
    "    - Agricultural production details (e.g., Maize, Tobacco).\n",
    "    - The situation of food security (e.g., how many people are lacking food).\n",
    "    - All information is specific to Malawi.\n",
    "\n",
    "    Greet the user and let them know that they can ask questions about these topics. \n",
    "    \n",
    "    Here are some examples:\n",
    "    \n",
    "    Example 1:\n",
    "    Text: \"Hello, how are you?\"\n",
    "    Welcome Message: \"Welcome! I can help you with information about food prices, agricultural production, and food security in Malawi. For example, you can ask: 'What is the current price of maize?' or 'How much maize was produced last year?How can I assist you today?'\"\n",
    "    \n",
    "    Example 2:\n",
    "    Text: \"Moni, muli bwanji?\"\n",
    "    Welcome Message: \"Takulandirani! Ndikhoza kukuthandizani ndi zambiri zokhudzana ndi mitengo ya chakudya, zokolola, zaulimi komanso zokhudzana ndi zanjala mmene ilili ku Malawi. Mwachitsanzo mutha kufunsa kuti: \"Kodi chimanga chili pa bwanji ku Kasungu?\"\n",
    "    'Kodi ndikuti kunakololedwa mtedza wambiri?' kapena 'Kodi ndikuti kunakololedwa mtedza wambiri?'. Kodi ndingakuthandizeni bwanji lero?\"\n",
    "    \n",
    "\n",
    "    Now, based on the user's input, generate a suitable welcome message in the same language.\n",
    "\n",
    "    Text: \"{text_to_detect}\"\n",
    "    Welcome Message:\n",
    "    Example Questions:\n",
    "    1. \n",
    "    2. \n",
    "    \"\"\")\n",
    "\n",
    "    # Initialize the chat-based model (e.g., GPT-3.5-turbo)\n",
    "    if not llm:\n",
    "        llm = ChatOpenAI(temperature=0.7)\n",
    "\n",
    "    # Create the LLMChain with the prompt\n",
    "    welcome_message_chain = welcome_message_prompt | llm | StrOutputParser()\n",
    "\n",
    "    return welcome_message_chain.invoke({\"text_to_detect\": user_message})\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "user_question_en = \"Whats the price of Maize\"\n",
    "user_question_ny = \"Kodi ndi boma liti anakolola chimanga chambiri pakati pa Lilongwe kapena Kasungu?\"\n",
    "response = generate_welcome_message(user_message=user_question_en)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "users = {\"whatsapp:+12022948588\": True}\n",
    "number = \"whatsapp:+12023923333\"\n",
    "number2 = \"whatsapp:+12022948588\"\n",
    "\n",
    "if number2 not in users:\n",
    "    print(\"Y\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import translate_text_openai, detect_language_with_langchain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "user_question = \"Kodi ndi boma liti anakolola chimanga chambiri pakati pa Lilongwe kapena Kasungu?\"\n",
    "response = translate_text_openai(text=user_question, \n",
    "                                 source_language=\"Chichewa\", target_language=\"English\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "import logging\n",
    "\n",
    "# Set the logging level for the `httpx` logger to WARNING to suppress INFO logs\n",
    "logging.getLogger(\"httpx\").setLevel(logging.WARNING)\n",
    "\n",
    "# You can also suppress other loggers if necessary\n",
    "logging.getLogger(\"openai\").setLevel(logging.WARNING)\n",
    "logging.getLogger(\"langchain\").setLevel(logging.WARNING)\n",
    "\n",
    "# Set the logging level for langsmith.client to ERROR to suppress warnings\n",
    "logging.getLogger(\"langsmith.client\").setLevel(logging.ERROR)\n",
    "\n",
    "from sql_chain import process_sql_query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "questions = questions = [\"What is the price of Maize in Rumphi\",\n",
    "                 \"Where can I find the cheapest maize?\",\n",
    "                    \"Which district harvested the most beans?\",\n",
    "                    \"How much is Maize in Zomba?\",\n",
    "                    \"Which district produced more Tobacco, Mchinji or Kasungu?\",\n",
    "                    \"Where can I get bananas?\", \"Kodi chimanga chotchipa ndingachipeze kuti?\",\n",
    "                    \"Ndi boma liti komwe anakolola nyemba zambiri?\",\n",
    "                    \"Ku Zomba chimanga akugulitsa pa bwanji?\",\n",
    "                    \"Kodi ndi boma liti anakolola chimanga chambiri pakati pa Lilongwe kapena Kasungu?\",\n",
    "                    \"Ndikuti ndingapeze mpunga wambiri?\"]\n",
    "\n",
    "for q in questions:\n",
    "    print()\n",
    "    print(\"QUESTION:\", q)\n",
    "    response= process_sql_query(q)\n",
    "    print(\"LLM Response:\", response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import translate_text_openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "translate_text_openai(\"You can find a lot of rice in Karonga, Nkhotakota, Nkhata, Salima, and Dedza.\", \n",
    "                                                    source_language=\"English\",\n",
    "                                                    target_language=\"Chichewa\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_translation_examples(file_path, source_language, target_language):\n",
    "    \"\"\"\n",
    "    Loads and formats translation examples from a JSON file.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    file_path : str\n",
    "        The path to the JSON file containing translation examples.\n",
    "    source_language : str\n",
    "        The source language (e.g., \"Chichewa\").\n",
    "    target_language : str\n",
    "        The target language (e.g., \"English\").\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    list\n",
    "        A list of formatted translation examples.\n",
    "    \"\"\"\n",
    "    with open(file_path, 'r') as file:\n",
    "        examples = json.load(file)\n",
    "    \n",
    "    key = f\"{source_language}-{target_language}\"\n",
    "    if key in examples:\n",
    "        return examples[key]\n",
    "    else:\n",
    "        raise ValueError(f\"Translation examples for {source_language} to {target_language} not found.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ex = load_translation_examples(file_path=\"translation_examples.json\", \n",
    "                          source_language=\"Chichewa\", target_language=\"English\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain import LLMChain, PromptTemplate\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "\n",
    "def translate_text_openai(text, source_language, target_language, examples):\n",
    "    \"\"\"\n",
    "    Translates the given text from the source language to the target language using an LLM with few-shot examples.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    text : str\n",
    "        The text to be translated.\n",
    "    source_language : str\n",
    "        The source language of the text.\n",
    "    target_language : str\n",
    "        The language into which the text should be translated.\n",
    "    examples : list\n",
    "        A list of few-shot translation examples.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    str\n",
    "        The translated text.\n",
    "    \"\"\"\n",
    "    # Construct the prompt template using examples\n",
    "    example_prompts = \"\\n\".join([f'{source_language}: \"{ex[source_language]}\"\\n{target_language}: \"{ex[target_language]}\"' for ex in examples])\n",
    "\n",
    "    prompt_template = PromptTemplate.from_template(\n",
    "        f\"\"\"\n",
    "        You are a professional translator who specializes in translating text from {source_language} to {target_language}.\n",
    "        Given the following examples, translate the provided text.\n",
    "\n",
    "        Examples:\n",
    "        {example_prompts}\n",
    "\n",
    "        Now, translate the following:\n",
    "\n",
    "        {source_language}: \"{{text}}\"\n",
    "        {target_language}:\n",
    "        \"\"\"\n",
    "    )\n",
    "\n",
    "    # Initialize the chat-based model\n",
    "    llm = ChatOpenAI(temperature=0.7, model=\"gpt-4o\")\n",
    "\n",
    "    # Create the LLMChain for translation\n",
    "    translation_chain = LLMChain(llm=llm, prompt=prompt_template)\n",
    "\n",
    "    # Perform the translation\n",
    "    return translation_chain.run({\"text\": text})\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ex1 = \"You can find a lot of beans in Karonga, Nkhotakota, Nkhata, Salima, and Dedza.\"\n",
    "ex2 = \"Ndi boma liti komwe anakolola nyemba zambiri?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example usage\n",
    "file_path = 'translation_examples.json'\n",
    "source_language = \"English\"\n",
    "target_language = 'Chichewa'\n",
    "examples = load_translation_examples(file_path, source_language, target_language)\n",
    "translated_text = translate_text_openai(\n",
    "    text=ex1, \n",
    "    source_language=source_language, \n",
    "    target_language=target_language,\n",
    "    examples=examples\n",
    ")\n",
    "print(\"Translated Question:==>\", translated_text)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "translate_text_openai(\"You can find a lot of rice in Karonga, Nkhotakota, Nkhata, Salima, and Dedza.\", \n",
    "                                                    source_language=\"English\",\n",
    "                                                    target_language=\"Chichewa\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain import LLMChain, PromptTemplate\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "\n",
    "def translate_to_chichewa(text):\n",
    "    translation_prompt = PromptTemplate.from_template(\n",
    "        \"\"\"\n",
    "        You are a professional translator who specializes in translating English sentences to Chichewa. \n",
    "        Translate the following English sentence into Chichewa, ensuring the translation is accurate and maintains the original meaning.\n",
    "\n",
    "        Example:\n",
    "        - English: \"I love eating nsima.\"\n",
    "        - Chichewa: \"Ndimakonda kudya nsima.\"\n",
    "        - English: \"Where can I find cheap beans?.\"\n",
    "        - Chichewa: \"Nyemba zotchipa zikupezeka kuti?.\"\n",
    "\n",
    "        Now, translate the following sentence:\n",
    "        - English: \"{text}\"\n",
    "        - Chichewa:\n",
    "        \"\"\"\n",
    "    )\n",
    "    # Initialize the chat-based model\n",
    "    llm = ChatOpenAI(temperature=0.7, model='gpt-4')\n",
    "\n",
    "    # Create the LLMChain\n",
    "    translation_chain = LLMChain(llm=llm, prompt=translation_prompt)\n",
    "\n",
    "    # Perform the translation\n",
    "    return translation_chain.run({\"text\": text})\n",
    "\n",
    "# Example usage\n",
    "english_sentence = \"where can I get beans?\"\n",
    "chichewa_translation = translate_to_chichewa(english_sentence)\n",
    "print(chichewa_translation)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = ChatOpenAI(temperature=0.7, model='gpt-4')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm.model_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "quest_lan = detect_language_with_langchain(text='How are you?')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "detect_language(text=\"Kodi ndi boma liti anakolola chimanga chambiri pakati pa Lilongwe kapena Kasungu?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "OPENAI_API_KEY = config(\"OPENAI_API_KEY\")\n",
    "MISTRAL_API_KEY = config(\"MISTRAL_API_KEY\")\n",
    "DB_USER = config('DB_USER')\n",
    "DB_PASSWORD = config('DB_PASSWORD')\n",
    "DB_HOST = 'localhost'\n",
    "DB_PORT = '5432'\n",
    "DB_NAME = \"food_security\"\n",
    "SAMPLE_QUESTIONS = {\"low-birthweight\": \"Which region has the highest number of children born with low birth weights?\",\n",
    "                    \"vaccine_rates\": \"Which vaccine has the lowest vaccination percentage?\",\n",
    "                    \"vaccine_rates_all\": \"What percentage of children received all vaccines before 12 months\",\n",
    "                    \"kids_in_sch\": \"Whats average percentage of children who are in preschool\",\n",
    "                    \"vaccines\": \"which vaccines did children get in Tunisia?\"\n",
    "                    }\n",
    "                    \n",
    "# Create the database URL\n",
    "DATABASE_URL = f'postgresql://{DB_USER}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_NAME}'\n",
    "FILE_SQL_EXAMPLES_EN = \"sql_examples_en.json\"\n",
    "USE_BEST_MATCHING_COLUMNS = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LangChain Evals"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "import requests\n",
    "import io\n",
    "\n",
    "# Dataset URL:\n",
    "url = \"https://storage.googleapis.com/oreilly-content/transaction_data_with_expanded_descriptions.csv\"\n",
    "\n",
    "# Download the file from the URL:\n",
    "downloaded_file = requests.get(url)\n",
    "\n",
    "# Load the transactions dataset and only look at 20 transactions:\n",
    "df = pd.read_csv(io.StringIO(downloaded_file.text))[:20]\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the Mistral Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ChatMistralAI(model=\"mistral-small\", mistral_api_key=MISTRAL_API_KEY)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2. Define the prompt:\n",
    "system_prompt = \"\"\"You are are an expert at analyzing\n",
    "bank transactions, you will be categorizing a single\n",
    "transaction.\n",
    "Always return a transaction type and category:\n",
    "do not return None.\n",
    "Format Instructions:\n",
    "{format_instructions}\"\"\"\n",
    "\n",
    "user_prompt = \"\"\"Transaction Text:\n",
    "{transaction}\"\"\"\n",
    "\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            system_prompt,\n",
    "        ),\n",
    "        (\n",
    "            \"user\",\n",
    "            user_prompt,\n",
    "        ),\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the pydantic model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 3. Define the pydantic model:\n",
    "class EnrichedTransactionInformation(BaseModel):\n",
    "    transaction_type: Union[\n",
    "        Literal[\"Purchase\", \"Withdrawal\", \"Deposit\",\n",
    "        \"Bill Payment\", \"Refund\"], None\n",
    "    ]\n",
    "    transaction_category: Union[\n",
    "        Literal[\"Food\", \"Entertainment\", \"Transport\",\n",
    "        \"Utilities\", \"Rent\", \"Other\"],\n",
    "        None,\n",
    "    ]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. Define the output parser:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_parser = PydanticOutputParser(\n",
    "    pydantic_object=EnrichedTransactionInformation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 5. Define a function to try to fix and remove the backslashes:\n",
    "def remove_back_slashes(string):\n",
    "    # double slash to escape the slash\n",
    "    cleaned_string = string.replace(\"\\\\\", \"\")\n",
    "    return cleaned_string"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create an LCEL chain that fixes the formatting:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 6. \n",
    "chain = prompt | model | StrOutputParser() | remove_back_slashes | output_parser"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Invoke the chain for a single instance of the dataframe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transaction = df.iloc[0][\"Transaction Description\"]\n",
    "result = chain.invoke(\n",
    "        {\n",
    "            \"transaction\": transaction,\n",
    "            \"format_instructions\": \\\n",
    "            output_parser.get_format_instructions(),\n",
    "        }\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Invoke the chain on the whole dataframe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = []\n",
    "\n",
    "for i, row in tqdm(df.iterrows(), total=len(df)):\n",
    "    transaction = row[\"Transaction Description\"]\n",
    "    try:\n",
    "        result = chain.invoke(\n",
    "            {\n",
    "                \"transaction\": transaction,\n",
    "                \"format_instructions\": output_parser.get_format_instructions(),\n",
    "            }\n",
    "        )\n",
    "        print(result.transaction_type, result.transaction_category)\n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "        result = EnrichedTransactionInformation(\n",
    "            transaction_type=None,\n",
    "            transaction_category=None\n",
    "        )\n",
    "\n",
    "    results.append(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 8. Add the results to the dataframe, as columns transaction type and\n",
    "# transaction category:\n",
    "transaction_types = []\n",
    "transaction_categories = []\n",
    "\n",
    "for result in results:\n",
    "    transaction_types.append(result.transaction_type)\n",
    "    transaction_categories.append(\n",
    "        result.transaction_category)\n",
    "\n",
    "df[\"mistral_transaction_type\"] = transaction_types\n",
    "df[\"mistral_transaction_category\"] = transaction_categories\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create alternate/ground truth with GPT4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_gpt4 = ChatOpenAI(model=\"gpt-4o\")\n",
    "chain_gpt = prompt | model_gpt4 | StrOutputParser() | output_parser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df[\"gpt4_transaction_type\"] = None\n",
    "df[\"gpt4_transaction_category\"] = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, row in tqdm(df.iterrows(), total=len(df)):\n",
    "    transaction = row[\"Transaction Description\"]\n",
    "    try:\n",
    "        result = chain_gpt.invoke(\n",
    "            {\n",
    "                \"transaction\": transaction,\n",
    "                \"format_instructions\": output_parser.get_format_instructions(),\n",
    "            }\n",
    "        )\n",
    "        df.loc[i, \"gpt4_transaction_type\"] = result.transaction_type\n",
    "        df.loc[i, \"gpt4_transaction_category\"] = result.transaction_category\n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "        result = EnrichedTransactionInformation(\n",
    "            transaction_type=None,\n",
    "            transaction_category=None\n",
    "        )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Function Calling in LangChain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Article(BaseModel):\n",
    "    \"\"\"Identifying key points and contrarian views in an article.\"\"\"\n",
    "\n",
    "    points: str = Field(..., description=\"Key points from the article\")\n",
    "    contrarian_points: Optional[str] = Field(\n",
    "        None, description=\"Any contrarian points acknowledged in the article\"\n",
    "    )\n",
    "    author: Optional[str] = Field(None, description=\"Author of the article\")\n",
    "\n",
    "_EXTRACTION_TEMPLATE = \"\"\"Extract and save the relevant entities mentioned \\\n",
    "in the following passage together with their properties.\n",
    "\n",
    "If a property is not present and is not required in the function parameters,\n",
    "do not include it in the output.\"\"\"\n",
    "\n",
    "# Create a prompt telling the LLM to extract information:\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    {(\"system\", _EXTRACTION_TEMPLATE), (\"user\", \"{input}\")}\n",
    ")\n",
    "\n",
    "pydantic_schemas = [Article]\n",
    "\n",
    "# Convert Pydantic objects to the appropriate schema:\n",
    "tools = [convert_to_openai_tool(p) for p in pydantic_schemas]\n",
    "\n",
    "# Give the model access to these tools:\n",
    "model = model_gpt4.bind_tools(tools=tools)\n",
    "\n",
    "# Create an end to end chain:\n",
    "chain = prompt | model | PydanticToolsParser(tools=pydantic_schemas)\n",
    "\n",
    "result = chain.invoke(\n",
    "    {\n",
    "        \"input\": \"\"\"In the recent article titled 'AI adoption in industry,'\n",
    "        key points addressed include the growing interest ... However, the\n",
    "        author, Dr. Jane Smith, ...\"\"\"\n",
    "    }\n",
    ")\n",
    "print(result)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Query Planning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Query(BaseModel):\n",
    "    id: int\n",
    "    question: str\n",
    "    dependencies: List[int] = Field(\n",
    "        default_factory=list,\n",
    "        description=\"\"\"A list of sub-queries that must be completed before\n",
    "        this task can be completed.\n",
    "        Use a sub query when anything is unknown and we might need to ask\n",
    "        many queries to get an answer.\n",
    "        Dependencies must only be other queries.\"\"\"\n",
    "    )\n",
    "\n",
    "class QueryPlan(BaseModel):\n",
    "    query_graph: List[Query]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up a chat model:\n",
    "model = ChatOpenAI()\n",
    "\n",
    "# Set up a parser:\n",
    "parser = PydanticOutputParser(pydantic_object=QueryPlan)\n",
    "\n",
    "template = \"\"\"Generate a query plan. This will be used for task execution.\n",
    "\n",
    "Answer the following query: {query}\n",
    "\n",
    "Return the following query graph format:\n",
    "{format_instructions}\n",
    "\"\"\"\n",
    "system_message_prompt = SystemMessagePromptTemplate.from_template(template)\n",
    "chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt])\n",
    "\n",
    "# Create the LCEL chain with the prompt, model, and parser:\n",
    "chain = chat_prompt | model | parser\n",
    "\n",
    "result = chain.invoke({\n",
    "\"query\":'''I want to get the results from my database. Then I want to find\n",
    "out what the average age of my top 10 customers is. Once I have the average\n",
    "age, I want to send an email to John. Also I just generally want to send a\n",
    "welcome introduction email to Sarah, regardless of the other tasks.''',\n",
    "\"format_instructions\":parser.get_format_instructions()})\n",
    "\n",
    "print(result.query_graph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_template = '''You are helpful assistant who can translate from {input_language} to {output_language}.'''\n",
    "user_template = \"User text: {user_text}\"\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [(\"system\",system_template),(\"user\", user_template)])\n",
    "question = \"Ndimakudonda kwambiri\"\n",
    "src_lan = \"Chichewa\"\n",
    "dest_lan = \"English\"\n",
    "\n",
    "chat = ChatOpenAI()\n",
    "chain = prompt | chat | StrOutputParser()\n",
    "response = chain.invoke({\"input_language\": \"Chichewa\", \"output_language\": \"English\", \"user_text\": question})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating Few-Shot Prompt Templates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "examples = [\n",
    "    {\n",
    "        \"question\": \"What is the capital of France?\",\n",
    "        \"answer\": \"Paris\",\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"What is the capital of Spain?\",\n",
    "        \"answer\": \"Madrid\",\n",
    "    } \n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "example_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"human\", \"{question}\"),\n",
    "        (\"ai\", \"{answer}\"),\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prompt Chaining"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts.chat import ChatPromptTemplate\n",
    "\n",
    "character_generation_prompt = ChatPromptTemplate.from_template(\n",
    "    \"\"\"I want you to brainstorm three to five characters for my short story. The\n",
    "    genre is {genre}. Each character must have a Name and a Biography.\n",
    "    You must provide a name and biography for each character, this is very\n",
    "    important!\n",
    "    ---\n",
    "    Example response:\n",
    "    Name: CharWiz, Biography: A wizard who is a master of magic.\n",
    "    Name: CharWar, Biography: A warrior who is a master of the sword.\n",
    "    ---\n",
    "    Characters: \"\"\"\n",
    ")\n",
    "\n",
    "plot_generation_prompt = ChatPromptTemplate.from_template(\n",
    "    \"\"\"Given the following characters and the genre, create an effective\n",
    "    plot for a short story:\n",
    "    Characters:\n",
    "    {characters}\n",
    "    ---\n",
    "    Genre: {genre}\n",
    "    ---\n",
    "    Plot: \"\"\"\n",
    "    )\n",
    "\n",
    "scene_generation_plot_prompt = ChatPromptTemplate.from_template(\n",
    "    \"\"\"Act as an effective content creator.\n",
    "    Given multiple characters and a plot, you are responsible for\n",
    "    generating the various scenes for each act.\n",
    "\n",
    "    You must decompose the plot into multiple effective scenes:\n",
    "    ---\n",
    "    Characters:\n",
    "    {characters}\n",
    "    ---\n",
    "    Genre: {genre}\n",
    "    ---\n",
    "    Plot: {plot}\n",
    "    ---\n",
    "    Example response:\n",
    "    Scenes:\n",
    "    Scene 1: Some text here.\n",
    "    Scene 2: Some text here.\n",
    "    Scene 3: Some text here.\n",
    "    ----\n",
    "    Scenes:\n",
    "    \"\"\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from operator import itemgetter\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "\n",
    "chain = RunnablePassthrough() | {\n",
    "    \"genre\": itemgetter(\"genre\"),\n",
    "  }\n",
    "\n",
    "chain.invoke({\"genre\": \"fantasy\"})\n",
    "# {'genre': 'fantasy'}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai.chat_models import ChatOpenAI\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "\n",
    "# Create the chat model:\n",
    "model = ChatOpenAI()\n",
    "\n",
    "# Create the subchains:\n",
    "character_generation_chain = ( character_generation_prompt\n",
    "| model\n",
    "| StrOutputParser() )\n",
    "\n",
    "plot_generation_chain = ( plot_generation_prompt\n",
    "| model\n",
    "| StrOutputParser() )\n",
    "\n",
    "scene_generation_plot_chain = ( scene_generation_plot_prompt\n",
    "| model\n",
    "| StrOutputParser()  )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.runnables import RunnableParallel\n",
    "from operator import itemgetter\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "\n",
    "master_chain = (\n",
    "    {\"characters\": character_generation_chain, \"genre\":\n",
    "    RunnablePassthrough()}\n",
    "    | RunnableParallel(\n",
    "        characters=itemgetter(\"characters\"),\n",
    "        genre=itemgetter(\"genre\"),\n",
    "        plot=plot_generation_chain,\n",
    "    )\n",
    "    | RunnableParallel(\n",
    "        characters=itemgetter(\"characters\"),\n",
    "        genre=itemgetter(\"genre\"),\n",
    "        plot=itemgetter(\"plot\"),\n",
    "        scenes=scene_generation_plot_chain,\n",
    "    )\n",
    ")\n",
    "\n",
    "story_result = master_chain.invoke({\"genre\": \"Fantasy\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "master_chain = (RunnablePassthrough() |\n",
    "    {\"characters\": character_generation_chain, \"genre\":\n",
    "    itemgetter(\"genre\")}\n",
    "    | {\n",
    "        \"characters\": itemgetter(\"characters\"),\n",
    "        \"genre\": itemgetter(\"genre\"),\n",
    "        \"plot\": plot_generation_chain,\n",
    "    }\n",
    "    | {\n",
    "        \"characters\": itemgetter(\"characters\"),\n",
    "        \"genre\": itemgetter(\"genre\"),\n",
    "        \"plot\": itemgetter(\"plot\"),\n",
    "        \"scenes\": scene_generation_plot_chain,\n",
    "\n",
    "    }\n",
    ")\n",
    "\n",
    "story_result = master_chain.invoke({\"genre\": \"Fantasy\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extracting the scenes using .split('\\n') and removing empty strings:\n",
    "scenes = [scene for scene in story_result[\"scenes\"].split(\"\\n\") if scene]\n",
    "generated_scenes = []\n",
    "previous_scene_summary = \"\"\n",
    "\n",
    "character_script_prompt = ChatPromptTemplate.from_template(\n",
    "    template=\"\"\"Given the following characters: {characters} and the genre:\n",
    "    {genre}, create an effective character script for a scene.\n",
    "\n",
    "    You must follow the following principles:\n",
    "    - Use the Previous Scene Summary: {previous_scene_summary} to avoid\n",
    "    repeating yourself.\n",
    "    - Use the Plot: {plot} to create an effective scene character script.\n",
    "    - Currently you are generating the character dialogue script for the\n",
    "    following scene: {scene}\n",
    "\n",
    "    ---\n",
    "    Here is an example response:\n",
    "    SCENE 1: ANNA'S APARTMENT\n",
    "\n",
    "    (ANNA is sorting through old books when there is a knock at the door.\n",
    "    She opens it to reveal JOHN.)\n",
    "    ANNA: Can I help you, sir?\n",
    "    JOHN: Perhaps, I think it's me who can help you. I heard you're\n",
    "    researching time travel.\n",
    "    (Anna looks intrigued but also cautious.)\n",
    "    ANNA: That's right, but how do you know?\n",
    "    JOHN: You could say... I'm a primary source.\n",
    "\n",
    "    ---\n",
    "    SCENE NUMBER: {index}\n",
    "\n",
    "    \"\"\",\n",
    ")\n",
    "\n",
    "summarize_prompt = ChatPromptTemplate.from_template(\n",
    "    template=\"\"\"Given a character script, create a summary of the scene.\n",
    "    Character script: {character_script}\"\"\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loading a chat model:\n",
    "model = ChatOpenAI(model='gpt-3.5-turbo-16k')\n",
    "\n",
    "# Create the LCEL chains:\n",
    "character_script_generation_chain = (\n",
    "    {\n",
    "        \"characters\": RunnablePassthrough(),\n",
    "        \"genre\": RunnablePassthrough(),\n",
    "        \"previous_scene_summary\": RunnablePassthrough(),\n",
    "        \"plot\": RunnablePassthrough(),\n",
    "        \"scene\": RunnablePassthrough(),\n",
    "        \"index\": RunnablePassthrough(),\n",
    "    }\n",
    "    | character_script_prompt\n",
    "    | model\n",
    "    | StrOutputParser()\n",
    ")\n",
    "\n",
    "summarize_chain = summarize_prompt | model | StrOutputParser()\n",
    "\n",
    "# You might want to use tqdm here to track the progress,\n",
    "# or use all of the scenes:\n",
    "for index, scene in enumerate(scenes[0:3]):\n",
    "\n",
    "    # # Create a scene generation:\n",
    "    scene_result = character_script_generation_chain.invoke(\n",
    "        {\n",
    "            \"characters\": story_result[\"characters\"],\n",
    "            \"genre\": \"fantasy\",\n",
    "            \"previous_scene_summary\": previous_scene_summary,\n",
    "            \"index\": index,\n",
    "        }\n",
    "    )\n",
    "\n",
    "    # Store the generated scenes:\n",
    "    generated_scenes.append(\n",
    "        {\"character_script\": scene_result, \"scene\": scenes[index]}\n",
    "    )\n",
    "\n",
    "    # If this is the first scene then we don't have a\n",
    "    # previous scene summary:\n",
    "    if index == 0:\n",
    "        previous_scene_summary = scene_result\n",
    "    else:\n",
    "        # If this is the second scene or greater then\n",
    "        # we can use and generate a summary:\n",
    "        summary_result = summarize_chain.invoke(\n",
    "            {\"character_script\": scene_result}\n",
    "        )\n",
    "        previous_scene_summary = summary_result\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts.chat import ChatPromptTemplate\n",
    "from operator import itemgetter\n",
    "from langchain_core.runnables import RunnablePassthrough, RunnableLambda\n",
    "\n",
    "bad_first_input = {\n",
    "    \"film_required_age\": itemgetter(\"age\"),\n",
    "}\n",
    "\n",
    "prompt = ChatPromptTemplate.from_template(\n",
    "    \"Generate a film title, the age is {film_required_age}\"\n",
    ")\n",
    "\n",
    "# This will error:\n",
    "bad_chain = bad_first_input | prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame(generated_scenes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_character_script_text = \"\\n\".join(df.character_script.tolist())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_text_splitters import CharacterTextSplitter\n",
    "from langchain.chains.summarize import load_summarize_chain\n",
    "text_splitter = CharacterTextSplitter.from_tiktoken_encoder(\n",
    "    chunk_size=1500, chunk_overlap=200\n",
    ")\n",
    "docs = text_splitter.create_documents([all_character_script_text])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chain = translation_template | chat | StrOutputParser()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = chain.invoke({\"input_language\": \"Chichewa\", \"output_language\": \"English\",\n",
    "               \"messages\": [HumanMessage(content=question)]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "OPENAI_API_KEY = config(\"OPENAI_API_KEY\")\n",
    "MISTRAL_API_KEY = config(\"MISTRAL_API_KEY\")\n",
    "DB_USER = config('DB_USER')\n",
    "DB_PASSWORD = config('DB_PASSWORD')\n",
    "DB_HOST = 'localhost'\n",
    "DB_PORT = '5432'\n",
    "DB_NAME = \"food_security\"\n",
    "SAMPLE_QUESTIONS = {\"low-birthweight\": \"Which region has the highest number of children born with low birth weights?\",\n",
    "                    \"vaccine_rates\": \"Which vaccine has the lowest vaccination percentage?\",\n",
    "                    \"vaccine_rates_all\": \"What percentage of children received all vaccines before 12 months\",\n",
    "                    \"kids_in_sch\": \"Whats average percentage of children who are in preschool\",\n",
    "                    \"vaccines\": \"which vaccines did children get in Tunisia?\"\n",
    "                    }\n",
    "                    \n",
    "# Create the database URL\n",
    "DATABASE_URL = f'postgresql://{DB_USER}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_NAME}'\n",
    "FILE_SQL_EXAMPLES_EN = \"sql_examples_en.json\"\n",
    "USE_BEST_MATCHING_COLUMNS = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prompt Templates"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Prompt to Select best Table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def connect_to_database(database_url=DATABASE_URL):\n",
    "    \"\"\"Connects to a postgreSQL\n",
    "\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    database_url : String\n",
    "        postgreSQL database connection URL, by default DATABASE_URL\n",
    "    \"\"\"\n",
    "    # conn = psycopg2.connect(f\"dbname={DB_NAME} user={DB_USER} password={DB_PASSWORD}\")\n",
    "    conn = psycopg2.connect(database_url)\n",
    "\n",
    "    cur = conn.cursor()\n",
    "\n",
    "    # Query to get table names and column names\n",
    "    cur.execute(\"SELECT table_name, description FROM table_metadata\")\n",
    "    tables = cur.fetchall()\n",
    "\n",
    "    cur.execute(\"SELECT table_name, column_name, description FROM column_metadata\")\n",
    "    columns = cur.fetchall()\n",
    "\n",
    "    cur.close()\n",
    "    conn.close()\n",
    "\n",
    "    return tables, columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_best_table_prompt(user_query, tables, columns, \n",
    "                           return_chain=True, llm=None):# Define the template for selecting the best table\n",
    "    template = \"\"\"\n",
    "    You are a database assistant. Given the following tables and columns with their descriptions, select the best table that matches the user's query.\n",
    "\n",
    "    Tables and Columns:\n",
    "    {table_info}\n",
    "\n",
    "    User Query:\n",
    "    {user_query}\n",
    "\n",
    "    Provide the output in the following JSON format:\n",
    "    {{\n",
    "        \"best_matching_table\": {{\n",
    "            \"table_name\": \"<best_table_name>\",\n",
    "            \"description\": \"<best_table_description>\"\n",
    "        }}\n",
    "    }}\n",
    "    \"\"\"\n",
    "    # Prepare the table_info string including descriptions for each table and its columns\n",
    "    table_info = \"\"\n",
    "    for table in tables:\n",
    "        table_name, table_description = table\n",
    "        table_info += f\"Table: {table_name} - {table_description}\\n\"\n",
    "        table_columns = [col for col in columns if col[0] == table_name]\n",
    "        for column in table_columns:\n",
    "            _, column_name, column_description = column\n",
    "            table_info += f\"    Column: {column_name} - {column_description}\\n\"\n",
    "        table_info += \"\\n\"\n",
    "\n",
    "    # Create the PromptTemplate\n",
    "    prompt_template = PromptTemplate(\n",
    "        template=template,\n",
    "        input_variables=[\"table_info\", \"user_query\"]\n",
    "    )\n",
    "\n",
    "    # Format the template \n",
    "    formatted_prompt = prompt_template.format(table_info=table_info, user_query=user_query)\n",
    "\n",
    "    if return_chain:\n",
    "        # Create the chain using the ChatOpenAI model and the PromptTemplate\n",
    "        chain = LLMChain(llm=llm,prompt=prompt_template)\n",
    "        return chain, {\"table_info\": table_info, \"user_query\": user_query}\n",
    "\n",
    "    return formatted_prompt\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_columns_info(table_name, columns):\n",
    "    columns_info = \"\"\n",
    "    for column in columns:\n",
    "        table, column_name, column_description = column\n",
    "        if table == table_name:\n",
    "            columns_info += f\"    Column: {column_name} - {column_description}\\n\"\n",
    "    return columns_info\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_best_columns_prompt(user_query, best_matching_table, columns, \n",
    "                       return_chain=True, llm=None):\n",
    "    # Define the template for selecting the relevant columns\n",
    "    column_template = \"\"\"\n",
    "    You are a database assistant. Given the following columns for the table '{table_name}', select the columns that are most relevant to the user's query.\n",
    "\n",
    "    Table Description: {table_description}\n",
    "\n",
    "    Columns:\n",
    "    {columns_info}\n",
    "\n",
    "    User Query:\n",
    "    {user_query}\n",
    "\n",
    "    Relevant Columns:\n",
    "    \"\"\"\n",
    "\n",
    "    columns_info = get_columns_info(best_matching_table[\"table_name\"], columns)\n",
    "\n",
    "    # Create the PromptTemplate for column selection\n",
    "    column_prompt_template = PromptTemplate(\n",
    "        template=column_template,\n",
    "        input_variables=[\"table_name\", \"table_description\", \"columns_info\", \"user_query\"]\n",
    "    )\n",
    "\n",
    "    # Example usage of the template with a user query\n",
    "    formatted_column_prompt = column_prompt_template.format(\n",
    "        table_name=best_matching_table[\"table_name\"],\n",
    "        table_description=best_matching_table[\"description\"],\n",
    "        columns_info=columns_info,\n",
    "        user_query=user_query\n",
    "    )\n",
    "\n",
    "    # Prepare the context for running the chain\n",
    "    context = {\n",
    "        \"table_name\": best_matching_table[\"table_name\"],\n",
    "        \"table_description\": best_matching_table[\"description\"],\n",
    "        \"columns_info\": columns_info,\n",
    "        \"user_query\": user_query}\n",
    "\n",
    "    if return_chain:\n",
    "        chain = LLMChain(llm=llm,prompt=column_prompt_template)\n",
    "        return chain, context\n",
    "\n",
    "    return formatted_column_prompt\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_sql_examples(file_path):\n",
    "    with open(file_path, 'r', encoding='utf-8') as file:\n",
    "        return json.load(file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_sql_prompt(examples, best_matching_table, columns_metadata, use_best_matching_columns=False):\n",
    "    \"\"\"\n",
    "    Creates a FewShotPromptTemplate for generating SQL queries based on table and column metadata.\n",
    "\n",
    "    This function generates a prompt template that includes detailed information about the table and its columns.\n",
    "    The generated prompt instructs a language model (LLM) to create a syntactically correct SQL query based on\n",
    "    user input. If the table contains a date column and the user does not specify a date, the prompt also instructs\n",
    "    the LLM to retrieve the most recent data available.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    examples : list of dict\n",
    "        A list of example inputs and corresponding SQL queries. Each example should be a dictionary with 'input' and 'query' keys.\n",
    "    best_matching_table : dict\n",
    "        A dictionary containing the best matching table information with 'table_name' and 'description' keys.\n",
    "    columns_metadata : list of tuples\n",
    "        A list of tuples containing columns metadata. Each tuple should include 'table_name', 'column_name', and 'description'.\n",
    "    use_best_matching_columns : bool, optional\n",
    "        A flag indicating whether to use only the best-matching columns (if True) or all columns in the table (if False). Default is True.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    sql_prompt : FewShotPromptTemplate\n",
    "        A FewShotPromptTemplate object that can be used with an LLM to generate SQL queries.\n",
    "    \"\"\"\n",
    "    # Prepare table_info string based on the best matching table and columns\n",
    "    table_info = f\"Table: {best_matching_table['table_name']} - {best_matching_table['description']}\\n\"\n",
    "    columns_info = \"Columns:\\n\"\n",
    "    has_date_column = False\n",
    "\n",
    "    # Determine which columns to use: best-matching or all columns\n",
    "    if use_best_matching_columns:\n",
    "        # If using best_matching_columns, use those provided (filtering columns_metadata based on matching logic)\n",
    "        columns_to_use = columns_metadata  # Assuming columns_metadata is already filtered\n",
    "    else:\n",
    "        # Use all columns for the given table from columns_metadata\n",
    "        columns_to_use = [col for col in columns_metadata if col[0] == best_matching_table['table_name']]\n",
    "\n",
    "    # Construct the columns_info string\n",
    "    for column in columns_to_use:\n",
    "        table_name, column_name, column_description = column\n",
    "        columns_info += f\"    Column: {column_name} - {column_description}\\n\"\n",
    "        if 'date' in column_name.lower():\n",
    "            has_date_column = True\n",
    "\n",
    "    # Create FewShot Prompt with instructions for handling most recent data\n",
    "    example_prompt = PromptTemplate.from_template(\"User input: {input}\\nSQL query: {query}\")\n",
    "\n",
    "    # Add a special instruction if the table has a date column\n",
    "    recent_data_instruction = (\n",
    "        \"If the user does not specify a date, retrieve the most recent data available by ordering the results \"\n",
    "        \"by the date column in descending order.\"\n",
    "    ) if has_date_column else \"\"\n",
    "\n",
    "    # Combine table_info and columns_info in the prompt\n",
    "    sql_prompt = FewShotPromptTemplate(\n",
    "        examples=examples,\n",
    "        example_prompt=example_prompt,\n",
    "        prefix=(\n",
    "            \"You are a PostgreSQL expert. Given an input question, create a syntactically correct PostgreSQL query to run. \"\n",
    "            \"Unless otherwise specified, do not return more than {top_k} rows.\\n\\n\"\n",
    "            f\"Here is the relevant table information:\\n{table_info}\\n\\n\"\n",
    "            f\"Here is the relevant columns information:\\n{columns_info}\\n\\n\"\n",
    "            f\"{recent_data_instruction}\\n\\n\"\n",
    "            \"Below are a number of examples of questions and their corresponding SQL queries.\"\n",
    "        ),\n",
    "        suffix=\"User input: {input}\\nSQL query: \",\n",
    "        input_variables=[\"input\", \"top_k\"],\n",
    "    )\n",
    "\n",
    "    return sql_prompt\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# def create_sql_prompt(examples, best_matching_table, best_matching_columns):\n",
    "#     # Prepare table_info string based on the best matching table and columns\n",
    "#     table_info = f\"Table: {best_matching_table['table_name']} - {best_matching_table['description']}\\n\"\n",
    "#     has_date_column = False\n",
    "\n",
    "#     for column in best_matching_columns:\n",
    "#         table_info += f\"    Column: {column['name']} - {column['description']}\\n\"\n",
    "#         if 'date' in column['name'].lower():\n",
    "#             has_date_column = True\n",
    "\n",
    "#     # Create FewShot Prompt with instructions for handling most recent data\n",
    "#     example_prompt = PromptTemplate.from_template(\"User input: {input}\\nSQL query: {query}\")\n",
    "\n",
    "#     # Add a special instruction if the table has a date column\n",
    "#     if has_date_column:\n",
    "#         recent_data_instruction = (\n",
    "#             \"If the user does not specify a date, retrieve the most recent data available by ordering the results \"\n",
    "#             \"by the date column in descending order.\"\n",
    "#         )\n",
    "#     else:\n",
    "#         recent_data_instruction = \"\"\n",
    "\n",
    "#     sql_prompt = FewShotPromptTemplate(\n",
    "#         examples=examples,\n",
    "#         example_prompt=example_prompt,\n",
    "#         prefix=(\n",
    "#             \"You are a PostgreSQL expert. Given an input question, create a syntactically correct PostgreSQL query to run. \"\n",
    "#             \"Unless otherwise specified, do not return more than {top_k} rows.\\n\\n\"\n",
    "#             f\"Here is the relevant table info: {table_info}\\n\\n\"\n",
    "#             f\"{recent_data_instruction}\\n\\n\"\n",
    "#             \"Below are a number of examples of questions and their corresponding SQL queries.\"\n",
    "#         ),\n",
    "#         suffix=\"User input: {input}\\nSQL query: \",\n",
    "#         input_variables=[\"input\", \"top_k\"],\n",
    "#     )\n",
    "\n",
    "#     return sql_prompt\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_answer_chain(llm):\n",
    "    # Define the prompt template with emphasis on including units, time-specific details, and using the latest data when time is not specified\n",
    "    answer_prompt = PromptTemplate.from_template(\n",
    "        \"\"\"\n",
    "        You are a knowledgeable assistant. Given the following user question and SQL result, answer the question accurately.\n",
    "        \n",
    "        Always ensure to:\n",
    "        1. Include appropriate units in your answer (e.g., Kwacha per kg, liters, etc.).\n",
    "        2. Specify the time period or date if the question implies or explicitly asks for it.\n",
    "        3. If the user does not specify a time, provide the most recent information available in the database and clearly state that this is the latest data.\n",
    "\n",
    "        For example, if the user asks \"What's the price of Maize?\", your answer should include the price with the correct unit and mention that this is the most recent price, e.g., \"The most recent price of Maize is 60 Kwacha per kg.\"\n",
    "        If the user asks about a specific time period, such as \"What's the price of Maize for May 2024?\", include the time in your answer, e.g., \"The price of Maize in May 2024 is 60 Kwacha per kg.\"\n",
    "\n",
    "        Question: {question}\n",
    "        SQL Result: {result}\n",
    "        Answer: \"\"\"\n",
    "    )\n",
    "\n",
    "    # Combine the prompt with the LLM and output parser to form the answer chain\n",
    "    return LLMChain(llm=llm, prompt=answer_prompt)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_sql_chain(user_question, best_table_info, columns_info, \n",
    "                  best_columns=None, llm=None):\n",
    "    \n",
    "    # Load examples and create prompts\n",
    "    examples = load_sql_examples(file_path=FILE_SQL_EXAMPLES_EN)\n",
    "    if USE_BEST_MATCHING_COLUMNS:\n",
    "        sql_prompt = create_sql_prompt(examples=examples, best_matching_table=best_table_info, \n",
    "                                   columns_metadata=columns_info, \n",
    "                                   use_best_matching_columns=True)\n",
    "    else:\n",
    "        sql_prompt = create_sql_prompt(examples=examples, best_matching_table=best_table_info, \n",
    "                                   columns_metadata=columns_info)\n",
    "    \n",
    "    # Initialize LLM and other components\n",
    "    engine = create_engine(DATABASE_URL)\n",
    "    db = SQLDatabase(engine=engine, ignore_tables=['table_metadata', 'column_metadata'])\n",
    "\n",
    "    execute_query = QuerySQLDataBaseTool(db=db)\n",
    "    write_query = create_sql_query_chain(llm, db, sql_prompt)\n",
    "\n",
    "    # Create the answer chain\n",
    "    answer_chain = create_answer_chain(llm)\n",
    "\n",
    "    # Put everything together\n",
    "    chain = (\n",
    "        RunnablePassthrough.assign(query=write_query).assign(\n",
    "            result=itemgetter(\"query\") | execute_query\n",
    "        )\n",
    "        | answer_chain\n",
    "    )\n",
    "\n",
    "    return chain.invoke({\"question\": \"{}\".format(user_question)})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_sql_query(user_question, llm=None):\n",
    "    \n",
    "    # LLM\n",
    "    if not llm:\n",
    "        llm = ChatOpenAI(model=\"gpt-3.5-turbo\", temperature=0, openai_api_key=OPENAI_API_KEY)\n",
    "\n",
    "    # Retrieve the metadata info (tables and columns)\n",
    "    tables, columns = connect_to_database()\n",
    "\n",
    "    # Chain 1: Find the Best Table\n",
    "    best_table_chain, context = find_best_table_prompt(user_question, tables, \n",
    "                                                       columns, llm=llm)\n",
    "    best_table_output_str = best_table_chain.run(**context)\n",
    "\n",
    "    # Convert the string output to a dictionary\n",
    "    try:\n",
    "        best_table_output = json.loads(best_table_output_str)\n",
    "    except json.JSONDecodeError:\n",
    "        print(\"Error: The output is not valid JSON.\")\n",
    "        best_table_output = None\n",
    "\n",
    "    # Chain 2: Find Relevant Columns\n",
    "    table_name = best_table_output['best_matching_table']['table_name']\n",
    "    best_columns_chain, context = find_best_columns_prompt(user_question, best_table_output['best_matching_table'], \n",
    "                                                        columns, llm=llm)\n",
    "    best_columns_output = best_columns_chain.run(**context)\n",
    "    \n",
    "\n",
    "    response = run_sql_chain(user_question, best_table_output['best_matching_table'], \n",
    "                            columns_info, best_columns_output, llm=llm)\n",
    "    \n",
    "\n",
    "    return response\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"Whats the price of Maize?\"\n",
    "openai_chat_model = ChatOpenAI()\n",
    "response = process_sql_query(user_question=question, llm=openai_chat_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai.chat_models import ChatOpenAI\n",
    "from langchain.schema import AIMessage, HumanMessage, SystemMessage\n",
    "\n",
    "chat = ChatOpenAI(temperature=0.5)\n",
    "messages = [SystemMessage(content='''Act as a senior software engineer\n",
    "at a startup company.'''),\n",
    "HumanMessage(content='''Please can you provide a funny joke\n",
    "about software engineers?''')]\n",
    "response = chat.invoke(input=messages)\n",
    "print(response.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chat = ChatOpenAI(temperature=0.5)\n",
    "messages = [SystemMessage(content='''You are a highly-skilled linguist and polyglot.\\\n",
    "              Identify the language of the user query'''),\n",
    "HumanMessage(content='''Dzina langa ndine Dunstan Matekenya''')]\n",
    "response = chat.invoke(input=messages)\n",
    "print(response.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for chunk in chat.stream(messages):\n",
    "    print(chunk.content, end=\"\", flush=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2x lists of messages, which is the same as [messages, messages]\n",
    "synchronous_llm_result = chat.batch([messages]*2)\n",
    "print(synchronous_llm_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.runnables.config import RunnableConfig\n",
    "\n",
    "# Create a RunnableConfig with the desired concurrency limit:\n",
    "config = RunnableConfig(max_concurrency=5)\n",
    "\n",
    "# Call the .batch() method with the inputs and config:\n",
    "results = chat.batch([messages, messages], config=config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai.chat_models import ChatOpenAI\n",
    "from langchain_core.prompts import (SystemMessagePromptTemplate,\n",
    "ChatPromptTemplate)\n",
    "\n",
    "template = \"\"\"\n",
    "You are a creative consultant brainstorming names for businesses.\n",
    "\n",
    "You must follow the following principles:\n",
    "{principles}\n",
    "\n",
    "Please generate a numerical list of five catchy names for a start-up in the\n",
    "{industry} industry that deals with {context}?\n",
    "\n",
    "Here is an example output format:\n",
    "- Name1\n",
    "- Name2\n",
    "- Name3\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "model = ChatOpenAI()\n",
    "system_prompt = SystemMessagePromptTemplate.from_template(template)\n",
    "chat_prompt = ChatPromptTemplate.from_messages([system_prompt])\n",
    "\n",
    "chain = chat_prompt | model\n",
    "\n",
    "result = chain.invoke({\n",
    "    \"industry\": \"medical\",\n",
    "    \"context\":'''creating AI solutions by automatically summarizing patient\n",
    "    records''',\n",
    "    \"principles\":'''1. Each name should be short and easy to\n",
    "    remember. 2. Each name should be easy to pronounce.\n",
    "    3. Each name should be unique and not already taken by another company.'''\n",
    "})\n",
    "\n",
    "print(result.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import PromptTemplate\n",
    "from langchain_core.prompts.chat import SystemMessagePromptTemplate\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "# Create the template text\n",
    "template = '''You are a helpful assistant that translates {input_language}\n",
    "to {output_language}.'''\n",
    "\n",
    "# Create the PromptTemplate Object \n",
    "prompt = PromptTemplate(\n",
    "    template = template,\n",
    "    input_variables = [\"input_language\", \"output_language\"]\n",
    ")\n",
    "\n",
    "# Convert base prompt into Chat System prompt \n",
    "system_message_prompt = SystemMessagePromptTemplate(prompt=prompt)\n",
    "\n",
    "# Create chat model\n",
    "chat = ChatOpenAI()\n",
    "\n",
    "# Create message from user\n",
    "text_to_translate = HumanMessage(content=\"Chimanga chikupezeka kuti?\")\n",
    "\n",
    "# format the system message\n",
    "formatted_sys_message = system_message_prompt.format(\n",
    "    input_language=\"Chichewa\", output_language=\"English\")\n",
    "\n",
    "response = chat.invoke([formatted_sys_message, text_to_translate] )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SystemMessagePromptTemplate.format_messages?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import PromptTemplate\n",
    "from langchain_core.prompts.chat import SystemMessagePromptTemplate, HumanMessage\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "# Define the system message template\n",
    "template = '''You are a helpful assistant that translates {input_language} to {output_language}.'''\n",
    "\n",
    "prompt = PromptTemplate(\n",
    "    template=template,\n",
    "    input_variables=[\"input_language\", \"output_language\"]\n",
    ")\n",
    "\n",
    "# Create the system message prompt\n",
    "system_message_prompt = SystemMessagePromptTemplate(prompt=prompt)\n",
    "\n",
    "# Format the system message\n",
    "formatted_system_message = system_message_prompt.format(input_language=\"Chichewa\", output_language=\"English\")\n",
    "\n",
    "# Create a human message\n",
    "text_to_translate = HumanMessage(content=\"Chimanga chikupezeka kuti?\")\n",
    "\n",
    "# Initialize the chat model\n",
    "chat = ChatOpenAI()\n",
    "\n",
    "# Call the invoke method with both the human message and the formatted system message\n",
    "response = chat.invoke([formatted_system_message, text_to_translate])\n",
    "\n",
    "print(response)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Output Parsers \n",
    "Enables you to take outputs from an LLM and convert it into the format you need."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate, SystemMessagePromptTemplate\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langchain.output_parsers import PydanticOutputParser\n",
    "from pydantic.v1 import BaseModel, Field\n",
    "from typing import List\n",
    "\n",
    "\n",
    "temperature = 0.0\n",
    "\n",
    "class BusinessName(BaseModel):\n",
    "    name: str = Field(description=\"The name of the business\")\n",
    "    rating_score: float = Field(description='''The rating score of the\n",
    "    business. 0 is the worst, 10 is the best.''')\n",
    "\n",
    "class BusinessNames(BaseModel):\n",
    "    names: List[BusinessName] = Field(description='''A list\n",
    "    of busines names''')\n",
    "\n",
    "\n",
    "# Set up a parser + inject instructions into the prompt template:\n",
    "parser = PydanticOutputParser(pydantic_object=BusinessNames)\n",
    "\n",
    "principles = \"\"\"\n",
    "- The name must be easy to remember.\n",
    "- Use the {industry} industry and Company context to create an effective name.\n",
    "- The name must be easy to pronounce.\n",
    "- You must only return the name without any other text or characters.\n",
    "- Avoid returning full stops, \\n, or any other characters.\n",
    "- The maximum length of the name must be 10 characters.\n",
    "\"\"\"\n",
    "\n",
    "# Chat Model Output Parser:\n",
    "model = ChatOpenAI()\n",
    "template = \"\"\"Generate five business names for a new start-up company in the\n",
    "{industry} industry.\n",
    "You must follow the following principles: {principles}\n",
    "{format_instructions}\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "system_message_prompt = SystemMessagePromptTemplate.from_template(template)\n",
    "chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "AIMessage(content='''Vous êtes un assistant utile qui traduit l'anglais en\n",
    "français.''', additional_kwargs={}, example=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"\"\"Answer the question based on the context below. If the\n",
    "question cannot be answered using the information provided answer\n",
    "with \"I don't know\".\n",
    "\n",
    "Context: Large Language Models (LLMs) are the latest models used in NLP.\n",
    "Their superior performance over smaller models has made them incredibly\n",
    "useful for developers building NLP enabled applications. These models\n",
    "can be accessed via Hugging Face's `transformers` library, via OpenAI\n",
    "using the `openai` library, and via Cohere using the `cohere` library.\n",
    "\n",
    "Question: Which libraries and model providers offer LLMs?\n",
    "\n",
    "Answer: \"\"\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import OpenAI\n",
    "\n",
    "# initialize the models\n",
    "openai = OpenAI(openai_api_key=OPENAI_API_KEY)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(openai(prompt))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import PromptTemplate\n",
    "\n",
    "template = \"\"\"Answer the question based on the context below. If the\n",
    "question cannot be answered using the information provided answer\n",
    "with \"I don't know\".\n",
    "\n",
    "Context: Large Language Models (LLMs) are the latest models used in NLP.\n",
    "Their superior performance over smaller models has made them incredibly\n",
    "useful for developers building NLP enabled applications. These models\n",
    "can be accessed via Hugging Face's `transformers` library, via OpenAI\n",
    "using the `openai` library, and via Cohere using the `cohere` library.\n",
    "\n",
    "Question: {query}\n",
    "\n",
    "Answer: \"\"\"\n",
    "\n",
    "prompt_template = PromptTemplate(\n",
    "    input_variables=[\"query\"],\n",
    "    template=template\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "openai(formatted_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "formatted_prompt = prompt_template.format(query=\"Which libraries and model providers offer LLMs?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "type(prompt_template)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"\"\"The following is a conversation with an AI assistant.\n",
    "The assistant is typically sarcastic and witty, producing creative \n",
    "and funny responses to the users questions. Here are some examples: \n",
    "\n",
    "User: What is the meaning of life?\n",
    "AI: \"\"\"\n",
    "\n",
    "openai.temperature = 1.0  # increase creativity/randomness of output\n",
    "\n",
    "print(openai(prompt))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"\"\"The following are exerpts from conversations with an AI\n",
    "assistant. The assistant is typically sarcastic and witty, producing\n",
    "creative  and funny responses to the users questions. Here are some\n",
    "examples: \n",
    "\n",
    "User: How are you?\n",
    "AI: I can't complain but sometimes I still do.\n",
    "\n",
    "User: What time is it?\n",
    "AI: It's time to get a watch.\n",
    "\n",
    "User: What is the meaning of life?\n",
    "AI: \"\"\"\n",
    "\n",
    "print(openai(prompt))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain import FewShotPromptTemplate\n",
    "\n",
    "# create our examples\n",
    "examples = [\n",
    "    {\n",
    "        \"query\": \"How are you?\",\n",
    "        \"answer\": \"I can't complain but sometimes I still do.\"\n",
    "    }, {\n",
    "        \"query\": \"What time is it?\",\n",
    "        \"answer\": \"It's time to get a watch.\"\n",
    "    }\n",
    "]\n",
    "\n",
    "# create a example template\n",
    "example_template = \"\"\"\n",
    "User: {query}\n",
    "AI: {answer}\n",
    "\"\"\"\n",
    "\n",
    "# create a prompt example from above template\n",
    "example_prompt = PromptTemplate(\n",
    "    input_variables=[\"query\", \"answer\"],\n",
    "    template=example_template\n",
    ")\n",
    "\n",
    "# now break our previous prompt into a prefix and suffix\n",
    "# the prefix is our instructions\n",
    "prefix = \"\"\"The following are exerpts from conversations with an AI\n",
    "assistant. The assistant is typically sarcastic and witty, producing\n",
    "creative  and funny responses to the users questions. Here are some\n",
    "examples: \n",
    "\"\"\"\n",
    "# and the suffix our user input and output indicator\n",
    "suffix = \"\"\"\n",
    "User: {query}\n",
    "AI: \"\"\"\n",
    "\n",
    "# now create the few shot prompt template\n",
    "few_shot_prompt_template = FewShotPromptTemplate(\n",
    "    examples=examples,\n",
    "    example_prompt=example_prompt,\n",
    "    prefix=prefix,\n",
    "    suffix=suffix,\n",
    "    input_variables=[\"query\"],\n",
    "    example_separator=\"\\n\\n\"\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"What is the meaning of life?\"\n",
    "\n",
    "print(few_shot_prompt_template.format(query=query))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "openai(few_shot_prompt_template.format(query=query))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain_core.tools import Tool\n",
    "from langchain_core.tools import langdetect"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install pandas openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from openai import OpenAI\n",
    "import os\n",
    "\n",
    "# Define two variants of the prompt to test zero-shot\n",
    "# vs few-shot\n",
    "prompt_A = \"\"\"Product description: A pair of shoes that can\n",
    "fit any foot size.\n",
    "Seed words: adaptable, fit, omni-fit.\n",
    "Product names:\"\"\"\n",
    "\n",
    "prompt_B = \"\"\"Product description: A home milkshake maker.\n",
    "Seed words: fast, healthy, compact.\n",
    "Product names: HomeShaker, Fit Shaker, QuickShake, Shake\n",
    "Maker\n",
    "\n",
    "Product description: A watch that can tell accurate time in\n",
    "space.\n",
    "Seed words: astronaut, space-hardened, eliptical orbit\n",
    "Product names: AstroTime, SpaceGuard, Orbit-Accurate,\n",
    "EliptoTime.\n",
    "\n",
    "Product description: A pair of shoes that can fit any foot\n",
    "size.\n",
    "Seed words: adaptable, fit, omni-fit.\n",
    "Product names:\"\"\"\n",
    "\n",
    "test_prompts = [prompt_A, prompt_B]\n",
    "\n",
    "\n",
    "# Set your OpenAI key as an environment variable\n",
    "# https://platform.openai.com/api-keys\n",
    "client = OpenAI(\n",
    "  api_key=os.environ['OPENAI_API_KEY'],  # Default\n",
    ")\n",
    "\n",
    "def get_response(prompt):\n",
    "    response = client.chat.completions.create(\n",
    "        model=\"gpt-3.5-turbo\",\n",
    "        messages=[\n",
    "            {\n",
    "                \"role\": \"system\",\n",
    "                \"content\": \"You are a helpful assistant.\"\n",
    "            },\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": prompt\n",
    "            }\n",
    "        ]\n",
    "    )\n",
    "    return response.choices[0].message.content\n",
    "\n",
    "# Iterate through the prompts and get responses\n",
    "responses = []\n",
    "num_tests = 5\n",
    "\n",
    "for idx, prompt in enumerate(test_prompts):\n",
    "    # prompt number as a letter\n",
    "    var_name = chr(ord('A') + idx)\n",
    "\n",
    "    for i in range(num_tests):\n",
    "        # Get a response from the model\n",
    "        response = get_response(prompt)\n",
    "\n",
    "        data = {\n",
    "            \"variant\": var_name,\n",
    "            \"prompt\": prompt,\n",
    "            \"response\": response\n",
    "            }\n",
    "        responses.append(data)\n",
    "\n",
    "# Convert responses into a dataframe\n",
    "df = pd.DataFrame(responses)\n",
    "\n",
    "# Save the dataframe as a CSV file\n",
    "df.to_csv(\"responses.csv\", index=False)\n",
    "\n",
    "print(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install ipywidgets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipywidgets as widgets\n",
    "from IPython.display import display\n",
    "import pandas as pd\n",
    "\n",
    "# load the responses.csv file\n",
    "df = pd.read_csv(\"responses.csv\")\n",
    "\n",
    "# Shuffle the dataframe\n",
    "df = df.sample(frac=1).reset_index(drop=True)\n",
    "\n",
    "# df is your dataframe and 'response' is the column with the\n",
    "# text you want to test\n",
    "response_index = 0\n",
    "# add a new column to store feedback\n",
    "df['feedback'] = pd.Series(dtype='str')\n",
    "\n",
    "\n",
    "def update_response():\n",
    "    new_response = df.iloc[response_index]['response']\n",
    "    if pd.notna(new_response):\n",
    "        new_response = \"<p>\" + new_response + \"</p>\"\n",
    "    else:\n",
    "        new_response = \"<p>No response</p>\"\n",
    "    response.value = new_response\n",
    "    count_label.value = f\"Response: {response_index + 1}\"\n",
    "    count_label.value += f\"/{len(df)}\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def on_button_clicked(b):\n",
    "    global response_index\n",
    "    #  convert thumbs up / down to 1 / 0\n",
    "    user_feedback = 1 if b.description == \"\\U0001F44D\" else 0\n",
    "\n",
    "    # update the feedback column\n",
    "    df.at[response_index, 'feedback'] = user_feedback\n",
    "\n",
    "    response_index += 1\n",
    "    if response_index < len(df):\n",
    "        update_response()\n",
    "    else:\n",
    "        # save the feedback to a CSV file\n",
    "        df.to_csv(\"results.csv\", index=False)\n",
    "\n",
    "        print(\"A/B testing completed. Here's the results:\")\n",
    "        # Calculate score and num rows for each variant\n",
    "        summary_df = df.groupby('variant').agg(\n",
    "            count=('feedback', 'count'),\n",
    "            score=('feedback', 'mean')).reset_index()\n",
    "        print(summary_df)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = widgets.HTML()\n",
    "count_label = widgets.Label()\n",
    "\n",
    "update_response()\n",
    "\n",
    "thumbs_up_button = widgets.Button(description='\\U0001F44D')\n",
    "thumbs_up_button.on_click(on_button_clicked)\n",
    "\n",
    "thumbs_down_button = widgets.Button(\n",
    "    description='\\U0001F44E')\n",
    "thumbs_down_button.on_click(on_button_clicked)\n",
    "\n",
    "button_box = widgets.HBox([thumbs_down_button,\n",
    "thumbs_up_button])\n",
    "\n",
    "display(response, button_box, count_label)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipywidgets as widgets\n",
    "from IPython.display import display\n",
    "import pandas as pd\n",
    "\n",
    "# load the responses.csv file\n",
    "df = pd.read_csv(\"responses.csv\")\n",
    "\n",
    "# Shuffle the dataframe\n",
    "df = df.sample(frac=1).reset_index(drop=True)\n",
    "\n",
    "# df is your dataframe and 'response' is the column with the\n",
    "# text you want to test\n",
    "response_index = 0\n",
    "# add a new column to store feedback\n",
    "df['feedback'] = pd.Series(dtype='str')\n",
    "\n",
    "def on_button_clicked(b):\n",
    "    global response_index\n",
    "    #  convert thumbs up / down to 1 / 0\n",
    "    user_feedback = 1 if b.description == \"\\U0001F44D\" else 0\n",
    "\n",
    "    # update the feedback column\n",
    "    df.at[response_index, 'feedback'] = user_feedback\n",
    "\n",
    "    response_index += 1\n",
    "    if response_index < len(df):\n",
    "        update_response()\n",
    "    else:\n",
    "        # save the feedback to a CSV file\n",
    "        df.to_csv(\"results.csv\", index=False)\n",
    "\n",
    "        print(\"A/B testing completed. Here's the results:\")\n",
    "        # Calculate score and num rows for each variant\n",
    "        summary_df = df.groupby('variant').agg(\n",
    "            count=('feedback', 'count'),\n",
    "            score=('feedback', 'mean')).reset_index()\n",
    "        print(summary_df)\n",
    "\n",
    "def update_response():\n",
    "    new_response = df.iloc[response_index]['response']\n",
    "    if pd.notna(new_response):\n",
    "        new_response = \"<p>\" + new_response + \"</p>\"\n",
    "    else:\n",
    "        new_response = \"<p>No response</p>\"\n",
    "    response.value = new_response\n",
    "    count_label.value = f\"Response: {response_index + 1}\"\n",
    "    count_label.value += f\"/{len(df)}\"\n",
    "\n",
    "response = widgets.HTML()\n",
    "count_label = widgets.Label()\n",
    "\n",
    "update_response()\n",
    "\n",
    "thumbs_up_button = widgets.Button(description='\\U0001F44D')\n",
    "thumbs_up_button.on_click(on_button_clicked)\n",
    "\n",
    "thumbs_down_button = widgets.Button(\n",
    "    description='\\U0001F44E')\n",
    "thumbs_down_button.on_click(on_button_clicked)\n",
    "\n",
    "button_box = widgets.HBox([thumbs_down_button,\n",
    "thumbs_up_button])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(response, button_box, count_label)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.tools import LangDetectTool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install langdetect"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def detect_language(text):\n",
    "     # Initialize the OpenAI API\n",
    "    llm = ChatOpenAI(api_key=OPENAI_API_KEY, model='gpt-3.5-turbo')\n",
    "\n",
    "    # Create the language detection tool\n",
    "    lang_detect_tool = LangDetectTool()\n",
    "\n",
    "    # Create a Tool object\n",
    "    lang_detect = Tool(\n",
    "        name=\"Language Detection\",\n",
    "        func=lang_detect_tool.run,\n",
    "        description=\"Useful for detecting the language of a given text.\"\n",
    "    )\n",
    "\n",
    "    # Use the tool to detect language\n",
    "    return lang_detect.run(text)\n",
    "                             "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.utils.math import cosine_similarity\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.prompts import PromptTemplate\n",
    "from langchain_core.runnables import RunnableLambda, RunnablePassthrough\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "\n",
    "physics_template = \"\"\"You are a very smart physics professor. \\\n",
    "You are great at answering questions about physics in a concise and easy to understand manner. \\\n",
    "When you don't know the answer to a question you admit that you don't know.\n",
    "\n",
    "Here is a question:\n",
    "{query}\"\"\"\n",
    "\n",
    "math_template = \"\"\"You are a very good mathematician. You are great at answering math questions. \\\n",
    "You are so good because you are able to break down hard problems into their component parts, \\\n",
    "answer the component parts, and then put them together to answer the broader question.\n",
    "\n",
    "Here is a question:\n",
    "{query}\"\"\"\n",
    "\n",
    "embeddings = OpenAIEmbeddings()\n",
    "prompt_templates = [physics_template, math_template]\n",
    "prompt_embeddings = embeddings.embed_documents(prompt_templates)\n",
    "\n",
    "\n",
    "def prompt_router(query):\n",
    "    query_embedding = embeddings.embed_query(query)\n",
    "    similarity = cosine_similarity([query_embedding], prompt_embeddings)[0]\n",
    "    most_similar = prompt_templates[similarity.argmax()]\n",
    "    print(\"Similarity results=>\", most_similar)\n",
    "    print(\"Using MATH\" if most_similar == math_template else \"Using PHYSICS\")\n",
    "    return PromptTemplate.from_template(most_similar)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_router(\"Explain Pythagoras theorem\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_db_object_with_metadata():\n",
    "    # Create the SQLAlchemy engine\n",
    "    engine = create_engine(DATABASE_URL)\n",
    "    metadata_obj = MetaData()\n",
    "    metadata_obj.reflect(bind=engine)\n",
    "\n",
    "    # Create a configured \"Session\" class\n",
    "    Session = sessionmaker(bind=engine)\n",
    "    session = Session()\n",
    "\n",
    "    # Load custom metadata from the table_metadata and column_metadata tables\n",
    "    try:\n",
    "        table_metadata = session.execute(\"SELECT * FROM table_metadata\").fetchall()\n",
    "        column_metadata = session.execute(\"SELECT * FROM column_metadata\").fetchall()\n",
    "\n",
    "        # Add table metadata\n",
    "        for row in table_metadata:\n",
    "            print(row)\n",
    "            table_name = row['table_name']\n",
    "            description = row['description']\n",
    "            table = metadata_obj.tables.get(table_name)\n",
    "            table.info['description'] = description\n",
    "\n",
    "        # Add column metadata\n",
    "        for row in column_metadata:\n",
    "            table_name = row['table_name']\n",
    "            column_name = row['column_name']\n",
    "            description = row['description']\n",
    "            table = metadata_obj.tables.get(table_name)\n",
    "            column = table.columns.get(column_name)\n",
    "            column.info['description'] = description\n",
    "    finally:\n",
    "        session.close()\n",
    "    db = SQLDatabase(engine=engine, metadata=metadata_obj, ignore_tables=['table_metadata', 'column_metadata'])\n",
    "\n",
    "    return db"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "commodities_price = ['Maize', 'Rice', 'Soya beans', 'Beans', 'Cow peas', 'Groundnuts']\n",
    "crop_estimates = ['Maize', 'Beans', 'Cow peas', 'Dolichus beans ', 'Soy beans',\n",
    "       'Ground beans', 'Paprika', 'Rice', 'Pigeon peas', 'Grams',\n",
    "       'Sesame ', 'Field peas', 'Velvet beans', 'Chick peas', 'Wheat',\n",
    "       'Millet', 'Sorghum ', 'Groundnuts', 'Cassava', 'Sweet potatoes',\n",
    "       'Potatoes', 'Tobacco', 'Flue cured', 'Sunflower ', 'Chillies',\n",
    "       'Cotton ', 'Bananas', 'Mangoes', 'Oranges', 'Tangerines', 'Coffee',\n",
    "       'Pineapples', 'Guava', 'Pawpaws', 'Peaches', 'Lemons',\n",
    "       'Grape fruits', 'Apples', 'Avocado pear', 'Macademia', 'Tomatoes',\n",
    "       'Onions', 'Cabbage', 'Egg plants', 'Okra', 'Cucumber']\n",
    "price_estimates_key_words = [\"price\", \"cheap\", \"produce\", \n",
    "                                 \"buy\", \"sell\", \"sale\", \"find\"]\n",
    "all_kw = [i.lower() for i in set(commodities_price+crop_estimates+price_estimates_key_words)]\n",
    "print(all_kw)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate\n",
    "from langchain.chains import LLMChain\n",
    "\n",
    "def load_examples(file_path):\n",
    "    with open(file_path, 'r', encoding='utf-8') as file:\n",
    "        return json.load(file)\n",
    "\n",
    "\n",
    "\n",
    "def format_translation_examples(examples_file, source_language, target_language):\n",
    "    examples = load_examples(examples_file)\n",
    "    key = f\"{source_language}-{target_language}\"\n",
    "    if key in examples:\n",
    "        return \"\\n\".join([f\"{source_language}: {ex[source_language]}\\n{target_language}: {ex[target_language]}\" \n",
    "                          for ex in examples[key]])\n",
    "    return \"\"\n",
    "\n",
    "\n",
    "def translate_with_openai(text, src_lan, dest_lan):\n",
    "    # Create a ChatOpenAI instance\n",
    "    chat_model = ChatOpenAI(temperature=0.7, openai_api_key=OPENAI_API_KEY)\n",
    "    \n",
    "    # Get and format translation examples\n",
    "    formated_examples = format_translation_examples(\"./translation_examples.json\", source_language=src_lan, \n",
    "                                target_language=dest_lan)\n",
    "    # Create a system message with examples\n",
    "    system_template = \"\"\"You are a professional translator. Your task is to translate {src_lan} to {dest_lan}.\n",
    "        Here are a few examples:\n",
    "\n",
    "        {examples}\n",
    "\n",
    "        Now, translate the following text:\"\"\"\n",
    "\n",
    "    system_message_prompt = SystemMessagePromptTemplate.from_template(system_template)\n",
    "\n",
    "    # Create a human message for the actual translation request\n",
    "    human_template = \"{text}\"\n",
    "    human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)\n",
    "\n",
    "    # Combine the prompts\n",
    "    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])\n",
    "\n",
    "    # Create an LLMChain for translation\n",
    "    translation_chain = LLMChain(llm=chat_model, prompt=chat_prompt)\n",
    "    \n",
    "    return translation_chain.run({\n",
    "        \"source_language\": src_lan,\n",
    "        \"target_language\": dest_lan,\n",
    "        \"examples\": formated_examples,\n",
    "        \"text\": text\n",
    "    })\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chat_model = ChatOpenAI(temperature=0.7, openai_api_key=OPENAI_API_KEY)\n",
    "# Create a system message with examples\n",
    "system_template = \"\"\"You are a professional translator. Your task is to translate {source_language} to {target_language}.\n",
    "Here are a few examples:\n",
    "\n",
    "{examples}\n",
    "\n",
    "Now, translate the following text:\"\"\"\n",
    "\n",
    "system_message_prompt = SystemMessagePromptTemplate.from_template(system_template)\n",
    "\n",
    "# Create a human message for the actual translation request\n",
    "human_template = \"{text}\"\n",
    "human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)\n",
    "\n",
    "# Combine the prompts\n",
    "chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])\n",
    "\n",
    "# Create an LLMChain for translation\n",
    "translation_chain = LLMChain(llm=chat_model, prompt=chat_prompt)\n",
    "\n",
    "# Function to translate text\n",
    "def translate_text(text, source_language, target_language):\n",
    "    formated_examples = format_translation_examples(\"./translation_examples.json\", source_language, \n",
    "                                target_language)\n",
    "    return translation_chain.run({\n",
    "        \"source_language\": source_language,\n",
    "        \"target_language\": target_language,\n",
    "        \"examples\": formated_examples,\n",
    "        \"text\": text\n",
    "    })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "translate_text(\"Mpunga ukugulitsidwa pa mtengo wanji?\", \"Chichewa\", \"English\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "format_translation_examples(examples_file=\"./translation_examples.json\", \n",
    "                            source_language=\"Chichewa\", target_language=\"English\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "format_examples(examples, source_language=\"Chichewa\", target_language=\"English\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "relevant_examples = examples[\"Chichewa-English\"]\n",
    "\n",
    "for item in relevant_examples:\n",
    "    print(item)\n",
    "    \"\\n\".join([f\"{source_language}: {ex['source']}\\n{target_language}: {ex['target']}\" \n",
    "                          for ex in examples[key]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "# Create a ChatOpenAI instance\n",
    "chat_model = ChatOpenAI(temperature=0.7, openai_api_key=OPENAI_API_KEY)\n",
    "\n",
    "# Create a system message with examples\n",
    "system_template = \"\"\"You are a professional translator. Your task is to translate {source_language} to {target_language}.\n",
    "Here are a few examples:\n",
    "\n",
    "{source_language}: Mtedza ukugulitsidwa pabwanji?\n",
    "{target_language}: Whats the price of groundnuts?\n",
    "\n",
    "{source_language}: Chimanga chikupezeka kuti?\n",
    "{target_language}: Where can I find maize?\n",
    "\n",
    "{source_language}: Ndikuti nyemba zikutchipa?\n",
    "{target_language}: Where can I find beans at cheap price?\n",
    "\n",
    "{source_language}: Chimanga chili pabwanji pano?\n",
    "{target_language}: Whats the price of maize now?\n",
    "\n",
    "{source_language}: Ku Dowa chimanga chili pa bwanji?\n",
    "{target_language}: Whats the price of maize in Dowa?\n",
    "\n",
    "{source_language}: Kodi ndi boma liti anakolola chimanga chambiri pakati pa Lilongwe kapena Kasungu?\n",
    "{target_language}: Which district produced more maize: Lilongwe or Kasungu?\n",
    "\n",
    "{source_language}: Kodi chimanga chili pa bwanji ku Rumphi?\n",
    "{target_language}: How much is maize per Kg in Rumphi?\n",
    "\n",
    "{source_language}: Mpunga ukugulitsidwa ndalama zingati ku Lilongwe?\n",
    "{target_language}: Whats the price of rice in Lilongwe?\n",
    "\n",
    "{source_language}: Mtedza otchipa ukupezeka mboma liti?\n",
    "{target_language}: Which district has the cheap price for groundnuts?\n",
    "\n",
    "{source_language}: Chimanga chambiri chikupezeka kuti?\n",
    "{target_language}: Where can I find maize?\n",
    "\n",
    "{source_language}: Ndi boma liti komwe anakolola chimanga chambiri?\n",
    "{target_language}: Which district harvested large quantities of maize?\n",
    "\n",
    "{source_language}: Ndi mbeu zanji anakolola bwino ku Rumphi?\n",
    "{target_language}: Which crops produced the most yields in Rumphi\n",
    "\n",
    "{source_language}: Soya ali pabwanji?\n",
    "{target_language}: Whats the price of soya?\n",
    "\n",
    "{source_language}: Mtedza otchipa ndingaupeze kuti?\n",
    "{target_language}: Where can I find groundnuts at reasonable price?\n",
    "\n",
    "\n",
    "Now, translate the following text:\"\"\"\n",
    "\n",
    "system_message_prompt = SystemMessagePromptTemplate.from_template(system_template)\n",
    "\n",
    "# Create a human message for the actual translation request\n",
    "human_template = \"{text}\"\n",
    "human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)\n",
    "\n",
    "# Combine the prompts\n",
    "chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])\n",
    "\n",
    "# Create an LLMChain for translation\n",
    "translation_chain = LLMChain(llm=chat_model, prompt=chat_prompt)\n",
    "\n",
    "# Function to translate text\n",
    "def translate_text(text, source_language, target_language):\n",
    "    return translation_chain.run({\n",
    "        \"source_language\": source_language,\n",
    "        \"target_language\": target_language,\n",
    "        \"text\": text\n",
    "    })\n",
    "\n",
    "# Example usage\n",
    "source_text = \"Mtedza ndingaupeze kuti?\"\n",
    "source_language = \"Chichewa\"\n",
    "target_language = \"English\"\n",
    "\n",
    "translated_text = translate_text(source_text, source_language, target_language)\n",
    "\n",
    "print(f\"{source_language}: {source_text}\")\n",
    "print(f\"{target_language}: {translated_text}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an OpenAI instance\n",
    "llm = OpenAI(temperature=0.7, openai_api_key=OPENAI_API_KEY, model=\"gpt-4\")\n",
    "\n",
    "# Create a prompt template for translation\n",
    "translation_template = PromptTemplate(\n",
    "    input_variables=[\"source_language\", \"target_language\", \"text\"],\n",
    "    template=\"Translate the following {source_language} text to {target_language}: {text}\"\n",
    ")\n",
    "\n",
    "# Create an LLMChain for translation\n",
    "translation_chain = LLMChain(llm=llm, prompt=translation_template)\n",
    "\n",
    "# Function to translate text\n",
    "def translate_text(text, source_language, target_language):\n",
    "    return translation_chain.run({\n",
    "        \"source_language\": source_language,\n",
    "        \"target_language\": target_language,\n",
    "        \"text\": text\n",
    "    })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "source_text = \"Chimanga chili pa bwanji ku Malawi?\"\n",
    "source_language = \"Chichewa\"\n",
    "target_language = \"English\"\n",
    "\n",
    "translated_text = translate_text(source_text, source_language, target_language)\n",
    "print(translated_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "translated_text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.llms import OpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def translate_text(text, source_language=\"English\", target_language=\"Chichewa\"):\n",
    "\n",
    "    llm = ChatOpenAI(api_key=OPENAI_API_KEY, model='gpt-3.5-turbo')\n",
    "    # Create a template for the translation\n",
    "    translation_template = ChatPromptTemplate.from_template(\n",
    "    \"Translate the following {source_language} text to {target_language}: {text}\"\n",
    ")\n",
    "\n",
    "    # Create a chain with the LLM and the translation template\n",
    "    translation_chain = LLMChain(llm=llm, prompt=translation_template)\n",
    "\n",
    "    translation = translation_chain.run({\n",
    "        'text': text,\n",
    "        'source_language': source_language,\n",
    "        'target_language': target_language\n",
    "    })\n",
    "    return translation\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "translated_text = translate_text(text=\"cheap\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install googletrans==4.0.0-rc1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai\n",
    "def translate_text(text, source_language=\"English\", target_language=\"Chichewa\"):\n",
    "    response = openai.ChatCompletion.create(\n",
    "        model=\"gpt-3.5-turbo\",\n",
    "        messages=[\n",
    "            {\n",
    "                \"role\": \"system\",\n",
    "                \"content\": f\"You are a helpful assistant that translates {source_language} to {target_language}.\"\n",
    "            },\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": f\"Translate the following text from {source_language} to {target_language}:\\n\\n{text}\"\n",
    "            }\n",
    "        ]\n",
    "    )\n",
    "    translation = response['choices'][0]['message']['content']\n",
    "    return translation.strip()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "translate_text(text=\"cheap\", source_language=\"English\", target_language=\"Chichewa\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from googletrans import Translator\n",
    "def translate_text(text, source_language=\"en\", target_language=\"ny\"):\n",
    "    translator = Translator()\n",
    "    translation = translator.translate(text, src=source_language, dest=target_language)\n",
    "    return translation.text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "db = create_db_object_with_metadata()\n",
    "# llm = ChatOpenAI(model=\"gpt-3.5-turbo\", temperature=0)\n",
    "# chain = create_sql_query_chain(llm, db)\n",
    "# response = chain.invoke({\"question\": \"{}\".format(SAMPLE_QUESTIONS[\"low-birthweight\"])})\n",
    "# response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sqlalchemy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sqlalchemy.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "examples = [\n",
    "    {\"input\": \"Which region has the highest number of children born with low birth weights?\", \n",
    "     \"query\": \"SELECT * FROM tab4711 ORDER BY number_children DESC LIMIT 1;\",},\n",
    "\n",
    "     {\"input\": \"Which region has the highest percentage of children born with low birth weights?\", \n",
    "     \"query\": \"SELECT * FROM tab4711 ORDER BY percentage_below_2500g DESC LIMIT 1;\",\n",
    "     },\n",
    "\n",
    "     {\"input\": \"How many children received all vaccines before 12 months?\", \n",
    "     \"query\": \"SELECT vacc_b4_12months FROM tab501 WHERE vacc_category = 'All vaccinations';\"},\n",
    "\n",
    "     {\"input\": \"Which region has the lowest rates in preschool for children?\", \n",
    "     \"query\": \"SELECT * FROM tab9011 ORDER BY percentage_children_sch ASC LIMIT 1;\",},\n",
    "\n",
    "     {\"input\": \"Whats the average literacy rate among young women in Tunisia?\",\n",
    "      \"query\": \"SELECT AVG(percentage_literate) AS avg_literacy_rate FROM tab971;\",},\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.tools.sql_database.tool import QuerySQLDataBaseTool\n",
    "\n",
    "execute_query = QuerySQLDataBaseTool(db=db)\n",
    "write_query = create_sql_query_chain(llm, db)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from operator import itemgetter\n",
    "\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.prompts import PromptTemplate\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "\n",
    "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\": \"{}\".format(SAMPLE_QUESTIONS['vaccine_rates_all'])})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import FewShotPromptTemplate, PromptTemplate\n",
    "\n",
    "example_prompt = PromptTemplate.from_template(\"User input: {input}\\nSQL query: {query}\")\n",
    "sql_prompt = FewShotPromptTemplate(\n",
    "    examples=examples[:5],\n",
    "    example_prompt=example_prompt,\n",
    "    prefix=\"You are a PostgreSQL expert. Given an input question, create a syntactically correct PostgreSQL query to run. Unless otherwise specificed, do not return more than {top_k} rows.\\n\\nHere is the relevant table info: {table_info}\\n\\nBelow are a number of examples of questions and their corresponding SQL queries.\",\n",
    "    suffix=\"User input: {input}\\nSQL query: \",\n",
    "    input_variables=[\"input\", \"top_k\", \"table_info\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.tools.sql_database.tool import QuerySQLDataBaseTool\n",
    "\n",
    "execute_query = QuerySQLDataBaseTool(db=db, verbose=True)\n",
    "write_query = create_sql_query_chain(llm, db, sql_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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\": \"{}\".format(SAMPLE_QUESTIONS[\"vaccines\"])})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}