Compare commits
No commits in common. "acb2d55ce35dc49fc843623381a90620182caafd" and "632acdcd2692de4074cac79be6210bf8714d9c9e" have entirely different histories.
acb2d55ce3
...
632acdcd26
2 changed files with 0 additions and 290 deletions
|
|
@ -1,91 +0,0 @@
|
||||||
"""OCRStep — run OCR, inject page tags, build the SegmentIndex (spec §6.2).
|
|
||||||
|
|
||||||
Runs after :class:`~ix.pipeline.setup_step.SetupStep`. Three things
|
|
||||||
happen:
|
|
||||||
|
|
||||||
1. Dispatch the flat page list to the injected :class:`OCRClient` and
|
|
||||||
write the raw :class:`~ix.contracts.OCRResult` onto the response.
|
|
||||||
2. Inject ``<page file="..." number="...">`` / ``</page>`` tag lines
|
|
||||||
around each page's content so the GenAIStep can ground citations
|
|
||||||
visually (spec §6.2c).
|
|
||||||
3. When provenance is on, build a :class:`SegmentIndex` over the
|
|
||||||
*non-tag* lines and stash it on the internal context.
|
|
||||||
|
|
||||||
Validation follows the triad-or-use_ocr rule from the spec: if any of
|
|
||||||
``include_geometries`` / ``include_ocr_text`` / ``ocr_only`` is set but
|
|
||||||
``context.files`` is empty, raise ``IX_000_004``. If OCR isn't wanted
|
|
||||||
(text-only request), return ``False`` to skip the step silently.
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
from ix.contracts import Line, RequestIX, ResponseIX
|
|
||||||
from ix.errors import IXErrorCode, IXException
|
|
||||||
from ix.ocr.client import OCRClient
|
|
||||||
from ix.pipeline.step import Step
|
|
||||||
from ix.segmentation import PageMetadata, SegmentIndex
|
|
||||||
|
|
||||||
|
|
||||||
class OCRStep(Step):
|
|
||||||
"""Inject-and-index OCR stage."""
|
|
||||||
|
|
||||||
def __init__(self, ocr_client: OCRClient) -> None:
|
|
||||||
self._client = ocr_client
|
|
||||||
|
|
||||||
async def validate(self, request_ix: RequestIX, response_ix: ResponseIX) -> bool:
|
|
||||||
opts = request_ix.options.ocr
|
|
||||||
ctx = response_ix.context
|
|
||||||
files = list(getattr(ctx, "files", [])) if ctx is not None else []
|
|
||||||
|
|
||||||
ocr_artifacts_requested = (
|
|
||||||
opts.include_geometries or opts.include_ocr_text or opts.ocr_only
|
|
||||||
)
|
|
||||||
if ocr_artifacts_requested and not files:
|
|
||||||
raise IXException(IXErrorCode.IX_000_004)
|
|
||||||
|
|
||||||
if not files:
|
|
||||||
return False
|
|
||||||
|
|
||||||
# OCR runs if use_ocr OR any of the artifact flags is set.
|
|
||||||
return bool(opts.use_ocr or ocr_artifacts_requested)
|
|
||||||
|
|
||||||
async def process(
|
|
||||||
self, request_ix: RequestIX, response_ix: ResponseIX
|
|
||||||
) -> ResponseIX:
|
|
||||||
ctx = response_ix.context
|
|
||||||
assert ctx is not None, "SetupStep must populate response_ix.context"
|
|
||||||
|
|
||||||
pages = list(getattr(ctx, "pages", []))
|
|
||||||
ocr_result = await self._client.ocr(pages)
|
|
||||||
|
|
||||||
# Inject page tags around each OCR page's content so the LLM can
|
|
||||||
# cross-reference the visual anchor without a separate prompt hack.
|
|
||||||
page_metadata: list[PageMetadata] = list(
|
|
||||||
getattr(ctx, "page_metadata", []) or []
|
|
||||||
)
|
|
||||||
for idx, ocr_page in enumerate(ocr_result.result.pages):
|
|
||||||
meta = page_metadata[idx] if idx < len(page_metadata) else PageMetadata()
|
|
||||||
file_idx = meta.file_index if meta.file_index is not None else 0
|
|
||||||
open_tag = Line(
|
|
||||||
text=f'<page file="{file_idx}" number="{ocr_page.page_no}">',
|
|
||||||
bounding_box=[],
|
|
||||||
)
|
|
||||||
close_tag = Line(text="</page>", bounding_box=[])
|
|
||||||
ocr_page.lines = [open_tag, *ocr_page.lines, close_tag]
|
|
||||||
|
|
||||||
response_ix.ocr_result = ocr_result
|
|
||||||
|
|
||||||
# Build SegmentIndex only when provenance is on. Segment IDs
|
|
||||||
# deliberately skip page-tag lines (see SegmentIndex.build).
|
|
||||||
if request_ix.options.provenance.include_provenance:
|
|
||||||
seg_idx = SegmentIndex.build(
|
|
||||||
ocr_result=ocr_result,
|
|
||||||
granularity="line",
|
|
||||||
pages_metadata=page_metadata,
|
|
||||||
)
|
|
||||||
ctx.segment_index = seg_idx
|
|
||||||
|
|
||||||
return response_ix
|
|
||||||
|
|
||||||
|
|
||||||
__all__ = ["OCRStep"]
|
|
||||||
|
|
@ -1,199 +0,0 @@
|
||||||
"""Tests for :class:`ix.pipeline.ocr_step.OCRStep` (spec §6.2)."""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import pytest
|
|
||||||
|
|
||||||
from ix.contracts import (
|
|
||||||
Context,
|
|
||||||
Line,
|
|
||||||
OCRDetails,
|
|
||||||
OCROptions,
|
|
||||||
OCRResult,
|
|
||||||
Options,
|
|
||||||
Page,
|
|
||||||
ProvenanceOptions,
|
|
||||||
RequestIX,
|
|
||||||
ResponseIX,
|
|
||||||
)
|
|
||||||
from ix.contracts.response import _InternalContext
|
|
||||||
from ix.errors import IXErrorCode, IXException
|
|
||||||
from ix.ocr import FakeOCRClient
|
|
||||||
from ix.pipeline.ocr_step import OCRStep
|
|
||||||
from ix.segmentation import PageMetadata, SegmentIndex
|
|
||||||
|
|
||||||
|
|
||||||
def _make_request(
|
|
||||||
*,
|
|
||||||
use_ocr: bool = True,
|
|
||||||
include_geometries: bool = False,
|
|
||||||
include_ocr_text: bool = False,
|
|
||||||
ocr_only: bool = False,
|
|
||||||
include_provenance: bool = True,
|
|
||||||
files: list | None = None,
|
|
||||||
texts: list[str] | None = None,
|
|
||||||
) -> RequestIX:
|
|
||||||
return RequestIX(
|
|
||||||
use_case="bank_statement_header",
|
|
||||||
ix_client_id="test",
|
|
||||||
request_id="r-1",
|
|
||||||
context=Context(files=files if files is not None else [], texts=texts or []),
|
|
||||||
options=Options(
|
|
||||||
ocr=OCROptions(
|
|
||||||
use_ocr=use_ocr,
|
|
||||||
include_geometries=include_geometries,
|
|
||||||
include_ocr_text=include_ocr_text,
|
|
||||||
ocr_only=ocr_only,
|
|
||||||
),
|
|
||||||
provenance=ProvenanceOptions(include_provenance=include_provenance),
|
|
||||||
),
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def _response_with_context(
|
|
||||||
*,
|
|
||||||
pages: list[Page] | None = None,
|
|
||||||
files: list | None = None,
|
|
||||||
texts: list[str] | None = None,
|
|
||||||
page_metadata: list[PageMetadata] | None = None,
|
|
||||||
) -> ResponseIX:
|
|
||||||
resp = ResponseIX()
|
|
||||||
resp.context = _InternalContext(
|
|
||||||
pages=pages or [],
|
|
||||||
files=files or [],
|
|
||||||
texts=texts or [],
|
|
||||||
page_metadata=page_metadata or [],
|
|
||||||
)
|
|
||||||
return resp
|
|
||||||
|
|
||||||
|
|
||||||
def _canned_ocr(pages: int = 1) -> OCRResult:
|
|
||||||
return OCRResult(
|
|
||||||
result=OCRDetails(
|
|
||||||
text="\n".join(f"text p{i+1}" for i in range(pages)),
|
|
||||||
pages=[
|
|
||||||
Page(
|
|
||||||
page_no=i + 1,
|
|
||||||
width=100.0,
|
|
||||||
height=200.0,
|
|
||||||
lines=[
|
|
||||||
Line(
|
|
||||||
text=f"line-content p{i+1}",
|
|
||||||
bounding_box=[0, 0, 10, 0, 10, 5, 0, 5],
|
|
||||||
)
|
|
||||||
],
|
|
||||||
)
|
|
||||||
for i in range(pages)
|
|
||||||
],
|
|
||||||
)
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
class TestValidate:
|
|
||||||
async def test_ocr_only_without_files_raises_IX_000_004(self) -> None:
|
|
||||||
step = OCRStep(ocr_client=FakeOCRClient(canned=_canned_ocr()))
|
|
||||||
req = _make_request(ocr_only=True, files=[], texts=["hi"])
|
|
||||||
resp = _response_with_context(files=[])
|
|
||||||
with pytest.raises(IXException) as ei:
|
|
||||||
await step.validate(req, resp)
|
|
||||||
assert ei.value.code is IXErrorCode.IX_000_004
|
|
||||||
|
|
||||||
async def test_include_ocr_text_without_files_raises_IX_000_004(self) -> None:
|
|
||||||
step = OCRStep(ocr_client=FakeOCRClient(canned=_canned_ocr()))
|
|
||||||
req = _make_request(include_ocr_text=True, files=[], texts=["hi"])
|
|
||||||
resp = _response_with_context(files=[])
|
|
||||||
with pytest.raises(IXException) as ei:
|
|
||||||
await step.validate(req, resp)
|
|
||||||
assert ei.value.code is IXErrorCode.IX_000_004
|
|
||||||
|
|
||||||
async def test_include_geometries_without_files_raises_IX_000_004(self) -> None:
|
|
||||||
step = OCRStep(ocr_client=FakeOCRClient(canned=_canned_ocr()))
|
|
||||||
req = _make_request(include_geometries=True, files=[], texts=["hi"])
|
|
||||||
resp = _response_with_context(files=[])
|
|
||||||
with pytest.raises(IXException) as ei:
|
|
||||||
await step.validate(req, resp)
|
|
||||||
assert ei.value.code is IXErrorCode.IX_000_004
|
|
||||||
|
|
||||||
async def test_text_only_skips_step(self) -> None:
|
|
||||||
step = OCRStep(ocr_client=FakeOCRClient(canned=_canned_ocr()))
|
|
||||||
req = _make_request(use_ocr=True, files=[], texts=["hi"])
|
|
||||||
resp = _response_with_context(files=[], texts=["hi"])
|
|
||||||
assert await step.validate(req, resp) is False
|
|
||||||
|
|
||||||
async def test_ocr_runs_when_files_and_use_ocr(self) -> None:
|
|
||||||
step = OCRStep(ocr_client=FakeOCRClient(canned=_canned_ocr()))
|
|
||||||
req = _make_request(use_ocr=True, files=["http://x"])
|
|
||||||
resp = _response_with_context(files=[("/tmp/x.pdf", "application/pdf")])
|
|
||||||
assert await step.validate(req, resp) is True
|
|
||||||
|
|
||||||
|
|
||||||
class TestProcess:
|
|
||||||
async def test_ocr_result_written_to_response(self) -> None:
|
|
||||||
canned = _canned_ocr(pages=2)
|
|
||||||
step = OCRStep(ocr_client=FakeOCRClient(canned=canned))
|
|
||||||
req = _make_request(use_ocr=True, files=["http://x"])
|
|
||||||
resp = _response_with_context(
|
|
||||||
pages=[Page(page_no=1, width=100.0, height=200.0, lines=[])] * 2,
|
|
||||||
files=[("/tmp/x.pdf", "application/pdf")],
|
|
||||||
page_metadata=[PageMetadata(file_index=0), PageMetadata(file_index=0)],
|
|
||||||
)
|
|
||||||
resp = await step.process(req, resp)
|
|
||||||
# Full OCR result written.
|
|
||||||
assert resp.ocr_result.result.text == canned.result.text
|
|
||||||
# Page tags injected: prepend + append around lines per page.
|
|
||||||
pages = resp.ocr_result.result.pages
|
|
||||||
assert len(pages) == 2
|
|
||||||
# First line is the opening <page> tag.
|
|
||||||
assert pages[0].lines[0].text is not None
|
|
||||||
assert pages[0].lines[0].text.startswith("<page ")
|
|
||||||
# Last line is the closing </page> tag.
|
|
||||||
assert pages[0].lines[-1].text == "</page>"
|
|
||||||
|
|
||||||
async def test_segment_index_built_when_provenance_on(self) -> None:
|
|
||||||
canned = _canned_ocr(pages=1)
|
|
||||||
step = OCRStep(ocr_client=FakeOCRClient(canned=canned))
|
|
||||||
req = _make_request(
|
|
||||||
use_ocr=True, include_provenance=True, files=["http://x"]
|
|
||||||
)
|
|
||||||
resp = _response_with_context(
|
|
||||||
pages=[Page(page_no=1, width=100.0, height=200.0, lines=[])],
|
|
||||||
files=[("/tmp/x.pdf", "application/pdf")],
|
|
||||||
page_metadata=[PageMetadata(file_index=0)],
|
|
||||||
)
|
|
||||||
resp = await step.process(req, resp)
|
|
||||||
seg_idx = resp.context.segment_index # type: ignore[union-attr]
|
|
||||||
assert isinstance(seg_idx, SegmentIndex)
|
|
||||||
# Page-tag lines are excluded; only the real line becomes a segment.
|
|
||||||
assert seg_idx._ordered_ids == ["p1_l0"]
|
|
||||||
pos = seg_idx.lookup_segment("p1_l0")
|
|
||||||
assert pos is not None
|
|
||||||
assert pos["text"] == "line-content p1"
|
|
||||||
|
|
||||||
async def test_segment_index_not_built_when_provenance_off(self) -> None:
|
|
||||||
canned = _canned_ocr(pages=1)
|
|
||||||
step = OCRStep(ocr_client=FakeOCRClient(canned=canned))
|
|
||||||
req = _make_request(
|
|
||||||
use_ocr=True, include_provenance=False, files=["http://x"]
|
|
||||||
)
|
|
||||||
resp = _response_with_context(
|
|
||||||
pages=[Page(page_no=1, width=100.0, height=200.0, lines=[])],
|
|
||||||
files=[("/tmp/x.pdf", "application/pdf")],
|
|
||||||
page_metadata=[PageMetadata(file_index=0)],
|
|
||||||
)
|
|
||||||
resp = await step.process(req, resp)
|
|
||||||
assert resp.context.segment_index is None # type: ignore[union-attr]
|
|
||||||
|
|
||||||
async def test_page_tags_include_file_index(self) -> None:
|
|
||||||
canned = _canned_ocr(pages=1)
|
|
||||||
step = OCRStep(ocr_client=FakeOCRClient(canned=canned))
|
|
||||||
req = _make_request(use_ocr=True, files=["http://x"])
|
|
||||||
resp = _response_with_context(
|
|
||||||
pages=[Page(page_no=1, width=100.0, height=200.0, lines=[])],
|
|
||||||
files=[("/tmp/x.pdf", "application/pdf")],
|
|
||||||
page_metadata=[PageMetadata(file_index=3)],
|
|
||||||
)
|
|
||||||
resp = await step.process(req, resp)
|
|
||||||
first_line = resp.ocr_result.result.pages[0].lines[0].text
|
|
||||||
assert first_line is not None
|
|
||||||
assert 'file="3"' in first_line
|
|
||||||
assert 'number="1"' in first_line
|
|
||||||
Loading…
Reference in a new issue