Repository: D-Star-AI/dsRAG

Issue:
feat: execute queries `get_all_ranked_results` concurrently
Hi, I've been testing `dsRAG` and it's working really well. However, right now `get_all_ranked_results` performs vector search sequentially for all queries. If we use something like `auto_query` and end up with > 10 search queries the search takes quite a while.

I've made a small change to use `ThreadPoolExecutor` like so:
```python
    def get_all_ranked_results(self, search_queries: list[str]):
        """
        - search_queries: list of search queries
        """
        with concurrent.futures.ThreadPoolExecutor() as executor:
            future_to_query = {executor.submit(self.search, query, 200): query for query in search_queries}
            all_ranked_results = []
            for future in concurrent.futures.as_completed(future_to_query):
                ranked_results = future.result()
                all_ranked_results.append(ranked_results)
        
        return all_ranked_results
```

...which cuts down the query time massively in case of many `search_queries`.

I'd be happy to make a PR for this if you find it useful!

Discussion:

Patch:
diff --git a/.github/workflows/manual-workflow.yml b/.github/workflows/manual-workflow.yml
new file mode 100644
index 0000000..56b64b9
--- /dev/null
+++ b/.github/workflows/manual-workflow.yml
@@ -0,0 +1,44 @@
+name: Manual Test Run
+
+on:
+  workflow_dispatch:
+    inputs:
+      repository:
+        description: 'Repository to run tests on (e.g., user/repo)'
+        required: true
+      branch:
+        description: 'Branch to run tests on'
+        required: true
+
+jobs:
+  run-tests:
+    runs-on: ubuntu-latest
+
+    env:
+      ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
+      CO_API_KEY: ${{ secrets.CO_API_KEY }}
+      OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+      VOYAGE_API_KEY: ${{ secrets.VOYAGE_API_KEY }}
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v2
+
+      - name: Set up Python
+        uses: actions/setup-python@v2
+        with:
+          python-version: '3.x'
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -r requirements.txt
+          pip install pytest
+
+      - name: Run unit tests
+        run: |
+          pytest tests/unit
+
+      - name: Run integration tests
+        run: |
+          pytest tests/integration
diff --git a/dsrag/database/chunk/basic_db.py b/dsrag/database/chunk/basic_db.py
index 1a620d5..1011c37 100644
--- a/dsrag/database/chunk/basic_db.py
+++ b/dsrag/database/chunk/basic_db.py
@@ -25,7 +25,7 @@ class BasicChunkDB(ChunkDB):
         )
         self.load()
 
-    def add_document(self, doc_id: str, chunks: dict[int, dict[str, Any]]) -> None:
+    def add_document(self, doc_id: str, chunks: dict[int, dict[str, Any]], supp_id: str = "", metadata: dict = {}) -> None:
         self.data[doc_id] = chunks
         self.save()
 
@@ -48,7 +48,7 @@ class BasicChunkDB(ChunkDB):
             full_document_string = ""
             if include_content:
                 # Concatenate the chunks into a single string
-                for chunk_index, chunk in document.items():
+                for _, chunk in document.items():
                     # Join each chunk text with a new line character
                     full_document_string += chunk["chunk_text"] + "\n"
 
diff --git a/dsrag/database/chunk/db.py b/dsrag/database/chunk/db.py
index b9ef4ae..4b69108 100644
--- a/dsrag/database/chunk/db.py
+++ b/dsrag/database/chunk/db.py
@@ -28,7 +28,7 @@ class ChunkDB(ABC):
             raise ValueError(f"Unknown subclass: {subclass_name}")
 
     @abstractmethod
-    def add_document(self, doc_id: str, chunks: dict[int, dict[str, Any]]) -> None:
+    def add_document(self, doc_id: str, chunks: dict[int, dict[str, Any]], supp_id: str = "", metadata: dict = {}) -> None:
         """
         Store all chunks for a given document.
         """
diff --git a/dsrag/database/chunk/sqlite_db.py b/dsrag/database/chunk/sqlite_db.py
index 2d4c441..8f86696 100644
--- a/dsrag/database/chunk/sqlite_db.py
+++ b/dsrag/database/chunk/sqlite_db.py
@@ -16,6 +16,19 @@ class SQLiteDB(ChunkDB):
             os.path.join(self.storage_directory, "chunk_storage"), exist_ok=True
         )
         self.db_path = os.path.join(self.storage_directory, "chunk_storage")
