Repository: D-Star-AI/dsRAG (language: python)

Issue:
Add Pinecone VectorDB class
The VectorDB component is a critical part of the system, and we want to support all of the major players. You can add support for a new one by subclassing the VectorDB base class and making sure you have all of the required methods.


Discussion / hints:

Fix patch:
diff --git a/dsrag/chunk_db.py b/dsrag/chunk_db.py
deleted file mode 100644
index 9f7e4a2..0000000
--- a/dsrag/chunk_db.py
+++ /dev/null
@@ -1,371 +0,0 @@
-from abc import ABC, abstractmethod
-import os
-import time
-import pickle
-import sqlite3
-
-class ChunkDB(ABC):
-    subclasses = {}
-
-    def __init_subclass__(cls, **kwargs):
-        super().__init_subclass__(**kwargs)
-        cls.subclasses[cls.__name__] = cls
-
-    def to_dict(self):
-        return {
-            'subclass_name': self.__class__.__name__,
-        }
-
-    @classmethod
-    def from_dict(cls, config):
-        subclass_name = config.pop('subclass_name', None)  # Remove subclass_name from config
-        subclass = cls.subclasses.get(subclass_name)
-        if subclass:
-            return subclass(**config)  # Pass the modified config without subclass_name
-        else:
-            raise ValueError(f"Unknown subclass: {subclass_name}")
-
-    @abstractmethod
-    def add_document(self, doc_id: str, chunks: dict[dict]):
-        """
-        Store all chunks for a given document.
-        """
-        pass
-
-    @abstractmethod
-    def remove_document(self, doc_id: str):
-        """
-        Remove all chunks and metadata associated with a given document ID.
-        """
-        pass
-
-    @abstractmethod
-    def get_chunk_text(self, doc_id: str, chunk_index: int) -> dict:
-        """
-        Retrieve a specific chunk from a given document ID.
-        """
-        pass
-
-    @abstractmethod
-    def get_document(self, doc_id: str) -> dict:
-        """
-        Retrieve all chunks from a given document ID.
-        """
-        pass
-
-    @abstractmethod
-    def get_document_title(self, doc_id: str, chunk_index: int) -> str:
-        """
-        Retrieve the document title of a specific chunk from a given document ID.
-        """
-        pass
-
-    def get_document_summary(self, doc_id: str, chunk_index: int) -> str:
-        """
-        Retrieve the document summary of a specific chunk from a given document ID.
-        """
-        pass
-
-    def get_section_title(self, doc_id: str, chunk_index: int) -> str:
-        """
-        Retrieve the section title of a specific chunk from a given document ID.
-        """
-        pass
-
-    def get_section_summary(self, doc_id: str, chunk_index: int) -> str:
-        """
-        Retrieve the section summary of a specific chunk from a given document ID.
-        """
-        pass
-
-    @abstractmethod
-    def get_all_doc_ids(self, supp_id: str = None) -> list:
-        """
-        Retrieve all document IDs.
-        """
-        pass
-
-    @abstractmethod
-    def delete(self):
-        """
-        Delete the chunk database.
-        """
-        pass
-
-
-class BasicChunkDB(ChunkDB):
-    """
-    This is a basic implementation of a ChunkDB that stores chunks in a nested dictionary and persists them to disk by pickling the dictionary.
-    """
-    def __init__(self, kb_id: str, storage_directory: str = '~/dsRAG'):
-        self.kb_id = kb_id
-        self.storage_directory = os.path.expanduser(storage_directory)  # Expand the user path
-        # Ensure the base directory and the chunk storage directory exist
-        os.makedirs(os.path.join(self.storage_directory, 'chunk_storage'), exist_ok=True)
-        self.storage_path = os.path.join(self.storage_directory, 'chunk_storage', f'{kb_id}.pkl')
-        self.load()
-
-    def add_document(self, doc_id: str, chunks: dict[dict]):
-        self.data[doc_id] = chunks
-        self.save()
-
-    def remove_document(self, doc_id: str):
-        self.data.pop(doc_id, None)
-        self.save()
-
-    def get_chunk_text(self, doc_id: str, chunk_index: int) -> str:
-        if doc_id in self.data and chunk_index in self.data[doc_id]:
-            return self.data[doc_id][chunk_index]['chunk_text']
-        return None
-    
-    def get_document(self, doc_id: str, include_content: bool = False) -> dict:
-        if doc_id in self.data:
-            document = self.data[doc_id]
-            formatted_document = {
-                'id': doc_id,
-                'title': document[0].get('document_title', "")
-            }
-            
-            if include_content:
-                # Concatenate the chunks into a single string
-                full_document_string = ""
-                for chunk_index, chunk in document.items():
-                    # Join each chunk text with a new line character
-                    full_document_string += chunk['chunk_text'] + "\n"
-                formatted_document["content"] = full_document_string
-            
-            return formatted_document
-        
-        else:
-            return None
-
-
-    def get_document_title(self, doc_id: str, chunk_index: int) -> str:
-        if doc_id in self.data and chunk_index in self.data[doc_id]:
-            if 'document_title' in self.data[doc_id][chunk_index]:
-                return self.data[doc_id][chunk_index]['document_title']
-            else:
-                return ""
-        return None
-    
-    def get_document_summary(self, doc_id: str, chunk_index: int) -> str:
-        if doc_id in self.data and chunk_index in self.data[doc_id]:
-            if 'document_summary' in self.data[doc_id][chunk_index]:
-                return self.data[doc_id][chunk_index]['document_summary']
-            else:
-                return ""
-        return None
-    
-    def get_section_title(self, doc_id: str, chunk_index: int) -> str:
-        if doc_id in self.data and chunk_index in self.data[doc_id]:
-            if 'section_title' in self.data[doc_id][chunk_index]:
-                return self.data[doc_id][chunk_index]['section_title']
-            else:
-                return ""
-        return None
-    
-    def get_section_summary(self, doc_id: str, chunk_index: int) -> str:
-        if doc_id in self.data and chunk_index in self.data[doc_id]:
-            if 'section_summary' in self.data[doc_id][chunk_index]:
-                return self.data[doc_id][chunk_index]['section_summary']
-            else:
-                return ""
-        return None
-    
-    def get_all_doc_ids(self, supp_id: str = None) -> list:
-        doc_ids = list(self.data.keys())
-        if supp_id:
-            doc_ids = [doc_id for doc_id in doc_ids if self.data[doc_id][0].get('supp_id', '') == supp_id]
-        return doc_ids
-
-    def load(self):
-        try:
-            with open(self.storage_path, 'rb') as f:
-                self.data = pickle.load(f)
-        except FileNotFoundError:
-            self.data = {}
-
-    def save(self):
-        with open(self.storage_path, 'wb') as f:
-            pickle.dump(self.data, f)
-    
-    def delete(self):
-        if os.path.exists(self.storage_path):
-            os.remove(self.storage_path)
-
-    def to_dict(self):
-        return {
-            **super().to_dict(),
-            'kb_id': self.kb_id,
-            'storage_directory': self.storage_directory,
-        }
-    
-
-class SQLiteDB(ChunkDB):
-
-    def __init__(self, kb_id: str, storage_directory: str = '~/dsRAG'):
-        self.kb_id = kb_id
-        self.storage_directory = os.path.expanduser(storage_directory)
-        os.makedirs(os.path.join(self.storage_directory, 'chunk_storage'), exist_ok=True)
-        self.db_path = os.path.join(self.storage_directory, 'chunk_storage')
-
-        # Create a table for this kb_id if it doesn't exist
-        conn = sqlite3.connect(os.path.join(self.db_path, f'{kb_id}.db'))
-        c = conn.cursor()
-        result = c.execute(f"SELECT name FROM sqlite_master WHERE type='table' AND name='documents'")
-        if not result.fetchone():
-            # Create a table for this kb_id
-            c.execute(f"CREATE TABLE documents (doc_id TEXT, document_title TEXT, document_summary TEXT, section_title TEXT, section_summary TEXT, chunk_text TEXT, chunk_index INT, created_on TEXT, supp_id TEXT)")
-            conn.commit()
-        else:
-            # Check if we need to add the columns to the table for the supp_id and created_on fields
-            c.execute("PRAGMA table_info(documents)")
-            columns = c.fetchall()
-            column_names = [column[1] for column in columns]
-            if 'supp_id' not in column_names:
-                c.execute("ALTER TABLE documents ADD COLUMN supp_id TEXT")
-            if 'created_on' not in column_names:
-                c.execute("ALTER TABLE documents ADD COLUMN created_on TEXT")
-        conn.close()
-        
-
-    def add_document(self, doc_id: str, chunks: dict[dict]):
-        # Add the docs to the sqlite table
-        conn = sqlite3.connect(os.path.join(self.db_path, f'{self.kb_id}.db'))
-        c = conn.cursor()
-        # Create a created on timestamp
-        created_on = str(int(time.time()))
-
-        # Get the data from the dictionary
-        for chunk_index, chunk in chunks.items():
-            document_title = chunk.get('document_title', "")
-            document_summary = chunk.get('document_summary', "")
-            section_title = chunk.get('section_title', "")
-            section_summary = chunk.get('section_summary', "")
-            chunk_text = chunk.get('chunk_text', "")
-            supp_id = chunk.get('supp_id', "")
-            c.execute(f"INSERT INTO documents (doc_id, document_title, document_summary, section_title, section_summary, chunk_text, chunk_index, created_on, supp_id) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", (doc_id, document_title, document_summary, section_title, section_summary, chunk_text, chunk_index, created_on, supp_id))
-
-        conn.commit()
-        conn.close()
-
-    def remove_document(self, doc_id: str):
-        # Remove the docs from the sqlite table
-        conn = sqlite3.connect(os.path.join(self.db_path, f'{self.kb_id}.db'))
-        c = conn.cursor()
-        c.execute(f"DELETE FROM documents WHERE doc_id='{doc_id}'")
-        conn.commit()
-        conn.close()
-    
-    def get_document(self, doc_id: str, include_content: bool = False) -> dict:
-        # Retrieve the document from the sqlite table
-        conn = sqlite3.connect(os.path.join(self.db_path, f'{self.kb_id}.db'))
-        c = conn.cursor()
-        columns = ["doc_id", "document_title", "document_summary", "created_on"]
-        if include_content:
-            columns += ["chunk_text", "chunk_index"]
-
-        query_statement = f"SELECT {', '.join(columns)} FROM documents WHERE doc_id='{doc_id}'"
-        c.execute(query_statement)
-        results = c.fetchall()
-        conn.close()
-
-        # If there are no results, return an empty dictionary
-        if not results:
-            return None
-
-        formatted_results = {}
-        # Turn the results into an object where the columns are keys
-        if include_content:
-            # Concatenate the chunks into a single string
-            full_document_string = ""
-            for result in results:
-                # Join each chunk text with a new line character
-                full_document_string += result[4] + "\n"
-            formatted_results["content"] = full_document_string
-
-        formatted_results["id"] = doc_id
-        formatted_results["created_on"] = results[0][3]
-        formatted_results["document_title"] = results[0][1]
-        formatted_results["document_summary"] = results[0][2]
-
-        return formatted_results
-
-    def get_chunk_text(self, doc_id: str, chunk_index: int) -> str:
-        # Retrieve the chunk text from the sqlite table
-        conn = sqlite3.connect(os.path.join(self.db_path, f'{self.kb_id}.db'))
-        c = conn.cursor()
-        c.execute(f"SELECT chunk_text FROM documents WHERE doc_id='{doc_id}' AND chunk_index={chunk_index}")
-        result = c.fetchone()
-        conn.close()
-        if result:
-            return result[0]
-        return ""
-
-    def get_document_title(self, doc_id: str, chunk_index: int) -> str:
-        # Retrieve the document title from the sqlite table
-        conn = sqlite3.connect(os.path.join(self.db_path, f'{self.kb_id}.db'))
-        c = conn.cursor()
-        c.execute(f"SELECT document_title FROM documents WHERE doc_id='{doc_id}' AND chunk_index={chunk_index}")
-        result = c.fetchone()
-        conn.close()
-        if result:
-            return result[0]
-        return ""
-
-    def get_document_summary(self, doc_id: str, chunk_index: int) -> str:
-        # Retrieve the document summary from the sqlite table
-        conn = sqlite3.connect(os.path.join(self.db_path, f'{self.kb_id}.db'))
-        c = conn.cursor()
-        c.execute(f"SELECT document_summary FROM documents WHERE doc_id='{doc_id}' AND chunk_index={chunk_index}")
-        result = c.fetchone()
-        conn.close()
-        if result:
-            return result[0]
-        return ""
-
-    def get_section_title(self, doc_id: str, chunk_index: int) -> str:
-        # Retrieve the section title from the sqlite table
-        conn = sqlite3.connect(os.path.join(self.db_path, f'{self.kb_id}.db'))
-        c = conn.cursor()
-        c.execute(f"SELECT section_title FROM documents WHERE doc_id='{doc_id}' AND chunk_index={chunk_index}")
-        result = c.fetchone()
-        conn.close()
-        if result:
-            return result[0]
-        return ""
-
-    def get_section_summary(self, doc_id: str, chunk_index: int) -> str:
-        # Retrieve the section summary from the sqlite table
-        conn = sqlite3.connect(os.path.join(self.db_path, f'{self.kb_id}.db'))
-        c = conn.cursor()
-        c.execute(f"SELECT section_summary FROM documents WHERE doc_id='{doc_id}' AND chunk_index={chunk_index}")
-        result = c.fetchone()
-        conn.close()
-        if result:
-            return result[0]
-        return ""
-
-    def get_all_doc_ids(self, supp_id: str = None) -> list:
-        # Retrieve all document IDs from the sqlite table
-        conn = sqlite3.connect(os.path.join(self.db_path, f'{self.kb_id}.db'))
-        c = conn.cursor()
-        query_statement = f"SELECT DISTINCT doc_id FROM documents"
-        if supp_id:
-            query_statement += f" WHERE supp_id='{supp_id}'"
-        c.execute(query_statement)
-        results = c.fetchall()
-        conn.close()
-        return [result[0] for result in results]
-
-    def delete(self):
-        # Delete the sqlite database
-        if os.path.exists(os.path.join(self.db_path, f'{self.kb_id}.db')):
-            os.remove(os.path.join(self.db_path, f'{self.kb_id}.db'))
-
-    def to_dict(self):
-        return {
-            **super().to_dict(),
-            'kb_id': self.kb_id,
-            'storage_directory': self.storage_directory
-        }
\ No newline at end of file
diff --git a/dsrag/database/__init__.py b/dsrag/database/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/dsrag/database/chunk/__init__.py b/dsrag/database/chunk/__init__.py
new file mode 100644
index 0000000..3533cb2
--- /dev/null
+++ b/dsrag/database/chunk/__init__.py
@@ -0,0 +1,4 @@
+from .basic_db import BasicChunkDB
+from .sqlite_db import SQLiteDB
+from .db import ChunkDB
+from .types import FormattedDocument
diff --git a/dsrag/database/chunk/basic_db.py b/dsrag/database/chunk/basic_db.py
new file mode 100644
index 0000000..13b6672
--- /dev/null
+++ b/dsrag/database/chunk/basic_db.py
@@ -0,0 +1,128 @@
+import os
+import pickle
+from typing import Any, Optional, cast
+
+from dsrag.database.chunk.db import ChunkDB
+from dsrag.database.chunk.types import FormattedDocument
+
+
+class BasicChunkDB(ChunkDB):
+    """
+    This is a basic implementation of a ChunkDB that stores chunks in a nested dictionary and persists them to disk by pickling the dictionary.
+    """
+
+    def __init__(self, kb_id: str, storage_directory: str = "~/dsRAG") -> None:
+        self.kb_id = kb_id
+        self.storage_directory = os.path.expanduser(
+            storage_directory
+        )  # Expand the user path
+        # Ensure the base directory and the chunk storage directory exist
+        os.makedirs(
+            os.path.join(self.storage_directory, "chunk_storage"), exist_ok=True
+        )
+        self.storage_path = os.path.join(
+            self.storage_directory, "chunk_storage", f"{kb_id}.pkl"
+        )
+        self.load()
+
+    def add_document(self, doc_id: str, chunks: dict[int, dict[str, Any]]) -> None:
+        self.data[doc_id] = chunks
+        self.save()
+
+    def remove_document(self, doc_id: str):
+        self.data.pop(doc_id, None)
+        self.save()
+
+    def get_chunk_text(self, doc_id: str, chunk_index: int) -> Optional[str]:
+        if doc_id in self.data and chunk_index in self.data[doc_id]:
+            return self.data[doc_id][chunk_index]["chunk_text"]
+        return None
+
+    def get_document(
+        self, doc_id: str, include_content: bool = False
+    ) -> FormattedDocument | None:
+        if doc_id in self.data:
+            document = self.data[doc_id]
+            title = cast(str, document[0].get("document_title", ""))
+
+            full_document_string = ""
+            if include_content:
+                # Concatenate the chunks into a single string
+                for chunk_index, chunk in document.items():
+                    # Join each chunk text with a new line character
+                    full_document_string += chunk["chunk_text"] + "\n"
+
+            return FormattedDocument(
+                id=doc_id,
+                title=title,
+                content=full_document_string if include_content else None,
+                summary=cast(str, document[0].get("document_summary", "")),
+                created_on=None,
+            )
+
+        else:
+            return None
+
+    def get_document_title(self, doc_id: str, chunk_index: int) -> Optional[str]:
+        if doc_id in self.data and chunk_index in self.data[doc_id]:
+            if "document_title" in self.data[doc_id][chunk_index]:
+                return self.data[doc_id][chunk_index]["document_title"]
+            else:
+                return None
+        return None
+
+    def get_document_summary(self, doc_id: str, chunk_index: int) -> Optional[str]:
+        if doc_id in self.data and chunk_index in self.data[doc_id]:
+            if "document_summary" in self.data[doc_id][chunk_index]:
+                return self.data[doc_id][chunk_index]["document_summary"]
+            else:
+                return None
+        return None
+
+    def get_section_title(self, doc_id: str, chunk_index: int) -> Optional[str]:
+        if doc_id in self.data and chunk_index in self.data[doc_id]:
+            if "section_title" in self.data[doc_id][chunk_index]:
+                return self.data[doc_id][chunk_index]["section_title"]
+            else:
+                return None
+        return None
+
+    def get_section_summary(self, doc_id: str, chunk_index: int) -> Optional[str]:
+        if doc_id in self.data and chunk_index in self.data[doc_id]:
+            if "section_summary" in self.data[doc_id][chunk_index]:
+                return self.data[doc_id][chunk_index]["section_summary"]
+            else:
+                return None
+        return None
+
+    def get_all_doc_ids(self, supp_id: Optional[str] = None) -> list[str]:
+        doc_ids = list(self.data.keys())
+        if supp_id:
+            doc_ids = [
+                doc_id
+                for doc_id in doc_ids
+                if self.data[doc_id][0].get("supp_id", "") == supp_id
+            ]
+        return doc_ids
+
+    def load(self):
+        try:
+            with open(self.storage_path, "rb") as f:
+                self.data = pickle.load(f)
+        except FileNotFoundError:
+            self.data = {}
+
+    def save(self):
+        with open(self.storage_path, "wb") as f:
+            pickle.dump(self.data, f)
+
+    def delete(self):
+        if os.path.exists(self.storage_path):
+            os.remove(self.storage_path)
+
+    def to_dict(self):
+        return {
+            **super().to_dict(),
+            "kb_id": self.kb_id,
+            "storage_directory": self.storage_directory,
+        }
diff --git a/dsrag/database/chunk/db.py b/dsrag/database/chunk/db.py
new file mode 100644
index 0000000..234a749
--- /dev/null
+++ b/dsrag/database/chunk/db.py
@@ -0,0 +1,95 @@
+from abc import ABC, abstractmethod
+from typing import Any, Optional
+
+from dsrag.database.chunk.types import FormattedDocument
+
+
+class ChunkDB(ABC):
+    subclasses = {}
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        cls.subclasses[cls.__name__] = cls
+
+    def to_dict(self):
+        return {
+            "subclass_name": self.__class__.__name__,
+        }
+
+    @classmethod
+    def from_dict(cls, config) -> "ChunkDB":
+        subclass_name = config.pop(
+            "subclass_name", None
+        )  # Remove subclass_name from config
+        subclass = cls.subclasses.get(subclass_name)
+        if subclass:
+            return subclass(**config)  # Pass the modified config without subclass_name
+        else:
+            raise ValueError(f"Unknown subclass: {subclass_name}")
+
+    @abstractmethod
+    def add_document(self, doc_id: str, chunks: dict[int, dict[str, Any]]) -> None:
+        """
+        Store all chunks for a given document.
+        """
+        pass
+
+    @abstractmethod
+    def remove_document(self, doc_id: str) -> None:
+        """
+        Remove all chunks and metadata associated with a given document ID.
+        """
+        pass
+
+    @abstractmethod
+    def get_chunk_text(self, doc_id: str, chunk_index: int) -> Optional[str]:
+        """
+        Retrieve a specific chunk from a given document ID.
+        """
+        pass
+
+    @abstractmethod
+    def get_document(self, doc_id: str) -> FormattedDocument | None:
+        """
+        Retrieve all chunks from a given document ID.
+        """
+        pass
+
+    @abstractmethod
+    def get_document_title(self, doc_id: str, chunk_index: int) -> Optional[str]:
+        """
+        Retrieve the document title of a specific chunk from a given document ID.
+        """
+        pass
+
+    def get_document_summary(self, doc_id: str, chunk_index: int) -> Optional[str]:
+        """
+        Retrieve the document summary of a specific chunk from a given document ID.
+        """
+        pass
+
+    def get_section_title(self, doc_id: str, chunk_index: int) -> Optional[str]:
+        """
+        Retrieve the section title of a specific chunk from a given document ID.
+        """
+        pass
+
+    def get_section_summary(self, doc_id: str, chunk_index: int) -> Optional[str]:
+        """
+        Retrieve the section summary of a specific chunk from a given document ID.
+        """
+        pass
+
+    @abstractmethod
+    def get_all_doc_ids(self, supp_id: Optional[str] = None) -> list[str]:
+        """
+        Retrieve all document IDs.
+        """
+        pass
+
+    @abstractmethod
+    def delete(self) -> None:
+        """
+        Delete the chunk database.
+        """
+        pass
diff --git a/dsrag/database/chunk/sqlite_db.py b/dsrag/database/chunk/sqlite_db.py
new file mode 100644
index 0000000..effba46
--- /dev/null
+++ b/dsrag/database/chunk/sqlite_db.py
@@ -0,0 +1,213 @@
+import os
+import time
+import sqlite3
+from typing import Any, Optional
+
+from dsrag.database.chunk.db import ChunkDB
+from dsrag.database.chunk.types import FormattedDocument
+
+
+class SQLiteDB(ChunkDB):
+
+    def __init__(self, kb_id: str, storage_directory: str = "~/dsRAG") -> None:
+        self.kb_id = kb_id
+        self.storage_directory = os.path.expanduser(storage_directory)
+        os.makedirs(
+            os.path.join(self.storage_directory, "chunk_storage"), exist_ok=True
+        )
+        self.db_path = os.path.join(self.storage_directory, "chunk_storage")
+
+        # Create a table for this kb_id if it doesn't exist
+        conn = sqlite3.connect(os.path.join(self.db_path, f"{kb_id}.db"))
+        c = conn.cursor()
+        result = c.execute(
+            "SELECT name FROM sqlite_master WHERE type='table' AND name='documents'"
+        )
+        if not result.fetchone():
+            # Create a table for this kb_id
+            c.execute(
+                "CREATE TABLE documents (doc_id TEXT, document_title TEXT, document_summary TEXT, section_title TEXT, section_summary TEXT, chunk_text TEXT, chunk_index INT, created_on TEXT, supp_id TEXT)"
+            )
+            conn.commit()
+        else:
+            # Check if we need to add the columns to the table for the supp_id and created_on fields
+            c.execute("PRAGMA table_info(documents)")
+            columns = c.fetchall()
+            column_names = [column[1] for column in columns]
+            if "supp_id" not in column_names:
+                c.execute("ALTER TABLE documents ADD COLUMN supp_id TEXT")
+            if "created_on" not in column_names:
+                c.execute("ALTER TABLE documents ADD COLUMN created_on TEXT")
+        conn.close()
+
+    def add_document(self, doc_id: str, chunks: dict[int, dict[str, Any]]) -> None:
+        # Add the docs to the sqlite table
+        conn = sqlite3.connect(os.path.join(self.db_path, f"{self.kb_id}.db"))
+        c = conn.cursor()
+        # Create a created on timestamp
+        created_on = str(int(time.time()))
+
+        # Get the data from the dictionary
+        for chunk_index, chunk in chunks.items():
+            document_title = chunk.get("document_title", "")
+            document_summary = chunk.get("document_summary", "")
+            section_title = chunk.get("section_title", "")
+            section_summary = chunk.get("section_summary", "")
+            chunk_text = chunk.get("chunk_text", "")
+            supp_id = chunk.get("supp_id", "")
+            c.execute(
+                "INSERT INTO documents (doc_id, document_title, document_summary, section_title, section_summary, chunk_text, chunk_index, created_on, supp_id) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
+                (
+                    doc_id,
+                    document_title,
+                    document_summary,
+                    section_title,
+                    section_summary,
+                    chunk_text,
+                    chunk_index,
+                    created_on,
+                    supp_id,
+                ),
+            )
+
+        conn.commit()
+        conn.close()
+
+    def remove_document(self, doc_id: str) -> None:
+        # Remove the docs from the sqlite table
+        conn = sqlite3.connect(os.path.join(self.db_path, f"{self.kb_id}.db"))
+        c = conn.cursor()
+        c.execute(f"DELETE FROM documents WHERE doc_id='{doc_id}'")
+        conn.commit()
+        conn.close()
+
+    def get_document(
+        self, doc_id: str, include_content: bool = False
+    ) -> FormattedDocument | None:
+        # Retrieve the document from the sqlite table
+        conn = sqlite3.connect(os.path.join(self.db_path, f"{self.kb_id}.db"))
+        c = conn.cursor()
+        columns = ["doc_id", "document_title", "document_summary", "created_on"]
+        if include_content:
+            columns += ["chunk_text", "chunk_index"]
+
+        query_statement = (
+            f"SELECT {', '.join(columns)} FROM documents WHERE doc_id='{doc_id}'"
+        )
+        c.execute(query_statement)
+        results = c.fetchall()
+        conn.close()
+
+        # If there are no results, return an empty dictionary
+        if not results:
+            return None
+
+        # Turn the results into an object where the columns are keys
+        full_document_string = ""
+        if include_content:
+            # Concatenate the chunks into a single string
+            for result in results:
+                # Join each chunk text with a new line character
+                full_document_string += result[4] + "\n"
+
+        title = results[0][1]
+        created_on = results[0][3]
+        title = results[0][1]
+        summary = results[0][2]
+
+        return FormattedDocument(
+            id=doc_id,
+            title=title,
+            content=full_document_string if include_content else None,
+            summary=summary,
+            created_on=created_on,
+        )
+
+    def get_chunk_text(self, doc_id: str, chunk_index: int) -> Optional[str]:
+        # Retrieve the chunk text from the sqlite table
+        conn = sqlite3.connect(os.path.join(self.db_path, f"{self.kb_id}.db"))
+        c = conn.cursor()
+        c.execute(
+            f"SELECT chunk_text FROM documents WHERE doc_id='{doc_id}' AND chunk_index={chunk_index}"
+        )
+        result = c.fetchone()
+        conn.close()
+        if result:
+            return result[0]
+        return None
+
+    def get_document_title(self, doc_id: str, chunk_index: int) -> Optional[str]:
+        # Retrieve the document title from the sqlite table
+        conn = sqlite3.connect(os.path.join(self.db_path, f"{self.kb_id}.db"))
+        c = conn.cursor()
+        c.execute(
+            f"SELECT document_title FROM documents WHERE doc_id='{doc_id}' AND chunk_index={chunk_index}"
+        )
+        result = c.fetchone()
+        conn.close()
+        if result:
+            return result[0]
+        return None
+
+    def get_document_summary(self, doc_id: str, chunk_index: int) -> Optional[str]:
+        # Retrieve the document summary from the sqlite table
+        conn = sqlite3.connect(os.path.join(self.db_path, f"{self.kb_id}.db"))
+        c = conn.cursor()
+        c.execute(
+            f"SELECT document_summary FROM documents WHERE doc_id='{doc_id}' AND chunk_index={chunk_index}"
+        )
+        result = c.fetchone()
+        conn.close()
+        if result:
+            return result[0]
+        return None
+
+    def get_section_title(self, doc_id: str, chunk_index: int) -> Optional[str]:
+        # Retrieve the section title from the sqlite table
+        conn = sqlite3.connect(os.path.join(self.db_path, f"{self.kb_id}.db"))
+        c = conn.cursor()
+        c.execute(
+            f"SELECT section_title FROM documents WHERE doc_id='{doc_id}' AND chunk_index={chunk_index}"
+        )
+        result = c.fetchone()
+        conn.close()
+        if result:
+            return result[0]
+        return None
+
+    def get_section_summary(self, doc_id: str, chunk_index: int) -> Optional[str]:
+        # Retrieve the section summary from the sqlite table
+        conn = sqlite3.connect(os.path.join(self.db_path, f"{self.kb_id}.db"))
+        c = conn.cursor()
+        c.execute(
+            f"SELECT section_summary FROM documents WHERE doc_id='{doc_id}' AND chunk_index={chunk_index}"
+        )
+        result = c.fetchone()
+        conn.close()
+        if result:
+            return result[0]
+        return None
+
+    def get_all_doc_ids(self, supp_id: Optional[str] = None) -> list[str]:
+        # Retrieve all document IDs from the sqlite table
+        conn = sqlite3.connect(os.path.join(self.db_path, f"{self.kb_id}.db"))
+        c = conn.cursor()
+        query_statement = "SELECT DISTINCT doc_id FROM documents"
+        if supp_id:
+            query_statement += f" WHERE supp_id='{supp_id}'"
+        c.execute(query_statement)
+        results = c.fetchall()
+        conn.close()
+        return [result[0] for result in results]
+
+    def delete(self) -> None:
+        # Delete the sqlite database
+        if os.path.exists(os.path.join(self.db_path, f"{self.kb_id}.db")):
+            os.remove(os.path.join(self.db_path, f"{self.kb_id}.db"))
+
+    def to_dict(self) -> dict[str, str]:
+        return {
+            **super().to_dict(),
+            "kb_id": self.kb_id,
+            "storage_directory": self.storage_directory,
+        }
diff --git a/dsrag/database/chunk/types.py b/dsrag/database/chunk/types.py
new file mode 100644
index 0000000..0ab8886
--- /dev/null
+++ b/dsrag/database/chunk/types.py
@@ -0,0 +1,11 @@
+from datetime import datetime
+from typing import Optional
+from typing_extensions import TypedDict
+
+
+class FormattedDocument(TypedDict):
+    id: str
+    title: str
+    content: Optional[str]
+    summary: Optional[str]
+    created_on: datetime | None
diff --git a/dsrag/database/vector/__init__.py b/dsrag/database/vector/__init__.py
new file mode 100644
index 0000000..fc3e091
--- /dev/null
+++ b/dsrag/database/vector/__init__.py
@@ -0,0 +1,5 @@
+from .db import VectorDB
+from .basic_db import BasicVectorDB
+from .chroma_db import ChromaDB
+from .types import VectorSearchResult, Vector, ChunkMetadata
+from .weaviate_db import WeaviateVectorDB
diff --git a/dsrag/database/vector/basic_db.py b/dsrag/database/vector/basic_db.py
new file mode 100644
index 0000000..d969ec1
--- /dev/null
+++ b/dsrag/database/vector/basic_db.py
@@ -0,0 +1,118 @@
+import pickle
+from dsrag.database.vector.db import VectorDB
+from typing import Sequence
+from dsrag.database.vector.types import ChunkMetadata, Vector, VectorSearchResult
+from sklearn.metrics.pairwise import cosine_similarity
+import os
+import numpy as np
+
+
+class BasicVectorDB(VectorDB):
+    def __init__(
+        self, kb_id: str, storage_directory: str = "~/dsRAG", use_faiss: bool = True
+    ) -> None:
+        self.kb_id = kb_id
+        self.storage_directory = storage_directory
+        self.use_faiss = use_faiss
+        self.vector_storage_path = os.path.join(
+            self.storage_directory, "vector_storage", f"{kb_id}.pkl"
+        )
+        self.load()
+
+    def add_vectors(
+        self, vectors: Sequence[Vector], metadata: Sequence[ChunkMetadata]
+    ) -> None:
+        try:
+            assert len(vectors) == len(metadata)
+        except AssertionError:
+            raise ValueError(
+                "Error in add_vectors: the number of vectors and metadata items must be the same."
+            )
+        self.vectors.extend(vectors)
+        self.metadata.extend(metadata)
+        self.save()
+
+    def search(self, query_vector, top_k=10) -> list[VectorSearchResult]:
+        if not self.vectors:
+            return []
+
+        if self.use_faiss:
+            return self.search_faiss(query_vector, top_k)
+
+        similarities = cosine_similarity([query_vector], self.vectors)[0]
+        indexed_similarities = sorted(
+            enumerate(similarities), key=lambda x: x[1], reverse=True
+        )
+        results: list[VectorSearchResult] = []
+        for i, similarity in indexed_similarities[:top_k]:
+            result = VectorSearchResult(
+                doc_id=None,
+                vector=None,
+                metadata=self.metadata[i],
+                similarity=similarity,
+            )
+            results.append(result)
+        return results
+
+    def search_faiss(self, query_vector, top_k=10) -> list[VectorSearchResult]:
+        from faiss.contrib.exhaustive_search import knn
+
+        # Limit top_k to the number of vectors we have - Faiss doesn't automatically handle this
+        top_k = min(top_k, len(self.vectors))
+
+        # faiss expects 2D arrays of vectors
+        vectors_array = (
+            np.array(self.vectors).astype("float32").reshape(len(self.vectors), -1)
+        )
+        query_vector_array = np.array(query_vector).astype("float32").reshape(1, -1)
+
+        _, I = knn(
+            query_vector_array, vectors_array, top_k
+        )  # I is a list of indices in the corpus_vectors array
+        results: list[VectorSearchResult] = []
+        for i in I[0]:
+            result = VectorSearchResult(
+                doc_id=None,
+                vector=None,
+                metadata=self.metadata[i],
+                similarity=cosine_similarity([query_vector], [self.vectors[i]])[0][0],
+            )
+            results.append(result)
+        return results
+
+    def remove_document(self, doc_id):
+        i = 0
+        while i < len(self.metadata):
+            if self.metadata[i]["doc_id"] == doc_id:
+                del self.vectors[i]
+                del self.metadata[i]
+            else:
+                i += 1
+        self.save()
+
+    def save(self):
+        os.makedirs(
+            os.path.dirname(self.vector_storage_path), exist_ok=True
+        )  # Ensure the directory exists
+        with open(self.vector_storage_path, "wb") as f:
+            pickle.dump((self.vectors, self.metadata), f)
+
+    def load(self):
+        if os.path.exists(self.vector_storage_path):
+            with open(self.vector_storage_path, "rb") as f:
+                self.vectors, self.metadata = pickle.load(f)
+        else:
+            self.vectors = []
+            self.metadata = []
+
+    def delete(self):
+        if os.path.exists(self.vector_storage_path):
+            os.remove(self.vector_storage_path)
+
+    def to_dict(self):
+        return {
+            **super().to_dict(),
+            "kb_id": self.kb_id,
+            "storage_directory": self.storage_directory,
+            "use_faiss": self.use_faiss,
+        }
diff --git a/dsrag/database/vector/chroma_db.py b/dsrag/database/vector/chroma_db.py
new file mode 100644
index 0000000..ce78311
--- /dev/null
+++ b/dsrag/database/vector/chroma_db.py
@@ -0,0 +1,81 @@
+from dsrag.database.vector.db import VectorDB
+from dsrag.database.vector.types import VectorSearchResult
+
+import os
+import chromadb
+
+
+class ChromaDB(VectorDB):
+
+    def __init__(self, kb_id: str, storage_directory: str = "~/dsRAG"):
+        self.kb_id = kb_id
+        self.storage_directory = os.path.expanduser(storage_directory)
+        self.vector_storage_path = os.path.join(
+            self.storage_directory, "vector_storage"
+        )
+        self.client = chromadb.PersistentClient(path=self.vector_storage_path)
+        self.collection = self.client.get_or_create_collection(
+            kb_id, metadata={"hnsw:space": "cosine"}
+        )
+
+    def get_num_vectors(self):
+        return self.collection.count()
+
+    def add_vectors(self, vectors, metadata):
+        try:
+            assert len(vectors) == len(metadata)
+        except AssertionError:
+            raise ValueError(
+                "Error in add_vectors: the number of vectors and metadata items must be the same."
+            )
+
+        # Create the ids from the metadata, defined as {metadata["doc_id"]}_{metadata["chunk_index"]}
+        ids = [f"{meta['doc_id']}_{meta['chunk_index']}" for meta in metadata]
+
+        self.collection.add(embeddings=vectors, metadatas=metadata, ids=ids)
+
+    def search(self, query_vector, top_k=10) -> list[VectorSearchResult]:
+
+        num_vectors = self.get_num_vectors()
+        if num_vectors == 0:
+            return []
+            # raise ValueError('No vectors stored in the database.')
+
+        query_results = self.collection.query(
+            query_embeddings=query_vector,
+            n_results=top_k,
+            include=["distances", "metadatas"],
+        )
+
+        metadata = query_results["metadatas"][0]
+        metadata = query_results["metadatas"][0]
+        distances = query_results["distances"][0]
+
+        results: list[VectorSearchResult] = []
+        for _, (distance, metadata) in enumerate(zip(distances, metadata)):
+            results.append(
+                VectorSearchResult(
+                    doc_id=metadata["doc_id"],
+                    vector=None,
+                    metadata=metadata,
+                    similarity=1 - distance,
+                )
+            )
+
+        results = sorted(results, key=lambda x: x["similarity"], reverse=True)
+
+        return results
+
+    def remove_document(self, doc_id):
+        data = self.collection.get(where={"doc_id": doc_id})
+        doc_ids = data["ids"]
+        self.collection.delete(ids=doc_ids)
+
+    def delete(self):
+        self.client.delete_collection(name=self.kb_id)
+
+    def to_dict(self):
+        return {
+            **super().to_dict(),
+            "kb_id": self.kb_id,
+        }
diff --git a/dsrag/database/vector/db.py b/dsrag/database/vector/db.py
new file mode 100644
index 0000000..30850aa
--- /dev/null
+++ b/dsrag/database/vector/db.py
@@ -0,0 +1,68 @@
+from abc import ABC, abstractmethod
+from typing import Sequence
+from dsrag.database.vector.types import ChunkMetadata, Vector, VectorSearchResult
+
+
+class VectorDB(ABC):
+    subclasses = {}
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        cls.subclasses[cls.__name__] = cls
+
+    def to_dict(self):
+        return {
+            "subclass_name": self.__class__.__name__,
+        }
+
+    @classmethod
+    def from_dict(cls, config):
+        subclass_name = config.pop(
+            "subclass_name", None
+        )  # Remove subclass_name from config
+        subclass = cls.subclasses.get(subclass_name)
+        if subclass:
+            return subclass(**config)  # Pass the modified config without subclass_name
+        else:
+            raise ValueError(f"Unknown subclass: {subclass_name}")
+
+    @abstractmethod
+    def add_vectors(
+        self, vectors: Sequence[Vector], metadata: Sequence[ChunkMetadata]
+    ) -> None:
+        """
+        Store a list of vectors with associated metadata.
+        """
+        pass
+
+    @abstractmethod
+    def remove_document(self, doc_id) -> None:
+        """
+        Remove all vectors and metadata associated with a given document ID.
+        """
+        pass
+
+    @abstractmethod
+    def search(self, query_vector, top_k=10) -> list[VectorSearchResult]:
+        """
+        Retrieve the top-k closest vectors to a given query vector.
+        - needs to return results as list of dictionaries in this format:
+        {
+            'metadata':
+                {
+                    'doc_id': doc_id,
+                    'chunk_index': chunk_index,
+                    'chunk_header': chunk_header,
+                    'chunk_text': chunk_text
+                },
+            'similarity': similarity,
+        }
+        """
+        pass
+
+    @abstractmethod
+    def delete(self) -> None:
+        """
+        Delete the vector database.
+        """
+        pass
diff --git a/dsrag/database/vector/types.py b/dsrag/database/vector/types.py
new file mode 100644
index 0000000..99b9857
--- /dev/null
+++ b/dsrag/database/vector/types.py
@@ -0,0 +1,23 @@
+from typing import Optional, Sequence
+from typing_extensions import TypedDict
+
+
+class ChunkMetadata(TypedDict):
+    doc_id: str
+    chunk_text: str
+    chunk_index: int
+    chunk_header: str
+    document_title: Optional[str]
+    document_summary: Optional[str]
+    section_title: Optional[str]
+    section_summary: Optional[str]
+
+
+Vector = Sequence[float] | Sequence[int]
+
+
+class VectorSearchResult(TypedDict):
+    doc_id: Optional[str]
+    vector: Vector | None
+    metadata: ChunkMetadata
+    similarity: float
diff --git a/dsrag/database/vector/weaviate_db.py b/dsrag/database/vector/weaviate_db.py
new file mode 100644
index 0000000..7fab9e9
--- /dev/null
+++ b/dsrag/database/vector/weaviate_db.py
@@ -0,0 +1,203 @@
+from typing import Sequence, cast
+from dsrag.database.vector.types import ChunkMetadata, Vector, VectorSearchResult
+from dsrag.database.vector.db import VectorDB
+import numpy as np
+import weaviate
+import weaviate.classes as wvc
+from weaviate.util import generate_uuid5
+
+
+class WeaviateVectorDB(VectorDB):
+    """
+    An implementation of the VectorDB interface for Weaviate using the Python v4 client.
+
+    This class provides methods for adding, removing, and searching for vectorized data
+    within a Weaviate instance.
+    """
+
+    def __init__(
+        self,
+        kb_id: str,
+        http_host="localhost",
+        http_port="8099",
+        http_secure=False,
+        grpc_host="localhost",
+        grpc_port="50052",
+        grpc_secure=False,
+        weaviate_secret="secr3tk3y",
+        init_timeout: int = 2,
+        query_timeout: int = 45,
+        insert_timeout: int = 120,
+        use_embedded_weaviate: bool = False,
+    ):
+        """
+        Initializes a WeaviateVectorDB instance.
+
+        Args:
+            http_host: The hostname of the Weaviate server.
+            http_port: The HTTP port of the Weaviate server.
+            http_secure: Whether to use HTTPS for the connection.
+            grpc_host: The hostname of the Weaviate server for gRPC connections.
+            grpc_port: The gRPC port of the Weaviate server.
+            grpc_secure: Whether to use gRPCs for the connection.
+            class_name: The name of the Weaviate class to use for storing data.
+            kb_id: An optional identifier for the knowledge base.
+        """
+
+        # save all of these parameters as attributes so they're easily accessible for the to_dict method
+        self.kb_id = kb_id
+        self.http_host = http_host
+        self.http_port = http_port
+        self.http_secure = http_secure
+        self.grpc_host = grpc_host
+        self.grpc_port = grpc_port
+        self.grpc_secure = grpc_secure
+        self.weaviate_secret = weaviate_secret
+        self.init_timeout = init_timeout
+        self.query_timeout = query_timeout
+        self.insert_timeout = insert_timeout
+        self.use_embedded_weaviate = use_embedded_weaviate
+
+        additional_headers = {}
+        if use_embedded_weaviate:
+            additional_headers["ENABLE_MODULES"] = (
+                "backup-filesystem,text2vec-openai,text2vec-cohere,text2vec-huggingface,ref2vec-centroid,generative-openai,qna-openai"
+            )
+            additional_headers["BACKUP_FILESYSTEM_PATH"] = "./weaviate/backups"
+            self.client = weaviate.WeaviateClient(
+                embedded_options=weaviate.embedded.EmbeddedOptions(
+                    persistence_data_path="./weaviate/data",
+                ),
+                additional_headers=additional_headers,
+            )
+        else:
+            connection_params = weaviate.connect.ConnectionParams.from_params(
+                http_host=http_host,
+                http_port=http_port,
+                http_secure=http_secure,
+                grpc_host=grpc_host,
+                grpc_port=grpc_port,
+                grpc_secure=grpc_secure,
+            )
+            self.client = weaviate.WeaviateClient(
+                connection_params=connection_params,
+                auth_client_secret=weaviate.auth.AuthApiKey(weaviate_secret),
+                additional_headers=additional_headers,
+                additional_config=weaviate.classes.init.AdditionalConfig(
+                    timeout=weaviate.classes.init.Timeout(
+                        init=init_timeout, query=query_timeout, insert=insert_timeout
+                    )
+                ),
+            )
+
+        self.client.connect()
+        self.collection = self.client.collections.get(
+            kb_id
+        )  # assume this creates a new collection if it doesn't exist
+
+    def close(self):
+        """
+        Closes the connection to Weaviate.
+        """
+        self.client.close()
+
+    def add_vectors(
+        self, vectors: Sequence[Vector], metadata: Sequence[ChunkMetadata]
+    ) -> None:
+        """
+        Adds a list of vectors with associated metadata to Weaviate.
+
+        Args:
+            vectors: A list of vector embeddings.
+            metadata: A list of dictionaries containing metadata for each vector.
+
+        Raises:
+            ValueError: If the number of vectors and metadata items do not match.
+        """
+        try:
+            assert len(vectors) == len(metadata)
+        except AssertionError as exc:
+            raise ValueError(
+                "Error in add_vectors: the number of vectors and metadata items must be the same."
+            ) from exc
+
+        with self.collection.batch.dynamic() as batch:
+            for vector, meta in zip(vectors, metadata):
+                doc_id = meta.get("doc_id", "")
+                chunk_text = meta.get("chunk_text", "")
+                chunk_index = meta.get("chunk_index", 0)
+                uuid = generate_uuid5(f"{doc_id}_{chunk_index}")
+                batch.add_object(
+                    properties={
+                        "content": chunk_text,
+                        "doc_id": doc_id,
+                        "chunk_index": chunk_index,
+                        "metadata": meta,
+                    },
+                    vector=vector,
+                    uuid=uuid,
+                )
+
+    def remove_document(self, doc_id) -> None:
+        """
+        Removes a document (data object) from Weaviate.
+
+        Args:
+            doc_id: The UUID of the document to remove.
+        """
+        self.collection.data.delete_many(
+            where=wvc.query.Filter.by_property("doc_id").contains_any([doc_id])
+        )
+
+    def search(self, query_vector, top_k=10) -> list[VectorSearchResult]:
+        """
+        Searches for the top-k closest vectors to the given query vector.
+
+        Args:
+            query_vector: The query vector embedding.
+            top_k: The number of results to return.
+
+        Returns:
+            A list of dictionaries containing the metadata and similarity scores of
+            the top-k results.
+        """
+        # convert the query vector to a list if it's not already
+        if isinstance(query_vector, np.ndarray):
+            query_vector = query_vector.tolist()
+
+        results: list[VectorSearchResult] = []
+        response = self.collection.query.near_vector(
+            near_vector=query_vector,
+            limit=top_k,
+            return_metadata=wvc.query.MetadataQuery(distance=True),
+        )
+        for obj in response.objects:
+            results.append(
+                VectorSearchResult(
+                    doc_id=cast(str, obj.properties["doc_id"]),
+                    metadata=cast(ChunkMetadata, obj.properties["metadata"]),
+                    similarity=cast(float, 1.0 - obj.metadata.distance),
+                    vector=cast(Vector, obj.vector),
+                )
+            )
+        return results
+
+    def to_dict(self):
+        return {
+            **super().to_dict(),
+            "kb_id": self.kb_id,
+            "http_host": self.http_host,
+            "http_port": self.http_port,
+            "http_secure": self.http_secure,
+            "grpc_host": self.grpc_host,
+            "grpc_port": self.grpc_port,
+            "grpc_secure": self.grpc_secure,
+            "weaviate_secret": self.weaviate_secret,
+            "init_timeout": self.init_timeout,
+            "query_timeout": self.query_timeout,
+            "insert_timeout": self.insert_timeout,
+            "use_embedded_weaviate": self.use_embedded_weaviate,
+        }
+
+    def delete(self):
+        pass
diff --git a/dsrag/embedding.py b/dsrag/embedding.py
index 51a125e..9782c61 100644
--- a/dsrag/embedding.py
+++ b/dsrag/embedding.py
@@ -1,5 +1,7 @@
 import os
 from abc import ABC, abstractmethod
