OpenSPP Genie.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1a909571-c9a2-4c2b-9ecb-6c197fcd910c",
   "metadata": {},
   "source": [
    "# OpenSPP Genie"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fad3035-1913-4258-bb7b-5e1997d182c6",
   "metadata": {},
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbc546f8-f21d-48e5-90f6-cbe87d70d8be",
   "metadata": {
    "tags": []
   },
   "source": [
    "## OpenSPP Architecture\n",
    "\n",
    "### High level Architecture\n",
    "\n",
    "![](OpenSPP_architecture_V2_Genie.jpg)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a5116ed-3e01-406a-bb42-e06cb7ada3ab",
   "metadata": {},
   "source": [
    "To run this, you need to:\n",
    "- Install [OpenSPP Docker](https://github.com/openspp/openspp-docker)\n",
    "- "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4faa6f60-9920-4c96-858a-b3951cee10a3",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Setting up the environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "09b9077f-3011-43c9-96f0-1c6f59b7d7b8",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "# update or install the necessary libraries\n",
    "# !pip install --upgrade openai\n",
    "# !pip install --upgrade langchain\n",
    "# !pip install --upgrade python-dotenv\n",
    "# !pip install --upgrade psycopg2\n",
    "# !pip install --upgrade aim\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ae2b0ea0-ba66-41a6-a965-f91e89a216ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai\n",
    "import os\n",
    "import IPython\n",
    "from langchain.llms import OpenAI\n",
    "import os\n",
    "from datetime import datetime\n",
    "from dotenv import load_dotenv\n",
    "load_dotenv()\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = os.getenv(\"OPENAI_API_KEY\")\n",
    "os.environ[\"OPENSPP_DB_LOGIN\"] = os.getenv(\"OPENSPP_DB_LOGIN\")\n",
    "os.environ[\"OPENSPP_DB_PASSWORD\"] = os.getenv(\"OPENSPP_DB_PASSWORD\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "88bf8576-6cfe-4b31-9df3-b8971e43f82e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sqlalchemy import create_engine\n",
    "from langchain.agents import AgentExecutor\n",
    "from langchain import OpenAI, SQLDatabase, SQLDatabaseChain\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain.prompts.prompt import PromptTemplate\n",
    "from langchain.callbacks import AimCallbackHandler, StdOutCallbackHandler\n",
    "from langchain.output_parsers import RetryWithErrorOutputParser\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "371bdd1a-ae9a-4522-a717-dcc9761f5472",
   "metadata": {},
   "outputs": [],
   "source": [
    "session_group = datetime.now().strftime(\"%m.%d.%Y_%H.%M.%S\")\n",
    "aim_callback = AimCallbackHandler(\n",
    "    repo=\".\",\n",
    "    experiment_name=\"OpenSPP Genie\",\n",
    ")\n",
    "\n",
    "callbacks = [StdOutCallbackHandler(), aim_callback]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51b49e2d-ad68-4a05-899c-8510962c4c49",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Initialising langchain\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "72e47e6e-194b-4101-a12b-04cc10ef994b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#\n",
    "engine=create_engine(\"postgresql://{login}:{password}@localhost:5432/devel?sslmode=disable\".format(login=os.environ[\"OPENSPP_DB_LOGIN\"], password=os.environ[\"OPENSPP_DB_PASSWORD\"]))\n",
    "include_tables=[\n",
    "    'res_partner', \n",
    "    'g2p_group_membership', \n",
    "    'g2p_group_membership_kind',\n",
    "    'g2p_group_membership_g2p_group_membership_kind_rel']\n",
    "\n",
    "db = SQLDatabase(engine, include_tables=include_tables) \n",
    "\n",
    "# llm = OpenAI(temperature=0, verbose=True)\n",
    "# llm = ChatOpenAI(temperature=0, model_name=\"gpt-3.5-turbo\")\n",
    "llm = ChatOpenAI(temperature=0, model_name=\"gpt-4\", callbacks=callbacks)\n",
    "llm_sql = ChatOpenAI(temperature=0, model_name=\"gpt-4\", callbacks=callbacks)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5055d3f1-8f8b-4c30-88c0-ea23253a4fcd",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Creating the SQL tools"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "901da5d6-2ca3-4338-bdad-e1adb3b46fb7",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Creating the prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2be8044a-1c58-4a70-b40c-2c3104af56bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_table_info = {\n",
    "    \"res_partner\": \"\"\"\n",
    "res_partner: table\n",
    "    + columns\n",
    "        id: integer NN default nextval('res_partner_id_seq'::regclass)\n",
    "            . references\n",
    "                sequence = res_partner_id_seq\n",
    "        name: varchar\n",
    "        date: date\n",
    "            --  Date\n",
    "        title: integer\n",
    "            --  Title\n",
    "        vat: varchar\n",
    "            --  Tax ID\n",
    "        street: varchar\n",
    "            --  Street\n",
    "        street2: varchar\n",
    "            --  Street2\n",
    "        zip: varchar\n",
    "            --  Zip\n",
    "        city: varchar\n",
    "            --  City\n",
    "        state_id: integer\n",
    "            --  State\n",
    "        country_id: integer\n",
    "            --  Country\n",
    "        partner_latitude: numeric\n",
    "            --  Geo Latitude\n",
    "        partner_longitude: numeric\n",
    "            --  Geo Longitude\n",
    "        email: varchar\n",
    "            --  Email\n",
    "        phone: varchar\n",
    "            --  Phone\n",
    "        mobile: varchar\n",
    "            --  Mobile\n",
    "        is_group: boolean\n",
    "            --  Group\n",
    "        registration_date: date\n",
    "            --  Registration Date\n",
    "        phone_sanitized: varchar\n",
    "            --  Sanitized Number\n",
    "        kind: integer\n",
    "            --  Kind\n",
    "        is_partial_group: boolean\n",
    "            --  Partial Group\n",
    "        family_name: varchar\n",
    "            --  Family Name\n",
    "        given_name: varchar\n",
    "            --  Given Name\n",
    "        addl_name: varchar\n",
    "            --  Additional Name\n",
    "        birth_place: varchar\n",
    "            --  Birth Place\n",
    "        birthdate_not_exact: boolean\n",
    "            --  Approximate Birthdate\n",
    "        birthdate: date\n",
    "            --  Date of Birth\n",
    "        gender: varchar\n",
    "            --  Gender (Female/Male)\n",
    "        z_ind_grp_num_individuals: integer\n",
    "            --  Number of individuals\n",
    "        area_id: integer\n",
    "            --  Area\n",
    "        program_membership_count: integer\n",
    "            --  # Program Memberships\n",
    "        entitlements_count: integer\n",
    "            --  # Cash Entitlements\n",
    "        cycles_count: integer\n",
    "            --  # Cycles\n",
    "        inkind_entitlements_count: integer\n",
    "            --  # In-kind Entitlements\n",
    "        z_cst_indv_disability_level: integer\n",
    "            --  Disability level\n",
    "        z_cst_indv_receive_government_benefits: boolean\n",
    "            --  Receive government benefits\n",
    "        z_cst_indv_locust_aug_2022_lost_livestock: boolean\n",
    "            --  Lost significant livestock during Locust Infestation Aug 2022\n",
    "        z_cst_indv_locust_aug_2022_lost_primary_source_income: boolean\n",
    "            --  Lost primary source income during Locust Infestation Aug 2022\n",
    "        z_cst_indv_has_birth_certificate: boolean\n",
    "            --  Has birth certificate\n",
    "        z_cst_indv_pregnancy_start_date: date\n",
    "            --  Pregnancy Start Date\n",
    "        z_cst_grp_is_farm: boolean\n",
    "            --  Is Farm\n",
    "        z_cst_grp_locust_affected_aug_2022: boolean\n",
    "            --  Locust Affected Aug 2022\n",
    "        z_cst_grp_house_damaged_during_locust_aug_2022: boolean\n",
    "            --  House damaged during locust infestation Aug 2022\n",
    "        z_cst_indv_is_farmer: boolean\n",
    "            --  Is a farmer\n",
    "        z_cst_indv_is_pregnant: boolean\n",
    "            --  Is pregnant\n",
    "        z_cst_grp_partial_crop_destruction_due_to_locust: boolean\n",
    "            --  Partial crop farm destruction due to locust infestation\n",
    "        z_cst_grp_repeated_crop_failure_due_lack_rainfall: boolean\n",
    "            --  Repeated crop failure due to consecutive lack of rainfall\n",
    "    + indices\n",
    "        #1: unique (id)\n",
    "        res_partner_name_index: index (name)\n",
    "        res_partner_display_name_index: index (display_name)\n",
    "        res_partner_date_index: index (date)\n",
    "        res_partner_parent_id_index: index (parent_id)\n",
    "        res_partner_ref_index: index (ref)\n",
    "        res_partner_vat_index: index (vat)\n",
    "        res_partner_commercial_partner_id_index: index (commercial_partner_id)\n",
    "        res_partner_message_main_attachment_id_index: index (message_main_attachment_id)\n",
    "    + keys\n",
    "        #1: PK (id) (underlying index #1)\n",
    "    + foreign-keys\n",
    "        #2: foreign key (title) -> res_partner_title (id) d:set_null\n",
    "        #5: foreign key (state_id) -> res_country_state (id) d:restrict\n",
    "        #6: foreign key (country_id) -> res_country (id) d:restrict\n",
    "        #13: foreign key (kind) -> g2p_group_kind (id) d:set_null\n",
    "        #14: foreign key (area_id) -> spp_area (id) d:set_null\n",
    "    + checks\n",
    "        res_partner_check_name: check ((((type)::text = 'contact'::text) AND (name IS NOT NULL)) OR ((type)::text <> 'contact'::text)) cols = [type, name]\n",
    "            --  CHECK( (type='contact' AND name IS NOT NULL) or (type!='contact') )\n",
    "\n",
    "\n",
    "/*\n",
    "3 rows from res_partner table:\n",
    "id\tname\tdate\ttitle\tvat\tstreet\tstreet2\tzip\tcity\tstate_id\tcountry_id\tpartner_latitude\tpartner_longitude\temail\tphone\tmobile\tis_group\tregistration_date\tphone_sanitized\tkind\tis_partial_group\tfamily_name\tgiven_name\taddl_name\tbirth_place\tbirthdate_not_exact\tbirthdate\tgender\tz_ind_grp_num_individuals\tarea_id\tprogram_membership_count\tentitlements_count\tcycles_count\tinkind_entitlements_count\tz_cst_indv_disability_level\tz_cst_indv_receive_government_benefits\tz_cst_indv_locust_aug_2022_lost_livestock\tz_cst_indv_locust_aug_2022_lost_primary_source_income\tz_cst_indv_has_birth_certificate\tz_cst_indv_pregnancy_start_date\tz_cst_grp_is_farm\tz_cst_grp_locust_affected_aug_2022\tz_cst_grp_house_damaged_during_locust_aug_2022\tz_cst_indv_is_farmer\tz_cst_indv_is_pregnant\tz_cst_grp_partial_crop_destruction_due_to_locust\tz_cst_grp_repeated_crop_failure_due_lack_rainfall\n",
    "61\tKayla Haryanto\t\t\t\tJalan Ciwastra No. 62\tGang Kebonjati\t25394\tPayakumbuh\t\t\t\t\t\t(090) 818-7257\t\tFALSE\t2023-05-31 00:00:00\t\t\t\tHaryanto\tKayla\t\t\t\t1941-09-19 00:00:00\tFemale\t\t\t0\t0\t0\t0\t0\tFALSE\tFALSE\tTRUE\t\t1977-12-24 00:00:00\t\t\t\t\t\t\t\n",
    "62\tArtawan Haryanto\t\t\t\tGg. Cikutra Timur No. 34\tJalan Dr. Djunjunan\t76926\tLubuklinggau\t\t\t\t\t\t(0984) 028 4823\t\tFALSE\t2023-05-31 00:00:00\t\t\t\tHaryanto\tArtawan\t\t\t\t2022-12-25 00:00:00\tMale\t\t\t0\t0\t0\t0\t0\tFALSE\tFALSE\tFALSE\t\t\t\t\t\t\t\t\t\n",
    "60\tHaryanto\t\t\t\tJalan Setiabudhi No. 828\tGg. Lembong\t72034\tProbolinggo\t\t\t\t\t\t\t\tTRUE\t2023-05-31 00:00:00\t\t1\t\t\t\t\t\t\t\t\t6\t6\t0\t0\t0\t0\t\t\t\t\t\t\tFALSE\tFALSE\tFALSE\t\t\t\t\n",
    "*/\n",
    "/*\n",
    "The possible values for \"gender\" are case sensitive: 'Male', 'Female'\n",
    "*/\n",
    "\"\"\",\n",
    "    \"g2p_group_membership\": \"\"\"\n",
    "    g2p_group_membership: table\n",
    "    --  Group Membership\n",
    "    + columns\n",
    "        id: integer NN default nextval('g2p_group_membership_id_seq'::regclass)\n",
    "            . references\n",
    "                sequence = g2p_group_membership_id_seq\n",
    "        group: integer NN\n",
    "            --  Group\n",
    "        individual: integer NN\n",
    "            --  Individual\n",
    "        start_date: timestamp\n",
    "            --  Start Date\n",
    "        ended_date: timestamp\n",
    "            --  Ended Date\n",
    "        is_ended: boolean\n",
    "            --  Is Ended\n",
    "    + indices\n",
    "        #1: unique (id)\n",
    "    + keys\n",
    "        #1: PK (id) (underlying index #1)\n",
    "    + foreign-keys\n",
    "        #1: foreign key (group) -> res_partner (id) d:restrict\n",
    "        #2: foreign key (individual) -> res_partner (id) d:restrict\n",
    "\n",
    "/*\n",
    "3 rows from g2p_group_membership table:\n",
    "id\tgroup\tindividual\tstart_date\tended_date\tis_ended\n",
    "1\t60\t61\t2023-06-09 16:17:47\tNone\tFalse\n",
    "2\t62\t63\t2023-06-09 16:17:47\tNone\tFalse\n",
    "3\t64\t65\t2023-06-09 16:17:47\tNone\tFalse\n",
    "*/\n",
    "\"\"\",\n",
    "    \"g2p_group_membership_kind\": \"\"\"g2p_group_membership_kind: table\n",
    "    --  Group Membership Kind\n",
    "    + columns\n",
    "        id: integer NN default nextval('g2p_group_membership_kind_id_seq'::regclass)\n",
    "            . references\n",
    "                sequence = g2p_group_membership_kind_id_seq\n",
    "        name: varchar\n",
    "            --  Kind\n",
    "        is_unique: boolean\n",
    "            --  Unique\n",
    "\n",
    "\n",
    "/*\n",
    "3 rows from g2p_group_membership_kind table:\n",
    "id\tname\tis_unique\tcreate_uid\tcreate_date\twrite_uid\twrite_date\n",
    "1\tHead\tTrue\t1\t2023-06-09 16:08:37.371047\t1\t2023-06-09 16:08:37.371047\n",
    "*/\n",
    "\"\"\",\n",
    "    \"g2p_group_membership_g2p_group_membership_kind_rel\":\"\"\"\n",
    "    g2p_group_membership_g2p_group_membership_kind_rel: table\n",
    "    --  RELATION BETWEEN g2p_group_membership AND g2p_group_membership_kind\n",
    "    + columns\n",
    "        g2p_group_membership_id: integer NN\n",
    "        g2p_group_membership_kind_id: integer NN\n",
    "    + indices\n",
    "        #1: unique (g2p_group_membership_id, g2p_group_membership_kind_id)\n",
    "        g2p_group_membership_g2p_grou_g2p_group_membership_kind_id__idx: index (g2p_group_membership_kind_id, g2p_group_membership_id)\n",
    "    + keys\n",
    "        #1: PK (g2p_group_membership_id, g2p_group_membership_kind_id) (underlying index #1)\n",
    "    + foreign-keys\n",
    "        g2p_group_membership_g2p_group_mem_g2p_group_membership_id_fkey: foreign key (g2p_group_membership_id) -> g2p_group_membership (id) d:cascade\n",
    "        g2p_group_membership_g2p_grou_g2p_group_membership_kind_id_fkey: foreign key (g2p_group_membership_kind_id) -> g2p_group_membership_kind (id) d:cascade\n",
    "\n",
    "/*\n",
    "3 rows from g2p_group_membership_g2p_group_membership_kind_rel table:\n",
    "g2p_group_membership_id\tg2p_group_membership_kind_id\n",
    "1\t1\n",
    "2\t1\n",
    "3\t1\n",
    "*/\n",
    "    \"\"\"\n",
    "}\n",
    "\n",
    "PROMPT_SUFFIX = \"\"\"Only use the following tables and fields:\n",
    "\n",
    "{table_info}\n",
    "\n",
    "Here is an example of a question and the corresponding SQL query:\n",
    "\n",
    "Example 1:\n",
    "Question: Get all households where the head is elderly.\n",
    "SQLQuery: \n",
    "-- Unsure - I consider elderly as 60 or above\n",
    "SELECT DISTINCT hh.id\n",
    "FROM res_partner as hh\n",
    "LEFT JOIN g2p_group_membership as gm ON hh.id = gm.group\n",
    "LEFT JOIN g2p_group_membership_g2p_group_membership_kind_rel as gmkr ON gm.id = gmkr.g2p_group_membership_id\n",
    "LEFT JOIN g2p_group_membership_kind as gmk ON gmkr.g2p_group_membership_kind_id = gmk.id\n",
    "LEFT JOIN res_partner as ind ON gm.individual = ind.id and ind.is_group = False\n",
    "where \n",
    "  hh.is_group = true and -- We make sure the household is a group\n",
    "  gmk.name = 'Head' and  -- We make sure the individual is the head of the household\n",
    "  ind.birthdate < current_date - interval '65 years' -- We make sure the individual is an elderly\n",
    "\n",
    "\n",
    "Now, try the following question:\n",
    "Question: {input}\"\"\"\n",
    "\n",
    "_postgres_prompt = \"\"\"\n",
    "You are a PostgreSQL expert and your task is to generate SQL queries based on given input questions.\n",
    "\n",
    "- Your task involves creating a PostgreSQL query and verifying it for correctness.\n",
    "- Use only the column names from the tables provided below and do not introduce any that do not exist.\n",
    "- When the question involves \"today\", use the `CURRENT_DATE` function.\n",
    "- In the database fields, `z_cst_` prefix signifies custom fields. If `z_cst_` is followed by `indv_`, it applies to individuals, and if followed by `grp_`, it applies to groups. Do not use fields starting with `z_ind_`.\n",
    "- The SQL queries you generate will be used in a social protection application to define enrollment criteria.\n",
    "- The only field you should query from the `res_partner` table is `id`.\n",
    "- The term 'Head' refers to the head of a household.\n",
    "- If criteria are listed without an explicit `AND`, consider them as `OR`.\n",
    "- Unless specified otherwise, your target should be groups.\n",
    "- Always use the `g2p_group_membership` table to get the members of a group.\n",
    "- Write your SQL queries as simple as possible without adding any unnecessary filters.\n",
    "- If a field necessary for the query is missing from the tables, insert a SQL comment at the exact point in the query where you'd need to use the field. Start the comment with \"Missing field: CRITERIA - \", followed by an explanation. Leave the field out of the query.\n",
    "- If you are uncertain about something in the query, insert a SQL comment at the exact point in the query where the uncertainty lies. Start the comment with \"Unsure - \", followed by an explanation.\n",
    "- Double-check that your SQL is valid\n",
    "\n",
    "Use the following format:\n",
    "\n",
    "Question: [Question here]\n",
    "Answer: [SQL Query to run]\n",
    "\n",
    "\"\"\"\n",
    "\n",
    "POSTGRES_PROMPT = PromptTemplate(\n",
    "    input_variables=[\"input\", \"table_info\"],\n",
    "    template=_postgres_prompt + PROMPT_SUFFIX,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1f9bd83-efe5-4e8a-bdba-71eae11a9dc8",
   "metadata": {},
   "source": [
    "### Configuring langchain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0fe10049-4358-4c59-9c87-51138fb5206c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# db_chain = SQLDatabaseChain.from_llm(llm, db, verbose=True, use_query_checker=True, return_direct=True)\n",
    "# db_chain = SQLDatabaseChain.from_llm(llm, db, verbose=True, return_intermediate_steps=True, prompt=POSTGRES_PROMPT)\n",
    "# db_chain = SQLDatabaseChain.from_llm(llm_sql, db, verbose=True, use_query_checker=True, prompt=POSTGRES_PROMPT, callbacks=callbacks)\n",
    "db_chain = SQLDatabaseChain.from_llm(llm_sql, db, verbose=True, prompt=POSTGRES_PROMPT, callbacks=callbacks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "db53de0b-ad55-461f-8f46-c656249e21f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import Tool\n",
    "\n",
    "sql_tool = Tool(\n",
    "    name='social_protection_db_assistant',\n",
    "    func=db_chain.run,\n",
    "    #description=\"Useful for when you need to write SQL queries for enrolment of beneficiaries and households. Give me all the details about the enrollment criteria, and I'll write the SQL query for you.\"\n",
    "    description=\"\"\"\\\n",
    "    Input to this tool is a detailed list of social protection targeting criteria; output is a SQL query.\n",
    "    Be very specific in the input to avoid any possible doubt. For example, when age is involved, for a child, say under 18 instead of a child. \n",
    "    If the request is not correct, an error message will be returned.\n",
    "    Example Input: Households with at least an individual under 36 months or headed by a woman \\\n",
    "    \"\"\"  # noqa: E501\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a2431b1d-881c-4061-9583-e22a217e15a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import load_tools\n",
    "\n",
    "tools = load_tools(\n",
    "    [\"llm-math\", \"human\"], \n",
    "    llm=llm\n",
    ")\n",
    "tools.append(sql_tool)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4dd23c9a-af93-4651-b65a-bf2f70da80f2",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Creating our main agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "6ad8e040-6961-4078-9761-c7111f575e5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import initialize_agent\n",
    "\n",
    "zero_shot_agent = initialize_agent(\n",
    "    agent=\"zero-shot-react-description\", \n",
    "    tools=tools, \n",
    "    llm=llm,\n",
    "    verbose=True,\n",
    "    max_iterations=5,\n",
    "    callbacks=callbacks,\n",
    "    output_parser = RetryWithErrorOutputParser\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "ad93b893-fc98-4744-8b69-a1e7b4550237",
   "metadata": {},
   "outputs": [],
   "source": [
    "QUERY = \"\"\"\n",
    "You are a Social Protection Officer responsible for defining enrollment criteria for a government welfare program. Your role also involves collaborating with data scientists to implement these criteria effectively. In this position, you have the power to shape the eligibility requirements for various social support programs, ensuring that they are fair, equitable, and targeted towards those who need them the most.\n",
    "\n",
    "Write the SQL query to enrol people in a new social protection program. The sql query will return the list of beneficiary ids.\n",
    "- Make sure to handle all the criteria listed.\n",
    "- If there is a list of criteria, and it is not explicitly an `AND` consider they are `OR`.\n",
    "- Unless explicitly specified, consider we target household.\n",
    "\n",
    "If you need more details about the targeting criteria, ask human precise questions about what else you want to know.\n",
    "\n",
    "Targeting criteria:\n",
    "{criteria}\n",
    "\n",
    "A sample of the data you have:\n",
    "\n",
    "3 rows from the group and individuals table:\n",
    "id\tname\tdate\ttitle\tvat\tstreet\tstreet2\tzip\tcity\tstate_id\tcountry_id\tpartner_latitude\tpartner_longitude\temail\tphone\tmobile\tis_group\tregistration_date\tphone_sanitized\tkind\tis_partial_group\tfamily_name\tgiven_name\taddl_name\tbirth_place\tbirthdate_not_exact\tbirthdate\tgender\tz_ind_grp_num_individuals\tarea_id\tprogram_membership_count\tentitlements_count\tcycles_count\tinkind_entitlements_count\tz_cst_indv_disability_level\tz_cst_indv_receive_government_benefits\tz_cst_indv_locust_aug_2022_lost_livestock\tz_cst_indv_locust_aug_2022_lost_primary_source_income\tz_cst_indv_has_birth_certificate\tz_cst_indv_pregnancy_start_date\tz_cst_grp_is_farm\tz_cst_grp_locust_affected_aug_2022\tz_cst_grp_house_damaged_during_locust_aug_2022\tz_cst_indv_is_farmer\tz_cst_indv_is_pregnant\tz_cst_grp_partial_crop_destruction_due_to_locust\tz_cst_grp_repeated_crop_failure_due_lack_rainfall\n",
    "61\tKayla Haryanto\t\t\t\tJalan Ciwastra No. 62\tGang Kebonjati\t25394\tPayakumbuh\t\t\t\t\t\t(090) 818-7257\t\tFALSE\t2023-05-31 00:00:00\t\t\t\tHaryanto\tKayla\t\t\t\t1941-09-19 00:00:00\tFemale\t\t\t0\t0\t0\t0\t0\tFALSE\tFALSE\tTRUE\t\t1977-12-24 00:00:00\t\t\t\t\t\t\t\n",
    "62\tArtawan Haryanto\t\t\t\tGg. Cikutra Timur No. 34\tJalan Dr. Djunjunan\t76926\tLubuklinggau\t\t\t\t\t\t(0984) 028 4823\t\tFALSE\t2023-05-31 00:00:00\t\t\t\tHaryanto\tArtawan\t\t\t\t2022-12-25 00:00:00\tMale\t\t\t0\t0\t0\t0\t0\tFALSE\tFALSE\tFALSE\t\t\t\t\t\t\t\t\t\n",
    "60\tHaryanto\t\t\t\tJalan Setiabudhi No. 828\tGg. Lembong\t72034\tProbolinggo\t\t\t\t\t\t\t\tTRUE\t2023-05-31 00:00:00\t\t1\t\t\t\t\t\t\t\t\t6\t6\t0\t0\t0\t0\t\t\t\t\t\t\tFALSE\tFALSE\tFALSE\t\t\t\t\n",
    "\n",
    "3 rows from g2p_group_membership table:\n",
    "id\tgroup\tindividual\tstart_date\tended_date\tis_ended\n",
    "1\t60\t61\t2023-06-09 16:17:47\tNone\tFalse\n",
    "2\t62\t63\t2023-06-09 16:17:47\tNone\tFalse\n",
    "3\t64\t65\t2023-06-09 16:17:47\tNone\tFalse\n",
    "\n",
    "3 rows from g2p_group_membership_kind table:\n",
    "id\tname\tis_unique\tcreate_uid\tcreate_date\twrite_uid\twrite_date\n",
    "1\tHead\tTrue\t1\t2023-06-09 16:08:37.371047\t1\t2023-06-09 16:08:37.371047\n",
    "\n",
    "3 rows from g2p_group_membership_g2p_group_membership_kind_rel table:\n",
    "g2p_group_membership_id\tg2p_group_membership_kind_id\n",
    "1\t1\n",
    "2\t1\n",
    "3\t1\n",
    "\n",
    "You also have a table that relates the group and individual together as well as their role in the group such as head\n",
    "\n",
    "Here's how you should proceed:\n",
    "- Verify the given targeting criteria against the sample data and identify any missing fields you have all the needed information, if any fields are missing, list the one missing only and suggest their addition to the dataset to the human.\n",
    "- ask for confirmation of your understanding of the criteria.\n",
    "- Craft an SQL query that adheres to the targeting criteria and works with the available data.\n",
    "\n",
    "Please remember to:\n",
    "- Avoid asking generic questions to the user, such as `Can you explain the targeting criteria in more detail`.\n",
    "- Follow the described format of 'Action:' after 'Thought:'.\n",
    "\"\"\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a393553b-57bf-4200-ba8b-06d3c13326a5",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Defining our targeting criteria"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "9014f9f1-6291-408c-93a8-83d63ca3ce95",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# criteria = \"\"\"\n",
    "# - Households (HH) with children\n",
    "# - HH with pregnant/lactating women\n",
    "# - HH with elderly (including single / elderly-headed HHs)\n",
    "# - HHs with disabled (mental or physical) members \n",
    "# - HHs with members that have chronic illness/medical conditions\n",
    "# \"\"\"\n",
    "\n",
    "# criteria = \"\"\"\n",
    "# 1. Farmers with total crop farm distraction/loss by locust,\n",
    "# 2. Farmers with partial farm distraction by locust,\n",
    "# 3. Households/Farmers previously faced repeated crop failure as a result of consecutive\n",
    "# rainfall failure,\n",
    "# 4. Locust Affected households with large family size/ dependent family member.\n",
    "# 5. Locust affected households with chronically ill, Elderly, and disable family member,\n",
    "# 6. Locust affected women headed household,\n",
    "# 7. Locust affected households with nutritionally impacted family members (pregnant and\n",
    "# lactating women, children under five)\n",
    "# \"\"\"\n",
    "\n",
    "criteria = \"\"\"\n",
    "1. Locust affected households with chronically ill, Elderly, and disable family member,\n",
    "2. Locust affected women headed household,\n",
    "3. Locust affected households with nutritionally impacted family members (pregnant and lactating women, children under five)\n",
    "\"\"\"\n",
    "\n",
    "# criteria = \"\"\"\n",
    "# - Woman-headed household with at least a kid\n",
    "# - households who already receive government benefits\n",
    "# \"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f003b6e5-20d8-459b-bc03-ea29ae4599bd",
   "metadata": {},
   "source": [
    "## Executing the chain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "edb5e693-4eb2-471e-ae6f-a7b4941845ef",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3mI need to verify if I have all the necessary fields in the sample data to target the specified criteria. I will list the fields required for each criterion and check if they are present in the sample data.\n",
      "\n",
      "1. Locust affected households with chronically ill, Elderly, and disable family member:\n",
      "   - z_cst_grp_locust_affected_aug_2022 (locust affected)\n",
      "   - chronically ill (missing)\n",
      "   - Elderly (can be derived from birthdate)\n",
      "   - z_cst_indv_disability_level (disable family member)\n",
      "\n",
      "2. Locust affected women headed household:\n",
      "   - z_cst_grp_locust_affected_aug_2022 (locust affected)\n",
      "   - gender (to identify women)\n",
      "   - g2p_group_membership_kind_id (to identify head of household)\n",
      "\n",
      "3. Locust affected households with nutritionally impacted family members (pregnant and lactating women, children under five):\n",
      "   - z_cst_grp_locust_affected_aug_2022 (locust affected)\n",
      "   - z_cst_indv_pregnancy_start_date (pregnant women)\n",
      "   - lactating women (missing)\n",
      "   - birthdate (children under five)\n",
      "\n",
      "Thought: I need to ask the human to add the missing fields to the dataset and confirm my understanding of the criteria.\n",
      "Action: Human\n",
      "Action Input: I have identified that the fields for chronically ill and lactating women are missing from the sample data. Can you please add these fields to the dataset? Also, please confirm if my understanding of the targeting criteria is correct.\u001b[0m\n",
      "\n",
      "I have identified that the fields for chronically ill and lactating women are missing from the sample data. Can you please add these fields to the dataset? Also, please confirm if my understanding of the targeting criteria is correct.\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      " ignore those fields and continue\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Observation: \u001b[33;1m\u001b[1;3mignore those fields and continue\u001b[0m\n",
      "Thought:"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:Error in on_chain_start callback: 'input'\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32;1m\u001b[1;3mI will proceed with crafting the SQL query without the missing fields and using the available data to target the specified criteria.\n",
      "\n",
      "1. Locust affected households with chronically ill, Elderly, and disable family member:\n",
      "   - z_cst_grp_locust_affected_aug_2022 (locust affected)\n",
      "   - Elderly (can be derived from birthdate)\n",
      "   - z_cst_indv_disability_level (disable family member)\n",
      "\n",
      "2. Locust affected women headed household:\n",
      "   - z_cst_grp_locust_affected_aug_2022 (locust affected)\n",
      "   - gender (to identify women)\n",
      "   - g2p_group_membership_kind_id (to identify head of household)\n",
      "\n",
      "3. Locust affected households with nutritionally impacted family members (pregnant and lactating women, children under five):\n",
      "   - z_cst_grp_locust_affected_aug_2022 (locust affected)\n",
      "   - z_cst_indv_pregnancy_start_date (pregnant women)\n",
      "   - birthdate (children under five)\n",
      "\n",
      "Action: social_protection_db_assistant\n",
      "Action Input: Households with at least an individual affected by locust and either elderly or with a disability level, or headed by a woman and affected by locust, or with at least an individual affected by locust and either pregnant or under 60 months\u001b[0m\n",
      "\n",
      "\u001b[1m> Entering new SQLDatabaseChain chain...\u001b[0m\n",
      "Households with at least an individual affected by locust and either elderly or with a disability level, or headed by a woman and affected by locust, or with at least an individual affected by locust and either pregnant or under 60 months\n",
      "SQLQuery:"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/jeremi/.virtualenvs/unicef/lib/python3.10/site-packages/langchain/sql_database.py:267: SAWarning: Cannot correctly sort tables; there are unresolvable cycles between tables \"account_account, account_account_template, account_account_type, account_analytic_account, account_analytic_group, account_analytic_line, account_bank_statement, account_bank_statement_cashbox, account_bank_statement_line, account_cash_rounding, account_chart_template, account_fiscal_position, account_full_reconcile, account_group, account_incoterms, account_journal, account_move, account_move_line, account_partial_reconcile, account_payment, account_payment_method, account_payment_method_line, account_payment_term, account_reconcile_model, account_tax, account_tax_group, account_tax_repartition_line, barcode_nomenclature, g2p_cycle, g2p_entitlement_inkind, g2p_group_kind, g2p_program, g2p_program_notification_manager_sms, ir_act_report_xml, ir_act_window, ir_attachment, ir_mail_server, ir_model, ir_module_category, ir_module_module, ir_sequence, ir_ui_view, mail_activity_type, mail_alias, mail_template, payment_acquirer, payment_token, payment_transaction, procurement_group, product_category, product_packaging, product_product, product_removal, product_template, report_paperformat, res_bank, res_company, res_country, res_country_group, res_country_state, res_currency, res_partner, res_partner_bank, res_partner_industry, res_partner_title, res_users, resource_calendar, sms_template, spp_area, spp_area_kind, spp_service_point, stock_location, stock_location_route, stock_move, stock_package_level, stock_package_type, stock_picking, stock_picking_type, stock_quant_package, stock_rule, stock_storage_category, stock_warehouse, stock_warehouse_orderpoint, uom_category, uom_uom\", which is usually caused by mutually dependent foreign key constraints.  Foreign key constraints involving these tables will not be considered; this warning may raise an error in a future release.\n",
      "  for tbl in self._metadata.sorted_tables\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32;1m\u001b[1;3mSELECT DISTINCT hh.id\n",
      "FROM res_partner as hh\n",
      "LEFT JOIN g2p_group_membership as gm ON hh.id = gm.group\n",
      "LEFT JOIN g2p_group_membership_g2p_group_membership_kind_rel as gmkr ON gm.id = gmkr.g2p_group_membership_id\n",
      "LEFT JOIN g2p_group_membership_kind as gmk ON gmkr.g2p_group_membership_kind_id = gmk.id\n",
      "LEFT JOIN res_partner as ind ON gm.individual = ind.id and ind.is_group = False\n",
      "WHERE \n",
      "  hh.is_group = true AND -- We make sure the household is a group\n",
      "  (\n",
      "    (\n",
      "      hh.z_ind_grp_is_hh_with_disabled = true OR\n",
      "      ind.birthdate < current_date - interval '65 years'\n",
      "    ) AND\n",
      "    ind.z_cst_indv_locust_aug_2022_injured = true\n",
      "  ) OR\n",
      "  (\n",
      "    gmk.name = 'Head' AND\n",
      "    ind.gender = 'female' AND\n",
      "    hh.z_cst_grp_locust_affected_aug_2022 = true\n",
      "  ) OR\n",
      "  (\n",
      "    ind.z_cst_indv_locust_aug_2022_injured = true AND\n",
      "    (\n",
      "      ind.z_cst_indv_is_pregnant = true OR\n",
      "      ind.birthdate >= current_date - interval '60 months'\n",
      "    )\n",
      "  )\u001b[0m\n",
      "SQLResult: \u001b[33;1m\u001b[1;3m[(95,), (99,), (106,), (167,), (181,), (201,), (360,), (457,), (465,), (470,), (507,), (543,), (578,), (600,), (614,), (641,), (678,), (688,), (693,), (739,), (749,), (756,), (781,), (794,), (813,), (892,), (926,), (939,), (956,)]\u001b[0m\n",
      "Answer:"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:Error in on_chain_end callback: 'output'\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32;1m\u001b[1;3mSELECT DISTINCT hh.id\n",
      "FROM res_partner as hh\n",
      "LEFT JOIN g2p_group_membership as gm ON hh.id = gm.group\n",
      "LEFT JOIN g2p_group_membership_g2p_group_membership_kind_rel as gmkr ON gm.id = gmkr.g2p_group_membership_id\n",
      "LEFT JOIN g2p_group_membership_kind as gmk ON gmkr.g2p_group_membership_kind_id = gmk.id\n",
      "LEFT JOIN res_partner as ind ON gm.individual = ind.id and ind.is_group = False\n",
      "WHERE \n",
      "  hh.is_group = true AND -- We make sure the household is a group\n",
      "  (\n",
      "    (\n",
      "      hh.z_ind_grp_is_hh_with_disabled = true OR\n",
      "      ind.birthdate < current_date - interval '65 years'\n",
      "    ) AND\n",
      "    ind.z_cst_indv_locust_aug_2022_injured = true\n",
      "  ) OR\n",
      "  (\n",
      "    gmk.name = 'Head' AND\n",
      "    ind.gender = 'female' AND\n",
      "    hh.z_cst_grp_locust_affected_aug_2022 = true\n",
      "  ) OR\n",
      "  (\n",
      "    ind.z_cst_indv_locust_aug_2022_injured = true AND\n",
      "    (\n",
      "      ind.z_cst_indv_is_pregnant = true OR\n",
      "      ind.birthdate >= current_date - interval '60 months'\n",
      "    )\n",
      "  )\u001b[0m\n",
      "\u001b[1m> Finished chain.\u001b[0m\n",
      "\n",
      "Observation: \u001b[38;5;200m\u001b[1;3mSELECT DISTINCT hh.id\n",
      "FROM res_partner as hh\n",
      "LEFT JOIN g2p_group_membership as gm ON hh.id = gm.group\n",
      "LEFT JOIN g2p_group_membership_g2p_group_membership_kind_rel as gmkr ON gm.id = gmkr.g2p_group_membership_id\n",
      "LEFT JOIN g2p_group_membership_kind as gmk ON gmkr.g2p_group_membership_kind_id = gmk.id\n",
      "LEFT JOIN res_partner as ind ON gm.individual = ind.id and ind.is_group = False\n",
      "WHERE \n",
      "  hh.is_group = true AND -- We make sure the household is a group\n",
      "  (\n",
      "    (\n",
      "      hh.z_ind_grp_is_hh_with_disabled = true OR\n",
      "      ind.birthdate < current_date - interval '65 years'\n",
      "    ) AND\n",
      "    ind.z_cst_indv_locust_aug_2022_injured = true\n",
      "  ) OR\n",
      "  (\n",
      "    gmk.name = 'Head' AND\n",
      "    ind.gender = 'female' AND\n",
      "    hh.z_cst_grp_locust_affected_aug_2022 = true\n",
      "  ) OR\n",
      "  (\n",
      "    ind.z_cst_indv_locust_aug_2022_injured = true AND\n",
      "    (\n",
      "      ind.z_cst_indv_is_pregnant = true OR\n",
      "      ind.birthdate >= current_date - interval '60 months'\n",
      "    )\n",
      "  )\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3mI now know the final answer\n",
      "Final Answer: SELECT DISTINCT hh.id\n",
      "FROM res_partner as hh\n",
      "LEFT JOIN g2p_group_membership as gm ON hh.id = gm.group\n",
      "LEFT JOIN g2p_group_membership_g2p_group_membership_kind_rel as gmkr ON gm.id = gmkr.g2p_group_membership_id\n",
      "LEFT JOIN g2p_group_membership_kind as gmk ON gmkr.g2p_group_membership_kind_id = gmk.id\n",
      "LEFT JOIN res_partner as ind ON gm.individual = ind.id and ind.is_group = False\n",
      "WHERE \n",
      "  hh.is_group = true AND -- We make sure the household is a group\n",
      "  (\n",
      "    (\n",
      "      hh.z_ind_grp_is_hh_with_disabled = true OR\n",
      "      ind.birthdate < current_date - interval '65 years'\n",
      "    ) AND\n",
      "    ind.z_cst_indv_locust_aug_2022_injured = true\n",
      "  ) OR\n",
      "  (\n",
      "    gmk.name = 'Head' AND\n",
      "    ind.gender = 'female' AND\n",
      "    hh.z_cst_grp_locust_affected_aug_2022 = true\n",
      "  ) OR\n",
      "  (\n",
      "    ind.z_cst_indv_locust_aug_2022_injured = true AND\n",
      "    (\n",
      "      ind.z_cst_indv_is_pregnant = true OR\n",
      "      ind.birthdate >= current_date - interval '60 months'\n",
      "    )\n",
      "  )\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "enrollment_query = zero_shot_agent.run(QUERY.format(custom_table_info=custom_table_info, criteria=criteria))\n",
    "aim_callback.flush_tracker(langchain_asset=zero_shot_agent, reset=False, finish=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3307550d-0895-41ce-a1bb-1b002fc2ae9d",
   "metadata": {},
   "source": [
    "## Pushing to OpenSPP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "47762c54-de5c-4a38-9b44-d96c41fe0864",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SELECT DISTINCT hh.id\n",
      "FROM res_partner as hh\n",
      "LEFT JOIN g2p_group_membership as gm ON hh.id = gm.group\n",
      "LEFT JOIN g2p_group_membership_g2p_group_membership_kind_rel as gmkr ON gm.id = gmkr.g2p_group_membership_id\n",
      "LEFT JOIN g2p_group_membership_kind as gmk ON gmkr.g2p_group_membership_kind_id = gmk.id\n",
      "LEFT JOIN res_partner as ind ON gm.individual = ind.id and ind.is_group = False\n",
      "WHERE \n",
      "  hh.is_group = true AND -- We make sure the household is a group\n",
      "  (\n",
      "    (\n",
      "      hh.z_ind_grp_is_hh_with_disabled = true OR\n",
      "      ind.birthdate < current_date - interval '65 years'\n",
      "    ) AND\n",
      "    ind.z_cst_indv_locust_aug_2022_injured = true\n",
      "  ) OR\n",
      "  (\n",
      "    gmk.name = 'Head' AND\n",
      "    ind.gender = 'female' AND\n",
      "    hh.z_cst_grp_locust_affected_aug_2022 = true\n",
      "  ) OR\n",
      "  (\n",
      "    ind.z_cst_indv_locust_aug_2022_injured = true AND\n",
      "    (\n",
      "      ind.z_cst_indv_is_pregnant = true OR\n",
      "      ind.birthdate >= current_date - interval '60 months'\n",
      "    )\n",
      "  )\n"
     ]
    }
   ],
   "source": [
    "print(enrollment_query)\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  },
  "toc-showmarkdowntxt": false
 },
 "nbformat": 4,
 "nbformat_minor": 5
}