danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Operate on DoclingDocument instead of Document

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-08-08 17:24:34 +02:00
parent 6a3c85e6ae
commit ad4b655a31
No known key found for this signature in database
7 changed files with 107 additions and 57 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
src/haiku/rag/chunker.py
View file

@ -1,11 +1,9 @@
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 docling_core.types.doc.document import DoclingDocument
from haiku.rag.config import Config
@ -33,27 +31,20 @@ class Chunker:
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.
async def chunk(self, document: DoclingDocument) -> list[str]:
"""Split the document into chunks using docling's structure-aware chunking.
Args:
text: The text to be split into chunks.
document: The DoclingDocument to be split into chunks.
Returns:
A list of text chunks with semantic boundaries.
"""
if not text:
if document is None:
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))
chunks = list(self.chunker.chunk(document))
return [self.chunker.contextualize(chunk) for chunk in chunks]

66
src/haiku/rag/client.py
View file

@ -16,6 +16,7 @@ from haiku.rag.store.models.chunk import Chunk
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
class HaikuRAG:
@ -49,6 +50,24 @@ class HaikuRAG:
self.close()
return False
async def _create_document_with_docling(
self,
docling_document,
uri: str | None = None,
metadata: dict | None = None,
chunks: list[Chunk] | None = None,
) -> Document:
"""Create a new document from DoclingDocument."""
content = docling_document.export_to_markdown()
document = Document(
content=content,
uri=uri,
metadata=metadata or {},
)
return await self.document_repository._create_with_docling(
document, docling_document, chunks
)
async def create_document(
self,
content: str,
@ -67,12 +86,17 @@ class HaikuRAG:
Returns:
The created Document instance.
"""
# Convert content to DoclingDocument for processing
docling_document = text_to_docling_document(content)
document = Document(
content=content,
uri=uri,
metadata=metadata or {},
)
return await self.document_repository.create(document, chunks)
return await self.document_repository._create_with_docling(
document, docling_document, chunks
)
async def create_document_from_source(
self, source: str | Path, metadata: dict = {}
@ -119,8 +143,7 @@ class HaikuRAG:
# MD5 unchanged, return existing document
return existing_doc
document = FileReader.parse_file(source_path)
content = document.export_to_markdown()
docling_document = FileReader.parse_file(source_path)
# Get content type from file extension
content_type, _ = mimetypes.guess_type(str(source_path))
@ -132,13 +155,15 @@ class HaikuRAG:
if existing_doc:
# Update existing document
existing_doc.content = content
existing_doc.content = docling_document.export_to_markdown()
existing_doc.metadata = metadata
return await self.update_document(existing_doc)
return await self.document_repository._update_with_docling(
existing_doc, docling_document
)
else:
# Create new document
return await self.create_document(
content=content, uri=uri, metadata=metadata
# Create new document using DoclingDocument
return await self._create_document_with_docling(
docling_document=docling_document, uri=uri, metadata=metadata
)
async def _create_or_update_document_from_url(
@ -194,19 +219,20 @@ class HaikuRAG:
temp_path = Path(temp_file.name)
# Parse the content using FileReader
document = FileReader.parse_file(temp_path)
content = document.export_to_markdown()
docling_document = FileReader.parse_file(temp_path)
# Merge metadata with contentType and md5
metadata.update({"contentType": content_type, "md5": md5_hash})
if existing_doc:
existing_doc.content = content
existing_doc.content = docling_document.export_to_markdown()
existing_doc.metadata = metadata
return await self.update_document(existing_doc)
return await self.document_repository._update_with_docling(
existing_doc, docling_document
)
else:
return await self.create_document(
content=content, uri=url, metadata=metadata
return await self._create_document_with_docling(
docling_document=docling_document, uri=url, metadata=metadata
)
def _get_extension_from_content_type_or_url(
@ -264,7 +290,12 @@ class HaikuRAG:
async def update_document(self, document: Document) -> Document:
"""Update an existing document."""
return await self.document_repository.update(document)
# Convert content to DoclingDocument
docling_document = text_to_docling_document(document.content)
return await self.document_repository._update_with_docling(
document, docling_document
)
async def delete_document(self, document_id: int) -> bool:
"""Delete a document by its ID."""
@ -346,8 +377,11 @@ class HaikuRAG:
for doc in documents:
if doc.id is not None:
# Convert content to DoclingDocument for rebuild
docling_document = text_to_docling_document(doc.content)
await self.chunk_repository.create_chunks_for_document(
doc.id, doc.content, commit=False
doc.id, docling_document, commit=False
)
yield doc.id

16
src/haiku/rag/reader.py
View file

@ -1,10 +1,10 @@
from io import BytesIO
from pathlib import Path
from typing import ClassVar
from docling.document_converter import DocumentConverter
from docling_core.types.doc.document import DoclingDocument
from docling_core.types.io import DocumentStream
from haiku.rag.utils import text_to_docling_document
class FileReader:
@ -106,18 +106,10 @@ class FileReader:
content = f"```{language}\n{content}\n```"
# Convert text to DoclingDocument by wrapping as markdown
bytes_io = BytesIO(content.encode("utf-8"))
doc_stream = DocumentStream(name=f"{path.stem}.md", stream=bytes_io)
converter = DocumentConverter()
result = converter.convert(doc_stream)
return result.document
return text_to_docling_document(content, name=f"{path.stem}.md")
else:
# Fallback: try to read as text and convert to DoclingDocument
content = path.read_text(encoding="utf-8")
bytes_io = BytesIO(content.encode("utf-8"))
doc_stream = DocumentStream(name=f"{path.stem}.md", stream=bytes_io)
converter = DocumentConverter()
result = converter.convert(doc_stream)
return result.document
return text_to_docling_document(content, name=f"{path.stem}.md")
except Exception:
raise ValueError(f"Failed to parse file: {path}")

8
src/haiku/rag/store/repositories/chunk.py
View file

@ -1,6 +1,8 @@
import json
import re
from docling_core.types.doc.document import DoclingDocument
from haiku.rag.chunker import chunker
from haiku.rag.embeddings import get_embedder
from haiku.rag.store.models.chunk import Chunk
@ -197,11 +199,11 @@ class ChunkRepository(BaseRepository[Chunk]):
]
async def create_chunks_for_document(
self, document_id: int, content: str, commit: bool = True
self, document_id: int, document: DoclingDocument, commit: bool = True
) -> list[Chunk]:
"""Create chunks and embeddings for a document."""
"""Create chunks and embeddings for a document from DoclingDocument."""
# Chunk the document content
chunk_texts = await chunker.chunk(content)
chunk_texts = await chunker.chunk(document)
created_chunks = []
# Create chunks with embeddings using the create method

36
src/haiku/rag/store/repositories/document.py
View file

@ -1,8 +1,11 @@
import json
from typing import TYPE_CHECKING
from docling_core.types.doc.document import DoclingDocument
from haiku.rag.store.models.document import Document
from haiku.rag.store.repositories.base import BaseRepository
from haiku.rag.utils import text_to_docling_document
if TYPE_CHECKING:
from haiku.rag.store.models.chunk import Chunk
@ -20,8 +23,11 @@ class DocumentRepository(BaseRepository[Document]):
chunk_repository = ChunkRepository(store)
self.chunk_repository = chunk_repository
async def create(
self, entity: Document, chunks: list["Chunk"] | None = None
async def _create_with_docling(
self,
entity: Document,
docling_document: DoclingDocument,
chunks: list["Chunk"] | None = None,
) -> Document:
"""Create a document with its chunks and embeddings."""
if self.store._connection is None:
@ -62,9 +68,9 @@ class DocumentRepository(BaseRepository[Document]):
chunk.metadata["order"] = order
await self.chunk_repository.create(chunk, commit=False)
else:
# Create chunks and embeddings using ChunkRepository
# Create chunks and embeddings using DoclingDocument
await self.chunk_repository.create_chunks_for_document(
document_id, entity.content, commit=False
document_id, docling_document, commit=False
)
cursor.execute("COMMIT")
@ -74,6 +80,13 @@ class DocumentRepository(BaseRepository[Document]):
cursor.execute("ROLLBACK")
raise
async def create(self, entity: Document) -> Document:
"""Create a document with its chunks and embeddings."""
# Convert content to DoclingDocument
docling_document = text_to_docling_document(entity.content)
return await self._create_with_docling(entity, docling_document)
async def get_by_id(self, entity_id: int) -> Document | None:
"""Get a document by its ID."""
if self.store._connection is None:
@ -134,7 +147,9 @@ class DocumentRepository(BaseRepository[Document]):
updated_at=updated_at,
)
async def update(self, entity: Document) -> Document:
async def _update_with_docling(
self, entity: Document, docling_document: DoclingDocument
) -> Document:
"""Update an existing document and regenerate its chunks and embeddings."""
if self.store._connection is None:
raise ValueError("Store connection is not available")
@ -163,10 +178,10 @@ class DocumentRepository(BaseRepository[Document]):
},
)
# Delete existing chunks and regenerate using ChunkRepository
# Delete existing chunks and regenerate using DoclingDocument
await self.chunk_repository.delete_by_document_id(entity.id, commit=False)
await self.chunk_repository.create_chunks_for_document(
entity.id, entity.content, commit=False
entity.id, docling_document, commit=False
)
cursor.execute("COMMIT")
@ -176,6 +191,13 @@ class DocumentRepository(BaseRepository[Document]):
cursor.execute("ROLLBACK")
raise
async def update(self, entity: Document) -> Document:
"""Update an existing document and regenerate its chunks and embeddings."""
# Convert content to DoclingDocument
docling_document = text_to_docling_document(entity.content)
return await self._update_with_docling(entity, docling_document)
async def delete(self, entity_id: int) -> bool:
"""Delete a document and all its associated chunks and embeddings."""
# Delete chunks and embeddings first

6
tests/test_chunk.py
View file

@ -6,6 +6,7 @@ from haiku.rag.store.models.chunk import Chunk
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.asyncio
@ -77,8 +78,11 @@ async def test_create_chunks_for_document(qa_corpus: Dataset):
assert document_id is not None, "Document ID should not be None"
# Convert text to DoclingDocument
docling_document = text_to_docling_document(document_text, name="test.md")
# Test creating chunks for the document
chunks = await chunk_repo.create_chunks_for_document(document_id, document_text)
chunks = await chunk_repo.create_chunks_for_document(document_id, docling_document)
# Verify chunks were created
assert len(chunks) > 0

11
tests/test_chunker.py
View file

@ -2,13 +2,18 @@ import pytest
from datasets import Dataset
from haiku.rag.chunker import Chunker
from haiku.rag.utils import text_to_docling_document
@pytest.mark.asyncio
async def test_chunker(qa_corpus: Dataset):
chunker = Chunker()
doc = qa_corpus[0]["document_extracted"]
chunks = await Chunker().chunk(doc)
doc_text = qa_corpus[0]["document_extracted"]
# Convert text to DoclingDocument
doc = text_to_docling_document(doc_text, name="test.md")
chunks = await chunker.chunk(doc)
# Ensure that the text is split into multiple chunks
assert len(chunks) > 1
@ -27,7 +32,7 @@ async def test_chunker(qa_corpus: Dataset):
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=()))
original_tokens = len(Chunker.encoder.encode(doc_text, disallowed_special=()))
# Due to structure-aware chunking, we might have some variation in token count
# but it should be reasonable