diff --git a/.gitignore b/.gitignore index 952f8b27..1b2e60fe 100644 --- a/.gitignore +++ b/.gitignore @@ -5,6 +5,7 @@ build/ dist/ wheels/ *.egg-info +**/.DS_Store # Virtual environments .venv diff --git a/haiku_rag_slim/haiku/rag/chunkers/base.py b/haiku_rag_slim/haiku/rag/chunkers/base.py index 01df87f6..783a2cb7 100644 --- a/haiku_rag_slim/haiku/rag/chunkers/base.py +++ b/haiku_rag_slim/haiku/rag/chunkers/base.py @@ -1,10 +1,20 @@ from abc import ABC, abstractmethod from typing import TYPE_CHECKING +from haiku.rag.store.models.chunk import ChunkMetadata + if TYPE_CHECKING: from docling_core.types.doc.document import DoclingDocument +class ChunkWithMetadata: + """A chunk with its text content and structured metadata.""" + + def __init__(self, text: str, metadata: ChunkMetadata): + self.text = text + self.metadata = metadata + + class DocumentChunker(ABC): """Abstract base class for document chunkers. @@ -13,14 +23,15 @@ class DocumentChunker(ABC): """ @abstractmethod - async def chunk(self, document: "DoclingDocument") -> list[str]: - """Split a document into chunks. + async def chunk(self, document: "DoclingDocument") -> list[ChunkWithMetadata]: + """Split a document into chunks with metadata. Args: document: The DoclingDocument to chunk. Returns: - List of text chunks with semantic boundaries preserved. + List of ChunkWithMetadata containing text and structured metadata + (doc_item_refs, headings, labels, page_numbers). Raises: ValueError: If chunking fails. diff --git a/haiku_rag_slim/haiku/rag/chunkers/docling_local.py b/haiku_rag_slim/haiku/rag/chunkers/docling_local.py index 7fa6d39b..31d05d74 100644 --- a/haiku_rag_slim/haiku/rag/chunkers/docling_local.py +++ b/haiku_rag_slim/haiku/rag/chunkers/docling_local.py @@ -1,9 +1,11 @@ -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast -from haiku.rag.chunkers.base import DocumentChunker +from haiku.rag.chunkers.base import ChunkWithMetadata, DocumentChunker from haiku.rag.config import AppConfig, Config +from haiku.rag.store.models.chunk import ChunkMetadata if TYPE_CHECKING: + from docling_core.transforms.chunker.doc_chunk import DocMeta from docling_core.types.doc.document import DoclingDocument @@ -93,18 +95,64 @@ class DoclingLocalChunker(DocumentChunker): "Must be 'hybrid' or 'hierarchical'." ) - async def chunk(self, document: "DoclingDocument") -> list[str]: - """Split the document into chunks using docling's structure-aware chunking. + async def chunk(self, document: "DoclingDocument") -> list[ChunkWithMetadata]: + """Split the document into chunks with metadata. + + Extracts structured metadata from each DocChunk including: + - doc_item_refs: JSON pointer references to DocItems (e.g., "#/texts/5") + - headings: Section heading hierarchy + - labels: Semantic labels for each doc_item (e.g., "paragraph", "table") + - page_numbers: Page numbers where content appears Args: document: The DoclingDocument to be split into chunks. Returns: - A list of text chunks with semantic boundaries. + List of ChunkWithMetadata containing text and structured metadata. """ if document is None: return [] - # Chunk using docling's hybrid chunker - chunks = list(self.chunker.chunk(document)) - return [self.chunker.contextualize(chunk) for chunk in chunks] + raw_chunks = list(self.chunker.chunk(document)) + result: list[ChunkWithMetadata] = [] + + for chunk in raw_chunks: + text = self.chunker.contextualize(chunk) + + # Extract metadata from DocChunk.meta (cast to DocMeta for type safety) + doc_item_refs: list[str] = [] + labels: list[str] = [] + page_numbers: list[int] = [] + headings: list[str] | None = None + + meta = cast("DocMeta | None", chunk.meta) + if meta and meta.doc_items: + for doc_item in meta.doc_items: + # Get JSON pointer reference + if doc_item.self_ref: + doc_item_refs.append(doc_item.self_ref) + # Get label + if doc_item.label: + labels.append(doc_item.label) + # Get page numbers from provenance + if doc_item.prov: + for prov in doc_item.prov: + if ( + prov.page_no is not None + and prov.page_no not in page_numbers + ): + page_numbers.append(prov.page_no) + + # Get headings from chunk metadata + if meta and meta.headings: + headings = list(meta.headings) + + metadata = ChunkMetadata( + doc_item_refs=doc_item_refs, + headings=headings, + labels=labels, + page_numbers=sorted(page_numbers), + ) + result.append(ChunkWithMetadata(text=text, metadata=metadata)) + + return result diff --git a/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py b/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py index d1bba864..4c4bdf9d 100644 --- a/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py +++ b/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py @@ -1,14 +1,48 @@ +import re from io import BytesIO from typing import TYPE_CHECKING import httpx -from haiku.rag.chunkers.base import DocumentChunker +from haiku.rag.chunkers.base import ChunkWithMetadata, DocumentChunker from haiku.rag.config import AppConfig, Config +from haiku.rag.store.models.chunk import ChunkMetadata if TYPE_CHECKING: from docling_core.types.doc.document import DoclingDocument +# Pattern to parse refs like "#/texts/5" or "#/tables/0" +REF_PATTERN = re.compile(r"^#/(\w+)/(\d+)$") + + +def _resolve_label_from_document(ref: str, document: "DoclingDocument") -> str | None: + """Resolve the label for a doc_item ref by looking it up in the document. + + The docling-serve API only returns ref strings in doc_items, not labels. + This function resolves actual labels from the DoclingDocument. + See: https://github.com/docling-project/docling-serve/issues/448 + + Args: + ref: JSON pointer reference like "#/texts/5" or "#/tables/0" + document: The DoclingDocument to look up the item in + + Returns: + The label string if found, None otherwise + """ + match = REF_PATTERN.match(ref) + if not match: + return None + + collection_name = match.group(1) + index = int(match.group(2)) + + collection = getattr(document, collection_name, None) + if collection is None or index >= len(collection): + return None + + item = collection[index] + return getattr(item, "label", None) + class DoclingServeChunker(DocumentChunker): """Remote document chunker using docling-serve API. @@ -27,54 +61,48 @@ class DoclingServeChunker(DocumentChunker): self.timeout = config.providers.docling_serve.timeout self.chunker_type = config.processing.chunker_type - async def chunk(self, document: "DoclingDocument") -> list[str]: - """Split the document into chunks via docling-serve. - - Exports the DoclingDocument to JSON and sends it to docling-serve's chunking - endpoint. The API will chunk the document and return the text chunks. + async def _call_chunk_api(self, document: "DoclingDocument") -> list[dict]: + """Call docling-serve chunking API and return raw chunk data. Args: document: The DoclingDocument to be split into chunks. Returns: - A list of text chunks with semantic boundaries. + List of chunk dictionaries from API response. Raises: ValueError: If chunking fails or service is unavailable. """ - if document is None: - return [] + # Determine endpoint based on chunker_type + if self.chunker_type == "hierarchical": + url = f"{self.base_url}/v1/chunk/hierarchical/file" + else: + url = f"{self.base_url}/v1/chunk/hybrid/file" + + # Export document to JSON + doc_json = document.model_dump_json() + doc_bytes = doc_json.encode("utf-8") + + # Prepare multipart request with DoclingDocument JSON + files = {"files": ("document.json", BytesIO(doc_bytes), "application/json")} + + # Build form data with chunking parameters + data = { + "chunking_max_tokens": str(self.config.processing.chunk_size), + "chunking_tokenizer": self.config.processing.chunking_tokenizer, + "chunking_merge_peers": str( + self.config.processing.chunking_merge_peers + ).lower(), + "chunking_use_markdown_tables": str( + self.config.processing.chunking_use_markdown_tables + ).lower(), + } + + headers = {} + if self.api_key: + headers["X-Api-Key"] = self.api_key try: - # Determine endpoint based on chunker_type - if self.chunker_type == "hierarchical": - url = f"{self.base_url}/v1/chunk/hierarchical/file" - else: - url = f"{self.base_url}/v1/chunk/hybrid/file" - - # Export document to JSON - doc_json = document.model_dump_json() - doc_bytes = doc_json.encode("utf-8") - - # Prepare multipart request with DoclingDocument JSON - files = {"files": ("document.json", BytesIO(doc_bytes), "application/json")} - - # Build form data with chunking parameters - data = { - "chunking_max_tokens": str(self.config.processing.chunk_size), - "chunking_tokenizer": self.config.processing.chunking_tokenizer, - "chunking_merge_peers": str( - self.config.processing.chunking_merge_peers - ).lower(), - "chunking_use_markdown_tables": str( - self.config.processing.chunking_use_markdown_tables - ).lower(), - } - - headers = {} - if self.api_key: - headers["X-Api-Key"] = self.api_key - async with httpx.AsyncClient(timeout=self.timeout) as client: response = await client.post( url, @@ -84,10 +112,7 @@ class DoclingServeChunker(DocumentChunker): ) response.raise_for_status() result = response.json() - - # Extract text from chunks - chunks = result.get("chunks", []) - return [chunk["text"] for chunk in chunks] + return result.get("chunks", []) except httpx.ConnectError as e: raise ValueError( @@ -107,3 +132,66 @@ class DoclingServeChunker(DocumentChunker): raise ValueError(f"HTTP error from docling-serve: {e}") except Exception as e: raise ValueError(f"Failed to chunk via docling-serve: {e}") + + async def chunk(self, document: "DoclingDocument") -> list[ChunkWithMetadata]: + """Split the document into chunks with metadata via docling-serve. + + Extracts structured metadata from the API response including: + - doc_item_refs: JSON pointer references to DocItems (e.g., "#/texts/5") + - headings: Section heading hierarchy + - labels: Semantic labels for each doc_item + - page_numbers: Page numbers where content appears + + Args: + document: The DoclingDocument to be split into chunks. + + Returns: + List of ChunkWithMetadata containing text and structured metadata. + + Raises: + ValueError: If chunking fails or service is unavailable. + """ + if document is None: + return [] + + raw_chunks = await self._call_chunk_api(document) + result: list[ChunkWithMetadata] = [] + + for chunk in raw_chunks: + text = chunk.get("text", "") + + # doc_items from docling-serve is a list of ref strings like ["#/texts/1", "#/tables/0"] + doc_items = chunk.get("doc_items", []) + doc_item_refs: list[str] = [] + labels: list[str] = [] + + for item in doc_items: + if isinstance(item, str): + # docling-serve returns refs as strings directly + doc_item_refs.append(item) + # Resolve label from the document using the ref + label = _resolve_label_from_document(item, document) + if label: + labels.append(label) + elif isinstance(item, dict): + # Handle dict format if API ever returns it + if "self_ref" in item: + doc_item_refs.append(item["self_ref"]) + if "label" in item: + labels.append(item["label"]) + + # Get headings directly from chunk + headings = chunk.get("headings") + + # Get page numbers directly from chunk + page_numbers = chunk.get("page_numbers", []) + + metadata = ChunkMetadata( + doc_item_refs=doc_item_refs, + headings=headings, + labels=labels, + page_numbers=sorted(page_numbers) if page_numbers else [], + ) + result.append(ChunkWithMetadata(text=text, metadata=metadata)) + + return result diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index 33623dd5..11055b5a 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py @@ -225,22 +225,28 @@ class ChunkRepository: ) raise e - chunk_texts = await chunker.chunk(processed_document) + chunks_with_metadata = await chunker.chunk(processed_document) + chunk_texts = [c.text for c in chunks_with_metadata] embeddings = await self.embedder.embed(chunk_texts) # Prepare all chunk records for batch insertion chunk_records = [] created_chunks = [] - for order, (chunk_text, embedding) in enumerate(zip(chunk_texts, embeddings)): + for order, (chunk_with_meta, embedding) in enumerate( + zip(chunks_with_metadata, embeddings) + ): chunk_id = str(uuid4()) + # Convert ChunkMetadata to dict for storage + metadata_dict = chunk_with_meta.metadata.model_dump() + chunk_record = self.store.ChunkRecord( id=chunk_id, document_id=document_id, - content=chunk_text, - metadata=json.dumps({}), + content=chunk_with_meta.text, + metadata=json.dumps(metadata_dict), order=order, vector=embedding, ) @@ -249,8 +255,8 @@ class ChunkRepository: chunk = Chunk( id=chunk_id, document_id=document_id, - content=chunk_text, - metadata={}, + content=chunk_with_meta.text, + metadata=metadata_dict, order=order, ) created_chunks.append(chunk) diff --git a/tests/data/doclaynet.pdf b/tests/data/doclaynet.pdf new file mode 100644 index 00000000..0fa7d1f8 Binary files /dev/null and b/tests/data/doclaynet.pdf differ diff --git a/tests/test_chunker.py b/tests/test_chunker.py index 1f956e3f..f75b0a5a 100644 --- a/tests/test_chunker.py +++ b/tests/test_chunker.py @@ -32,7 +32,7 @@ async def test_local_chunker(qa_corpus: Dataset): # Ensure that chunks are reasonably sized (allowing more flexibility for structure-aware chunking) total_tokens = 0 for chunk in chunks: - encoded_tokens = tokenizer.encode(chunk, add_special_tokens=False) + encoded_tokens = tokenizer.encode(chunk.text, add_special_tokens=False) token_count = len(encoded_tokens) total_tokens += token_count @@ -95,7 +95,7 @@ async def test_local_chunker_hierarchical(qa_corpus: Dataset): assert len(chunks) > 0 # Each chunk should be non-empty for chunk in chunks: - assert len(chunk.strip()) > 0 + assert len(chunk.text.strip()) > 0 def test_local_chunker_invalid_type(): @@ -127,8 +127,8 @@ async def test_local_chunker_markdown_tables(): chunks_md = await chunker_md.chunk(doc) # Should contain markdown table format - assert any("|" in chunk for chunk in chunks_md) - assert any("Column 1" in chunk for chunk in chunks_md) + assert any("|" in chunk.text for chunk in chunks_md) + assert any("Column 1" in chunk.text for chunk in chunks_md) # Test with markdown tables disabled (narrative format) config_narrative = AppConfig() @@ -137,11 +137,59 @@ async def test_local_chunker_markdown_tables(): chunks_narrative = await chunker_narrative.chunk(doc) # Should contain narrative format (no pipe characters in table) - table_content = [chunk for chunk in chunks_narrative if "Value" in chunk][0] + table_content = [chunk.text for chunk in chunks_narrative if "Value" in chunk.text][ + 0 + ] # Narrative format uses commas, not pipes for table structure assert "," in table_content and "|" not in table_content +@pytest.mark.asyncio +async def test_local_chunker_metadata_extraction(): + """Test that DoclingLocalChunker extracts metadata correctly.""" + sample_md = """# Chapter 1: Introduction + +This is the first paragraph of the introduction. + +## Section 1.1: Background + +Here is some background information. + +| Header 1 | Header 2 | +|----------|----------| +| Value 1 | Value 2 | +""" + converter = get_converter(Config) + doc = await converter.convert_text(sample_md, name="test.md") + + chunker = DoclingLocalChunker() + chunks = await chunker.chunk(doc) + + assert len(chunks) > 0 + + # Check that at least one chunk has doc_item_refs + all_refs = [] + all_labels = [] + all_headings = [] + for chunk in chunks: + all_refs.extend(chunk.metadata.doc_item_refs) + all_labels.extend(chunk.metadata.labels) + if chunk.metadata.headings: + all_headings.extend(chunk.metadata.headings) + + # Should have JSON pointer refs like #/texts/0, #/tables/0 + assert len(all_refs) > 0 + assert any(ref.startswith("#/") for ref in all_refs) + + # Should have labels + assert len(all_labels) > 0 + assert "text" in all_labels or "table" in all_labels + + # Should have headings + assert len(all_headings) > 0 + assert any("Chapter" in h or "Section" in h for h in all_headings) + + def test_get_chunker_docling_serve(): """Test factory returns DoclingServeChunker for docling-serve.""" config = AppConfig() @@ -193,8 +241,8 @@ class TestDoclingServeChunker: chunks = await chunker.chunk(doc) assert len(chunks) == 2 - assert chunks[0] == "Chunk 1" - assert chunks[1] == "Chunk 2" + assert chunks[0].text == "Chunk 1" + assert chunks[1].text == "Chunk 2" mock_client.post.assert_called_once() @pytest.mark.asyncio @@ -332,3 +380,156 @@ class TestDoclingServeChunker: with pytest.raises(ValueError, match="Authentication failed"): await chunker.chunk(doc) + + @pytest.mark.asyncio + @patch("haiku.rag.chunkers.docling_serve.httpx.AsyncClient") + async def test_chunk_metadata_extraction(self, mock_client_class, chunker): + """Test that metadata is correctly extracted from API response. + + Labels are resolved from the DoclingDocument using the refs, so we need + to create a document with matching structure for the mocked API response. + """ + mock_response = Mock() + mock_response.status_code = 200 + # docling-serve returns doc_items as list of ref strings + # We'll reference texts[0], texts[1], and tables[0] + mock_response.json.return_value = { + "chunks": [ + { + "text": "Chapter 1\nThis is content.", + "doc_items": ["#/texts/0", "#/texts/1"], + "headings": ["Chapter 1"], + "page_numbers": [1], + }, + { + "text": "Table content here.", + "doc_items": ["#/tables/0"], + "headings": ["Chapter 1", "Section 1.1"], + "page_numbers": [1, 2], + }, + ] + } + mock_response.raise_for_status = Mock() + + mock_client = AsyncMock() + mock_client.post.return_value = mock_response + mock_client_class.return_value.__aenter__.return_value = mock_client + + # Create a document with texts and tables that match the mocked refs + converter = get_converter(Config) + doc = await converter.convert_text( + """# Chapter 1 + +This is content. + +| Col1 | Col2 | +|------|------| +| A | B | +""", + name="test.md", + ) + + chunks = await chunker.chunk(doc) + + assert len(chunks) == 2 + + # First chunk - labels resolved from document + assert chunks[0].text == "Chapter 1\nThis is content." + assert chunks[0].metadata.doc_item_refs == ["#/texts/0", "#/texts/1"] + # texts[0] is title (# heading), texts[1] is text (paragraph) + assert chunks[0].metadata.labels == ["title", "text"] + assert chunks[0].metadata.headings == ["Chapter 1"] + assert chunks[0].metadata.page_numbers == [1] + + # Second chunk - label resolved from document + assert chunks[1].text == "Table content here." + assert chunks[1].metadata.doc_item_refs == ["#/tables/0"] + assert chunks[1].metadata.labels == ["table"] + assert chunks[1].metadata.headings == ["Chapter 1", "Section 1.1"] + assert chunks[1].metadata.page_numbers == [1, 2] + + +def is_docling_serve_available(base_url: str = "http://localhost:5001") -> bool: + """Check if docling-serve is running and accessible.""" + import requests + + try: + response = requests.get(f"{base_url}/health", timeout=2) + return response.status_code == 200 + except Exception: + return False + + +@pytest.mark.skipif( + not is_docling_serve_available(), + reason="docling-serve not available at http://localhost:5001", +) +@pytest.mark.integration +@pytest.mark.asyncio +async def test_local_and_serve_chunkers_produce_same_output(): + """Test that local and serve chunkers produce identical output for the same document. + + This integration test requires docling-serve to be running on localhost:5001. + + Note: Labels are resolved from the DoclingDocument since docling-serve API + only returns ref strings, not labels. See: + https://github.com/docling-project/docling-serve/issues/448 + """ + from pathlib import Path + + from haiku.rag.chunkers.docling_local import DoclingLocalChunker + from haiku.rag.chunkers.docling_serve import DoclingServeChunker + from haiku.rag.converters.docling_serve import DoclingServeConverter + + # Use docling-serve to convert the PDF (ensures same conversion for both chunkers) + converter = DoclingServeConverter(Config) + pdf_path = Path("tests/data/doclaynet.pdf") + doc = await converter.convert_file(pdf_path) + + # Create both chunkers with same config + config = AppConfig() + config.processing.chunk_size = 256 + config.processing.chunker_type = "hybrid" + config.processing.chunking_merge_peers = True + config.processing.chunking_use_markdown_tables = True + + local_chunker = DoclingLocalChunker(config) + serve_chunker = DoclingServeChunker(config) + + # Chunk with both + local_chunks = await local_chunker.chunk(doc) + serve_chunks = await serve_chunker.chunk(doc) + + # Same number of chunks + assert len(local_chunks) == len(serve_chunks), ( + f"Chunk count mismatch: local={len(local_chunks)}, serve={len(serve_chunks)}" + ) + + # Compare each chunk + for i, (local, serve) in enumerate(zip(local_chunks, serve_chunks)): + # Text should match + assert local.text == serve.text, f"Chunk {i} text mismatch" + + # doc_item_refs should match + assert local.metadata.doc_item_refs == serve.metadata.doc_item_refs, ( + f"Chunk {i} doc_item_refs mismatch: " + f"local={local.metadata.doc_item_refs}, serve={serve.metadata.doc_item_refs}" + ) + + # Labels should match (now that serve resolves from document) + assert local.metadata.labels == serve.metadata.labels, ( + f"Chunk {i} labels mismatch: " + f"local={local.metadata.labels}, serve={serve.metadata.labels}" + ) + + # Headings should match + assert local.metadata.headings == serve.metadata.headings, ( + f"Chunk {i} headings mismatch: " + f"local={local.metadata.headings}, serve={serve.metadata.headings}" + ) + + # Page numbers should match + assert local.metadata.page_numbers == serve.metadata.page_numbers, ( + f"Chunk {i} page_numbers mismatch: " + f"local={local.metadata.page_numbers}, serve={serve.metadata.page_numbers}" + )