final-agent.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Wrapper class for the locally hosted LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests as rq\n",
    "import json\n",
    "\n",
    "class LLM():\n",
    "    def __init__(self, hostname, port):\n",
    "        self.url = f\"http://{hostname}:{port}/api/v1/generate\"\n",
    "    \n",
    "    def get_response(self, user_request):\n",
    "        # Sends a request to the LLM with the complete prompt\n",
    "        # The preset used for the parameters was Divine Intellect\n",
    "        response = rq.post(self.url, json.dumps({\n",
    "            \"max_context_length\": 2048, \n",
    "            \"max_length\": 120, \n",
    "            \"prompt\": user_request,\n",
    "            \"temperature\": 1.31,\n",
    "            \"top_k\": 49,\n",
    "            \"top_p\": 0.14,\n",
    "            \"stop_sequence\": [\"</SQLQuery>\\n\"]}))\n",
    "        return response.json()[\"results\"][0][\"text\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = LLM(\"localhost\", 5001)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Wrapper class for the Vector Dabatase that holds the context added through RAG "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chromadb\n",
    "from sentence_transformers import SentenceTransformer\n",
    "import regex as re\n",
    "\n",
    "class ContextVectorDB():\n",
    "    def __init__(self, collection_name):\n",
    "        # Initialization of the embedding model\n",
    "        self.embedding_model = SentenceTransformer('sentence-transformers/multi-qa-MiniLM-L6-cos-v1')\n",
    "        self.chroma_client = chromadb.Client()\n",
    "        self.collection = self.chroma_client.create_collection(name=collection_name)\n",
    "\n",
    "    def parse_string_to_json(self, input_string):\n",
    "        # Uses Regex to parse the raw input string\n",
    "        question_pattern = r'<question>\\s*Question:\\s*(.*?)\\s*\\n</question>'\n",
    "        query_pattern = r'<SQLQuery>\\s*SQLQuery:\\s*(.*?)\\s*\\n</SQLQuery>'\n",
    "\n",
    "        question_match = re.search(question_pattern, input_string, re.DOTALL)\n",
    "        query_match = re.search(query_pattern, input_string, re.DOTALL)\n",
    "\n",
    "        question = question_match.group(1).strip() if question_match else None\n",
    "        query = query_match.group(1).strip() if query_match else None\n",
    "\n",
    "        result = {\n",
    "            \"question\": question,\n",
    "            \"query_result\": query\n",
    "        }\n",
    "\n",
    "        return result\n",
    "\n",
    "    def load_examples(self, examples):\n",
    "        examples_parsed = []\n",
    "        for example in examples:\n",
    "            examples_parsed.append(self.parse_string_to_json(example))\n",
    "        self.populate_vectors(examples_parsed)\n",
    "\n",
    "    def populate_vectors(self, dataset):\n",
    "        # Iterates over the dataset and adds the embeddings to the ChromaDB collection\n",
    "        for i, item in enumerate(dataset):\n",
    "            user_input = item['question']\n",
    "            query = item['query_result']\n",
    "            combined_text = f\"\"\"\n",
    "<question>\n",
    "Question: {user_input}\n",
    "</question>\n",
    "<SQLquery>\n",
    "SQLQuery: {query}\n",
    "</SQLquery>\"\"\"\n",
    "            embeddings = self.embedding_model.encode(user_input).tolist()\n",
    "            self.collection.add(embeddings=[embeddings], documents=[combined_text], ids=[f\"id_{i}\"])\n",
    "\n",
    "    def search_context(self, user_input, n_results=1):\n",
    "        # Method to search the ChromaDB collection for relevant context based on a user input\n",
    "        user_embeddings = self.embedding_model.encode(user_input).tolist()\n",
    "        return self.collection.query(query_embeddings=user_embeddings, n_results=n_results, include=['documents'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Agent class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sqlalchemy as db\n",
    "from sqlalchemy.sql import text\n",
    "import random as rd\n",
    "import csv\n",
    "\n",
    "class QueriesTranslationAgent():\n",
    "    def __init__(self, llm, _db, examples_file):\n",
    "        self.llm = llm\n",
    "        self.engine = db.create_engine(f\"sqlite:///{_db}\")\n",
    "        self.data_retreived = []\n",
    "\n",
    "        with open(examples_file, \"r\") as file:\n",
    "            examples = file.read()\n",
    "        examples = examples.split(\"\\n\\n\")\n",
    "        self.vector_db = ContextVectorDB(\"user_inputs\")\n",
    "        self.vector_db.load_examples(examples)\n",
    "    \n",
    "    def build_prompt(self, user_request):\n",
    "        # For the prompt, two shot prompting was used, along with a template that explained the task\n",
    "        examples_raw = self.vector_db.search_context(user_request, n_results=2)\n",
    "        example1 = examples_raw[\"documents\"][0][0]\n",
    "        example2 = examples_raw[\"documents\"][0][1]\n",
    "\n",
    "        template = \"\"\"\n",
    "You are an agent designed to translate questions about a database into SQL queries.\n",
    "You also have a tool that can export the results to a csv file. \n",
    "You have a table called apple_daily with the following columns:\n",
    "\n",
    "(Date DATE, Time TIME, Open REAL, High REAL, Low REAL, Close REAL, Volume INTEGER)\n",
    "\n",
    "Never query for all the columns from a specific table, only ask for the relevant columns given the question.\n",
    "DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database.\n",
    "If the users asks for either results to be exported or a csv file with results, you should call your tool making your answer be \"SELECT export data;\".\n",
    "If you think the user wants to export data you shouldn't say anything else other than \"SELECT export data;\".\n",
    "You should always think about what to do step by step.\n",
    "Translate:\n",
    "\n",
    "{example1}\n",
    "\n",
    "{example2}\n",
    "\n",
    "<question>\n",
    "Question: {user_request}\n",
    "<question>\n",
    "<SQLQuery>\n",
    "SQLQuery: \"\"\"\n",
    "        prompt = template.format(example1=example1, example2=example2, user_request=user_request)\n",
    "        return prompt\n",
    "\n",
    "    def parse_raw_response(self, response):\n",
    "        query = response.split(\"\\n\")[0].lstrip()\n",
    "        return query\n",
    "    \n",
    "    def execute_query(self, query):\n",
    "        query_result = self.engine.connect().execute(text(query)).fetchall()\n",
    "        self.engine.dispose()\n",
    "        self.data_retreived.append([query_result])\n",
    "        return query_result\n",
    "\n",
    "    def export_data(self):\n",
    "        with open(\"data.csv\", \"w\") as file:\n",
    "            writer = csv.writer(file, delimiter=\",\", quotechar='\"', lineterminator=\"\\n\")\n",
    "            for data in self.data_retreived:\n",
    "                writer.writerows(data)\n",
    "        self.data_retreived = []\n",
    "        return \"Data exported to file data.csv\"\n",
    "\n",
    "    def invoke(self, user_request):\n",
    "        prompt = self.build_prompt(user_request)\n",
    "        response = self.llm.get_response(prompt)\n",
    "        query = self.parse_raw_response(response)\n",
    "        if query == \"SELECT export data;\":\n",
    "            return self.export_data()\n",
    "        return self.execute_query(query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = QueriesTranslationAgent(llm, \"apple.db\", \"examples.txt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[('2023-01-04', '01:00:00', 126.89, 128.6557, 125.08, 126.36, 89113633)]\n"
     ]
    }
   ],
   "source": [
    "print(agent.invoke(\"Give me all the rows corresponding to the date 2023-01-04\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(250,)]\n"
     ]
    }
   ],
   "source": [
    "print(agent.invoke(\"Count how many rows in the table are from the year 2023\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(23,)]\n"
     ]
    }
   ],
   "source": [
    "print(agent.invoke(\"Count how many rows in the table are from month 03\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(172.25632000000002, 1.0)]\n"
     ]
    }
   ],
   "source": [
    "print(agent.invoke(\"What's the average value of the columns Open and Time in the table apple_daily?\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data exported to file data.csv\n"
     ]
    }
   ],
   "source": [
    "print(agent.invoke(\"I want to export results to a csv file\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sandbox for testing the Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(172.25632000000002,)]\n",
      "Data exported to file data.csv\n"
     ]
    }
   ],
   "source": [
    "output = \"Begin!\"\n",
    "while (output != \"Data exported to file data.csv\"):\n",
    "    user_request = input(\"Enter your question: \")\n",
    "    output = agent.invoke(user_request)\n",
    "    print(output)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}