Coverage for src / quber / providers / landing / client.py: 25%

1"""Landing.AI ADE parse client. 

2 

3`parse_document` submits one PDF or image through the async Parse Jobs API, 

4polls the job to completion, and returns the full parse response as a dict. 

5Polling is the only completion signal the platform offers; there is no 

6webhook. Results over the inline size limit arrive as a presigned URL that 

7expires shortly after the poll that returned it, so the download happens 

8immediately. 

9 

10Two API generations are in play. The generally-available dpt-2 model line 

11rides the v1 jobs API; the preview dpt-3 line rides the v2 jobs API. The 

12model name picks the route. 

13 

14`custom_prompts` passes per-chunk-type instruction to the v1 model line, keyed 

15by chunk type (`figure`, `text`, ...). It cannot change the shape of what comes 

16back — an instruction to emit a chart as an HTML table returns the same prose 

17description — so it is for content-level instruction only, and it is the 

18caller's to supply. 

19 

20A parse that reports failed pages raises instead of returning: a silently 

21missing page in a financial document is a liability, so a partial parse 

22never looks like a success. 

23""" 

24 

25from __future__ import annotations 

26 

27import json 

28import time 

29import urllib.request 

30from pathlib import Path 

31from typing import Any, Optional 

32 

33from loguru import logger 

34 

35from quber.settings import get_settings 

36 

37#: Generally-available default. The dpt-3 preview names route to the v2 API. 

38DEFAULT_MODEL = "dpt-2" 

39 

40POLL_SECONDS = 5.0 

41TIMEOUT_SECONDS = 1800.0 

42 

43 

44def parse_document( 

45 source: Path, 

46 model: str = DEFAULT_MODEL, 

47 custom_prompts: Optional[dict[str, str]] = None, 

48 poll_seconds: float = POLL_SECONDS, 

49) -> dict[str, Any]: 

50 """Parse one local PDF or image through ADE; return the raw parse response. 

51 

52 `custom_prompts` reaches the v1 model line only; the v2 API takes no such 

53 argument, so it is dropped for a dpt-3 model. `poll_seconds` sets how often 

54 the job is polled — a single-page submission finishes in seconds and wants 

55 a shorter interval than a whole document. 

56 """ 

57 client = build_client() 

58 if is_v2_model(model): 

59 jobs = client.v2.parse_jobs 

60 job = jobs.create(document=source, model=model) 

61 else: 

62 jobs = client.parse_jobs 

63 extra = {"custom_prompts": custom_prompts} if custom_prompts else {} 

64 job = jobs.create(document=source, model=model, **extra) 

65 logger.info( 

66 "ADE parse job {} submitted: {} ({} bytes, model={})", 

67 job.job_id, 

68 source.name, 

69 source.stat().st_size, 

70 model, 

71 ) 

72 return _await_result(jobs, job.job_id, poll_seconds) 

73 

74 

75def is_v2_model(model: str) -> bool: 

76 """True for the dpt-3 model line, which lives on the v2 parse API.""" 

77 return model.startswith("dpt-3") 

78 

79 

80def build_client() -> Any: 

81 """An authenticated SDK client, keyed from settings (ADE_API_KEY).""" 

82 from landingai_ade import LandingAIADE 

83 

84 key = get_settings().landing.api_key 

85 if not key: 

86 raise RuntimeError("No Landing.AI key: set ADE_API_KEY in the environment or .env") 

87 return LandingAIADE(apikey=key) 

88 

89 

90def _await_result(jobs: Any, job_id: str, poll_seconds: float = POLL_SECONDS) -> dict[str, Any]: 

91 """Poll one parse job to completion and return its response dict.""" 

92 t0 = time.time() 

93 last = None 

94 while True: 

95 r = jobs.get(job_id) 

96 state = (r.status, r.progress) 

97 if state != last: 

98 logger.info( 

99 "ADE job {}: status={} progress={:.0%} elapsed={:.0f}s", 

100 job_id, 

101 r.status, 

102 r.progress or 0.0, 

103 time.time() - t0, 

104 ) 

105 last = state 

106 if r.status in ("completed", "failed", "cancelled"): 

107 break 

108 if time.time() - t0 > TIMEOUT_SECONDS: 

109 raise RuntimeError(f"ADE job {job_id} timed out after {TIMEOUT_SECONDS:.0f}s") 

110 time.sleep(poll_seconds) 

111 

112 if r.status != "completed": 

113 # v1 jobs report a failure_reason string; v2 jobs report an error object. 

114 reason = getattr(r, "failure_reason", None) or getattr(r, "error", None) 

115 raise RuntimeError(f"ADE job {job_id} {r.status}: {reason}") 

116 

117 # The completed payload is `data` on a v1 job and `result` on a v2 job. 

118 payload = getattr(r, "data", None) or getattr(r, "result", None) 

119 output_url = getattr(r, "output_url", None) or (getattr(r, "raw", None) or {}).get("output_url") 

120 if payload is not None: 

121 doc = json.loads(payload.model_dump_json()) 

122 elif output_url: 

123 logger.info("ADE job {}: result over inline limit, downloading from output_url", job_id) 

124 with urllib.request.urlopen(output_url) as resp: 

125 doc = json.loads(resp.read()) 

126 else: 

127 raise RuntimeError(f"ADE job {job_id} completed with neither inline payload nor output_url") 

128 

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

130 failed = meta.get("failed_pages") or [] 

131 if failed: 

132 raise RuntimeError(f"ADE parse failed on pages {failed}; refusing to return a partial parse") 

133 # v1 metadata reports credit_usage and version; v2 reports 

134 # billing.total_credits and model_version. 

135 credits = meta.get("credit_usage") 

136 if credits is None: 

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

138 logger.info( 

139 "ADE parse complete: pages={} credits={} version={}", 

140 meta.get("page_count"), 

141 credits, 

142 meta.get("version") or meta.get("model_version"), 

143 ) 

144 return doc