1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359"""Document Processing Service — orchestrates pipeline processing with job infrastructure.
Integrates PipelineViewerService (versioned step-by-step processing) with:
- Redis job state updates
- S3 storage for markdown results, figures, and change ledger
- Prometheus metrics
"""
from __future__ import annotations
import asyncio
import base64
import json
import logging
import time
from typing import Any, Literal
from redis.asyncio import Redis
from ..agents.events import (
EventBus,
PipelinePhaseEvent,
ProcessingCompleteEvent,
ProcessingErrorEvent,
register_event_bus,
unregister_event_bus,
)
from ..config import settings
from .metrics_service import job_duration_seconds, jobs_completed_total
from .pipeline_viewer import PipelineViewerService
from .pipeline_viewer_models import FigureData, PipelineViewerResult
logger = logging.getLogger(__name__)
class DocumentProcessingService:
"""Orchestrates document processing with job tracking and storage.
Downloads PDF from S3, runs PipelineViewerService for processing,
stores results back to S3, and updates Redis job state throughout.
Usage:
service = DocumentProcessingService(redis_client, storage_service)
await service.process_document(
job_id="abc-123",
s3_key="temp/file.pdf",
filename="document.pdf",
)
"""
def __init__(
self,
redis_client: Redis[Any] | None,
storage_service: Any, # StorageService
s3_url_service: Any, # S3URLService
) -> None:
self.redis = redis_client
self.storage = storage_service
self.s3_url = s3_url_service
self.status_prefix = settings.job_status_prefix
async def process_document(
self,
job_id: str,
s3_key: str,
filename: str,
review_mode: Literal["auto", "human"] = "auto",
max_rounds: int = 1,
ocr_languages: list[str] | None = None,
) -> PipelineViewerResult:
"""Process a document through the versioned pipeline.
1. Downloads PDF from S3
2. Runs PipelineViewerService (docling → structure → corrections → boundaries → cleanup)
3. Stores markdown and figures to S3
4. Updates Redis job state throughout
Args:
job_id: Unique job identifier
s3_key: S3 key where PDF is stored
filename: Original filename
review_mode: 'auto' or 'human' (preserved for future use)
max_rounds: Unused (kept for interface compatibility)
ocr_languages: Tesseract OCR language codes for scanned documents
Returns:
PipelineViewerResult with versioned markdowns and step data
"""
await self._update_job_state(
job_id,
status="processing",
processing_phase="docling",
review_mode=review_mode,
)
# Track jobs in processing for CloudWatch scaling metrics
try:
if self.redis is not None:
await self.redis.incr("eq-pdf:metrics:jobs_in_processing")
except Exception:
logger.warning("Failed to increment jobs_in_processing counter")
# Set up event bus for SSE streaming
event_bus = EventBus()
register_event_bus(job_id, event_bus)
try:
processing_start_time = time.time()
# 1. Download PDF from S3
logger.info(f"Job {job_id}: Downloading PDF from S3: {s3_key}")
file_content = await self.storage.download_temp_file(s3_key)
# 2. Run pipeline processing
await self._update_job_state(job_id, processing_phase="pipeline")
def _on_phase(phase: str, display_name: str, step: int, total: int) -> None:
event_bus.publish(PipelinePhaseEvent(
phase=phase,
display_name=display_name,
step_number=step,
total_steps=total,
))
service = PipelineViewerService()
result = await service.process(
file_content=file_content,
filename=filename,
images_scale=2.0,
do_table_structure=True,
enable_structure=True,
enable_page_content=True,
enable_boundaries=True,
ocr_languages=ocr_languages,
on_phase=_on_phase,
)
# 2a. Check for classification errors (unsupported document type)
if not result.versions:
# Classification blocked processing — no versions were produced
classification_step = next(
(s for s in result.steps if s.name == "classification" and s.error),
None,
)
error_msg = (
classification_step.error
if classification_step
else "PDF classification rejected this document"
)
await self._update_job_state(
job_id, status="failed", error=error_msg
)
jobs_completed_total.labels(status="failed").inc()
return result
# 3. Extract final markdown (latest version available)
final_markdown = (
result.versions.get("v3")
or result.versions.get("v2")
or result.versions.get("v1")
or result.versions.get("v0")
or ""
)
# 4. Store results to S3
markdown_s3_key = await self._store_markdown(job_id, final_markdown)
ledger_s3_key = await self._store_change_ledger(job_id, result)
stored_figures = await self._store_figures_from_pipeline(job_id, result.figures)
# 5. Aggregate cost/token info from steps
total_input_tokens = sum(s.input_tokens for s in result.steps)
total_output_tokens = sum(s.output_tokens for s in result.steps)
total_tokens = total_input_tokens + total_output_tokens
cost_cents = sum(s.cost_cents for s in result.steps)
total_edits = sum(len(s.changes) for s in result.steps)
# 6. Update final job state
update_fields: dict[str, Any] = dict(
status="completed",
processing_phase="complete",
result_url=markdown_s3_key,
ledger_s3_key=ledger_s3_key,
total_edits=total_edits,
total_pages=result.total_pages,
llm_input_tokens=total_input_tokens,
llm_output_tokens=total_output_tokens,
llm_total_tokens=total_tokens,
llm_cost_cents=cost_cents,
stored_figures=stored_figures,
)
if result.warnings:
update_fields["warnings"] = result.warnings
await self._update_job_state(job_id, **update_fields)
# 7. Record Prometheus metrics
processing_duration = time.time() - processing_start_time
jobs_completed_total.labels(status="success").inc()
job_duration_seconds.labels(stage="processing").observe(processing_duration)
logger.info(
f"Job {job_id} complete: {total_edits} edits, "
f"{total_tokens} tokens, ${cost_cents / 100:.4f}"
)
event_bus.publish(ProcessingCompleteEvent())
return result
except Exception as e:
logger.error(f"Job {job_id} failed: {e}", exc_info=True)
await self._update_job_state(
job_id,
status="failed",
error=str(e),
)
jobs_completed_total.labels(status="failed").inc()
event_bus.publish(ProcessingErrorEvent(error=str(e)))
raise
finally:
# Decrement jobs-in-processing counter
try:
if self.redis is not None:
await self.redis.decr("eq-pdf:metrics:jobs_in_processing")
except Exception:
logger.warning("Failed to decrement jobs_in_processing counter")
# Let SSE subscribers flush before unregistering
await asyncio.sleep(0.5)
unregister_event_bus(job_id)
async def _update_job_state(self, job_id: str, **fields: Any) -> None:
"""Update job state in Redis."""
from datetime import UTC, datetime
update_data: dict[str | bytes, bytes | float | int | str] = {
"updated_at": datetime.now(UTC).isoformat(),
}
for key, value in fields.items():
if isinstance(value, (dict, list)):
update_data[key] = json.dumps(value)
elif value is not None:
update_data[key] = str(value)
assert self.redis is not None, "Redis client required for _update_job_state"
await self.redis.hset(f"{self.status_prefix}{job_id}", mapping=update_data)
if "status" in fields:
await self._set_job_ttl(job_id, fields["status"])
async def _set_job_ttl(self, job_id: str, status: str) -> None:
"""Set TTL for job based on status."""
if status == "completed":
ttl = settings.job_ttl_completed
elif status == "failed":
ttl = settings.job_ttl_failed
else:
ttl = settings.job_ttl_active
assert self.redis is not None, "Redis client required for _set_job_ttl"
await self.redis.expire(f"{self.status_prefix}{job_id}", ttl)
async def _store_markdown(self, job_id: str, markdown: str) -> str:
"""Store final markdown to S3. Returns S3 key."""
s3_key = f"results/{job_id}/result.md"
await self.storage.upload_file(
key=s3_key,
content=markdown.encode("utf-8"),
content_type="text/markdown",
)
return s3_key
async def _store_change_ledger(self, job_id: str, result: PipelineViewerResult) -> str:
"""Store pipeline change ledger to S3 as JSON. Returns S3 key."""
ledger_data = {
"document_id": job_id,
"steps": [
{
"name": step.name,
"display_name": step.display_name,
"elapsed_ms": step.elapsed_ms,
"input_tokens": step.input_tokens,
"output_tokens": step.output_tokens,
"cost_cents": step.cost_cents,
"changes": [
{
"page": change.page,
"old_text": change.old_text,
"new_text": change.new_text,
"reasoning": change.reasoning,
"stage": change.stage,
}
for change in step.changes
],
}
for step in result.steps
if not step.skipped
],
"total_edits": sum(len(s.changes) for s in result.steps),
}
s3_key = f"results/{job_id}/ledger.json"
await self.storage.upload_file(
key=s3_key,
content=json.dumps(ledger_data, indent=2).encode("utf-8"),
content_type="application/json",
)
return s3_key
async def _store_figures_from_pipeline(
self,
job_id: str,
figures: list[FigureData],
) -> list[dict[str, str | int]]:
"""Store pipeline figures to S3 (base64 → PNG files).
Returns:
List of stored figure metadata dicts for Redis (figure_id, s3_key,
page_num, caption).
"""
stored: list[dict[str, str | int]] = []
if not figures:
return stored
for fig in figures:
if not fig.image_base64:
continue
dest_key = f"results/{job_id}/figures/{fig.ref_id}.png"
try:
image_data = base64.b64decode(fig.image_base64)
await self.storage.upload_file(
key=dest_key,
content=image_data,
content_type="image/png",
)
stored.append({
"figure_id": fig.ref_id,
"s3_key": dest_key,
"page_num": fig.page_number,
"alt_text": "",
"caption": fig.caption,
})
except Exception as e:
logger.warning(f"Job {job_id}: Failed to store figure {fig.ref_id}: {e}")
logger.info(f"Job {job_id}: Stored {len(stored)} figures to results/{job_id}/figures/")
return stored
async def get_ledger(self, job_id: str) -> dict[str, Any] | None:
"""Retrieve ledger from S3."""
s3_key = f"results/{job_id}/ledger.json"
try:
content = await self.storage.download_file(s3_key)
result: dict[str, Any] = json.loads(content.decode("utf-8"))
return result
except Exception as e:
logger.warning(f"Failed to retrieve ledger for job {job_id}: {e}")
return None