+from typing import Optional
+from dsrag.database.vector.types import Vector
 from openai import OpenAI
 import cohere
 import voyageai
@@ -20,10 +22,11 @@
     "nomic-embed-text": 768,
 }
 
+
 class Embedding(ABC):
     subclasses = {}
 
-    def __init__(self, dimension=None):
+    def __init__(self, dimension: int | None = None):
         self.dimension = dimension
 
     def __init_subclass__(cls, **kwargs):
@@ -31,24 +34,24 @@ def __init_subclass__(cls, **kwargs):
         cls.subclasses[cls.__name__] = cls
 
     def to_dict(self):
-        return {
-            'subclass_name': self.__class__.__name__,
-            'dimension': self.dimension
-        }
+        return {"subclass_name": self.__class__.__name__, "dimension": self.dimension}
 
     @classmethod
-    def from_dict(cls, config):
-        subclass_name = config.pop('subclass_name', None)  # Remove subclass_name from config
+    def from_dict(cls, config) -> "Embedding":
+        subclass_name = config.pop(
+            "subclass_name", None
+        )  # Remove subclass_name from config
         subclass = cls.subclasses.get(subclass_name)
         if subclass:
             return subclass(**config)  # Pass the modified config without subclass_name
         else:
             raise ValueError(f"Unknown subclass: {subclass_name}")
