feat: refactory, tests and logger

Author: Daniele Briggi

src/sqlite_rag/logger.py

@@ -0,0 +1,35 @@
+from enum import Enum
+from typing import Callable, Optional
+
+
+class LogLevel(Enum):
+    DEBUG = 1
+    INFO = 2
+    WARNING = 3
+    ERROR = 4
+
+
+class Logger:
+    def __init__(
+        self,
+        level: LogLevel = LogLevel.INFO,
+        output_fn: Optional[Callable[[str], None]] = None,
+    ):
+        self.level = level
+        self.output_fn = output_fn or print
+
+    def debug(self, message: str):
+        if self.level.value <= LogLevel.DEBUG.value:
+            self.output_fn(f"DEBUG: {message}")
+
+    def info(self, message: str):
+        if self.level.value <= LogLevel.INFO.value:
+            self.output_fn(f"{message}")
+
+    def warning(self, message: str):
+        if self.level.value <= LogLevel.WARNING.value:
+            self.output_fn(f"WARNING: {message}")
+
+    def error(self, message: str):
+        if self.level.value <= LogLevel.ERROR.value:
+            self.output_fn(f"ERROR: {message}")

src/sqlite_rag/reader.py

@@ -51,3 +51,27 @@ def parse_file(path: Path) -> str:
             return reader.convert(path).text_content
         except Exception as exc:
             raise ValueError(f"Failed to parse file {path}") from exc
+
+    @staticmethod
+    def collect_files(path: Path, recursive: bool = False) -> list[Path]:
+        """Collect files from the path, optionally recursively"""
+        if not path.exists():
+            raise FileNotFoundError(f"{path} does not exist.")
+
+        if path.is_file() and FileReader.is_supported(path):
+            return [path]
+
+        files_to_process = []
+        if path.is_dir():
+            if recursive:
+                files_to_process = list(path.rglob("*"))
+            else:
+                files_to_process = list(path.glob("*"))
+
+            files_to_process = [
+                f
+                for f in files_to_process
+                if f.is_file() and FileReader.is_supported(f)
+            ]
+
+        return files_to_process

src/sqlite_rag/sqliterag.py

@@ -1,8 +1,10 @@
-import glob
-from pathlib import Path
+import re
 import sqlite3
+from pathlib import Path
 from typing import Optional
 
+from sqlite_rag.logger import Logger
+
 from .chunker import Chunker
 from .database import Database
 from .engine import Engine
@@ -17,10 +19,12 @@ def __init__(self, settings: Optional[Settings] = None):
         if settings is None:
             # TODO: load defaults or from the database
             settings = Settings(
-                model_path_or_name="./all-MiniLM-L6-v2.e4ce9877.q8_0.gguf", db_path="sqliterag.db"
+                model_path_or_name="./all-MiniLM-L6-v2.e4ce9877.q8_0.gguf",
+                db_path="sqliterag.db",
             )
 
         self.settings = settings
+        self._logger = Logger()
 
         self._conn = sqlite3.connect(settings.db_path)
 
@@ -37,54 +41,39 @@ def _ensure_initialized(self):
 
         self.ready = True
 
-    def add(self, path: str, recursively: bool = False):
+    def add(self, path: str, recursive: bool = False) -> int:
         """Add the file content into the database"""
         self._ensure_initialized()
 
         if not Path(path).exists():
             raise FileNotFoundError(f"{path} does not exist.")
 
-        files_to_process: list[Path] = []
-        path_obj = Path(path)
-
-        if path_obj.is_file():
-            files_to_process.append(path_obj)
-        elif path_obj.is_dir():
-            if recursively:
-                files_to_process = list(path_obj.rglob("*"))
-            else:
-                files_to_process = list(path_obj.glob("*"))
-
-            files_to_process = [
-                f
-                for f in files_to_process
-                if f.is_file() and FileReader.is_supported(f)
-            ]
+        files_to_process = FileReader.collect_files(Path(path), recursive=recursive)
 
+        self._logger.info(f"Processing {len(files_to_process)} files...")
         for file_path in files_to_process:
-            # TODO: check the file extension
-            content = FileReader.parse_file(file_path)
             # TODO: include metadata extraction and mdx options (see our docsearch)
-            document = Document(content=content, uri=str(file_path))
-            chunks = self._chunker.chunk(document.content)
-            chunks = self._engine.generate_embedding(chunks)
-            document.chunks = chunks
+            content = FileReader.parse_file(file_path)
+            document = self._engine.process(
+                Document(content=content, uri=str(file_path))
+            )
 
             self._repository.add_document(document)
+            self._logger.info(str(file_path))
 
         if self.settings.quantize_scan:
             self._engine.quantize()
 
+        return len(files_to_process)
+
     def add_text(
         self, text: str, uri: Optional[str] = None, metadata: dict = {}
     ) -> None:
         """Add a text content into the database"""
         self._ensure_initialized()
 
         document = Document(content=text, uri=uri, metadata=metadata)
-        chunks = self._chunker.chunk(document.content)
-        chunks = self._engine.generate_embedding(chunks)
-        document.chunks = chunks
+        document = self._engine.process(document)
 
         self._repository.add_document(document)
 
@@ -105,16 +94,16 @@ def find_document(self, identifier: str) -> Document | None:
     def remove_document(self, identifier: str) -> bool:
         """Remove document by ID or URI"""
         self._ensure_initialized()
