💾 Adding a vectorstore (first implementation)

.gitignore

@@ -6,4 +6,6 @@ dist
 build
 **.pickle
 **.pkl
-.env
+.env
+volumes
+docker-compose.yml

example.ipynb

@@ -1,31 +1,40 @@
 {
  "cells": [
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Examples\n",
+    "\n",
+    "Below you can find some examples of how to use the 🤖 `Megabots` library."
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 13,
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "True"
-      ]
-     },
-     "execution_count": 8,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
     "from megabots import bot\n",
-    "from dotenv import load_dotenv\n",
+    "from dotenv import load_dotenv"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Creating a bot\n",
+    "\n",
+    "The `bot` object is the main object of the library. It is used to create a bot and to interact with it.\n",
     "\n",
-    "load_dotenv()"
+    "The `index` argument specifies the index to use for the bot. It can either be a saved index file (e.g., `index.pkl`) or a directory of documents (e.g., `./index`). In the case of the directory the index will be automatically created. If no index is specified `bot` will look for `index.pkl` or `./index`"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 14,
    "metadata": {},
    "outputs": [
     {
@@ -33,7 +42,7 @@
      "output_type": "stream",
      "text": [
       "Using model: gpt-3.5-turbo\n",
-      "Loading path from disk...\n"
+      "Loading path from pickle file:  ./index.pkl ...\n"
      ]
     },
     {
@@ -42,7 +51,7 @@
        "'The first roster of the Avengers included Iron Man, Thor, Hulk, Ant-Man, and the Wasp.'"
       ]
      },
-     "execution_count": 9,
+     "execution_count": 14,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -52,9 +61,19 @@
     "qnabot.ask(\"what was the first roster of the avengers?\")"
    ]
   },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Changing the bot's prompt\n",
+    "\n",
+    "You can change the bots promnpt to customize it to your needs."
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 15,
    "metadata": {},
    "outputs": [
     {
@@ -62,7 +81,7 @@
      "output_type": "stream",
      "text": [
       "Using model: gpt-3.5-turbo\n",
-      "Loading path from disk...\n"
+      "Loading path from pickle file:  ./index.pkl ...\n"
      ]
     },
     {
@@ -71,7 +90,7 @@
        "\"Hmmm! Let me think about that... Ah yes, the original Avengers lineup included Iron Man, Thor, Hulk, Ant-Man, and the Wasp. They were like the ultimate superhero squad, except for maybe the Teenage Mutant Ninja Turtles. But let's be real, they were just a bunch of turtles who liked pizza.\""
       ]
      },
-     "execution_count": 10,
+     "execution_count": 15,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -97,6 +116,59 @@
     ")\n",
     "qnabot.ask(\"what was the first roster of the avengers?\")\n"
    ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Using Megabots with Milvus\n",
+    "\n",
+    "Megabots `bot` can also use Milvus as a backend for its search engine. You can find an example of how to do it below.\n",
+    "\n",
+    "In order to run Milvus you need to follow [this guide](https://milvus.io/docs/example_code.md) to download a docker compose file and run it.\n",
+    "The command is:\n",
+    "  \n",
+    "```bash\n",
+    "wget https://raw.githubusercontent.com/milvus-io/pymilvus/v2.2.7/examples/hello_milvus.py\n",
+    "```\n",
+    "You can then [install Attu](https://milvus.io/docs/attu_install-docker.md) as a management tool for Milvus"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Using model: gpt-3.5-turbo\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "'The first roster of the Avengers included Iron Man, Thor, Hulk, Ant-Man, and the Wasp.'"
+      ]
+     },
+     "execution_count": 11,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from megabots import bot, vectorstore\n",
+    "\n",
+    "# Create a vectorstore object. Default port is 19530 and default host is localhost\n",
+    "milvus = vectorstore(\"milvus\")\n",
+    "\n",
+    "# Point it to your files directory so that it can index the files and add them to the vectorstore\n",
+    "bot = bot(\"qna-over-docs\", index=\"./examples/files/\", vectorstore=milvus)\n",
+    "\n",
+    "bot.ask(\"what was the first roster of the avengers?\")\n"
+   ]
   }
  ],
  "metadata": {

megabots/__init__.py

@@ -1,7 +1,7 @@
+from typing import Any
 from langchain.llms import OpenAI
 from langchain.chat_models import ChatOpenAI
 from langchain.embeddings import OpenAIEmbeddings
-from langchain.document_loaders import DirectoryLoader, S3DirectoryLoader
 from langchain.chains.qa_with_sources import load_qa_with_sources_chain
 from langchain.vectorstores.faiss import FAISS
 import gradio as gr
@@ -11,10 +11,9 @@ import os
 from dotenv import load_dotenv
 from langchain.prompts import PromptTemplate
 from langchain.chains.question_answering import load_qa_chain
-from langchain.chains.conversational_retrieval.prompts import (
-    CONDENSE_QUESTION_PROMPT,
-    QA_PROMPT,
-)
+from langchain.chains.conversational_retrieval.prompts import QA_PROMPT
+from langchain.document_loaders import DirectoryLoader
+from megabots.vectorstores import vectorstore
 
 load_dotenv()
 
@@ -25,15 +24,17 @@ class Bot:
         model: str | None = None,
         prompt_template: str | None = None,
         prompt_variables: list[str] | None = None,
-        memory: str | None = None,
         index: str | None = None,
         sources: bool | None = False,
+        # TODO: Fix this typing
+        vectorstore: Any | None = None,
+        memory: str | None = None,
         verbose: bool = False,
         temperature: int = 0,
     ):
         self.select_model(model, temperature)
         self.create_loader(index)
