diff --git a/CHANGELOG.md b/CHANGELOG.md index 87f3c16a..5a36ca1d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 8f22316e..dd4741b9 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -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, diff --git a/tests/test_client.py b/tests/test_client.py index ce78bff8..478bd067 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -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