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

# #region TranslationPreview [C:4] [TYPE Module] [SEMANTICS sqlalchemy, translate, preview, llm, review]
# @BRIEF Preview session management: fetch sample rows from Superset, send to LLM with context + filtered dictionary, return side-by-side results.
# @LAYER Domain
# @RELATION DEPENDS_ON -> [TranslationJob:Class]
# @RELATION DEPENDS_ON -> [TranslationPreviewSession:Class]
# @RELATION DEPENDS_ON -> [TranslationPreviewRecord:Class]
# @RELATION DEPENDS_ON -> [LLMProviderService]
# @RELATION DEPENDS_ON -> [DictionaryManager]
# @RELATION DEPENDS_ON -> [SupersetClient]
# @RELATION DEPENDS_ON -> [render_prompt]
# @RELATION DEPENDS_ON -> [ConfigManager]
# @PRE Database session and config manager are available.
# @POST Preview sessions are created with LLM-translated rows; records can be approved/edited/rejected.
# @SIDE_EFFECT Fetches sample data from Superset; calls LLM provider; creates DB rows.
# @RATIONALE C4 because preview is stateful with approve/edit/reject lifecycle and LLM API calls.
# @REJECTED Transient preview state (in-memory/session-scoped) — would lose decisions on restart and cannot gate execution reliably.

import time
import uuid
from datetime import UTC, datetime, timedelta
from typing import Any

from sqlalchemy.orm import Session

from ...core.config_manager import ConfigManager
from ._lang_detect import detect_language, get_detector
from .preview_constants import DEFAULT_EXECUTION_PROMPT_TEMPLATE, DEFAULT_PREVIEW_PROMPT_TEMPLATE  # noqa: F401
from ...core.logger import belief_scope, logger
from ...models.translate import (
    TranslationJob,
    TranslationPreviewLanguage,
    TranslationPreviewRecord,
    TranslationPreviewSession,
)
from .preview_executor import PreviewExecutor
from .preview_prompt_builder import PreviewPromptBuilder
from .preview_review import PreviewSessionManager
from .preview_token_estimator import TokenEstimator
from .preview_response_parser import (
    compute_config_hash as _compute_config_hash_module,
    parse_llm_response as _parse_llm_response_module,
)


