fix(tests): engine tests

Author: Daniele Briggi

src/sqlite_rag/engine.py

@@ -22,12 +22,12 @@ def __init__(
         conn: sqlite3.Connection,
         settings: Settings,
         chunker: Chunker,
-        sentence_chunker: SentenceSplitter,
+        sentence_splitter: SentenceSplitter,
     ):
         self._conn = conn
         self._settings = settings
         self._chunker = chunker
-        self._sentence_chunker = sentence_chunker
+        self._sentence_splitter = sentence_splitter
         self._logger = Logger()
 
     def load_model(self):
@@ -55,7 +55,7 @@ def process(self, document: Document) -> Document:
             chunk.title = document.get_title()
             chunk.embedding = self.generate_embedding(chunk.get_embedding_text())
 
-            sentences = self._sentence_chunker.split(chunk)
+            sentences = self._sentence_splitter.split(chunk)
             for sentence in sentences:
                 sentence.embedding = self.generate_embedding(sentence.content)
             chunk.sentences = sentences
@@ -225,7 +225,7 @@ def search(
         return results
 
     def search_sentences(
-        self, query: str, chunk_id: int, k: int
+        self, query: str, chunk_id: int, top_k: int
     ) -> List[SentenceResult]:
         query_embedding = self.generate_embedding(query)
 
@@ -250,7 +250,8 @@ def search_sentences(
                 FROM {vector_scan_type}('sentences', 'embedding', :query_embedding) AS v
                     JOIN sentences ON sentences.rowid = v.rowid
                 WHERE sentences.chunk_id = :chunk_id
-                LIMIT :k
+                ORDER BY rank_number ASC
+                LIMIT :top_k
             )
             SELECT
                 sentence_id,
@@ -260,11 +261,10 @@ def search_sentences(
                 rank_number,
                 distance
             FROM vec_matches
-            ORDER BY rank_number ASC
             """,  # nosec B608
             {
                 "query_embedding": query_embedding,
-                "k": k,
+                "top_k": top_k,
                 "chunk_id": chunk_id,
             },
         )
@@ -283,7 +283,7 @@ def search_sentences(
                 )
             )
 
-        return sentences[:k]
+        return sentences[:top_k]
 
     def versions(self) -> dict:
         """Get versions of the loaded extensions."""

src/sqlite_rag/sentence_splitter.py

@@ -10,23 +10,23 @@ class SentenceSplitter:
 
     def split(self, chunk: Chunk) -> List[Sentence]:
         """Split chunk into sentences."""
-        sentence_chunks = []
+        sentences = []
 
-        sentences = self._split_into_sentences(chunk.content)
+        sentences_text = self._split_into_sentences(chunk.content)
         start_offset = 0
         end_offset = 0
-        for sentence in sentences:
-            start_offset = chunk.content.index(sentence, end_offset)
-            end_offset = start_offset + len(sentence)
+        for sentence_text in sentences_text:
+            start_offset = chunk.content.index(sentence_text, end_offset)
+            end_offset = start_offset + len(sentence_text)
 
-            sentence_chunk = Sentence(
-                content=sentence,
+            sentence = Sentence(
+                content=sentence_text,
                 start_offset=start_offset,
                 end_offset=end_offset,
             )
-            sentence_chunks.append(sentence_chunk)
+            sentences.append(sentence)
 
-        return sentence_chunks
+        return sentences
 
     def _split_into_sentences(self, text: str) -> List[str]:
         """Split into focused segments for semantic matching."""

src/sqlite_rag/sqliterag.py

@@ -31,7 +31,7 @@ def __init__(self, connection: sqlite3.Connection, settings: Settings):
             self._conn,
             settings,
             chunker=self._chunker,
-            sentence_chunker=SentenceSplitter(),
+            sentence_splitter=SentenceSplitter(),
         )
         self._extractor = Extractor()
 
@@ -332,7 +332,7 @@ def search(
         # Refine chunks with top sentences
         for result in results:
             result.sentences = self._engine.search_sentences(
-                semantic_query, result.chunk_id, k=self._settings.top_k_sentences
+                semantic_query, result.chunk_id, top_k=self._settings.top_k_sentences
             )
 
         return results

tests/conftest.py

@@ -6,6 +6,7 @@
 from sqlite_rag.chunker import Chunker
 from sqlite_rag.database import Database
 from sqlite_rag.engine import Engine
+from sqlite_rag.sentence_splitter import SentenceSplitter
 from sqlite_rag.settings import Settings
 
 
@@ -28,7 +29,12 @@ def db_conn():
 def engine(db_conn) -> Engine:
     conn, settings = db_conn
 
-    engine = Engine(conn, settings, chunker=Chunker(conn, settings))
+    engine = Engine(
+        conn,
+        settings,
+        chunker=Chunker(conn, settings),
+        sentence_splitter=SentenceSplitter(),
+    )
     engine.load_model()
     engine.quantize()
     engine.create_new_context()