+        self.columns = [
+            {"name": "doc_id", "type": "TEXT"},
+            {"name": "document_title", "type": "TEXT"},
+            {"name": "document_summary", "type": "TEXT"},
+            {"name": "section_title", "type": "TEXT"},
+            {"name": "section_summary", "type": "TEXT"},
+            {"name": "chunk_text", "type": "TEXT"},
+            {"name": "chunk_index", "type": "INT"},
+            {"name": "chunk_length", "type": "INT"},
+            {"name": "created_on", "type": "TEXT"},
+            {"name": "supp_id", "type": "TEXT"},
+            {"name": "metadata", "type": "TEXT"},
+        ]
 
         # 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"))
@@ -25,28 +38,33 @@ class SQLiteDB(ChunkDB):
         )
         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)"
-            )
+            query_statement = "CREATE TABLE documents ("
+            for column in self.columns:
+                query_statement += f"{column['name']} {column['type']}, "
+            query_statement = query_statement[:-2] + ")"
+            c.execute(query_statement)
             conn.commit()
         else:
-            # Check if we need to add the columns to the table for the supp_id and created_on fields
+            # Check if we need to add any columns to the table. This happens if the columns have been updated
             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")
+            for column in self.columns:
+                if column["name"] not in column_names:
+                    # Add the column to the table
+                    c.execute("ALTER TABLE documents ADD COLUMN {} {}".format(column["name"], column["type"]))
         conn.close()
 
-    def add_document(self, doc_id: str, chunks: dict[int, dict[str, Any]]) -> None:
+    def add_document(self, doc_id: str, chunks: dict[int, dict[str, Any]], supp_id: str = "", metadata: dict = {}) -> 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()))
 
+        # Turn the metadata object into a string
+        metadata = str(metadata)
+
         # Get the data from the dictionary
         for chunk_index, chunk in chunks.items():
             document_title = chunk.get("document_title", "")
@@ -54,9 +72,9 @@ class SQLiteDB(ChunkDB):
             section_title = chunk.get("section_title", "")
             section_summary = chunk.get("section_summary", "")
             chunk_text = chunk.get("chunk_text", "")
-            supp_id = chunk.get("supp_id", "")
+            chunk_length = len(chunk_text)
             c.execute(
-                "INSERT INTO documents (doc_id, document_title, document_summary, section_title, section_summary, chunk_text, chunk_index, created_on, supp_id) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
+                "INSERT INTO documents (doc_id, document_title, document_summary, section_title, section_summary, chunk_text, chunk_index, chunk_length, created_on, supp_id, metadata) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
                 (
                     doc_id,
                     document_title,
@@ -65,8 +83,10 @@ class SQLiteDB(ChunkDB):
                     section_summary,
                     chunk_text,
                     chunk_index,
+                    chunk_length,
                     created_on,
                     supp_id,
+                    metadata
                 ),
             )
 
@@ -87,7 +107,7 @@ class SQLiteDB(ChunkDB):
         # 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"]
+        columns = ["supp_id", "document_title", "document_summary", "created_on", "metadata"]
         if include_content:
             columns += ["chunk_text", "chunk_index"]
 
@@ -110,17 +130,24 @@ class SQLiteDB(ChunkDB):
                 # Join each chunk text with a new line character
                 full_document_string += result[4] + "\n"
 
-        title = results[0][1]
-        created_on = results[0][3]
+        supp_id = results[0][0]
         title = results[0][1]
         summary = results[0][2]
+        created_on = results[0][3]
+        metadata = results[0][4]
+
+        # Convert the metadata string back into a dictionary
+        if metadata:
+            metadata = eval(metadata)
 
         return FormattedDocument(
             id=doc_id,
+            supp_id=supp_id,
             title=title,
             content=full_document_string if include_content else None,
             summary=summary,
             created_on=created_on,
+            metadata=metadata
         )
 
     def get_chunk_text(self, doc_id: str, chunk_index: int) -> Optional[str]:
@@ -199,6 +226,28 @@ class SQLiteDB(ChunkDB):
         results = c.fetchall()
         conn.close()
         return [result[0] for result in results]
