😇 Test your code people

.gitignore

@@ -4,4 +4,5 @@ __pycache__
 qnabot.egg-info
 dist
 build
-**.pickle
+**.pickle
+.env

README.md

@@ -31,6 +31,7 @@ bot = QnABot(directory="./mydata", index="index.pickle")
 - [x] Save / load index to reduce costs (Open AI embedings are used to create the index)
 - [x] Local data source (directory of documents) or S3 data source
 - [x] FAISS for storing vectors / index
+- [ ] Support for other vector databases (e.g. Weaviate, Pinecone)
 - [ ] Customise prompt
 - [ ] Expose API
 - [ ] Support for LLaMA model

examples/example.ipynb

@@ -18,7 +18,7 @@
     "from qnabot.QnABot import QnABot\n",
     "import os, sys\n",
     "\n",
-    "os.environ[\"OPENAI_API_KEY\"] = \"you api key\"\n",
+    "os.environ[\"OPENAI_API_KEY\"] = \"your api key\"\n",
     "\n",
     "bot = QnABot(directory=\"./files\")"
    ]
@@ -40,7 +40,7 @@
     }
    ],
    "source": [
-    "bot.print_answer(\"what was the first roster of angers in comics?\")\n",
+    "bot.print_answer(\"what was the first roster of avengers in comics?\")\n",
     "bot.print_answer(\"Who is Vision?\")"
    ]
   }

qnabot/QnABot.py

@@ -73,7 +73,7 @@ class QnABot:
             )["output_text"]
         )
 
-    def get_answer(self, question, k=1):
+    def get_answer(self, question, k=1) -> str:
         # Retrieve the answer to the given question and return it
         input_documents = self.search_index.similarity_search(question, k=k)
         return self.chain(

qnabot/tests/test_QnABot.py

@@ -0,0 +1,44 @@
+import os
+import tempfile
+from qnabot import QnABot
+import pickle
+from langchain.vectorstores.faiss import FAISS
+from dotenv import load_dotenv
+
+load_dotenv()
+
+# Define test data
+test_directory = "./examples/files"
+test_question = "what was the first roster of avengers in comics?"
+correct_answer = "Iron Man, Thor, Hulk, Ant-Man"
+sources = "SOURCES:"
+
+
+def test_get_answer():
+    bot = QnABot(directory=test_directory)
+    answer = bot.get_answer(test_question)
+
+    # Assert that the answer contains the correct answer
+    assert correct_answer in answer
+    # Assert that the answer contains the sources
+    assert sources in answer
+
+
+def test_save_load_index():
+    # Create a temporary directory and file path for the test index
+    with tempfile.TemporaryDirectory() as temp_dir:
+        index_path = os.path.join(temp_dir, "test_index.pkl")
+
+        # Create a bot and save the index to the temporary file path
+        bot = QnABot(directory=test_directory, index=index_path)
+        bot.save_index(index_path)
+
+        # Load the saved index and assert that it is the same as the original index
+        with open(index_path, "rb") as f:
+            saved_index = pickle.load(f)
+        assert isinstance(saved_index, FAISS)
+
+        bot_with_predefined_index = QnABot(directory=test_directory, index=index_path)
+
+        # Assert that the bot returns the correct answer to the test question
+        assert correct_answer in bot_with_predefined_index.get_answer(test_question)