wehub-resource-sync
982 lines
30 KiB
Python
982 lines
30 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Tests for the MarkItDown integration."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
import logging
|
|
import sys
|
|
import types
|
|
import warnings
|
|
from dataclasses import dataclass
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
from starlette.responses import StreamingResponse
|
|
|
|
import omlx.server as server_module
|
|
from omlx.api.markitdown import (
|
|
MARKITDOWN_EMPTY_PDF_MESSAGE,
|
|
MARKITDOWN_MODEL_ID,
|
|
MarkItDownFile,
|
|
MarkItDownRequestError,
|
|
convert_file_to_markdown,
|
|
parse_file_part,
|
|
preprocess_markitdown_file_parts,
|
|
preprocess_markitdown_file_parts_async,
|
|
quiet_pdf_parser_loggers,
|
|
)
|
|
from omlx.api.markitdown_pdf_fallback import (
|
|
convert_pdf_with_ocr_engine,
|
|
resolve_pdf_ocr_model,
|
|
stream_pdf_with_ocr_engine,
|
|
)
|
|
from omlx.api.openai_models import ChatCompletionRequest, Message
|
|
from omlx.engine_pool import EngineEntry, EnginePool
|
|
from omlx.server import ServerState, app
|
|
from omlx.settings import GlobalSettings
|
|
|
|
|
|
def _data_uri(payload: bytes = b"doc", mime_type: str = "application/pdf") -> str:
|
|
encoded = base64.b64encode(payload).decode("ascii")
|
|
return f"data:{mime_type};base64,{encoded}"
|
|
|
|
|
|
def _file_part(
|
|
filename: str = "sample.pdf",
|
|
data: str | None = None,
|
|
mime_type: str = "application/pdf",
|
|
) -> dict:
|
|
return {
|
|
"type": "file",
|
|
"file": {
|
|
"filename": filename,
|
|
"mime_type": mime_type,
|
|
"file_data": data or _data_uri(mime_type=mime_type),
|
|
},
|
|
}
|
|
|
|
|
|
class _EmptyPool:
|
|
def get_status(self) -> dict:
|
|
return {
|
|
"final_ceiling": 0,
|
|
"current_model_memory": 0,
|
|
"model_count": 0,
|
|
"loaded_count": 0,
|
|
"models": [],
|
|
}
|
|
|
|
|
|
def _settings_with_markitdown_model() -> GlobalSettings:
|
|
settings = GlobalSettings()
|
|
settings.integrations.markitdown_expose_model = True
|
|
return settings
|
|
|
|
|
|
def test_openai_models_hides_markitdown_by_default():
|
|
state = ServerState()
|
|
state.engine_pool = _EmptyPool()
|
|
state.global_settings = GlobalSettings()
|
|
|
|
with patch("omlx.server._server_state", state):
|
|
client = TestClient(app, raise_server_exceptions=False)
|
|
response = client.get("/v1/models")
|
|
|
|
assert response.status_code == 200
|
|
ids = [m["id"] for m in response.json()["data"]]
|
|
assert MARKITDOWN_MODEL_ID not in ids
|
|
|
|
|
|
def test_openai_models_includes_markitdown_when_exposed():
|
|
state = ServerState()
|
|
state.engine_pool = _EmptyPool()
|
|
state.global_settings = _settings_with_markitdown_model()
|
|
|
|
with patch("omlx.server._server_state", state):
|
|
client = TestClient(app, raise_server_exceptions=False)
|
|
response = client.get("/v1/models")
|
|
|
|
assert response.status_code == 200
|
|
ids = [m["id"] for m in response.json()["data"]]
|
|
assert "MarkItDown" in ids
|
|
|
|
|
|
def test_openai_models_hides_markitdown_when_disabled():
|
|
state = ServerState()
|
|
state.engine_pool = _EmptyPool()
|
|
state.global_settings = _settings_with_markitdown_model()
|
|
state.global_settings.integrations.markitdown_enabled = False
|
|
|
|
with patch("omlx.server._server_state", state):
|
|
client = TestClient(app, raise_server_exceptions=False)
|
|
response = client.get("/v1/models")
|
|
|
|
assert response.status_code == 200
|
|
ids = [m["id"] for m in response.json()["data"]]
|
|
assert MARKITDOWN_MODEL_ID not in ids
|
|
|
|
|
|
def test_openai_models_hides_markitdown_when_not_exposed():
|
|
state = ServerState()
|
|
state.engine_pool = _EmptyPool()
|
|
state.global_settings = GlobalSettings()
|
|
state.global_settings.integrations.markitdown_expose_model = False
|
|
|
|
with patch("omlx.server._server_state", state):
|
|
client = TestClient(app, raise_server_exceptions=False)
|
|
response = client.get("/v1/models")
|
|
|
|
assert response.status_code == 200
|
|
ids = [m["id"] for m in response.json()["data"]]
|
|
assert MARKITDOWN_MODEL_ID not in ids
|
|
|
|
|
|
def test_markitdown_chat_completion_converts_file(monkeypatch):
|
|
state = ServerState()
|
|
state.engine_pool = _EmptyPool()
|
|
state.global_settings = _settings_with_markitdown_model()
|
|
|
|
def fake_convert(file: MarkItDownFile, **kwargs) -> str:
|
|
assert file.filename == "sample.pdf"
|
|
return "# Converted"
|
|
|
|
monkeypatch.setattr("omlx.api.markitdown.convert_file_to_markdown", fake_convert)
|
|
|
|
with patch("omlx.server._server_state", state):
|
|
client = TestClient(app, raise_server_exceptions=False)
|
|
response = client.post(
|
|
"/v1/chat/completions",
|
|
json={
|
|
"model": MARKITDOWN_MODEL_ID,
|
|
"messages": [{"role": "user", "content": [_file_part()]}],
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
content = response.json()["choices"][0]["message"]["content"]
|
|
assert "## Attached file: sample.pdf" in content
|
|
assert "# Converted" in content
|
|
|
|
|
|
def test_markitdown_chat_completion_uses_latest_user_turn(monkeypatch):
|
|
state = ServerState()
|
|
state.engine_pool = _EmptyPool()
|
|
state.global_settings = _settings_with_markitdown_model()
|
|
|
|
def fake_convert(file: MarkItDownFile, **kwargs) -> str:
|
|
return f"# Converted {file.filename}"
|
|
|
|
monkeypatch.setattr("omlx.api.markitdown.convert_file_to_markdown", fake_convert)
|
|
|
|
with patch("omlx.server._server_state", state):
|
|
client = TestClient(app, raise_server_exceptions=False)
|
|
response = client.post(
|
|
"/v1/chat/completions",
|
|
json={
|
|
"model": MARKITDOWN_MODEL_ID,
|
|
"messages": [
|
|
{"role": "user", "content": [_file_part("file1.pdf")]},
|
|
{
|
|
"role": "assistant",
|
|
"content": (
|
|
"## Attached file: file1.pdf\n\n# Converted file1.pdf"
|
|
),
|
|
},
|
|
{"role": "user", "content": [_file_part("file2.pdf")]},
|
|
],
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
content = response.json()["choices"][0]["message"]["content"]
|
|
assert "file2.pdf" in content
|
|
assert "file1.pdf" not in content
|
|
|
|
|
|
def test_markitdown_chat_completion_disabled_returns_404():
|
|
state = ServerState()
|
|
state.engine_pool = _EmptyPool()
|
|
state.global_settings = GlobalSettings()
|
|
state.global_settings.integrations.markitdown_enabled = False
|
|
|
|
with patch("omlx.server._server_state", state):
|
|
client = TestClient(app, raise_server_exceptions=False)
|
|
response = client.post(
|
|
"/v1/chat/completions",
|
|
json={
|
|
"model": MARKITDOWN_MODEL_ID,
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 404
|
|
assert "markitdown" in response.json()["error"]["message"].lower()
|
|
|
|
|
|
def test_markitdown_chat_completion_hidden_model_returns_404():
|
|
state = ServerState()
|
|
state.engine_pool = _EmptyPool()
|
|
state.global_settings = GlobalSettings()
|
|
state.global_settings.integrations.markitdown_expose_model = False
|
|
|
|
with patch("omlx.server._server_state", state):
|
|
client = TestClient(app, raise_server_exceptions=False)
|
|
response = client.post(
|
|
"/v1/chat/completions",
|
|
json={
|
|
"model": MARKITDOWN_MODEL_ID,
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 404
|
|
assert "markitdown" in response.json()["error"]["message"].lower()
|
|
|
|
|
|
def test_markitdown_stream_response_starts_before_conversion(monkeypatch):
|
|
state = ServerState()
|
|
state.engine_pool = _EmptyPool()
|
|
state.global_settings = _settings_with_markitdown_model()
|
|
started = False
|
|
|
|
async def fake_stream_messages(*args, **kwargs):
|
|
nonlocal started
|
|
started = True
|
|
yield "Converted markdown"
|
|
|
|
monkeypatch.setattr(
|
|
server_module,
|
|
"stream_messages_to_markdown_async",
|
|
fake_stream_messages,
|
|
)
|
|
request = ChatCompletionRequest(
|
|
model=MARKITDOWN_MODEL_ID,
|
|
messages=[{"role": "user", "content": "hello"}],
|
|
stream=True,
|
|
)
|
|
|
|
async def exercise():
|
|
with patch("omlx.server._server_state", state):
|
|
response = await server_module._create_markitdown_chat_completion(
|
|
request,
|
|
None,
|
|
)
|
|
|
|
assert isinstance(response, StreamingResponse)
|
|
assert started is False
|
|
|
|
iterator = response.body_iterator.__aiter__()
|
|
first = await iterator.__anext__()
|
|
assert first.startswith("data: ")
|
|
assert started is False
|
|
|
|
role_chunk = await iterator.__anext__()
|
|
assert '"role":"assistant"' in role_chunk
|
|
assert started is False
|
|
|
|
content_chunk = await iterator.__anext__()
|
|
assert "Converted markdown" in content_chunk
|
|
assert started is True
|
|
await iterator.aclose()
|
|
|
|
asyncio.run(exercise())
|
|
|
|
|
|
def test_markitdown_non_stream_response_starts_before_conversion(monkeypatch):
|
|
state = ServerState()
|
|
state.engine_pool = _EmptyPool()
|
|
state.global_settings = _settings_with_markitdown_model()
|
|
started = False
|
|
|
|
async def fake_convert_messages(*args, **kwargs):
|
|
nonlocal started
|
|
started = True
|
|
return "Converted markdown"
|
|
|
|
monkeypatch.setattr(
|
|
server_module,
|
|
"convert_messages_to_markdown_async",
|
|
fake_convert_messages,
|
|
)
|
|
request = ChatCompletionRequest(
|
|
model=MARKITDOWN_MODEL_ID,
|
|
messages=[{"role": "user", "content": "hello"}],
|
|
)
|
|
|
|
async def exercise():
|
|
with patch("omlx.server._server_state", state):
|
|
response = await server_module._create_markitdown_chat_completion(
|
|
request,
|
|
None,
|
|
)
|
|
|
|
assert isinstance(response, StreamingResponse)
|
|
assert started is False
|
|
|
|
iterator = response.body_iterator.__aiter__()
|
|
first = await iterator.__anext__()
|
|
assert first == " "
|
|
chunk = await iterator.__anext__()
|
|
assert "Converted markdown" in chunk
|
|
assert started is True
|
|
await iterator.aclose()
|
|
|
|
asyncio.run(exercise())
|
|
|
|
|
|
def test_preprocess_file_parts_for_llm(monkeypatch):
|
|
def fake_convert(file: MarkItDownFile, **kwargs) -> str:
|
|
return "Converted text"
|
|
|
|
monkeypatch.setattr("omlx.api.markitdown.convert_file_to_markdown", fake_convert)
|
|
messages = [
|
|
Message(
|
|
role="user",
|
|
content=[
|
|
{"type": "text", "text": "Summarize this."},
|
|
_file_part("paper.pdf"),
|
|
],
|
|
)
|
|
]
|
|
|
|
processed = preprocess_markitdown_file_parts(
|
|
messages,
|
|
global_settings=GlobalSettings(),
|
|
)
|
|
|
|
parts = processed[0].content
|
|
assert isinstance(parts, list)
|
|
assert parts[0].type == "text"
|
|
assert parts[1].type == "text"
|
|
assert "## Attached file: paper.pdf" in (parts[1].text or "")
|
|
assert "Converted text" in (parts[1].text or "")
|
|
|
|
|
|
def test_preprocess_file_parts_works_when_model_not_exposed(monkeypatch):
|
|
def fake_convert(file: MarkItDownFile, **kwargs) -> str:
|
|
return "Converted text"
|
|
|
|
monkeypatch.setattr("omlx.api.markitdown.convert_file_to_markdown", fake_convert)
|
|
settings = GlobalSettings()
|
|
settings.integrations.markitdown_expose_model = False
|
|
|
|
processed = preprocess_markitdown_file_parts(
|
|
[Message(role="user", content=[_file_part("paper.pdf")])],
|
|
global_settings=settings,
|
|
)
|
|
|
|
parts = processed[0].content
|
|
assert isinstance(parts, list)
|
|
assert parts[0].type == "text"
|
|
assert "Converted text" in (parts[0].text or "")
|
|
|
|
|
|
def test_text_and_markdown_file_parts_are_inlined_without_converter(monkeypatch):
|
|
called = False
|
|
|
|
def fake_convert(file: MarkItDownFile, **kwargs) -> str:
|
|
nonlocal called
|
|
called = True
|
|
raise AssertionError("plain text attachments should not use MarkItDown")
|
|
|
|
monkeypatch.setattr("omlx.api.markitdown.convert_file_to_markdown", fake_convert)
|
|
|
|
processed = preprocess_markitdown_file_parts(
|
|
[
|
|
Message(
|
|
role="user",
|
|
content=[
|
|
_file_part(
|
|
"notes.txt",
|
|
data=_data_uri(b"Plain notes", mime_type="text/plain"),
|
|
mime_type="text/plain",
|
|
),
|
|
_file_part(
|
|
"guide.md",
|
|
data=_data_uri(b"# Guide", mime_type="text/markdown"),
|
|
mime_type="text/markdown",
|
|
),
|
|
],
|
|
)
|
|
],
|
|
global_settings=GlobalSettings(),
|
|
)
|
|
|
|
parts = processed[0].content
|
|
assert isinstance(parts, list)
|
|
rendered = "\n".join(part.text or "" for part in parts)
|
|
assert "## Attached file: notes.txt" in rendered
|
|
assert "Plain notes" in rendered
|
|
assert "## Attached file: guide.md" in rendered
|
|
assert "# Guide" in rendered
|
|
assert called is False
|
|
|
|
|
|
def test_preprocess_file_parts_does_not_create_mixed_content_warning(monkeypatch):
|
|
def fake_convert(file: MarkItDownFile, **kwargs) -> str:
|
|
return "Converted text"
|
|
|
|
monkeypatch.setattr("omlx.api.markitdown.convert_file_to_markdown", fake_convert)
|
|
messages = [
|
|
Message(
|
|
role="user",
|
|
content=[
|
|
{"type": "text", "text": "Summarize this."},
|
|
_file_part("paper.pdf"),
|
|
],
|
|
)
|
|
]
|
|
|
|
processed = preprocess_markitdown_file_parts(
|
|
messages,
|
|
global_settings=GlobalSettings(),
|
|
)
|
|
|
|
with warnings.catch_warnings(record=True) as caught:
|
|
warnings.simplefilter("always")
|
|
processed[0].model_dump()
|
|
|
|
assert not [
|
|
warning
|
|
for warning in caught
|
|
if "Pydantic serializer warnings" in str(warning.message)
|
|
]
|
|
|
|
|
|
def test_preprocess_allows_missing_historical_file_parts(monkeypatch):
|
|
def fake_convert(file: MarkItDownFile, **kwargs) -> str:
|
|
raise AssertionError("missing historical files should not be converted")
|
|
|
|
monkeypatch.setattr("omlx.api.markitdown.convert_file_to_markdown", fake_convert)
|
|
|
|
processed = preprocess_markitdown_file_parts(
|
|
[
|
|
Message(
|
|
role="user",
|
|
content=[
|
|
{"type": "text", "text": "Summarize this."},
|
|
{
|
|
"type": "file",
|
|
"file": {
|
|
"filename": "paper.pdf",
|
|
"mime_type": "application/pdf",
|
|
"data": "",
|
|
},
|
|
},
|
|
],
|
|
),
|
|
Message(role="assistant", content="Done."),
|
|
Message(role="user", content="Continue."),
|
|
],
|
|
global_settings=GlobalSettings(),
|
|
allow_missing_historical_files=True,
|
|
)
|
|
|
|
parts = processed[0].content
|
|
assert isinstance(parts, list)
|
|
assert parts[0].type == "text"
|
|
assert parts[1].type == "text"
|
|
assert "Attached file unavailable: paper.pdf" in (parts[1].text or "")
|
|
assert processed[2].content == "Continue."
|
|
|
|
|
|
def test_preprocess_still_rejects_missing_latest_file_part():
|
|
with pytest.raises(MarkItDownRequestError, match="file_data"):
|
|
preprocess_markitdown_file_parts(
|
|
[
|
|
Message(
|
|
role="user",
|
|
content=[
|
|
{
|
|
"type": "file",
|
|
"file": {
|
|
"filename": "paper.pdf",
|
|
"mime_type": "application/pdf",
|
|
"data": "",
|
|
},
|
|
}
|
|
],
|
|
)
|
|
],
|
|
global_settings=GlobalSettings(),
|
|
allow_missing_historical_files=True,
|
|
)
|
|
|
|
|
|
def test_async_preprocess_allows_missing_historical_file_parts(monkeypatch):
|
|
def fake_convert(file: MarkItDownFile, **kwargs) -> str:
|
|
raise AssertionError("missing historical files should not be converted")
|
|
|
|
monkeypatch.setattr("omlx.api.markitdown.convert_file_to_markdown", fake_convert)
|
|
|
|
async def exercise():
|
|
return await preprocess_markitdown_file_parts_async(
|
|
[
|
|
Message(
|
|
role="user",
|
|
content=[
|
|
{
|
|
"type": "file",
|
|
"file": {
|
|
"filename": "paper.pdf",
|
|
"mime_type": "application/pdf",
|
|
"data": "",
|
|
},
|
|
}
|
|
],
|
|
),
|
|
Message(role="assistant", content="Done."),
|
|
Message(role="user", content="Continue."),
|
|
],
|
|
global_settings=GlobalSettings(),
|
|
allow_missing_historical_files=True,
|
|
)
|
|
|
|
processed = asyncio.run(exercise())
|
|
|
|
parts = processed[0].content
|
|
assert isinstance(parts, list)
|
|
assert parts[0].type == "text"
|
|
assert "Attached file unavailable: paper.pdf" in (parts[0].text or "")
|
|
|
|
|
|
def test_server_llm_preprocess_allows_stored_document_placeholders(monkeypatch):
|
|
def fake_convert(file: MarkItDownFile, **kwargs) -> str:
|
|
raise AssertionError("stored document placeholders should not be converted")
|
|
|
|
monkeypatch.setattr("omlx.api.markitdown.convert_file_to_markdown", fake_convert)
|
|
|
|
state = ServerState()
|
|
state.engine_pool = _EmptyPool()
|
|
state.global_settings = GlobalSettings()
|
|
request = ChatCompletionRequest(
|
|
model="test-model",
|
|
messages=[
|
|
Message(
|
|
role="user",
|
|
content=[
|
|
{
|
|
"type": "file",
|
|
"file": {
|
|
"filename": "paper.pdf",
|
|
"mime_type": "application/pdf",
|
|
"data": "",
|
|
},
|
|
}
|
|
],
|
|
),
|
|
Message(role="assistant", content="Done."),
|
|
Message(role="user", content="Continue."),
|
|
],
|
|
)
|
|
|
|
async def exercise():
|
|
with patch("omlx.server._server_state", state):
|
|
return await server_module._preprocess_markitdown_files_for_llm(request)
|
|
|
|
processed = asyncio.run(exercise())
|
|
|
|
parts = processed.messages[0].content
|
|
assert isinstance(parts, list)
|
|
assert parts[0].type == "text"
|
|
assert "Attached file unavailable: paper.pdf" in (parts[0].text or "")
|
|
|
|
|
|
def test_preprocess_file_parts_rejects_when_disabled():
|
|
settings = GlobalSettings()
|
|
settings.integrations.markitdown_enabled = False
|
|
|
|
with pytest.raises(MarkItDownRequestError, match="disabled"):
|
|
preprocess_markitdown_file_parts(
|
|
[Message(role="user", content=[_file_part()])],
|
|
global_settings=settings,
|
|
)
|
|
|
|
|
|
def test_xlsx_is_rejected_without_pandas_dependency():
|
|
part = {
|
|
"type": "file",
|
|
"file": {
|
|
"filename": "sheet.xlsx",
|
|
"mime_type": (
|
|
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
|
|
),
|
|
"file_data": _data_uri(),
|
|
},
|
|
}
|
|
|
|
with pytest.raises(MarkItDownRequestError, match="Spreadsheet"):
|
|
parse_file_part(part, max_file_size_mb=25)
|
|
|
|
|
|
def test_openai_file_data_without_filename_infers_from_data_uri():
|
|
parsed = parse_file_part(
|
|
{
|
|
"type": "file",
|
|
"file": {
|
|
"file_data": _data_uri(
|
|
b"Plain notes",
|
|
mime_type="text/plain",
|
|
),
|
|
},
|
|
},
|
|
max_file_size_mb=25,
|
|
)
|
|
|
|
assert parsed.filename == "attachment.txt"
|
|
assert parsed.mime_type == "text/plain"
|
|
assert parsed.data == b"Plain notes"
|
|
|
|
|
|
def test_file_id_is_rejected():
|
|
with pytest.raises(MarkItDownRequestError, match="file_id"):
|
|
parse_file_part(
|
|
{"type": "file", "file": {"file_id": "file_123", "filename": "x.pdf"}},
|
|
max_file_size_mb=25,
|
|
)
|
|
|
|
|
|
def test_empty_pdf_conversion_logs_warning(monkeypatch, caplog):
|
|
@dataclass(frozen=True)
|
|
class FakeStreamInfo:
|
|
extension: str | None = None
|
|
mimetype: str | None = None
|
|
filename: str | None = None
|
|
|
|
class FakeResult:
|
|
markdown = ""
|
|
|
|
class FakeConverter:
|
|
def convert_stream(self, stream, stream_info=None):
|
|
return FakeResult()
|
|
|
|
fake_markitdown = types.ModuleType("markitdown")
|
|
fake_markitdown.StreamInfo = FakeStreamInfo
|
|
monkeypatch.setitem(sys.modules, "markitdown", fake_markitdown)
|
|
monkeypatch.setattr("omlx.api.markitdown._converter", FakeConverter())
|
|
|
|
caplog.set_level("WARNING")
|
|
with pytest.raises(MarkItDownRequestError) as exc_info:
|
|
convert_file_to_markdown(
|
|
MarkItDownFile(
|
|
filename="scan.pdf",
|
|
mime_type="application/pdf",
|
|
data=b"%PDF",
|
|
)
|
|
)
|
|
|
|
assert exc_info.value.detail == MARKITDOWN_EMPTY_PDF_MESSAGE
|
|
assert "no extractable text" in caplog.text.lower()
|
|
|
|
|
|
def test_pdf_parser_debug_loggers_are_quieted():
|
|
logger_names = ("pdfminer", "pdfminer.psparser", "pdfminer.pdfinterp")
|
|
original_levels = {name: logging.getLogger(name).level for name in logger_names}
|
|
try:
|
|
for name in logger_names:
|
|
logging.getLogger(name).setLevel(logging.DEBUG)
|
|
|
|
quiet_pdf_parser_loggers()
|
|
|
|
for name in logger_names:
|
|
assert logging.getLogger(name).level == logging.WARNING
|
|
finally:
|
|
for name, level in original_levels.items():
|
|
logging.getLogger(name).setLevel(level)
|
|
|
|
|
|
def test_ocr_pdf_engine_converts_pages_in_order_limits_concurrency_and_unloads(
|
|
monkeypatch,
|
|
):
|
|
class FakePool:
|
|
def __init__(self):
|
|
self.engine = FakeEngine()
|
|
self.unloaded = []
|
|
|
|
def resolve_model_id(self, model_id, settings_manager):
|
|
return model_id
|
|
|
|
def get_entry(self, model_id):
|
|
return types.SimpleNamespace(
|
|
config_model_type="dots_ocr",
|
|
engine_type="vlm",
|
|
)
|
|
|
|
def acquire(self, model_id):
|
|
assert model_id == "OCR-Model"
|
|
pool = self
|
|
|
|
class Lease:
|
|
async def __aenter__(self):
|
|
return pool.engine
|
|
|
|
async def __aexit__(self, exc_type, exc, tb):
|
|
return None
|
|
|
|
return Lease()
|
|
|
|
async def unload_if_idle_unpinned(self, model_id):
|
|
self.unloaded.append(model_id)
|
|
return True
|
|
|
|
active = 0
|
|
max_active = 0
|
|
|
|
class FakeEngine:
|
|
async def chat(self, messages, **kwargs):
|
|
nonlocal active, max_active
|
|
active += 1
|
|
max_active = max(max_active, active)
|
|
try:
|
|
await asyncio.sleep(0.01)
|
|
url = messages[0]["content"][0]["image_url"]["url"]
|
|
return types.SimpleNamespace(text=f"markdown for {url[-1]}")
|
|
finally:
|
|
active -= 1
|
|
|
|
settings = GlobalSettings()
|
|
settings.scheduler.max_concurrent_requests = 2
|
|
pool = FakePool()
|
|
|
|
monkeypatch.setattr(
|
|
"omlx.api.markitdown_pdf_fallback.render_pdf_pages_to_image_data_uris",
|
|
lambda file: ["page1", "page2", "page3"],
|
|
)
|
|
|
|
markdown = asyncio.run(
|
|
convert_pdf_with_ocr_engine(
|
|
MarkItDownFile(
|
|
filename="scan.pdf",
|
|
mime_type="application/pdf",
|
|
data=b"%PDF",
|
|
),
|
|
engine_model_id="OCR-Model",
|
|
engine_pool=pool,
|
|
settings_manager=None,
|
|
global_settings=settings,
|
|
get_sampling_params=lambda *_args: (
|
|
0.0,
|
|
1.0,
|
|
0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
4096,
|
|
0.0,
|
|
0.1,
|
|
),
|
|
)
|
|
)
|
|
|
|
assert "### Page 1\n\nmarkdown for 1" in markdown
|
|
assert "### Page 2\n\nmarkdown for 2" in markdown
|
|
assert "### Page 3\n\nmarkdown for 3" in markdown
|
|
assert max_active == 2
|
|
assert pool.unloaded == ["OCR-Model"]
|
|
|
|
|
|
def test_ocr_pdf_engine_streams_ready_prefix_in_page_order(monkeypatch):
|
|
class FakePool:
|
|
def __init__(self):
|
|
self.engine = FakeEngine()
|
|
|
|
def resolve_model_id(self, model_id, settings_manager):
|
|
return model_id
|
|
|
|
def get_entry(self, model_id):
|
|
return types.SimpleNamespace(
|
|
config_model_type="dots_ocr",
|
|
engine_type="vlm",
|
|
)
|
|
|
|
def acquire(self, model_id):
|
|
pool = self
|
|
|
|
class Lease:
|
|
async def __aenter__(self):
|
|
return pool.engine
|
|
|
|
async def __aexit__(self, exc_type, exc, tb):
|
|
return None
|
|
|
|
return Lease()
|
|
|
|
async def unload_if_idle_unpinned(self, model_id):
|
|
return True
|
|
|
|
class FakeEngine:
|
|
async def chat(self, messages, **kwargs):
|
|
url = messages[0]["content"][0]["image_url"]["url"]
|
|
delay_by_page = {"page1": 0.03, "page2": 0.0, "page3": 0.01}
|
|
await asyncio.sleep(delay_by_page[url])
|
|
return types.SimpleNamespace(text=f"markdown for {url[-1]}")
|
|
|
|
settings = GlobalSettings()
|
|
settings.scheduler.max_concurrent_requests = 3
|
|
monkeypatch.setattr(
|
|
"omlx.api.markitdown_pdf_fallback.render_pdf_pages_to_image_data_uris",
|
|
lambda file: ["page1", "page2", "page3"],
|
|
)
|
|
|
|
async def collect_chunks():
|
|
chunks = []
|
|
async for chunk in stream_pdf_with_ocr_engine(
|
|
MarkItDownFile(
|
|
filename="scan.pdf",
|
|
mime_type="application/pdf",
|
|
data=b"%PDF",
|
|
),
|
|
engine_model_id="OCR-Model",
|
|
engine_pool=FakePool(),
|
|
settings_manager=None,
|
|
global_settings=settings,
|
|
get_sampling_params=lambda *_args: (
|
|
0.0,
|
|
1.0,
|
|
0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
4096,
|
|
0.0,
|
|
0.1,
|
|
),
|
|
):
|
|
chunks.append(chunk)
|
|
return chunks
|
|
|
|
chunks = asyncio.run(collect_chunks())
|
|
|
|
assert [chunk.split("\n", 1)[0] for chunk in chunks] == [
|
|
"### Page 1",
|
|
"### Page 2",
|
|
"### Page 3",
|
|
]
|
|
|
|
|
|
def test_ocr_pdf_engine_rejects_non_ocr_config_model_type():
|
|
class FakePool:
|
|
def resolve_model_id(self, model_id, settings_manager):
|
|
return model_id
|
|
|
|
def get_entry(self, model_id):
|
|
return types.SimpleNamespace(
|
|
config_model_type="qwen3_vl",
|
|
engine_type="vlm",
|
|
)
|
|
|
|
with pytest.raises(MarkItDownRequestError, match="OCR"):
|
|
resolve_pdf_ocr_model(
|
|
"plain-model",
|
|
engine_pool=FakePool(),
|
|
settings_manager=None,
|
|
)
|
|
|
|
|
|
def test_ocr_pdf_engine_ignores_model_id_ocr_without_ocr_config_model_type():
|
|
class FakePool:
|
|
def resolve_model_id(self, model_id, settings_manager):
|
|
return model_id
|
|
|
|
def get_entry(self, model_id):
|
|
return types.SimpleNamespace(
|
|
model_id="Invoice-OCR",
|
|
source_repo_id=None,
|
|
config_model_type="qwen2_vl",
|
|
engine_type="vlm",
|
|
)
|
|
|
|
with pytest.raises(MarkItDownRequestError, match="config model_type"):
|
|
resolve_pdf_ocr_model(
|
|
"Invoice-OCR",
|
|
engine_pool=FakePool(),
|
|
settings_manager=None,
|
|
)
|
|
|
|
|
|
def test_async_preprocess_uses_ocr_pdf_processing_engine(monkeypatch):
|
|
class FakePool:
|
|
def resolve_model_id(self, model_id, settings_manager):
|
|
return model_id
|
|
|
|
def get_entry(self, model_id):
|
|
return types.SimpleNamespace(
|
|
config_model_type="dots_ocr",
|
|
engine_type="vlm",
|
|
)
|
|
|
|
def acquire(self, model_id):
|
|
class Lease:
|
|
async def __aenter__(self):
|
|
return types.SimpleNamespace(
|
|
chat=lambda **kwargs: None,
|
|
)
|
|
|
|
async def __aexit__(self, exc_type, exc, tb):
|
|
return None
|
|
|
|
return Lease()
|
|
|
|
async def fake_convert(file, **kwargs):
|
|
assert file.filename == "paper.pdf"
|
|
assert kwargs["engine_model_id"] == "OCR-Model"
|
|
return "OCR markdown"
|
|
|
|
settings = GlobalSettings()
|
|
settings.integrations.markitdown_pdf_processing_engine = "OCR-Model"
|
|
monkeypatch.setattr(
|
|
"omlx.api.markitdown_pdf_fallback.convert_pdf_with_ocr_engine",
|
|
fake_convert,
|
|
)
|
|
|
|
processed = asyncio.run(
|
|
preprocess_markitdown_file_parts_async(
|
|
[Message(role="user", content=[_file_part("paper.pdf")])],
|
|
global_settings=settings,
|
|
engine_pool=FakePool(),
|
|
settings_manager=None,
|
|
get_sampling_params=lambda *_args: None,
|
|
)
|
|
)
|
|
|
|
parts = processed[0].content
|
|
assert isinstance(parts, list)
|
|
assert "OCR markdown" in (parts[0].text or "")
|
|
|
|
|
|
def test_unload_if_idle_unpinned_skips_pinned_and_unloads_idle(monkeypatch):
|
|
class FakeEngine:
|
|
def has_active_requests(self):
|
|
return False
|
|
|
|
pool = EnginePool()
|
|
pool._entries["OCR-Model"] = EngineEntry(
|
|
model_id="OCR-Model",
|
|
model_path="/models/OCR-Model",
|
|
model_type="vlm",
|
|
engine_type="vlm",
|
|
estimated_size=1,
|
|
config_model_type="dots_ocr",
|
|
is_pinned=True,
|
|
engine=FakeEngine(),
|
|
)
|
|
|
|
unloaded = []
|
|
|
|
async def fake_unload(model_id):
|
|
unloaded.append(model_id)
|
|
|
|
monkeypatch.setattr(pool, "_unload_engine", fake_unload)
|
|
|
|
assert asyncio.run(pool.unload_if_idle_unpinned("OCR-Model")) is False
|
|
assert unloaded == []
|
|
|
|
pool._entries["OCR-Model"].is_pinned = False
|
|
assert asyncio.run(pool.unload_if_idle_unpinned("OCR-Model")) is True
|
|
assert unloaded == ["OCR-Model"]
|