Support for ChatGPT (#65)

* Rename Chatbot class to Buster

* move gradio app to .py file

* add new completers class

* add OOD test

* update apps

* update OOD prompt

* Update prompt engineering

buster/apps/gradio_app.ipynb

@@ -1,138 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "4a6b2b70",
-   "metadata": {
-    "scrolled": true
-   },
-   "outputs": [],
-   "source": [
-    "%load_ext autoreload\n",
-    "%autoreload 2\n",
-    "\n",
-    "import gradio as gr\n",
-    "\n",
-    "from buster.chatbot import Chatbot, ChatbotConfig\n",
-    "\n",
-    "hf_transformers_cfg = ChatbotConfig(\n",
-    "    documents_file=\"../data/document_embeddings_huggingface.tar.gz\",\n",
-    "    unknown_prompt=\"This doesn't seem to be related to the huggingface library. I am not sure how to answer.\",\n",
-    "    embedding_model=\"text-embedding-ada-002\",\n",
-    "    top_k=3,\n",
-    "    thresh=0.7,\n",
-    "    max_words=3000,\n",
-    "    completion_kwargs={\n",
-    "        \"engine\": \"text-davinci-003\",\n",
-    "        \"max_tokens\": 500,\n",
-    "    },\n",
-    "    link_format=\"gradio\",\n",
-    "    response_footnote=\"I'm a bot 🤖 trained to answer huggingface 🤗 transformers questions. My answers aren't always perfect.\",\n",
-    "    text_before_prompt=\"\"\"You are a slack chatbot assistant answering technical questions about huggingface transformers, a library to train transformers in python.\n",
-    "Make sure to format your answers in Markdown format, including code block and snippets.\n",
-    "Do not include any links to urls or hyperlinks in your answers.\n",
-    "\n",
-    "If you do not know the answer to a question, or if it is completely irrelevant to the library usage, simply reply with:\n",
-    "\n",
-    "'This doesn't seem to be related to the huggingface library.'\n",
-    "\n",
-    "For example:\n",
-    "\n",
-    "What is the meaning of life for huggingface?\n",
-    "\n",
-    "This doesn't seem to be related to the huggingface library.\n",
-    "\n",
-    "Now answer the following question:\n",
-    "\"\"\",\n",
-    ")\n",
-    "hf_transformers_chatbot = Chatbot(hf_transformers_cfg)\n",
-    "\n",
-    "def chat(question, history):\n",
-    "    history = history or []\n",
-    "    \n",
-    "    answer = hf_transformers_chatbot.process_input(question)\n",
-    "    \n",
-    "    # formatting hack for code blocks to render properly every time\n",
-    "    answer = answer.replace(\"```\", \"\\n```\\n\")\n",
-    "\n",
-    "    history.append((question, answer))\n",
-    "    return history, history\n",
-    "\n",
-    "\n",
-    "\n",
-    "block = gr.Blocks(css=\".gradio-container {background-color: lightgray}\")\n",
-    "\n",
-    "with block:\n",
-    "    with gr.Row():\n",
-    "        gr.Markdown(\"<h3><center>Buster 🤖: A Question-Answering Bot for Huggingface 🤗 Transformers </center></h3>\")\n",
-    "\n",
-    "\n",
-    "    chatbot = gr.Chatbot()\n",
-    "\n",
-    "    with gr.Row():\n",
-    "        message = gr.Textbox(\n",
-    "            label=\"What's your question?\",\n",
-    "            placeholder=\"What kind of model should I use for sentiment analysis?\",\n",
-    "            lines=1,\n",
-    "        )\n",
-    "        submit = gr.Button(value=\"Send\", variant=\"secondary\").style(full_width=False)\n",
-    "\n",
-    "    gr.Examples(\n",
-    "        examples=[\n",
-    "            \"What kind of models should I use for images and text?\",\n",
-    "            \"When should I finetune a model vs. training it form scratch?\",\n",
-    "            \"How can I deploy my trained huggingface model?\",\n",
-    "            \"Can you give me some python code to quickly finetune a model on my sentiment analysis dataset?\",\n",
-    "        ],\n",
-    "        inputs=message,\n",
-    "    )\n",
-    "\n",
-    "    gr.Markdown(\n",
-    "    \"\"\"This simple application uses GPT to search the huggingface 🤗 transformers docs and answer questions.\n",
-    "    For more info on huggingface transformers view the [full documentation.](https://huggingface.co/docs/transformers/index).\"\"\" \n",
-    "    )\n",
-    "\n",
-    "\n",
-    "    gr.HTML(\n",
-    "        \"️<center> Created with ❤️ by @jerpint and @hadrienbertrand\"\n",
-    "    )\n",
-    "\n",
-    "    state = gr.State()\n",
-    "    agent_state = gr.State()\n",
-    "\n",
-    "    submit.click(chat, inputs=[message, state], outputs=[chatbot, state])\n",
-    "    message.submit(chat, inputs=[message, state], outputs=[chatbot, state])\n",
-    "\n",
-    "\n",
-    "block.launch(debug=True)"
-   ]
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "buster",
-   "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.10.9"
-  },
-  "vscode": {
-   "interpreter": {
-    "hash": "bfa91706490f6a3314a87f4853806d905e46027cd889e58fcad4739e8600f624"
-   }
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}

