PR: DocumentsManager interface (#57)

buster/chatbot.py

@@ -9,7 +9,7 @@ import pandas as pd
 import promptlayer
 from openai.embeddings_utils import cosine_similarity, get_embedding
 
-from buster.docparser import read_documents
+from buster.documents import get_documents_manager_from_extension
 from buster.formatter import Formatter, HTMLFormatter, MarkdownFormatter, SlackFormatter
 from buster.formatter.base import Response, Source
 
@@ -47,7 +47,7 @@ class ChatbotConfig:
     text_after_response: Generic response to add the the chatbot's reply.
     """
 
-    documents_file: str = "buster/data/document_embeddings.csv"
+    documents_file: str = "buster/data/document_embeddings.tar.gz"
     embedding_model: str = "text-embedding-ada-002"
     top_k: int = 3
     thresh: float = 0.7
@@ -82,7 +82,7 @@ class Chatbot:
     def _init_documents(self):
         filepath = self.cfg.documents_file
         logger.info(f"loading embeddings from {filepath}...")
-        self.documents = read_documents(filepath)
+        self.documents = get_documents_manager_from_extension(filepath)(filepath)
         logger.info(f"embeddings loaded.")
 
     def _init_unk_embedding(self):
@@ -94,7 +94,6 @@ class Chatbot:
 
     def rank_documents(
         self,
-        documents: pd.DataFrame,
         query: str,
         top_k: float,
         thresh: float,
@@ -108,14 +107,7 @@ class Chatbot:
             query,
             engine=engine,
         )
-        documents["similarity"] = documents.embedding.apply(lambda x: cosine_similarity(x, query_embedding))
-
-        # sort the matched_documents by score
-        matched_documents = documents.sort_values("similarity", ascending=False)
-
-        # limit search to top_k matched_documents.
-        top_k = len(matched_documents) if top_k == -1 else top_k
-        matched_documents = matched_documents.head(top_k)
+        matched_documents = self.documents.retrieve(query_embedding, top_k)
 
         # log matched_documents to the console
         logger.info(f"matched documents before thresh: {matched_documents}")
@@ -236,7 +228,6 @@ class Chatbot:
             question += "\n"
 
         matched_documents = self.rank_documents(
-            documents=self.documents,
             query=question,
             top_k=self.cfg.top_k,
             thresh=self.cfg.thresh,

buster/docparser.py

@@ -8,16 +8,13 @@ import tiktoken
 from bs4 import BeautifulSoup
 from openai.embeddings_utils import get_embedding
 
-from buster.db import DocumentsDB
+from buster.documents import get_documents_manager_from_extension
 from buster.parser import HuggingfaceParser, Parser, SphinxParser
 
 EMBEDDING_MODEL = "text-embedding-ada-002"
 EMBEDDING_ENCODING = "cl100k_base"  # this the encoding for text-embedding-ada-002
 
 
-PICKLE_EXTENSIONS = [".gz", ".bz2", ".zip", ".xz", ".zst", ".tar", ".tar.gz", ".tar.xz", ".tar.bz2"]
-
-
 supported_docs = {
     "mila": {
         "base_url": "https://docs.mila.quebec/",
@@ -77,46 +74,6 @@ def get_all_documents(
     return documents_df
 
 
-def get_file_extension(filepath: str) -> str:
-    return os.path.splitext(filepath)[1]
-
-
-def write_documents(filepath: str, documents_df: pd.DataFrame, source: str = ""):
-    ext = get_file_extension(filepath)
-
-    if ext == ".csv":
-        documents_df.to_csv(filepath, index=False)
-    elif ext in PICKLE_EXTENSIONS:
-        documents_df.to_pickle(filepath)
-    elif ext == ".db":
-        db = DocumentsDB(filepath)
-        db.write_documents(source, documents_df)
-    else:
-        raise ValueError(f"Unsupported format: {ext}.")
-
-
-def read_documents(filepath: str, source: str = "") -> pd.DataFrame:
-    ext = get_file_extension(filepath)
-
-    if ext == ".csv":
-        df = pd.read_csv(filepath)
-
-        if "embedding" in df.columns:
-            df["embedding"] = df.embedding.apply(eval).apply(np.array)
-    elif ext in PICKLE_EXTENSIONS:
-        df = pd.read_pickle(filepath)
-
-        if "embedding" in df.columns:
-            df["embedding"] = df.embedding.apply(np.array)
-    elif ext == ".db":
-        db = DocumentsDB(filepath)
-        df = db.get_documents(source)
-    else:
-        raise ValueError(f"Unsupported format: {ext}.")
-
-    return df
-
-
 def compute_n_tokens(df: pd.DataFrame) -> pd.DataFrame:
     encoding = tiktoken.get_encoding(EMBEDDING_ENCODING)
     # TODO are there unexpected consequences of allowing endoftext?
@@ -129,10 +86,13 @@ def precompute_embeddings(df: pd.DataFrame) -> pd.DataFrame:
     return df
 
 
-def generate_embeddings(filepath: str, output_file: str, source: str) -> pd.DataFrame:
+def generate_embeddings(root_dir: str, output_filepath: str, source: str) -> pd.DataFrame:
     # Get all documents and precompute their embeddings
-    df = read_documents(filepath, source)
-    df = compute_n_tokens(df)
-    df = precompute_embeddings(df)
-    write_documents(filepath=output_file, documents_df=df, source=source)
-    return df
+    documents = get_all_documents(root_dir, supported_docs[source]["base_url"], supported_docs[source]["parser"])
+    documents = compute_n_tokens(documents)
+    documents = precompute_embeddings(documents)
+
+    documents_manager = get_documents_manager_from_extension(output_filepath)(output_filepath)
+    documents_manager.add(source, documents)
+
+    return documents

buster/documents/__init__.py

@@ -0,0 +1,6 @@
+from .base import DocumentsManager
+from .pickle import DocumentsPickle
+from .sqlite import DocumentsDB
+from .utils import get_documents_manager_from_extension
+
+__all__ = [DocumentsManager, DocumentsPickle, DocumentsDB, get_documents_manager_from_extension]

buster/documents/base.py

@@ -0,0 +1,30 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+import pandas as pd
+from openai.embeddings_utils import cosine_similarity
+
+
+@dataclass
+class DocumentsManager(ABC):
+    @abstractmethod
+    def add(self, source: str, df: pd.DataFrame):
+        ...
+
+    @abstractmethod
+    def get_documents(self, source: str) -> pd.DataFrame:
+        ...
+
+    def retrieve(self, query_embedding: list[float], top_k: int, source: str = None) -> pd.DataFrame:
+        documents = self.get_documents(source)
+
+        documents["similarity"] = documents.embedding.apply(lambda x: cosine_similarity(x, query_embedding))
+
+        # sort the matched_documents by score
+        matched_documents = documents.sort_values("similarity", ascending=False)
+
+        # limit search to top_k matched_documents.
+        top_k = len(matched_documents) if top_k == -1 else top_k
+        matched_documents = matched_documents.head(top_k)
+
+        return matched_documents

buster/documents/pickle.py

@@ -0,0 +1,38 @@
+import os
+
+import pandas as pd
+
+from buster.documents.base import DocumentsManager
+
+
+class DocumentsPickle(DocumentsManager):
+    def __init__(self, filepath: str):
+        self.filepath = filepath
+
+        if os.path.exists(filepath):
+            self.documents = pd.read_pickle(filepath)
+        else:
+            self.documents = None
+
+    def add(self, source: str, df: pd.DataFrame):
+        if source is not None:
+            df["source"] = source
+
+        df["current"] = 1
+
+        if self.documents is not None:
+            self.documents.loc[self.documents.source == source, "current"] = 0
+            self.documents = pd.concat([self.documents, df])
+        else:
+            self.documents = df
+
+        self.documents.to_pickle(self.filepath)
+
+    def get_documents(self, source: str) -> pd.DataFrame:
+        documents = self.documents.copy()
+        documents = documents[documents.current == 1]
+
+        if source is not None and "source" in documents.columns:
+            documents = documents[documents.source == source]
+
+        return documents

buster/{db.py → documents/sqlite.py}

@@ -5,6 +5,8 @@ import zlib
 import numpy as np
 import pandas as pd
 
+from buster.documents.base import DocumentsManager
+
 documents_table = """CREATE TABLE IF NOT EXISTS documents (
     id INTEGER PRIMARY KEY AUTOINCREMENT,
     source TEXT NOT NULL,
@@ -33,7 +35,7 @@ qa_table = """CREATE TABLE IF NOT EXISTS qa (
 )"""
 
 
-class DocumentsDB:
+class DocumentsDB(DocumentsManager):
     """Simple SQLite database for storing documents and questions/answers.
 
     The database is just a file on disk. It can store documents from different sources, and it can store multiple versions of the same document (e.g. if the document is updated).
@@ -41,13 +43,13 @@ class DocumentsDB:
 
     Example:
         >>> db = DocumentsDB("/path/to/the/db.db")
-        >>> db.write_documents("source", df)  # df is a DataFrame containing the documents from a given source, obtained e.g. by using buster.docparser.generate_embeddings
+        >>> db.add("source", df)  # df is a DataFrame containing the documents from a given source, obtained e.g. by using buster.docparser.generate_embeddings
         >>> df = db.get_documents("source")
     """
 
-    def __init__(self, db_path):
-        self.db_path = db_path
-        self.conn = sqlite3.connect(db_path)
+    def __init__(self, filepath: str):
+        self.db_path = filepath
+        self.conn = sqlite3.connect(filepath)
         self.cursor = self.conn.cursor()
 
         self.__initialize()
@@ -61,7 +63,7 @@ class DocumentsDB:
         self.cursor.execute(qa_table)
         self.conn.commit()
 
-    def write_documents(self, source: str, df: pd.DataFrame):
+    def add(self, source: str, df: pd.DataFrame):
         """Write all documents from the dataframe into the db. All previous documents from that source will be set to `current = 0`."""
         df = df.copy()
 
@@ -102,7 +104,10 @@ class DocumentsDB:
     def get_documents(self, source: str) -> pd.DataFrame:
         """Get all current documents from a given source."""
         # Execute the SQL statement and fetch the results
-        results = self.cursor.execute("SELECT * FROM documents WHERE source = ? AND current = 1", (source,))
+        if source is not None:
+            results = self.cursor.execute("SELECT * FROM documents WHERE source = ? AND current = 1", (source,))
+        else:
+            results = self.cursor.execute("SELECT * FROM documents WHERE current = 1")
         rows = results.fetchall()
 
         # Convert the results to a pandas DataFrame

buster/documents/utils.py

@@ -0,0 +1,23 @@
+import os
+from typing import Type
+
+from buster.documents.base import DocumentsManager
+from buster.documents.pickle import DocumentsPickle
+from buster.documents.sqlite import DocumentsDB
+
+PICKLE_EXTENSIONS = [".gz", ".bz2", ".zip", ".xz", ".zst", ".tar", ".tar.gz", ".tar.xz", ".tar.bz2"]
+
+
+def get_file_extension(filepath: str) -> str:
+    return os.path.splitext(filepath)[1]
+
+
+def get_documents_manager_from_extension(filepath: str) -> Type[DocumentsManager]:
+    ext = get_file_extension(filepath)
+
+    if ext in PICKLE_EXTENSIONS:
+        return DocumentsPickle
+    elif ext == ".db":
+        return DocumentsDB
+    else:
+        raise ValueError(f"Unsupported format: {ext}.")

tests/test_docparser.py

@@ -1,26 +1,24 @@
 import numpy as np
 import pandas as pd
 
-from buster.docparser import generate_embeddings, read_documents, write_documents
+from buster.docparser import generate_embeddings
+from buster.documents import get_documents_manager_from_extension
 
 
 def test_generate_embeddings(tmp_path, monkeypatch):
-    # Patch the get_embedding function to return a fixed embedding
-    monkeypatch.setattr("buster.docparser.get_embedding", lambda x, engine: [-0.005, 0.0018])
-
     # Create fake data
     data = pd.DataFrame.from_dict({"title": ["test"], "url": ["http://url.com"], "content": ["cool text"]})
 
-    # Write the data to a file
-    filepath = tmp_path / "test_document.csv"
-    write_documents(filepath=filepath, documents_df=data, source="test")
+    # Patch the get_embedding function to return a fixed embedding
+    monkeypatch.setattr("buster.docparser.get_embedding", lambda x, engine: [-0.005, 0.0018])
+    monkeypatch.setattr("buster.docparser.get_all_documents", lambda a, b, c: data)
 
     # Generate embeddings, store in a file
     output_file = tmp_path / "test_document_embeddings.tar.gz"
-    df = generate_embeddings(filepath=filepath, output_file=output_file, source="test")
+    df = generate_embeddings(tmp_path, output_file, source="mila")
 
     # Read the embeddings from the file
-    read_df = read_documents(output_file, "test")
+    read_df = get_documents_manager_from_extension(output_file)(output_file).get_documents("mila")
 
     # Check all the values are correct across the files
     assert df["title"].iloc[0] == data["title"].iloc[0] == read_df["title"].iloc[0]

tests/{test_db.py → test_documents.py}

@@ -1,11 +1,13 @@
 import numpy as np
 import pandas as pd
+import pytest
 
-from buster.db import DocumentsDB
+from buster.documents import DocumentsDB, DocumentsPickle
 
 
-def test_write_read():
-    db = DocumentsDB(":memory:")
+@pytest.mark.parametrize("documents_manager, extension", [(DocumentsDB, "db"), (DocumentsPickle, "tar.gz")])
+def test_write_read(tmp_path, documents_manager, extension):
+    db = documents_manager(tmp_path / f"test.{extension}")
 
     data = pd.DataFrame.from_dict(
         {
@@ -16,7 +18,7 @@ def test_write_read():
             "n_tokens": [10],
         }
     )
-    db.write_documents(source="test", df=data)
+    db.add(source="test", df=data)
 
     db_data = db.get_documents("test")
 
@@ -27,8 +29,9 @@ def test_write_read():
     assert db_data["n_tokens"].iloc[0] == data["n_tokens"].iloc[0]
 
 
-def test_write_write_read():
-    db = DocumentsDB(":memory:")
+@pytest.mark.parametrize("documents_manager, extension", [(DocumentsDB, "db"), (DocumentsPickle, "tar.gz")])
+def test_write_write_read(tmp_path, documents_manager, extension):
+    db = documents_manager(tmp_path / f"test.{extension}")
 
     data_1 = pd.DataFrame.from_dict(
         {
@@ -39,7 +42,7 @@ def test_write_write_read():
             "n_tokens": [10],
         }
     )
-    db.write_documents(source="test", df=data_1)
+    db.add(source="test", df=data_1)
 
     data_2 = pd.DataFrame.from_dict(
         {
@@ -50,7 +53,7 @@ def test_write_write_read():
             "n_tokens": [20],
         }
     )
-    db.write_documents(source="test", df=data_2)
+    db.add(source="test", df=data_2)
 
     db_data = db.get_documents("test")