feat(chunks): set limit per document

Author: Daniele Briggi

src/sqlite_rag/cli.py

@@ -162,6 +162,10 @@ def configure_settings(
         None,
         help="Template for retrieval query prompts, use `{content}` as placeholder",
     ),
+    max_chunks_per_document: Optional[int] = typer.Option(
+        None,
+        help="Maximum number of chunks to generate per document (0 for no limit)",
+    ),
 ):
     """Configure settings for the RAG system.
 
@@ -190,6 +194,7 @@ def configure_settings(
         ),  # Set only if True
         "prompt_template_retrieval_document": prompt_template_retrieval_document,
         "prompt_template_retrieval_query": prompt_template_retrieval_query,
+        "max_chunks_per_document": max_chunks_per_document,
     }
     print(updates)
     # Filter out None values (unset options)

src/sqlite_rag/engine.py

@@ -36,6 +36,10 @@ def load_model(self):
 
     def process(self, document: Document) -> Document:
         chunks = self._chunker.chunk(document.content, document.metadata)
+
+        if self._settings.max_chunks_per_document > 0:
+            chunks = chunks[: self._settings.max_chunks_per_document]
+
         chunks = self.generate_embedding(chunks)
         document.chunks = chunks
         return document

src/sqlite_rag/settings.py

@@ -63,6 +63,13 @@ class Settings:
     prompt_template_retrieval_document: str = "title: {title} | text: {content}"
     prompt_template_retrieval_query: str = "task: search result | query: {content}"
 
+    #
+    # Index settings
+    #
+
+    # Zero means no limit
+    max_chunks_per_document: int = 1000
+
 
 class SettingsManager:
     def __init__(self, connection: sqlite3.Connection):

tests/test_engine.py

@@ -53,6 +53,74 @@ def test_generate_embedding_with_prompt_template(
             "SELECT llm_embed_generate(?) AS embedding", (expected_content,)
         )
 
+    def test_extract_document_title(self):
+        text = """# This is the Title
+        This is the content of the document.
+        It has multiple lines.
+        """
+
+        engine = Engine(None, Settings(), None)  # type: ignore
+
+        title = engine.extract_document_title(text)
+        assert title == "This is the Title"
+
+    @pytest.mark.parametrize(
+        "fallback, expected_title",
+        [
+            (True, "This is the first line of the document without a title."),
+            (False, None),
+        ],
+    )
+    def test_extract_document_title_from_first_line(self, fallback, expected_title):
+        text = """
+        This is the first line of the document without a title.
+        It has multiple lines.
+        """
+
+        engine = Engine(None, Settings(), None)  # type: ignore
+
+        title = engine.extract_document_title(text, fallback)
+        assert title == expected_title
+
+    @pytest.mark.parametrize(
+        "max_chunks_per_document, expected_chunk_count",
+        [(0, 2), (1, 1), (4, 2)],
+    )
+    def test_process_with_max_chunks_per_document(
+        self, mocker, max_chunks_per_document, expected_chunk_count
+    ):
+        # Arrange
+        chunks = [
+            Chunk(content="Chunk 1"),
+            Chunk(content="Chunk 2"),
+            Chunk(content="Chunk 3"),
+        ]
+
+        mock_conn = mocker.Mock()
+        settings = Settings(max_chunks_per_document=max_chunks_per_document)
+        mock_chunker = mocker.Mock()
+        mock_chunker.chunk.return_value = chunks
+
+        engine = Engine(mock_conn, settings, mock_chunker)
+
+        mock_generate_embedding = mocker.patch.object(engine, "generate_embedding")
+        mock_generate_embedding = mocker.spy(
+            mock_generate_embedding, "generate_embedding"
+        )
+        mock_generate_embedding.return_value = chunks
+
+        document = Document(content="Test document content")
+
+        # Act
+        engine.process(document)
+
+        # Assert
+        for call_args in mock_generate_embedding.call_args_list:
+            chunks = call_args[0][0]  # First argument
+            assert len(chunks) == expected_chunk_count
+
+
+class TestEngineSearch:
     def test_search_with_empty_database(self, engine):
         results = engine.search("nonexistent query", top_k=5)
 
@@ -230,32 +298,3 @@ def test_search_exact_match(self, db_conn):
         assert len(results) > 0
         assert doc1_id == results[0].document.id
         assert 0.0 == results[0].vec_distance
-
-    def test_extract_document_title(self):
-        text = """# This is the Title
-        This is the content of the document.
-        It has multiple lines.
-        """
-
-        engine = Engine(None, Settings(), None)  # type: ignore
-
-        title = engine.extract_document_title(text)
-        assert title == "This is the Title"
-
-    @pytest.mark.parametrize(
-        "fallback, expected_title",
-        [
-            (True, "This is the first line of the document without a title."),
-            (False, None),
-        ],
-    )
-    def test_extract_document_title_from_first_line(self, fallback, expected_title):
-        text = """
-        This is the first line of the document without a title.
-        It has multiple lines.
-        """
-
-        engine = Engine(None, Settings(), None)  # type: ignore
-
-        title = engine.extract_document_title(text, fallback)
-        assert title == expected_title