106 lines
3.2 KiB
Python
106 lines
3.2 KiB
Python
"""Unit coverage for canonical text/code parsing and pinned Docling selection."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
|
|
from ragcore.domain.sources import SourceFingerprint
|
|
from ragcore.infrastructure.parsers import (
|
|
DoclingConversion,
|
|
DoclingParser,
|
|
PinnedDoclingError,
|
|
TextParser,
|
|
)
|
|
from ragcore.ingestion.parsing import (
|
|
CanonicalBlockKind,
|
|
ParseRequest,
|
|
ParserRouter,
|
|
UnsupportedParserError,
|
|
)
|
|
|
|
|
|
def parse_request(content: bytes, name: str, media_type: str) -> ParseRequest:
|
|
return ParseRequest(uuid4(), SourceFingerprint("a" * 64), content, media_type, name)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_markdown_preserves_headings_tables_code_and_locations() -> None:
|
|
parser = TextParser()
|
|
request = parse_request(
|
|
b"# Title\n\nIntro\n\n| A | B |\n|---|---|\n| 1 | 2 |\n\n```python\nprint(1)\n```\n",
|
|
"runbook.md",
|
|
"text/markdown",
|
|
)
|
|
|
|
artifact = await parser.parse(request)
|
|
|
|
assert [block.kind for block in artifact.blocks] == [
|
|
CanonicalBlockKind.HEADING,
|
|
CanonicalBlockKind.PARAGRAPH,
|
|
CanonicalBlockKind.TABLE,
|
|
CanonicalBlockKind.CODE,
|
|
]
|
|
assert artifact.blocks[0].source_location == {"line_start": 1, "line_end": 1}
|
|
assert artifact.blocks[2].content.startswith("| A | B |")
|
|
assert artifact.blocks[3].metadata["language"] == "python"
|
|
assert artifact.blocks[3].source_location == {"line_start": 9, "line_end": 11}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_router_uses_code_parser_before_text_fallback() -> None:
|
|
from ragcore.infrastructure.parsers import CodeParser
|
|
|
|
router = ParserRouter((CodeParser(), TextParser()))
|
|
artifact = await router.parse(
|
|
parse_request(b"def f():\n return 1\n", "main.py", "text/plain")
|
|
)
|
|
|
|
assert artifact.extraction_report.normalizer_profile == "code-v1"
|
|
assert artifact.blocks[0].kind is CanonicalBlockKind.CODE
|
|
with pytest.raises(UnsupportedParserError):
|
|
await ParserRouter((TextParser(),)).parse(
|
|
parse_request(b"%PDF-1.7", "file.pdf", "application/pdf")
|
|
)
|
|
|
|
|
|
class Backend:
|
|
version = "2.117.0"
|
|
|
|
def __init__(self) -> None:
|
|
self.ocr_values: list[bool] = []
|
|
|
|
async def convert(self, content: bytes, *, enable_ocr: bool) -> DoclingConversion:
|
|
self.ocr_values.append(enable_ocr)
|
|
return DoclingConversion(())
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_docling_ocr_is_selected_only_without_a_text_layer() -> None:
|
|
backend = Backend()
|
|
parser = DoclingParser(backend)
|
|
await parser.parse(parse_request(b"pdf", "scan.pdf", "application/pdf"))
|
|
request = parse_request(b"pdf", "scan.pdf", "application/pdf")
|
|
request = ParseRequest(
|
|
request.document_version_id,
|
|
request.source_fingerprint,
|
|
request.content,
|
|
request.media_type,
|
|
request.original_name,
|
|
ocr_required=True,
|
|
text_layer_present=False,
|
|
)
|
|
artifact = await parser.parse(request)
|
|
|
|
assert backend.ocr_values == [False, True]
|
|
assert artifact.extraction_report.ocr_used is True
|
|
|
|
|
|
def test_docling_minor_version_is_pinned() -> None:
|
|
class OldBackend(Backend):
|
|
version = "2.47.9"
|
|
|
|
with pytest.raises(PinnedDoclingError):
|
|
DoclingParser(OldBackend())
|