🎉 Bot memory, many fixes, some refactoring

README.md

@@ -2,7 +2,6 @@
 
 [![Tests](https://github.com/momegas/qnabot/actions/workflows/python-package.yml/badge.svg)](https://github.com/momegas/qnabot/actions/workflows/python-package.yml)
 ![](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 🫵
 
@@ -58,13 +57,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.
@@ -75,6 +176,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.
@@ -89,14 +192,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"
    ]
   },
   {
@@ -29,14 +19,22 @@
     "\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",
-    "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`"
+    "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`\n"
    ]
   },
   {
    "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 bots promnpt to customize it to your needs."
+    "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).\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,6 +1,7 @@
 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
 
 

megabots/bot.py

@@ -10,6 +10,7 @@ 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
@@ -19,8 +20,7 @@ class Bot:
     def __init__(
         self,
         model: str | None = None,
-        prompt_template: str | None = None,
-        prompt_variables: list[str] | None = None,
+        prompt: PromptTemplate | None = None,
         index: str | None = None,
         sources: bool | None = False,
         vectorstore: VectorStore | None = None,
@@ -28,36 +28,36 @@ class Bot:
         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)
-        self.vectorstore = vectorstore
-        self.memory = memory
+
         # Load the question-answering chain for the selected model
-        self.chain = self.create_chain(
-            prompt_template, prompt_variables, sources=sources, verbose=verbose
-        )
+        self.chain = self.create_chain(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
+                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=prompt
+            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):
@@ -132,6 +132,7 @@ SUPPORTED_TASKS = {
             "model": "gpt-3.5-turbo",
             "temperature": 0,
             "index": "./index",
+            "input_variables": ["context", "question"],
         },
     }
 }
@@ -141,10 +142,10 @@ SUPPORTED_MODELS = {}
 
 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,
+    prompt: str | None = None,
     memory: str | Memory | None = None,
     vectorstore: str | VectorStore | None = None,
     verbose: bool = False,
@@ -154,13 +155,21 @@ def bot(
 
     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.
+
+        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:
@@ -178,11 +187,17 @@ def bot(
 
     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_template=prompt_template,
-        prompt_variables=prompt_variables,
+        prompt=None
+        if prompt is None
+        else PromptTemplate(
+            template=prompt, input_variables=task_defaults["input_variables"]
+        ),
         temperature=temperature,
         verbose=verbose,
         vectorstore=megabots.vectorstore(vectorstore)

megabots/memory.py

@@ -1,31 +1,15 @@
-from langchain.memory import (
-    ConversationBufferMemory,
-    ConversationBufferWindowMemory,
-    ConversationSummaryMemory,
-    ConversationSummaryBufferMemory,
-)
+from langchain.memory import ConversationBufferMemory, ConversationBufferWindowMemory
 
 
 class ConversationBuffer:
     def __init__(self):
-        self.memory = ConversationBufferMemory
+        self.memory = ConversationBufferMemory(input_key="question")
 
 
 class ConversationBufferWindow:
-    def __init__(self, memory_window: int):
-        self.memory_window: int = memory_window
-        self.memory = ConversationBufferWindowMemory
-
-
-class ConversationSummary:
-    def __init__(self):
-        self.memory = ConversationSummaryMemory
-
-
-class ConversationSummaryBuffer:
-    def __init__(self, max_token_limit: int):
-        self.max_token_limit: int = max_token_limit
-        self.memory = ConversationSummaryBufferMemory
+    def __init__(self, k: int):
+        self.k: int = k
+        self.memory = ConversationBufferWindowMemory(k=self.k, input_key="question")
 
 
 SUPPORTED_MEMORY = {
@@ -35,35 +19,17 @@ SUPPORTED_MEMORY = {
     },
     "conversation-buffer-window": {
         "impl": ConversationBufferWindow,
-        "default": {"memory_window": 3},
-    },
-    "conversation-summary": {
-        "impl": ConversationSummary,
-        "default": {},
-    },
-    "conversation-summary-buffer": {
-        "impl": ConversationSummaryBuffer,
-        "default": {"max_token_limit": 40},
+        "default": {"k": 3},
     },
 }
 
 
-Memory = type(
-    "Memory",
-    (
-        ConversationBuffer,
-        ConversationBufferWindow,
-        ConversationSummary,
-        ConversationSummaryBuffer,
-    ),
-    {},
-)
+Memory = type("Memory", (ConversationBuffer, ConversationBufferWindow), {})
 
 
 def memory(
     name: str = "conversation-buffer-window",
-    memory_window: int | None = None,
-    max_token_limit: int | None = None,
+    k: int | None = None,
 ) -> Memory:
     if name is None:
         raise RuntimeError("Impossible to instantiate memory without a name.")
@@ -74,13 +40,6 @@ def memory(
     cl = SUPPORTED_MEMORY[name]["impl"]
 
     if name == "conversation-buffer-window":
-        if max_token_limit != None:
-            raise ValueError(f"max_token_limit cannot be set for {name} memory")
-        return cl(memory_window=memory_window)
-
-    if name == "conversation-summary-buffer":
-        if max_token_limit != None:
-            raise ValueError(f"memory_window cannot be set for {name} memory")
-        return cl(max_token_limit=max_token_limit)
+        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)

tests/test_memory.py

@@ -1,42 +1,25 @@
-import pytest
-from megabots.memory import (
-    ConversationBufferWindow,
-    ConversationSummaryBuffer,
-    memory,
-    Memory,
-    SUPPORTED_MEMORY,
-)
-
-
-def test_memory_name_none():
-    with pytest.raises(RuntimeError):
-        memory(name=None)
+from pytest import raises
+from megabots import memory
+from megabots.memory import ConversationBuffer, ConversationBufferWindow
 
 
-def test_memory_not_supported():
-    with pytest.raises(ValueError):
-        memory(name="unsupported_memory_type")
+def test_memory_conversation_buffer():
+    mem = memory(name="conversation-buffer")
+    assert isinstance(mem, ConversationBuffer)
 
 
 def test_memory_conversation_buffer_window():
-    mem_obj = memory(name="conversation-buffer-window", memory_window=5)
-    assert isinstance(mem_obj, ConversationBufferWindow)
-    assert mem_obj.memory_window == 5
-    assert mem_obj.__class__ == SUPPORTED_MEMORY["conversation-buffer-window"]["impl"]
-
+    mem = memory(name="conversation-buffer-window", k=10)
+    assert isinstance(mem, ConversationBufferWindow)
 
-def test_memory_conversation_buffer_window_invalid_max_token_limit():
-    with pytest.raises(ValueError):
-        memory(name="conversation-buffer-window", memory_window=5, max_token_limit=10)
 
+def test_memory_unsupported_name():
+    with raises(ValueError, match=r"Memory invalid-name is not supported."):
+        memory(name="invalid-name")
 
-def test_memory_conversation_summary_buffer():
-    mem_obj = memory(name="conversation-summary-buffer", max_token_limit=10)
-    assert isinstance(mem_obj, ConversationSummaryBuffer)
-    assert mem_obj.max_token_limit == 10
-    assert mem_obj.__class__ == SUPPORTED_MEMORY["conversation-summary-buffer"]["impl"]
 
-
-def test_memory_conversation_summary_buffer_invalid_memory_window():
-    with pytest.raises(ValueError):
-        memory(name="conversation-summary-buffer", memory_window=5, max_token_limit=10)
+def test_memory_no_name():
+    with raises(
+        RuntimeError, match=r"Impossible to instantiate memory without a name."
+    ):
+        memory(name=None)