Switch from requests to httpx.AsyncClient for docling-serve
This commit is contained in:
Yiorgis Gozadinos
parent
14bdb7cb27
commit
e40c352026
2 changed files with 173 additions and 137 deletions
|
|
@ -4,7 +4,7 @@ import asyncio
|
|||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, ClassVar
|
||||
|
||||
import requests
|
||||
import httpx
|
||||
|
||||
from haiku.rag.config import AppConfig
|
||||
from haiku.rag.converters.base import DocumentConverter
|
||||
|
|
@ -62,44 +62,11 @@ class DoclingServeConverter(DocumentConverter):
|
|||
"""Return list of file extensions supported by this converter."""
|
||||
return self.docling_serve_extensions + TextFileHandler.text_extensions
|
||||
|
||||
def _sync_make_request(
|
||||
self, files: dict, name: str, data: dict, headers: dict
|
||||
) -> "DoclingDocument":
|
||||
"""Synchronous HTTP request to docling-serve."""
|
||||
from docling_core.types.doc.document import DoclingDocument
|
||||
|
||||
url = f"{self.base_url}/v1/convert/file"
|
||||
response = requests.post(
|
||||
url,
|
||||
files=files,
|
||||
data=data,
|
||||
headers=headers,
|
||||
timeout=self.timeout,
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
|
||||
result = response.json()
|
||||
|
||||
if result["status"] not in ("success", "partial_success"):
|
||||
errors = result.get("errors", [])
|
||||
raise ValueError(f"Conversion failed: {errors}")
|
||||
|
||||
json_content = result["document"]["json_content"]
|
||||
|
||||
if json_content is None:
|
||||
raise ValueError(
|
||||
f"docling-serve did not return JSON content for {name}. "
|
||||
"This may indicate an unsupported file format."
|
||||
)
|
||||
|
||||
return DoclingDocument.model_validate(json_content)
|
||||
|
||||
async def _make_request(self, files: dict, name: str) -> "DoclingDocument":
|
||||
"""Make a request to docling-serve and return the DoclingDocument.
|
||||
|
||||
Args:
|
||||
files: Dictionary with files parameter for requests
|
||||
files: Dictionary with files parameter for httpx
|
||||
name: Name of the document being converted (for error messages)
|
||||
|
||||
Returns:
|
||||
|
|
@ -108,17 +75,19 @@ class DoclingServeConverter(DocumentConverter):
|
|||
Raises:
|
||||
ValueError: If conversion fails or service is unavailable
|
||||
"""
|
||||
from docling_core.types.doc.document import DoclingDocument
|
||||
|
||||
try:
|
||||
opts = self.config.processing.conversion_options
|
||||
|
||||
data = {
|
||||
"to_formats": ["json"],
|
||||
"do_ocr": opts.do_ocr,
|
||||
"force_ocr": opts.force_ocr,
|
||||
"do_table_structure": opts.do_table_structure,
|
||||
data: dict[str, str | list[str]] = {
|
||||
"to_formats": "json",
|
||||
"do_ocr": str(opts.do_ocr).lower(),
|
||||
"force_ocr": str(opts.force_ocr).lower(),
|
||||
"do_table_structure": str(opts.do_table_structure).lower(),
|
||||
"table_mode": opts.table_mode,
|
||||
"table_cell_matching": opts.table_cell_matching,
|
||||
"images_scale": opts.images_scale,
|
||||
"table_cell_matching": str(opts.table_cell_matching).lower(),
|
||||
"images_scale": str(opts.images_scale),
|
||||
}
|
||||
|
||||
if opts.ocr_lang:
|
||||
|
|
@ -128,26 +97,50 @@ class DoclingServeConverter(DocumentConverter):
|
|||
if self.api_key:
|
||||
headers["X-Api-Key"] = self.api_key
|
||||
|
||||
return await asyncio.to_thread(
|
||||
self._sync_make_request, files, name, data, headers
|
||||
)
|
||||
url = f"{self.base_url}/v1/convert/file"
|
||||
|
||||
except requests.exceptions.ConnectionError as e:
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
response = await client.post(
|
||||
url,
|
||||
files=files,
|
||||
data=data,
|
||||
headers=headers,
|
||||
)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
if result["status"] not in ("success", "partial_success"):
|
||||
errors = result.get("errors", [])
|
||||
raise ValueError(f"Conversion failed: {errors}")
|
||||
|
||||
json_content = result["document"]["json_content"]
|
||||
|
||||
if json_content is None:
|
||||
raise ValueError(
|
||||
f"docling-serve did not return JSON content for {name}. "
|
||||
"This may indicate an unsupported file format."
|
||||
)
|
||||
|
||||
return DoclingDocument.model_validate(json_content)
|
||||
|
||||
except httpx.ConnectError as e:
|
||||
raise ValueError(
|
||||
f"Could not connect to docling-serve at {self.base_url}. "
|
||||
f"Ensure the service is running and accessible. Error: {e}"
|
||||
)
|
||||
except requests.exceptions.Timeout as e:
|
||||
except httpx.TimeoutException as e:
|
||||
raise ValueError(
|
||||
f"Request to docling-serve timed out after {self.timeout}s. "
|
||||
f"Consider increasing the timeout in configuration. Error: {e}"
|
||||
)
|
||||
except requests.exceptions.HTTPError as e:
|
||||
except httpx.HTTPStatusError as e:
|
||||
if e.response.status_code == 401:
|
||||
raise ValueError(
|
||||
"Authentication failed. Check your API key configuration."
|
||||
)
|
||||
raise ValueError(f"HTTP error from docling-serve: {e}")
|
||||
except ValueError:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise ValueError(f"Failed to convert via docling-serve: {e}")
|
||||
|
||||
|
|
@ -175,11 +168,12 @@ class DoclingServeConverter(DocumentConverter):
|
|||
except Exception as e:
|
||||
raise ValueError(f"Failed to read text file {path}: {e}")
|
||||
|
||||
def read_and_prepare_files():
|
||||
def read_file():
|
||||
with open(path, "rb") as f:
|
||||
return {"files": (path.name, f.read(), "application/octet-stream")}
|
||||
return f.read()
|
||||
|
||||
files = await asyncio.to_thread(read_and_prepare_files)
|
||||
file_content = await asyncio.to_thread(read_file)
|
||||
files = {"files": (path.name, file_content, "application/octet-stream")}
|
||||
return await self._make_request(files, path.name)
|
||||
|
||||
async def convert_text(
|
||||
|
|
@ -199,8 +193,6 @@ class DoclingServeConverter(DocumentConverter):
|
|||
Raises:
|
||||
ValueError: If the text cannot be converted.
|
||||
"""
|
||||
from io import BytesIO
|
||||
|
||||
text_bytes = text.encode("utf-8")
|
||||
files = {"files": (name, BytesIO(text_bytes), "text/markdown")}
|
||||
files = {"files": (name, text_bytes, "text/markdown")}
|
||||
return await self._make_request(files, name)
|
||||
|
|
|
|||
|
|
@ -2,8 +2,9 @@
|
|||
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from unittest.mock import Mock, patch
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import requests
|
||||
from docling_core.types.doc.document import DoclingDocument
|
||||
|
|
@ -201,8 +202,7 @@ class TestDoclingServeConverter:
|
|||
assert ".md" in extensions
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("haiku.rag.converters.docling_serve.requests.post")
|
||||
async def test_convert_text_success(self, mock_post, converter):
|
||||
async def test_convert_text_success(self, converter):
|
||||
"""Test successful text conversion via docling-serve."""
|
||||
mock_response = Mock()
|
||||
mock_response.status_code = 200
|
||||
|
|
@ -210,16 +210,22 @@ class TestDoclingServeConverter:
|
|||
"status": "success",
|
||||
"document": {"json_content": create_mock_docling_document_json("test")},
|
||||
}
|
||||
mock_post.return_value = mock_response
|
||||
mock_response.raise_for_status = Mock()
|
||||
|
||||
doc = await converter.convert_text("# Test", name="test.md")
|
||||
assert isinstance(doc, DoclingDocument)
|
||||
assert doc.version == "1.8.0"
|
||||
mock_post.assert_called_once()
|
||||
with patch("httpx.AsyncClient") as mock_client_class:
|
||||
mock_client = AsyncMock()
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
||||
mock_client.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_client_class.return_value = mock_client
|
||||
|
||||
doc = await converter.convert_text("# Test", name="test.md")
|
||||
assert isinstance(doc, DoclingDocument)
|
||||
assert doc.version == "1.8.0"
|
||||
mock_client.post.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("haiku.rag.converters.docling_serve.requests.post")
|
||||
async def test_convert_text_with_api_key(self, mock_post, config):
|
||||
async def test_convert_text_with_api_key(self, config):
|
||||
"""Test that API key is included in request headers."""
|
||||
config.providers.docling_serve.api_key = "test-key"
|
||||
converter = DoclingServeConverter(config)
|
||||
|
|
@ -230,17 +236,23 @@ class TestDoclingServeConverter:
|
|||
"status": "success",
|
||||
"document": {"json_content": create_mock_docling_document_json("test")},
|
||||
}
|
||||
mock_post.return_value = mock_response
|
||||
mock_response.raise_for_status = Mock()
|
||||
|
||||
await converter.convert_text("# Test")
|
||||
with patch("httpx.AsyncClient") as mock_client_class:
|
||||
mock_client = AsyncMock()
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
||||
mock_client.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_client_class.return_value = mock_client
|
||||
|
||||
call_kwargs = mock_post.call_args.kwargs
|
||||
assert "headers" in call_kwargs
|
||||
assert call_kwargs["headers"]["X-Api-Key"] == "test-key"
|
||||
await converter.convert_text("# Test")
|
||||
|
||||
call_kwargs = mock_client.post.call_args.kwargs
|
||||
assert "headers" in call_kwargs
|
||||
assert call_kwargs["headers"]["X-Api-Key"] == "test-key"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("haiku.rag.converters.docling_serve.requests.post")
|
||||
async def test_conversion_options_passed_to_api(self, mock_post, config):
|
||||
async def test_conversion_options_passed_to_api(self, config):
|
||||
"""Test that conversion options are passed to docling-serve API."""
|
||||
config.processing.conversion_options.do_ocr = False
|
||||
config.processing.conversion_options.force_ocr = True
|
||||
|
|
@ -257,63 +269,78 @@ class TestDoclingServeConverter:
|
|||
"status": "success",
|
||||
"document": {"json_content": create_mock_docling_document_json("test")},
|
||||
}
|
||||
mock_post.return_value = mock_response
|
||||
mock_response.raise_for_status = Mock()
|
||||
|
||||
await converter.convert_text("# Test")
|
||||
with patch("httpx.AsyncClient") as mock_client_class:
|
||||
mock_client = AsyncMock()
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
||||
mock_client.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_client_class.return_value = mock_client
|
||||
|
||||
call_kwargs = mock_post.call_args.kwargs
|
||||
assert "data" in call_kwargs
|
||||
data = call_kwargs["data"]
|
||||
assert data["do_ocr"] is False
|
||||
assert data["force_ocr"] is True
|
||||
assert data["ocr_lang"] == ["en", "fr"]
|
||||
assert "pdf_backend" not in data
|
||||
assert data["table_mode"] == "fast"
|
||||
assert data["table_cell_matching"] is False
|
||||
assert data["do_table_structure"] is False
|
||||
assert data["images_scale"] == 3.0
|
||||
await converter.convert_text("# Test")
|
||||
|
||||
call_kwargs = mock_client.post.call_args.kwargs
|
||||
assert "data" in call_kwargs
|
||||
data = call_kwargs["data"]
|
||||
assert data["do_ocr"] == "false"
|
||||
assert data["force_ocr"] == "true"
|
||||
assert data["ocr_lang"] == ["en", "fr"]
|
||||
assert data["table_mode"] == "fast"
|
||||
assert data["table_cell_matching"] == "false"
|
||||
assert data["do_table_structure"] == "false"
|
||||
assert data["images_scale"] == "3.0"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("haiku.rag.converters.docling_serve.requests.post")
|
||||
async def test_convert_text_connection_error(self, mock_post, converter):
|
||||
async def test_convert_text_connection_error(self, converter):
|
||||
"""Test handling of connection errors."""
|
||||
import requests
|
||||
with patch("httpx.AsyncClient") as mock_client_class:
|
||||
mock_client = AsyncMock()
|
||||
mock_client.post = AsyncMock(
|
||||
side_effect=httpx.ConnectError("Connection failed")
|
||||
)
|
||||
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
||||
mock_client.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_client_class.return_value = mock_client
|
||||
|
||||
mock_post.side_effect = requests.exceptions.ConnectionError("Connection failed")
|
||||
|
||||
with pytest.raises(ValueError, match="Could not connect to docling-serve"):
|
||||
await converter.convert_text("# Test")
|
||||
with pytest.raises(ValueError, match="Could not connect to docling-serve"):
|
||||
await converter.convert_text("# Test")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("haiku.rag.converters.docling_serve.requests.post")
|
||||
async def test_convert_text_timeout_error(self, mock_post, converter):
|
||||
async def test_convert_text_timeout_error(self, converter):
|
||||
"""Test handling of timeout errors."""
|
||||
import requests
|
||||
with patch("httpx.AsyncClient") as mock_client_class:
|
||||
mock_client = AsyncMock()
|
||||
mock_client.post = AsyncMock(side_effect=httpx.TimeoutException("Timeout"))
|
||||
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
||||
mock_client.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_client_class.return_value = mock_client
|
||||
|
||||
mock_post.side_effect = requests.exceptions.Timeout("Timeout")
|
||||
|
||||
with pytest.raises(ValueError, match="timed out"):
|
||||
await converter.convert_text("# Test")
|
||||
with pytest.raises(ValueError, match="timed out"):
|
||||
await converter.convert_text("# Test")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("haiku.rag.converters.docling_serve.requests.post")
|
||||
async def test_convert_text_auth_error(self, mock_post, converter):
|
||||
async def test_convert_text_auth_error(self, converter):
|
||||
"""Test handling of authentication errors."""
|
||||
import requests
|
||||
|
||||
mock_response = Mock()
|
||||
mock_response.status_code = 401
|
||||
mock_response.raise_for_status.side_effect = requests.exceptions.HTTPError(
|
||||
response=mock_response
|
||||
)
|
||||
mock_post.return_value = mock_response
|
||||
|
||||
with pytest.raises(ValueError, match="Authentication failed"):
|
||||
await converter.convert_text("# Test")
|
||||
with patch("httpx.AsyncClient") as mock_client_class:
|
||||
mock_client = AsyncMock()
|
||||
mock_client.post = AsyncMock(
|
||||
side_effect=httpx.HTTPStatusError(
|
||||
"Auth failed", request=Mock(), response=mock_response
|
||||
)
|
||||
)
|
||||
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
||||
mock_client.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_client_class.return_value = mock_client
|
||||
|
||||
with pytest.raises(ValueError, match="Authentication failed"):
|
||||
await converter.convert_text("# Test")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("haiku.rag.converters.docling_serve.requests.post")
|
||||
async def test_convert_text_no_json_content(self, mock_post, converter):
|
||||
async def test_convert_text_no_json_content(self, converter):
|
||||
"""Test handling when docling-serve returns no JSON content."""
|
||||
mock_response = Mock()
|
||||
mock_response.status_code = 200
|
||||
|
|
@ -321,14 +348,20 @@ class TestDoclingServeConverter:
|
|||
"status": "success",
|
||||
"document": {"json_content": None},
|
||||
}
|
||||
mock_post.return_value = mock_response
|
||||
mock_response.raise_for_status = Mock()
|
||||
|
||||
with pytest.raises(ValueError, match="did not return JSON content"):
|
||||
await converter.convert_text("# Test")
|
||||
with patch("httpx.AsyncClient") as mock_client_class:
|
||||
mock_client = AsyncMock()
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
||||
mock_client.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_client_class.return_value = mock_client
|
||||
|
||||
with pytest.raises(ValueError, match="did not return JSON content"):
|
||||
await converter.convert_text("# Test")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("haiku.rag.converters.docling_serve.requests.post")
|
||||
async def test_convert_file_pdf(self, mock_post, converter):
|
||||
async def test_convert_file_pdf(self, converter):
|
||||
"""Test converting PDF file via docling-serve."""
|
||||
mock_response = Mock()
|
||||
mock_response.status_code = 200
|
||||
|
|
@ -336,20 +369,26 @@ class TestDoclingServeConverter:
|
|||
"status": "success",
|
||||
"document": {"json_content": create_mock_docling_document_json("test")},
|
||||
}
|
||||
mock_post.return_value = mock_response
|
||||
mock_response.raise_for_status = Mock()
|
||||
|
||||
with tempfile.NamedTemporaryFile(suffix=".pdf") as f:
|
||||
f.write(b"fake pdf content")
|
||||
f.flush()
|
||||
temp_path = Path(f.name)
|
||||
doc = await converter.convert_file(temp_path)
|
||||
with patch("httpx.AsyncClient") as mock_client_class:
|
||||
mock_client = AsyncMock()
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
||||
mock_client.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_client_class.return_value = mock_client
|
||||
|
||||
assert isinstance(doc, DoclingDocument)
|
||||
mock_post.assert_called_once()
|
||||
with tempfile.NamedTemporaryFile(suffix=".pdf") as f:
|
||||
f.write(b"fake pdf content")
|
||||
f.flush()
|
||||
temp_path = Path(f.name)
|
||||
doc = await converter.convert_file(temp_path)
|
||||
|
||||
assert isinstance(doc, DoclingDocument)
|
||||
mock_client.post.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("haiku.rag.converters.docling_serve.requests.post")
|
||||
async def test_convert_file_text(self, mock_post, converter):
|
||||
async def test_convert_file_text(self, converter):
|
||||
"""Test converting text file (reads locally, sends to docling-serve)."""
|
||||
mock_response = Mock()
|
||||
mock_response.status_code = 200
|
||||
|
|
@ -357,20 +396,25 @@ class TestDoclingServeConverter:
|
|||
"status": "success",
|
||||
"document": {"json_content": create_mock_docling_document_json("test")},
|
||||
}
|
||||
mock_post.return_value = mock_response
|
||||
mock_response.raise_for_status = Mock()
|
||||
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".py") as f:
|
||||
f.write("def hello():\n pass")
|
||||
f.flush()
|
||||
temp_path = Path(f.name)
|
||||
doc = await converter.convert_file(temp_path)
|
||||
with patch("httpx.AsyncClient") as mock_client_class:
|
||||
mock_client = AsyncMock()
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
||||
mock_client.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_client_class.return_value = mock_client
|
||||
|
||||
assert isinstance(doc, DoclingDocument)
|
||||
# Should call docling-serve for conversion
|
||||
mock_post.assert_called_once()
|
||||
# Check that code was wrapped in code block
|
||||
call_kwargs = mock_post.call_args.kwargs
|
||||
assert "files" in call_kwargs
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".py") as f:
|
||||
f.write("def hello():\n pass")
|
||||
f.flush()
|
||||
temp_path = Path(f.name)
|
||||
doc = await converter.convert_file(temp_path)
|
||||
|
||||
assert isinstance(doc, DoclingDocument)
|
||||
mock_client.post.assert_called_once()
|
||||
call_kwargs = mock_client.post.call_args.kwargs
|
||||
assert "files" in call_kwargs
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
|
|
|
|||
Loading…
Reference in a new issue