-        
+
         # First find the document to get its ID
         document = self._repository.find_document_by_id_or_uri(identifier)
         if not document:
             return False
-        
+
         return self._repository.remove_document(document.id or "")
 
     # def search(
     #     self, query: str, top_k: int = 10
     # ) -> list[Document]:
     #     """Search for documents matching the query"""
-    #     self._ensure_initialized()
+    #     self._ensure_initialized()

tests/test_reader.py

@@ -0,0 +1,78 @@
+import tempfile
+from pathlib import Path
+
+from sqlite_rag.reader import FileReader
+
+
+class TestFileReader:
+    def test_collect_files_empty_directory(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            files = FileReader.collect_files(Path(temp_dir), recursive=False)
+            assert len(files) == 0
+
+    def test_collect_files_single_file(self):
+        with tempfile.NamedTemporaryFile(suffix=".txt", delete=False) as f:
+            f.write(b"This is a test document.")
+            temp_file_path = f.name
+
+        files = FileReader.collect_files(Path(temp_file_path), recursive=False)
+        assert len(files) == 1
+        assert Path(temp_file_path) in files
+
+    def test_collect_files(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            temp_dir = Path(temp_dir)
+            file1 = temp_dir / "test1.txt"
+            file2 = temp_dir / "test2.txt"
+            file1.write_text("This is the first test document.")
+            file2.write_text("This is the second test document.")
+
+            files = FileReader.collect_files(temp_dir, recursive=False)
+            assert len(files) == 2
+            assert file1 in files
+            assert file2 in files
+
+    def test_collect_ignore_unsupported_files(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            sub_dir = Path(temp_dir) / "subdir"
+            sub_dir.mkdir()
+            file1 = Path(temp_dir) / "test1.txt"
+            file2 = sub_dir / "test2.unsupported"
+            file1.write_text("This is a test document.")
+            file2.write_text("This is an unsupported file type.")
+
+            files = FileReader.collect_files(Path(temp_dir), recursive=True)
+            assert len(files) == 1
+            assert file1 in files
+            assert file2 not in files
+
+    def test_collect_files_recursive_directory(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            sub_dir = Path(temp_dir) / "subdir"
+            sub_dir.mkdir()
+            file1 = Path(temp_dir) / "file1.txt"
+            file2 = sub_dir / "file2.txt"
+
+            file1.write_text("This is the first test document.")
+            file2.write_text("This is the second test document.")
+
+            files = FileReader.collect_files(Path(temp_dir), recursive=True)
+            assert len(files) == 2
+            assert file1 in files
+            assert file2 in files
+
+    def test_is_supported(self):
+        unsupported_extensions = [".exe", ".bin", ".jpg", ".png"]
+
+        for ext in FileReader.extensions:
+            assert FileReader.is_supported(Path(f"test{ext}"))
+
+        for ext in unsupported_extensions:
+            assert not FileReader.is_supported(Path(f"test{ext}"))
+
+    def test_parse_html_into_markdown(self):
+        with tempfile.NamedTemporaryFile(suffix=".html", delete=False) as f:
+            f.write(b"<html><body><h1>This is a test markdown file.</h1></body></html>")
+
+        content = FileReader.parse_file(Path(f.name))
+        assert "# This is a test markdown file." in content

tests/test_sqlite_rag.py

@@ -26,6 +26,18 @@ def test_add_simple_text_file(self, db_settings):
         chunk_count = cursor.fetchone()[0]
         assert chunk_count > 0
 
+    def test_add_unsupported_file_type(self, db_settings):
+        # Create a temporary file with an unsupported extension
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".unsupported", delete=False) as f:
+            f.write("This is a test document with unsupported file type.")
+
+        rag = SQLiteRag(db_settings)
+
+        # Attempt to add the unsupported file
+        processed = rag.add(f.name)
+
+        assert processed == 0
+
     def test_add_directory(self, db_settings):
         with tempfile.TemporaryDirectory() as temp_dir:
             file1 = Path(temp_dir) / "file1.txt"
@@ -34,9 +46,6 @@ def test_add_directory(self, db_settings):
             file1.write_text("This is the first test document.")
             file2.write_text("This is the second test document.")
 
-            db_settings.chunk_size = 100
-            db_settings.chunk_overlap = 10
-
             rag = SQLiteRag(db_settings)
 
             rag.add(temp_dir)
@@ -50,6 +59,30 @@ def test_add_directory(self, db_settings):
             chunk_count = cursor.fetchone()[0]
             assert chunk_count > 0
 
+    def test_add_directory_recursively(self, db_settings):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            sub_dir = Path(temp_dir) / "subdir"
+            sub_dir.mkdir()
+
+            file1 = Path(temp_dir) / "file1.txt"
+            file2 = sub_dir / "file2.txt"
+
+            file1.write_text("This is the first test document.")
+            file2.write_text("This is the second test document in a subdirectory.")
+
+            rag = SQLiteRag(db_settings)
+
+            rag.add(temp_dir, recursive=True)
+
+            conn = rag._conn
+            cursor = conn.execute("SELECT COUNT(*) FROM documents")
+            doc_count = cursor.fetchone()[0]
+            assert doc_count == 2
+
+            cursor = conn.execute("SELECT COUNT(*) FROM chunks")
+            chunk_count = cursor.fetchone()[0]
+            assert chunk_count > 0
+
     def test_add_text(self, db_settings):
         rag = SQLiteRag(db_settings)