HaikuRag.convert() utility method
This commit is contained in:
Yiorgis Gozadinos
parent
ab57c55eaa
commit
ace915473c
3 changed files with 163 additions and 1 deletions
|
|
@ -38,6 +38,11 @@
|
|||
- When `docling_document_json` is provided without `chunks`, content is extracted and document is rechunked
|
||||
- When `docling_document_json` is provided with `chunks`, both are stored (chunks used as-is)
|
||||
- `content` and `docling_document_json` are mutually exclusive to avoid ambiguity
|
||||
- **New `convert()` Method**: Convert files, URLs, or text to DoclingDocument
|
||||
- `client.convert(Path(...))` - convert local file
|
||||
- `client.convert("https://...")` - download and convert URL
|
||||
- `client.convert("text content")` - convert plain text
|
||||
- Supports `file://` URIs
|
||||
|
||||
### Changed
|
||||
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ from collections.abc import AsyncGenerator
|
|||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
from typing import TYPE_CHECKING, overload
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import httpx
|
||||
|
|
@ -24,6 +24,8 @@ from haiku.rag.store.repositories.document import DocumentRepository
|
|||
from haiku.rag.store.repositories.settings import SettingsRepository
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from docling_core.types.doc.document import DoclingDocument
|
||||
|
||||
from haiku.rag.graph.common.models import Citation
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
|
@ -90,6 +92,86 @@ class HaikuRAG:
|
|||
self.close()
|
||||
return False
|
||||
|
||||
# =========================================================================
|
||||
# Processing Primitives
|
||||
# =========================================================================
|
||||
|
||||
@overload
|
||||
async def convert(self, source: Path) -> "DoclingDocument": ...
|
||||
|
||||
@overload
|
||||
async def convert(self, source: str) -> "DoclingDocument": ...
|
||||
|
||||
async def convert(self, source: Path | str) -> "DoclingDocument":
|
||||
"""Convert a file, URL, or text to DoclingDocument.
|
||||
|
||||
Args:
|
||||
source: One of:
|
||||
- Path: Local file path to convert
|
||||
- str (URL): HTTP/HTTPS URL to download and convert
|
||||
- str (text): Raw text content to convert
|
||||
|
||||
Returns:
|
||||
DoclingDocument from the converted source.
|
||||
|
||||
Raises:
|
||||
ValueError: If the file doesn't exist or has unsupported extension.
|
||||
httpx.RequestError: If URL download fails.
|
||||
"""
|
||||
converter = get_converter(self._config)
|
||||
|
||||
# Path object - convert file directly
|
||||
if isinstance(source, Path):
|
||||
if not source.exists():
|
||||
raise ValueError(f"File does not exist: {source}")
|
||||
if source.suffix.lower() not in converter.supported_extensions:
|
||||
raise ValueError(f"Unsupported file extension: {source.suffix}")
|
||||
return await converter.convert_file(source)
|
||||
|
||||
# String - check if URL or text
|
||||
parsed = urlparse(source)
|
||||
|
||||
if parsed.scheme in ("http", "https"):
|
||||
# URL - download and convert
|
||||
async with httpx.AsyncClient() as http:
|
||||
response = await http.get(source)
|
||||
response.raise_for_status()
|
||||
|
||||
content_type = response.headers.get("content-type", "").lower()
|
||||
file_extension = self._get_extension_from_content_type_or_url(
|
||||
source, content_type
|
||||
)
|
||||
|
||||
if file_extension not in converter.supported_extensions:
|
||||
raise ValueError(
|
||||
f"Unsupported content type/extension: {content_type}/{file_extension}"
|
||||
)
|
||||
|
||||
with tempfile.NamedTemporaryFile(
|
||||
mode="wb", suffix=file_extension, delete=False
|
||||
) as temp_file:
|
||||
temp_file.write(response.content)
|
||||
temp_file.flush()
|
||||
temp_path = Path(temp_file.name)
|
||||
|
||||
try:
|
||||
return await converter.convert_file(temp_path)
|
||||
finally:
|
||||
temp_path.unlink(missing_ok=True)
|
||||
|
||||
elif parsed.scheme == "file":
|
||||
# file:// URI
|
||||
file_path = Path(parsed.path)
|
||||
if not file_path.exists():
|
||||
raise ValueError(f"File does not exist: {file_path}")
|
||||
if file_path.suffix.lower() not in converter.supported_extensions:
|
||||
raise ValueError(f"Unsupported file extension: {file_path.suffix}")
|
||||
return await converter.convert_file(file_path)
|
||||
|
||||
else:
|
||||
# Treat as text content
|
||||
return await converter.convert_text(source)
|
||||
|
||||
async def _create_document_with_docling(
|
||||
self,
|
||||
docling_document,
|
||||
|
|
|
|||
|
|
@ -1564,3 +1564,78 @@ async def test_client_visualize_chunk_with_pdf(temp_db_path):
|
|||
# Verify returned objects are PIL Images
|
||||
for img in images:
|
||||
assert isinstance(img, PILImage)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# convert() method tests
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_convert_text(temp_db_path):
|
||||
"""Test convert() with plain text content."""
|
||||
from docling_core.types.doc.document import DoclingDocument
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
text = "This is some test content for conversion."
|
||||
docling_doc = await client.convert(text)
|
||||
|
||||
assert isinstance(docling_doc, DoclingDocument)
|
||||
# Check the content is preserved in markdown export
|
||||
markdown = docling_doc.export_to_markdown()
|
||||
assert "test content" in markdown
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_convert_file(temp_db_path):
|
||||
"""Test convert() with a file path."""
|
||||
from docling_core.types.doc.document import DoclingDocument
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
temp_path = Path(temp_dir) / "test.txt"
|
||||
temp_path.write_text("File content for conversion test.")
|
||||
|
||||
docling_doc = await client.convert(temp_path)
|
||||
|
||||
assert isinstance(docling_doc, DoclingDocument)
|
||||
markdown = docling_doc.export_to_markdown()
|
||||
assert "File content" in markdown
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_convert_file_not_found(temp_db_path):
|
||||
"""Test convert() raises ValueError for non-existent file."""
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
with pytest.raises(ValueError, match="File does not exist"):
|
||||
await client.convert(Path("/nonexistent/path/file.txt"))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_convert_unsupported_extension(temp_db_path):
|
||||
"""Test convert() raises ValueError for unsupported file extension."""
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
temp_path = Path(temp_dir) / "test.xyz"
|
||||
temp_path.write_text("content")
|
||||
|
||||
with pytest.raises(ValueError, match="Unsupported file extension"):
|
||||
await client.convert(temp_path)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_convert_file_uri(temp_db_path):
|
||||
"""Test convert() with a file:// URI string."""
|
||||
from docling_core.types.doc.document import DoclingDocument
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
temp_path = Path(temp_dir) / "test.txt"
|
||||
temp_path.write_text("URI file content.")
|
||||
file_uri = temp_path.as_uri()
|
||||
|
||||
docling_doc = await client.convert(file_uri)
|
||||
|
||||
assert isinstance(docling_doc, DoclingDocument)
|
||||
markdown = docling_doc.export_to_markdown()
|
||||
assert "URI file content" in markdown
|
||||
|
|
|
|||
Loading…
Reference in a new issue