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

Switch from requests to httpx.AsyncClient for docling-serve

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-11-26 18:00:29 +02:00
parent 14bdb7cb27
commit e40c352026
No known key found for this signature in database
2 changed files with 173 additions and 137 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

102
haiku_rag_slim/haiku/rag/converters/docling_serve.py
View file

@ -4,7 +4,7 @@ import asyncio
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING, ClassVar from typing import TYPE_CHECKING, ClassVar
import requests import httpx
from haiku.rag.config import AppConfig from haiku.rag.config import AppConfig
from haiku.rag.converters.base import DocumentConverter from haiku.rag.converters.base import DocumentConverter
@ -62,23 +62,51 @@ class DoclingServeConverter(DocumentConverter):
"""Return list of file extensions supported by this converter.""" """Return list of file extensions supported by this converter."""
return self.docling_serve_extensions + TextFileHandler.text_extensions return self.docling_serve_extensions + TextFileHandler.text_extensions
def _sync_make_request( async def _make_request(self, files: dict, name: str) -> "DoclingDocument":
self, files: dict, name: str, data: dict, headers: dict """Make a request to docling-serve and return the DoclingDocument.
) -> "DoclingDocument":
"""Synchronous HTTP request to docling-serve.""" Args:
files: Dictionary with files parameter for httpx
name: Name of the document being converted (for error messages)
Returns:
DoclingDocument representation
Raises:
ValueError: If conversion fails or service is unavailable
"""
from docling_core.types.doc.document import DoclingDocument from docling_core.types.doc.document import DoclingDocument
try:
opts = self.config.processing.conversion_options
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": str(opts.table_cell_matching).lower(),
"images_scale": str(opts.images_scale),
}
if opts.ocr_lang:
data["ocr_lang"] = opts.ocr_lang
headers = {}
if self.api_key:
headers["X-Api-Key"] = self.api_key
url = f"{self.base_url}/v1/convert/file" url = f"{self.base_url}/v1/convert/file"
response = requests.post(
async with httpx.AsyncClient(timeout=self.timeout) as client:
response = await client.post(
url, url,
files=files, files=files,
data=data, data=data,
headers=headers, headers=headers,
timeout=self.timeout,
) )
response.raise_for_status() response.raise_for_status()
result = response.json() result = response.json()
if result["status"] not in ("success", "partial_success"): if result["status"] not in ("success", "partial_success"):
@ -95,59 +123,24 @@ class DoclingServeConverter(DocumentConverter):
return DoclingDocument.model_validate(json_content) return DoclingDocument.model_validate(json_content)
async def _make_request(self, files: dict, name: str) -> "DoclingDocument": except httpx.ConnectError as e:
"""Make a request to docling-serve and return the DoclingDocument.
Args:
files: Dictionary with files parameter for requests
name: Name of the document being converted (for error messages)
Returns:
DoclingDocument representation
Raises:
ValueError: If conversion fails or service is unavailable
"""
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,
"table_mode": opts.table_mode,
"table_cell_matching": opts.table_cell_matching,
"images_scale": opts.images_scale,
}
if opts.ocr_lang:
data["ocr_lang"] = opts.ocr_lang
headers = {}
if self.api_key:
headers["X-Api-Key"] = self.api_key
return await asyncio.to_thread(
self._sync_make_request, files, name, data, headers
)
except requests.exceptions.ConnectionError as e:
raise ValueError( raise ValueError(
f"Could not connect to docling-serve at {self.base_url}. " f"Could not connect to docling-serve at {self.base_url}. "
f"Ensure the service is running and accessible. Error: {e}" f"Ensure the service is running and accessible. Error: {e}"
) )
except requests.exceptions.Timeout as e: except httpx.TimeoutException as e:
raise ValueError( raise ValueError(
f"Request to docling-serve timed out after {self.timeout}s. " f"Request to docling-serve timed out after {self.timeout}s. "
f"Consider increasing the timeout in configuration. Error: {e}" 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: if e.response.status_code == 401:
raise ValueError( raise ValueError(
"Authentication failed. Check your API key configuration." "Authentication failed. Check your API key configuration."
) )
raise ValueError(f"HTTP error from docling-serve: {e}") raise ValueError(f"HTTP error from docling-serve: {e}")
except ValueError:
raise
except Exception as e: except Exception as e:
raise ValueError(f"Failed to convert via docling-serve: {e}") raise ValueError(f"Failed to convert via docling-serve: {e}")
@ -175,11 +168,12 @@ class DoclingServeConverter(DocumentConverter):
except Exception as e: except Exception as e:
raise ValueError(f"Failed to read text file {path}: {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: 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) return await self._make_request(files, path.name)
async def convert_text( async def convert_text(
@ -199,8 +193,6 @@ class DoclingServeConverter(DocumentConverter):
Raises: Raises:
ValueError: If the text cannot be converted. ValueError: If the text cannot be converted.
""" """
from io import BytesIO
text_bytes = text.encode("utf-8") 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) return await self._make_request(files, name)

144
tests/test_converters.py
View file

@ -2,8 +2,9 @@
import tempfile import tempfile
from pathlib import Path from pathlib import Path
from unittest.mock import Mock, patch from unittest.mock import AsyncMock, Mock, patch
import httpx
import pytest import pytest
import requests import requests
from docling_core.types.doc.document import DoclingDocument from docling_core.types.doc.document import DoclingDocument
@ -201,8 +202,7 @@ class TestDoclingServeConverter:
assert ".md" in extensions assert ".md" in extensions
@pytest.mark.asyncio @pytest.mark.asyncio
@patch("haiku.rag.converters.docling_serve.requests.post") async def test_convert_text_success(self, converter):
async def test_convert_text_success(self, mock_post, converter):
"""Test successful text conversion via docling-serve.""" """Test successful text conversion via docling-serve."""
mock_response = Mock() mock_response = Mock()
mock_response.status_code = 200 mock_response.status_code = 200
@ -210,16 +210,22 @@ class TestDoclingServeConverter:
"status": "success", "status": "success",
"document": {"json_content": create_mock_docling_document_json("test")}, "document": {"json_content": create_mock_docling_document_json("test")},
} }
mock_post.return_value = mock_response mock_response.raise_for_status = Mock()
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") doc = await converter.convert_text("# Test", name="test.md")
assert isinstance(doc, DoclingDocument) assert isinstance(doc, DoclingDocument)
assert doc.version == "1.8.0" assert doc.version == "1.8.0"
mock_post.assert_called_once() mock_client.post.assert_called_once()
@pytest.mark.asyncio @pytest.mark.asyncio
@patch("haiku.rag.converters.docling_serve.requests.post") async def test_convert_text_with_api_key(self, config):
async def test_convert_text_with_api_key(self, mock_post, config):
"""Test that API key is included in request headers.""" """Test that API key is included in request headers."""
config.providers.docling_serve.api_key = "test-key" config.providers.docling_serve.api_key = "test-key"
converter = DoclingServeConverter(config) converter = DoclingServeConverter(config)
@ -230,17 +236,23 @@ class TestDoclingServeConverter:
"status": "success", "status": "success",
"document": {"json_content": create_mock_docling_document_json("test")}, "document": {"json_content": create_mock_docling_document_json("test")},
} }
mock_post.return_value = mock_response mock_response.raise_for_status = Mock()
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
await converter.convert_text("# Test") await converter.convert_text("# Test")
call_kwargs = mock_post.call_args.kwargs call_kwargs = mock_client.post.call_args.kwargs
assert "headers" in call_kwargs assert "headers" in call_kwargs
assert call_kwargs["headers"]["X-Api-Key"] == "test-key" assert call_kwargs["headers"]["X-Api-Key"] == "test-key"
@pytest.mark.asyncio @pytest.mark.asyncio
@patch("haiku.rag.converters.docling_serve.requests.post") async def test_conversion_options_passed_to_api(self, config):
async def test_conversion_options_passed_to_api(self, mock_post, config):
"""Test that conversion options are passed to docling-serve API.""" """Test that conversion options are passed to docling-serve API."""
config.processing.conversion_options.do_ocr = False config.processing.conversion_options.do_ocr = False
config.processing.conversion_options.force_ocr = True config.processing.conversion_options.force_ocr = True
@ -257,63 +269,78 @@ class TestDoclingServeConverter:
"status": "success", "status": "success",
"document": {"json_content": create_mock_docling_document_json("test")}, "document": {"json_content": create_mock_docling_document_json("test")},
} }
mock_post.return_value = mock_response mock_response.raise_for_status = Mock()
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
await converter.convert_text("# Test") await converter.convert_text("# Test")
call_kwargs = mock_post.call_args.kwargs call_kwargs = mock_client.post.call_args.kwargs
assert "data" in call_kwargs assert "data" in call_kwargs
data = call_kwargs["data"] data = call_kwargs["data"]
assert data["do_ocr"] is False assert data["do_ocr"] == "false"
assert data["force_ocr"] is True assert data["force_ocr"] == "true"
assert data["ocr_lang"] == ["en", "fr"] assert data["ocr_lang"] == ["en", "fr"]
assert "pdf_backend" not in data
assert data["table_mode"] == "fast" assert data["table_mode"] == "fast"
assert data["table_cell_matching"] is False assert data["table_cell_matching"] == "false"
assert data["do_table_structure"] is False assert data["do_table_structure"] == "false"
assert data["images_scale"] == 3.0 assert data["images_scale"] == "3.0"
@pytest.mark.asyncio @pytest.mark.asyncio
@patch("haiku.rag.converters.docling_serve.requests.post") async def test_convert_text_connection_error(self, converter):
async def test_convert_text_connection_error(self, mock_post, converter):
"""Test handling of connection errors.""" """Test handling of connection errors."""
import requests with patch("httpx.AsyncClient") as mock_client_class:
mock_client = AsyncMock()
mock_post.side_effect = requests.exceptions.ConnectionError("Connection failed") 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
with pytest.raises(ValueError, match="Could not connect to docling-serve"): with pytest.raises(ValueError, match="Could not connect to docling-serve"):
await converter.convert_text("# Test") await converter.convert_text("# Test")
@pytest.mark.asyncio @pytest.mark.asyncio
@patch("haiku.rag.converters.docling_serve.requests.post") async def test_convert_text_timeout_error(self, converter):
async def test_convert_text_timeout_error(self, mock_post, converter):
"""Test handling of timeout errors.""" """Test handling of timeout errors."""
import requests with patch("httpx.AsyncClient") as mock_client_class:
mock_client = AsyncMock()
mock_post.side_effect = requests.exceptions.Timeout("Timeout") 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
with pytest.raises(ValueError, match="timed out"): with pytest.raises(ValueError, match="timed out"):
await converter.convert_text("# Test") await converter.convert_text("# Test")
@pytest.mark.asyncio @pytest.mark.asyncio
@patch("haiku.rag.converters.docling_serve.requests.post") async def test_convert_text_auth_error(self, converter):
async def test_convert_text_auth_error(self, mock_post, converter):
"""Test handling of authentication errors.""" """Test handling of authentication errors."""
import requests
mock_response = Mock() mock_response = Mock()
mock_response.status_code = 401 mock_response.status_code = 401
mock_response.raise_for_status.side_effect = requests.exceptions.HTTPError(
response=mock_response 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_post.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="Authentication failed"): with pytest.raises(ValueError, match="Authentication failed"):
await converter.convert_text("# Test") await converter.convert_text("# Test")
@pytest.mark.asyncio @pytest.mark.asyncio
@patch("haiku.rag.converters.docling_serve.requests.post") async def test_convert_text_no_json_content(self, converter):
async def test_convert_text_no_json_content(self, mock_post, converter):
"""Test handling when docling-serve returns no JSON content.""" """Test handling when docling-serve returns no JSON content."""
mock_response = Mock() mock_response = Mock()
mock_response.status_code = 200 mock_response.status_code = 200
@ -321,14 +348,20 @@ class TestDoclingServeConverter:
"status": "success", "status": "success",
"document": {"json_content": None}, "document": {"json_content": None},
} }
mock_post.return_value = mock_response mock_response.raise_for_status = Mock()
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"): with pytest.raises(ValueError, match="did not return JSON content"):
await converter.convert_text("# Test") await converter.convert_text("# Test")
@pytest.mark.asyncio @pytest.mark.asyncio
@patch("haiku.rag.converters.docling_serve.requests.post") async def test_convert_file_pdf(self, converter):
async def test_convert_file_pdf(self, mock_post, converter):
"""Test converting PDF file via docling-serve.""" """Test converting PDF file via docling-serve."""
mock_response = Mock() mock_response = Mock()
mock_response.status_code = 200 mock_response.status_code = 200
@ -336,7 +369,14 @@ class TestDoclingServeConverter:
"status": "success", "status": "success",
"document": {"json_content": create_mock_docling_document_json("test")}, "document": {"json_content": create_mock_docling_document_json("test")},
} }
mock_post.return_value = mock_response mock_response.raise_for_status = Mock()
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 tempfile.NamedTemporaryFile(suffix=".pdf") as f: with tempfile.NamedTemporaryFile(suffix=".pdf") as f:
f.write(b"fake pdf content") f.write(b"fake pdf content")
@ -345,11 +385,10 @@ class TestDoclingServeConverter:
doc = await converter.convert_file(temp_path) doc = await converter.convert_file(temp_path)
assert isinstance(doc, DoclingDocument) assert isinstance(doc, DoclingDocument)
mock_post.assert_called_once() mock_client.post.assert_called_once()
@pytest.mark.asyncio @pytest.mark.asyncio
@patch("haiku.rag.converters.docling_serve.requests.post") async def test_convert_file_text(self, converter):
async def test_convert_file_text(self, mock_post, converter):
"""Test converting text file (reads locally, sends to docling-serve).""" """Test converting text file (reads locally, sends to docling-serve)."""
mock_response = Mock() mock_response = Mock()
mock_response.status_code = 200 mock_response.status_code = 200
@ -357,7 +396,14 @@ class TestDoclingServeConverter:
"status": "success", "status": "success",
"document": {"json_content": create_mock_docling_document_json("test")}, "document": {"json_content": create_mock_docling_document_json("test")},
} }
mock_post.return_value = mock_response mock_response.raise_for_status = Mock()
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 tempfile.NamedTemporaryFile(mode="w", suffix=".py") as f: with tempfile.NamedTemporaryFile(mode="w", suffix=".py") as f:
f.write("def hello():\n pass") f.write("def hello():\n pass")
@ -366,10 +412,8 @@ class TestDoclingServeConverter:
doc = await converter.convert_file(temp_path) doc = await converter.convert_file(temp_path)
assert isinstance(doc, DoclingDocument) assert isinstance(doc, DoclingDocument)
# Should call docling-serve for conversion mock_client.post.assert_called_once()
mock_post.assert_called_once() call_kwargs = mock_client.post.call_args.kwargs
# Check that code was wrapped in code block
call_kwargs = mock_post.call_args.kwargs
assert "files" in call_kwargs assert "files" in call_kwargs