Merge branch 'main' into general_fixes

README.md

@@ -5,7 +5,6 @@
 [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
 [![License](https://img.shields.io/badge/License-MIT-informational.svg)](https://github.com/momegas/megabots/blob/main/LICENCE)
 ![](https://dcbadge.vercel.app/api/server/zkqDWk5S7P?style=flat&n&compact=true)
-<a href="https://www.producthunt.com/posts/megabots-2?utm_source=badge-featured&utm_medium=badge&utm_souce=badge-megabots&#0045;2" target="_blank"><img src="https://api.producthunt.com/widgets/embed-image/v1/featured.svg?post_id=390033&theme=light" alt="Megabots - 🤖&#0032;Production&#0032;ready&#0032;full&#0045;stack&#0032;LLM&#0032;apps&#0032;made&#0032;mega&#0045;easy | Product Hunt" style="width: 150px; height: 34px;" width="250" height="54" /></a>
 
 
 🤖 Megabots provides State-of-the-art, production ready LLM apps made mega-easy, so you don't have to build them from scratch 🤯 Create a bot, now 🫵
@@ -63,13 +62,115 @@ qnabot = bot("qna-over-docs", index="./index")
 
 # Change the model
 qnabot = bot("qna-over-docs", model="text-davinci-003")
+```
+
+## Changing the bot's prompt
+
+You can change the bots promnpt to customize it to your needs. In the `qna-over-docs` type of bot you will need to pass 2 variables for the `context` (knwoledge searched from the index) and the `question` (the human question).
+
+```python
+from megabots import bot
+
+prompt = """
+Use the following pieces of context to answer the question at the end.
+If you don't know the answer, just say that you don't know, don't try to make up an answer.
+Answer in the style of Tony Stark.
+
+{context}
+
+Question: {question}
+Helpful humorous answer:"""
+
+qnabot = bot("qna-over-docs", index="./index.pkl", prompt=prompt)
+
+qnabot.ask("what was the first roster of the avengers?")
+```
+
+## Working with memory
+
+You can easily add memory to your `bot` using the `memory` parameter. It accepts a string with the type of the memory to be used. This defaults to some sane dafaults.
+Should you need more configuration, you can use the `memory` function and pass the type of memory and the configuration you need.
+
+```python
+from megabots import bot
+
+qnabot = bot("qna-over-docs", index="./index.pkl", memory="conversation-buffer")
+
+print(qnabot.ask("who is iron man?"))
+print(qnabot.ask("was he in the first roster?"))
+# Bot should understand who "he" refers to.
+```
+
+Or using the `memory`factory function
+
+```python
+from megabots import bot, memory
+
+mem("conversation-buffer-window", k=5)
+
+qnabot = bot("qna-over-docs", index="./index.pkl", memory=mem)
+
+print(qnabot.ask("who is iron man?"))
+print(qnabot.ask("was he in the first roster?"))
+```
+
+NOTE: For the `qna-over-docs` bot, when using memory and passing your custom prompt, it is important to remember to pass one more variable to your custom prompt to facilitate for chat history. The variable name is `history`.
+
+```python
+from megabots import bot
+
+prompt = """
+Use the following pieces of context to answer the question at the end.
+If you don't know the answer, just say that you don't know, don't try to make up an answer.
+
+{context}
 
-# Change the prompt
-prompt_template = "Be humourous in your responses. Question: {question}\nContext: {context}, Answer:"
-prompt_variables=["question", "context"]
-qnabot = bot("qna-over-docs", prompt_template=prompt_template, prompt_variables=prompt_variables)
+{history}
+Human: {question}
+AI:"""
+
+qnabot = bot("qna-over-docs", prompt=prompt, index="./index.pkl", memory="conversation-buffer")
+
+print(qnabot.ask("who is iron man?"))
+print(qnabot.ask("was he in the first roster?"))
 ```
 
+## Using Megabots with Milvus (more DBs comming soon)
+
+Megabots `bot` can also use Milvus as a backend for its search engine. You can find an example of how to do it below.
+
+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.
+The command is:
+
+```bash
+wget https://raw.githubusercontent.com/milvus-io/pymilvus/v2.2.7/examples/hello_milvus.py
+```
+
+You can then [install Attu](https://milvus.io/docs/attu_install-docker.md) as a management tool for Milvus
+
+```python
+from megabots import bot
+
+# Attach a vectorstore by passing the name of the database. Default port for milvus is 19530 and default host is localhost
+# Point it to your files directory so that it can index the files and add them to the vectorstore
+bot = bot("qna-over-docs", index="./examples/files/", vectorstore="milvus")
+
+bot.ask("what was the first roster of the avengers?")
+```
+
+Or use the `vectorstore` factory function for more customisation
+
+```python
+
+from megabots import bot, vectorstore
+
+milvus = vectorstore("milvus", host="localhost", port=19530)
+
+bot = bot("qna-over-docs", index="./examples/files/", vectorstore=milvus)
+```
+
+## Exposing an API with FastAPI
+
 You can also create a FastAPI app that will expose the bot as an API using the create_app function.
 Assuming you file is called `main.py` run `uvicorn main:app --reload` to run the API locally.
 You should then be able to visit `http://localhost:8000/docs` to see the API documentation.
@@ -80,6 +181,8 @@ from megabots import bot, create_api
 app = create_app(bot("qna-over-docs"))
 ```
 
+## Exposing a Gradio chat-like interface
+
 You can expose a gradio UI for the bot using `create_interface` function.
 Assuming your file is called `ui.py` run `gradio qnabot/ui.py` to run the UI locally.
 You should then be able to visit `http://127.0.0.1:7860` to see the API documentation.
@@ -94,14 +197,16 @@ demo = create_interface(bot("qna-over-docs"))
 
 The `bot` function should serve as the starting point for creating and customising your bot. Below is a list of the available arguments in `bot`.
 
-| Argument         | Description                                                                                                                                                                                                                                                                                |
-| ---------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
-| task             | The type of bot to create. Available options: `qna-over-docs`. More comming soon                                                                                                                                                                                                           |
-| index            | 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` |
-| model            | The name of the model to use for the bot. You can specify a different model by providing its name, like "text-davinci-003". Supported models: `gpt-3.5-turbo` (default),`text-davinci-003` More comming soon.                                                                              |
-| prompt_template  | A string template for the prompt, which defines the format of the question and context passed to the model. The template should include placeholders for the variables specified in `prompt_variables`.                                                                                    |
-| prompt_variables | A list of variables to be used in the prompt template. These variables are replaced with actual values when the bot processes a query.                                                                                                                                                     |
-| sources          | When `sources` is `True` the bot will also include sources in the response. A known [issue](https://github.com/hwchase17/langchain/issues/2858) exists, where if you pass a custom prompt with sources the code breaks.                                                                    |
+| Argument    | Description                                                                                                                                                                                                                                                                                |
+| ----------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
+| task        | The type of bot to create. Available options: `qna-over-docs`. More comming soon                                                                                                                                                                                                           |
+| index       | 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` |
+| model       | The name of the model to use for the bot. You can specify a different model by providing its name, like "text-davinci-003". Supported models: `gpt-3.5-turbo` (default),`text-davinci-003` More comming soon.                                                                              |
+| prompt      | A string template for the prompt, which defines the format of the question and context passed to the model. The template should include placeholder variables like so: `context`, `{question}` and in the case of using memory `history`.                                                  |
+| memory      | The type of memory to be used by the bot. Can be a string with the type of the memory or you can use `memory` factory function. Supported memories: `conversation-buffer`, `conversation-buffer-window`                                                                                    |
+| vectorstore | The vectorstore to be used for the index. Can be a string with the name of the databse or you can use `vectorstore` factory function. Supported DBs: `milvus`.                                                                                                                             |
+
+| sources | When `sources` is `True` the bot will also include sources in the response. A known [issue](https://github.com/hwchase17/langchain/issues/2858) exists, where if you pass a custom prompt with sources the code breaks. |
 
 ## How QnA bot works
 

example.ipynb

@@ -7,17 +7,7 @@
    "source": [
     "# Examples\n",
     "\n",
-    "Below you can find some examples of how to use the 🤖 `Megabots` library."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 13,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from megabots import bot\n",
-    "from dotenv import load_dotenv"
+    "Below you can find some examples of how to use the 🤖 `Megabots` library.\n"
    ]
   },
   {
@@ -34,9 +24,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/momegas/Desktop/qnabot/.venv/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    },
     {
      "name": "stdout",
      "output_type": "stream",
@@ -51,14 +49,17 @@
        "'The first roster of the Avengers included Iron Man, Thor, Hulk, Ant-Man, and the Wasp.'"
       ]
      },
-     "execution_count": 14,
+     "execution_count": 1,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
+    "from megabots import bot\n",
+    "\n",
     "qnabot = bot(\"qna-over-docs\", index=\"./index.pkl\")\n",
-    "qnabot.ask(\"what was the first roster of the avengers?\")"
+    "\n",
+    "qnabot.ask(\"what was the first roster of the avengers?\")\n"
    ]
   },
   {
@@ -68,12 +69,12 @@
    "source": [
     "### Changing the bot's prompt\n",
     "\n",
-    "You can change the bot's prompt to customize it to your needs."
+    "You can change the bot's promnpt to customize it to your needs. In the `qna-over-docs` type of bot you will need to pass 2 variables for the `context` (knwoledge searched from the index) and the `question` (the human question).\n"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [
     {
@@ -87,33 +88,29 @@
     {
      "data": {
       "text/plain": [
-       "\"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.\""
+       "'The first roster of the Avengers included Iron Man, Thor, Hulk, Ant-Man, and the Wasp.'"
       ]
      },
-     "execution_count": 15,
+     "execution_count": 2,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "prompt_template = \"\"\"\n",
+    "from megabots import bot\n",
+    "\n",
+    "prompt = \"\"\"\n",
     "Use the following pieces of context to answer the question at the end. \n",
     "If you don't know the answer, just say that you don't know, don't try to make up an answer.\n",
-    "Be very playfull and humourous in your responses. always try to make the user laugh.\n",
-    "Always start your answers with 'Hmmm! Let me think about that...'\n",
+    "Answer in the style of Tony Stark.\n",
+    "\n",
     "{context}\n",
     "\n",
     "Question: {question}\n",
     "Helpful humorous answer:\"\"\"\n",
     "\n",
-    "load_dotenv()\n",
+    "qnabot = bot(\"qna-over-docs\", index=\"./index.pkl\", prompt=prompt)\n",
     "\n",
-    "qnabot = bot(\n",
-    "    \"qna-over-docs\",\n",
-    "    index=\"./index.pkl\",\n",
-    "    prompt_template=prompt_template,\n",
-    "    prompt_variables=[\"context\", \"question\"],\n",
-    ")\n",
     "qnabot.ask(\"what was the first roster of the avengers?\")\n"
    ]
   },
@@ -128,16 +125,17 @@
     "\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",
+    "\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"
+    "\n",
+    "You can then [install Attu](https://milvus.io/docs/attu_install-docker.md) as a management tool for Milvus\n"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [
     {
@@ -153,22 +151,170 @@
        "'The first roster of the Avengers included Iron Man, Thor, Hulk, Ant-Man, and the Wasp.'"
       ]
      },
-     "execution_count": 11,
+     "execution_count": 3,
      "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",
+    "from megabots import bot\n",
     "\n",
+    "# Attach a vectorstore by passing the name of the database. Default port for milvus is 19530 and default host is localhost\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",
+    "bot = bot(\"qna-over-docs\", index=\"./examples/files/\", vectorstore=\"milvus\")\n",
     "\n",
     "bot.ask(\"what was the first roster of the avengers?\")\n"
    ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Or use the `vectorstore` factory function for more customisation\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Using model: gpt-3.5-turbo\n"
+     ]
+    }
+   ],
+   "source": [
+    "from megabots import bot, vectorstore\n",
+    "\n",
+    "milvus = vectorstore(\"milvus\", host=\"localhost\", port=19530)\n",
+    "\n",
+    "bot = bot(\"qna-over-docs\", index=\"./examples/files/\", vectorstore=milvus)\n"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Working with memory\n",
+    "\n",
+    "You can easily add memory to your `bot` using the `memory` parameter. It accepts a string with the type of the memory to be used. This defaults to some sane dafaults.\n",
+    "Should you need more configuration, you can use the `memory` function and pass the type of memory and the configuration you need.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Using model: gpt-3.5-turbo\n",
+      "Loading path from pickle file:  ./index.pkl ...\n",
+      "Iron Man is a superhero character who is a member of the Avengers. He is known for his high-tech suit of armor and his alter ego, Tony Stark.\n",
+      "Yes, Iron Man was part of the original Avengers lineup.\n"
+     ]
+    }
+   ],
+   "source": [
+    "from megabots import bot\n",
+    "\n",
+    "qnabot = bot(\"qna-over-docs\", index=\"./index.pkl\", memory=\"conversation-buffer\")\n",
+    "\n",
+    "print(qnabot.ask(\"who is iron man?\"))\n",
+    "print(qnabot.ask(\"was he in the first roster?\"))\n"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Or using the `memory`factory function"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Using model: gpt-3.5-turbo\n",
+      "Loading path from pickle file:  ./index.pkl ...\n",
+      "Iron Man is a superhero character who is a member of the Avengers. He is known for his high-tech suit of armor and his alter ego, Tony Stark.\n",
+      "Yes, Iron Man was part of the original Avengers lineup.\n"
+     ]
+    }
+   ],
+   "source": [
+    "from megabots import bot, memory\n",
+    "\n",
+    "qnabot = bot(\n",
+    "    \"qna-over-docs\",\n",
+    "    index=\"./index.pkl\",\n",
+    "    memory=memory(\"conversation-buffer-window\", k=5),\n",
+    ")\n",
+    "\n",
+    "print(qnabot.ask(\"who is iron man?\"))\n",
+    "print(qnabot.ask(\"was he in the first roster?\"))\n"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "NOTE: For the `qna-over-docs` bot, when using memory and passing your custom prompt, it is important to remember to pass one more variable to your custom prompt to facilitate for chat history. The variable name is `history`.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Using model: gpt-3.5-turbo\n",
+      "Loading path from pickle file:  ./index.pkl ...\n",
+      "Iron Man is a superhero character who is a member of the Avengers. He is a wealthy businessman named Tony Stark who uses his advanced technology to create a suit of armor that gives him superhuman abilities.\n",
+      "Yes, Iron Man was part of the original Avengers lineup.\n"
+     ]
+    }
+   ],
+   "source": [
+    "from megabots import bot\n",
+    "\n",
+    "prompt = \"\"\"\n",
+    "Use the following pieces of context to answer the question at the end. \n",
+    "If you don't know the answer, just say that you don't know, don't try to make up an answer.\n",
+    "\n",
+    "{context}\n",
+    "\n",
+    "{history}\n",
+    "Human: {question}\n",
+    "AI:\"\"\"\n",
+    "\n",
+    "qnabot = bot(\n",
+    "    \"qna-over-docs\",\n",
+    "    prompt=prompt,\n",
+    "    index=\"./index.pkl\",\n",
+    "    memory=\"conversation-buffer\",\n",
+    ")\n",
+    "\n",
+    "print(qnabot.ask(\"who is iron man?\"))\n",
+    "print(qnabot.ask(\"was he in the first roster?\"))"
+   ]
   }
  ],
  "metadata": {

megabots/__init__.py

@@ -1,225 +1,10 @@
-from typing import Any
-from langchain.llms import OpenAI
-from langchain.chat_models import ChatOpenAI
-from langchain.embeddings import OpenAIEmbeddings
-from langchain.chains.qa_with_sources import load_qa_with_sources_chain
-from langchain.vectorstores.faiss import FAISS
-import gradio as gr
-from fastapi import FastAPI
-import pickle
-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 QA_PROMPT
-from langchain.document_loaders import DirectoryLoader
-from megabots.vectorstores import VectorStore, vectorstore
-
-load_dotenv()
-
-
-class Bot:
-    def __init__(
-        self,
-        model: str | None = None,
-        prompt_template: str | None = None,
-        prompt_variables: list[str] | None = None,
-        index: str | None = None,
-        sources: bool | None = False,
-        vectorstore: VectorStore | 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, vectorstore)
-
-        # Load the question-answering chain for the selected model
-        self.chain = self.create_chain(
-            prompt_template, prompt_variables, sources=sources, verbose=verbose
-        )
-
-    def create_chain(
-        self,
-        prompt_template: str | None = None,
-        prompt_variables: list[str] | None = None,
-        sources: bool | None = False,
-        verbose: bool = False,
-    ):
-        prompt = (
-            PromptTemplate(template=prompt_template, input_variables=prompt_variables)
-            if prompt_template is not None and prompt_variables is not None
-            else QA_PROMPT
-        )
-        # TODO: Changing the prompt here is not working. Leave it as is for now.
-        # Reference: https://github.com/hwchase17/langchain/issues/2858
-        if sources:
-            return load_qa_with_sources_chain(
-                self.llm, chain_type="stuff", verbose=verbose
-            )
-        return load_qa_chain(
-            self.llm, chain_type="stuff", verbose=verbose, prompt=prompt
-        )
-
-    def select_model(self, model: str | None, temperature: float):
-        # Select and set the appropriate model based on the provided input
-        if model is None or model == "gpt-3.5-turbo":
-            print("Using model: gpt-3.5-turbo")
-            self.llm = ChatOpenAI(temperature=temperature)
-
-        if model == "text-davinci-003":
-            print("Using model: text-davinci-003")
-            self.llm = OpenAI(temperature=temperature)
-
-    def create_loader(self, index: str | None):
-        # Create a loader based on the provided directory (either local or S3)
-        if index is None:
-            raise RuntimeError(
-                """
-            Impossible to find a valid index. 
-            Either provide a valid path to a pickle file or a directory.               
-            """
-            )
-        self.loader = DirectoryLoader(index, recursive=True)
-
-    def load_or_create_index(self, index: str, vectorstore: VectorStore | None = 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 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 is not None and os.path.isdir(index):
-            print("Creating index...")
-            self.search_index = FAISS.from_documents(
-                self.loader.load_and_split(), OpenAIEmbeddings()
-            )
-            return
-
-        raise RuntimeError(
-            """
-            Impossible to find a valid index. 
-            Either provide a valid path to a pickle file or a directory.               
-            """
-        )
-
-    def save_index(self, index_path: str):
-        # Save the index to the specified path
-        with open(index_path, "wb") as f:
-            pickle.dump(self.search_index, f)
-
-    def ask(self, question: str, k=1) -> str:
-        # Retrieve the answer to the given question and return it
-        input_documents = self.search_index.similarity_search(question, k=k)
-        answer = self.chain.run(input_documents=input_documents, question=question)
-        return answer
-
-
-SUPPORTED_TASKS = {
-    "qna-over-docs": {
-        "impl": Bot,
-        "default": {
-            "model": "gpt-3.5-turbo",
-            "temperature": 0,
-            "index": "./index",
-        },
-    }
-}
+from megabots.vectorstore import VectorStore, vectorstore
+from megabots.memory import Memory, memory
+from megabots.bot import Bot, bot
+from megabots.prompt import prompt
+from megabots.utils import create_api, create_interface
 
-SUPPORTED_MODELS = {}
 
+from dotenv import load_dotenv
 
-def bot(
-    task: str | None = None,
-    model: str | None = None,
-    index: str | None = None,
-    prompt_template: str | None = None,
-    prompt_variables: list[str] | None = None,
-    verbose: bool = False,
-    temperature: int = 0,
-    **kwargs,
-) -> Bot:
-    """Instanciate a bot based on the provided task. Each supported tasks has it's own default sane defaults.
-
-    Args:
-        task (str | None, optional): The given task. Can be one of the SUPPORTED_TASKS.
-        model (str | None, optional): Model to be used. Can be one of the SUPPORTED_MODELS.
-        index (str | None, optional): Data that the model will load and store index info.
-        Can be either a local file path, a pickle file, or a url of a vector database.
-        By default it will look for a local directory called "files" in the current working directory.
-        prompt_template (str | None, optional): Prompt template to be used. Specify variables with {var_name}.
-        prompt_variables (list[str] | None, optional): Prompt variables to be used in the prompt template.
-        verbose (bool, optional): Verbocity. Defaults to False.
-        temperature (int, optional): Temperature. Defaults to 0.
-
-    Raises:
-        RuntimeError: _description_
-        ValueError: _description_
-
-    Returns:
-        Bot: Bot instance
-    """
-
-    if task is None:
-        raise RuntimeError("Impossible to instantiate a bot without a task.")
-    if task not in SUPPORTED_TASKS:
-        raise ValueError(f"Task {task} is not supported.")
-
-    task_defaults = SUPPORTED_TASKS[task]["default"]
-    return SUPPORTED_TASKS[task]["impl"](
-        model=model or task_defaults["model"],
-        index=index or task_defaults["index"],
-        prompt_template=prompt_template,
-        prompt_variables=prompt_variables,
-        temperature=temperature,
-        verbose=verbose,
-        **kwargs,
-    )
-
-
-def create_api(bot: Bot):
-    app = FastAPI()
-
-    @app.get("/v1/ask/{question}")
-    async def ask(question: str):
-        answer = bot.ask(question)
-        return {"answer": answer}
-
-    return app
-
-
-def create_interface(bot_instance: Bot, examples: list[list[str]] = []):
-    with gr.Blocks() as interface:
-        chatbot = gr.Chatbot([], elem_id="chatbot").style(height=750)
-        msg = gr.Textbox(
-            show_label=False,
-            placeholder="Enter text and press enter, or upload an image",
-        ).style(container=False)
-        clear = gr.Button("Clear")
-
-        def user(user_message, history):
-            return "", history + [[user_message, None]]
-
-        def bot(history):
-            response = bot_instance.ask(history[-1][0])
-            history[-1][1] = response
-            return history
-
-        msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-            bot, chatbot, chatbot
-        )
-        clear.click(lambda: None, None, chatbot, queue=False)
-
-    return interface
+load_dotenv()

megabots/bot.py

@@ -0,0 +1,207 @@
+from typing import Any
+from langchain.llms import OpenAI
+from langchain.chat_models import ChatOpenAI
+from langchain.embeddings import OpenAIEmbeddings
+from langchain.chains.qa_with_sources import load_qa_with_sources_chain
+from langchain.vectorstores.faiss import FAISS
+import pickle
+import os
+from langchain.prompts import PromptTemplate
+from langchain.chains.question_answering import load_qa_chain
+from langchain.chains.conversational_retrieval.prompts import QA_PROMPT
+from langchain.document_loaders import DirectoryLoader
+from megabots.prompt import QA_MEMORY_PROMPT
+from megabots.vectorstore import VectorStore
+from megabots.memory import Memory
+import megabots
+
+
+class Bot:
+    def __init__(
+        self,
+        model: str | None = None,
+        prompt: PromptTemplate | None = None,
+        index: str | None = None,
+        sources: bool | None = False,
+        vectorstore: VectorStore | None = None,
+        memory: Memory | None = None,
+        verbose: bool = False,
+        temperature: int = 0,
+    ):
+        self.vectorstore = vectorstore
+        self.memory = memory
+        self.prompt = prompt or QA_MEMORY_PROMPT if self.memory else QA_PROMPT
+        self.select_model(model, temperature)
+        self.create_loader(index)
+        self.load_or_create_index(index, vectorstore)
+
+        # Load the question-answering chain for the selected model
+        self.chain = self.create_chain(sources=sources, verbose=verbose)
+
+    def create_chain(
+        self,
+        sources: bool | None = False,
+        verbose: bool = False,
+    ):
+        # TODO: Changing the prompt here is not working. Leave it as is for now.
+        # Reference: https://github.com/hwchase17/langchain/issues/2858
+        if sources:
+            return load_qa_with_sources_chain(
+                self.llm,
+                chain_type="stuff",
+                memory=self.memory.memory if self.memory else None,
+                verbose=verbose,
+            )
+        return load_qa_chain(
+            self.llm,
+            chain_type="stuff",
+            verbose=verbose,
+            prompt=self.prompt,
+            memory=self.memory.memory if self.memory else None,
+        )
+
+    def select_model(self, model: str | None, temperature: float):
+        # Select and set the appropriate model based on the provided input
+        if model is None or model == "gpt-3.5-turbo":
+            print("Using model: gpt-3.5-turbo")
+            self.llm = ChatOpenAI(temperature=temperature)
+
+        if model == "text-davinci-003":
+            print("Using model: text-davinci-003")
+            self.llm = OpenAI(temperature=temperature)
+
+    def create_loader(self, index: str | None):
+        # Create a loader based on the provided directory (either local or S3)
+        if index is None:
+            raise RuntimeError(
+                """
+            Impossible to find a valid index. 
+            Either provide a valid path to a pickle file or a directory.               
+            """
+            )
+        self.loader = DirectoryLoader(index, recursive=True)
+
+    def load_or_create_index(self, index: str, vectorstore: VectorStore | None = 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 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 is not None and os.path.isdir(index):
+            print("Creating index...")
+            self.search_index = FAISS.from_documents(
+                self.loader.load_and_split(), OpenAIEmbeddings()
+            )
+            return
+
+        raise RuntimeError(
+            """
+            Impossible to find a valid index. 
+            Either provide a valid path to a pickle file or a directory.               
+            """
+        )
+
+    def save_index(self, index_path: str):
+        # Save the index to the specified path
+        with open(index_path, "wb") as f:
+            pickle.dump(self.search_index, f)
+
+    def ask(self, question: str, k=1) -> str:
+        # Retrieve the answer to the given question and return it
+        input_documents = self.search_index.similarity_search(question, k=k)
+        answer = self.chain.run(input_documents=input_documents, question=question)
+        return answer
+
+
+SUPPORTED_TASKS = {
+    "qna-over-docs": {
+        "impl": Bot,
+        "default": {
+            "model": "gpt-3.5-turbo",
+            "temperature": 0,
+            "index": "./index",
+            "input_variables": ["context", "question"],
+        },
+    }
+}
+
+SUPPORTED_MODELS = {}
+
+
+def bot(
+    task: str | None = None,
+    *,
+    model: str | None = None,
+    index: str | None = None,
+    prompt: str | None = None,
+    memory: str | Memory | None = None,
+    vectorstore: str | VectorStore | None = None,
+    verbose: bool = False,
+    temperature: int = 0,
+) -> Bot:
+    """Instanciate a bot based on the provided task. Each supported tasks has it's own default sane defaults.
+
+    Args:
+        task (str | None, optional): The given task. Can be one of the SUPPORTED_TASKS.
+
+        model (str | None, optional): Model to be used. Can be one of the SUPPORTED_MODELS.
+
+        index (str | None, optional): Data that the model will load and store index info.
+        Can be either a local file path, a pickle file, or a url of a vector database.
+        By default it will look for a local directory called "files" in the current working directory.
+
+        prompt (str | None, optional): The prompt that the bot will take in. Mark variables like this: {variable}.
+        Variables are context, question, and history if the bot has memory.
+
+        vectorstore: (str | VectorStore | None, optional): The vectorstore that the bot will save the index to.
+        If only a string is passed, the defaults values willl be used.
+
+        verbose (bool, optional): Verbocity. Defaults to False.
+
+        temperature (int, optional): Temperature. Defaults to 0.
+
+    Raises:
+        RuntimeError: _description_
+        ValueError: _description_
+
+    Returns:
+        Bot: Bot instance
+    """
+
+    if task is None:
+        raise RuntimeError("Impossible to instantiate a bot without a task.")
+    if task not in SUPPORTED_TASKS:
+        raise ValueError(f"Task {task} is not supported.")
+
+    task_defaults = SUPPORTED_TASKS[task]["default"]
+
+    if memory is not None:
+        task_defaults["input_variables"].append("history")
+
+    return SUPPORTED_TASKS[task]["impl"](
+        model=model or task_defaults["model"],
+        index=index or task_defaults["index"],
+        prompt=None
+        if prompt is None
+        else PromptTemplate(
+            template=prompt, input_variables=task_defaults["input_variables"]
+        ),
+        temperature=temperature,
+        verbose=verbose,
+        vectorstore=megabots.vectorstore(vectorstore)
+        if isinstance(vectorstore, str)
+        else vectorstore,
+        memory=megabots.memory(memory) if isinstance(memory, str) else memory,
+    )

megabots/memory.py

@@ -0,0 +1,45 @@
+from langchain.memory import ConversationBufferMemory, ConversationBufferWindowMemory
+
+
+class ConversationBuffer:
+    def __init__(self):
+        self.memory = ConversationBufferMemory(input_key="question")
+
+
+class ConversationBufferWindow:
+    def __init__(self, k: int):
+        self.k: int = k
+        self.memory = ConversationBufferWindowMemory(k=self.k, input_key="question")
+
+
+SUPPORTED_MEMORY = {
+    "conversation-buffer": {
+        "impl": ConversationBuffer,
+        "default": {},
+    },
+    "conversation-buffer-window": {
+        "impl": ConversationBufferWindow,
+        "default": {"k": 3},
+    },
+}
+
+
+Memory = type("Memory", (ConversationBuffer, ConversationBufferWindow), {})
+
+
+def memory(
+    name: str = "conversation-buffer-window",
+    k: int | None = None,
+) -> Memory:
+    if name is None:
+        raise RuntimeError("Impossible to instantiate memory without a name.")
+
+    if name not in SUPPORTED_MEMORY:
+        raise ValueError(f"Memory {name} is not supported.")
+
+    cl = SUPPORTED_MEMORY[name]["impl"]
+
+    if name == "conversation-buffer-window":
+        return cl(k=k)
+
+    return SUPPORTED_MEMORY[name]["impl"]()

megabots/prompt.py

@@ -0,0 +1,18 @@
+from typing import List
+from langchain import PromptTemplate
+
+QNA_TEMPLATE = """Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.
+
+{context}
+
+{history}
+Human: {question}
+AI:"""
+
+QA_MEMORY_PROMPT = PromptTemplate(
+    template=QNA_TEMPLATE, input_variables=["context", "history", "question"]
+)
+
+
+def prompt(template: str, variables: List[str]):
+    return PromptTemplate(template=template, input_variables=variables)

megabots/utils.py

@@ -0,0 +1,39 @@
+import gradio as gr
+from fastapi import FastAPI
+from megabots.bot import Bot
+
+
+def create_api(bot: Bot):
+    app = FastAPI()
+
+    @app.get("/v1/ask/{question}")
+    async def ask(question: str):
+        answer = bot.ask(question)
+        return {"answer": answer}
+
+    return app
+
+
+def create_interface(bot_instance: Bot, examples: list[list[str]] = []):
+    with gr.Blocks() as interface:
+        chatbot = gr.Chatbot([], elem_id="chatbot").style(height=750)
+        msg = gr.Textbox(
+            show_label=False,
+            placeholder="Enter text and press enter, or upload an image",
+        ).style(container=False)
+        clear = gr.Button("Clear")
+
+        def user(user_message, history):
+            return "", history + [[user_message, None]]
+
+        def bot(history):
+            response = bot_instance.ask(history[-1][0])
+            history[-1][1] = response
+            return history
+
+        msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+            bot, chatbot, chatbot
+        )
+        clear.click(lambda: None, None, chatbot, queue=False)
+
+    return interface

megabots/{vectorstores.py → vectorstore.py}

@@ -1,5 +1,5 @@
 from typing import Type, TypeVar
-from langchain.vectorstores import Milvus, Qdrant
+from langchain.vectorstores import Milvus
 from abc import ABC
 
 
@@ -26,7 +26,9 @@ SUPPORTED_VECTORSTORES = {
 }
 
 
-def vectorstore(name: str) -> VectorStore:
+def vectorstore(
+    name: str, host: str | None = None, port: int | None = None
+) -> VectorStore:
     """Return a vectorstore object."""
 
     if name is None:
@@ -36,6 +38,6 @@ def vectorstore(name: str) -> VectorStore:
         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"],
+        host=host or SUPPORTED_VECTORSTORES[name]["default"]["host"],
+        port=port or SUPPORTED_VECTORSTORES[name]["default"]["port"],
     )

setup.py

@@ -1,6 +1,6 @@
 from setuptools import setup, find_packages
 
-VERSION = "0.0.10"
+VERSION = "0.0.11"
 
 setup(
     name="megabots",

tests/test_memory.py

@@ -0,0 +1,25 @@
+from pytest import raises
+from megabots import memory
+from megabots.memory import ConversationBuffer, ConversationBufferWindow
+
+
+def test_memory_conversation_buffer():
+    mem = memory(name="conversation-buffer")
+    assert isinstance(mem, ConversationBuffer)
+
+
+def test_memory_conversation_buffer_window():
+    mem = memory(name="conversation-buffer-window", k=10)
+    assert isinstance(mem, ConversationBufferWindow)
+
+
+def test_memory_unsupported_name():
+    with raises(ValueError, match=r"Memory invalid-name is not supported."):
+        memory(name="invalid-name")
+
+
+def test_memory_no_name():
+    with raises(
+        RuntimeError, match=r"Impossible to instantiate memory without a name."
+    ):
+        memory(name=None)