Create-knowledge-find-decode/pipeline/acquisition_runner.py

"""Pipeline adapter for the acquisition stage."""
from __future__ import annotations

from dataclasses import dataclass
from typing import Any
from uuid import UUID

from acquisition.repositories.base import AcquisitionRepository
from acquisition.runner import RunBatchResult, run_batch
from core.config import Settings
from pipeline.models import PipelineJob, PipelineRun
from pipeline.repository import PipelineRepository


@dataclass(frozen=True)
class AcquisitionStageResult:
    pipeline_job: PipelineJob | None
    acquisition: RunBatchResult


def run_acquisition_stage(
    *,
    acquisition_repo: AcquisitionRepository,
    batch_id: UUID,
    settings: Settings,
    pipeline_repo: PipelineRepository | None = None,
    pipeline_run: PipelineRun | None = None,
    **kwargs: Any,
) -> AcquisitionStageResult:
    job = None
    if pipeline_repo and pipeline_run and pipeline_run.id:
        job = pipeline_repo.save_pipeline_job(
            run_id=pipeline_run.id,
            stage="search",
            target_id=batch_id,
            status="running",
            metadata={"batch_id": str(batch_id)},
        )
    try:
        result = run_batch(acquisition_repo, batch_id=batch_id, settings=settings, **kwargs)
        if pipeline_repo and job and job.id:
            job = pipeline_repo.mark_job_status(
                job.id,
                status="done" if result.failed == 0 else "partial",
                metadata=result.__dict__,
            )
        return AcquisitionStageResult(pipeline_job=job, acquisition=result)
    except Exception as exc:
        if pipeline_repo and job and job.id:
            pipeline_repo.mark_job_status(job.id, status="failed", error_message=str(exc)[:300])
        raise