🤝 Add custom prompt template

example.ipynb

@@ -2,7 +2,30 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "True"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from megabots import bot\n",
+    "from dotenv import load_dotenv\n",
+    "\n",
+    "load_dotenv()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
    "metadata": {},
    "outputs": [
     {
@@ -12,35 +35,67 @@
       "Using model: gpt-3.5-turbo\n",
       "Loading path from disk...\n"
      ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "'The first roster of the Avengers included Iron Man, Thor, Hulk, Ant-Man, and the Wasp.'"
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
     }
    ],
    "source": [
-    "from megabots import bot\n",
-    "from dotenv import load_dotenv\n",
-    "\n",
-    "load_dotenv()\n",
-    "\n",
-    "qnabot = bot(\"qna-over-docs\", index=\"./index.pkl\")"
+    "qnabot = bot(\"qna-over-docs\", index=\"./index.pkl\")\n",
+    "qnabot.ask(\"what was the first roster of the avengers?\")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 10,
    "metadata": {},
    "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Using model: gpt-3.5-turbo\n",
+      "Loading path from disk...\n"
+     ]
+    },
     {
      "data": {
       "text/plain": [
-       "'The document does not provide an answer to the question \"What is the meaning of life?\".\\nSOURCES:'"
+       "\"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": 2,
+     "execution_count": 10,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "qnabot.ask(\"What is the meaning of life?\")"
+    "prompt_template = \"\"\"\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",
+    "{context}\n",
+    "\n",
+    "Question: {question}\n",
+    "Helpful humorous answer:\"\"\"\n",
+    "\n",
+    "load_dotenv()\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"
    ]
   }
  ],

megabots/__init__.py

@@ -9,6 +9,12 @@ 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 (
+    CONDENSE_QUESTION_PROMPT,
+    QA_PROMPT,
+)
 
 load_dotenv()
 
@@ -17,10 +23,11 @@ class Bot:
     def __init__(
         self,
         model: str | None = None,
-        prompt: str | None = None,
+        prompt_template: str | None = None,
+        prompt_variables: list[str] | None = None,
         memory: str | None = None,
         index: str | None = None,
-        source: str | None = None,
+        sources: bool | None = False,
         verbose: bool = False,
         temperature: int = 0,
     ):
@@ -29,7 +36,24 @@ class Bot:
         self.load_or_create_index(index)
 
         # Load the question-answering chain for the selected model
-        self.chain = load_qa_with_sources_chain(self.llm, verbose=verbose)
+        self.chain = self.create_chain(
+            prompt_template, prompt_variables, verbose=verbose
+        )
+
+    def create_chain(
+        self,
+        prompt_template: str | None = None,
+        prompt_variables: list[str] | None = None,
+        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
+        )
+        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
@@ -43,6 +67,13 @@ class Bot:
 
     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_path: str):
@@ -100,10 +131,9 @@ SUPPORTED_MODELS = {}
 def bot(
     task: str | None = None,
     model: str | None = None,
-    prompt: str | None = None,
-    memory: str | None = None,
+    prompt_template: str | None = None,
+    prompt_variables: list[str] | None = None,
     index: str | None = None,
-    source: str | None = None,
     verbose: bool = False,
     temperature: int = 0,
     **kwargs,
@@ -135,6 +165,9 @@ def bot(
     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,
     )

tests/test_bots.py

@@ -19,7 +19,7 @@ def test_ask():
     # Assert that the answer contains the correct answer
     assert correct_answer in answer
     # Assert that the answer contains the sources
-    assert sources in answer
+    assert sources not in answer
 
 
 def test_save_load_index():