+    
+    def get_document_count(self) -> int:
+        # Retrieve the number of documents in the sqlite table
+        conn = sqlite3.connect(os.path.join(self.db_path, f'{self.kb_id}.db'))
+        c = conn.cursor()
+        c.execute(f"SELECT COUNT(DISTINCT doc_id) FROM documents")
+        result = c.fetchone()
+        conn.close()
+        if result is None:
+            return 0
+        return result[0]
+
+    def get_total_num_characters(self) -> int:
+        # Retrieve the total number of characters in the sqlite table
+        conn = sqlite3.connect(os.path.join(self.db_path, f'{self.kb_id}.db'))
+        c = conn.cursor()
+        c.execute(f"SELECT SUM(chunk_length) FROM documents")
+        result = c.fetchone()
+        conn.close()
+        if result is None or result[0] is None:
+            return 0
+        return result[0]
 
     def delete(self) -> None:
         # Delete the sqlite database
diff --git a/dsrag/database/chunk/types.py b/dsrag/database/chunk/types.py
index eac2f93..725a328 100644
--- a/dsrag/database/chunk/types.py
+++ b/dsrag/database/chunk/types.py
@@ -9,3 +9,5 @@ class FormattedDocument(TypedDict):
     content: Optional[str]
     summary: Optional[str]
     created_on: Optional[datetime]
+    supp_id: Optional[str]
+    metadata: Optional[dict]
diff --git a/dsrag/database/vector/chroma_db.py b/dsrag/database/vector/chroma_db.py
index ce78311..f97b16e 100644
--- a/dsrag/database/vector/chroma_db.py
+++ b/dsrag/database/vector/chroma_db.py
@@ -47,7 +47,6 @@ class ChromaDB(VectorDB):
             include=["distances", "metadatas"],
         )
 
-        metadata = query_results["metadatas"][0]
         metadata = query_results["metadatas"][0]
         distances = query_results["distances"][0]
 
diff --git a/dsrag/database/vector/types.py b/dsrag/database/vector/types.py
index f8dc873..4a318a7 100644
--- a/dsrag/database/vector/types.py
+++ b/dsrag/database/vector/types.py
@@ -7,10 +7,6 @@ class ChunkMetadata(TypedDict):
     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 = Union[Sequence[float], Sequence[int]]
diff --git a/dsrag/knowledge_base.py b/dsrag/knowledge_base.py
index 4634175..c44ad32 100644
--- a/dsrag/knowledge_base.py
+++ b/dsrag/knowledge_base.py
@@ -4,6 +4,7 @@ import os
 import time
 import json
 from typing import Optional, Union, Dict
+import concurrent.futures
 from dsrag.auto_context import (
     get_document_title,
     get_document_summary,
@@ -176,6 +177,8 @@ class KnowledgeBase:
         semantic_sectioning_config: dict = {},
         chunk_size: int = 800,
         min_length_for_chunking: int = 1600,
+        supp_id: str = "",
+        metadata: dict = {},
     ):
         """
         Inputs:
@@ -195,6 +198,9 @@ class KnowledgeBase:
             - use_semantic_sectioning: if False, semantic sectioning will be skipped (default is True)
         - chunk_size: the maximum number of characters to include in each chunk
         - min_length_for_chunking: the minimum length of text to allow chunking (measured in number of characters); if the text is shorter than this, it will be added as a single chunk. If semantic sectioning is used, this parameter will be applied to each section. Setting this to a higher value than the chunk_size can help avoid unnecessary chunking of short documents or sections.
+        - document_type: the type of document being added (Can be any string you like. Useful for filtering documents later on.)
+        - supp_id: supplementary ID for the document (Can be any string you like. Useful for filtering documents later on.)
+        - file_name: the name of the file that the document came from (if applicable)
         """
 
         # verify that the document does not already exist in the KB - the doc_id should be unique
@@ -335,6 +341,8 @@ class KnowledgeBase:
                 }
                 for i, chunk in enumerate(chunks)
             },
+            supp_id,
+            metadata
         )
 
         # create metadata list to add to the vector database
@@ -412,10 +420,14 @@ class KnowledgeBase:
         """
         - search_queries: list of search queries
         """
-        all_ranked_results = []
-        for query in search_queries:
-            ranked_results = self.search(query, 200)
-            all_ranked_results.append(ranked_results)
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            futures = [executor.submit(self.search, query, 200) for query in search_queries]
+            
+            all_ranked_results = []
+            for future in futures:
+                ranked_results = future.result()
+                all_ranked_results.append(ranked_results)
+        
         return all_ranked_results
 
     def get_segment_text_from_database(
@@ -550,4 +562,4 @@ class KnowledgeBase:
                 segment_info["chunk_end"],
             )
 
-        return relevant_segment_info
\ No newline at end of file
+        return relevant_segment_info