buster/apps/gradio_app.py

@@ -0,0 +1,87 @@
+import gradio as gr
+
+from buster.buster import Buster, BusterConfig
+
+buster_cfg = BusterConfig(
+    documents_file="../data/document_embeddings_huggingface.tar.gz",
+    unknown_prompt="I'm sorry, but I am an AI language model trained to assist with questions related to the huggingface transformers library. I cannot answer that question as it is not relevant to the library or its usage. Is there anything else I can assist you with?",
+    embedding_model="text-embedding-ada-002",
+    top_k=3,
+    thresh=0.7,
+    max_words=3000,
+    completer_cfg={
+        "name": "ChatGPT",
+        "text_before_prompt": (
+            """You are a slack chatbot assistant answering technical questions about huggingface transformers, a library to train transformers in python. """
+            """Make sure to format your answers in Markdown format, including code block and snippets. """
+            """Do not include any links to urls or hyperlinks in your answers. """
+            """If you do not know the answer to a question, or if it is completely irrelevant to the library usage, let the user know you cannot answer with this response:\n"""
+            """'I'm sorry, but I am an AI language model trained to assist with questions related to the huggingface transformers library. I cannot answer that question as it is not relevant to the library or its usage. Is there anything else I can assist you with?'"""
+            """For example:\n"""
+            """What is the meaning of life for huggingface?\n"""
+            """I'm sorry, but I am an AI language model trained to assist with questions related to the huggingface transformers library. I cannot answer that question as it is not relevant to the library or its usage. Is there anything else I can assist you with?"""
+            """Now answer the following question:\n"""
+        ),
+        "text_before_documents": "Only use these documents as reference:\n",
+        "completion_kwargs": {
+            "model": "gpt-3.5-turbo",
+        },
+    },
+    response_format="gradio",
+)
+buster = Buster(buster_cfg)
+
+
+def chat(question, history):
+    history = history or []
+
+    answer = buster.process_input(question)
+
+    # formatting hack for code blocks to render properly every time
+    answer = answer.replace("```", "\n```\n")
+
+    history.append((question, answer))
+    return history, history
+
+
+block = gr.Blocks(css=".gradio-container {background-color: lightgray}")
+
+with block:
+    with gr.Row():
+        gr.Markdown("<h3><center>Buster 🤖: A Question-Answering Bot for Huggingface 🤗 Transformers </center></h3>")
+
+    chatbot = gr.Chatbot()
+
+    with gr.Row():
+        message = gr.Textbox(
+            label="What's your question?",
+            placeholder="What kind of model should I use for sentiment analysis?",
+            lines=1,
+        )
+        submit = gr.Button(value="Send", variant="secondary").style(full_width=False)
+
+    gr.Examples(
+        examples=[
+            "What kind of models should I use for images and text?",
+            "When should I finetune a model vs. training it form scratch?",
+            "How can I deploy my trained huggingface model?",
+            "Can you give me some python code to quickly finetune a model on my sentiment analysis dataset?",
+        ],
+        inputs=message,
+    )
+
+    gr.Markdown(
+        """This simple application uses GPT to search the huggingface 🤗 transformers docs and answer questions.
+    For more info on huggingface transformers view the [full documentation.](https://huggingface.co/docs/transformers/index)."""
+    )
+
+    gr.HTML("️<center> Created with ❤️ by @jerpint and @hadrienbertrand")
+
+    state = gr.State()
+    agent_state = gr.State()
+
+    submit.click(chat, inputs=[message, state], outputs=[chatbot, state])
+    message.submit(chat, inputs=[message, state], outputs=[chatbot, state])
+
+
+block.launch(debug=True)

buster/apps/slackbot.py

@@ -3,7 +3,7 @@ import os
 
 from slack_bolt import App
 
-from buster.chatbot import Chatbot, ChatbotConfig
+from buster.buster import Buster, BusterConfig
 
 logger = logging.getLogger(__name__)
 logging.basicConfig(level=logging.INFO)
@@ -14,137 +14,122 @@ ORION_CHANNEL = "C04LYHGUYB0"
 PYTORCH_CHANNEL = "C04MEK6N882"
 HF_TRANSFORMERS_CHANNEL = "C04NJNCJWHE"
 
