Coverage for src / quber / providers / landing / orchestrator.py: 45%

1"""Run one ADE parse end to end: submit, write artifacts, summarize. 

2 

3`run_parse` takes a source PDF or image (local path or ``s3://`` URI), an 

4output destination (local directory or ``s3://`` prefix), a model name, and 

5an optional 1-based page number. It parses the document through ADE, writes 

6the artifacts through the same output contract the other extraction routines 

7honor, and returns a run summary with the credit usage the platform reported. 

8 

9Artifacts per run: ``<base>.ade.json`` (raw parse response), ``<base>.ade.md`` 

10(the parse markdown), and ``<base>.ade.complete.json`` written strictly last as 

11the completion marker. 

12""" 

13 

14from __future__ import annotations 

15 

16import json 

17import tempfile 

18from dataclasses import dataclass, field 

19from pathlib import Path 

20from typing import Optional 

21 

22from loguru import logger 

23 

24from quber.files.cache import resolve_document 

25from quber.files.output import output_sink, write_completion_marker 

26from quber.files.pdf import slice_page 

27from quber.providers.landing.client import DEFAULT_MODEL, parse_document 

28 

29 

30@dataclass 

31class RunSummary: 

32 """What one ADE run produced, for the CLI to print and callers to log.""" 

33 

34 base: str 

35 pages: Optional[int] 

36 #: Credits the platform billed for the run (v2 bills fractional credits). 

37 credits: Optional[float] 

38 model_version: Optional[str] 

39 artifacts: list[str] = field(default_factory=list) 

40 

41 

42def run_parse( 

43 source: str, 

44 output_dir: str = "output", 

45 model: str = DEFAULT_MODEL, 

46 page: Optional[int] = None, 

47) -> RunSummary: 

48 """Parse ``source`` through ADE and write artifacts to ``output_dir``.""" 

49 local = resolve_document(source) 

50 base = local.stem 

51 if page is not None: 

52 base = f"{base}.p{page}" 

53 

54 with tempfile.TemporaryDirectory(prefix="quber-ade-") as tmp: 

55 submitted = local 

56 if page is not None: 

57 submitted = slice_page(local, page, Path(tmp) / f"{base}.pdf") 

58 doc = parse_document(submitted, model=model) 

59 

60 meta = doc.get("metadata") or {} 

61 

62 # v1 metadata reports credit_usage and version; v2 reports 

63 # billing.total_credits and model_version. 

64 credits = meta.get("credit_usage") 

65 if credits is None: 

66 credits = (meta.get("billing") or {}).get("total_credits") 

67 summary = RunSummary( 

68 base=base, 

69 pages=meta.get("page_count"), 

70 credits=credits, 

71 model_version=meta.get("version") or meta.get("model_version"), 

72 ) 

73 

74 with output_sink(output_dir) as out: 

75 json_path = out / f"{base}.ade.json" 

76 json_path.write_text(json.dumps(doc, indent=1), encoding="utf-8") 

77 md_path = out / f"{base}.ade.md" 

78 md_path.write_text(doc.get("markdown") or "", encoding="utf-8") 

79 summary.artifacts = [json_path.name, md_path.name] 

80 

81 write_completion_marker( 

82 output_dir, 

83 f"{base}.ade.complete.json", 

84 { 

85 "artifacts": summary.artifacts, 

86 "page_count": summary.pages, 

87 "credit_usage": summary.credits, 

88 "version": summary.model_version, 

89 "job_id": meta.get("job_id"), 

90 }, 

91 ) 

92 logger.info( 

93 "ADE run complete: base={} pages={} credits={}", 

94 summary.base, 

95 summary.pages, 

96 summary.credits, 

97 ) 

98 return summary