Test patch:
diff --git a/tests/unit/test_chunk_db.py b/tests/unit/test_chunk_db.py
index d72b219..1da5844 100644
--- a/tests/unit/test_chunk_db.py
+++ b/tests/unit/test_chunk_db.py
@@ -86,22 +86,6 @@ class TestChunkDB(unittest.TestCase):
         summary = db.get_section_summary(doc_id, 0)
         self.assertEqual(summary, "Summary 1")
 
-    def test__get_by_supp_id(self):
-        db = SQLiteDB(self.kb_id, self.storage_directory)
-        doc_id = "doc1"
-        chunks = {
-            0: {"supp_id": "Supp ID 1", "chunk_text": "Content of chunk 1"},
-        }
-        db.add_document(doc_id, chunks)
-        doc_id = "doc2"
-        chunks = {
-            0: {"chunk_text": "Content of chunk 2"},
-        }
-        db.add_document(doc_id, chunks)
-        docs = db.get_all_doc_ids("Supp ID 1")
-        # There should only be one document with the supp_id 'Supp ID 1'
-        self.assertEqual(len(docs), 1)
-
     def test__remove_document(self):
         db = BasicChunkDB(self.kb_id, self.storage_directory)
         doc_id = "doc1"
@@ -225,15 +209,16 @@ class TestSQLiteDB(unittest.TestCase):
     def test__get_by_supp_id(self):
         db = SQLiteDB(self.kb_id, self.storage_directory)
         doc_id = "doc1"
+        supp_id = "Supp ID 1"
         chunks = {
-            0: {"supp_id": "Supp ID 1", "chunk_text": "Content of chunk 1"},
+            0: {"chunk_text": "Content of chunk 1"},
         }
-        db.add_document(doc_id, chunks)
+        db.add_document(doc_id=doc_id, chunks=chunks, supp_id=supp_id)
         doc_id = "doc2"
         chunks = {
             0: {"chunk_text": "Content of chunk 2"},
         }
-        db.add_document(doc_id, chunks)
+        db.add_document(doc_id=doc_id, chunks=chunks)
         docs = db.get_all_doc_ids("Supp ID 1")
         # There should only be one document with the supp_id 'Supp ID 1'
         self.assertEqual(len(docs), 1)
diff --git a/tests/unit/test_vector_db.py b/tests/unit/test_vector_db.py
index 8ed3646..92a3a1c 100644
--- a/tests/unit/test_vector_db.py
+++ b/tests/unit/test_vector_db.py
@@ -32,20 +32,12 @@ class TestVectorDB(unittest.TestCase):
                 "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,
             },
         ]
 
@@ -66,20 +58,12 @@ class TestVectorDB(unittest.TestCase):
                 "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,
             },
         ]
 
@@ -106,20 +90,12 @@ class TestVectorDB(unittest.TestCase):
                 "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,
             },
         ]
 
@@ -159,20 +135,12 @@ class TestVectorDB(unittest.TestCase):
                 "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,
             },
         ]
 
@@ -192,20 +160,12 @@ class TestVectorDB(unittest.TestCase):
                 "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,
             },
         ]
 
@@ -229,20 +189,12 @@ class TestVectorDB(unittest.TestCase):
                 "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,
             },
         ]
 
@@ -263,20 +215,12 @@ class TestVectorDB(unittest.TestCase):
                 "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,
             },
         ]
 
@@ -304,7 +248,7 @@ class TestWeaviateVectorDB(unittest.TestCase):
         self.db.close()
         return super().tearDown()
 
-    def test_add_vectors_and_search(self):
+    def test__add_vectors_and_search(self):
         vectors = [np.array([1, 0]), np.array([0, 1]), np.array([1, 1])]
         metadata: Sequence[ChunkMetadata] = [
             {
@@ -312,30 +256,18 @@ class TestWeaviateVectorDB(unittest.TestCase):
                 "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)
@@ -350,7 +282,7 @@ class TestWeaviateVectorDB(unittest.TestCase):
         self.assertEqual(results[0]["metadata"]["chunk_text"], "Text1")
         self.assertGreaterEqual(results[0]["similarity"], 0.99)
 
-    def test_remove_document(self):
+    def test__remove_document(self):
         vectors = [np.array([1, 0]), np.array([0, 1]), np.array([1, 1])]
         metadata: Sequence[ChunkMetadata] = [
             {
@@ -358,30 +290,18 @@ class TestWeaviateVectorDB(unittest.TestCase):
                 "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)
@@ -403,20 +323,12 @@ class TestWeaviateVectorDB(unittest.TestCase):
                 "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,
             },
         ]