Merge pull request #31 from ggozad/fix/docling-chunking
Refactor reading & chunking around docling documents.
This commit is contained in:
Yiorgis Gozadinos
GitHub
commit
17608844b5
12 changed files with 300 additions and 57 deletions
|
|
@ -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]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import asyncio
|
||||
import warnings
|
||||
from importlib.metadata import version
|
||||
from pathlib import Path
|
||||
|
||||
|
|
@ -9,6 +10,9 @@ from haiku.rag.app import HaikuRAGApp
|
|||
from haiku.rag.config import Config
|
||||
from haiku.rag.utils import is_up_to_date
|
||||
|
||||
if not Config.ENV == "development":
|
||||
warnings.filterwarnings("ignore")
|
||||
|
||||
cli = typer.Typer(
|
||||
context_settings={"help_option_names": ["-h", "--help"]}, no_args_is_help=True
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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 = {}
|
||||
|
|
@ -101,9 +125,12 @@ class HaikuRAG:
|
|||
parsed_url = urlparse(source_str)
|
||||
if parsed_url.scheme in ("http", "https"):
|
||||
return await self._create_or_update_document_from_url(source_str, metadata)
|
||||
|
||||
# Handle as file path
|
||||
source_path = Path(source) if isinstance(source, str) else source
|
||||
elif parsed_url.scheme == "file":
|
||||
# Handle file:// URI by converting to path
|
||||
source_path = Path(parsed_url.path)
|
||||
else:
|
||||
# Handle as regular file path
|
||||
source_path = Path(source) if isinstance(source, str) else source
|
||||
if source_path.suffix.lower() not in FileReader.extensions:
|
||||
raise ValueError(f"Unsupported file extension: {source_path.suffix}")
|
||||
|
||||
|
|
@ -119,7 +146,7 @@ class HaikuRAG:
|
|||
# MD5 unchanged, return existing document
|
||||
return existing_doc
|
||||
|
||||
content = FileReader.parse_file(source_path)
|
||||
docling_document = FileReader.parse_file(source_path)
|
||||
|
||||
# Get content type from file extension
|
||||
content_type, _ = mimetypes.guess_type(str(source_path))
|
||||
|
|
@ -131,13 +158,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(
|
||||
|
|
@ -193,18 +222,20 @@ class HaikuRAG:
|
|||
temp_path = Path(temp_file.name)
|
||||
|
||||
# Parse the content using FileReader
|
||||
content = FileReader.parse_file(temp_path)
|
||||
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(
|
||||
|
|
@ -262,7 +293,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."""
|
||||
|
|
@ -328,6 +364,13 @@ class HaikuRAG:
|
|||
async def rebuild_database(self) -> AsyncGenerator[int, None]:
|
||||
"""Rebuild the database by deleting all chunks and re-indexing all documents.
|
||||
|
||||
For documents with URIs:
|
||||
- Deletes the document and re-adds it from source if source exists
|
||||
- Skips documents where source no longer exists
|
||||
|
||||
For documents without URIs:
|
||||
- Re-creates chunks from existing content
|
||||
|
||||
Yields:
|
||||
int: The ID of the document currently being processed
|
||||
"""
|
||||
|
|
@ -343,9 +386,36 @@ class HaikuRAG:
|
|||
documents = await self.list_documents()
|
||||
|
||||
for doc in documents:
|
||||
if doc.id is not None:
|
||||
assert doc.id is not None, "Document ID should not be None"
|
||||
if doc.uri:
|
||||
# Document has a URI - delete and try to re-add from source
|
||||
try:
|
||||
# Delete the old document first
|
||||
await self.delete_document(doc.id)
|
||||
|
||||
# Try to re-create from source (this creates the document with chunks)
|
||||
new_doc = await self.create_document_from_source(
|
||||
doc.uri, doc.metadata or {}
|
||||
)
|
||||
|
||||
assert new_doc.id is not None, "New document ID should not be None"
|
||||
yield new_doc.id
|
||||
|
||||
except (FileNotFoundError, ValueError, OSError) as e:
|
||||
# Source doesn't exist or can't be accessed - document already deleted, skip
|
||||
print(f"Skipping document with URI {doc.uri}: {e}")
|
||||
continue
|
||||
except Exception as e:
|
||||
# Unexpected error - log it and skip
|
||||
print(
|
||||
f"Unexpected error processing document with URI {doc.uri}: {e}"
|
||||
)
|
||||
continue
|
||||
else:
|
||||
# Document without URI - re-create chunks from existing content
|
||||
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
|
||||
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ load_dotenv()
|
|||
|
||||
|
||||
class AppConfig(BaseModel):
|
||||
ENV: str = "development"
|
||||
ENV: str = "production"
|
||||
|
||||
DEFAULT_DATA_DIR: Path = get_default_data_dir()
|
||||
MONITOR_DIRECTORIES: list[Path] = []
|
||||
|
|
|
|||
|
|
@ -2,6 +2,9 @@ from pathlib import Path
|
|||
from typing import ClassVar
|
||||
|
||||
from docling.document_converter import DocumentConverter
|
||||
from docling_core.types.doc.document import DoclingDocument
|
||||
|
||||
from haiku.rag.utils import text_to_docling_document
|
||||
|
||||
|
||||
class FileReader:
|
||||
|
|
@ -84,7 +87,7 @@ class FileReader:
|
|||
extensions: ClassVar[list[str]] = docling_extensions + text_extensions
|
||||
|
||||
@staticmethod
|
||||
def parse_file(path: Path) -> str:
|
||||
def parse_file(path: Path) -> DoclingDocument:
|
||||
try:
|
||||
file_extension = path.suffix.lower()
|
||||
|
||||
|
|
@ -92,7 +95,7 @@ class FileReader:
|
|||
# Use docling for complex document formats
|
||||
converter = DocumentConverter()
|
||||
result = converter.convert(path)
|
||||
return result.document.export_to_markdown()
|
||||
return result.document
|
||||
elif file_extension in FileReader.text_extensions:
|
||||
# Read plain text files directly
|
||||
content = path.read_text(encoding="utf-8")
|
||||
|
|
@ -100,11 +103,13 @@ class FileReader:
|
|||
# Wrap code files (but not plain txt) in markdown code blocks for better presentation
|
||||
if file_extension in FileReader.code_markdown_identifier:
|
||||
language = FileReader.code_markdown_identifier[file_extension]
|
||||
return f"```{language}\n{content}\n```"
|
||||
content = f"```{language}\n{content}\n```"
|
||||
|
||||
return content
|
||||
# Convert text to DoclingDocument by wrapping as markdown
|
||||
return text_to_docling_document(content, name=f"{path.stem}.md")
|
||||
else:
|
||||
# Fallback: try to read as text
|
||||
return path.read_text(encoding="utf-8")
|
||||
# Fallback: try to read as text and convert to DoclingDocument
|
||||
content = path.read_text(encoding="utf-8")
|
||||
return text_to_docling_document(content, name=f"{path.stem}.md")
|
||||
except Exception:
|
||||
raise ValueError(f"Failed to parse file: {path}")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -1,8 +1,12 @@
|
|||
import sys
|
||||
from importlib import metadata
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
|
||||
import httpx
|
||||
from docling.document_converter import DocumentConverter
|
||||
from docling_core.types.doc.document import DoclingDocument
|
||||
from docling_core.types.io import DocumentStream
|
||||
from packaging.version import Version, parse
|
||||
|
||||
|
||||
|
|
@ -77,3 +81,20 @@ async def is_up_to_date() -> tuple[bool, Version, Version]:
|
|||
# If no network connection, do not raise alarms.
|
||||
pypi_version = running_version
|
||||
return running_version >= pypi_version, running_version, pypi_version
|
||||
|
||||
|
||||
def text_to_docling_document(text: str, name: str = "content.md") -> DoclingDocument:
|
||||
"""Convert text content to a DoclingDocument.
|
||||
|
||||
Args:
|
||||
text: The text content to convert.
|
||||
name: The name to use for the document stream (defaults to "content.md").
|
||||
|
||||
Returns:
|
||||
A DoclingDocument created from the text content.
|
||||
"""
|
||||
bytes_io = BytesIO(text.encode("utf-8"))
|
||||
doc_stream = DocumentStream(name=name, stream=bytes_io)
|
||||
converter = DocumentConverter()
|
||||
result = converter.convert(doc_stream)
|
||||
return result.document
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -15,8 +15,9 @@ def test_code_file_wrapped_in_code_block():
|
|||
f.flush()
|
||||
temp_path = Path(f.name)
|
||||
|
||||
result = FileReader.parse_file(temp_path)
|
||||
document = FileReader.parse_file(temp_path)
|
||||
result = document.export_to_markdown()
|
||||
|
||||
assert result.startswith("```python\n")
|
||||
assert result.startswith("```\n")
|
||||
assert result.endswith("\n```")
|
||||
assert "def hello_world():" in result
|
||||
|
|
|
|||
|
|
@ -1,4 +1,8 @@
|
|||
from haiku.rag.utils import int_to_semantic_version, semantic_version_to_int
|
||||
from haiku.rag.utils import (
|
||||
int_to_semantic_version,
|
||||
semantic_version_to_int,
|
||||
text_to_docling_document,
|
||||
)
|
||||
|
||||
|
||||
def test_sqlite_user_version():
|
||||
|
|
@ -13,3 +17,117 @@ def test_sqlite_user_version():
|
|||
version = "255.255.255"
|
||||
assert semantic_version_to_int(version) == 16777215
|
||||
assert int_to_semantic_version(16777215) == version
|
||||
|
||||
|
||||
def test_text_to_docling_document():
|
||||
"""Test the text_to_docling_document utility function."""
|
||||
# Test basic text conversion
|
||||
simple_text = "This is a simple text document."
|
||||
doc = text_to_docling_document(simple_text)
|
||||
|
||||
# Verify it returns a DoclingDocument
|
||||
from docling_core.types.doc.document import DoclingDocument
|
||||
|
||||
assert isinstance(doc, DoclingDocument)
|
||||
|
||||
# Verify the content can be exported back to markdown
|
||||
markdown = doc.export_to_markdown()
|
||||
assert "This is a simple text document." in markdown
|
||||
|
||||
|
||||
def test_text_to_docling_document_with_custom_name():
|
||||
"""Test text_to_docling_document with custom name parameter."""
|
||||
code_text = """# Python Code
|
||||
|
||||
```python
|
||||
def hello():
|
||||
print("Hello, World!")
|
||||
return True
|
||||
```"""
|
||||
|
||||
doc = text_to_docling_document(code_text, name="hello.md")
|
||||
|
||||
# Verify it's a valid DoclingDocument
|
||||
from docling_core.types.doc.document import DoclingDocument
|
||||
|
||||
assert isinstance(doc, DoclingDocument)
|
||||
|
||||
# Verify the content is preserved
|
||||
markdown = doc.export_to_markdown()
|
||||
assert "def hello():" in markdown
|
||||
assert "Hello, World!" in markdown
|
||||
|
||||
|
||||
def test_text_to_docling_document_markdown_content():
|
||||
"""Test text_to_docling_document with markdown content."""
|
||||
markdown_text = """# Test Document
|
||||
|
||||
This is a test document with:
|
||||
|
||||
- List item 1
|
||||
- List item 2
|
||||
|
||||
## Code Example
|
||||
|
||||
```python
|
||||
def test():
|
||||
return "Hello"
|
||||
```
|
||||
|
||||
**Bold text** and *italic text*."""
|
||||
|
||||
doc = text_to_docling_document(markdown_text, name="test.md")
|
||||
|
||||
# Verify it's a DoclingDocument
|
||||
from docling_core.types.doc.document import DoclingDocument
|
||||
|
||||
assert isinstance(doc, DoclingDocument)
|
||||
|
||||
# Verify the markdown structure is preserved
|
||||
result_markdown = doc.export_to_markdown()
|
||||
assert "# Test Document" in result_markdown
|
||||
assert "List item 1" in result_markdown
|
||||
assert "def test():" in result_markdown
|
||||
|
||||
|
||||
def test_text_to_docling_document_empty_content():
|
||||
"""Test text_to_docling_document with empty content."""
|
||||
doc = text_to_docling_document("")
|
||||
|
||||
# Should still create a valid DoclingDocument
|
||||
from docling_core.types.doc.document import DoclingDocument
|
||||
|
||||
assert isinstance(doc, DoclingDocument)
|
||||
|
||||
# Export should work even with empty content
|
||||
markdown = doc.export_to_markdown()
|
||||
assert isinstance(markdown, str)
|
||||
|
||||
|
||||
def test_text_to_docling_document_unicode_content():
|
||||
"""Test text_to_docling_document with unicode content."""
|
||||
unicode_text = """# 测试文档
|
||||
|
||||
这是一个包含中文的测试文档。
|
||||
|
||||
## Código en Español
|
||||
```javascript
|
||||
function saludar() {
|
||||
return "¡Hola mundo!";
|
||||
}
|
||||
```
|
||||
|
||||
Emoji test: 🚀 ✅ 📝"""
|
||||
|
||||
doc = text_to_docling_document(unicode_text, name="unicode.md")
|
||||
|
||||
# Verify it's a DoclingDocument
|
||||
from docling_core.types.doc.document import DoclingDocument
|
||||
|
||||
assert isinstance(doc, DoclingDocument)
|
||||
|
||||
# Verify unicode content is preserved
|
||||
result_markdown = doc.export_to_markdown()
|
||||
assert "测试文档" in result_markdown
|
||||
assert "¡Hola mundo!" in result_markdown
|
||||
assert "🚀" in result_markdown
|
||||
|
|
|
|||
Loading…
Reference in a new issue