Merge pull request #56 from ggozad/feat/pre-process

Allow for custom markdown pre-processing before chunking/embedding.

docs/configuration.md

@@ -223,3 +223,35 @@ CHUNK_SIZE=256
 # into single chunks with continuous content to eliminate duplication
 CONTEXT_CHUNK_RADIUS=0
 ```
+
+#### Markdown Preprocessor
+
+Optionally preprocess Markdown before chunking by pointing to a callable that receives and returns Markdown text. This is useful for normalizing content, stripping boilerplate, or applying custom transformations before chunk boundaries are computed.
+
+```bash
+# A callable path in one of these formats:
+# - package.module:func
+# - package.module.func
+# - /abs/or/relative/path/to/file.py:func
+MARKDOWN_PREPROCESSOR="my_pkg.preprocess:clean_md"
+```
+
+!!! note
+    - The function signature should be `def clean_md(text: str) -> str` or `async def clean_md(text: str) -> str`.
+    - If the function raises or returns a non-string, haiku.rag logs a warning and proceeds without preprocessing.
+    - The preprocessor affects only the chunking pipeline. The stored document content remains unchanged.
+
+Example implementation:
+
+```python
+# my_pkg/preprocess.py
+def clean_md(text: str) -> str:
+    # strip HTML comments and collapse multiple blank lines
+    lines = [line for line in text.splitlines() if not line.strip().startswith("<!--")]
+    out = []
+    for line in lines:
+        if line.strip() == "" and (out and out[-1] == ""):
+            continue
+        out.append(line)
+    return "\n".join(out)
+```

src/haiku/rag/config.py

@@ -32,6 +32,10 @@ class AppConfig(BaseModel):
     CHUNK_SIZE: int = 256
     CONTEXT_CHUNK_RADIUS: int = 0
 
+    # Optional dotted path or file path to a callable that preprocesses
+    # markdown content before chunking. Examples:
+    MARKDOWN_PREPROCESSOR: str = ""
+
     OLLAMA_BASE_URL: str = "http://localhost:11434"
     VLLM_EMBEDDINGS_BASE_URL: str = ""
     VLLM_RERANK_BASE_URL: str = ""

src/haiku/rag/store/repositories/chunk.py

@@ -1,4 +1,5 @@
 import asyncio
+import inspect
 import json
 import logging
 from uuid import uuid4
@@ -11,6 +12,7 @@ from haiku.rag.config import Config
 from haiku.rag.embeddings import get_embedder
 from haiku.rag.store.engine import DocumentRecord, Store
 from haiku.rag.store.models.chunk import Chunk
+from haiku.rag.utils import load_callable, text_to_docling_document
 
 logger = logging.getLogger(__name__)
 
@@ -152,7 +154,28 @@ class ChunkRepository:
         self, document_id: str, document: DoclingDocument
     ) -> list[Chunk]:
         """Create chunks and embeddings for a document from DoclingDocument."""
-        chunk_texts = await chunker.chunk(document)
+        # Optionally preprocess markdown before chunking
+        processed_document = document
+        preprocessor_path = Config.MARKDOWN_PREPROCESSOR
+        if preprocessor_path:
+            try:
+                pre_fn = load_callable(preprocessor_path)
+                markdown = document.export_to_markdown()
+                result = pre_fn(markdown)
+                if inspect.isawaitable(result):
+                    result = await result  # type: ignore[assignment]
+                processed_markdown = result
+                if not isinstance(processed_markdown, str):
+                    raise ValueError("Preprocessor must return a markdown string")
+                processed_document = text_to_docling_document(
+                    processed_markdown, name="content.md"
+                )
+            except Exception as e:
+                logger.warning(
+                    f"Failed to apply MARKDOWN_PREPROCESSOR '{preprocessor_path}': {e}. Proceeding without preprocessing."
+                )
+
+        chunk_texts = await chunker.chunk(processed_document)
 
         embeddings = await self.embedder.embed(chunk_texts)
 

src/haiku/rag/utils.py

@@ -1,10 +1,13 @@
 import asyncio
+import importlib
+import importlib.util
 import sys
 from collections.abc import Callable
 from functools import wraps
 from importlib import metadata
 from io import BytesIO
 from pathlib import Path
+from types import ModuleType
 
 import httpx
 from docling.document_converter import DocumentConverter
@@ -106,3 +109,54 @@ def text_to_docling_document(text: str, name: str = "content.md") -> DoclingDocu
     converter = DocumentConverter()
     result = converter.convert(doc_stream)
     return result.document
+
+
+def load_callable(path: str):
+    """Load a callable from a dotted path or file path.
+
+    Supported formats:
+    - "package.module:func" or "package.module.func"
+    - "path/to/file.py:func"
+
+    Returns the loaded callable. Raises ValueError on failure.
+    """
+    if not path:
+        raise ValueError("Empty callable path provided")
+
+    module_part = None
+    func_name = None
+
+    if ":" in path:
+        module_part, func_name = path.split(":", 1)
+    else:
+        # split by last dot for module.attr
+        if "." in path:
+            module_part, func_name = path.rsplit(".", 1)
+        else:
+            raise ValueError(
+                "Invalid callable path format. Use 'module:func' or 'module.func' or 'file.py:func'."
+            )
+
+    # Try file path first
+    mod: ModuleType | None = None
+    module_path = Path(module_part)
+    if module_path.suffix == ".py" and module_path.exists():
+        spec = importlib.util.spec_from_file_location(module_path.stem, module_path)
+        if spec and spec.loader:
+            mod = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(mod)
+    else:
+        # Import as a module path
+        try:
+            mod = importlib.import_module(module_part)
+        except Exception as e:
+            raise ValueError(f"Failed to import module '{module_part}': {e}")
+
+    if not hasattr(mod, func_name):
+        raise ValueError(f"Callable '{func_name}' not found in module '{module_part}'")
+    func = getattr(mod, func_name)
+    if not callable(func):
+        raise ValueError(
+            f"Attribute '{func_name}' in module '{module_part}' is not callable"
+        )
+    return func

tests/test_preprocessor.py

@@ -0,0 +1,71 @@
+from pathlib import Path
+
+import pytest
+
+from haiku.rag.config import Config
+from haiku.rag.store.engine import Store
+from haiku.rag.store.models.document import Document
+from haiku.rag.store.repositories.chunk import ChunkRepository
+from haiku.rag.store.repositories.document import DocumentRepository
+from haiku.rag.utils import text_to_docling_document
+
+
+@pytest.mark.parametrize(
+    "is_async, marker",
+    [
+        (False, "MARKER_LINE"),
+        (True, "ASYNC_MARKER"),
+    ],
+)
+@pytest.mark.asyncio
+async def test_markdown_preprocessor_applied_parametrized(
+    is_async: bool, marker: str, tmp_path: Path, temp_db_path: Path
+):
+    """Ensure MARKDOWN_PREPROCESSOR (sync or async) transforms markdown before chunking."""
+    pre_file = tmp_path / ("pre_async.py" if is_async else "pre.py")
+    if is_async:
+        pre_file.write_text(
+            """
+import asyncio
+
+async def add_marker(text: str) -> str:
+    await asyncio.sleep(0)
+    return text + "\\n\\nASYNC_MARKER\\n"
+"""
+        )
+    else:
+        pre_file.write_text(
+            """
+def add_marker(text: str) -> str:
+    return text + "\\n\\nMARKER_LINE\\n"
+"""
+        )
+
+    original_pre = Config.MARKDOWN_PREPROCESSOR
+    try:
+        Config.MARKDOWN_PREPROCESSOR = f"{pre_file}:add_marker"
+
+        store = Store(temp_db_path)
+        chunk_repo = ChunkRepository(store)
+        doc_repo = DocumentRepository(store)
+
+        document = Document(content="Hello world")
+        created_doc = await doc_repo.create(document)
+        assert created_doc.id is not None
+
+        # Stub embeddings to avoid network
+        dim = chunk_repo.embedder._vector_dim
+
+        async def fake_embed(x):  # type: ignore[override]
+            if isinstance(x, list):
+                return [[0.0] * dim for _ in x]
+            return [0.0] * dim
+
+        chunk_repo.embedder.embed = fake_embed  # type: ignore[assignment]
+
+        docling = text_to_docling_document(document.content, name="test.md")
+        chunks = await chunk_repo.create_chunks_for_document(created_doc.id, docling)
+
+        assert any(marker in c.content for c in chunks)
+    finally:
+        Config.MARKDOWN_PREPROCESSOR = original_pre