ss-tools/backend/src/plugins/translate/orchestrator_exec.py

# #region TranslationExecutionEngine [C:4] [TYPE Class] [SEMANTICS translate, execution, run, orchestration]
# @defgroup Translate Module group.
# @BRIEF Execute translation runs: dispatch executor, manage completion and failure paths.
# @PRE Database session and config manager are available. Run is in valid state.
# @POST Run is executed with SQL generated; events recorded.
# @SIDE_EFFECT LLM calls, DB writes, Superset API calls.
# @RELATION DEPENDS_ON -> [TranslationRun]
# @RELATION DEPENDS_ON -> [TranslationExecutor]
# @RELATION DEPENDS_ON -> [SQLInsertService]
# @RELATION DEPENDS_ON -> [TranslationResultAggregator]
# @RELATION DEPENDS_ON -> [TranslationEventLog]
# @RELATION DEPENDS_ON -> [orchestrator_run_completion]

from collections.abc import Callable
import uuid

from sqlalchemy.orm import Session

from ...core.config_manager import ConfigManager
from ...core.logger import belief_scope, logger
from ...models.translate import (
    TranslationJob,
    TranslationRun,
    TranslationRunLanguageStats,
)
from .events import TranslationEventLog
from .executor import TranslationExecutor
from .orchestrator_aggregator import TranslationResultAggregator
from .orchestrator_run_completion import (
    complete_cancelled as _complete_cancelled,
    complete_success as _complete_success,
    handle_executor_failure as _handle_executor_failure,
)
from .orchestrator_sql import SQLInsertService


class TranslationExecutionEngine:
    """Execute translation runs: dispatch executor, handle failure/completion paths."""

    def __init__(
        self,
        db: Session,
        config_manager: ConfigManager,
        event_log: TranslationEventLog,
        current_user: str | None = None,
    ):
        self.db = db
        self.config_manager = config_manager
        self.event_log = event_log
        self.current_user = current_user
        self._sql_service = SQLInsertService(db, config_manager, event_log)
        self._aggregator = TranslationResultAggregator(db, event_log)

    # #region TranslationExecutionEngine.execute_run [C:4] [TYPE Function] [SEMANTICS translate,execution,run]
    # @BRIEF Execute a translation run: dispatch executor, handle outcomes.
    # @PRE  run is in PENDING status.
    # @POST  Run executed, SQL generated, Superset submission attempted.
    # @SIDE_EFFECT  LLM calls, DB writes, Superset API calls.
    async def execute_run(
        self,
        run: TranslationRun,
        on_batch_progress: Callable[[str, int, int, int, int], None] | None = None,
        skip_insert: bool = False,
    ) -> TranslationRun:
        with belief_scope("TranslationExecutionEngine.execute_run"):
            job = self.db.query(TranslationJob).filter(TranslationJob.id == run.job_id).first()
            if not job:
                raise ValueError(f"Job '{run.job_id}' not found")
            if run.status != "PENDING":
                raise ValueError(
                    f"Cannot execute run in status '{run.status}'. Run must be in PENDING status."
                )

            logger.reason("Executing run", {"run_id": run.id, "job_id": job.id, "skip_insert": skip_insert})
            self.event_log.log_event(
                job_id=job.id, run_id=run.id, event_type="TRANSLATION_PHASE_STARTED",
                payload={}, created_by=self.current_user,
            )

            language_stats_map = self._init_language_stats(run, job)
            executor = TranslationExecutor(
                self.db, self.config_manager, self.current_user,
                on_batch_progress=on_batch_progress,
            )
            try:
                run = await executor.execute_run(run, llm_progress_callback=None, language_stats_map=language_stats_map)
            except Exception as e:
                return _handle_executor_failure(self.db, self.event_log, run, job, e, self.current_user)

            if run.status == "CANCELLED":
                return _complete_cancelled(self.db, self.event_log, self._aggregator, run, language_stats_map, self.current_user)

            return await _complete_success(
                self.db, self.event_log, self._aggregator, self._sql_service,
                run, job, skip_insert, language_stats_map, self.current_user,
            )
    # #endregion TranslationExecutionEngine.execute_run

    # #region TranslationExecutionEngine._init_language_stats [C:2] [TYPE Function] [SEMANTICS translate,language,stats]
    # @BRIEF Initialize per-language stats for a run.
    def _init_language_stats(
        self,
        run: TranslationRun,
        job: TranslationJob,
    ) -> dict[str, TranslationRunLanguageStats]:
        target_languages = job.target_languages or [job.target_dialect or "en"]
        if not isinstance(target_languages, list):
            target_languages = [str(target_languages)]
        language_stats_map: dict[str, TranslationRunLanguageStats] = {}
        for lang_code in target_languages:
            lang_stat = TranslationRunLanguageStats(
                id=str(uuid.uuid4()), run_id=run.id, language_code=lang_code,
                total_rows=0, translated_rows=0, failed_rows=0, skipped_rows=0,
                token_count=0, estimated_cost=0.0,
            )
            self.db.add(lang_stat)
            language_stats_map[lang_code] = lang_stat
        self.db.flush()
        return language_stats_map
    # #endregion TranslationExecutionEngine._init_language_stats

# #endregion TranslationExecutionEngine