-        
+
     @abstractmethod
-    def get_embeddings(self, text, input_type=None):
+    def get_embeddings(self, text, input_type=None) -> list[Vector]:
         pass
 
+
 class OpenAIEmbedding(Embedding):
     def __init__(self, model: str = "text-embedding-3-small", dimension: int = 768):
         """
@@ -58,50 +61,56 @@ def __init__(self, model: str = "text-embedding-3-small", dimension: int = 768):
         self.model = model
         self.client = OpenAI()
 
-    def get_embeddings(self, text, input_type=None):
-        response = self.client.embeddings.create(input=text, model=self.model, dimensions=int(self.dimension))
+    def get_embeddings(self, text, input_type: Optional[str]) -> list[Vector]:
+        response = self.client.embeddings.create(
+            input=text, model=self.model, dimensions=self.dimension
+        )
         embeddings = [embedding_item.embedding for embedding_item in response.data]
         return embeddings[0] if isinstance(text, str) else embeddings
-    
+
     def to_dict(self):
         base_dict = super().to_dict()
-        base_dict.update({
-            'model': self.model
-        })
+        base_dict.update({"model": self.model})
         return base_dict
 
+
 class CohereEmbedding(Embedding):
-    def __init__(self, model: str = "embed-english-v3.0", dimension: int = None):
+    def __init__(self, model: str = "embed-english-v3.0", dimension: int | None = None):
         super().__init__()
         self.model = model
-        self.client = cohere.Client(os.environ['CO_API_KEY'])
-        
+        self.client = cohere.Client(os.environ["CO_API_KEY"])
+
         # Set dimension if not provided
         if dimension is None:
             try:
                 self.dimension = dimensionality[model]
             except KeyError:
-                raise ValueError(f"Dimension for model {model} is unknown. Please provide the dimension manually.")
+                raise ValueError(
+                    f"Dimension for model {model} is unknown. Please provide the dimension manually."
+                )
         else:
             self.dimension = dimension
 
-    def get_embeddings(self, text, input_type=None):
+    def get_embeddings(self, text, input_type: Optional[str] = None):
         if input_type == "query":
             input_type = "search_query"
         elif input_type == "document":
             input_type = "search_document"
-        response = self.client.embed(texts=[text] if isinstance(text, str) else text, input_type=input_type, model=self.model)
+        response = self.client.embed(
+            texts=[text] if isinstance(text, str) else text,
+            input_type=input_type,
+            model=self.model,
+        )
         return response.embeddings[0] if isinstance(text, str) else response.embeddings
-    
+
     def to_dict(self):
         base_dict = super().to_dict()
-        base_dict.update({
-            'model': self.model
-        })
+        base_dict.update({"model": self.model})
         return base_dict
 
+
 class VoyageAIEmbedding(Embedding):
-    def __init__(self, model: str = "voyage-large-2", dimension: int = None):
+    def __init__(self, model: str = "voyage-large-2", dimension: int | None = None):
         super().__init__()
         self.model = model
         self.client = voyageai.Client()
@@ -111,19 +120,23 @@ def __init__(self, model: str = "voyage-large-2", dimension: int = None):
             try:
                 self.dimension = dimensionality[model]
             except KeyError:
-                raise ValueError(f"Dimension for model {model} is unknown. Please provide the dimension manually.")
+                raise ValueError(
+                    f"Dimension for model {model} is unknown. Please provide the dimension manually."
+                )
         else:
             self.dimension = dimension
 
     def get_embeddings(self, text, input_type=None):
-        response = self.client.embed(texts=[text] if isinstance(text, str) else text, model=self.model, input_type=input_type)
+        response = self.client.embed(
+            texts=[text] if isinstance(text, str) else text,
+            model=self.model,
+            input_type=input_type,
+        )
         return response.embeddings[0] if isinstance(text, str) else response.embeddings
-        
+
     def to_dict(self):
         base_dict = super().to_dict()
-        base_dict.update({
-            'model': self.model
-        })
+        base_dict.update({"model": self.model})
         base_dict.update({"model": self.model})
         return base_dict
 
@@ -131,7 +144,10 @@ def to_dict(self):
 class OllamaEmbedding(Embedding):
 
     def __init__(
-        self, model: str = "llama3", dimension: int = None, client: ollama.Client = None
+        self,
+        model: str = "llama3",
+        dimension: int | None = None,
+        client: ollama.Client = None,
     ):
         super().__init__(dimension)
         self.model = model
diff --git a/dsrag/knowledge_base.py b/dsrag/knowledge_base.py
index 636cefa..b823d72 100644
--- a/dsrag/knowledge_base.py
+++ b/dsrag/knowledge_base.py
@@ -3,11 +3,22 @@
 import os
 import time
 import json
-from typing import Union, Dict
-from dsrag.auto_context import get_document_title, get_document_summary, get_section_summary, get_chunk_header, get_segment_header
-from dsrag.rse import get_relevance_values, get_best_segments, get_meta_document, RSE_PARAMS_PRESETS
-from dsrag.vector_db import VectorDB, BasicVectorDB
-from dsrag.chunk_db import ChunkDB, BasicChunkDB
+from typing import Optional, Union, Dict
+from dsrag.auto_context import (
+    get_document_title,
+    get_document_summary,
+    get_section_summary,
+    get_chunk_header,
+    get_segment_header,
+)
+from dsrag.rse import (
+    get_relevance_values,
+    get_best_segments,
+    get_meta_document,
+    RSE_PARAMS_PRESETS,
+)
+from dsrag.database.vector import Vector, VectorDB, BasicVectorDB
+from dsrag.database.chunk import ChunkDB, BasicChunkDB
 from dsrag.embedding import Embedding, OpenAIEmbedding
 from dsrag.reranker import Reranker, CohereReranker
 from dsrag.llm import LLM, OpenAIChatAPI
@@ -15,74 +26,133 @@
 
 
 class KnowledgeBase:
-    def __init__(self, kb_id: str, title: str = "", supp_id: str = "", description: str = "", language: str = "en", storage_directory: str = '~/dsRAG', embedding_model: Embedding = None, reranker: Reranker = None, auto_context_model: LLM = None, vector_db: VectorDB = None, chunk_db: ChunkDB = None, exists_ok: bool = True):
+    def __init__(
+        self,
+        kb_id: str,
+        title: str = "",
+        supp_id: str = "",
+        description: str = "",
+        language: str = "en",
+        storage_directory: str = "~/dsRAG",
+        embedding_model: Embedding | None = None,
+        reranker: Reranker | None = None,
+        auto_context_model: LLM | None = None,
+        vector_db: VectorDB | None = None,
+        chunk_db: ChunkDB | None = None,
+        exists_ok: bool = True,
+        save_metadata_to_disk: bool = True,
+    ):
         self.kb_id = kb_id
         self.storage_directory = os.path.expanduser(storage_directory)
 
-        # load the KB if it exists; otherwise, initialize it and save it to disk
-        metadata_path = self.get_metadata_path()
-        if os.path.exists(metadata_path) and exists_ok:
-            self.load(auto_context_model, reranker) # allow the user to override the auto_context_model and reranker
-        elif os.path.exists(metadata_path) and not exists_ok:
-            raise ValueError(f"Knowledge Base with ID {kb_id} already exists. Use exists_ok=True to load it.")
+        if save_metadata_to_disk:
+            # load the KB if it exists; otherwise, initialize it and save it to disk
+            metadata_path = self.get_metadata_path()
+            if os.path.exists(metadata_path) and exists_ok:
+                self.load(
+                    auto_context_model, reranker
+                )  # allow the user to override the auto_context_model and reranker
+            elif os.path.exists(metadata_path) and not exists_ok:
+                raise ValueError(
+                    f"Knowledge Base with ID {kb_id} already exists. Use exists_ok=True to load it."
+                )
+            else:
+                created_time = time.time()
+                # We don't care about the milliseconds
+                created_time = int(created_time)
+                self.kb_metadata = {
+                    "title": title,
+                    "description": description,
+                    "language": language,
+                    "supp_id": supp_id,
+                    "created_on": created_time,
+                }
+                self.initialize_components(
+                    embedding_model, reranker, auto_context_model, vector_db, chunk_db
+                )
+                self.save()  # save the config for the KB to disk
         else:
-            created_time = time.time()
-            # We don't care about the milliseconds
-            created_time = int(created_time)
             self.kb_metadata = {
-                'title': title,
-                'description': description,
-                'language': language,
-                'supp_id': supp_id,
-                'created_on': created_time,
+                "title": title,
+                "description": description,
+                "language": language,
+                "supp_id": supp_id,
             }
-            self.initialize_components(embedding_model, reranker, auto_context_model, vector_db, chunk_db)
-            self.save() # save the config for the KB to disk
+            self.initialize_components(
+                embedding_model, reranker, auto_context_model, vector_db, chunk_db
+            )
 
     def get_metadata_path(self):
-        return os.path.join(self.storage_directory, 'metadata', f'{self.kb_id}.json')
-
-    def initialize_components(self, embedding_model, reranker, auto_context_model, vector_db, chunk_db):
+        return os.path.join(self.storage_directory, "metadata", f"{self.kb_id}.json")
+
+    def initialize_components(
+        self,
+        embedding_model: Embedding | None,
+        reranker: Reranker | None,
+        auto_context_model: LLM | None,
+        vector_db: VectorDB | None,
+        chunk_db: ChunkDB | None,
+    ):
         self.embedding_model = embedding_model if embedding_model else OpenAIEmbedding()
         self.reranker = reranker if reranker else CohereReranker()
-        self.auto_context_model = auto_context_model if auto_context_model else OpenAIChatAPI()
-        self.vector_db = vector_db if vector_db else BasicVectorDB(self.kb_id, self.storage_directory)
-        self.chunk_db = chunk_db if chunk_db else BasicChunkDB(self.kb_id, self.storage_directory)
+        self.auto_context_model = (
+            auto_context_model if auto_context_model else OpenAIChatAPI()
+        )
+        self.vector_db = (
+            vector_db
+            if vector_db
+            else BasicVectorDB(self.kb_id, self.storage_directory)
+        )
+        self.chunk_db = (
+            chunk_db if chunk_db else BasicChunkDB(self.kb_id, self.storage_directory)
+        )
         self.vector_dimension = self.embedding_model.dimension
 
     def save(self):
         # Serialize components
         components = {
-            'embedding_model': self.embedding_model.to_dict(),
-            'reranker': self.reranker.to_dict(),
-            'auto_context_model': self.auto_context_model.to_dict(),
-            'vector_db': self.vector_db.to_dict(),
-            'chunk_db': self.chunk_db.to_dict()
+            "embedding_model": self.embedding_model.to_dict(),
+            "reranker": self.reranker.to_dict(),
+            "auto_context_model": self.auto_context_model.to_dict(),
+            "vector_db": self.vector_db.to_dict(),
+            "chunk_db": self.chunk_db.to_dict(),
         }
         # Combine metadata and components
-        full_data = {**self.kb_metadata, 'components': components}
+        full_data = {**self.kb_metadata, "components": components}
 
-        metadata_dir = os.path.join(self.storage_directory, 'metadata')
+        metadata_dir = os.path.join(self.storage_directory, "metadata")
         if not os.path.exists(metadata_dir):
             os.makedirs(metadata_dir)
 
-        with open(self.get_metadata_path(), 'w') as f:
+        with open(self.get_metadata_path(), "w") as f:
             json.dump(full_data, f, indent=4)
 
     def load(self, auto_context_model=None, reranker=None):
         """
         Note: auto_context_model and reranker can be passed in to override the models in the metadata file. The other components are not overridable because that would break things.
         """
-        with open(self.get_metadata_path(), 'r') as f:
+        with open(self.get_metadata_path(), "r") as f:
             data = json.load(f)
-            self.kb_metadata = {key: value for key, value in data.items() if key != 'components'}
-            components = data.get('components', {})
+            self.kb_metadata = {
+                key: value for key, value in data.items() if key != "components"
+            }
+            components = data.get("components", {})
             # Deserialize components
-            self.embedding_model = Embedding.from_dict(components.get('embedding_model', {}))
-            self.reranker = reranker if reranker else Reranker.from_dict(components.get('reranker', {}))
-            self.auto_context_model = auto_context_model if auto_context_model else LLM.from_dict(components.get('auto_context_model', {}))
-            self.vector_db = VectorDB.from_dict(components.get('vector_db', {}))
-            self.chunk_db = ChunkDB.from_dict(components.get('chunk_db', {}))
+            self.embedding_model = Embedding.from_dict(
+                components.get("embedding_model", {})
+            )
+            self.reranker = (
+                reranker
+                if reranker
+                else Reranker.from_dict(components.get("reranker", {}))
+            )
+            self.auto_context_model = (
+                auto_context_model
+                if auto_context_model
+                else LLM.from_dict(components.get("auto_context_model", {}))
+            )
+            self.vector_db = VectorDB.from_dict(components.get("vector_db", {}))
+            self.chunk_db = ChunkDB.from_dict(components.get("chunk_db", {}))
             self.vector_dimension = self.embedding_model.dimension
 
     def delete(self):
@@ -90,14 +160,23 @@ def delete(self):
         doc_ids_to_delete = self.chunk_db.get_all_doc_ids()
         for doc_id in doc_ids_to_delete:
             self.delete_document(doc_id)
-        
+
         self.chunk_db.delete()
         self.vector_db.delete()
 
         # delete the metadata file
         os.remove(self.get_metadata_path())
 
-    def add_document(self, doc_id: str, text: str, document_title: str = "", auto_context_config: dict = {}, semantic_sectioning_config: dict = {}, chunk_size: int = 800, min_length_for_chunking: int = 1600):
+    def add_document(
+        self,
+        doc_id: str,
+        text: str,
+        document_title: str = "",
+        auto_context_config: dict = {},
+        semantic_sectioning_config: dict = {},
+        chunk_size: int = 800,
+        min_length_for_chunking: int = 1600,
+    ):
         """
         Inputs:
         - doc_id: unique identifier for the document; a file name or path is a good choice
