haiku.rag 0.4.3__tar.gz → 0.5.0__tar.gz

{haiku_rag-0.4.3 → haiku_rag-0.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haiku.rag
-Version: 0.4.3
+Version: 0.5.0
 Summary: Retrieval Augmented Generation (RAG) with SQLite
 Author-email: Yiorgis Gozadinos <ggozadinos@gmail.com>
 License: MIT

{haiku_rag-0.4.3 → haiku_rag-0.5.0}/docs/configuration.md

@@ -150,7 +150,4 @@ DEFAULT_DATA_DIR="/path/to/data"
 ```bash
 # Chunk size for document processing
 CHUNK_SIZE=256
-
-# Chunk overlap for better context
-CHUNK_OVERLAP=32
 ```

{haiku_rag-0.4.3 → haiku_rag-0.5.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "haiku.rag"
-version = "0.4.3"
+version = "0.5.0"
 description = "Retrieval Augmented Generation (RAG) with SQLite"
 authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
 license = { text = "MIT" }

haiku_rag-0.5.0/src/haiku/rag/chunker.py

@@ -0,0 +1,60 @@
+from io import BytesIO
+from typing import ClassVar
+
+import tiktoken
+from docling.chunking import HybridChunker  # type: ignore
+from docling.document_converter import DocumentConverter
+from docling_core.transforms.chunker.tokenizer.openai import OpenAITokenizer
+from docling_core.types.io import DocumentStream
+
+from haiku.rag.config import Config
+
+
+class Chunker:
+    """A class that chunks text into smaller pieces for embedding and retrieval.
+
+    Uses docling's structure-aware chunking to create semantically meaningful chunks
+    that respect document boundaries.
+
+    Args:
+        chunk_size: The maximum size of a chunk in tokens.
+    """
+
+    encoder: ClassVar[tiktoken.Encoding] = tiktoken.encoding_for_model("gpt-4o")
+
+    def __init__(
+        self,
+        chunk_size: int = Config.CHUNK_SIZE,
+    ):
+        self.chunk_size = chunk_size
+        tokenizer = OpenAITokenizer(
+            tokenizer=tiktoken.encoding_for_model("gpt-4o"), max_tokens=chunk_size
+        )
+
+        self.chunker = HybridChunker(tokenizer=tokenizer)  # type: ignore
+
+    async def chunk(self, text: str) -> list[str]:
+        """Split the text into chunks using docling's structure-aware chunking.
+
+        Args:
+            text: The text to be split into chunks.
+
+        Returns:
+            A list of text chunks with semantic boundaries.
+        """
+        if not text:
+            return []
+
+        # Convert to docling document
+        bytes_io = BytesIO(text.encode("utf-8"))
+        doc_stream = DocumentStream(name="text.md", stream=bytes_io)
+        converter = DocumentConverter()
+        result = converter.convert(doc_stream)
+        doc = result.document
+
+        # Chunk using docling's hybrid chunker
+        chunks = list(self.chunker.chunk(doc))
+        return [self.chunker.contextualize(chunk) for chunk in chunks]
+
+
+chunker = Chunker()

{haiku_rag-0.4.3 → haiku_rag-0.5.0}/src/haiku/rag/config.py

@@ -27,7 +27,6 @@ class AppConfig(BaseModel):
     QA_MODEL: str = "qwen3"
 
     CHUNK_SIZE: int = 256
-    CHUNK_OVERLAP: int = 32
 
     OLLAMA_BASE_URL: str = "http://localhost:11434"
 

{haiku_rag-0.4.3 → haiku_rag-0.5.0}/src/haiku/rag/reader.py

@@ -16,7 +16,8 @@ class FileReader:
         ".jpeg",
         ".jpg",
         ".md",
-        ".pdf.png",
+        ".pdf",
+        ".png",
         ".pptx",
         ".tiff",
         ".xlsx",

{haiku_rag-0.4.3 → haiku_rag-0.5.0}/src/haiku/rag/store/repositories/settings.py

@@ -63,7 +63,6 @@ class SettingsRepository:
             "EMBEDDINGS_MODEL",
             "EMBEDDINGS_VECTOR_DIM",
             "CHUNK_SIZE",
-            "CHUNK_OVERLAP",
         ]
 
         errors = []

haiku_rag-0.5.0/tests/test_chunker.py

@@ -0,0 +1,34 @@
+import pytest
+from datasets import Dataset
+
+from haiku.rag.chunker import Chunker
+
+
+@pytest.mark.asyncio
+async def test_chunker(qa_corpus: Dataset):
+    chunker = Chunker()
+    doc = qa_corpus[0]["document_extracted"]
+    chunks = await Chunker().chunk(doc)
+
+    # Ensure that the text is split into multiple chunks
+    assert len(chunks) > 1
+
+    # Ensure that chunks are reasonably sized (allowing more flexibility for structure-aware chunking)
+    total_tokens = 0
+    for chunk in chunks:
+        encoded_tokens = Chunker.encoder.encode(chunk, disallowed_special=())
+        token_count = len(encoded_tokens)
+        total_tokens += token_count
+
+        # Each chunk should be reasonably sized (allowing more flexibility than the old strict limits)
+        assert (
+            token_count <= chunker.chunk_size * 1.2
+        )  # Allow some flexibility for semantic boundaries
+        assert token_count > 5  # Ensure chunks aren't too small
+
+    # Ensure that all chunks together contain roughly the same content as original
+    original_tokens = len(Chunker.encoder.encode(doc, disallowed_special=()))
+
+    # Due to structure-aware chunking, we might have some variation in token count
+    # but it should be reasonable
+    assert abs(total_tokens - original_tokens) <= original_tokens * 0.1

haiku_rag-0.5.0/tests/test_reader.py

@@ -0,0 +1,22 @@
+import tempfile
+from pathlib import Path
+
+from haiku.rag.reader import FileReader
+
+
+def test_code_file_wrapped_in_code_block():
+    """Test that code files are wrapped in markdown code blocks."""
+    python_code = '''def hello_world():
+    print("Hello, World!")
+    return "success"'''
+
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".py") as f:
+        f.write(python_code)
+        f.flush()
+        temp_path = Path(f.name)
+
+        result = FileReader.parse_file(temp_path)
+
+        assert result.startswith("```python\n")
+        assert result.endswith("\n```")
+        assert "def hello_world():" in result