# #region TranslationPreview [C:4] [TYPE Class]
# @BRIEF Manages preview lifecycle: fetch sample rows, call LLM, manage row-level approve/edit/reject, accept gate.
# @PRE Database session and config manager are available.
# @POST Preview sessions created with persisted records; full execution gates on accepted session.
# @SIDE_EFFECT Fetches sample data from Superset; calls LLM provider; creates DB rows.
class TranslationPreview:

    def __init__(
        self,
        db: Session,
        config_manager: ConfigManager,
        current_user: str | None = None,
    ):
        self.db = db
        self.config_manager = config_manager
        self.current_user = current_user
        self._executor = PreviewExecutor(db, config_manager)
        self._prompt_builder = PreviewPromptBuilder(db)
        self._session_mgr = PreviewSessionManager(db)

    # region preview_rows [TYPE Function]
    # @PURPOSE: Fetch sample rows, send to LLM, create preview session with per-language records.
    # @SIDE_EFFECT: Fetches data from Superset; calls LLM; creates DB rows.
    def preview_rows(
        self,
        job_id: str,
        sample_size: int = 10,
        prompt_template: str | None = None,
        env_id: str | None = None,
    ) -> dict[str, Any]:
        with belief_scope("TranslationPreview.preview_rows"):
            job = self.db.query(TranslationJob).filter(TranslationJob.id == job_id).first()
            if not job:
                raise ValueError(f"Translation job '{job_id}' not found")
            if not job.source_datasource_id:
                raise ValueError("Job must have a source datasource configured for preview")
            if not job.translation_column:
                raise ValueError("Job must have a translation column configured for preview")
            if not job.target_languages:
                raise ValueError("Job must have at least one target language configured for preview")
            if not job.provider_id:
                raise ValueError("Job must have an LLM provider configured for preview")

            target_languages = job.target_languages or [job.target_dialect or "en"]
            if not isinstance(target_languages, list):
                target_languages = [str(target_languages)]

            config_hash = self._executor.compute_config_hash(job)
            dict_snapshot_hash = self._executor.compute_dict_snapshot_hash(job_id)
            source_rows = self._executor.fetch_sample_rows(job=job, sample_size=sample_size, env_id=env_id)
            if not source_rows:
                raise ValueError("No rows returned from datasource for preview")

            provider_model = self._executor.resolve_provider_model(job)
            budget_result = self._prompt_builder.estimate_token_budget_for_rows(
                source_rows=source_rows, target_languages=target_languages, job=job, provider_model=provider_model,
            )
            source_rows, actual_row_count, token_budget = (
                budget_result["source_rows"], budget_result["actual_row_count"], budget_result["token_budget"],
            )

            if token_budget.get("warning"):
                logger.explore("Token budget warning", {"warning": token_budget["warning"], "sample_size": actual_row_count})

            prompt_data = self._prompt_builder.build_prompt_from_rows(
                job=job, source_rows=source_rows, sample_size=sample_size, prompt_template=prompt_template,
            )

            # ★ NEW: Run local language detection on row_meta (no LLM)
            detector = get_detector(target_languages)
            for meta in prompt_data["row_meta"]:
                meta["_detected_lang"] = detect_language(meta["source_text"], detector)

            t_llm = time.monotonic()
            llm_response = self._executor.call_llm(
                job=job, prompt=prompt_data["prompt"], max_tokens=token_budget["max_output_needed"],
            )
            logger.reason(f"TIMING: LLM call: {time.monotonic() - t_llm:.2f}s", {})

            translations = self._executor.parse_llm_response(
                llm_response, prompt_data["actual_row_count"], target_languages=target_languages,
            )

            session = TranslationPreviewSession(
                id=str(uuid.uuid4()), job_id=job_id, status="ACTIVE",
                created_by=self.current_user, created_at=datetime.now(UTC),
                expires_at=datetime.now(UTC) + timedelta(hours=24),
            )
            self.db.add(session)
            self.db.flush()

            records = self._create_preview_records(
                job=job, row_meta=prompt_data["row_meta"], translations=translations,
                target_languages=target_languages, session=session,
            )

            result = {
                "id": session.id, "job_id": job_id, "status": "ACTIVE",
                "created_by": self.current_user, "created_at": session.created_at.isoformat(),
                "expires_at": session.expires_at.isoformat() if session.expires_at else None,
                "records": records, "target_languages": target_languages,
                "cost_estimate": {
                    "sample_size": prompt_data["actual_row_count"],
                    "num_languages": prompt_data["num_languages"],
                    "sample_prompt_tokens": prompt_data["sample_prompt_tokens"],
                    "sample_output_tokens": prompt_data["sample_output_tokens"],
                    "sample_total_tokens": prompt_data["sample_total_tokens"],
                    "sample_cost": prompt_data["sample_cost"],
                    "estimated_total_rows": prompt_data["total_est_rows"],
                    "estimated_tokens": prompt_data["total_est_tokens"],
                    "estimated_cost": prompt_data["total_est_cost"],
                    "warning": prompt_data["cost_warning"],
                },
                "token_budget": {
                    "batch_size_adjusted": token_budget["batch_size_adjusted"],
                    "estimated_input_tokens": token_budget["estimated_input_tokens"],
                    "estimated_output_tokens": token_budget["estimated_output_tokens"],
                    "max_output_needed": token_budget["max_output_needed"],
                    "warning": token_budget["warning"],
                },
                "config_hash": config_hash, "dict_snapshot_hash": dict_snapshot_hash,
            }
            logger.reflect("Preview completed", {"session_id": session.id, "row_count": actual_row_count})
            return result
    # endregion preview_rows

    # region _create_preview_records [TYPE Function]
    def _create_preview_records(
        self,
        job: TranslationJob, row_meta: list[dict[str, Any]],
        translations: dict[str, dict[str, str]], target_languages: list[str],
        session: TranslationPreviewSession,
    ) -> list[dict[str, Any]]:
        records = []
        for meta in row_meta:
            idx = meta["row_index"]
            source_text = meta["source_text"]
            translation_data = translations.get(str(idx), {})
            # ★ Local detection (from lingua) takes priority over LLM response
            detected_lang = meta.get("_detected_lang", "und") or "und"
            if detected_lang == "und" and isinstance(translation_data, dict):
                detected_lang = translation_data.get("detected_source_language", "und") or "und"

            lang_entries: list[TranslationPreviewLanguage] = []
            for lang_code in target_languages:
                lang_translation = None
                if isinstance(translation_data, dict):
                    lang_translation = translation_data.get(lang_code) or translation_data.get("translation", "")
                elif isinstance(translation_data, str):
                    lang_translation = translation_data
                if str(lang_code).lower() == str(detected_lang).lower() and detected_lang != "und":
                    lang_translation = source_text
                if not lang_translation:
                    lang_translation = ""

                lang_entries.append(TranslationPreviewLanguage(
                    id=str(uuid.uuid4()), preview_record_id="",
                    language_code=lang_code, source_language_detected=detected_lang,
                    translated_value=str(lang_translation), final_value=str(lang_translation),
                    status="pending", needs_review=(detected_lang == "und"),
                ))

            source_row = meta.get("source_row", {})
            source_data = None
            if job.target_key_cols:
                source_data = {k: source_row.get(k) for k in job.target_key_cols if k in source_row}
            elif source_row:
                source_data = dict(source_row)

            record = TranslationPreviewRecord(
                id=str(uuid.uuid4()), session_id=session.id, source_sql=source_text,
                target_sql=lang_entries[0].translated_value if lang_entries else "",
                source_object_type="table_row", source_object_id=str(idx),
                source_object_name=f"Row {idx + 1}", source_data=source_data,
                status="PENDING", feedback=None, created_at=datetime.now(UTC),
            )
            self.db.add(record)
            self.db.flush()

            serialized_langs = []
            for le in lang_entries:
                le.preview_record_id = record.id
                self.db.add(le)
                serialized_langs.append({
                    "language_code": le.language_code, "source_language_detected": le.source_language_detected,
                    "translated_value": le.translated_value, "final_value": le.final_value,
                    "status": le.status, "needs_review": le.needs_review,
                })

            records.append({
                "id": record.id, "source_sql": record.source_sql, "target_sql": record.target_sql,
                "source_object_type": record.source_object_type, "source_object_id": record.source_object_id,
                "source_object_name": record.source_object_name, "status": record.status,
                "feedback": record.feedback, "source_language_detected": detected_lang,
                "needs_review": any(le.needs_review for le in lang_entries),
                "languages": serialized_langs,
            })

        self.db.commit()
        return records
    # endregion _create_preview_records

    # Backward-compatible static methods (delegates to preview_response_parser)
    @staticmethod
    def _parse_llm_response(
        response_text: str, expected_count: int,
        target_languages: list[str] | None = None,
        finish_reason: str | None = None,
    ) -> dict[str, dict[str, str]]:
        return _parse_llm_response_module(response_text, expected_count, target_languages, finish_reason)

    @staticmethod
    def _compute_config_hash(job: TranslationJob) -> str:
        return _compute_config_hash_module(job)

    # Delegated methods
    def update_preview_row(self, job_id: str, row_id: str, action: str, translation: str | None = None,
                           feedback: str | None = None, language_code: str | None = None) -> dict[str, Any]:
        return self._session_mgr.update_preview_row(job_id, row_id, action, translation, feedback, language_code)

    def accept_preview_session(self, job_id: str) -> dict[str, Any]:
        return self._session_mgr.accept_preview_session(job_id)

    def get_preview_session(self, job_id: str) -> dict[str, Any]:
        return self._session_mgr.get_preview_session(job_id)

# #endregion TranslationPreview

# Re-export for backward compatibility
from .preview_token_estimator import TokenEstimator  # noqa: E402, F401
# #endregion TranslationPreview