@@ -120,141 +199,215 @@ def add_document(self, doc_id: str, text: str, document_title: str = "", auto_co
 
         # verify that the document does not already exist in the KB - the doc_id should be unique
         if doc_id in self.chunk_db.get_all_doc_ids():
-            print (f"Document with ID {doc_id} already exists in the KB. Skipping...")
+            print(f"Document with ID {doc_id} already exists in the KB. Skipping...")
             return
 
         # semantic sectioning
-        if semantic_sectioning_config.get('use_semantic_sectioning', True):
-            llm_provider = semantic_sectioning_config.get('llm_provider', 'openai')
-            model = semantic_sectioning_config.get('model', 'gpt-4o-mini')
+        if semantic_sectioning_config.get("use_semantic_sectioning", True):
+            llm_provider = semantic_sectioning_config.get("llm_provider", "openai")
+            model = semantic_sectioning_config.get("model", "gpt-4o-mini")
             sections = get_sections(text, llm_provider=llm_provider, model=model)
         else:
             sections = [
                 {
-                    'title': '',
-                    'content': text,
+                    "title": "",
+                    "content": text,
                 }
             ]
 
         # document title and summary
-        if not document_title and auto_context_config.get('use_generated_title', True):
-            document_title_guidance = auto_context_config.get('document_title_guidance', '')
-            document_title = get_document_title(self.auto_context_model, text, document_title_guidance=document_title_guidance)
+        if not document_title and auto_context_config.get("use_generated_title", True):
+            document_title_guidance = auto_context_config.get(
+                "document_title_guidance", ""
+            )
+            document_title = get_document_title(
+                self.auto_context_model,
+                text,
+                document_title_guidance=document_title_guidance,
+            )
         elif not document_title:
             document_title = doc_id
 
-        if auto_context_config.get('get_document_summary', True):
-            document_summarization_guidance = auto_context_config.get('document_summarization_guidance', '')
-            document_summary = get_document_summary(self.auto_context_model, text, document_title=document_title, document_summarization_guidance=document_summarization_guidance)
+        if auto_context_config.get("get_document_summary", True):
+            document_summarization_guidance = auto_context_config.get(
+                "document_summarization_guidance", ""
+            )
+            document_summary = get_document_summary(
+                self.auto_context_model,
+                text,
+                document_title=document_title,
+                document_summarization_guidance=document_summarization_guidance,
+            )
         else:
-            document_summary = ''
+            document_summary = ""
 
         # split the document into chunks
-        chunks = [] # chunks is a list of dictionaries with keys 'chunk_text', 'document_title', 'document_summary', 'section_title', 'section_summary'
+        chunks = (
+            []
+        )  # chunks is a list of dictionaries with keys 'chunk_text', 'document_title', 'document_summary', 'section_title', 'section_summary'
         for section in sections:
-            section_text = section['content']
-            section_title = section['title']
-            
+            section_text = section["content"]
+            section_title = section["title"]
+
             # section summary
-            get_section_summaries = auto_context_config.get('get_section_summaries', False)
+            get_section_summaries = auto_context_config.get(
+                "get_section_summaries", False
+            )
             if get_section_summaries and len(sections) > 1:
-                section_summarization_guidance = auto_context_config.get('section_summarization_guidance', '')
-                section_summary = get_section_summary(self.auto_context_model, section_text, document_title=document_title, section_title=section_title, section_summarization_guidance=section_summarization_guidance)
+                section_summarization_guidance = auto_context_config.get(
+                    "section_summarization_guidance", ""
+                )
+                section_summary = get_section_summary(
+                    self.auto_context_model,
+                    section_text,
+                    document_title=document_title,
+                    section_title=section_title,
+                    section_summarization_guidance=section_summarization_guidance,
+                )
             else:
-                section_summary = ''
+                section_summary = ""
 
             # break section into chunks
             if len(section_text) < min_length_for_chunking:
-                chunks.append({
-                    'chunk_text': section_text,
-                    'document_title': document_title,
-                    'document_summary': document_summary,
-                    'section_title': section_title,
-                    'section_summary': section_summary,
-                })
+                chunks.append(
+                    {
+                        "chunk_text": section_text,
+                        "document_title": document_title,
+                        "document_summary": document_summary,
+                        "section_title": section_title,
+                        "section_summary": section_summary,
+                    }
+                )
             else:
-                section_chunks = self.split_into_chunks(section_text, chunk_size=chunk_size)
+                section_chunks = self.split_into_chunks(
+                    section_text, chunk_size=chunk_size
+                )
                 for chunk in section_chunks:
-                    chunks.append({
-                        'chunk_text': chunk,
-                        'document_title': document_title,
-                        'document_summary': document_summary,
-                        'section_title': section_title,
-                        'section_summary': section_summary,
-                    })
-        
-        print (f'Adding {len(chunks)} chunks to the database')
+                    chunks.append(
+                        {
+                            "chunk_text": chunk,
+                            "document_title": document_title,
+                            "document_summary": document_summary,
+                            "section_title": section_title,
+                            "section_summary": section_summary,
+                        }
+                    )
+
+        print(f"Adding {len(chunks)} chunks to the database")
 
         # prepare the chunks for embedding by prepending the chunk headers
         chunks_to_embed = []
         for i, chunk in enumerate(chunks):
-            chunk_header = get_chunk_header(document_title=chunk['document_title'], document_summary=chunk['document_summary'], section_title=chunk['section_title'], section_summary=chunk['section_summary'])
+            chunk_header = get_chunk_header(
+                document_title=chunk["document_title"],
+                document_summary=chunk["document_summary"],
+                section_title=chunk["section_title"],
+                section_summary=chunk["section_summary"],
+            )
             chunk_to_embed = f"{chunk_header}\n\n{chunk['chunk_text']}"
             chunks_to_embed.append(chunk_to_embed)
 
         # embed the chunks
         if len(chunks) <= 50:
             # if the document is short, we can get all the embeddings at once
-            chunk_embeddings = self.get_embeddings(chunks_to_embed, input_type="document")
+            chunk_embeddings = self.get_embeddings(
+                chunks_to_embed, input_type="document"
+            )
         else:
             # if the document is long, we need to get the embeddings in chunks
             chunk_embeddings = []
             for i in range(0, len(chunks), 50):
-                chunk_embeddings += self.get_embeddings(chunks_to_embed[i:i+50], input_type="document")
+                chunk_embeddings += self.get_embeddings(
+                    chunks_to_embed[i : i + 50], input_type="document"
+                )
 
         # add the chunks to the chunk database
         assert len(chunks) == len(chunk_embeddings) == len(chunks_to_embed)
-        self.chunk_db.add_document(doc_id, {i: {'chunk_text': chunk['chunk_text'], 'document_title': chunk['document_title'], 'document_summary': chunk['document_summary'], 'section_title': chunk['section_title'], 'section_summary': chunk['section_summary']} for i, chunk in enumerate(chunks)})
+        self.chunk_db.add_document(
+            doc_id,
+            {
+                i: {
+                    "chunk_text": chunk["chunk_text"],
+                    "document_title": chunk["document_title"],
+                    "document_summary": chunk["document_summary"],
+                    "section_title": chunk["section_title"],
+                    "section_summary": chunk["section_summary"],
+                }
+                for i, chunk in enumerate(chunks)
+            },
+        )
 
         # create metadata list to add to the vector database
         metadata = []
         for i, chunk in enumerate(chunks):
-            metadata.append({'doc_id': doc_id, 'chunk_index': i, 'chunk_text': chunk['chunk_text'], 'chunk_header': get_chunk_header(document_title=chunk['document_title'], document_summary=chunk['document_summary'], section_title=chunk['section_title'], section_summary=chunk['section_summary'])})
+            metadata.append(
+                {
+                    "doc_id": doc_id,
+                    "chunk_index": i,
+                    "chunk_text": chunk["chunk_text"],
+                    "chunk_header": get_chunk_header(
+                        document_title=chunk["document_title"],
+                        document_summary=chunk["document_summary"],
+                        section_title=chunk["section_title"],
+                        section_summary=chunk["section_summary"],
+                    ),
+                }
+            )
 
         # add the vectors and metadata to the vector database
         self.vector_db.add_vectors(vectors=chunk_embeddings, metadata=metadata)
 
-        self.save() # save to disk after adding a document
+        self.save()  # save to disk after adding a document
 
     def delete_document(self, doc_id: str):
         self.chunk_db.remove_document(doc_id)
         self.vector_db.remove_document(doc_id)
 
-    def get_chunk_text(self, doc_id: str, chunk_index: int) -> str:
+    def get_chunk_text(self, doc_id: str, chunk_index: int) -> Optional[str]:
         return self.chunk_db.get_chunk_text(doc_id, chunk_index)
-    
+
     def get_segment_header(self, doc_id: str, chunk_index: int) -> str:
-        document_title = self.chunk_db.get_document_title(doc_id, chunk_index)
-        document_summary = self.chunk_db.get_document_summary(doc_id, chunk_index)
-        return get_segment_header(document_title=document_title, document_summary=document_summary)
+        document_title = self.chunk_db.get_document_title(doc_id, chunk_index) or ""
+        document_summary = self.chunk_db.get_document_summary(doc_id, chunk_index) or ""
+        return get_segment_header(
+            document_title=document_title, document_summary=document_summary
+        )
 
-    def get_embeddings(self, text: str or list[str], input_type: str = ""):
+    def get_embeddings(self, text: list[str], input_type: str = "") -> list[Vector]:
         return self.embedding_model.get_embeddings(text, input_type)
-    
+
     def split_into_chunks(self, text: str, chunk_size: int):
         """
         Note: it's very important that chunk overlap is set to 0 here, since results are created by concatenating chunks.
         """
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=0, length_function=len)
+        text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=chunk_size, chunk_overlap=0, length_function=len
+        )
         texts = text_splitter.create_documents([text])
         chunks = [text.page_content for text in texts]
         return chunks
 
     def cosine_similarity(self, v1, v2):
