Coverage for src / quber / core / extractors / set_of_mark / orchestrator.py: 53%

1"""SetOfMarkExtractor (Set-of-Mark): vision identifies tables, Camelot fills them in-region. 

2 

3Thin wrapper that renders the document's pages once, wires `SetOfMarkDeps`, and 

4runs the vision-guided extraction graph built in `pipeline`. The graph does the 

5work: locate every table, capture each in its own region, clean structure under 

6the no-number guard, and reading-order the result. Emits one `ExtractedTable` 

7per visual table; a region where Camelot finds nothing is still emitted with 

8its visual identity so a table is never silently dropped. 

9""" 

10 

11from __future__ import annotations 

12 

13import asyncio 

14import tempfile 

15from pathlib import Path 

16from typing import List, Optional 

17 

18import fitz 

19 

20from quber.agents.capture_advisor import CaptureAdvisor, get_capture_advisor 

21from quber.agents.grid_locator import GridLocator, get_grid_locator 

22from quber.agents.llm_client import LLMClient, get_llm_client 

23from quber.agents.status_inspector import StatusInspector, get_status_inspector 

24from quber.core.extractors.base import ExtractedTable 

25from quber.core.extractors.camelot.acquire import render_pages 

26from quber.core.extractors.set_of_mark.pipeline import SetOfMarkDeps, SetOfMarkState, build_set_of_mark_graph 

27 

28 

29class SetOfMarkExtractor: 

30 """Set-of-Mark extractor: vision locates, Camelot fills in-region, correction cleans.""" 

31 

32 def __init__( 

33 self, 

34 locator: Optional[GridLocator] = None, 

35 llm: Optional[LLMClient] = None, 

36 dpi: int = 200, 

37 page_concurrency: int = 16, 

38 correct_concurrency: int = 16, 

39 advisor: Optional[CaptureAdvisor] = None, 

40 inspector: Optional[StatusInspector] = None, 

41 ) -> None: 

42 self.locator = locator or get_grid_locator() 

43 self.llm = llm or get_llm_client(None) 

44 self.dpi = dpi 

45 self.page_concurrency = page_concurrency 

46 self.correct_concurrency = correct_concurrency 

47 self.advisor = advisor if advisor is not None else get_capture_advisor() 

48 self.inspector = inspector if inspector is not None else get_status_inspector() 

49 

50 async def extract_tables(self, source: Path) -> List[ExtractedTable]: 

51 doc = fitz.open(str(source)) 

52 page_dims = {i + 1: (doc[i].rect.width, doc[i].rect.height) for i in range(doc.page_count)} 

53 with tempfile.TemporaryDirectory(prefix="quber-som-") as tmp: 

54 images = await asyncio.to_thread(render_pages, source, self.dpi, Path(tmp)) 

55 page_images = {i + 1: img for i, img in enumerate(images)} 

56 deps = SetOfMarkDeps( 

57 locator=self.locator, 

58 llm=self.llm, 

59 source=source, 

60 page_images=page_images, 

61 page_dims=page_dims, 

62 dpi=self.dpi, 

63 correct_sem=asyncio.Semaphore(self.correct_concurrency), 

64 page_sem=asyncio.Semaphore(self.page_concurrency), 

65 advisor=self.advisor, 

66 inspector=self.inspector, 

67 ) 

68 return await build_set_of_mark_graph().run(state=SetOfMarkState(), deps=deps, inputs=source) 

69 

70 def extract_tables_sync(self, source: Path) -> List[ExtractedTable]: 

71 return asyncio.run(self.extract_tables(source))