🤖 Enter Megabots!

.github/ISSUE_TEMPLATE/bug_report.md

@@ -1,10 +1,9 @@
 ---
 name: Bug report
 about: Create a report to help us improve
-title: ''
+title: ""
 labels: bug
 assignees: momegas
-
 ---
 
 **Describe the bug**
@@ -12,8 +11,9 @@ A clear and concise description of what the bug is.
 
 **To Reproduce**
 Steps to reproduce the behavior:
-1.  pip install qnabot
-2. ...
+
+1.  pip install megabots
+2.  ...
 
 **Code that breaks**
 If applicable, add the code that produced the bug to help explain your problem.
@@ -22,9 +22,10 @@ If applicable, add the code that produced the bug to help explain your problem.
 A clear and concise description of what you expected to happen.
 
 **Python version (please complete the following information):**
- - [ ] <= 3.9
- - [ ] 3.10
- - [ ] 3.11
+
+- [ ] <= 3.9
+- [ ] 3.10
+- [ ] 3.11
 
 **Additional context**
 Add any other context about the problem here.

.gitignore

@@ -1,7 +1,7 @@
 .venv
 __pycache__
 .pytest_cache
-qnabot.egg-info
+**.egg-info
 dist
 build
 **.pickle

Makefile

@@ -1,7 +1,7 @@
 # Define variables
 PYTHON=python
 PIP=pip
-PACKAGE=qnabot
+PACKAGE=megabots
 
 .PHONY: install test clean build publish
 

README.md

@@ -1,34 +1,35 @@
-# 🦾🤖🤳 Mega Bots
+# 🤖 Megabots
 
 [![Tests](https://github.com/momegas/qnabot/actions/workflows/python-package.yml/badge.svg)](https://github.com/momegas/qnabot/actions/workflows/python-package.yml)
 
-Here is an example of what you build with this library: [Demo](https://huggingface.co/spaces/momegas/megas-bot)
+🤖 Megabots provides State-of-the-art, production ready bots made mega-easy, so you don't have to build them from scratch 🤯 Create a bot, now 🫵
 
-🦾🤖🤳 Megabots provides ready made production ready bots so you don't have to build them from scratch 🤯
-
-Note: This is a work in progress. The API is not stable and will change.
+Note: This is a work in progress. The API might change.
 
 ```bash
 pip install megabots
 ```
 
 ```python
-from qnabot import QnABot
+from megabots import bot
 import os
 
 os.environ["OPENAI_API_KEY"] = "my key"
 
-# Create a bot 👉 with one line of code
-bot = QnABot(directory="./mydata")
+# Create a bot 👉 with one line of code. Automatically loads your data from ./index or index.pkl.
+qnabot = bot("qna-over-docs")
 
 # Ask a question
 answer = bot.ask("How do I use this bot?")
 
 # Save the index to save costs (GPT is used to create the index)
-bot.save_index("index.pickle")
+bot.save_index("index.pkl")
 
 # Load the index from a previous run
-bot = QnABot(directory="./mydata", index="index.pickle")
+qnabot = bot("qna-over-docs", index="./index.pkl")
+
+# Or create the index from a directory of documents
+qnabot = bot("qna-over-docs", index="./index")
 ```
 
 You can also create a FastAPI app that will expose the bot as an API using the create_app function.
@@ -36,9 +37,9 @@ Assuming you file is called `main.py` run `uvicorn main:app --reload` to run the
 You should then be able to visit `http://localhost:8000/docs` to see the API documentation.
 
 ```python
-from qnabot import QnABot, create_app
+from megabots import bot, create_api
 
-app = create_app(QnABot("./mydata"))
+app = create_app(bot("qna-over-docs"))
 ```
 
 You can expose a gradio UI for the bot using `create_interface` function.
@@ -46,9 +47,9 @@ Assuming your file is called `ui.py` run `gradio qnabot/ui.py` to run the UI loc
 You should then be able to visit `http://127.0.0.1:7860` to see the API documentation.
 
 ```python
-from qnabot import QnABot, create_interface
+from megabots import bot, create_interface
 
-demo = create_interface(QnABot("./mydata"))
+demo = create_interface(QnABot("qna-over-docs"))
 ```
 
 ### Features

example.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [
     {
@@ -15,13 +15,12 @@
     }
    ],
    "source": [
-    "from qnabot import QnABot\n",
+    "from megabots import bot\n",
     "from dotenv import load_dotenv\n",
     "\n",
     "load_dotenv()\n",
     "\n",
-    "bot = QnABot(directory=\"./examples/files\", index=\"./index.pkl\", verbose=True)\n",
-    "# bot.save_index(\"./index.pkl\")"
+    "qnabot = bot(\"qna-over-docs\", index=\"./index.pkl\")"
    ]
   },
   {
@@ -30,19 +29,18 @@
    "metadata": {},
    "outputs": [
     {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "The first roster of Avengers in comics included Iron Man, Thor, Hulk, Ant-Man, and the Wasp.\n",
-      "SOURCES: examples/files/facts.txt\n",
-      "Vision is an android superhero who was created by Ultron but ultimately joined the Avengers and became an important member of the team.\n",
-      "SOURCES: examples/files/facts.txt\n"
-     ]
+     "data": {
+      "text/plain": [
+       "'The document does not provide an answer to the question \"What is the meaning of life?\".\\nSOURCES:'"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
     }
    ],
    "source": [
-    "bot.print_answer(\"what was the first roster of avengers in comics?\")\n",
-    "bot.print_answer(\"Who is Vision?\")"
+    "qnabot.ask(\"What is the meaning of life?\")"
    ]
   }
  ],

qnabot/QnABot.py → megabots/__init__.py

@@ -4,23 +4,28 @@ from langchain.embeddings import OpenAIEmbeddings
 from langchain.document_loaders import DirectoryLoader, S3DirectoryLoader
 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
 
+load_dotenv()
 
-class QnABot:
+
+class Bot:
     def __init__(
         self,
-        directory: str | None = None,
-        index: str | None = None,
         model: str | None = None,
+        prompt: str | None = None,
+        memory: str | None = None,
+        index: str | None = None,
+        source: str | None = None,
         verbose: bool = False,
         temperature: int = 0,
     ):
-        # Initialize the QnABot by selecting a model, creating a loader,
-        # and loading or creating an index
         self.select_model(model, temperature)
-        self.create_loader(directory)
+        self.create_loader(index)
         self.load_or_create_index(index)
 
         # Load the question-answering chain for the selected model
@@ -36,40 +41,125 @@ class QnABot:
             print("Using model: text-davinci-003")
             self.llm = OpenAI(temperature=temperature)
 
-    def create_loader(self, directory: str | None):
-        if directory is None:
-            return
+    def create_loader(self, index: str | None):
         # Create a loader based on the provided directory (either local or S3)
-        if directory.startswith("s3://"):
-            self.loader = S3DirectoryLoader(directory)
-        else:
-            self.loader = DirectoryLoader(directory, recursive=True)
+        self.loader = DirectoryLoader(index, recursive=True)
 
-    def load_or_create_index(self, index_path: str | None):
+    def load_or_create_index(self, index_path: str):
         # Load an existing index from disk or create a new one if not available
-        if index_path is not None and os.path.exists(index_path):
+
+        # Is pickle
+        if index_path is not None and "pkl" in index_path or "pickle" in index_path:
             print("Loading path from disk...")
             with open(index_path, "rb") as f:
                 self.search_index = pickle.load(f)
-        else:
+            return
+
+        # Is directory
+        if index_path is not None and os.path.isdir(index_path):
             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 print_answer(self, question: str, k=1):
-        # Retrieve and print the answer to the given question
-        input_documents = self.search_index.similarity_search(question, k=k)
-        a = self.chain.run(input_documents=input_documents, question=question)
-        print(a)
-
     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",
+            "prompt": "",
+            "temperature": 0,
+            "index": "./files",
+        },
+    }
+}
+
+SUPPORTED_MODELS = {}
+
+
+def bot(
+    task: str | None = None,
+    model: str | None = None,
+    prompt: str | None = None,
+    memory: str | None = None,
+    index: str | None = None,
+    source: 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.
+        verbose (bool, optional): Verbocity. Defaults to False.
+
+    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"],
+        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: Bot, examples: list[list[str]] = []):
+    def ask(question: str):
+        return bot.ask(question)
+
+    interface = gr.Interface(
+        fn=ask,
+        inputs=gr.components.Textbox(lines=5, label="Question"),
+        outputs=gr.components.Textbox(lines=5, label="Answer"),
+        examples=examples,
+    )
+
+    return interface