{haiku_rag-0.4.3 → haiku_rag-0.5.0}/tests/test_search.py

@@ -36,7 +36,7 @@ async def test_search_qa_corpus(qa_corpus: Dataset):
         created_document = await doc_repo.create(document)
         documents.append((created_document, doc_data))
 
-    for i in range(num_documents):  # Test with first few documents
+    for i in range(5):  # Test with first few documents
         target_document, doc_data = documents[i]
         question = doc_data["question"]
 
@@ -50,6 +50,10 @@ async def test_search_qa_corpus(qa_corpus: Dataset):
         target_document_ids = {chunk.document_id for chunk, _ in fts_results}
         assert target_document.id in target_document_ids
 
+    for i in range(num_documents):  # Test with first few documents
+        target_document, doc_data = documents[i]
+        question = doc_data["question"]
+
         # Test hybrid search
         hybrid_results = await chunk_repo.search_chunks_hybrid(question, limit=5)
         target_document_ids = {chunk.document_id for chunk, _ in hybrid_results}

{haiku_rag-0.4.3 → haiku_rag-0.5.0}/uv.lock

@@ -990,7 +990,7 @@ wheels = [
 
 [[package]]
 name = "haiku-rag"
-version = "0.4.3"
+version = "0.5.0"
 source = { editable = "." }
 dependencies = [
     { name = "docling" },

haiku_rag-0.4.3/src/haiku/rag/chunker.py

@@ -1,67 +0,0 @@
-from typing import ClassVar
-
-import tiktoken
-
-from haiku.rag.config import Config
-
-
-class Chunker:
-    """A class that chunks text into smaller pieces for embedding and retrieval.
-
-    Args:
-        chunk_size: The maximum size of a chunk in tokens.
-        chunk_overlap: The number of tokens of overlap between chunks.
-    """
-
-    encoder: ClassVar[tiktoken.Encoding] = tiktoken.encoding_for_model("gpt-4o")
-
-    def __init__(
-        self,
-        chunk_size: int = Config.CHUNK_SIZE,
-        chunk_overlap: int = Config.CHUNK_OVERLAP,
-    ):
-        self.chunk_size = chunk_size
-        self.chunk_overlap = chunk_overlap
-
-    async def chunk(self, text: str) -> list[str]:
-        """Split the text into chunks based on token boundaries.
-
-        Args:
-            text: The text to be split into chunks.
-
-        Returns:
-            A list of text chunks with token-based boundaries and overlap.
-        """
-        if not text:
-            return []
-
-        encoded_tokens = self.encoder.encode(text, disallowed_special=())
-
-        if self.chunk_size > len(encoded_tokens):
-            return [text]
-
-        chunks = []
-        i = 0
-        split_id_counter = 0
-        while i < len(encoded_tokens):
-            # Overlap
-            start_i = i
-            end_i = min(i + self.chunk_size, len(encoded_tokens))
-
-            chunk_tokens = encoded_tokens[start_i:end_i]
-            chunk_text = self.encoder.decode(chunk_tokens)
-
-            chunks.append(chunk_text)
-            split_id_counter += 1
-
-            # Exit loop if this was the last possible chunk
-            if end_i == len(encoded_tokens):
-                break
-
-            i += (
-                self.chunk_size - self.chunk_overlap
-            )  # Step forward, considering overlap
-        return chunks
-
-
-chunker = Chunker()

haiku_rag-0.4.3/tests/test_chunker.py

@@ -1,44 +0,0 @@
-import pytest
-from datasets import Dataset
-
-from haiku.rag.chunker import Chunker
-
-
-@pytest.mark.asyncio
-async def test_chunker(qa_corpus: Dataset):
-    chunker = Chunker()
-    doc = qa_corpus[0]["document_extracted"]
-    chunks = await Chunker().chunk(doc)
-
-    # Ensure that the text is split into multiple chunks
-    assert len(chunks) > 1
-
-    # Ensure that each chunk corresponds to roughly Config.CHUNK_SIZE tokens
-    for chunk in chunks[:-1]:
-        encoded_tokens = Chunker.encoder.encode(chunk, disallowed_special=())
-        assert len(encoded_tokens) <= Chunker().chunk_size
-        assert len(encoded_tokens) > Chunker().chunk_size * 0.9
-
-    # Ensure that the last chunk is less than Config.CHUNK_SIZE tokens
-    assert (
-        len(Chunker.encoder.encode(chunks[-1], disallowed_special=()))
-        < Chunker().chunk_size
-    )
-
-    # Test overlap between consecutive chunks
-    for i in range(len(chunks) - 1):
-        current_chunk = chunks[i]
-        next_chunk = chunks[i + 1]
-
-        current_tokens = Chunker.encoder.encode(current_chunk, disallowed_special=())
-        next_tokens = Chunker.encoder.encode(next_chunk, disallowed_special=())
-
-        overlap_size = min(chunker.chunk_overlap, len(current_tokens))
-        current_overlap_tokens = current_tokens[-overlap_size:]
-        next_overlap_tokens = next_tokens[:overlap_size]
-
-        # The overlapping tokens should be identical
-        assert current_overlap_tokens == next_overlap_tokens
-        assert len(current_overlap_tokens) == min(
-            chunker.chunk_overlap, len(current_tokens)
-        )