-        self.load_or_create_index(index)
+        self.load_or_create_index(index, vectorstore)
 
         # Load the question-answering chain for the selected model
         self.chain = self.create_chain(
@@ -83,18 +84,25 @@ class Bot:
             )
         self.loader = DirectoryLoader(index, recursive=True)
 
-    def load_or_create_index(self, index_path: str):
+    def load_or_create_index(self, index: str, vectorstore=None):
         # Load an existing index from disk or create a new one if not available
+        if vectorstore is not None:
+            self.search_index = vectorstore.client.from_documents(
+                self.loader.load_and_split(),
+                OpenAIEmbeddings(),
+                connection_args={"host": vectorstore.host, "port": vectorstore.port},
+            )
+            return
 
         # Is pickle
-        if index_path is not None and "pkl" in index_path or "pickle" in index_path:
-            print("Loading path from disk...")
-            with open(index_path, "rb") as f:
+        if index is not None and "pkl" in index or "pickle" in index:
+            print("Loading path from pickle file: ", index, "...")
+            with open(index, "rb") as f:
                 self.search_index = pickle.load(f)
             return
 
         # Is directory
-        if index_path is not None and os.path.isdir(index_path):
+        if index is not None and os.path.isdir(index):
             print("Creating index...")
             self.search_index = FAISS.from_documents(
                 self.loader.load_and_split(), OpenAIEmbeddings()
@@ -125,9 +133,8 @@ SUPPORTED_TASKS = {
         "impl": Bot,
         "default": {
             "model": "gpt-3.5-turbo",
-            "prompt": "",
             "temperature": 0,
-            "index": "./files",
+            "index": "./index",
         },
     }
 }

megabots/vectorstores.py

@@ -0,0 +1,33 @@
+from typing import Type, TypeVar
+from langchain.vectorstores import Milvus, Qdrant
+from abc import ABC
+
+
+class MilvusVectorStore:
+    def __init__(self, host: str, port: int):
+        self.host = host
+        self.port = port
+        self.client = Milvus
+
+
+SUPPORTED_VECTORSTORES = {
+    "milvus": {
+        "impl": MilvusVectorStore,
+        "default": {"host": "localhost", "port": 19530},
+    }
+}
+
+
+def vectorstore(name: str) -> MilvusVectorStore:
+    """Return a vectorstore object."""
+
+    if name is None:
+        raise RuntimeError("Impossible to instantiate a vectorstore without a name.")
+
+    if name not in SUPPORTED_VECTORSTORES:
+        raise ValueError(f"Vectorstore {name} is not supported.")
+
+    return SUPPORTED_VECTORSTORES[name]["impl"](
+        host=SUPPORTED_VECTORSTORES[name]["default"]["host"],
+        port=SUPPORTED_VECTORSTORES[name]["default"]["port"],
+    )