{qnabot → megabots}/tests/test_api.py

@@ -1,9 +1,9 @@
 import json
 from fastapi.testclient import TestClient
-from qnabot import QnABot, create_app
+from megabots import bot, create_api
 
-bot = QnABot(directory="./examples/files")
-app = create_app(bot)
+qnabot = bot("qna-over-docs", index="./examples/files")
+app = create_api(qnabot)
 
 client = TestClient(app)
 

qnabot/tests/test_QnABot.py → megabots/tests/test_bots.py

@@ -1,6 +1,6 @@
 import os
 import tempfile
-from qnabot import QnABot
+from megabots import bot
 import pickle
 from langchain.vectorstores.faiss import FAISS
 from dotenv import load_dotenv
@@ -15,8 +15,8 @@ sources = "SOURCES:"
 
 
 def test_ask():
-    bot = QnABot(directory=test_directory)
-    answer = bot.ask(test_question)
+    qnabot = bot("qna-over-docs", index=test_directory)
+    answer = qnabot.ask(test_question)
 
     # Assert that the answer contains the correct answer
     assert correct_answer in answer
@@ -30,15 +30,15 @@ def test_save_load_index():
         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)
+        qnabot = bot("qna-over-docs", index=test_directory)
+        qnabot.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)
+        bot_with_predefined_index = bot("qna-over-docs", index=index_path)
 
         # Assert that the bot returns the correct answer to the test question
         assert correct_answer in bot_with_predefined_index.ask(test_question)

{qnabot → megabots}/tests/test_ui.py

@@ -1,9 +1,9 @@
 import gradio as gr
-from qnabot import create_interface
+from megabots import create_interface
 
 
 def test_create_interface():
-    # create a mock QnABot object
+    # create a mock Bot object
     class MockBot:
         def ask(self, question: str):
             return "Answer"

qnabot/__init__.py

@@ -1,3 +0,0 @@
-from .QnABot import QnABot
-from .api import create_app
-from .ui import create_interface

qnabot/api.py

@@ -1,13 +0,0 @@
-from fastapi import FastAPI
-from qnabot import QnABot
-
-
-def create_app(bot: QnABot):
-    app = FastAPI()
-
-    @app.get("/v1/ask/{question}")
-    async def ask(question: str):
-        answer = bot.ask(question)
-        return {"answer": answer}
-
-    return app

qnabot/ui.py

@@ -1,16 +0,0 @@
-import gradio as gr
-from qnabot import QnABot
-
-
-def create_interface(bot: QnABot, examples: list[list[str]] = []):
-    def ask(question: str):
-        return bot.ask(question)
-
-    interface = gr.Interface(
-        fn=ask,
-        inputs=gr.components.Textbox(lines=5, label="Question"),
-        outputs=gr.components.Textbox(lines=5, label="Answer"),
-        examples=examples,
-    )
-
-    return interface