-mila_doc_cfg = ChatbotConfig(
+mila_doc_cfg = BusterConfig(
     documents_file="../data/document_embeddings_mila.tar.gz",
     unknown_prompt="This doesn't seem to be related to cluster usage.",
     embedding_model="text-embedding-ada-002",
     top_k=3,
     thresh=0.7,
     max_words=3000,
-    completion_kwargs={
-        "engine": "text-davinci-003",
-        "max_tokens": 200,
-    },
-    separator="\n",
     response_format="slack",
-    response_footnote="""I'm a bot 🤖 and not always perfect.
-    For more info, view the full documentation here (https://docs.mila.quebec/) or contact [email protected]
-    """,
-    text_before_prompt="""
-    You are a slack chatbot assistant answering technical questions about a cluster.
-    Make sure to format your answers in Markdown format, including code block and snippets.
-    Do not include any links to urls or hyperlinks in your answers.
-
-    If you do not know the answer to a question, or if it is completely irrelevant to cluster usage, simply reply with:
-
-    'This doesn't seem to be related to cluster usage.'
-
-    For example:
-
-    What is the meaning of life on the cluster?
-
-    This doesn't seem to be related to cluster usage.
-
-    Now answer the following question:
-    """,
+    completer_cfg={
+        "name": "ChatGPT",
+        "text_before_prompt": (
+            """You are a slack chatbot assistant answering technical questions about the mila cluster. """
+            """Make sure to format your answers in Markdown format, including code block and snippets. """
+            """Do not include any links to urls or hyperlinks in your answers. """
+            """If you do not know the answer to a question, or if it is completely irrelevant to the library usage, simply reply with: """
+            """'This doesn't seem to be related to the pytorch library.'\n"""
+            """For example:\n"""
+            """What is the meaning of life for pytorch?\n"""
+            """This doesn't seem to be related to the pytorch library.\n"""
+            """Now answer the following question:\n"""
+        ),
+        "text_before_documents": "Only use these documents as reference:\n",
+        "completion_kwargs": {
+            "model": "gpt-3.5-turbo",
+        },
+    },
 )
-mila_doc_chatbot = Chatbot(mila_doc_cfg)
+mila_doc_chatbot = Buster(mila_doc_cfg)
 
-orion_cfg = ChatbotConfig(
+orion_cfg = BusterConfig(
     documents_file="../data/document_embeddings_orion.tar.gz",
     unknown_prompt="This doesn't seem to be related to the orion library. I am not sure how to answer.",
     embedding_model="text-embedding-ada-002",
     top_k=3,
     thresh=0.7,
     max_words=3000,
-    completion_kwargs={
-        "engine": "text-davinci-003",
-        "max_tokens": 200,
+    completer_cfg={
+        "name": "ChatGPT",
+        "text_before_prompt": (
+            """You are a slack chatbot assistant answering technical questions about orion, a hyperparameter optimization library written in python. """
+            """Make sure to format your answers in Markdown format, including code block and snippets. """
+            """Do not include any links to urls or hyperlinks in your answers. """
+            """If you do not know the answer to a question, or if it is completely irrelevant to the library usage, simply reply with: """
+            """'This doesn't seem to be related to the pytorch library.'\n"""
+            """For example:\n"""
+            """What is the meaning of life for pytorch?\n"""
+            """This doesn't seem to be related to the pytorch library.\n"""
+            """Now answer the following question:\n"""
+        ),
+        "text_before_documents": "Only use these documents as reference:\n",
+        "completion_kwargs": {
+            "model": "gpt-3.5-turbo",
+        },
     },
-    separator="\n",
     response_format="slack",
-    text_before_prompt="""You are a slack chatbot assistant answering technical questions about orion, a hyperparameter optimization library written in python.
-    Make sure to format your answers in Markdown format, including code block and snippets.
-    Do not include any links to urls or hyperlinks in your answers.
-
-    If you do not know the answer to a question, or if it is completely irrelevant to the library usage, simply reply with:
-
-    'This doesn't seem to be related to the orion library.'
-
-    For example:
-
-    What is the meaning of life for orion?
-
-    This doesn't seem to be related to the orion library.
-
-    Now answer the following question:
-    """,
 )
-orion_chatbot = Chatbot(orion_cfg)
+orion_chatbot = Buster(orion_cfg)
 
-pytorch_cfg = ChatbotConfig(
+pytorch_cfg = BusterConfig(
     documents_file="../data/document_embeddings_pytorch.tar.gz",
     unknown_prompt="This doesn't seem to be related to the pytorch library. I am not sure how to answer.",
     embedding_model="text-embedding-ada-002",
     top_k=3,
     thresh=0.7,
     max_words=3000,
-    completion_kwargs={
-        "engine": "text-davinci-003",
-        "max_tokens": 500,
+    completer_cfg={
+        "name": "ChatGPT",
+        "text_before_prompt": (
+            """You are a slack chatbot assistant answering technical questions about pytorch, a library to train neural networks written in python. """
+            """Make sure to format your answers in Markdown format, including code block and snippets. """
+            """Do not include any links to urls or hyperlinks in your answers. """
+            """If you do not know the answer to a question, or if it is completely irrelevant to the library usage, simply reply with: """
+            """'This doesn't seem to be related to the pytorch library.'\n"""
+            """For example:\n"""
+            """What is the meaning of life for pytorch?\n"""
+            """This doesn't seem to be related to the pytorch library.\n"""
+            """Now answer the following question:\n"""
+        ),
+        "text_before_documents": "Only use these documents as reference:\n",
+        "completion_kwargs": {
+            "model": "gpt-3.5-turbo",
+        },
     },
-    separator="\n",
     response_format="slack",
-    text_before_prompt="""You are a slack chatbot assistant answering technical questions about pytorch, a library to train neural networks written in python.
-    Make sure to format your answers in Markdown format, including code block and snippets.
-    Do not include any links to urls or hyperlinks in your answers.
-
-    If you do not know the answer to a question, or if it is completely irrelevant to the library usage, simply reply with:
-
-    'This doesn't seem to be related to the pytorch library.'
-
-    For example:
-
-    What is the meaning of life for pytorch?
-
-    This doesn't seem to be related to the pytorch library.
-
-    Now answer the following question:
-    """,
 )
-pytorch_chatbot = Chatbot(pytorch_cfg)
+pytorch_chatbot = Buster(pytorch_cfg)
 
-hf_transformers_cfg = ChatbotConfig(
+hf_transformers_cfg = BusterConfig(
     documents_file="../data/document_embeddings_huggingface.tar.gz",
-    unknown_prompt="This doesn't seem to be related to the huggingface library. I am not sure how to answer.",
+    unknown_prompt="I'm sorry, but I am an AI language model trained to assist with questions related to the huggingface transformers library. I cannot answer that question as it is not relevant to the library or its usage. Is there anything else I can assist you with?",
     embedding_model="text-embedding-ada-002",
     top_k=3,
     thresh=0.7,
     max_words=3000,
-    completion_kwargs={
-        "engine": "text-davinci-003",
-        "max_tokens": 500,
+    completer_cfg={
+        "name": "ChatGPT",
+        "text_before_prompt": (
+            """You are a slack chatbot assistant answering technical questions about huggingface transformers, a library to train transformers in python. """
+            """Make sure to format your answers in Markdown format, including code block and snippets. """
+            """Do not include any links to urls or hyperlinks in your answers. """
+            """If you do not know the answer to a question, or if it is completely irrelevant to the library usage, let the user know you cannot answer. """
+            """For example:\n"""
+            """What is the meaning of life for huggingface?\n"""
+            """This doesn't seem to be related to the huggingface library.\n"""
+            """I'm sorry, but I am an AI language model trained to assist with questions related to the huggingface transformers library. I cannot answer that question as it is not relevant to the library or its usage. Is there anything else I can assist you with?"""
+            """"""
+            """Now answer the following question:\n"""
+        ),
+        "text_before_documents": "Only use these documents as reference:\n",
+        "completion_kwargs": {
+            "model": "gpt-3.5-turbo",
+        },
     },
-    separator="\n",
     response_format="slack",
-    text_before_prompt="""You are a slack chatbot assistant answering technical questions about huggingface transformers, a library to train transformers in python.
-    Make sure to format your answers in Markdown format, including code block and snippets.
-    Do not include any links to urls or hyperlinks in your answers.
-
-    If you do not know the answer to a question, or if it is completely irrelevant to the library usage, simply reply with:
-
-    'This doesn't seem to be related to the huggingface library.'
-
-    For example:
-
-    What is the meaning of life for huggingface?
-
-    This doesn't seem to be related to the huggingface library.
-
-    Now answer the following question:
-    """,
 )
-hf_transformers_chatbot = Chatbot(hf_transformers_cfg)
+hf_transformers_chatbot = Buster(hf_transformers_cfg)
 
 # TODO: eventually move this to a factory of sorts
 # Put all the bots in a dict by channel

buster/{chatbot.py → buster.py}

@@ -1,14 +1,11 @@
 import logging
-import os
 from dataclasses import dataclass, field
-from typing import Iterable
 
 import numpy as np
-import openai
 import pandas as pd
-import promptlayer
 from openai.embeddings_utils import cosine_similarity, get_embedding
 
+from buster.completers import get_completer
 from buster.documents import get_documents_manager_from_extension
 from buster.formatter import (
     Response,
@@ -20,19 +17,9 @@ from buster.formatter import (
 logger = logging.getLogger(__name__)
 logging.basicConfig(level=logging.INFO)
 
-# Check if an API key exists for promptlayer, if it does, use it
-promptlayer_api_key = os.environ.get("PROMPTLAYER_API_KEY")
-if promptlayer_api_key:
-    logger.info("Enabling prompt layer...")
-    promptlayer.api_key = promptlayer_api_key
-
-    # replace openai with the promptlayer wrapper
-    openai = promptlayer.openai
-    openai.api_key = os.environ.get("OPENAI_API_KEY")
-
 
 @dataclass
-class ChatbotConfig:
+class BusterConfig:
     """Configuration object for a chatbot.
 
     documents_csv: Path to the csv file containing the documents and their embeddings.
@@ -54,29 +41,32 @@ class ChatbotConfig:
     thresh: float = 0.7
     max_words: int = 3000
     unknown_threshold: float = 0.9  # set to 0 to deactivate
-
-    completion_kwargs: dict = field(
+    completer_cfg: dict = field(
+        # TODO: Put all this in its own config with sane defaults?
         default_factory=lambda: {
-            "engine": "text-davinci-003",
-            "max_tokens": 200,
-            "temperature": None,
-            "top_p": None,
-            "frequency_penalty": 1,
-            "presence_penalty": 1,
+            "name": "GPT3",
+            "text_before_documents": "You are a chatbot answering questions.\n",
+            "text_before_prompt": "Answer the following question:\n",
+            "completion_kwargs": {
+                "engine": "text-davinci-003",
+                "max_tokens": 200,
+                "temperature": None,
+                "top_p": None,
+                "frequency_penalty": 1,
+                "presence_penalty": 1,
+            },
         }
     )
-    separator: str = "\n"
     response_format: str = "slack"
     unknown_prompt: str = "I Don't know how to answer your question."
-    text_before_documents: str = "You are a chatbot answering questions.\n"
-    text_before_prompt: str = "Answer the following question:\n"
     response_footnote: str = "I'm a bot 🤖 and not always perfect."
 
 
-class Chatbot:
-    def __init__(self, cfg: ChatbotConfig):
+class Buster:
+    def __init__(self, cfg: BusterConfig):
         # TODO: right now, the cfg is being passed as an omegaconf, is this what we want?
         self.cfg = cfg
+        self.completer = get_completer(cfg.completer_cfg)
         self._init_documents()
         self._init_unk_embedding()
         self._init_response_formatter()
@@ -141,67 +131,21 @@ class Chatbot:
 
         return documents_str
 
-    def prepare_prompt(
+    def add_sources(
         self,
-        question: str,
+        response,
         matched_documents: pd.DataFrame,
-        text_before_prompt: str,
-        text_before_documents: str,
-    ) -> str:
-        """
-        Prepare the prompt with prompt engineering.
-        """
-        documents_str: str = self.prepare_documents(matched_documents, max_words=self.cfg.max_words)
-        return text_before_documents + documents_str + text_before_prompt + question
-
-    def get_gpt_response(self, **completion_kwargs) -> Response:
-        # Call the API to generate a response
-        logger.info(f"querying GPT...")
-        try:
-            response = openai.Completion.create(**completion_kwargs)
-        except Exception as e:
-            # log the error and return a generic response instead.
-            logger.exception("Error connecting to OpenAI API. See traceback:")
-            return Response("", True, "We're having trouble connecting to OpenAI right now... Try again soon!")
-
-        text = response["choices"][0]["text"]
-        return Response(text)
-
-    def generate_response(
-        self, prompt: str, matched_documents: pd.DataFrame, unknown_prompt: str
-    ) -> tuple[Response, Iterable[Source]]:
-        """
-        Generate a response based on the retrieved documents.
-        """
-        if len(matched_documents) == 0:
-            # No matching documents were retrieved, return
-            sources = tuple()
-            return Response(unknown_prompt), sources
-
-        logger.info(f"Prompt:  {prompt}")
-        response = self.get_gpt_response(prompt=prompt, **self.cfg.completion_kwargs)
-        if response:
-            logger.info(f"GPT Response:\n{response.text}")
-            relevant = self.check_response_relevance(
-                response=response.text,
-                engine=self.cfg.embedding_model,
-                unk_embedding=self.unk_embedding,
-                unk_threshold=self.cfg.unknown_threshold,
-            )
-            if relevant:
-                sources = (
-                    Source(dct["source"], dct["url"], dct["similarity"])
-                    for dct in matched_documents.to_dict(orient="records")
-                )
-            else:
-                # Override the answer with a generic unknown prompt, without sources.
-                response = Response(text=self.cfg.unknown_prompt)
-                sources = tuple()
-
-        return response, sources
+        unknown_prompt: str,
+    ):
+        logger.info(f"GPT Response:\n{response.text}")
+        sources = (
+            Source(dct["source"], dct["url"], dct["similarity"]) for dct in matched_documents.to_dict(orient="records")
+        )
+
+        return sources
 
     def check_response_relevance(
-        self, response: str, engine: str, unk_embedding: np.array, unk_threshold: float
+        self, completion: str, engine: str, unk_embedding: np.array, unk_threshold: float
     ) -> bool:
         """Check to see if a response is relevant to the chatbot's knowledge or not.
 
@@ -211,7 +155,7 @@ class Chatbot:
         set the unk_threshold to 0 to essentially turn off this feature.
         """
         response_embedding = get_embedding(
-            response,
+            completion,
             engine=engine,
         )
         score = cosine_similarity(response_embedding, unk_embedding)
@@ -220,29 +164,45 @@ class Chatbot:
         # Likely that the answer is meaningful, add the top sources
         return score < unk_threshold
 
-    def process_input(self, question: str, formatter: ResponseFormatter = None) -> str:
+    def process_input(self, user_input: str, formatter: ResponseFormatter = None) -> str:
         """
         Main function to process the input question and generate a formatted output.
         """
 
-        logger.info(f"User Question:\n{question}")
+        logger.info(f"User Input:\n{user_input}")
 
         # We make sure there is always a newline at the end of the question to avoid completing the question.
-        if not question.endswith("\n"):
-            question += "\n"
+        if not user_input.endswith("\n"):
+            user_input += "\n"
 
         matched_documents = self.rank_documents(
-            query=question,
+            query=user_input,
             top_k=self.cfg.top_k,
             thresh=self.cfg.thresh,
             engine=self.cfg.embedding_model,
         )
-        prompt = self.prepare_prompt(
-            question=question,
-            matched_documents=matched_documents,
-            text_before_prompt=self.cfg.text_before_prompt,
-            text_before_documents=self.cfg.text_before_documents,
+
+        if len(matched_documents) == 0:
+            response = Response("I did not find any sources to answer your question.")
+            sources = tuple()
+            return self.response_formatter(response, sources)
+
+        # generate a completion
+        documents: str = self.prepare_documents(matched_documents, max_words=self.cfg.max_words)
+        response = self.completer.generate_response(user_input, documents)
+        sources = self.add_sources(response, matched_documents, self.cfg.unknown_prompt)
+
+        # check for relevance
+        relevant = self.check_response_relevance(
+            completion=response.text,
+            engine=self.cfg.embedding_model,
+            unk_embedding=self.unk_embedding,
+            unk_threshold=self.cfg.unknown_threshold,
         )
-        response, sources = self.generate_response(prompt, matched_documents, self.cfg.unknown_prompt)
+        if not relevant:
+            # answer generated was the chatbot saying it doesn't know how to answer
+            # override completion with generic "I don't know"
+            response = Response(text=self.cfg.unknown_prompt)
+            sources = tuple()
 
         return self.response_formatter(response, sources)

buster/completers/__init__.py

@@ -0,0 +1,7 @@
+from .base import ChatGPTCompleter, GPT3Completer, get_completer
+
+__all__ = [
+    get_completer,
+    GPT3Completer,
+    ChatGPTCompleter,
+]

buster/completers/base.py

@@ -0,0 +1,99 @@
+import logging
+import os
+from abc import ABC, abstractmethod
+
+import openai
+import promptlayer
+
+from buster.formatter.base import Response
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+# Check if an API key exists for promptlayer, if it does, use it
+promptlayer_api_key = os.environ.get("PROMPTLAYER_API_KEY")
+if promptlayer_api_key:
+    logger.info("Enabling prompt layer...")
+    promptlayer.api_key = promptlayer_api_key
+
+    # replace openai with the promptlayer wrapper
+    openai = promptlayer.openai
+    openai.api_key = os.environ.get("OPENAI_API_KEY")
+
+
+class Completer(ABC):
+    def __init__(self, cfg):
+        self.cfg = cfg
+
+    @abstractmethod
+    def complete(self, prompt) -> str:
+        ...
+
+    def generate_response(self, user_input, documents) -> Response:
+        # Call the API to generate a response
+        prompt = self.prepare_prompt(user_input, documents)
+        name = self.cfg["name"]
+        logger.info(f"querying model {name}...")
+        logger.info(f"{prompt=}")
+        try:
+            completion_kwargs = self.cfg["completion_kwargs"]
+            completion = self.complete(prompt=prompt, **completion_kwargs)
+        except Exception as e:
+            # log the error and return a generic response instead.
+            logger.exception("Error connecting to OpenAI API. See traceback:")
+            return Response("", True, "We're having trouble connecting to OpenAI right now... Try again soon!")
+
+        return Response(completion)
+
+
+class GPT3Completer(Completer):
+    def prepare_prompt(
+        self,
+        user_input: str,
+        documents: str,
+    ) -> str:
+        """
+        Prepare the prompt with prompt engineering.
+        """
+        text_before_docs = self.cfg["text_before_documents"]
+        text_before_prompt = self.cfg["text_before_prompt"]
+        return text_before_docs + documents + text_before_prompt + user_input
+
+    def complete(self, prompt, **completion_kwargs):
+        response = openai.Completion.create(prompt=prompt, **completion_kwargs)
+        return response["choices"][0]["text"]
+
+
+class ChatGPTCompleter(Completer):
+    def prepare_prompt(
+        self,
+        user_input: str,
+        documents: str,
+    ) -> list:
+        """
+        Prepare the prompt with prompt engineering.
+        """
+        text_before_docs = self.cfg["text_before_documents"]
+        text_before_prompt = self.cfg["text_before_prompt"]
+        prompt = [
+            {"role": "system", "content": text_before_docs + documents + text_before_prompt},
+            {"role": "user", "content": user_input},
+        ]
+        return prompt
+
+    def complete(self, prompt, **completion_kwargs) -> str:
+        response = openai.ChatCompletion.create(
+            messages=prompt,
+            **completion_kwargs,
+        )
+
+        return response["choices"][0]["message"]["content"]
+
+
+def get_completer(completer_cfg):
+    name = completer_cfg["name"]
+    completers = {
+        "GPT3": GPT3Completer,
+        "ChatGPT": ChatGPTCompleter,
+    }
+    return completers[name](completer_cfg)

tests/test_chatbot.py

@@ -4,7 +4,7 @@ from pathlib import Path
 import numpy as np
 import pandas as pd
 
-from buster.chatbot import Chatbot, ChatbotConfig
+from buster.buster import Buster, BusterConfig
 from buster.documents import DocumentsManager
 
 TEST_DATA_DIR = Path(__file__).resolve().parent / "data"
@@ -39,61 +39,136 @@ class DocumentsMock(DocumentsManager):
         return self.documents
 
 
-def test_chatbot_real_data():
-    hf_transformers_cfg = ChatbotConfig(
-        documents_file=DOCUMENTS_FILE,
+def test_chatbot_mock_data(tmp_path, monkeypatch):
+    gpt_expected_answer = "this is GPT answer"
+    monkeypatch.setattr("buster.buster.get_documents_manager_from_extension", lambda filepath: DocumentsMock)
+    monkeypatch.setattr("buster.buster.get_embedding", lambda x, engine: get_fake_embedding())
+    monkeypatch.setattr("openai.Completion.create", lambda **kwargs: {"choices": [{"text": gpt_expected_answer}]})
+
+    hf_transformers_cfg = BusterConfig(
+        documents_file=tmp_path / "not_a_real_file.tar.gz",
         unknown_prompt="This doesn't seem to be related to the huggingface library. I am not sure how to answer.",
         embedding_model="text-embedding-ada-002",
         top_k=3,
         thresh=0.7,
         max_words=3000,
-        completion_kwargs={
-            "temperature": 0,
-            "engine": "text-davinci-003",
-            "max_tokens": 100,
+        response_format="slack",
+        completer_cfg={
+            "name": "GPT3",
+            "text_before_prompt": (
+                """You are a slack chatbot assistant answering technical questions about huggingface transformers, a library to train transformers in python.\n"""
+                """Make sure to format your answers in Markdown format, including code block and snippets.\n"""
+                """Do not include any links to urls or hyperlinks in your answers.\n\n"""
+                """Now answer the following question:\n"""
+            ),
+            "text_before_documents": "",
+            "completion_kwargs": {
+                "engine": "text-davinci-003",
+                "max_tokens": 200,
+                "temperature": None,
+                "top_p": None,
+                "frequency_penalty": 1,
+                "presence_penalty": 1,
+            },
         },
+    )
+    buster = Buster(hf_transformers_cfg)
+    answer = buster.process_input("What is a transformer?")
+    assert isinstance(answer, str)
+    assert answer.startswith(gpt_expected_answer)
+
+
+def test_chatbot_real_data__chatGPT():
+    hf_transformers_cfg = BusterConfig(
+        documents_file=DOCUMENTS_FILE,
+        unknown_prompt="I'm sorry, but I am an AI language model trained to assist with questions related to the huggingface transformers library. I cannot answer that question as it is not relevant to the library or its usage. Is there anything else I can assist you with?",
+        embedding_model="text-embedding-ada-002",
+        top_k=3,
+        thresh=0.7,
+        max_words=3000,
         response_format="slack",
-        text_before_prompt=(
-            """You are a slack chatbot assistant answering technical questions about huggingface transformers, a library to train transformers in python.\n"""
-            """Make sure to format your answers in Markdown format, including code block and snippets.\n"""
-            """Do not include any links to urls or hyperlinks in your answers.\n\n"""
-            """Now answer the following question:\n"""
-        ),
+        completer_cfg={
+            "name": "ChatGPT",
+            "text_before_prompt": (
+                """You are a slack chatbot assistant answering technical questions about huggingface transformers, a library to train transformers in python.\n"""
+                """Make sure to format your answers in Markdown format, including code block and snippets.\n"""
+                """Do not include any links to urls or hyperlinks in your answers.\n\n"""
+                """Now answer the following question:\n"""
+            ),
+            "text_before_documents": "Only use these documents as reference:\n",
+            "completion_kwargs": {
+                "model": "gpt-3.5-turbo",
+            },
+        },
     )
-    chatbot = Chatbot(hf_transformers_cfg)
-    answer = chatbot.process_input("What is a transformer?")
+    buster = Buster(hf_transformers_cfg)
+    answer = buster.process_input("What is a transformer?")
     assert isinstance(answer, str)
 
 
-def test_chatbot_mock_data(tmp_path, monkeypatch):
-    gpt_expected_answer = "this is GPT answer"
-    monkeypatch.setattr("buster.chatbot.get_documents_manager_from_extension", lambda filepath: DocumentsMock)
-    monkeypatch.setattr("buster.chatbot.get_embedding", lambda x, engine: get_fake_embedding())
-    monkeypatch.setattr(
-        "buster.chatbot.openai.Completion.create", lambda **kwargs: {"choices": [{"text": gpt_expected_answer}]}
+def test_chatbot_real_data__chatGPT_OOD():
+    buster_cfg = BusterConfig(
+        documents_file=DOCUMENTS_FILE,
+        unknown_prompt="I'm sorry, but I am an AI language model trained to assist with questions related to the huggingface transformers library. I cannot answer that question as it is not relevant to the library or its usage. Is there anything else I can assist you with?",
+        embedding_model="text-embedding-ada-002",
+        top_k=3,
+        thresh=0.7,
+        max_words=3000,
+        completer_cfg={
+            "name": "ChatGPT",
+            "text_before_prompt": (
+                """You are a slack chatbot assistant answering technical questions about huggingface transformers, a library to train transformers in python. """
+                """Make sure to format your answers in Markdown format, including code block and snippets. """
+                """Do not include any links to urls or hyperlinks in your answers. """
+                """If you do not know the answer to a question, or if it is completely irrelevant to the library usage, let the user know you cannot answer. """
+                """Use this response: """
+                """I'm sorry, but I am an AI language model trained to assist with questions related to the huggingface transformers library. I cannot answer that question as it is not relevant to the library or its usage. Is there anything else I can assist you with?"""
+                """For example:\n"""
+                """What is the meaning of life for huggingface?\n"""
+                """I'm sorry, but I am an AI language model trained to assist with questions related to the huggingface transformers library. I cannot answer that question as it is not relevant to the library or its usage. Is there anything else I can assist you with?"""
+                """Now answer the following question:\n"""
+            ),
+            "text_before_documents": "Only use these documents as reference:\n",
+            "completion_kwargs": {
+                "model": "gpt-3.5-turbo",
+            },
+        },
+        response_format="gradio",
     )
+    buster = Buster(buster_cfg)
+    answer = buster.process_input("What is a good recipe for brocolli soup?")
+    assert isinstance(answer, str)
+    assert buster_cfg.unknown_prompt in answer
 
-    hf_transformers_cfg = ChatbotConfig(
-        documents_file=tmp_path / "not_a_real_file.tar.gz",
+
+def test_chatbot_real_data__GPT():
+    hf_transformers_cfg = BusterConfig(
+        documents_file=DOCUMENTS_FILE,
         unknown_prompt="This doesn't seem to be related to the huggingface library. I am not sure how to answer.",
         embedding_model="text-embedding-ada-002",
         top_k=3,
         thresh=0.7,
         max_words=3000,
-        completion_kwargs={
-            "temperature": 0,
-            "engine": "text-davinci-003",
-            "max_tokens": 100,
-        },
         response_format="slack",
-        text_before_prompt=(
-            """You are a slack chatbot assistant answering technical questions about huggingface transformers, a library to train transformers in python.\n"""
-            """Make sure to format your answers in Markdown format, including code block and snippets.\n"""
-            """Do not include any links to urls or hyperlinks in your answers.\n\n"""
-            """Now answer the following question:\n"""
-        ),
+        completer_cfg={
+            "name": "GPT3",
+            "text_before_prompt": (
+                """You are a slack chatbot assistant answering technical questions about huggingface transformers, a library to train transformers in python.\n"""
+                """Make sure to format your answers in Markdown format, including code block and snippets.\n"""
+                """Do not include any links to urls or hyperlinks in your answers.\n\n"""
+                """Now answer the following question:\n"""
+            ),
+            "text_before_documents": "",
+            "completion_kwargs": {
+                "engine": "text-davinci-003",
+                "max_tokens": 200,
+                "temperature": None,
+                "top_p": None,
+                "frequency_penalty": 1,
+                "presence_penalty": 1,
+            },
+        },
     )
-    chatbot = Chatbot(hf_transformers_cfg)
-    answer = chatbot.process_input("What is a transformer?")
+    buster = Buster(hf_transformers_cfg)
+    answer = buster.process_input("What is a transformer?")
     assert isinstance(answer, str)
-    assert answer.startswith(gpt_expected_answer)