{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "d2ff9d0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.document_loaders import TextLoader\n",
    "sota_loader = TextLoader(\"../../modules/state_of_the_union.txt\")\n",
    "pg_loader = TextLoader(\"../../../../gpt_index/examples/paul_graham_essay/data/paul_graham_essay.txt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "57794791",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.indexes import VectorstoreIndexCreator\n",
    "from langchain.vectorstores import FAISS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "5fa10ffb",
   "metadata": {},
   "outputs": [],
   "source": [
    "sota_index = VectorstoreIndexCreator(vectorstore_cls=FAISS).from_loaders([sota_loader])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "34ceb9c6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running Chroma using direct local API.\n",
      "Using DuckDB in-memory for database. Data will be transient.\n"
     ]
    }
   ],
   "source": [
    "pg_index = VectorstoreIndexCreator(vectorstore_kwargs={\"collection_name\": \"paul-graham\"}).from_loaders([pg_loader])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "9c3d06e4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\" The President nominated Circuit Court of Appeals Judge Ketanji Brown Jackson to serve on the United States Supreme Court. He said she is one of the nation's top legal minds and will continue Justice Breyer's legacy of excellence.\""
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sota_index.query(\"what did the president about kentaji brown jackson?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "94be0f0f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\" Kentaji Brown Jackson was not mentioned in the context, so I don't know.\""
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pg_index.query(\"what did the president about kentaji brown jackson?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "fb7a1185",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import initialize_agent, Tool\n",
    "from langchain.tools import BaseTool\n",
    "from langchain.llms import OpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "b853ce83",
   "metadata": {},
   "outputs": [],
   "source": [
    "tools = [\n",
    "    Tool(\n",
    "        name = \"State of Union QA System\",\n",
    "        func=sota_index.query,\n",
    "        description=\"useful for when you need to answer questions about the most recent state of the union address. Input should be a fully formed question.\"\n",
    "    ),\n",
    "    Tool(\n",
    "        name = \"Paul Graham QA System\",\n",
    "        func=pg_index.query,\n",
    "        description=\"useful for when you need to answer questions about Paul Graham. Input should be a fully formed question.\"\n",
    "    ),\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "a959727e",
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = initialize_agent(tools, OpenAI(temperature=0), agent=\"zero-shot-react-description\", verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "20440754",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "59d1547f",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"../../../notebooks/state_of_union_qa.json\") as f:\n",
    "    sota_qa = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "c3c457df",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"../../../notebooks/paul_graham_qa.json\") as f:\n",
    "    pg_qa = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "36e1ddc2",
   "metadata": {},
   "outputs": [],
   "source": [
    "for d in sota_qa:\n",
    "    d['steps'] = [{\"tool\": \"State of Union QA System\"}, {\"tool_input\": d[\"question\"]}]\n",
    "for d in pg_qa:\n",
    "    d['steps'] = [{\"tool\": \"Paul Graham QA System\"}, {\"tool_input\": d[\"question\"]}]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "59069433",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_vectorstore_routing = sota_qa + pg_qa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "157a27bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"vectorstore_sota_pg.json\", \"w\") as f:\n",
    "    json.dump(all_vectorstore_routing, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe86c9d2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}