Merge pull request #7 from sqliteai/extract-frontmatter

feat(extractor): file specific extractor

Author: danielebriggi

.gitignore

@@ -29,6 +29,8 @@ share/python-wheels/
 .installed.cfg
 *.egg
 MANIFEST
+.venv/
+venv/
 
 # PyInstaller
 *.manifest

pyproject.toml

@@ -25,6 +25,7 @@ dependencies = [
     "markitdown[pptx]",
     "markitdown[xls]",
     "markitdown[xlsx]",
+    "python-frontmatter",
     "prompt-toolkit",
     "sqlite-ai",
     "sqliteai-vector"
@@ -34,7 +35,7 @@ dependencies = [
 dev = [
     "pytest",
     "pytest-mock",
-    "pytest-cov==6.3.0",
+    "pytest-cov",
     "black",
     "flake8",
     "bandit",

src/sqlite_rag/extractor.py

@@ -0,0 +1,37 @@
+from pathlib import Path
+from typing import Dict, Optional, Tuple
+
+from sqlite_rag.extractors.base import MetadataExtractor
+from sqlite_rag.extractors.frontmatter import FrontmatterExtractor
+
+
+class Extractor:
+    extractors = [
+        FrontmatterExtractor(),
+    ]
+
+    def get_extractor(self, file_extension: str) -> Optional[MetadataExtractor]:
+        """Get the appropriate extractor based on file type."""
+        for extractor in self.extractors:
+            if extractor.supports_file_type(file_extension):
+                return extractor
+
+        return None
+
+    def extract_metadata(self, content: str, file_path: Path) -> Tuple[str, Dict]:
+        """Extract metadata and clean content based on file type.
+
+        Args:
+            content: Raw content to extract metadata from
+            file_path: Path to the file for context
+
+        Returns:
+            Tuple of (clean_content, metadata_dict)
+        """
+        file_extension = file_path.suffix
+
+        extractor = self.get_extractor(file_extension)
+        if extractor:
+            return extractor.extract(content, file_path)
+
+        return content, {}

src/sqlite_rag/extractors/__init__.py

@@ -0,0 +1,32 @@
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Dict, Optional, Tuple
+
+
+class MetadataExtractor(ABC):
+    """Base interface for metadata extractors."""
+
+    @abstractmethod
+    def extract(
+        self, content: str, file_path: Optional[Path] = None
+    ) -> Tuple[str, Dict]:
+        """Extract metadata from content.
+
+        Args:
+            content: The raw content to extract metadata from
+            file_path: Optional file path for context
+
+        Returns:
+            Tuple of (clean_content, metadata_dict)
+        """
+
+    @abstractmethod
+    def supports_file_type(self, file_extension: str) -> bool:
+        """Check if this extractor supports the given file type.
+
+        Args:
+            file_extension: File extension (e.g., '.md', '.pdf')
+
+        Returns:
+            True if this extractor can handle the file type
+        """

src/sqlite_rag/extractors/base.py

@@ -0,0 +1,27 @@
+from pathlib import Path
+from typing import Dict, Optional, Tuple
+
+import frontmatter
+
+from sqlite_rag.extractors.base import MetadataExtractor
+
+
+class FrontmatterExtractor(MetadataExtractor):
+    """Extracts frontmatter from markdown files."""
+
+    def extract(
+        self, content: str, file_path: Optional[Path] = None
+    ) -> Tuple[str, Dict]:
+        """Extract frontmatter from markdown content."""
+        try:
+            post = frontmatter.loads(content)
+            clean_content = post.content
+            metadata = dict(post.metadata)
+            return clean_content, metadata
+        except Exception:
+            # If frontmatter parsing fails, return original content
+            return content, {}
+
+    def supports_file_type(self, file_extension: str) -> bool:
+        """Support markdown files."""
+        return file_extension.lower() in [".md", ".mdx", ".txt"]

src/sqlite_rag/extractors/frontmatter.py

@@ -47,6 +47,7 @@ def is_supported(path: Path) -> bool:
 
     @staticmethod
     def parse_file(path: Path, max_document_size_bytes: Optional[int] = None) -> str:
+        """Read the file and convert into Markdown text."""
         try:
             converter = MarkItDown()
             text = converter.convert(

src/sqlite_rag/reader.py

@@ -3,6 +3,7 @@
 from pathlib import Path
 from typing import Any, Optional
 
+from sqlite_rag.extractor import Extractor
 from sqlite_rag.logger import Logger
 from sqlite_rag.models.document_result import DocumentResult
 
@@ -25,6 +26,7 @@ def __init__(self, connection: sqlite3.Connection, settings: Settings):
         self._repository = Repository(self._conn, settings)
         self._chunker = Chunker(self._conn, settings)
         self._engine = Engine(self._conn, settings, chunker=self._chunker)
+        self._extractor = Extractor()
 
         self.ready = False
 
@@ -103,7 +105,16 @@ def add(
                     if use_relative_paths
                     else str(file_path.absolute())
                 )
-                document = Document(content=content, uri=uri, metadata=metadata.copy())
+
+                content, file_metadata = self._extractor.extract_metadata(
+                    content, file_path
+                )
+
+                merged_metadata = metadata.copy()
+                if file_metadata:
+                    merged_metadata["extracted"] = file_metadata
+
+                document = Document(content=content, uri=uri, metadata=merged_metadata)
 
                 exists = self._repository.document_exists_by_hash(document.hash())
                 if exists:
@@ -187,7 +198,21 @@ def rebuild(self, remove_missing: bool = False) -> dict:
                     content = FileReader.parse_file(
                         Path(doc.uri), self._settings.max_document_size_bytes
                     )
+
+                    if not content:
+                        self._logger.warning(
+                            f"{i+1}/{total_docs} Skipping empty file: {doc.uri}"
+                        )
+                        not_found += 1
+                        continue
+
+                    content, file_metadata = self._extractor.extract_metadata(
+                        content, Path(doc.uri)
+                    )
+
                     doc.content = content
+                    if file_metadata:
+                        doc.metadata["extracted"] = file_metadata
 
                     self._repository.remove_document(doc_id)
                     processed_doc = self._engine.process(doc)

src/sqlite_rag/sqliterag.py

@@ -0,0 +1,52 @@
+from sqlite_rag.extractors.frontmatter import FrontmatterExtractor
+
+
+class TestFrontmatterExtractor:
+    def test_extract_with_frontmatter(self):
+        content = """---
+title: Test Document
+author: John Doe
+---
+# Heading 1
+This is a test document.
+"""
+        extractor = FrontmatterExtractor()
+        clean_content, metadata = extractor.extract(content)
+        assert "title" in metadata
+        assert metadata["title"] == "Test Document"
+        assert "author" in metadata
+        assert metadata["author"] == "John Doe"
+        assert "# Heading 1" in clean_content
+        assert "This is a test document." in clean_content
+
+    def test_extract_without_frontmatter(self):
+        content = """# Heading 1
+This is a test document without frontmatter.
+"""
+        extractor = FrontmatterExtractor()
+        clean_content, metadata = extractor.extract(content)
+        assert metadata == {}
+        assert "# Heading 1" in clean_content
+        assert "This is a test document without frontmatter." in clean_content
+
+    def test_supports_file_type(self):
+        extractor = FrontmatterExtractor()
+        assert extractor.supports_file_type(".md")
+        assert extractor.supports_file_type(".MDX")
+        assert extractor.supports_file_type(".txt")
+        assert not extractor.supports_file_type(".pdf")
+        assert not extractor.supports_file_type(".html")
+
+    def test_extract_malformed_frontmatter(self):
+        content = """---
+title: Test Document
+author John Doe
+---
+# Heading 1
+"""
+        extractor = FrontmatterExtractor()
+        clean_content, metadata = extractor.extract(content)
+        # Should return original content and empty metadata on failure
+        assert metadata == {}
+        assert content == clean_content
+        assert "# Heading 1" in clean_content

tests/extractors/test_frontmatter.py

@@ -0,0 +1,40 @@
+from pathlib import Path
+
+from sqlite_rag.extractor import Extractor
+
+
+class TestExtractor:
+    def test_extract_metadata_from_md(self):
+        extractor = Extractor()
+        content = """---
+title: Sample Document
+author: Test Author
+---
+# Heading 1
+This is a sample markdown document.
+"""
+        file_path = Path("sample.md")
+        clean_content, metadata = extractor.extract_metadata(content, file_path)
+        assert "title" in metadata
+        assert metadata["title"] == "Sample Document"
+        assert "author" in metadata
+        assert metadata["author"] == "Test Author"
+        assert "# Heading 1" in clean_content
+        assert "This is a sample markdown document." in clean_content
+
+    def test_no_extractor_for_unsupported_file(self):
+        extractor = Extractor()
+        content = "<html><body>This is HTML content.</body></html>"
+        file_path = Path("sample.html")
+        clean_content, metadata = extractor.extract_metadata(content, file_path)
+        assert clean_content == content
+        assert metadata == {}
+
+    def test_get_extractor(self):
+        extractor = Extractor()
+        md_extractor = extractor.get_extractor(".md")
+        assert md_extractor is not None
+        assert md_extractor.supports_file_type(".md")
+
+        html_extractor = extractor.get_extractor(".html")
+        assert html_extractor is None