-        return np.dot(v1, v2) # since the embeddings are normalized
+        return np.dot(v1, v2)  # since the embeddings are normalized
 
     def search(self, query: str, top_k: int) -> list:
         """
         Get top k most relevant chunks for a given query. This is where we interface with the vector database.
         - returns a list of dictionaries, where each dictionary has the following keys: `metadata` (which contains 'doc_id', 'chunk_index', 'chunk_text', and 'chunk_header') and `similarity`
         """
-        query_vector = self.get_embeddings(query, input_type="query") # embed the query
-        search_results = self.vector_db.search(query_vector, top_k) # do a vector database search
+        query_vector = self.get_embeddings(
+            [query], input_type="query"
+        )  # embed the query
+        search_results = self.vector_db.search(
+            query_vector, top_k
+        )  # do a vector database search
         if len(search_results) == 0:
             return []
-        search_results = self.reranker.rerank_search_results(query, search_results) # rerank search results using a reranker
+        search_results = self.reranker.rerank_search_results(
+            query, search_results
+        )  # rerank search results using a reranker
         return search_results
-    
+
     def get_all_ranked_results(self, search_queries: list[str]):
         """
         - search_queries: list of search queries
@@ -264,15 +417,24 @@ def get_all_ranked_results(self, search_queries: list[str]):
             ranked_results = self.search(query, 200)
             all_ranked_results.append(ranked_results)
         return all_ranked_results
-    
-    def get_segment_text_from_database(self, doc_id: str, chunk_start: int, chunk_end: int) -> str:
-        segment = f"{self.get_segment_header(doc_id=doc_id, chunk_index=chunk_start)}\n\n" # initialize the segment with the segment header
-        for chunk_index in range(chunk_start, chunk_end): # NOTE: end index is non-inclusive
-            chunk_text = self.get_chunk_text(doc_id, chunk_index)
+
+    def get_segment_text_from_database(
+        self, doc_id: str, chunk_start: int, chunk_end: int
+    ) -> str:
+        segment = f"{self.get_segment_header(doc_id=doc_id, chunk_index=chunk_start)}\n\n"  # initialize the segment with the segment header
+        for chunk_index in range(
+            chunk_start, chunk_end
+        ):  # NOTE: end index is non-inclusive
+            chunk_text = self.get_chunk_text(doc_id, chunk_index) or ""
             segment += chunk_text
         return segment.strip()
-    
-    def query(self, search_queries: list[str], rse_params: Union[Dict, str] = "balanced", latency_profiling: bool = False) -> list[dict]:
+
+    def query(
+        self,
+        search_queries: list[str],
+        rse_params: Union[Dict, str] = "balanced",
+        latency_profiling: bool = False,
+    ) -> list[dict]:
         """
         Inputs:
         - search_queries: list of search queries
