EAP – Document AI Parser
# ── STEP 1: Parse files via DocAI ─────────────────────────
parsed_prefix = f"{tenant_id}/{connector_id}/parsed"
logger.info("📄 Parsing Files With DocAI")
t0 = perf_counter()
parse_result = await parse_files(
input_bucket=bucket,
input_prefix=base_path,
engine="docai",
output_prefix_base=parsed_prefix,
)
timings["parsing"] = perf_counter() - t0
logger.info(
f"Parsing complete [{_fmt(timings['parsing'])}]: "
f"processed={parse_result['processed']} | "
f"success={parse_result['succeeded']} | "
f"failed={parse_result['failed']}"
)
async def parse_files(
input_bucket: str,
engine: str = "mistral",
input_prefix: str | None = None,
output_bucket: str | None = None,
output_prefix_base: str | None = "parsed",
) -> dict:
"""
Parse files from an input GCS bucket and upload parsed outputs.
Files are discovered by listing blobs in ``input_bucket`` and optionally
filtering by ``input_prefix``. This is the single async entry point for
batch parsing.
For ``engine='docai'``, a single DocumentAIProcessor is created once and
reused across all files — this avoids per-file SSL handshakes and reuses
the persistent retry session.
"""
logger = logging.getLogger(__name__)
if not str(output_prefix_base or "").strip() and not output_bucket:
raise ValueError(
"output_bucket is required when output_prefix_base is not provided"
)
cfg = get_settings()
concurrency = int(cfg.get("parsing", {}).get("batch", {}).get("max_concurrency", 5))
semaphore = asyncio.Semaphore(max(1, int(concurrency)))
destination_bucket = output_bucket or input_bucket
gcs = GCSBucketConnector()
object_names = await gcs.list_files(input_bucket, prefix=input_prefix)
object_names = [name for name in object_names if not name.endswith("/")]
# ── FIX: Create DocumentAIProcessor ONCE, reuse across all files ──────────
# Previously created inside _run_one per file — caused new SSL session
# and new credentials object for every single file processed.
docai_processor: DocumentAIProcessor | None = None
if engine == "docai":
logger.info("🔌 [PARSER] Initializing DocumentAIProcessor (shared across all files)")
docai_processor = DocumentAIProcessor()
logger.info("✅ [PARSER] DocumentAIProcessor ready")
logger.info(
f"🚀 [PARSER] Starting batch parsing | "
f"input_bucket={input_bucket} | "
f"input_prefix={input_prefix} | "
f"output_bucket={destination_bucket} | "
f"engine={engine} | "
f"total_files={len(object_names)} | "
f"concurrency={concurrency}"
)
async def _run_one(object_name: str) -> dict:
result = {
"gcs_uri": f"gs://{input_bucket}/{object_name}",
"output_gcs_uri": None,
"status": "failed",
"error": None,
}
async with semaphore:
try:
ext = Path(object_name).suffix.lower()
destination_object = _build_output_object_name(object_name, output_prefix_base)
logger.info(
f"📄 [PARSER] Processing file | "
f"gcs_uri=gs://{input_bucket}/{object_name} | "
f"ext={ext} | "
f"engine={engine}"
)
# ── Download file from GCS ────────────────────────────────────
payload = await gcs.download(input_bucket, object_name)
if not isinstance(payload, bytes):
raise TypeError(
f"Downloaded object is not bytes for "
f"'gs://{input_bucket}/{object_name}'"
)
logger.info(
f"⬇️ [PARSER] Downloaded | "
f"file={object_name} | "
f"size={len(payload) / 1024:.1f}KB"
)
# ── Parse file content ────────────────────────────────────────
parsed_output = await asyncio.to_thread(
_parse_content_bytes,
payload,
ext,
engine,
docai_processor, # ← FIX: pass shared processor
)
# ── Serialize output ──────────────────────────────────────────
if isinstance(parsed_output, (dict, list)):
destination_object = str(Path(destination_object).with_suffix(".json"))
raw_json = cast(dict | list[dict], parsed_output)
output_bytes = json.dumps(raw_json, ensure_ascii=True, indent=2).encode("utf-8")
content_type = "application/json"
else:
destination_object = str(Path(destination_object).with_suffix(".md"))
markdown_text = cast(str, parsed_output)
output_bytes = markdown_text.encode("utf-8")
content_type = "text/markdown"
# ── Upload parsed output to GCS ───────────────────────────────
await gcs.save(
destination_bucket,
destination_object,
data=output_bytes,
content_type=content_type,
)
logger.info(
f"⬆️ [PARSER] Uploaded | "
f"output=gs://{destination_bucket}/{destination_object} | "
f"size={len(output_bytes) / 1024:.1f}KB"
)
result["status"] = "success"
result["output_gcs_uri"] = f"gs://{destination_bucket}/{destination_object}"
logger.info(
f"✅ [PARSER] File complete | "
f"gcs_uri=gs://{input_bucket}/{object_name}"
)
return result
except Exception as exc:
err = _error_payload(exc)
result["error"] = err
logger.error(
f"❌ [PARSER] Failed | "
f"gcs_uri=gs://{input_bucket}/{object_name} | "
f"error_type={err['type']} | "
f"message={err['message']}\n"
f"{err['traceback']}"
)
return result
# ── Run all files concurrently ────────────────────────────────────────────
results = (
await asyncio.gather(*[_run_one(name) for name in object_names])
if object_names
else []
)
succeeded = sum(1 for item in results if item.get("status") == "success")
failed = len(results) - succeeded
logger.info(
f"\n{'─' * 50}\n"
f" 📊 PARSING SUMMARY\n"
f"{'─' * 50}\n"
f" 📁 Total files : {len(results)}\n"
f" ✅ Succeeded : {succeeded}\n"
f" ❌ Failed : {failed}\n"
f" 🪣 Input bucket : {input_bucket}\n"
f" 🪣 Output bucket : {destination_bucket}\n"
f"{'─' * 50}"
)
return {
"input_bucket": input_bucket,
"input_prefix": input_prefix,
"output_bucket": destination_bucket,
"processed": len(results),
"succeeded": succeeded,
"failed": failed,
"results": list(results),
}
class DocumentAIProcessor:
"""Document AI processor via REST API, accepting bytes for GCS-native pipelines."""
def __init__(self):
self.logger = logging.getLogger(__name__)
cfg = get_settings()
docai_cfg = cfg.get("parsing", {}).get("docai", {})
self.max_pages = docai_cfg.get("max_pages_per_request", 15)
self.max_mb = docai_cfg.get("max_mb_per_request", 30.0)
self.timeout_seconds = docai_cfg.get("timeout_seconds", 600)
self.max_retries = docai_cfg.get("max_retries", 3)
self.retry_backoff_factor = docai_cfg.get("retry_backoff_factor", 2.0)
self.mime_map: dict[str, str] = cfg.get("extensions", {}).get("docai_mimes", {})
project_id = get_project_id()
region = os.getenv("DOCAI_LOCATION") or get_region()
processor_id = os.getenv("DOCAI_PROCESSOR_ID")
if not processor_id:
raise RuntimeError("DOCAI_PROCESSOR_ID environment variable not set")
self.endpoint = (
f"https://{region}-documentai.googleapis.com/v1/"
f"projects/{project_id}/locations/{region}/processors/{processor_id}:process"
)
self.logger.info(f"🔌 [DOCAI INIT] Endpoint: {self.endpoint}")
# ── Credentials (ADC — works on Cloud Run automatically) ──────────────
self._credentials, _ = google.auth.default(
scopes=["https://www.googleapis.com/auth/cloud-platform"]
)
self.logger.info("✅ [DOCAI INIT] Credentials loaded via ADC")
# ── Persistent session for token refresh (handles SSL drops) ──────────
self._auth_session = _build_retry_session(retries=5, backoff_factor=2.0)
self._auth_request = Request(session=self._auth_session)
# ── Persistent session for DocAI API calls ────────────────────────────
self._docai_session = _build_retry_session(retries=5, backoff_factor=2.0)
self.logger.info("✅ [DOCAI INIT] Retry sessions initialized")