@@ -298,50 +460,94 @@ def query(self, search_queries: list[str], rse_params: Union[Dict, str] = "balan
             raise ValueError(f"Invalid rse_params preset name: {rse_params}")
 
         # set the RSE parameters
-        default_rse_params = RSE_PARAMS_PRESETS['balanced'] # use the 'balanced' preset as the default for any missing parameters
-        max_length = rse_params.get('max_length', default_rse_params['max_length'])
-        overall_max_length = rse_params.get('overall_max_length', default_rse_params['overall_max_length'])
-        minimum_value = rse_params.get('minimum_value', default_rse_params['minimum_value'])
-        irrelevant_chunk_penalty = rse_params.get('irrelevant_chunk_penalty', default_rse_params['irrelevant_chunk_penalty'])
-        overall_max_length_extension = rse_params.get('overall_max_length_extension', default_rse_params['overall_max_length_extension'])
-        decay_rate = rse_params.get('decay_rate', default_rse_params['decay_rate'])
-        top_k_for_document_selection = rse_params.get('top_k_for_document_selection', default_rse_params['top_k_for_document_selection'])
-
-        overall_max_length += (len(search_queries) - 1) * overall_max_length_extension # increase the overall max length for each additional query
+        default_rse_params = RSE_PARAMS_PRESETS[
+            "balanced"
+        ]  # use the 'balanced' preset as the default for any missing parameters
+        max_length = rse_params.get("max_length", default_rse_params["max_length"])
+        overall_max_length = rse_params.get(
+            "overall_max_length", default_rse_params["overall_max_length"]
+        )
+        minimum_value = rse_params.get(
+            "minimum_value", default_rse_params["minimum_value"]
+        )
+        irrelevant_chunk_penalty = rse_params.get(
+            "irrelevant_chunk_penalty", default_rse_params["irrelevant_chunk_penalty"]
+        )
+        overall_max_length_extension = rse_params.get(
+            "overall_max_length_extension",
+            default_rse_params["overall_max_length_extension"],
+        )
+        decay_rate = rse_params.get("decay_rate", default_rse_params["decay_rate"])
+        top_k_for_document_selection = rse_params.get(
+            "top_k_for_document_selection",
+            default_rse_params["top_k_for_document_selection"],
+        )
+
+        overall_max_length += (
+            len(search_queries) - 1
+        ) * overall_max_length_extension  # increase the overall max length for each additional query
 
         start_time = time.time()
         all_ranked_results = self.get_all_ranked_results(search_queries=search_queries)
         if latency_profiling:
-            print(f"get_all_ranked_results took {time.time() - start_time} seconds to run for {len(search_queries)} queries")
+            print(
+                f"get_all_ranked_results took {time.time() - start_time} seconds to run for {len(search_queries)} queries"
+            )
 
-        document_splits, document_start_points, unique_document_ids = get_meta_document(all_ranked_results=all_ranked_results, top_k_for_document_selection=top_k_for_document_selection)
+        document_splits, document_start_points, unique_document_ids = get_meta_document(
+            all_ranked_results=all_ranked_results,
+            top_k_for_document_selection=top_k_for_document_selection,
+        )
 
         # verify that we have a valid meta-document - otherwise return an empty list of segments
         if len(document_splits) == 0:
             return []
-        
+
         # get the length of the meta-document so we don't have to pass in the whole list of splits
         meta_document_length = document_splits[-1]
 
         # get the relevance values for each chunk in the meta-document and use those to find the best segments
-        all_relevance_values = get_relevance_values(all_ranked_results=all_ranked_results, meta_document_length=meta_document_length, document_start_points=document_start_points, unique_document_ids=unique_document_ids, irrelevant_chunk_penalty=irrelevant_chunk_penalty, decay_rate=decay_rate)
-        best_segments, scores = get_best_segments(all_relevance_values=all_relevance_values, document_splits=document_splits, max_length=max_length, overall_max_length=overall_max_length, minimum_value=minimum_value)
-        
+        all_relevance_values = get_relevance_values(
+            all_ranked_results=all_ranked_results,
+            meta_document_length=meta_document_length,
+            document_start_points=document_start_points,
+            unique_document_ids=unique_document_ids,
+            irrelevant_chunk_penalty=irrelevant_chunk_penalty,
+            decay_rate=decay_rate,
+        )
+        best_segments, scores = get_best_segments(
+            all_relevance_values=all_relevance_values,
+            document_splits=document_splits,
+            max_length=max_length,
+            overall_max_length=overall_max_length,
+            minimum_value=minimum_value,
+        )
+
         # convert the best segments into a list of dictionaries that contain the document id and the start and end of the chunk
         relevant_segment_info = []
         for segment_index, (start, end) in enumerate(best_segments):
             # find the document that this segment starts in
             for i, split in enumerate(document_splits):
-                if start < split: # splits represent the end of each document
-                    doc_start = document_splits[i-1] if i > 0 else 0
-                    relevant_segment_info.append({"doc_id": unique_document_ids[i], "chunk_start": start - doc_start, "chunk_end": end - doc_start}) # NOTE: end index is non-inclusive
+                if start < split:  # splits represent the end of each document
+                    doc_start = document_splits[i - 1] if i > 0 else 0
+                    relevant_segment_info.append(
+                        {
+                            "doc_id": unique_document_ids[i],
+                            "chunk_start": start - doc_start,
+                            "chunk_end": end - doc_start,
+                        }
+                    )  # NOTE: end index is non-inclusive
                     break
-        
+
             score = scores[segment_index]
             relevant_segment_info[-1]["score"] = score
 
         # retrieve the actual text (including segment header) for each of the segments
         for segment_info in relevant_segment_info:
-            segment_info["text"] = (self.get_segment_text_from_database(segment_info["doc_id"], segment_info["chunk_start"], segment_info["chunk_end"]))
+            segment_info["text"] = self.get_segment_text_from_database(
+                segment_info["doc_id"],
+                segment_info["chunk_start"],
+                segment_info["chunk_end"],
+            )
 
-        return relevant_segment_info
\ No newline at end of file
+        return relevant_segment_info
diff --git a/dsrag/vector_db.py b/dsrag/vector_db.py
deleted file mode 100644
index 395101c..0000000
--- a/dsrag/vector_db.py
+++ /dev/null
@@ -1,424 +0,0 @@
-from abc import ABC, abstractmethod
-from sklearn.metrics.pairwise import cosine_similarity
-import pickle
-import os
-import numpy as np
-import weaviate
-import weaviate.classes as wvc
-from weaviate.util import generate_uuid5
-import chromadb
-
-
-class VectorDB(ABC):
-    subclasses = {}
-
-    def __init_subclass__(cls, **kwargs):
-        super().__init_subclass__(**kwargs)
-        cls.subclasses[cls.__name__] = cls
-
-    def to_dict(self):
-        return {
-            'subclass_name': self.__class__.__name__,
-        }
-
-    @classmethod
-    def from_dict(cls, config):
-        subclass_name = config.pop('subclass_name', None)  # Remove subclass_name from config
-        subclass = cls.subclasses.get(subclass_name)
-        if subclass:
-            return subclass(**config)  # Pass the modified config without subclass_name
-        else:
-            raise ValueError(f"Unknown subclass: {subclass_name}")
-
-    @abstractmethod
-    def add_vectors(self, vector, metadata):
-        """
-        Store a list of vectors with associated metadata.
-        """
-        pass
-
-    @abstractmethod
-    def remove_document(self, doc_id):
-        """
-        Remove all vectors and metadata associated with a given document ID.
-        """
-        pass
-
-    @abstractmethod
-    def search(self, query_vector, top_k=10):
-        """
-        Retrieve the top-k closest vectors to a given query vector.
-        - needs to return results as list of dictionaries in this format: 
-        {
-            'metadata': 
-                {
-                    'doc_id': doc_id, 
-                    'chunk_index': chunk_index, 
-                    'chunk_header': chunk_header,
-                    'chunk_text': chunk_text
-                }, 
-            'similarity': similarity, 
-        }
-        """
-        pass
-
-    @abstractmethod
-    def delete(self):
-        """
-        Delete the vector database.
-        """
-        pass
-
-
-class BasicVectorDB(VectorDB):
-    def __init__(self, kb_id: str, storage_directory: str = '~/dsRAG', use_faiss: bool = True):
-        self.kb_id = kb_id
-        self.storage_directory = storage_directory
-        self.use_faiss = use_faiss
-        self.vector_storage_path = os.path.join(self.storage_directory, 'vector_storage', f'{kb_id}.pkl')
-        self.load()
-
-    def add_vectors(self, vectors, metadata):
-        try:
-            assert len(vectors) == len(metadata)
-        except AssertionError:
-            raise ValueError('Error in add_vectors: the number of vectors and metadata items must be the same.')
-        self.vectors.extend(vectors)
-        self.metadata.extend(metadata)
-        self.save()
-
-    def search(self, query_vector, top_k=10):
-        if not self.vectors:
-            return []
-        
-        if self.use_faiss:
-            return self.search_faiss(query_vector, top_k)
-
-        similarities = cosine_similarity([query_vector], self.vectors)[0]
-        indexed_similarities = sorted(enumerate(similarities), key=lambda x: x[1], reverse=True)
-        results = []
-        for i, similarity in indexed_similarities[:top_k]:
-            result = {
-                'metadata': self.metadata[i],
-                'similarity': similarity,
-            }
-            results.append(result)
-        return results
-    
-    def search_faiss(self, query_vector, top_k=10):
-        from faiss.contrib.exhaustive_search import knn
-        import numpy as np
-
-        # Limit top_k to the number of vectors we have - Faiss doesn't automatically handle this
-        top_k = min(top_k, len(self.vectors))
-
-        # faiss expects 2D arrays of vectors
-        vectors_array = np.array(self.vectors).astype('float32').reshape(len(self.vectors), -1)
-        query_vector_array = np.array(query_vector).astype('float32').reshape(1, -1)
-        
-        _, I = knn(query_vector_array, vectors_array, top_k) # I is a list of indices in the corpus_vectors array
-        results = []
-        for i in I[0]:
-            result = {
-                'metadata': self.metadata[i],
-                'similarity': cosine_similarity([query_vector], [self.vectors[i]])[0][0],
-            }
-            results.append(result)
-        return results
-
-    def remove_document(self, doc_id):
-        i = 0
-        while i < len(self.metadata):
-            if self.metadata[i]['doc_id'] == doc_id:
-                del self.vectors[i]
-                del self.metadata[i]
-            else:
-                i += 1
-        self.save()
-
-    def save(self):
-        os.makedirs(os.path.dirname(self.vector_storage_path), exist_ok=True)  # Ensure the directory exists
-        with open(self.vector_storage_path, 'wb') as f:
-            pickle.dump((self.vectors, self.metadata), f)
-
-    def load(self):
-        if os.path.exists(self.vector_storage_path):
-            with open(self.vector_storage_path, 'rb') as f:
-                self.vectors, self.metadata = pickle.load(f)
-        else:
-            self.vectors = []
-            self.metadata = []
-    
-    def delete(self):
-        if os.path.exists(self.vector_storage_path):
-            os.remove(self.vector_storage_path)
-
-    def to_dict(self):
-        return {
-            **super().to_dict(),
-            'kb_id': self.kb_id,
-            'storage_directory': self.storage_directory,
-            'use_faiss': self.use_faiss,
-        }
-
-
-class WeaviateVectorDB(VectorDB):
-    """
-    An implementation of the VectorDB interface for Weaviate using the Python v4 client.
-
-    This class provides methods for adding, removing, and searching for vectorized data
-    within a Weaviate instance.
-    """
-
-    def __init__(
-        self,
-        kb_id: str,
-        http_host="localhost",
-        http_port="8099",
-        http_secure=False,
-        grpc_host="localhost",
-        grpc_port="50052",
-        grpc_secure=False,
-        weaviate_secret="secr3tk3y",
-        init_timeout: int = 2,
-        query_timeout: int = 45,
-        insert_timeout: int = 120,
-        use_embedded_weaviate: bool = False,
-    ):
-        """
-        Initializes a WeaviateVectorDB instance.
-
-        Args:
-            http_host: The hostname of the Weaviate server.
-            http_port: The HTTP port of the Weaviate server.
-            http_secure: Whether to use HTTPS for the connection.
-            grpc_host: The hostname of the Weaviate server for gRPC connections.
-            grpc_port: The gRPC port of the Weaviate server.
-            grpc_secure: Whether to use gRPCs for the connection.
-            class_name: The name of the Weaviate class to use for storing data.
-            kb_id: An optional identifier for the knowledge base.
-        """
-        
-        # save all of these parameters as attributes so they're easily accessible for the to_dict method
-        self.kb_id = kb_id
-        self.http_host = http_host
-        self.http_port = http_port
-        self.http_secure = http_secure
-        self.grpc_host = grpc_host
-        self.grpc_port = grpc_port
-        self.grpc_secure = grpc_secure
-        self.weaviate_secret = weaviate_secret
-        self.init_timeout = init_timeout
-        self.query_timeout = query_timeout
-        self.insert_timeout = insert_timeout
-        self.use_embedded_weaviate = use_embedded_weaviate
-
-        additional_headers = {}
-        if use_embedded_weaviate:
-            additional_headers["ENABLE_MODULES"] = (
-                "backup-filesystem,text2vec-openai,text2vec-cohere,text2vec-huggingface,ref2vec-centroid,generative-openai,qna-openai"
-            )
-            additional_headers["BACKUP_FILESYSTEM_PATH"] = "./weaviate/backups"
-            self.client = weaviate.WeaviateClient(
-                embedded_options=weaviate.embedded.EmbeddedOptions(
-                    persistence_data_path="./weaviate/data",
-                ),
-                additional_headers=additional_headers,
-            )
-        else:
-            connection_params = weaviate.connect.ConnectionParams.from_params(
-                http_host=http_host,
-                http_port=http_port,
-                http_secure=http_secure,
-                grpc_host=grpc_host,
-                grpc_port=grpc_port,
-                grpc_secure=grpc_secure,
-            )
-            self.client = weaviate.WeaviateClient(
-                connection_params=connection_params,
-                auth_client_secret=weaviate.auth.AuthApiKey(weaviate_secret),
-                additional_headers=additional_headers,
-                additional_config=weaviate.classes.init.AdditionalConfig(
-                    timeout=weaviate.classes.init.Timeout(
-                        init=init_timeout, query=query_timeout, insert=insert_timeout
-                    )
-                ),
-            )
-        
-        self.client.connect()
-        self.collection = self.client.collections.get(kb_id) # assume this creates a new collection if it doesn't exist
-
-    def close(self):
-        """
-        Closes the connection to Weaviate.
-        """
-        self.client.close()
-
-    def add_vectors(self, vectors, metadata):
-        """
-        Adds a list of vectors with associated metadata to Weaviate.
-
-        Args:
-            vectors: A list of vector embeddings.
-            metadata: A list of dictionaries containing metadata for each vector.
-
-        Raises:
-            ValueError: If the number of vectors and metadata items do not match.
-        """
-        try:
-            assert len(vectors) == len(metadata)
-        except AssertionError as exc:
-            raise ValueError(
-                "Error in add_vectors: the number of vectors and metadata items must be the same."
-            ) from exc
-
-        with self.collection.batch.dynamic() as batch:
-            for vector, meta in zip(vectors, metadata):
-                doc_id = meta.get("doc_id", "")
-                chunk_text = meta.get("chunk_text", "")
-                chunk_index = meta.get("chunk_index", 0)
-                uuid = generate_uuid5(f"{doc_id}_{chunk_index}")
-                batch.add_object(
-                    properties={
-                        "content": chunk_text,
-                        "doc_id": doc_id,
-                        "chunk_index": chunk_index,
-                        "metadata": meta,
-                    },
-                    vector=vector,
-                    uuid=uuid,
-                )
-
-    def remove_document(self, doc_id):
-        """
-        Removes a document (data object) from Weaviate.
-
-        Args:
-            doc_id: The UUID of the document to remove.
-        """
-        self.collection.data.delete_many(
-            where=wvc.query.Filter.by_property("doc_id").contains_any([doc_id])
-        )
-
-    def search(self, query_vector, top_k=10):
-        """
-        Searches for the top-k closest vectors to the given query vector.
-
-        Args:
-            query_vector: The query vector embedding.
-            top_k: The number of results to return.
-
-        Returns:
-            A list of dictionaries containing the metadata and similarity scores of
-            the top-k results.
-        """
-        # convert the query vector to a list if it's not already
-        if isinstance(query_vector, np.ndarray):
-            query_vector = query_vector.tolist()
-
-        results = []
-        response = self.collection.query.near_vector(
-            near_vector=query_vector,
-            limit=top_k,
-            return_metadata=wvc.query.MetadataQuery(distance=True),
-        )
-        for obj in response.objects:
-            results.append(
-                {
-                    "doc_id": obj.properties["doc_id"],
-                    "metadata": obj.properties["metadata"],
-                    "similarity": 1.0 - obj.metadata.distance,
-                    "vector": obj.vector,
-                }
-            )
-        return results
-
-    def to_dict(self):
-        return {
-            **super().to_dict(),
-            'kb_id': self.kb_id,
-            'http_host': self.http_host,
-            'http_port': self.http_port,
-            'http_secure': self.http_secure,
-            'grpc_host': self.grpc_host,
-            'grpc_port': self.grpc_port,
-            'grpc_secure': self.grpc_secure,
-            'weaviate_secret': self.weaviate_secret,
-            'init_timeout': self.init_timeout,
-            'query_timeout': self.query_timeout,
-            'insert_timeout': self.insert_timeout,
-            'use_embedded_weaviate': self.use_embedded_weaviate,
-        }
-    
-    def delete(self):
-        pass
-
-
-class ChromaDB(VectorDB):
-
-    def __init__(self, kb_id: str, storage_directory: str = '~/dsRAG'):
-        self.kb_id = kb_id
-        self.storage_directory = os.path.expanduser(storage_directory)
-        self.vector_storage_path = os.path.join(self.storage_directory, 'vector_storage')
-        self.client = chromadb.PersistentClient(path=self.vector_storage_path)
-        self.collection = self.client.get_or_create_collection(kb_id, metadata={"hnsw:space": "cosine"})
-    
-    def get_num_vectors(self):
-        return self.collection.count()
-    
-    def add_vectors(self, vectors, metadata):
-        try:
-            assert len(vectors) == len(metadata)
-        except AssertionError:
-            raise ValueError('Error in add_vectors: the number of vectors and metadata items must be the same.')
-        
-        # Create the ids from the metadata, defined as {metadata["doc_id"]}_{metadata["chunk_index"]}
-        ids = [f"{meta['doc_id']}_{meta['chunk_index']}" for meta in metadata]
-
-        self.collection.add(
-            embeddings=vectors,
-            metadatas=metadata,
-            ids=ids
-        )
-    
-    def search(self, query_vector, top_k=10):
-
-        num_vectors = self.get_num_vectors()
-        if num_vectors == 0:
-            return []
-            #raise ValueError('No vectors stored in the database.')
-        
-        query_results = self.collection.query(
-            query_embeddings=query_vector,
-            n_results=top_k,
-            include=["distances", "metadatas"]
-        )
-        
-        metadata = query_results["metadatas"][0]
-        distances = query_results["distances"][0]
-
-        results = []
-        for _, (distance, metadata) in enumerate(zip(distances, metadata)):
-            results.append({
-                'metadata': metadata,
-                'similarity': 1 - distance,
-            })
-        
-        results = sorted(results, key=lambda x: x['similarity'], reverse=True)
-
-        return results
-        
-    
-    def remove_document(self, doc_id):
-        data = self.collection.get(where={"doc_id": doc_id})
-        doc_ids = data["ids"]
-        self.collection.delete(ids=doc_ids)
-
-    def delete(self):
-        self.client.delete_collection(name=self.kb_id)
-    
-    def to_dict(self):
-        return {
-            **super().to_dict(),
-            'kb_id': self.kb_id,
-        }
\ No newline at end of file


Test patch:
diff --git a/tests/integration/test_create_kb_and_query_chromadb.py b/tests/integration/test_create_kb_and_query_chromadb.py
index 35a2dec..d14cb7b 100644
--- a/tests/integration/test_create_kb_and_query_chromadb.py
+++ b/tests/integration/test_create_kb_and_query_chromadb.py
@@ -3,38 +3,45 @@
 import unittest
 
 # add ../../ to sys.path
-sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
 
 from dsrag.knowledge_base import KnowledgeBase
 from dsrag.document_parsing import extract_text_from_pdf
-from dsrag.vector_db import ChromaDB
+from dsrag.database.vector import ChromaDB
 
 
-class TestCreateKB(unittest.TestCase):    
+class TestCreateKB(unittest.TestCase):
     def test__001_create_kb_and_query(self):
-        self.cleanup() # delete the KnowledgeBase object if it exists so we can start fresh
-        
+        self.cleanup()  # delete the KnowledgeBase object if it exists so we can start fresh
+
         # Get the absolute path of the script file
         script_dir = os.path.dirname(os.path.abspath(__file__))
 
         # Construct the absolute path to the dataset file
         file_path = os.path.join(script_dir, "../data/levels_of_agi.pdf")
-        
+
         kb_id = "levels_of_agi"
         document_text = extract_text_from_pdf(file_path)
 
         vector_db = ChromaDB(kb_id=kb_id)
         kb = KnowledgeBase(kb_id=kb_id, vector_db=vector_db, exists_ok=False)
-        kb.add_document(doc_id="levels_of_agi.pdf", text=document_text, semantic_sectioning_config={"use_semantic_sectioning": False})
+        kb.add_document(
+            doc_id="levels_of_agi.pdf",
+            text=document_text,
+            semantic_sectioning_config={"use_semantic_sectioning": False},
+        )
 
         # verify that the document is in the chunk db
         self.assertEqual(len(kb.chunk_db.get_all_doc_ids()), 1)
 
         # run a query and verify results are returned
-        search_queries = ["What are the levels of AGI?", "What is the highest level of AGI?"]
+        search_queries = [
+            "What are the levels of AGI?",
+            "What is the highest level of AGI?",
+        ]
         segment_info = kb.query(search_queries)
         self.assertGreater(len(segment_info[0]), 0)
-    
+
     def test__002_remove_document(self):
         kb = KnowledgeBase(kb_id="levels_of_agi", exists_ok=True)
         kb.delete_document(doc_id="levels_of_agi.pdf")
@@ -42,7 +49,7 @@ def test__002_remove_document(self):
 
         # delete the KnowledgeBase object
         kb.delete()
-    
+
     def test__003_query_empty_kb(self):
         kb = KnowledgeBase(kb_id="levels_of_agi", exists_ok=True)
         results = kb.query(["What are the levels of AGI?"])
@@ -53,5 +60,6 @@ def cleanup(self):
         kb = KnowledgeBase(kb_id="levels_of_agi", exists_ok=True)
         kb.delete()
 
+
 if __name__ == "__main__":
-    unittest.main()
\ No newline at end of file
+    unittest.main()
diff --git a/tests/integration/test_create_kb_and_query_weaviate.py b/tests/integration/test_create_kb_and_query_weaviate.py
index 1247a5d..e757359 100644
--- a/tests/integration/test_create_kb_and_query_weaviate.py
+++ b/tests/integration/test_create_kb_and_query_weaviate.py
@@ -3,23 +3,23 @@
 import unittest
 
 # add ../../ to sys.path
-sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
 
 from dsrag.knowledge_base import KnowledgeBase
 from dsrag.document_parsing import extract_text_from_pdf
-from dsrag.vector_db import WeaviateVectorDB
+from dsrag.database.vector import WeaviateVectorDB
 
 
-class TestCreateKB(unittest.TestCase):    
+class TestCreateKB(unittest.TestCase):
     def test_create_kb_and_query(self):
-        self.cleanup() # delete the KnowledgeBase object if it exists so we can start fresh
-        
+        self.cleanup()  # delete the KnowledgeBase object if it exists so we can start fresh
+
         # Get the absolute path of the script file
         script_dir = os.path.dirname(os.path.abspath(__file__))
 
         # Construct the absolute path to the dataset file
         file_path = os.path.join(script_dir, "../data/levels_of_agi.pdf")
-        
+
         kb_id = "levels_of_agi"
         document_text = extract_text_from_pdf(file_path)
 
@@ -31,7 +31,10 @@ def test_create_kb_and_query(self):
         self.assertEqual(len(kb.chunk_db.get_all_doc_ids()), 1)
 
         # run a query and verify results are returned
-        search_queries = ["What are the levels of AGI?", "What is the highest level of AGI?"]
+        search_queries = [
+            "What are the levels of AGI?",
+            "What is the highest level of AGI?",
+        ]
         segment_info = kb.query(search_queries)
         self.assertGreater(len(segment_info[0]), 0)
 
@@ -42,5 +45,6 @@ def cleanup(self):
         kb = KnowledgeBase(kb_id="levels_of_agi", exists_ok=True)
         kb.delete()
 
+
 if __name__ == "__main__":
-    unittest.main()
\ No newline at end of file
+    unittest.main()
diff --git a/tests/unit/test_chunk_db.py b/tests/unit/test_chunk_db.py
index 77c592b..2389ca8 100644
--- a/tests/unit/test_chunk_db.py
+++ b/tests/unit/test_chunk_db.py
@@ -1,17 +1,19 @@
 import os
 import sys
 import unittest
+from dsrag.database.chunk.basic_db import BasicChunkDB
+from dsrag.database.chunk.sqlite_db import SQLiteDB
 
-sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
 
-from dsrag.chunk_db import BasicChunkDB, ChunkDB, SQLiteDB
+from dsrag.database.chunk.db import ChunkDB
 import shutil
 
 
 class TestChunkDB(unittest.TestCase):
     def setUp(self):
-        self.storage_directory = '~/test__chunk_db_dsRAG'
-        self.kb_id = 'test_kb'
+        self.storage_directory = "~/test__chunk_db_dsRAG"
+        self.kb_id = "test_kb"
         resolved_test_storage_directory = os.path.expanduser(self.storage_directory)
         if os.path.exists(resolved_test_storage_directory):
             shutil.rmtree(resolved_test_storage_directory)
@@ -19,74 +21,81 @@ def setUp(self):
 
     @classmethod
     def tearDownClass(cls):
-        test_storage_directory = '~/test__chunk_db_dsRAG'
+        test_storage_directory = "~/test__chunk_db_dsRAG"
         resolved_test_storage_directory = os.path.expanduser(test_storage_directory)
         if os.path.exists(resolved_test_storage_directory):
             shutil.rmtree(resolved_test_storage_directory)
         return super().tearDownClass()
 
-
     def test__add_and_get_chunk_text(self):
         db = BasicChunkDB(self.kb_id, self.storage_directory)
-        doc_id = 'doc1'
+        doc_id = "doc1"
         chunks = {
-            0: {'chunk_text': 'Content of chunk 1', 'document_title': 'Title of document 1', 'document_summary': 'Summary of document 1', 'section_title': 'Section title 1', 'section_summary': 'Section summary 1'},
-            1: {'chunk_text': 'Content of chunk 2', 'document_title': 'Title of document 2', 'document_summary': 'Summary of document 2', 'section_title': 'Section title 2', 'section_summary': 'Section summary 2'},
+            0: {
+                "chunk_text": "Content of chunk 1",
+                "document_title": "Title of document 1",
+                "document_summary": "Summary of document 1",
+                "section_title": "Section title 1",
+                "section_summary": "Section summary 1",
+            },
+            1: {
+                "chunk_text": "Content of chunk 2",
+                "document_title": "Title of document 2",
+                "document_summary": "Summary of document 2",
+                "section_title": "Section title 2",
+                "section_summary": "Section summary 2",
+            },
         }
         db.add_document(doc_id, chunks)
         retrieved_chunk = db.get_chunk_text(doc_id, 0)
-        self.assertEqual(retrieved_chunk, chunks[0]['chunk_text'])
+        self.assertEqual(retrieved_chunk, chunks[0]["chunk_text"])
 
     def test__get_document_title(self):
         db = BasicChunkDB(self.kb_id, self.storage_directory)
-        doc_id = 'doc1'
-        chunks = {
-            0: {'document_title': 'Title 1', 'chunk_text': 'Content of chunk 1'}
-        }
+        doc_id = "doc1"
+        chunks = {0: {"document_title": "Title 1", "chunk_text": "Content of chunk 1"}}
         db.add_document(doc_id, chunks)
         title = db.get_document_title(doc_id, 0)
-        self.assertEqual(title, 'Title 1')
+        self.assertEqual(title, "Title 1")
 
     def test__get_document_summary(self):
         db = BasicChunkDB(self.kb_id, self.storage_directory)
-        doc_id = 'doc1'
+        doc_id = "doc1"
         chunks = {
-            0: {'document_summary': 'Summary 1', 'chunk_text': 'Content of chunk 1'}
+            0: {"document_summary": "Summary 1", "chunk_text": "Content of chunk 1"}
         }
         db.add_document(doc_id, chunks)
         summary = db.get_document_summary(doc_id, 0)
-        self.assertEqual(summary, 'Summary 1')
+        self.assertEqual(summary, "Summary 1")
 
     def test__get_section_title(self):
         db = BasicChunkDB(self.kb_id, self.storage_directory)
-        doc_id = 'doc1'
-        chunks = {
-            0: {'section_title': 'Title 1', 'chunk_text': 'Content of chunk 1'}
-        }
+        doc_id = "doc1"
+        chunks = {0: {"section_title": "Title 1", "chunk_text": "Content of chunk 1"}}
         db.add_document(doc_id, chunks)
         title = db.get_section_title(doc_id, 0)
-        self.assertEqual(title, 'Title 1')
+        self.assertEqual(title, "Title 1")
 
     def test__get_section_summary(self):
         db = BasicChunkDB(self.kb_id, self.storage_directory)
-        doc_id = 'doc1'
+        doc_id = "doc1"
         chunks = {
-            0: {'section_summary': 'Summary 1', 'chunk_text': 'Content of chunk 1'}
+            0: {"section_summary": "Summary 1", "chunk_text": "Content of chunk 1"}
         }
         db.add_document(doc_id, chunks)
         summary = db.get_section_summary(doc_id, 0)
-        self.assertEqual(summary, 'Summary 1')
-    
+        self.assertEqual(summary, "Summary 1")
+
     def test__get_by_supp_id(self):
         db = SQLiteDB(self.kb_id, self.storage_directory)
-        doc_id = 'doc1'
+        doc_id = "doc1"
         chunks = {
-            0: {'supp_id': 'Supp ID 1', 'chunk_text': 'Content of chunk 1'},
+            0: {"supp_id": "Supp ID 1", "chunk_text": "Content of chunk 1"},
         }
         db.add_document(doc_id, chunks)
-        doc_id = 'doc2'
+        doc_id = "doc2"
         chunks = {
-            0: {'chunk_text': 'Content of chunk 2'},
+            0: {"chunk_text": "Content of chunk 2"},
         }
         db.add_document(doc_id, chunks)
         docs = db.get_all_doc_ids("Supp ID 1")
@@ -95,19 +104,23 @@ def test__get_by_supp_id(self):
 
     def test__remove_document(self):
         db = BasicChunkDB(self.kb_id, self.storage_directory)
-        doc_id = 'doc1'
-        chunks = {
-            0: {'chunk_text': 'Content of chunk 1'}
-        }
+        doc_id = "doc1"
+        chunks = {0: {"chunk_text": "Content of chunk 1"}}
         db.add_document(doc_id, chunks)
         db.remove_document(doc_id)
         self.assertNotIn(doc_id, db.data)
 
     def test__persistence(self):
         db = BasicChunkDB(self.kb_id, self.storage_directory)
-        doc_id = 'doc1'
+        doc_id = "doc1"
         chunks = {
-            0: {'chunk_text': 'Content of chunk 1', 'document_title': 'Title of document 1', 'document_summary': 'Summary of document 1', 'section_title': 'Section title 1', 'section_summary': 'Section summary 1'},
+            0: {
+                "chunk_text": "Content of chunk 1",
+                "document_title": "Title of document 1",
+                "document_summary": "Summary of document 1",
+                "section_title": "Section title 1",
+                "section_summary": "Section summary 1",
+            },
         }
         db.add_document(doc_id, chunks)
         db2 = BasicChunkDB(self.kb_id, self.storage_directory)
@@ -119,100 +132,106 @@ def test__save_and_load_from_dict(self):
         db2 = ChunkDB.from_dict(config)
         assert db2.kb_id == db.kb_id, "Failed to load kb_id from dict."
         self.assertEqual(db2.kb_id, db.kb_id)
-    
+
     def test__delete(self):
         db = BasicChunkDB(self.kb_id, self.storage_directory)
-        doc_id = 'doc1'
-        chunks = {
-            0: {'chunk_text': 'Content of chunk 1'}
-        }
+        doc_id = "doc1"
+        chunks = {0: {"chunk_text": "Content of chunk 1"}}
         db.add_document(doc_id, chunks)
         # Make sure the storage directory exists before deleting it
         self.assertTrue(os.path.exists(db.storage_path))
         db.delete()
         # Make sure the storage directory does not exist
         self.assertFalse(os.path.exists(db.storage_path))
-        
+
 
 class TestSQLiteDB(unittest.TestCase):
 
     def setUp(self):
-        self.storage_directory = '~/test__sqlite_db_dsRAG'
-        self.kb_id = 'test_kb'
+        self.storage_directory = "~/test__sqlite_db_dsRAG"
+        self.kb_id = "test_kb"
         resolved_test_storage_directory = os.path.expanduser(self.storage_directory)
         if os.path.exists(resolved_test_storage_directory):
             shutil.rmtree(resolved_test_storage_directory)
         return super().setUp()
-    
+
     @classmethod
     def tearDownClass(cls):
-        test_storage_directory = '~/test__chunk_db_dsRAG'
+        test_storage_directory = "~/test__chunk_db_dsRAG"
         resolved_test_storage_directory = os.path.expanduser(test_storage_directory)
         if os.path.exists(resolved_test_storage_directory):
             shutil.rmtree(resolved_test_storage_directory)
         return super().tearDownClass()
-    
+
     def test__add_and_get_chunk_text(self):
         db = SQLiteDB(self.kb_id, self.storage_directory)
-        doc_id = 'doc1'
+        doc_id = "doc1"
         chunks = {
-            0: {'chunk_text': 'Content of chunk 1', 'document_title': 'Title of document 1', 'document_summary': 'Summary of document 1', 'section_title': 'Section title 1', 'section_summary': 'Section summary 1'},
-            1: {'chunk_text': 'Content of chunk 2', 'document_title': 'Title of document 2', 'document_summary': 'Summary of document 2', 'section_title': 'Section title 2', 'section_summary': 'Section summary 2'},
+            0: {
+                "chunk_text": "Content of chunk 1",
+                "document_title": "Title of document 1",
+                "document_summary": "Summary of document 1",
+                "section_title": "Section title 1",
+                "section_summary": "Section summary 1",
+            },
+            1: {
+                "chunk_text": "Content of chunk 2",
+                "document_title": "Title of document 2",
+                "document_summary": "Summary of document 2",
+                "section_title": "Section title 2",
+                "section_summary": "Section summary 2",
+            },
         }
         db.add_document(doc_id, chunks)
         retrieved_chunk = db.get_chunk_text(doc_id, 0)
-        self.assertEqual(retrieved_chunk, chunks[0]['chunk_text'])
+        self.assertEqual(retrieved_chunk, chunks[0]["chunk_text"])
 
     def test__get_document_title(self):
         db = SQLiteDB(self.kb_id, self.storage_directory)
-        doc_id = 'doc1'
-        chunks = {
-            0: {'document_title': 'Title 1', 'chunk_text': 'Content of chunk 1'}
-        }
+        doc_id = "doc1"
+        chunks = {0: {"document_title": "Title 1", "chunk_text": "Content of chunk 1"}}
         db.add_document(doc_id, chunks)
         title = db.get_document_title(doc_id, 0)
-        self.assertEqual(title, 'Title 1')
+        self.assertEqual(title, "Title 1")
 
     def test__get_document_summary(self):
         db = SQLiteDB(self.kb_id, self.storage_directory)
-        doc_id = 'doc1'
+        doc_id = "doc1"
         chunks = {
-            0: {'document_summary': 'Summary 1', 'chunk_text': 'Content of chunk 1'}
+            0: {"document_summary": "Summary 1", "chunk_text": "Content of chunk 1"}
         }
         db.add_document(doc_id, chunks)
         summary = db.get_document_summary(doc_id, 0)
-        self.assertEqual(summary, 'Summary 1')
+        self.assertEqual(summary, "Summary 1")
 
     def test__get_section_title(self):
         db = SQLiteDB(self.kb_id, self.storage_directory)
-        doc_id = 'doc1'
-        chunks = {
-            0: {'section_title': 'Title 1', 'chunk_text': 'Content of chunk 1'}
-        }
+        doc_id = "doc1"
+        chunks = {0: {"section_title": "Title 1", "chunk_text": "Content of chunk 1"}}
         db.add_document(doc_id, chunks)
         title = db.get_section_title(doc_id, 0)
-        self.assertEqual(title, 'Title 1')
-    
+        self.assertEqual(title, "Title 1")
+
     def test__get_section_summary(self):
         db = SQLiteDB(self.kb_id, self.storage_directory)
-        doc_id = 'doc1'
+        doc_id = "doc1"
         chunks = {
-            0: {'section_summary': 'Summary 1', 'chunk_text': 'Content of chunk 1'}
+            0: {"section_summary": "Summary 1", "chunk_text": "Content of chunk 1"}
         }
         db.add_document(doc_id, chunks)
         summary = db.get_section_summary(doc_id, 0)
-        self.assertEqual(summary, 'Summary 1')
-    
+        self.assertEqual(summary, "Summary 1")
+
     def test__get_by_supp_id(self):
         db = SQLiteDB(self.kb_id, self.storage_directory)
-        doc_id = 'doc1'
+        doc_id = "doc1"
         chunks = {
-            0: {'supp_id': 'Supp ID 1', 'chunk_text': 'Content of chunk 1'},
+            0: {"supp_id": "Supp ID 1", "chunk_text": "Content of chunk 1"},
         }
         db.add_document(doc_id, chunks)
-        doc_id = 'doc2'
+        doc_id = "doc2"
         chunks = {
-            0: {'chunk_text': 'Content of chunk 2'},
+            0: {"chunk_text": "Content of chunk 2"},
         }
         db.add_document(doc_id, chunks)
         docs = db.get_all_doc_ids("Supp ID 1")
@@ -221,10 +240,8 @@ def test__get_by_supp_id(self):
 
     def test__remove_document(self):
         db = SQLiteDB(self.kb_id, self.storage_directory)
-        doc_id = 'doc1'
-        chunks = {
-            0: {'chunk_text': 'Content of chunk 1'}
-        }
+        doc_id = "doc1"
+        chunks = {0: {"chunk_text": "Content of chunk 1"}}
         db.add_document(doc_id, chunks)
         db.remove_document(doc_id)
         results = db.get_document(doc_id)
@@ -237,22 +254,19 @@ def test__save_and_load_from_dict(self):
         db2 = ChunkDB.from_dict(config)
         assert db2.kb_id == db.kb_id, "Failed to load kb_id from dict."
         self.assertEqual(db2.kb_id, db.kb_id)
-    
+
     def test__delete(self):
         db = SQLiteDB(self.kb_id, self.storage_directory)
-        doc_id = 'doc1'
-        chunks = {
-            0: {'chunk_text': 'Content of chunk 1'}
-        }
+        doc_id = "doc1"
+        chunks = {0: {"chunk_text": "Content of chunk 1"}}
         db.add_document(doc_id, chunks)
         # Make sure the storage directory exists before deleting it
-        self.assertTrue(os.path.exists(os.path.join(db.db_path, f'{self.kb_id}.db')))
+        self.assertTrue(os.path.exists(os.path.join(db.db_path, f"{self.kb_id}.db")))
         db.delete()
         # Make sure the storage directory does not exist
-        self.assertFalse(os.path.exists(os.path.join(db.db_path, f'{self.kb_id}.db')))
-
+        self.assertFalse(os.path.exists(os.path.join(db.db_path, f"{self.kb_id}.db")))
 
 
 # Run all tests
-if __name__ == '__main__':
-    unittest.main()
\ No newline at end of file
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/unit/test_vector_db.py b/tests/unit/test_vector_db.py
index 6dcfec9..547d59c 100644
--- a/tests/unit/test_vector_db.py
+++ b/tests/unit/test_vector_db.py
@@ -1,20 +1,24 @@
+from typing import Sequence
+from dsrag.database.vector.types import ChunkMetadata
 import numpy as np
 import os
 import sys
 import unittest
 
-sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
-from dsrag.vector_db import BasicVectorDB, VectorDB, WeaviateVectorDB
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
+from dsrag.database.vector import BasicVectorDB, VectorDB, WeaviateVectorDB
 
 
 class TestVectorDB(unittest.TestCase):
     def setUp(self):
-        self.storage_directory = '~/test__vector_db_dsRAG'
-        self.kb_id = 'test_db'
+        self.storage_directory = "~/test__vector_db_dsRAG"
+        self.kb_id = "test_db"
         return super().setUp()
 
     def tearDown(self):
-        storage_path = os.path.join(self.storage_directory, 'vector_storage', f'{self.kb_id}.pkl')
+        storage_path = os.path.join(
+            self.storage_directory, "vector_storage", f"{self.kb_id}.pkl"
+        )
         if os.path.exists(storage_path):
             os.remove(storage_path)
         return super().tearDown()
@@ -22,104 +26,226 @@ def tearDown(self):
     def test__add_vectors_and_search(self):
         db = BasicVectorDB(self.kb_id, self.storage_directory)
         vectors = [np.array([1, 0]), np.array([0, 1])]
-        metadata = [{'doc_id': '1', 'chunk_index': 0, 'chunk_header': 'Header1', 'chunk_text': 'Text1'},
-                    {'doc_id': '2', 'chunk_index': 1, 'chunk_header': 'Header2', 'chunk_text': 'Text2'}]
-        
+        metadata: Sequence[ChunkMetadata] = [
+            {
+                "doc_id": "1",
+                "chunk_index": 0,
+                "chunk_header": "Header1",
+                "chunk_text": "Text1",
+                "document_title": None,
+                "document_summary": None,
+                "section_title": None,
+                "section_summary": None,
+            },
+            {
+                "doc_id": "2",
+                "chunk_index": 1,
+                "chunk_header": "Header2",
+                "chunk_text": "Text2",
+                "document_title": None,
+                "document_summary": None,
+                "section_title": None,
+                "section_summary": None,
+            },
+        ]
+
         db.add_vectors(vectors, metadata)
         query_vector = np.array([1, 0])
-        results = db.search(query_vector, top_k=1)      
+        results = db.search(query_vector, top_k=1)
 
         self.assertEqual(len(results), 1)
-        self.assertEqual(results[0]['metadata']['doc_id'], '1')
-        self.assertGreaterEqual(results[0]['similarity'], 0.99)
+        self.assertEqual(results[0]["metadata"]["doc_id"], "1")
+        self.assertGreaterEqual(results[0]["similarity"], 0.99)
 
     def test__remove_document(self):
         db = BasicVectorDB(self.kb_id, self.storage_directory)
         vectors = [np.array([1, 0]), np.array([0, 1])]
-        metadata = [{'doc_id': '1', 'chunk_index': 0, 'chunk_header': 'Header1', 'chunk_text': 'Text1'},
-                    {'doc_id': '2', 'chunk_index': 1, 'chunk_header': 'Header2', 'chunk_text': 'Text2'}]
-        
+        metadata: Sequence[ChunkMetadata] = [
+            {
+                "doc_id": "1",
+                "chunk_index": 0,
+                "chunk_header": "Header1",
+                "chunk_text": "Text1",
+                "document_title": None,
+                "document_summary": None,
+                "section_title": None,
+                "section_summary": None,
+            },
+            {
+                "doc_id": "2",
+                "chunk_index": 1,
+                "chunk_header": "Header2",
+                "chunk_text": "Text2",
+                "document_title": None,
+                "document_summary": None,
+                "section_title": None,
+                "section_summary": None,
+            },
+        ]
+
         db.add_vectors(vectors, metadata)
-        db.remove_document('1')
-        
+        db.remove_document("1")
+
         print(db.metadata)
         self.assertEqual(len(db.metadata), 1)
-        self.assertEqual(db.metadata[0]['doc_id'], '2')
+        self.assertEqual(db.metadata[0]["doc_id"], "2")
 
     def test__empty_search(self):
         db = BasicVectorDB(self.kb_id, self.storage_directory)
         query_vector = np.array([1, 0])
         results = db.search(query_vector)
-        
+
         self.assertEqual(len(results), 0)
 
     def test__save_and_load(self):
         db = BasicVectorDB(self.kb_id, self.storage_directory)
         vectors = [np.array([1, 0]), np.array([0, 1])]
-        metadata = [{'doc_id': '1', 'chunk_index': 0, 'chunk_header': 'Header1', 'chunk_text': 'Text1'},
-                    {'doc_id': '2', 'chunk_index': 1, 'chunk_header': 'Header2', 'chunk_text': 'Text2'}]
-        
+        metadata: Sequence[ChunkMetadata] = [
+            {
+                "doc_id": "1",
+                "chunk_index": 0,
+                "chunk_header": "Header1",
+                "chunk_text": "Text1",
+                "document_title": None,
+                "document_summary": None,
+                "section_title": None,
+                "section_summary": None,
+            },
+            {
+                "doc_id": "2",
+                "chunk_index": 1,
+                "chunk_header": "Header2",
+                "chunk_text": "Text2",
+                "document_title": None,
+                "document_summary": None,
+                "section_title": None,
+                "section_summary": None,
+            },
+        ]
+
         db.add_vectors(vectors, metadata)
         db.save()
-        
+
         new_db = BasicVectorDB(self.kb_id, self.storage_directory)
         new_db.load()
-        
+
         self.assertEqual(len(new_db.metadata), 2)
-        self.assertEqual(new_db.metadata[0]['doc_id'], '1')
-        self.assertEqual(new_db.metadata[1]['doc_id'], '2')
+        self.assertEqual(new_db.metadata[0]["doc_id"], "1")
+        self.assertEqual(new_db.metadata[1]["doc_id"], "2")
 
     def test__load_from_dict(self):
         config = {
-            'subclass_name': 'BasicVectorDB',
-            'kb_id': 'test_db',
-            'storage_directory': '/tmp'
+            "subclass_name": "BasicVectorDB",
+            "kb_id": "test_db",
+            "storage_directory": "/tmp",
         }
         vector_db_instance = VectorDB.from_dict(config)
         self.assertIsInstance(vector_db_instance, BasicVectorDB)
-        self.assertEqual(vector_db_instance.kb_id, 'test_db')
+        self.assertEqual(vector_db_instance.kb_id, "test_db")
 
     def test__save_and_load_from_dict(self):
         db = BasicVectorDB(self.kb_id, self.storage_directory)
         config = db.to_dict()
         vector_db_instance = VectorDB.from_dict(config)
         self.assertIsInstance(vector_db_instance, BasicVectorDB)
-        self.assertEqual(vector_db_instance.kb_id, 'test_db')
+        self.assertEqual(vector_db_instance.kb_id, "test_db")
 
     def test__assertion_error_on_mismatched_input_lengths(self):
         db = BasicVectorDB(self.kb_id, self.storage_directory)
         vectors = [np.array([1, 0])]
-        metadata = [{'doc_id': '1', 'chunk_index': 0, 'chunk_header': 'Header1', 'chunk_text': 'Text1'},
-                    {'doc_id': '2', 'chunk_index': 1, 'chunk_header': 'Header2', 'chunk_text': 'Text2'}]
-        
+        metadata: Sequence[ChunkMetadata] = [
+            {
+                "doc_id": "1",
+                "chunk_index": 0,
+                "chunk_header": "Header1",
+                "chunk_text": "Text1",
+                "document_title": None,
+                "document_summary": None,
+                "section_title": None,
+                "section_summary": None,
+            },
+            {
+                "doc_id": "2",
+                "chunk_index": 1,
+                "chunk_header": "Header2",
+                "chunk_text": "Text2",
+                "document_title": None,
+                "document_summary": None,
+                "section_title": None,
+                "section_summary": None,
+            },
+        ]
+
         with self.assertRaises(ValueError) as context:
             db.add_vectors(vectors, metadata)
-        self.assertTrue('Error in add_vectors: the number of vectors and metadata items must be the same.' in str(context.exception))
+        self.assertTrue(
+            "Error in add_vectors: the number of vectors and metadata items must be the same."
+            in str(context.exception)
+        )
 
     def test__faiss_search(self):
         db = BasicVectorDB(self.kb_id, self.storage_directory, use_faiss=True)
         vectors = [np.array([1, 0]), np.array([0, 1])]
-        metadata = [{'doc_id': '1', 'chunk_index': 0, 'chunk_header': 'Header1', 'chunk_text': 'Text1'},
-                    {'doc_id': '2', 'chunk_index': 1, 'chunk_header': 'Header2', 'chunk_text': 'Text2'}]
-        
+        metadata: Sequence[ChunkMetadata] = [
+            {
+                "doc_id": "1",
+                "chunk_index": 0,
+                "chunk_header": "Header1",
+                "chunk_text": "Text1",
+                "document_title": None,
+                "document_summary": None,
+                "section_title": None,
+                "section_summary": None,
+            },
+            {
+                "doc_id": "2",
+                "chunk_index": 1,
+                "chunk_header": "Header2",
+                "chunk_text": "Text2",
+                "document_title": None,
+                "document_summary": None,
+                "section_title": None,
+                "section_summary": None,
+            },
+        ]
+
         db.add_vectors(vectors, metadata)
         query_vector = np.array([1, 0])
-        
+
         faiss_results = db.search(query_vector, top_k=1)
 
         db.use_faiss = False
         non_faiss_results = db.search(query_vector, top_k=1)
 
         self.assertEqual(faiss_results, non_faiss_results)
-    
 
     def test__delete(self):
 
         db = BasicVectorDB(self.kb_id, self.storage_directory, use_faiss=True)
         vectors = [np.array([1, 0]), np.array([0, 1])]
-        metadata = [{'doc_id': '1', 'chunk_index': 0, 'chunk_header': 'Header1', 'chunk_text': 'Text1'},
-                    {'doc_id': '2', 'chunk_index': 1, 'chunk_header': 'Header2', 'chunk_text': 'Text2'}]
-        
+        metadata: Sequence[ChunkMetadata] = [
+            {
+                "doc_id": "1",
+                "chunk_index": 0,
+                "chunk_header": "Header1",
+                "chunk_text": "Text1",
+                "document_title": None,
+                "document_summary": None,
+                "section_title": None,
+                "section_summary": None,
+            },
+            {
+                "doc_id": "2",
+                "chunk_index": 1,
+                "chunk_header": "Header2",
+                "chunk_text": "Text2",
+                "document_title": None,
+                "document_summary": None,
+                "section_title": None,
+                "section_summary": None,
+            },
+        ]
+
         db.add_vectors(vectors, metadata)
 
         # Make sure the storage directory exists before deleting it
@@ -131,9 +257,29 @@ def test__delete(self):
     def test__top_k_greater_than_num_vectors(self):
         db = BasicVectorDB(self.kb_id, self.storage_directory)
         vectors = [np.array([1, 0]), np.array([0, 1])]
-        metadata = [{'doc_id': '1', 'chunk_index': 0, 'chunk_header': 'Header1', 'chunk_text': 'Text1'},
-                    {'doc_id': '2', 'chunk_index': 1, 'chunk_header': 'Header2', 'chunk_text': 'Text2'}]
-        
+        metadata: Sequence[ChunkMetadata] = [
+            {
+                "doc_id": "1",
+                "chunk_index": 0,
+                "chunk_header": "Header1",
+                "chunk_text": "Text1",
+                "document_title": None,
+                "document_summary": None,
+                "section_title": None,
+                "section_summary": None,
+            },
+            {
+                "doc_id": "2",
+                "chunk_index": 1,
+                "chunk_header": "Header2",
+                "chunk_text": "Text2",
+                "document_title": None,
+                "document_summary": None,
+                "section_title": None,
+                "section_summary": None,
+            },
+        ]
+
         db.add_vectors(vectors, metadata)
         query_vector = np.array([1, 0])
 
@@ -160,9 +306,38 @@ def tearDown(self):
 
     def test_add_vectors_and_search(self):
         vectors = [np.array([1, 0]), np.array([0, 1]), np.array([1, 1])]
-        metadata = [{'doc_id': '1', 'chunk_index': 0, 'chunk_header': 'Header1', 'chunk_text': 'Text1'},
-                    {'doc_id': '1', 'chunk_index': 1, 'chunk_header': 'Header2', 'chunk_text': 'Text2'},
-                    {'doc_id': '2', 'chunk_index': 0, 'chunk_header': 'Header3', 'chunk_text': 'Text3'}]
+        metadata: Sequence[ChunkMetadata] = [
+            {
+                "doc_id": "1",
+                "chunk_index": 0,
+                "chunk_header": "Header1",
+                "chunk_text": "Text1",
+                "document_title": None,
+                "document_summary": None,
+                "section_title": None,
+                "section_summary": None,
+            },
+            {
+                "doc_id": "1",
+                "chunk_index": 1,
+                "chunk_header": "Header2",
+                "chunk_text": "Text2",
+                "document_title": None,
+                "document_summary": None,
+                "section_title": None,
+                "section_summary": None,
+            },
+            {
+                "doc_id": "2",
+                "chunk_index": 0,
+                "chunk_header": "Header3",
+                "chunk_text": "Text3",
+                "document_title": None,
+                "document_summary": None,
+                "section_title": None,
+                "section_summary": None,
+            },
+        ]
         self.db.add_vectors(vectors, metadata)
 
         query_vector = np.array([1, 0])
@@ -177,9 +352,38 @@ def test_add_vectors_and_search(self):
 
     def test_remove_document(self):
         vectors = [np.array([1, 0]), np.array([0, 1]), np.array([1, 1])]
-        metadata = [{'doc_id': '1', 'chunk_index': 0, 'chunk_header': 'Header1', 'chunk_text': 'Text1'},
-                    {'doc_id': '1', 'chunk_index': 1, 'chunk_header': 'Header2', 'chunk_text': 'Text2'},
-                    {'doc_id': '2', 'chunk_index': 0, 'chunk_header': 'Header3', 'chunk_text': 'Text3'}]
+        metadata: Sequence[ChunkMetadata] = [
+            {
+                "doc_id": "1",
+                "chunk_index": 0,
+                "chunk_header": "Header1",
+                "chunk_text": "Text1",
+                "document_title": None,
+                "document_summary": None,
+                "section_title": None,
+                "section_summary": None,
+            },
+            {
+                "doc_id": "1",
+                "chunk_index": 1,
+                "chunk_header": "Header2",
+                "chunk_text": "Text2",
+                "document_title": None,
+                "document_summary": None,
+                "section_title": None,
+                "section_summary": None,
+            },
+            {
+                "doc_id": "2",
+                "chunk_index": 0,
+                "chunk_header": "Header3",
+                "chunk_text": "Text3",
+                "document_title": None,
+                "document_summary": None,
+                "section_title": None,
+                "section_summary": None,
+            },
+        ]
         self.db.add_vectors(vectors, metadata)
 
         self.db.remove_document("1")
@@ -193,15 +397,35 @@ def test_remove_document(self):
 
     def test__save_and_load(self):
         vectors = [np.array([1, 0]), np.array([0, 1])]
-        metadata = [{'doc_id': '1', 'chunk_index': 0, 'chunk_header': 'Header1', 'chunk_text': 'Text1'},
-                    {'doc_id': '2', 'chunk_index': 1, 'chunk_header': 'Header2', 'chunk_text': 'Text2'}]
-        
+        metadata: Sequence[ChunkMetadata] = [
+            {
+                "doc_id": "1",
+                "chunk_index": 0,
+                "chunk_header": "Header1",
+                "chunk_text": "Text1",
+                "document_title": None,
+                "document_summary": None,
+                "section_title": None,
+                "section_summary": None,
+            },
+            {
+                "doc_id": "2",
+                "chunk_index": 1,
+                "chunk_header": "Header2",
+                "chunk_text": "Text2",
+                "document_title": None,
+                "document_summary": None,
+                "section_title": None,
+                "section_summary": None,
+            },
+        ]
+
         self.db.add_vectors(vectors, metadata)
         self.db.close()
-        
+
         # load the saved db
         self.db = WeaviateVectorDB(kb_id=self.kb_id, use_embedded_weaviate=True)
-        
+
         # Verify data existence indirectly (Weaviate doesn't provide a direct way to list documents)
         query_vector = np.array([1, 0])
         results = self.db.search(query_vector, top_k=2)
@@ -217,5 +441,5 @@ def test__save_and_load_from_dict(self):
         self.assertEqual(self.db.kb_id, self.kb_id)
 
 
-if __name__ == '__main__':
-    unittest.main()
\ No newline at end of file
+if __name__ == "__main__":
+    unittest.main()