diff --git a/backend/app/api/docs/tts_evaluation/start_evaluation.md b/backend/app/api/docs/tts_evaluation/start_evaluation.md index 40a7c3ff2..fda71407f 100644 --- a/backend/app/api/docs/tts_evaluation/start_evaluation.md +++ b/backend/app/api/docs/tts_evaluation/start_evaluation.md @@ -9,7 +9,10 @@ Optional fields: The evaluation will: 1. Process each text sample through the specified TTS models -2. Generate speech audio using Gemini Batch API -3. Store WAV audio files in S3 for human review +2. Generate speech audio: Gemini models via the Gemini Batch API; Sarvam and ElevenLabs models via direct API calls in background workers +3. Store 24 kHz mono WAV audio files in S3 for human review -**Supported models:** `gemini-2.5-pro-preview-tts` +Synthesis uses the dataset's language; datasets without a language are synthesized as `en-IN`. +The response's `run_metadata.voices` lists the voice used per model. + +**Supported models:** `gemini-2.5-pro-preview-tts`, `gemini-3.1-flash-tts-preview`, `bulbul:v3` (Sarvam), `eleven_v3`, `eleven_v4` (ElevenLabs) diff --git a/backend/app/api/routes/tts_evaluations/evaluation.py b/backend/app/api/routes/tts_evaluations/evaluation.py index 6751d8219..b51f59793 100644 --- a/backend/app/api/routes/tts_evaluations/evaluation.py +++ b/backend/app/api/routes/tts_evaluations/evaluation.py @@ -25,6 +25,10 @@ from app.services.tts_evaluations.constants import ( DEFAULT_STYLE_PROMPT, DEFAULT_VOICE_NAME, + ELEVENLABS_TTS_MODELS, + ELEVENLABS_VOICE_ID, + SARVAM_SPEAKER, + SARVAM_TTS_MODELS, ) from app.utils import APIResponse, load_description @@ -114,11 +118,22 @@ def start_tts_evaluation( detail=f"Failed to queue batch submission: {e}", ) + voices: dict[str, str] = {} + for model in run_create.models: + if model in SARVAM_TTS_MODELS: + voices[model] = SARVAM_SPEAKER + elif model in ELEVENLABS_TTS_MODELS: + voices[model] = ELEVENLABS_VOICE_ID + else: + voices[model] = DEFAULT_VOICE_NAME + return APIResponse.success_response( data=TTSEvaluationRunPublic.from_model( run, run_metadata={ + # Legacy key kept for existing clients; `voices` is per model. "voice_name": DEFAULT_VOICE_NAME, + "voices": voices, "style_prompt": DEFAULT_STYLE_PROMPT, }, ) diff --git a/backend/app/celery/tasks/job_execution.py b/backend/app/celery/tasks/job_execution.py index 483384364..375ebc89d 100644 --- a/backend/app/celery/tasks/job_execution.py +++ b/backend/app/celery/tasks/job_execution.py @@ -429,6 +429,28 @@ def run_tts_result_processing( ) +@celery_app.task(bind=True, queue="default", priority=2) +@gevent_timeout(settings.CELERY_TASK_SOFT_TIME_LIMIT, "run_tts_sync_generation") +def run_tts_sync_generation( + self, project_id: int, job_id: str, trace_id: str, **kwargs +): + from app.services.tts_evaluations.sync_generation import ( + execute_tts_sync_generation, + ) + + _set_trace(trace_id) + return _run_with_otel_parent( + self, + lambda: execute_tts_sync_generation( + project_id=project_id, + job_id=job_id, + task_id=current_task.request.id, + task_instance=self, + **kwargs, + ), + ) + + @celery_app.task(bind=True, queue="default", priority=6) @gevent_timeout(settings.CELERY_TASK_SOFT_TIME_LIMIT, "run_evaluation_fast_chunk") def run_evaluation_fast_chunk( diff --git a/backend/app/celery/utils.py b/backend/app/celery/utils.py index 3bd205d22..df9e9643c 100644 --- a/backend/app/celery/utils.py +++ b/backend/app/celery/utils.py @@ -270,6 +270,24 @@ def start_tts_result_processing( return task_id +def start_tts_sync_generation( + project_id: int, job_id: str, trace_id: str = "N/A", **kwargs: Any +) -> str: + from app.celery.tasks.job_execution import run_tts_sync_generation + + task_id = _enqueue_with_trace_context( + run_tts_sync_generation, + project_id=project_id, + job_id=job_id, + trace_id=trace_id, + **kwargs, + ) + logger.info( + f"[start_tts_sync_generation] Started job {job_id} with Celery task {task_id}" + ) + return task_id + + def start_fast_evaluation_chunk( eval_run_id: int, chunk_index: int, trace_id: str = "N/A" ) -> str: diff --git a/backend/app/core/batch/operations.py b/backend/app/core/batch/operations.py index d8db4151b..04f1e76e6 100644 --- a/backend/app/core/batch/operations.py +++ b/backend/app/core/batch/operations.py @@ -49,6 +49,11 @@ def start_batch_job( batch_job = create_batch_job(session=session, batch_job_create=batch_job_create) + # create_batch_job's refresh() re-opens a transaction; end it so no pooled + # connection sits idle-in-transaction during the (slow) upload + create. + # Nothing is pending at this point, so this only releases the connection. + session.commit() + try: batch_result = provider.create_batch(jsonl_data=jsonl_data, config=config) diff --git a/backend/app/crud/tts_evaluations/cron.py b/backend/app/crud/tts_evaluations/cron.py index 2f3c67aaf..62b0d5d33 100644 --- a/backend/app/crud/tts_evaluations/cron.py +++ b/backend/app/crud/tts_evaluations/cron.py @@ -1,43 +1,57 @@ """Cron processing functions for TTS evaluations. -This module provides functions that are called periodically to process -pending TTS evaluations - polling batch status and dispatching result -processing to Celery workers. - -Follows the same pattern as STT evaluations: single query to fetch all -processing runs, grouped by project_id for credential management. +The cron is the only code that sets a TTS run's final status. Each tick it polls +Gemini batches (dispatching result processing on success), fails the rows of +batches that ended badly or of sync models whose worker is gone, then recomputes +the run status from its result rows via `finalize_tts_run_status`. """ import logging +from datetime import timedelta from typing import Any -from sqlmodel import Session +from sqlmodel import Session, select from app.celery.utils import start_tts_result_processing -from app.core.batch import GeminiBatchProvider +from app.core.batch import BatchJobState, GeminiBatchProvider, GeminiClient +from app.core.config import settings +from app.core.util import now from app.crud.evaluations.cron_utils import ( + TERMINAL_STATES, get_batch_jobs_for_run, + make_empty_summary, make_poll_result, - poll_all_pending_evaluations_by_type, poll_batch_jobs, ) from app.crud.tts_evaluations.result import ( - count_results_by_status, get_pending_results_for_run, + update_tts_result, ) -from app.crud.tts_evaluations.run import update_tts_run +from app.crud.tts_evaluations.run import finalize_tts_run_status from app.models import EvaluationRun from app.models.batch_job import BatchJob, BatchJobType from app.models.job import JobStatus from app.models.stt_evaluation import EvaluationType +from app.services.tts_evaluations.constants import ( + GEMINI_TTS_MODELS, + SYNC_TTS_MODELS, + TTS_SYNC_STALE_GRACE_SECONDS, +) logger = logging.getLogger(__name__) +EVAL_TYPE_LABEL = "tts" +SUCCESS_ACTIONS = ("completed", "dispatched") + async def poll_all_pending_tts_evaluations( session: Session, ) -> dict[str, Any]: - """Poll all pending TTS evaluations across all organizations. + """Poll every processing TTS run, including runs with only Sarvam/ElevenLabs models. + + Not built on `poll_all_pending_evaluations_by_type`: that helper skips runs + without a `batch_job_id` and fails a whole project when its Gemini client + can't be built, which would break runs that don't use Gemini at all. Args: session: Database session @@ -45,14 +59,67 @@ async def poll_all_pending_tts_evaluations( Returns: Summary dict with total, processed, failed, still_processing counts """ - return await poll_all_pending_evaluations_by_type( - session, - eval_type="tts", - eval_type_enum_value=EvaluationType.TTS.value, - update_run_fn=update_tts_run, - poll_run_fn=poll_tts_run, - success_actions=("completed", "processed", "dispatched"), + logger.info("[poll_all_pending_tts_evaluations] Starting TTS evaluation polling") + + statement = select(EvaluationRun).where( + EvaluationRun.type == EvaluationType.TTS.value, + EvaluationRun.status == "processing", ) + pending_runs = list(session.exec(statement).all()) + + if not pending_runs: + logger.info("[poll_all_pending_tts_evaluations] No pending TTS runs found") + return make_empty_summary() + + all_results: list[dict[str, Any]] = [] + total_processed = 0 + total_failed = 0 + total_still_processing = 0 + + for run in pending_runs: + try: + result = await poll_tts_run( + session=session, run=run, org_id=run.organization_id + ) + except Exception as e: + # Left in `processing` so the next tick retries; a transient poll error + # must not fail models that may still succeed. + logger.error( + f"[poll_all_pending_tts_evaluations] Failed to poll TTS run | " + f"run_id={run.id} | {e}", + exc_info=True, + ) + session.rollback() + result = make_poll_result( + run=run, + eval_type=EVAL_TYPE_LABEL, + previous_status=run.status, + current_status=run.status, + action="no_change", + error=f"Polling failed: {str(e)}", + ) + all_results.append(result) + + if result["action"] in SUCCESS_ACTIONS: + total_processed += 1 + elif result["action"] == "failed": + total_failed += 1 + else: + total_still_processing += 1 + + logger.info( + f"[poll_all_pending_tts_evaluations] Polling summary | " + f"processed={total_processed} | failed={total_failed} | " + f"still_processing={total_still_processing}" + ) + + return { + "total": len(pending_runs), + "processed": total_processed, + "failed": total_failed, + "still_processing": total_still_processing, + "details": all_results, + } def _dispatch_tts_result_processing( @@ -93,18 +160,13 @@ def _dispatch_tts_result_processing( async def poll_tts_run( session: Session, run: EvaluationRun, - batch_provider: GeminiBatchProvider, org_id: int, ) -> dict[str, Any]: - """Poll a single TTS evaluation run's batch status. - - Finds all batch jobs for this run (one per provider) and polls each. - When a batch reaches SUCCEEDED, dispatches result processing to Celery. + """Advance a single TTS run and recompute its status. Args: session: Database session run: The evaluation run to poll - batch_provider: Initialized GeminiBatchProvider org_id: Organization ID Returns: @@ -116,102 +178,157 @@ async def poll_tts_run( logger.info(f"{log_prefix} Polling run") previous_status = run.status - - batch_jobs = get_batch_jobs_for_run( - session=session, run=run, job_type=BatchJobType.TTS_EVALUATION + models = run.providers or [] + batch_models = [model for model in models if model in GEMINI_TTS_MODELS] + sync_models = [model for model in models if model in SYNC_TTS_MODELS] + # Both timestamps are naive UTC (`now()` strips tzinfo before storing). + stale_after = timedelta( + seconds=settings.CELERY_TASK_TIME_LIMIT + TTS_SYNC_STALE_GRACE_SECONDS ) - - if not batch_jobs: - logger.warning(f"{log_prefix} No batch jobs found") - update_tts_run( - session=session, - run_id=run.id, - status="failed", - error_message="No batch jobs found", - ) - return make_poll_result( - run=run, - eval_type="tts", - previous_status=previous_status, - current_status="failed", - action="failed", - error="No batch jobs found", - ) - - async def _on_batch_succeeded(batch_job: BatchJob, provider_name: str) -> bool: - _dispatch_tts_result_processing(run, batch_job, org_id, provider_name) - return True - - async def _on_already_succeeded(batch_job: BatchJob, provider_name: str) -> bool: - pending = get_pending_results_for_run( - session=session, run_id=run.id, provider=provider_name - ) - if pending: - logger.info( - f"{log_prefix} Dispatching reprocessing for " - f"{len(pending)} pending results | " - f"batch_job_id={batch_job.id}" + is_stale = now() - run.inserted_at > stale_after + any_dispatched = False + + batch_jobs: list[BatchJob] = [] + if batch_models: + # A job whose creation failed has no provider_batch_id and nothing to poll; + # submission already marked its rows FAILED. + batch_jobs = [ + batch_job + for batch_job in get_batch_jobs_for_run( + session=session, run=run, job_type=BatchJobType.TTS_EVALUATION + ) + if batch_job.provider_batch_id + ] + + # Submission marks rows FAILED when a batch can't be created, so leftover PENDING + # rows without a batch job mean the submission task died part-way. + if batch_models and not batch_jobs and is_stale: + logger.warning(f"{log_prefix} No batch jobs found for Gemini models") + for model in batch_models: + _fail_pending_results( + session=session, + run_id=run.id, + model=model, + error_message="No batch jobs found", ) - _dispatch_tts_result_processing(run, batch_job, org_id, provider_name) - return True - return False - - result = await poll_batch_jobs( - session=session, - batch_jobs=batch_jobs, - batch_provider=batch_provider, - provider_config_key="tts_provider", - log_prefix=log_prefix, - on_succeeded=_on_batch_succeeded, - on_already_succeeded=_on_already_succeeded, - ) - - if not result.all_terminal: - return make_poll_result( - run=run, - eval_type="tts", - previous_status=previous_status, - current_status=run.status, - action="no_change", - ) - # If we dispatched processing to Celery, keep the run as "processing". - # The Celery task will finalize the run status when done. - if result.any_dispatched: - return make_poll_result( - run=run, - eval_type="tts", - previous_status=previous_status, - current_status="processing", - action="dispatched", - ) + if batch_jobs: + batch_provider: GeminiBatchProvider | None = None + try: + gemini_client = GeminiClient.from_credentials( + session=session, org_id=org_id, project_id=run.project_id + ) + batch_provider = GeminiBatchProvider(client=gemini_client.client) + except Exception as client_err: + logger.error( + f"{log_prefix} Failed to get Gemini client | error={client_err}" + ) + for model in batch_models: + _fail_pending_results( + session=session, + run_id=run.id, + model=model, + error_message=f"Gemini client initialization failed: {str(client_err)}", + ) + + if batch_provider is not None: + + async def _on_batch_succeeded( + batch_job: BatchJob, provider_name: str + ) -> bool: + _dispatch_tts_result_processing(run, batch_job, org_id, provider_name) + return True + + async def _on_already_succeeded( + batch_job: BatchJob, provider_name: str + ) -> bool: + pending = get_pending_results_for_run( + session=session, run_id=run.id, provider=provider_name + ) + if not pending: + return False + logger.info( + f"{log_prefix} Dispatching reprocessing for " + f"{len(pending)} pending results | batch_job_id={batch_job.id}" + ) + _dispatch_tts_result_processing(run, batch_job, org_id, provider_name) + return True + + poll_result = await poll_batch_jobs( + session=session, + batch_jobs=batch_jobs, + batch_provider=batch_provider, + provider_config_key="tts_provider", + log_prefix=log_prefix, + on_succeeded=_on_batch_succeeded, + on_already_succeeded=_on_already_succeeded, + ) + any_dispatched = poll_result.any_dispatched + + for batch_job in batch_jobs: + if batch_job.provider_status not in TERMINAL_STATES: + continue + if batch_job.provider_status == BatchJobState.SUCCEEDED.value: + continue + _fail_pending_results( + session=session, + run_id=run.id, + model=batch_job.config.get("tts_provider", ""), + error_message=( + f"Gemini batch ended: " + f"{batch_job.error_message or batch_job.provider_status}" + ), + ) + + # Sync workers write rows as they go; past the stale window the worker is gone. + if is_stale: + for model in sync_models: + _fail_pending_results( + session=session, + run_id=run.id, + model=model, + error_message="Synthesis did not complete", + ) - # All batch jobs are done and no dispatching needed - finalize the run - status_counts = count_results_by_status(session=session, run_id=run.id) - pending = status_counts.get(JobStatus.PENDING.value, 0) - failed_count = status_counts.get(JobStatus.FAILED.value, 0) - - final_status = "completed" if pending == 0 else "processing" - error_message = None - if result.any_failed: - error_message = "; ".join(result.errors) - elif failed_count > 0: - error_message = f"{failed_count} synthesis(es) failed" - - update_tts_run( - session=session, - run_id=run.id, - status=final_status, - error_message=error_message, - ) + finalized_run = finalize_tts_run_status(session=session, run_id=run.id) + final_status = finalized_run.status if finalized_run else run.status + error_message = finalized_run.error_message if finalized_run else None - action = "completed" if not result.any_failed else "failed" + if final_status in ("completed", "failed"): + action = final_status + elif any_dispatched: + action = "dispatched" + else: + action = "no_change" return make_poll_result( run=run, - eval_type="tts", + eval_type=EVAL_TYPE_LABEL, previous_status=previous_status, current_status=final_status, action=action, error=error_message, ) + + +def _fail_pending_results( + *, session: Session, run_id: int, model: str, error_message: str +) -> None: + """Mark one model's PENDING rows of a run FAILED.""" + pending = get_pending_results_for_run( + session=session, run_id=run_id, provider=model + ) + if not pending: + return + logger.warning( + f"[_fail_pending_results] Failing pending results | run_id={run_id}, " + f"model={model}, count={len(pending)}, error={error_message}" + ) + for result in pending: + update_tts_result( + session=session, + result_id=result.id, + status=JobStatus.FAILED.value, + error_message=error_message, + ) + session.commit() diff --git a/backend/app/crud/tts_evaluations/run.py b/backend/app/crud/tts_evaluations/run.py index 881491d29..2f9831288 100644 --- a/backend/app/crud/tts_evaluations/run.py +++ b/backend/app/crud/tts_evaluations/run.py @@ -6,7 +6,9 @@ from sqlmodel import Session, func, select from app.core.util import now +from app.crud.tts_evaluations.result import count_results_by_status from app.models import EvaluationRun +from app.models.job import JobStatus from app.models.stt_evaluation import EvaluationType from app.models.tts_evaluation import TTSEvaluationRunPublic @@ -219,3 +221,55 @@ def update_tts_run( ) return run + + +def finalize_tts_run_status( + *, + session: Session, + run_id: int, + error_message: str | None = None, +) -> EvaluationRun | None: + """Recompute a TTS run's status from its result rows and persist it. + + `processing` while any row is PENDING, `failed` when no row succeeded (every + model failed), otherwise `completed`. Only the cron calls this. + """ + statement = ( + select(EvaluationRun).where(EvaluationRun.id == run_id).with_for_update() + ) + run = session.exec(statement).one_or_none() + + if not run: + logger.warning(f"[finalize_tts_run_status] Run not found | run_id: {run_id}") + return None + + status_counts = count_results_by_status(session=session, run_id=run_id) + pending = status_counts.get(JobStatus.PENDING.value, 0) + success = status_counts.get(JobStatus.SUCCESS.value, 0) + failed = status_counts.get(JobStatus.FAILED.value, 0) + + if pending > 0: + run.status = "processing" + elif success == 0: + run.status = "failed" + else: + run.status = "completed" + + if error_message is None and failed > 0: + error_message = f"{failed} synthesis(es) failed" + if error_message is not None: + run.error_message = error_message + + run.updated_at = now() + + session.add(run) + session.commit() + session.refresh(run) + + logger.info( + f"[finalize_tts_run_status] Run finalized | run_id: {run_id}, " + f"status: {run.status}, pending: {pending}, success: {success}, " + f"failed: {failed}" + ) + + return run diff --git a/backend/app/services/tts_evaluations/batch_job.py b/backend/app/services/tts_evaluations/batch_job.py index ad27d8d51..97147b0b8 100644 --- a/backend/app/services/tts_evaluations/batch_job.py +++ b/backend/app/services/tts_evaluations/batch_job.py @@ -1,16 +1,38 @@ -"""Celery task function for TTS evaluation batch submission.""" +"""Celery task function for TTS evaluation submission. + +Gemini models are submitted to the Gemini Batch API; each Sarvam/ElevenLabs model +gets its own `run_tts_sync_generation` Celery task. The cron sets the final run status. +""" import logging +from typing import Any -from gevent import Timeout +from asgi_correlation_id import correlation_id +from celery import Task from celery.exceptions import SoftTimeLimitExceeded +from gevent import Timeout from sqlmodel import Session +from app.celery.utils import start_tts_sync_generation +from app.core.cloud.storage import get_cloud_storage from app.core.db import engine from app.crud.tts_evaluations.batch import start_tts_evaluation_batch from app.crud.tts_evaluations.dataset import get_tts_dataset_by_id -from app.crud.tts_evaluations.result import create_tts_results +from app.crud.tts_evaluations.result import ( + create_tts_results, + get_pending_results_for_run, + update_tts_result, +) from app.crud.tts_evaluations.run import get_tts_run_by_id, update_tts_run +from app.models import Language +from app.models.job import JobStatus +from app.models.tts_evaluation import TTSResult +from app.services.tts_evaluations.constants import ( + DEFAULT_TTS_LANGUAGE, + GEMINI_TTS_MODELS, + SYNC_TTS_MODELS, + LOCALE_TO_BCP47, +) from app.services.tts_evaluations.dataset import get_sample_texts_from_dataset logger = logging.getLogger(__name__) @@ -20,16 +42,17 @@ def execute_batch_submission( project_id: int, job_id: str, task_id: str, - task_instance, + task_instance: Task, organization_id: int, dataset_id: int, models: list[str], - **kwargs, -) -> dict: - """Execute TTS evaluation batch submission in a Celery worker. + **kwargs: Any, +) -> dict[str, Any]: + """Execute TTS evaluation submission in a Celery worker. - Handles result record creation, JSONL creation, Gemini file upload, - and batch job creation. + Creates result rows, submits the Gemini batch and queues one sync task per + Sarvam/ElevenLabs model. A model that fails to start has its rows marked + FAILED; the other models carry on. Args: project_id: Project ID @@ -44,6 +67,8 @@ def execute_batch_submission( dict: Result summary with batch job info """ run_id = int(job_id) + batch_models = [model for model in models if model in GEMINI_TTS_MODELS] + sync_models = [model for model in models if model in SYNC_TTS_MODELS] logger.info( f"[execute_batch_submission] Starting | " @@ -51,15 +76,14 @@ def execute_batch_submission( f"celery_task_id: {task_id}" ) - with Session(engine) as session: - try: + try: + with Session(engine) as session: run = get_tts_run_by_id( session=session, run_id=run_id, org_id=organization_id, project_id=project_id, ) - if not run: logger.warning( f"[execute_batch_submission] Run not found | run_id: {run_id}" @@ -72,7 +96,6 @@ def execute_batch_submission( org_id=organization_id, project_id=project_id, ) - if not dataset: logger.warning( f"[execute_batch_submission] Dataset not found | " @@ -86,70 +109,141 @@ def execute_batch_submission( ) return {"success": False, "error": "Dataset not found"} - sample_texts = get_sample_texts_from_dataset(session, dataset, project_id) + language_code = DEFAULT_TTS_LANGUAGE + if dataset.language_id is not None: + language = session.get(Language, dataset.language_id) + if language: + language_code = LOCALE_TO_BCP47.get( + language.locale, DEFAULT_TTS_LANGUAGE + ) - if not sample_texts: - logger.warning( - f"[execute_batch_submission] No samples found | " - f"run_id: {run_id}, dataset_id: {dataset_id}" - ) - update_tts_run( - session=session, - run_id=run_id, - status="failed", - error_message="No samples found for dataset", - ) - return {"success": False, "error": "No samples found"} + storage = get_cloud_storage(session=session, project_id=project_id) + + # Read from S3 storage with no db session open + sample_texts = get_sample_texts_from_dataset(storage, dataset) + if not sample_texts: + logger.warning( + f"[execute_batch_submission] No samples found | " + f"run_id: {run_id}, dataset_id: {dataset_id}" + ) + _mark_run_failed( + run_id=run_id, error_message="No samples found for dataset" + ) + return {"success": False, "error": "No samples found"} - # Create result records for each sample text and model + with Session(engine) as session: results = create_tts_results( session=session, sample_texts=sample_texts, - evaluation_run_id=run.id, + evaluation_run_id=run_id, org_id=organization_id, project_id=project_id, models=models, ) - - batch_result = start_tts_evaluation_batch( - session=session, - run=run, - results=results, - org_id=organization_id, - project_id=project_id, - ) - - logger.info( - f"[execute_batch_submission] Batch submitted | " - f"run_id: {run_id}, " - f"batch_jobs: {list(batch_result.get('batch_jobs', {}).keys())}" - ) - - return batch_result - - except (Timeout, SoftTimeLimitExceeded) as err: - timeout_err = TimeoutError("Task exceeded soft time limit") - logger.warning( - f"[execute_batch_submission] TTS batch submission timed out | run_id={run_id}" - ) - update_tts_run( - session=session, - run_id=run_id, - status="failed", - error_message=str(timeout_err), - ) - raise - - except Exception as e: - logger.error( - f"[execute_batch_submission] Batch submission failed | " - f"run_id: {run_id}, error: {str(e)}", - exc_info=True, - ) - update_tts_run( - session=session, - run_id=run_id, - status="failed", - error_message=str(e), - ) - return {"success": False, "error": str(e)} + result_count = len(results) + update_tts_run(session=session, run_id=run_id, status="processing") + + # model -> error; these models' rows are marked FAILED once all models are started. + failed_models: dict[str, str] = {} + + # Queue sync tasks first: queuing is instant, while Gemini submission + # (JSONL upload + batch create) takes seconds per model. + trace_id = correlation_id.get() or "N/A" + for model in sync_models: + try: + start_tts_sync_generation( + project_id=project_id, + job_id=str(run_id), + trace_id=trace_id, + organization_id=organization_id, + model=model, + language_code=language_code, + ) + except Exception as e: + logger.error( + f"[execute_batch_submission] Failed to queue sync generation | " + f"run_id: {run_id}, model: {model}, error: {str(e)}", + exc_info=True, + ) + failed_models[model] = f"Failed to queue synthesis: {str(e)}" + + batch_result: dict[str, Any] = {} + if batch_models: + with Session(engine) as session: + batch_results: list[TTSResult] = [] + for model in batch_models: + batch_results.extend( + get_pending_results_for_run( + session=session, run_id=run_id, provider=model + ) + ) + try: + batch_result = start_tts_evaluation_batch( + session=session, + run=run, + results=batch_results, + org_id=organization_id, + project_id=project_id, + ) + except Exception as e: + logger.error( + f"[execute_batch_submission] Gemini batch submission failed | " + f"run_id: {run_id}, error: {str(e)}", + exc_info=True, + ) + for model in batch_models: + failed_models[model] = f"Batch submission failed: {str(e)}" + + if failed_models: + with Session(engine) as session: + for model, error_message in failed_models.items(): + pending = get_pending_results_for_run( + session=session, run_id=run_id, provider=model + ) + for result in pending: + update_tts_result( + session=session, + result_id=result.id, + status=JobStatus.FAILED.value, + error_message=error_message, + ) + session.commit() + + logger.info( + f"[execute_batch_submission] Submitted | run_id: {run_id}, " + f"batch_jobs: {list(batch_result.get('batch_jobs', {}).keys())}, " + f"sync_models: {sync_models}, failed_models: {list(failed_models.keys())}" + ) + + return { + "success": True, + "run_id": run_id, + "batch_jobs": batch_result.get("batch_jobs", {}), + "result_count": result_count, + } + + except (Timeout, SoftTimeLimitExceeded): + logger.warning( + f"[execute_batch_submission] TTS batch submission timed out | run_id={run_id}" + ) + _mark_run_failed(run_id=run_id, error_message="Task exceeded soft time limit") + raise + + except Exception as e: + logger.error( + f"[execute_batch_submission] Batch submission failed | " + f"run_id: {run_id}, error: {str(e)}", + exc_info=True, + ) + _mark_run_failed(run_id=run_id, error_message=str(e)) + return {"success": False, "error": str(e)} + + +def _mark_run_failed(*, run_id: int, error_message: str) -> None: + with Session(engine) as session: + update_tts_run( + session=session, + run_id=run_id, + status="failed", + error_message=error_message, + ) diff --git a/backend/app/services/tts_evaluations/batch_result_processing.py b/backend/app/services/tts_evaluations/batch_result_processing.py index a2dc3c208..3613c58bb 100644 --- a/backend/app/services/tts_evaluations/batch_result_processing.py +++ b/backend/app/services/tts_evaluations/batch_result_processing.py @@ -11,18 +11,21 @@ from gevent import Timeout from celery.exceptions import SoftTimeLimitExceeded -from sqlmodel import Session, select +from sqlmodel import Session from app.core.batch import BATCH_KEY, GeminiBatchProvider, GeminiClient from app.core.cloud.storage import get_cloud_storage from app.core.db import engine from app.core.storage_utils import upload_to_object_store from app.core.util import now -from app.crud.tts_evaluations.result import count_results_by_status -from app.crud.tts_evaluations.run import update_tts_run +from app.crud.tts_evaluations.result import ( + get_pending_results_for_run, + update_tts_result, +) from app.models.job import JobStatus from app.models.tts_evaluation import TTSResult from app.core.audio_utils import calculate_duration, pcm_to_wav +from app.services.tts_evaluations.constants import TTS_AUDIO_SUBDIRECTORY logger = logging.getLogger(__name__) @@ -62,211 +65,192 @@ def execute_tts_result_processing( f"provider={tts_provider}, celery_task_id={task_id}" ) - with Session(engine) as session: - try: - # Initialize Gemini client and batch provider + try: + with Session(engine) as session: gemini_client = GeminiClient.from_credentials( session=session, org_id=organization_id, project_id=project_id, ) - batch_provider = GeminiBatchProvider(client=gemini_client.client) - - # Get cloud storage for S3 uploads storage = get_cloud_storage(session=session, project_id=project_id) - # Download batch results - results = batch_provider.download_batch_results(provider_batch_id) - + # Only PENDING rows, so reprocessing a batch skips rows already finished. + pending = get_pending_results_for_run( + session=session, run_id=evaluation_run_id, provider=tts_provider + ) + result_map: dict[int, TTSResult] = {} + for result in pending: + result_map[result.id] = result + + logger.info( + f"[execute_tts_result_processing] Pre-fetched result records | " + f"run_id={evaluation_run_id}, provider={tts_provider}, " + f"count={len(result_map)}" + ) + + # Download, conversion and uploads run with no session open; the detached + # rows are modified in memory and written back in one short session below. + batch_provider = GeminiBatchProvider(client=gemini_client.client) + results = batch_provider.download_batch_results(provider_batch_id) + + logger.info( + f"[execute_tts_result_processing] Got batch results | " + f"run_id={evaluation_run_id}, result_count={len(results)}" + ) + + if results: + first = results[0] + resp = first.get("response") or {} + resp_keys = ( + list(resp.keys()) if isinstance(resp, dict) else type(resp).__name__ + ) logger.info( - f"[execute_tts_result_processing] Got batch results | " - f"run_id={evaluation_run_id}, result_count={len(results)}" + f"[execute_tts_result_processing] First result structure | " + f"keys={list(first.keys())}, response_keys={resp_keys}" ) - if results: - first = results[0] - resp = first.get("response") or {} - resp_keys = ( - list(resp.keys()) if isinstance(resp, dict) else type(resp).__name__ - ) - logger.info( - f"[execute_tts_result_processing] First result structure | " - f"keys={list(first.keys())}, response_keys={resp_keys}" + processed_count = 0 + failed_count = 0 + + for batch_result in results: + custom_id = batch_result[BATCH_KEY] + try: + result_id = int(custom_id) + except (ValueError, TypeError): + logger.warning( + f"[execute_tts_result_processing] Invalid {BATCH_KEY} | " + f"run_id={evaluation_run_id}, {BATCH_KEY}={custom_id}" ) + failed_count += 1 + continue - # Pre-fetch all TTSResult records for this run+provider - prefetch_stmt = select(TTSResult).where( - TTSResult.evaluation_run_id == evaluation_run_id, - TTSResult.provider == tts_provider, - ) - result_map: dict[int, TTSResult] = { - r.id: r for r in session.exec(prefetch_stmt).all() # type: ignore[misc] - } - - logger.info( - f"[execute_tts_result_processing] Pre-fetched result records | " - f"run_id={evaluation_run_id}, provider={tts_provider}, " - f"count={len(result_map)}" - ) + result_record = result_map.get(result_id) - processed_count = 0 - failed_count = 0 + if not result_record: + logger.warning( + f"[execute_tts_result_processing] Result record not found or " + f"not pending | result_id={result_id}" + ) + failed_count += 1 + continue - for batch_result in results: - custom_id = batch_result[BATCH_KEY] + if batch_result.get("response"): try: - result_id = int(custom_id) - except (ValueError, TypeError): - logger.warning( - f"[execute_tts_result_processing] Invalid {BATCH_KEY} | " - f"run_id={evaluation_run_id}, {BATCH_KEY}={custom_id}" - ) - failed_count += 1 - continue - - # Look up result record from pre-fetched map - result_record = result_map.get(result_id) + audio_b64 = _extract_audio_from_response(batch_result["response"]) - if not result_record: - logger.warning( - f"[execute_tts_result_processing] Result record not found | " - f"result_id={result_id}" - ) - failed_count += 1 - continue + if not audio_b64: + result_record.status = JobStatus.FAILED.value + result_record.error_message = "No audio data in response" + result_record.updated_at = now() + failed_count += 1 + continue - if batch_result.get("response"): - try: - audio_b64 = _extract_audio_from_response( - batch_result["response"] - ) + # Without validate=True, garbage decodes to b"" and is stored as a 0s "success". + pcm_data = base64.b64decode(audio_b64, validate=True) - if not audio_b64: - result_record.status = JobStatus.FAILED.value - result_record.error_message = "No audio data in response" - result_record.updated_at = now() - failed_count += 1 - continue - - # Decode base64 -> raw PCM bytes - pcm_data = base64.b64decode(audio_b64) - - # Wrap in WAV container - wav_data = pcm_to_wav(pcm_data) - - # Calculate duration - duration = calculate_duration(len(pcm_data)) - - # Upload WAV to S3 - audio_filename = f"{uuid.uuid4()}.wav" - audio_url = upload_to_object_store( - storage=storage, - content=wav_data, - filename=audio_filename, - subdirectory="evaluations/tts/audio", - content_type="audio/wav", - ) + wav_data = pcm_to_wav(pcm_data) + duration = calculate_duration(len(pcm_data)) - if not audio_url: - result_record.status = JobStatus.FAILED.value - result_record.error_message = ( - "Audio upload to object store failed" - ) - result_record.updated_at = now() - failed_count += 1 - continue - - # Update result - result_record.object_store_url = audio_url - result_record.metadata_ = { - "duration_seconds": round(duration, 3), - "size_bytes": len(wav_data), - } - result_record.status = JobStatus.SUCCESS.value - result_record.updated_at = now() - processed_count += 1 + audio_filename = f"{uuid.uuid4()}.wav" + audio_url = upload_to_object_store( + storage=storage, + content=wav_data, + filename=audio_filename, + subdirectory=TTS_AUDIO_SUBDIRECTORY, + content_type="audio/wav", + ) - except Exception as audio_err: - logger.warning( - f"[execute_tts_result_processing] Audio processing failed | " - f"result_id={result_id}, error={str(audio_err)}" - ) + if not audio_url: result_record.status = JobStatus.FAILED.value result_record.error_message = ( - f"Audio processing failed: {str(audio_err)}" + "Audio upload to object store failed" ) result_record.updated_at = now() failed_count += 1 - else: + continue + + result_record.object_store_url = audio_url + result_record.metadata_ = { + "duration_seconds": round(duration, 3), + "size_bytes": len(wav_data), + } + result_record.status = JobStatus.SUCCESS.value + result_record.updated_at = now() + processed_count += 1 + + except Exception as audio_err: + logger.warning( + f"[execute_tts_result_processing] Audio processing failed | " + f"result_id={result_id}, error={str(audio_err)}" + ) result_record.status = JobStatus.FAILED.value - result_record.error_message = batch_result.get( - "error", "Unknown error" + result_record.error_message = ( + f"Audio processing failed: {str(audio_err)}" ) result_record.updated_at = now() failed_count += 1 - + else: + result_record.status = JobStatus.FAILED.value + result_record.error_message = batch_result.get("error", "Unknown error") + result_record.updated_at = now() + failed_count += 1 + + with Session(engine) as session: + session.add_all(list(result_map.values())) session.commit() - # Finalize run status - status_counts = count_results_by_status( - session=session, run_id=evaluation_run_id - ) - pending = status_counts.get(JobStatus.PENDING.value, 0) - total_failed = status_counts.get(JobStatus.FAILED.value, 0) - - final_status = "completed" if pending == 0 else "processing" - error_message = ( - f"{total_failed} synthesis(es) failed" if total_failed > 0 else None - ) - - update_tts_run( + logger.info( + f"[execute_tts_result_processing] Completed | " + f"run_id={evaluation_run_id}, provider={tts_provider}, " + f"processed={processed_count}, failed={failed_count}" + ) + + return { + "success": True, + "run_id": evaluation_run_id, + "processed": processed_count, + "failed": failed_count, + } + + except (Timeout, SoftTimeLimitExceeded): + logger.warning( + f"[execute_tts_result_processing] TTS result processing timed out | run_id={evaluation_run_id}" + ) + _fail_pending_results( + run_id=evaluation_run_id, + model=tts_provider, + error_message="Result processing exceeded soft time limit", + ) + raise + + except Exception as e: + logger.error( + f"[execute_tts_result_processing] Failed | " + f"run_id={evaluation_run_id}, error={str(e)}", + exc_info=True, + ) + _fail_pending_results( + run_id=evaluation_run_id, + model=tts_provider, + error_message=f"Result processing failed: {str(e)}", + ) + return {"success": False, "error": str(e)} + + +def _fail_pending_results(*, run_id: int, model: str, error_message: str) -> None: + """Fail only this model's leftover rows so the other models still decide the run.""" + with Session(engine) as session: + pending = get_pending_results_for_run( + session=session, run_id=run_id, provider=model + ) + for result in pending: + update_tts_result( session=session, - run_id=evaluation_run_id, - status=final_status, + result_id=result.id, + status=JobStatus.FAILED.value, error_message=error_message, ) - - logger.info( - f"[execute_tts_result_processing] Completed | " - f"run_id={evaluation_run_id}, provider={tts_provider}, " - f"processed={processed_count}, failed={failed_count}, " - f"run_status={final_status}" - ) - - return { - "success": True, - "run_id": evaluation_run_id, - "processed": processed_count, - "failed": failed_count, - "run_status": final_status, - } - - except (Timeout, SoftTimeLimitExceeded) as err: - timeout_err = TimeoutError("Task exceeded soft time limit") - logger.warning( - f"[execute_tts_result_processing] TTS result processing timed out | run_id={evaluation_run_id}" - ) - update_tts_run( - session=session, - run_id=evaluation_run_id, - status="failed", - error_message=str(timeout_err), - ) - raise - - except Exception as e: - logger.error( - f"[execute_tts_result_processing] Failed | " - f"run_id={evaluation_run_id}, error={str(e)}", - exc_info=True, - ) - update_tts_run( - session=session, - run_id=evaluation_run_id, - status="failed", - error_message=f"Result processing failed: {str(e)}", - ) - return {"success": False, "error": str(e)} + session.commit() def _extract_audio_from_response(response: dict[str, Any]) -> str | None: diff --git a/backend/app/services/tts_evaluations/constants.py b/backend/app/services/tts_evaluations/constants.py index 169b9ff1e..aa483800f 100644 --- a/backend/app/services/tts_evaluations/constants.py +++ b/backend/app/services/tts_evaluations/constants.py @@ -1,13 +1,54 @@ -"""Shared constants for TTS evaluation services.""" +"""Shared constants for TTS evaluation services. + +Kept free of `app.models` imports: `app.models.tts_evaluation` imports this module, +so importing models here would create a cycle. +""" + +# Gemini models go through the Gemini Batch API; Sarvam/ElevenLabs are called directly. +GEMINI_TTS_MODELS = ["gemini-2.5-pro-preview-tts", "gemini-3.1-flash-tts-preview"] +SARVAM_TTS_MODELS = ["bulbul:v3"] +ELEVENLABS_TTS_MODELS = ["eleven_v4", "eleven_v3"] + +SYNC_TTS_MODELS = SARVAM_TTS_MODELS + ELEVENLABS_TTS_MODELS # Supported TTS models for evaluation -SUPPORTED_TTS_MODELS = ["gemini-2.5-pro-preview-tts"] +SUPPORTED_TTS_MODELS = GEMINI_TTS_MODELS + SYNC_TTS_MODELS # Default TTS model DEFAULT_TTS_MODEL = "gemini-2.5-pro-preview-tts" # Default voice configuration DEFAULT_VOICE_NAME = "Kore" +SARVAM_SPEAKER = "shubh" +# "Rahul S - Hindi Conversational": closest match to Sarvam's male speaker for side-by-side review. +ELEVENLABS_VOICE_ID = "2cdvnKJ5TZi631y5PN1s" # Default style prompt for TTS synthesis DEFAULT_STYLE_PROMPT = "Read in a calm, professional customer service tone" + +DEFAULT_TTS_LANGUAGE = "en-IN" + +# Dataset locale (ISO 639-1) -> BCP-47 tag; Sarvam spells Odia "od-IN", not "or-IN". +LOCALE_TO_BCP47 = { + "en": "en-IN", + "hi": "hi-IN", + "bn": "bn-IN", + "ta": "ta-IN", + "te": "te-IN", + "mr": "mr-IN", + "gu": "gu-IN", + "kn": "kn-IN", + "ml": "ml-IN", + "pa": "pa-IN", + "or": "od-IN", +} + +# Rows synthesized at once per sync task; kept low because eleven_v3 and eleven_v4 +# share one ElevenLabs key and its concurrent-request limit. +TTS_SYNC_MAX_WORKERS = 2 + +TTS_AUDIO_SUBDIRECTORY = "evaluations/tts/audio" + +# Once a run is older than the Celery hard time limit plus this grace, no sync worker can +# still be writing its rows, so the cron fails any that are left PENDING. +TTS_SYNC_STALE_GRACE_SECONDS = 600 diff --git a/backend/app/services/tts_evaluations/dataset.py b/backend/app/services/tts_evaluations/dataset.py index 9db8aae5b..02d56d0ec 100644 --- a/backend/app/services/tts_evaluations/dataset.py +++ b/backend/app/services/tts_evaluations/dataset.py @@ -5,9 +5,11 @@ import logging from typing import Any +from celery.exceptions import SoftTimeLimitExceeded from sqlmodel import Session from app.core.cloud import get_cloud_storage +from app.core.cloud.storage import CloudStorage from app.core.storage_utils import ( generate_timestamped_filename, upload_to_object_store, @@ -163,16 +165,16 @@ def _samples_to_csv(samples: list[TTSSampleCreate]) -> bytes: def get_sample_texts_from_dataset( - session: Session, + storage: CloudStorage, dataset: "EvaluationDataset", - project_id: int, ) -> list[str]: """Extract sample texts from a TTS dataset's CSV in S3. + Takes storage rather than a session so callers can read S3 with no DB session open. + Args: - session: Database session + storage: Cloud storage for the dataset's project dataset: The evaluation dataset record - project_id: Project ID Returns: List of text strings @@ -185,10 +187,12 @@ def get_sample_texts_from_dataset( return [] try: - storage = get_cloud_storage(session=session, project_id=project_id) csv_bytes = storage.stream(dataset.object_store_url).read() samples = parse_tts_samples_from_csv(csv_bytes) return [s["text"] for s in samples] + except SoftTimeLimitExceeded: + # Let the task's timeout handler report it instead of "No samples found". + raise except Exception as e: logger.error( f"[get_sample_texts_from_dataset] Failed to load CSV | " diff --git a/backend/app/services/tts_evaluations/sync_generation.py b/backend/app/services/tts_evaluations/sync_generation.py new file mode 100644 index 000000000..7557c38dc --- /dev/null +++ b/backend/app/services/tts_evaluations/sync_generation.py @@ -0,0 +1,287 @@ +"""Celery task function for synchronous TTS synthesis (Sarvam, ElevenLabs). + +One task synthesizes every PENDING row of one model, TTS_SYNC_MAX_WORKERS rows at a +time. Provider calls and uploads run with no DB session open so a slow provider can't +pin pooled connections. The task only writes result rows; the cron sets the run status. +""" + +import base64 +import io +import logging +import uuid +import wave +from concurrent.futures import Future, ThreadPoolExecutor, as_completed +from typing import Any, Required, TypedDict + +from celery import Task +from celery.exceptions import SoftTimeLimitExceeded +from elevenlabs import ElevenLabs +from gevent import Timeout +from sarvamai import SarvamAI +from sqlmodel import Session + +from app.core.audio_utils import calculate_duration, pcm_to_wav +from app.core.cloud.storage import get_cloud_storage +from app.core.db import engine +from app.core.providers import Provider +from app.core.storage_utils import upload_to_object_store +from app.crud.credentials import get_provider_credential +from app.crud.tts_evaluations.result import ( + get_pending_results_for_run, + update_tts_result, +) +from app.models.job import JobStatus +from app.models.llm.constants import BCP47_TO_ELEVENLABS_LANG +from app.services.llm.providers.eleven_ai import ElevenlabsAIProvider +from app.services.llm.providers.sarvam_ai import SarvamAIProvider +from app.services.tts_evaluations.constants import ( + ELEVENLABS_TTS_MODELS, + ELEVENLABS_VOICE_ID, + SARVAM_SPEAKER, + SARVAM_TTS_MODELS, + TTS_AUDIO_SUBDIRECTORY, + TTS_SYNC_MAX_WORKERS, +) + +logger = logging.getLogger(__name__) + +# Every stored clip is 24 kHz mono 16-bit +TTS_SAMPLE_RATE = 24000 + + +class TTSResultFields(TypedDict, total=False): + """Keyword arguments for `update_tts_result` produced for one synthesized row.""" + + status: Required[str] + object_store_url: str + metadata: dict[str, float | int] + error_message: str + + +def synthesize_sarvam( + client: SarvamAI, text: str, model: str, language_code: str +) -> bytes: + """Synthesize text with Sarvam; returns raw 24 kHz mono 16-bit PCM.""" + response = client.text_to_speech.convert( + text=text, + model=model, + target_language_code=language_code, + speaker=SARVAM_SPEAKER, + speech_sample_rate=TTS_SAMPLE_RATE, + ) + if not response.audios: + raise ValueError("Sarvam returned no audio") + + # Long inputs come back as several base64 chunks, each a full WAV file; strip + # every header so later RIFF headers don't end up in the PCM as clicks. + pcm = b"" + for audio_b64 in response.audios: + with wave.open(io.BytesIO(base64.b64decode(audio_b64)), "rb") as wav_file: + pcm += wav_file.readframes(wav_file.getnframes()) + return pcm + + +def synthesize_elevenlabs( + client: ElevenLabs, text: str, model: str, language_code: str +) -> bytes: + """Synthesize text with ElevenLabs; returns raw 24 kHz mono 16-bit PCM.""" + chunks = client.text_to_speech.convert( + voice_id=ELEVENLABS_VOICE_ID, + text=text, + model_id=model, + output_format=f"pcm_{TTS_SAMPLE_RATE}", + # Unmapped languages (e.g. Odia) are left to ElevenLabs' auto-detect. + language_code=BCP47_TO_ELEVENLABS_LANG.get(language_code), + ) + pcm = b"".join(chunks) + if not pcm: + raise ValueError("ElevenLabs returned no audio") + return pcm + + +def execute_tts_sync_generation( + project_id: int, + job_id: str, + task_id: str, + task_instance: Task, + organization_id: int, + model: str, + language_code: str, + **kwargs: Any, +) -> dict[str, Any]: + """Synthesize every PENDING Sarvam/ElevenLabs result of one model in a Celery worker. + + Only PENDING rows are picked up, so a redelivered task skips rows already finished. + + Args: + project_id: Project ID + job_id: Evaluation run ID (as string) + task_id: Celery task ID + task_instance: Celery task instance + organization_id: Organization ID + model: Sarvam or ElevenLabs model + language_code: BCP-47 tag resolved from the dataset (e.g. hi-IN) + + Returns: + dict: Summary with processed/failed counts + """ + run_id = int(job_id) + if model in SARVAM_TTS_MODELS: + provider = Provider.SARVAMAI + elif model in ELEVENLABS_TTS_MODELS: + provider = Provider.ELEVENLABS + else: + error_message = f"Unsupported sync TTS model: {model}" + logger.error(f"[execute_tts_sync_generation] {error_message} | run_id={run_id}") + _fail_pending_results(run_id=run_id, model=model, error_message=error_message) + return {"success": False, "run_id": run_id, "error": error_message} + + logger.info( + f"[execute_tts_sync_generation] Starting | run_id={run_id}, model={model}, " + f"language={language_code}, celery_task_id={task_id}" + ) + + try: + with Session(engine) as session: + pending = get_pending_results_for_run( + session=session, run_id=run_id, provider=model + ) + # Plain tuples: the ORM rows detach when this session closes. + pending_rows: list[tuple[int, str]] = [] + for result in pending: + pending_rows.append((result.id, result.sample_text)) + + credentials = get_provider_credential( + session=session, + org_id=organization_id, + project_id=project_id, + provider=provider.value, + ) + storage = get_cloud_storage(session=session, project_id=project_id) + + if not isinstance(credentials, dict): + raise ValueError(f"{provider.value} credentials are not configured") + + if model in SARVAM_TTS_MODELS: + client = SarvamAIProvider.create_client(credentials) + else: + client = ElevenlabsAIProvider.create_client(credentials) + + # keep the function local to use it easily with + # internal threadpool exuector + + def synthesize_and_upload(result_id: int, text: str) -> TTSResultFields: + """Return `update_tts_result` fields; never raises, so one row can't stop the rest.""" + try: + if model in SARVAM_TTS_MODELS: + pcm = synthesize_sarvam(client, text, model, language_code) + else: + pcm = synthesize_elevenlabs(client, text, model, language_code) + wav = pcm_to_wav(pcm) + audio_url = upload_to_object_store( + storage=storage, + content=wav, + filename=f"{uuid.uuid4()}.wav", + subdirectory=TTS_AUDIO_SUBDIRECTORY, + content_type="audio/wav", + ) + except Exception as e: + logger.warning( + f"[execute_tts_sync_generation] Synthesis failed | " + f"run_id={run_id}, model={model}, result_id={result_id}, " + f"error={str(e)}" + ) + return { + "status": JobStatus.FAILED.value, + "error_message": f"{provider.value} synthesis failed: {str(e)}", + } + + if not audio_url: + return { + "status": JobStatus.FAILED.value, + "error_message": "Audio upload to object store failed", + } + return { + "status": JobStatus.SUCCESS.value, + "object_store_url": audio_url, + "metadata": { + "duration_seconds": round(calculate_duration(len(pcm)), 3), + "size_bytes": len(wav), + }, + } + + processed = 0 + failed = 0 + # Shut down by hand: a `with` block would wait for in-flight provider calls on timeout. + executor = ThreadPoolExecutor(max_workers=TTS_SYNC_MAX_WORKERS) + try: + futures: dict[Future[TTSResultFields], int] = {} + for result_id, text in pending_rows: + future = executor.submit(synthesize_and_upload, result_id, text) + futures[future] = result_id + + for future in as_completed(futures): + fields = future.result() + with Session(engine) as session: + update_tts_result( + session=session, result_id=futures[future], **fields + ) + session.commit() + + if fields["status"] == JobStatus.SUCCESS.value: + processed += 1 + else: + failed += 1 + finally: + executor.shutdown(wait=False, cancel_futures=True) + + logger.info( + f"[execute_tts_sync_generation] Completed | run_id={run_id}, model={model}, " + f"processed={processed}, failed={failed}" + ) + return { + "success": True, + "run_id": run_id, + "processed": processed, + "failed": failed, + } + + except (Timeout, SoftTimeLimitExceeded): + logger.warning( + f"[execute_tts_sync_generation] Timed out | run_id={run_id}, model={model}" + ) + _fail_pending_results( + run_id=run_id, + model=model, + error_message="Synthesis timed out before this sample was processed", + ) + raise + + except Exception as e: + logger.error( + f"[execute_tts_sync_generation] Failed | run_id={run_id}, model={model}, " + f"error={str(e)}", + exc_info=True, + ) + _fail_pending_results( + run_id=run_id, + model=model, + error_message=f"{provider.value} synthesis failed: {str(e)}", + ) + return {"success": False, "run_id": run_id, "error": str(e)} + + +def _fail_pending_results(*, run_id: int, model: str, error_message: str) -> None: + """Mark this model's leftover PENDING rows FAILED; nothing re-drives a sync task.""" + with Session(engine) as session: + pending = get_pending_results_for_run( + session=session, run_id=run_id, provider=model + ) + for result in pending: + update_tts_result( + session=session, + result_id=result.id, + status=JobStatus.FAILED.value, + error_message=error_message, + ) + session.commit() diff --git a/backend/app/tests/api/routes/test_tts_evaluation.py b/backend/app/tests/api/routes/test_tts_evaluation.py index ae9ade1c8..1add26532 100644 --- a/backend/app/tests/api/routes/test_tts_evaluation.py +++ b/backend/app/tests/api/routes/test_tts_evaluation.py @@ -603,6 +603,40 @@ def test_start_evaluation_multiple_models_total_items( # 5 samples × 1 model assert data["total_items"] == 5 + @patch("app.api.routes.tts_evaluations.evaluation.start_tts_batch_submission") + def test_start_evaluation_mixed_models_reports_voice_per_model( + self, + mock_start_job: MagicMock, + client: TestClient, + user_api_key_header: dict[str, str], + test_dataset_with_samples: EvaluationDataset, + ) -> None: + mock_start_job.return_value = "mock-celery-task-id" + models = ["gemini-2.5-pro-preview-tts", "bulbul:v3", "eleven_v3"] + + response = client.post( + "/api/v1/evaluations/tts/runs", + json={ + "run_name": "mixed_models_run", + "dataset_id": test_dataset_with_samples.id, + "models": models, + }, + headers=user_api_key_header, + ) + + assert response.status_code == 200, response.text + data = response.json()["data"] + assert data["models"] == models + assert data["total_items"] == 9 + assert data["run_metadata"]["voices"] == { + "gemini-2.5-pro-preview-tts": "Kore", + "bulbul:v3": "shubh", + "eleven_v3": "2cdvnKJ5TZi631y5PN1s", + } + # legacy single-voice key kept for existing clients + assert data["run_metadata"]["voice_name"] == "Kore" + assert mock_start_job.call_args.kwargs["models"] == models + @patch("app.api.routes.tts_evaluations.evaluation.start_tts_batch_submission") def test_start_evaluation_celery_failure( self, diff --git a/backend/app/tests/celery/test_tts_tasks.py b/backend/app/tests/celery/test_tts_tasks.py new file mode 100644 index 000000000..803e134d6 --- /dev/null +++ b/backend/app/tests/celery/test_tts_tasks.py @@ -0,0 +1,64 @@ +from typing import Any +from unittest.mock import MagicMock, patch + +from asgi_correlation_id import correlation_id + +from app.celery.tasks import job_execution +from app.celery.utils import start_tts_sync_generation + +SYNC_KWARGS: dict[str, Any] = { + "organization_id": 3, + "model": "bulbul:v3", + "language_code": "hi-IN", +} + + +def test_start_tts_sync_generation_enqueues_task_kwargs() -> None: + # apply_async is the broker boundary; nothing is published in tests. + with patch.object( + job_execution.run_tts_sync_generation, + "apply_async", + return_value=MagicMock(id="celery-123"), + ) as apply_async: + task_id = start_tts_sync_generation( + project_id=7, job_id="42", trace_id="trace-abc", **SYNC_KWARGS + ) + + assert task_id == "celery-123" + assert apply_async.call_args.kwargs["kwargs"] == { + "project_id": 7, + "job_id": "42", + "trace_id": "trace-abc", + **SYNC_KWARGS, + } + + +def test_run_tts_sync_generation_forwards_kwargs_and_trace() -> None: + seen: dict[str, Any] = {} + + def fake_execute(**kwargs: Any) -> dict[str, Any]: + seen.update(kwargs, correlation_id=correlation_id.get()) + return {"success": True, "processed": 2} + + with patch( + "app.services.tts_evaluations.sync_generation.execute_tts_sync_generation", + side_effect=fake_execute, + ): + result = job_execution.run_tts_sync_generation.apply( + kwargs={ + "project_id": 7, + "job_id": "42", + "trace_id": "trace-abc", + **SYNC_KWARGS, + }, + task_id="task-xyz", + ).get() + + assert result == {"success": True, "processed": 2} + assert {k: seen[k] for k in ("project_id", "job_id", "task_id")} == { + "project_id": 7, + "job_id": "42", + "task_id": "task-xyz", + } + assert {k: seen[k] for k in SYNC_KWARGS} == SYNC_KWARGS + assert seen["correlation_id"] == "trace-abc" diff --git a/backend/app/tests/core/batch/test_operations.py b/backend/app/tests/core/batch/test_operations.py new file mode 100644 index 000000000..013069fcf --- /dev/null +++ b/backend/app/tests/core/batch/test_operations.py @@ -0,0 +1,102 @@ +from collections.abc import Iterator +from typing import Any + +import pytest +from sqlmodel import Session, delete, select + +from app.core.batch.operations import start_batch_job +from app.core.db import engine +from app.models.batch_job import BatchJob, BatchJobType +from app.tests.utils.auth import TestAuthContext +from app.tests.utils.utils import random_lower_string + + +class _RecordingProvider: + """Fake batch provider that snapshots DB state while the 'network' call runs.""" + + def __init__( + self, session: Session, marker: str, error: Exception | None = None + ) -> None: + self._session = session + self._marker = marker + self._error = error + self.in_transaction_during_call: bool | None = None + self.job_visible_to_other_connections: bool | None = None + + def create_batch( + self, jsonl_data: list[dict[str, Any]], config: dict[str, Any] + ) -> dict[str, Any]: + self.in_transaction_during_call = self._session.in_transaction() + with Session(engine) as other: + jobs = other.exec( + select(BatchJob).where(BatchJob.provider == self._marker) + ).all() + self.job_visible_to_other_connections = len(jobs) == 1 + if self._error is not None: + raise self._error + return { + "provider_batch_id": "batches/abc", + "provider_file_id": "files/abc", + "provider_status": "JOB_STATE_PENDING", + } + + +@pytest.fixture +def committed_session(user_api_key: TestAuthContext) -> Iterator[tuple[Session, str]]: + # The conftest `db` session always sits inside a savepoint, so it can never + # report "no open transaction"; this needs a real session, cleaned up by hand. + marker = f"test-provider-{random_lower_string()}" + session = Session(engine) + try: + yield session, marker + finally: + session.rollback() + session.exec(delete(BatchJob).where(BatchJob.provider == marker)) + session.commit() + session.close() + + +def _start( + session: Session, provider: _RecordingProvider, marker: str, auth: TestAuthContext +) -> BatchJob: + return start_batch_job( + session=session, + provider=provider, + provider_name=marker, + job_type=BatchJobType.TTS_EVALUATION.value, + organization_id=auth.organization_id, + project_id=auth.project_id, + jsonl_data=[{"key": "1"}, {"key": "2"}], + config={"model": "m"}, + ) + + +class TestStartBatchJob: + def test_no_transaction_is_held_during_provider_call( + self, committed_session: tuple[Session, str], user_api_key: TestAuthContext + ) -> None: + session, marker = committed_session + provider = _RecordingProvider(session, marker) + + job = _start(session, provider, marker, user_api_key) + + assert provider.in_transaction_during_call is False + assert provider.job_visible_to_other_connections is True + assert job.provider_batch_id == "batches/abc" + assert job.provider_status == "JOB_STATE_PENDING" + assert job.total_items == 2 + + def test_provider_failure_marks_job_failed_and_reraises( + self, committed_session: tuple[Session, str], user_api_key: TestAuthContext + ) -> None: + session, marker = committed_session + provider = _RecordingProvider( + session, marker, error=RuntimeError("upload refused") + ) + + with pytest.raises(RuntimeError, match="upload refused"): + _start(session, provider, marker, user_api_key) + + [job] = session.exec(select(BatchJob).where(BatchJob.provider == marker)).all() + assert job.provider_status == "failed" + assert job.error_message == "Batch creation failed: upload refused" diff --git a/backend/app/tests/crud/tts_evaluations/__init__.py b/backend/app/tests/crud/tts_evaluations/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/backend/app/tests/crud/tts_evaluations/test_cron.py b/backend/app/tests/crud/tts_evaluations/test_cron.py new file mode 100644 index 000000000..5e82a1275 --- /dev/null +++ b/backend/app/tests/crud/tts_evaluations/test_cron.py @@ -0,0 +1,294 @@ +from collections.abc import Iterator +from datetime import timedelta +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest +from sqlmodel import Session + +from app.core.batch import BatchJobState +from app.crud.tts_evaluations.cron import ( + poll_all_pending_tts_evaluations, + poll_tts_run, +) +from app.models import EvaluationRun +from app.models.batch_job import BatchJob +from app.models.job import JobStatus +from app.models.tts_evaluation import TTSResult +from app.core.config import settings +from app.services.tts_evaluations.constants import TTS_SYNC_STALE_GRACE_SECONDS +from app.tests.utils.auth import TestAuthContext +from app.tests.utils.tts_evaluation import ( + create_test_tts_batch_job, + create_test_tts_result_row, + create_test_tts_run_with_dataset, +) + +MODULE = "app.crud.tts_evaluations.cron" +GEMINI = "gemini-2.5-pro-preview-tts" +SARVAM = "bulbul:v3" + + +def _processing_run( + db: Session, auth: TestAuthContext, models: list[str], stale: bool = False +) -> EvaluationRun: + run = create_test_tts_run_with_dataset( + db, + organization_id=auth.organization_id, + project_id=auth.project_id, + models=models, + ) + run.status = "processing" + if stale: + run.inserted_at -= timedelta( + seconds=settings.CELERY_TASK_TIME_LIMIT + TTS_SYNC_STALE_GRACE_SECONDS + 60 + ) + db.add(run) + db.commit() + return run + + +@pytest.fixture +def dispatch() -> Iterator[MagicMock]: + # Celery enqueue of result processing is the boundary. + with patch( + f"{MODULE}.start_tts_result_processing", return_value="celery-task" + ) as dispatch: + yield dispatch + + +@pytest.fixture +def gemini_client() -> Iterator[MagicMock]: + with patch(f"{MODULE}.GeminiClient") as client: + yield client + + +def _gemini_reports(state: BatchJobState | Exception) -> Any: + """Stand-in for the Gemini status call made while polling a batch job.""" + + def fake_poll(*, session: Session, batch_job: BatchJob, **_kwargs: Any) -> None: + if isinstance(state, Exception): + raise state + batch_job.provider_status = state.value + session.add(batch_job) + session.commit() + + return patch( + "app.crud.evaluations.cron_utils.poll_batch_status", side_effect=fake_poll + ) + + +async def _poll(db: Session, run: EvaluationRun) -> dict[str, Any]: + return await poll_tts_run(session=db, run=run, org_id=run.organization_id) + + +def _refresh(db: Session, *rows: TTSResult | EvaluationRun) -> None: + for row in rows: + db.refresh(row) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("gemini_client") +class TestPollTTSRunGeminiBatches: + @pytest.mark.parametrize( + "initial_state", [BatchJobState.RUNNING, BatchJobState.SUCCEEDED] + ) + async def test_succeeded_batch_dispatches_result_processing( + self, + db: Session, + user_api_key: TestAuthContext, + dispatch: MagicMock, + initial_state: BatchJobState, + ) -> None: + run = _processing_run(db, user_api_key, [GEMINI]) + batch_job = create_test_tts_batch_job( + db, run=run, provider_status=initial_state.value + ) + create_test_tts_result_row(db, run=run, provider=GEMINI) + + with _gemini_reports(BatchJobState.SUCCEEDED): + outcome = await _poll(db, run) + + assert (outcome["action"], outcome["current_status"]) == ( + "dispatched", + "processing", + ) + dispatched = dispatch.call_args.kwargs + assert ( + dispatched["evaluation_run_id"], + dispatched["tts_provider"], + dispatched["provider_batch_id"], + ) == (run.id, GEMINI, batch_job.provider_batch_id) + + async def test_succeeded_batch_with_no_pending_rows_is_not_reprocessed( + self, db: Session, user_api_key: TestAuthContext, dispatch: MagicMock + ) -> None: + run = _processing_run(db, user_api_key, [GEMINI]) + create_test_tts_batch_job( + db, run=run, provider_status=BatchJobState.SUCCEEDED.value + ) + create_test_tts_result_row( + db, run=run, provider=GEMINI, status=JobStatus.SUCCESS + ) + + outcome = await _poll(db, run) + + dispatch.assert_not_called() + assert outcome["action"] == "completed" + _refresh(db, run) + assert run.status == "completed" + + @pytest.mark.parametrize( + ("state", "batch_error", "expected_error"), + [ + ( + BatchJobState.FAILED, + "quota exceeded", + "Gemini batch ended: quota exceeded", + ), + (BatchJobState.CANCELLED, None, "Gemini batch ended: JOB_STATE_CANCELLED"), + (BatchJobState.EXPIRED, None, "Gemini batch ended: JOB_STATE_EXPIRED"), + ], + ) + async def test_unsuccessful_batch_fails_only_that_models_rows( + self, + db: Session, + user_api_key: TestAuthContext, + dispatch: MagicMock, + state: BatchJobState, + batch_error: str | None, + expected_error: str, + ) -> None: + run = _processing_run(db, user_api_key, [GEMINI, SARVAM]) + create_test_tts_batch_job( + db, run=run, provider_status=state.value, error_message=batch_error + ) + gemini_row = create_test_tts_result_row(db, run=run, provider=GEMINI) + sarvam_row = create_test_tts_result_row(db, run=run, provider=SARVAM) + + outcome = await _poll(db, run) + + _refresh(db, gemini_row, sarvam_row, run) + assert (gemini_row.status, gemini_row.error_message) == ( + JobStatus.FAILED.value, + expected_error, + ) + assert sarvam_row.status == JobStatus.PENDING.value + assert run.status == outcome["current_status"] == "processing" + + async def test_running_batch_leaves_rows_pending( + self, db: Session, user_api_key: TestAuthContext, dispatch: MagicMock + ) -> None: + run = _processing_run(db, user_api_key, [GEMINI]) + create_test_tts_batch_job( + db, run=run, provider_status=BatchJobState.PENDING.value + ) + row = create_test_tts_result_row(db, run=run, provider=GEMINI) + + with _gemini_reports(BatchJobState.RUNNING): + outcome = await _poll(db, run) + + assert outcome["action"] == "no_change" + dispatch.assert_not_called() + _refresh(db, row) + assert row.status == JobStatus.PENDING.value + + async def test_gemini_client_failure_fails_only_gemini_rows( + self, + db: Session, + user_api_key: TestAuthContext, + gemini_client: MagicMock, + ) -> None: + gemini_client.from_credentials.side_effect = ValueError("no gemini key") + run = _processing_run(db, user_api_key, [GEMINI, SARVAM]) + create_test_tts_batch_job( + db, run=run, provider_status=BatchJobState.RUNNING.value + ) + gemini_row = create_test_tts_result_row(db, run=run, provider=GEMINI) + sarvam_row = create_test_tts_result_row(db, run=run, provider=SARVAM) + + await _poll(db, run) + + _refresh(db, gemini_row, sarvam_row) + assert (gemini_row.status, gemini_row.error_message) == ( + JobStatus.FAILED.value, + "Gemini client initialization failed: no gemini key", + ) + assert sarvam_row.status == JobStatus.PENDING.value + + async def test_batch_job_never_created_is_skipped_and_run_finalizes( + self, db: Session, user_api_key: TestAuthContext + ) -> None: + # Regression: a job whose creation failed (no provider_batch_id) was polled + # each tick, and the poll error kept the run stuck in processing. + run = _processing_run(db, user_api_key, [GEMINI, SARVAM]) + create_test_tts_batch_job( + db, run=run, provider_status="failed", provider_batch_id=None + ) + create_test_tts_result_row( + db, run=run, provider=GEMINI, status=JobStatus.FAILED + ) + create_test_tts_result_row( + db, run=run, provider=SARVAM, status=JobStatus.SUCCESS + ) + + with _gemini_reports(ValueError("batch name must not be None")): + outcome = await _poll(db, run) + + assert outcome["action"] == "completed" + _refresh(db, run) + assert (run.status, run.error_message) == ( + "completed", + "1 synthesis(es) failed", + ) + + +@pytest.mark.asyncio +class TestPollTTSRunStaleModels: + @pytest.mark.parametrize( + ("model", "expected_error"), + [(SARVAM, "Synthesis did not complete"), (GEMINI, "No batch jobs found")], + ) + @pytest.mark.parametrize("stale", [False, True], ids=["fresh", "stale"]) + async def test_pending_rows_fail_only_after_stale_window( + self, + db: Session, + user_api_key: TestAuthContext, + model: str, + expected_error: str, + stale: bool, + ) -> None: + run = _processing_run(db, user_api_key, [model], stale=stale) + row = create_test_tts_result_row(db, run=run, provider=model) + + await _poll(db, run) + + _refresh(db, row, run) + if stale: + assert (row.status, row.error_message) == ( + JobStatus.FAILED.value, + expected_error, + ) + assert run.status == "failed" + else: + assert row.status == JobStatus.PENDING.value + assert run.status == "processing" + + +@pytest.mark.asyncio +class TestPollAllPendingTTSEvaluations: + async def test_finalizes_sync_only_run_without_batch_job( + self, db: Session, user_api_key: TestAuthContext + ) -> None: + run = _processing_run(db, user_api_key, [SARVAM]) + assert run.batch_job_id is None + create_test_tts_result_row( + db, run=run, provider=SARVAM, status=JobStatus.SUCCESS + ) + + summary = await poll_all_pending_tts_evaluations(session=db) + + [detail] = [d for d in summary["details"] if d["run_id"] == run.id] + assert detail["action"] == "completed" + _refresh(db, run) + assert run.status == "completed" diff --git a/backend/app/tests/crud/tts_evaluations/test_run.py b/backend/app/tests/crud/tts_evaluations/test_run.py new file mode 100644 index 000000000..3e2e7254a --- /dev/null +++ b/backend/app/tests/crud/tts_evaluations/test_run.py @@ -0,0 +1,68 @@ +import pytest +from sqlmodel import Session + +from app.crud.tts_evaluations.run import finalize_tts_run_status +from app.models import EvaluationRun +from app.models.job import JobStatus +from app.tests.utils.auth import TestAuthContext +from app.tests.utils.tts_evaluation import ( + create_test_tts_result_row, + create_test_tts_run_with_dataset, +) +from app.tests.utils.utils import get_non_existent_id + +S, P, F = JobStatus.SUCCESS, JobStatus.PENDING, JobStatus.FAILED + + +@pytest.fixture +def run(db: Session, user_api_key: TestAuthContext) -> EvaluationRun: + return create_test_tts_run_with_dataset( + db, + organization_id=user_api_key.organization_id, + project_id=user_api_key.project_id, + ) + + +class TestFinalizeTTSRunStatus: + @pytest.mark.parametrize( + ("row_statuses", "expected_status", "expected_error"), + [ + ([S, S], "completed", None), + ([S, P, F], "processing", "1 synthesis(es) failed"), + ([S, F, F], "completed", "2 synthesis(es) failed"), + ([F, F], "failed", "2 synthesis(es) failed"), + ([], "failed", None), + ], + ids=["all_success", "pending", "partial_failure", "all_failed", "no_rows"], + ) + def test_status_derived_from_rows( + self, + db: Session, + run: EvaluationRun, + row_statuses: list[JobStatus], + expected_status: str, + expected_error: str | None, + ) -> None: + for status in row_statuses: + create_test_tts_result_row(db, run=run, status=status) + + finalize_tts_run_status(session=db, run_id=run.id) + + db.refresh(run) + assert (run.status, run.error_message) == (expected_status, expected_error) + + def test_explicit_error_message_wins_over_failed_count( + self, db: Session, run: EvaluationRun + ) -> None: + create_test_tts_result_row(db, run=run, status=JobStatus.FAILED) + + finalized = finalize_tts_run_status( + session=db, run_id=run.id, error_message="gemini: JOB_STATE_FAILED" + ) + + assert finalized.error_message == "gemini: JOB_STATE_FAILED" + + def test_missing_run_returns_none(self, db: Session) -> None: + missing_id = get_non_existent_id(db, EvaluationRun) + + assert finalize_tts_run_status(session=db, run_id=missing_id) is None diff --git a/backend/app/tests/services/tts_evaluations/__init__.py b/backend/app/tests/services/tts_evaluations/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/backend/app/tests/services/tts_evaluations/test_batch_job.py b/backend/app/tests/services/tts_evaluations/test_batch_job.py new file mode 100644 index 000000000..27823fc1c --- /dev/null +++ b/backend/app/tests/services/tts_evaluations/test_batch_job.py @@ -0,0 +1,249 @@ +import io +from collections.abc import Iterator +from types import SimpleNamespace +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest +from celery.exceptions import SoftTimeLimitExceeded +from gevent import Timeout +from sqlmodel import Session, select + +from app.crud.language import get_language_by_locale +from app.models import EvaluationDataset, EvaluationRun +from app.models.job import JobStatus +from app.models.tts_evaluation import TTSResult +from app.services.tts_evaluations.batch_job import execute_batch_submission +from app.tests.utils.auth import TestAuthContext +from app.tests.utils.tts_evaluation import ( + create_test_tts_run_with_dataset, + fake_gemini_batch, + use_test_session, +) +from app.tests.utils.utils import get_non_existent_id + +MODULE = "app.services.tts_evaluations.batch_job" +CRUD_BATCH = "app.crud.tts_evaluations.batch" +GEMINI = "gemini-2.5-pro-preview-tts" +SARVAM = "bulbul:v3" +ELEVEN = "eleven_v3" + + +def _csv(count: int) -> bytes: + return ("text\n" + "".join(f"sample {i}\n" for i in range(count))).encode() + + +@pytest.fixture +def boundaries(db: Session) -> Iterator[SimpleNamespace]: + state = SimpleNamespace(csv=_csv(2), enqueued=[], enqueue_errors={}) + storage = MagicMock() + storage.stream.side_effect = lambda _url: io.BytesIO(state.csv) + + def enqueue(**kwargs: Any) -> str: + if kwargs["model"] in state.enqueue_errors: + raise state.enqueue_errors[kwargs["model"]] + state.enqueued.append(kwargs) + return f"task-{len(state.enqueued)}" + + with ( + use_test_session(MODULE, db), + patch(f"{MODULE}.get_cloud_storage", return_value=storage), + patch(f"{MODULE}.start_tts_sync_generation", side_effect=enqueue), + ): + yield state + + +def _run( + db: Session, auth: TestAuthContext, models: list[str], **kwargs: Any +) -> EvaluationRun: + return create_test_tts_run_with_dataset( + db, + organization_id=auth.organization_id, + project_id=auth.project_id, + models=models, + **kwargs, + ) + + +def _execute(run: EvaluationRun, dataset_id: int | None = None) -> dict[str, Any]: + return execute_batch_submission( + project_id=run.project_id, + job_id=str(run.id), + task_id="celery-task", + task_instance=MagicMock(), + organization_id=run.organization_id, + dataset_id=dataset_id or run.dataset_id, + models=run.providers, + ) + + +def _results(db: Session, run: EvaluationRun, model: str) -> list[TTSResult]: + stmt = ( + select(TTSResult) + .where(TTSResult.evaluation_run_id == run.id, TTSResult.provider == model) + .order_by(TTSResult.id) + ) + return list(db.exec(stmt).all()) + + +def _row_states( + db: Session, run: EvaluationRun, model: str +) -> set[tuple[str, str | None]]: + return {(r.status, r.error_message) for r in _results(db, run, model)} + + +@pytest.mark.usefixtures("boundaries") +class TestExecuteBatchSubmission: + def test_mixed_run_creates_rows_and_starts_every_model( + self, db: Session, user_api_key: TestAuthContext, boundaries: SimpleNamespace + ) -> None: + run = _run(db, user_api_key, [GEMINI, SARVAM, ELEVEN]) + + with fake_gemini_batch(CRUD_BATCH) as gemini: + outcome = _execute(run) + + assert (outcome["success"], outcome["result_count"]) == (True, 6) + for model in (GEMINI, SARVAM, ELEVEN): + assert [r.sample_text for r in _results(db, run, model)] == [ + "sample 0", + "sample 1", + ] + assert _row_states(db, run, model) == {(JobStatus.PENDING.value, None)} + assert {r["key"] for r in gemini.submitted[f"models/{GEMINI}"]} == { + str(r.id) for r in _results(db, run, GEMINI) + } + assert [ + (c["model"], c["job_id"], c["language_code"]) for c in boundaries.enqueued + ] == [(SARVAM, str(run.id), "en-IN"), (ELEVEN, str(run.id), "en-IN")] + db.refresh(run) + assert run.status == "processing" + + @pytest.mark.parametrize( + ("locale", "expected_language"), + [(None, "en-IN"), ("hi", "hi-IN"), ("or", "od-IN")], + ) + def test_sync_language_comes_from_dataset_locale( + self, + db: Session, + user_api_key: TestAuthContext, + boundaries: SimpleNamespace, + locale: str | None, + expected_language: str, + ) -> None: + language_id = ( + get_language_by_locale(session=db, locale=locale).id if locale else None + ) + run = _run(db, user_api_key, [SARVAM], language_id=language_id) + + _execute(run) + + assert [c["language_code"] for c in boundaries.enqueued] == [expected_language] + + @pytest.mark.parametrize("models", [[GEMINI], [GEMINI, SARVAM]]) + def test_gemini_submission_failure_fails_only_gemini_rows( + self, + db: Session, + user_api_key: TestAuthContext, + boundaries: SimpleNamespace, + models: list[str], + ) -> None: + run = _run(db, user_api_key, models) + + with patch(f"{CRUD_BATCH}.GeminiClient") as gemini_client: + gemini_client.from_credentials.side_effect = ValueError("no gemini key") + outcome = _execute(run) + + assert outcome["success"] is True + assert _row_states(db, run, GEMINI) == { + (JobStatus.FAILED.value, "Batch submission failed: no gemini key") + } + if SARVAM in models: + assert _row_states(db, run, SARVAM) == {(JobStatus.PENDING.value, None)} + assert [c["model"] for c in boundaries.enqueued] == [SARVAM] + # Left for the cron to finalize. + db.refresh(run) + assert run.status == "processing" + + def test_enqueue_failure_fails_only_that_models_rows( + self, db: Session, user_api_key: TestAuthContext, boundaries: SimpleNamespace + ) -> None: + boundaries.enqueue_errors[SARVAM] = ConnectionError("broker unreachable") + run = _run(db, user_api_key, [SARVAM, ELEVEN]) + + outcome = _execute(run) + + assert outcome["success"] is True + assert _row_states(db, run, SARVAM) == { + (JobStatus.FAILED.value, "Failed to queue synthesis: broker unreachable") + } + assert _row_states(db, run, ELEVEN) == {(JobStatus.PENDING.value, None)} + assert [c["model"] for c in boundaries.enqueued] == [ELEVEN] + db.refresh(run) + assert run.status == "processing" + + @pytest.mark.parametrize( + ("case", "expected_error"), + [ + ("missing_dataset", "Dataset not found"), + ("empty_csv", "No samples found for dataset"), + ("no_object_store_url", "No samples found for dataset"), + ], + ) + def test_unstartable_run_is_failed_without_rows( + self, + db: Session, + user_api_key: TestAuthContext, + boundaries: SimpleNamespace, + case: str, + expected_error: str, + ) -> None: + if case == "empty_csv": + boundaries.csv = _csv(0) + run = _run( + db, + user_api_key, + [SARVAM], + object_store_url=None if case == "no_object_store_url" else "s3://b/x.csv", + ) + dataset_id = ( + get_non_existent_id(db, EvaluationDataset) + if case == "missing_dataset" + else None + ) + + outcome = _execute(run, dataset_id=dataset_id) + + assert outcome["success"] is False + db.refresh(run) + assert (run.status, run.error_message) == ("failed", expected_error) + assert _results(db, run, SARVAM) == [] + assert boundaries.enqueued == [] + + @pytest.mark.parametrize( + ("error", "expected_message"), + [ + (Timeout(), "Task exceeded soft time limit"), + (SoftTimeLimitExceeded(), "Task exceeded soft time limit"), + (RuntimeError("no bucket"), "no bucket"), + ], + ids=["gevent_timeout", "soft_timeout", "unexpected"], + ) + def test_worker_error_fails_run( + self, + db: Session, + user_api_key: TestAuthContext, + error: BaseException, + expected_message: str, + ) -> None: + run = _run(db, user_api_key, [SARVAM]) + + with patch(f"{MODULE}.get_cloud_storage", side_effect=error): + if isinstance(error, RuntimeError): + assert _execute(run)["success"] is False + else: + # Timeouts re-raise so Celery records the task as timed out. + with pytest.raises(type(error)): + _execute(run) + + db.refresh(run) + assert (run.status, run.error_message) == ("failed", expected_message) diff --git a/backend/app/tests/services/tts_evaluations/test_batch_result_processing.py b/backend/app/tests/services/tts_evaluations/test_batch_result_processing.py new file mode 100644 index 000000000..2dbd087f3 --- /dev/null +++ b/backend/app/tests/services/tts_evaluations/test_batch_result_processing.py @@ -0,0 +1,237 @@ +import base64 +from collections.abc import Iterator +from dataclasses import dataclass +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest +from celery.exceptions import SoftTimeLimitExceeded +from gevent import Timeout +from sqlmodel import Session + +from app.models import EvaluationRun +from app.models.job import JobStatus +from app.models.tts_evaluation import TTSResult +from app.services.tts_evaluations.batch_result_processing import ( + execute_tts_result_processing, +) +from app.tests.utils.auth import TestAuthContext +from app.tests.utils.tts_evaluation import ( + create_test_tts_result_row, + create_test_tts_run_with_dataset, + use_test_session, +) + +MODULE = "app.services.tts_evaluations.batch_result_processing" +GEMINI = "gemini-2.5-pro-preview-tts" +SARVAM = "bulbul:v3" +# 24 kHz / 16-bit / mono: 48000 bytes is one second; the WAV adds a 44-byte header. +ONE_SECOND_PCM_B64 = base64.b64encode(b"\x00\x01" * 24000).decode() + + +def _audio_response( + data: str = ONE_SECOND_PCM_B64, key: str = "inlineData" +) -> dict[str, Any]: + return { + "candidates": [ + {"content": {"parts": [{key: {"mimeType": "audio/L16", "data": data}}]}} + ] + } + + +def _ok(result_id: int | str, **kwargs: Any) -> dict[str, Any]: + return {"custom_id": str(result_id), "response": _audio_response(**kwargs)} + + +@dataclass +class Boundaries: + gemini_client: MagicMock + download: MagicMock + upload: MagicMock + + +def _fake_upload(*, subdirectory: str, filename: str, **_kwargs: Any) -> str: + return f"s3://bucket/{subdirectory}/{filename}" + + +@pytest.fixture +def boundaries(db: Session) -> Iterator[Boundaries]: + with ( + use_test_session(MODULE, db), + patch(f"{MODULE}.GeminiClient") as gemini_client, + patch(f"{MODULE}.GeminiBatchProvider") as batch_provider, + patch(f"{MODULE}.get_cloud_storage"), + patch(f"{MODULE}.upload_to_object_store", side_effect=_fake_upload) as upload, + ): + download = batch_provider.return_value.download_batch_results + download.return_value = [] + yield Boundaries(gemini_client, download, upload) + + +@pytest.fixture +def run(db: Session, user_api_key: TestAuthContext) -> EvaluationRun: + return create_test_tts_run_with_dataset( + db, + organization_id=user_api_key.organization_id, + project_id=user_api_key.project_id, + models=[GEMINI, SARVAM], + ) + + +def _execute(run: EvaluationRun) -> dict[str, Any]: + return execute_tts_result_processing( + project_id=run.project_id, + job_id="1", + task_id="celery-task", + task_instance=MagicMock(), + organization_id=run.organization_id, + evaluation_run_id=run.id, + tts_provider=GEMINI, + provider_batch_id="batches/abc", + ) + + +def _refresh(db: Session, *rows: TTSResult | EvaluationRun) -> None: + for row in rows: + db.refresh(row) + + +@pytest.mark.usefixtures("boundaries") +class TestExecuteTTSResultProcessing: + @pytest.mark.parametrize("key", ["inlineData", "inline_data"]) + def test_success_stores_wav_without_finalizing_run( + self, db: Session, run: EvaluationRun, boundaries: Boundaries, key: str + ) -> None: + rows = [create_test_tts_result_row(db, run=run) for _ in range(2)] + boundaries.download.return_value = [_ok(r.id, key=key) for r in rows] + + outcome = _execute(run) + + assert (outcome["processed"], outcome["failed"]) == (2, 0) + _refresh(db, run, *rows) + for row in rows: + assert row.status == JobStatus.SUCCESS.value + assert row.object_store_url.startswith("s3://bucket/evaluations/tts/audio/") + assert row.metadata_ == {"duration_seconds": 1.0, "size_bytes": 48044} + assert run.status == "pending" + + def test_only_this_models_pending_rows_are_processed( + self, db: Session, run: EvaluationRun, boundaries: Boundaries + ) -> None: + done = create_test_tts_result_row( + db, + run=run, + status=JobStatus.SUCCESS, + object_store_url="s3://bucket/keep.wav", + ) + other_model = create_test_tts_result_row(db, run=run, provider=SARVAM) + pending = create_test_tts_result_row(db, run=run) + boundaries.download.return_value = [ + _ok(done.id), + _ok(other_model.id), + _ok(pending.id), + ] + + _execute(run) + + _refresh(db, done, other_model, pending) + assert done.object_store_url == "s3://bucket/keep.wav" + assert other_model.status == JobStatus.PENDING.value + assert pending.status == JobStatus.SUCCESS.value + + def test_bad_results_fail_their_row_and_others_still_succeed( + self, db: Session, run: EvaluationRun, boundaries: Boundaries + ) -> None: + errored, no_error_text, silent, bad_padding, not_base64, odd_shape, ok = ( + create_test_tts_result_row(db, run=run) for _ in range(7) + ) + boundaries.download.return_value = [ + {"custom_id": "not-a-number", "response": _audio_response()}, + {"custom_id": str(errored.id), "response": None, "error": "SAFETY block"}, + {"custom_id": str(no_error_text.id)}, + { + "custom_id": str(silent.id), + "response": {"candidates": [{"content": {"parts": [{"text": "hi"}]}}]}, + }, + _ok(bad_padding.id, data="abc"), + # Without strict decoding this became b"" and was stored as a 0s success. + _ok(not_base64.id, data="###"), + {"custom_id": str(odd_shape.id), "response": ["not", "a", "dict"]}, + _ok(ok.id), + ] + + outcome = _execute(run) + + assert (outcome["processed"], outcome["failed"]) == (1, 7) + _refresh( + db, errored, no_error_text, silent, bad_padding, not_base64, odd_shape, ok + ) + assert ok.status == JobStatus.SUCCESS.value + assert errored.error_message == "SAFETY block" + assert no_error_text.error_message == "Unknown error" + assert silent.error_message == "No audio data in response" + for row in (bad_padding, not_base64, odd_shape): + assert row.error_message.startswith("Audio processing failed:") + for row in (errored, no_error_text, silent, bad_padding, not_base64, odd_shape): + assert (row.status, row.object_store_url) == (JobStatus.FAILED.value, None) + + def test_upload_failure_fails_row( + self, db: Session, run: EvaluationRun, boundaries: Boundaries + ) -> None: + boundaries.upload.side_effect = None + boundaries.upload.return_value = None + row = create_test_tts_result_row(db, run=run) + boundaries.download.return_value = [_ok(row.id)] + + _execute(run) + + _refresh(db, row) + assert (row.status, row.error_message) == ( + JobStatus.FAILED.value, + "Audio upload to object store failed", + ) + + @pytest.mark.parametrize( + "timeout", [Timeout(), SoftTimeLimitExceeded()], ids=["gevent", "soft"] + ) + def test_timeout_fails_only_this_models_pending_rows_and_reraises( + self, + db: Session, + run: EvaluationRun, + boundaries: Boundaries, + timeout: BaseException, + ) -> None: + boundaries.download.side_effect = timeout + gemini_row = create_test_tts_result_row(db, run=run) + sarvam_row = create_test_tts_result_row(db, run=run, provider=SARVAM) + + with pytest.raises(type(timeout)): + _execute(run) + + _refresh(db, run, gemini_row, sarvam_row) + assert (gemini_row.status, gemini_row.error_message) == ( + JobStatus.FAILED.value, + "Result processing exceeded soft time limit", + ) + assert sarvam_row.status == JobStatus.PENDING.value + assert run.status == "pending" + + def test_unexpected_error_fails_only_this_models_pending_rows( + self, db: Session, run: EvaluationRun, boundaries: Boundaries + ) -> None: + boundaries.gemini_client.from_credentials.side_effect = ValueError( + "no gemini key" + ) + gemini_row = create_test_tts_result_row(db, run=run) + sarvam_row = create_test_tts_result_row(db, run=run, provider=SARVAM) + + outcome = _execute(run) + + assert outcome == {"success": False, "error": "no gemini key"} + _refresh(db, run, gemini_row, sarvam_row) + assert (gemini_row.status, gemini_row.error_message) == ( + JobStatus.FAILED.value, + "Result processing failed: no gemini key", + ) + assert sarvam_row.status == JobStatus.PENDING.value + assert run.status == "pending" diff --git a/backend/app/tests/services/tts_evaluations/test_dataset.py b/backend/app/tests/services/tts_evaluations/test_dataset.py new file mode 100644 index 000000000..cbc1b5c43 --- /dev/null +++ b/backend/app/tests/services/tts_evaluations/test_dataset.py @@ -0,0 +1,59 @@ +import io +from unittest.mock import MagicMock + +import pytest +from celery.exceptions import SoftTimeLimitExceeded +from sqlmodel import Session + +from app.services.tts_evaluations.dataset import get_sample_texts_from_dataset +from app.tests.utils.auth import TestAuthContext +from app.tests.utils.tts_evaluation import TEST_DATASET_URL, create_test_tts_dataset + + +def _storage(stream: object) -> MagicMock: + storage = MagicMock() + storage.stream.side_effect = stream + return storage + + +@pytest.mark.parametrize( + ("object_store_url", "stream", "expected"), + [ + ( + TEST_DATASET_URL, + lambda url: io.BytesIO(b"text\nHello world\n \n Good morning \n"), + ["Hello world", "Good morning"], + ), + (None, lambda url: io.BytesIO(b"text\nunused\n"), []), + (TEST_DATASET_URL, OSError("s3 down"), []), + ], + ids=["reads_csv", "no_url", "stream_error"], +) +def test_get_sample_texts_from_dataset( + db: Session, + user_api_key: TestAuthContext, + object_store_url: str | None, + stream: object, + expected: list[str], +) -> None: + dataset = create_test_tts_dataset( + db, + organization_id=user_api_key.organization_id, + project_id=user_api_key.project_id, + object_store_url=object_store_url, + ) + + assert get_sample_texts_from_dataset(_storage(stream), dataset) == expected + + +def test_get_sample_texts_from_dataset_reraises_soft_timeout( + db: Session, user_api_key: TestAuthContext +) -> None: + dataset = create_test_tts_dataset( + db, + organization_id=user_api_key.organization_id, + project_id=user_api_key.project_id, + ) + + with pytest.raises(SoftTimeLimitExceeded): + get_sample_texts_from_dataset(_storage(SoftTimeLimitExceeded()), dataset) diff --git a/backend/app/tests/services/tts_evaluations/test_sync_generation.py b/backend/app/tests/services/tts_evaluations/test_sync_generation.py new file mode 100644 index 000000000..346107084 --- /dev/null +++ b/backend/app/tests/services/tts_evaluations/test_sync_generation.py @@ -0,0 +1,304 @@ +import base64 +from collections.abc import Iterator +from types import SimpleNamespace +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest +from celery.exceptions import SoftTimeLimitExceeded +from gevent import Timeout +from sqlmodel import Session + +from app.core.audio_utils import pcm_to_wav +from app.models import EvaluationRun +from app.models.job import JobStatus +from app.models.tts_evaluation import TTSResult +from app.services.tts_evaluations.sync_generation import ( + execute_tts_sync_generation, + synthesize_elevenlabs, + synthesize_sarvam, +) +from app.tests.utils.auth import TestAuthContext +from app.tests.utils.tts_evaluation import ( + create_test_tts_result_row, + create_test_tts_run_with_dataset, + use_test_session, +) + +MODULE = "app.services.tts_evaluations.sync_generation" +SARVAM = "bulbul:v3" +ELEVEN = "eleven_v3" +# 24 kHz / 16-bit / mono: 48000 bytes is one second; the WAV adds a 44-byte header. +ONE_SECOND_PCM = b"\x00\x01" * 24000 +FAIL_TEXT = "provider rejects this" + + +def _sarvam_client(audios: list[str]) -> MagicMock: + client = MagicMock() + client.text_to_speech.convert.return_value = SimpleNamespace(audios=audios) + return client + + +class TestSynthesizeSarvam: + def test_joins_chunks_without_wav_headers(self) -> None: + chunks = [base64.b64encode(pcm_to_wav(pcm)).decode() for pcm in (b"ab", b"cd")] + client = _sarvam_client(chunks) + + pcm = synthesize_sarvam(client, "namaste", SARVAM, "hi-IN") + + assert pcm == b"abcd" + assert ( + client.text_to_speech.convert.call_args.kwargs["target_language_code"] + == "hi-IN" + ) + + def test_no_audio_raises(self) -> None: + with pytest.raises(ValueError, match="Sarvam returned no audio"): + synthesize_sarvam(_sarvam_client([]), "namaste", SARVAM, "hi-IN") + + +class TestSynthesizeElevenLabs: + # od-IN has no ElevenLabs code; None lets ElevenLabs auto-detect. + @pytest.mark.parametrize( + ("language_code", "expected"), [("hi-IN", "hi"), ("od-IN", None)] + ) + def test_maps_language_and_joins_pcm( + self, language_code: str, expected: str | None + ) -> None: + client = MagicMock() + client.text_to_speech.convert.return_value = iter([b"ab", b"cd"]) + + pcm = synthesize_elevenlabs(client, "namaste", ELEVEN, language_code) + + assert pcm == b"abcd" + sent = client.text_to_speech.convert.call_args.kwargs + assert (sent["language_code"], sent["output_format"]) == ( + expected, + "pcm_24000", + ) + + def test_no_audio_raises(self) -> None: + client = MagicMock() + client.text_to_speech.convert.return_value = iter([]) + + with pytest.raises(ValueError, match="ElevenLabs returned no audio"): + synthesize_elevenlabs(client, "namaste", ELEVEN, "hi-IN") + + +def _sarvam_convert(*, text: str, **_kwargs: Any) -> Any: + if text == FAIL_TEXT: + raise RuntimeError("bad request") + return SimpleNamespace( + audios=[base64.b64encode(pcm_to_wav(ONE_SECOND_PCM)).decode()] + ) + + +def _fake_upload(*, subdirectory: str, filename: str, **_kwargs: Any) -> str: + return f"s3://bucket/{subdirectory}/{filename}" + + +@pytest.fixture +def boundaries(db: Session) -> Iterator[SimpleNamespace]: + sarvam = MagicMock() + sarvam.text_to_speech.convert.side_effect = _sarvam_convert + eleven = MagicMock() + eleven.text_to_speech.convert.side_effect = lambda **_kw: iter([ONE_SECOND_PCM]) + with ( + use_test_session(MODULE, db), + patch(f"{MODULE}.SarvamAIProvider.create_client", return_value=sarvam), + patch(f"{MODULE}.ElevenlabsAIProvider.create_client", return_value=eleven), + patch( + f"{MODULE}.get_provider_credential", return_value={"api_key": "k"} + ) as credentials, + patch(f"{MODULE}.get_cloud_storage") as storage, + patch(f"{MODULE}.upload_to_object_store", side_effect=_fake_upload) as upload, + ): + yield SimpleNamespace( + sarvam=sarvam, credentials=credentials, storage=storage, upload=upload + ) + + +@pytest.fixture +def run(db: Session, user_api_key: TestAuthContext) -> EvaluationRun: + return create_test_tts_run_with_dataset( + db, + organization_id=user_api_key.organization_id, + project_id=user_api_key.project_id, + models=[SARVAM, ELEVEN], + ) + + +def _row( + db: Session, + run: EvaluationRun, + model: str = SARVAM, + status: JobStatus = JobStatus.PENDING, + text: str | None = None, +) -> TTSResult: + return create_test_tts_result_row( + db, run=run, provider=model, status=status, sample_text=text + ) + + +def _execute(run: EvaluationRun, model: str = SARVAM) -> dict[str, Any]: + return execute_tts_sync_generation( + project_id=run.project_id, + job_id=str(run.id), + task_id="celery-task", + task_instance=MagicMock(), + organization_id=run.organization_id, + model=model, + language_code="hi-IN", + ) + + +def _refresh(db: Session, *rows: TTSResult | EvaluationRun) -> None: + for row in rows: + db.refresh(row) + + +@pytest.mark.usefixtures("boundaries") +class TestExecuteTTSSyncGeneration: + @pytest.mark.parametrize("model", [SARVAM, ELEVEN]) + def test_success_stores_audio_without_finalizing_run( + self, db: Session, run: EvaluationRun, model: str + ) -> None: + rows = [_row(db, run, model), _row(db, run, model)] + + outcome = _execute(run, model) + + assert (outcome["success"], outcome["processed"], outcome["failed"]) == ( + True, + 2, + 0, + ) + _refresh(db, run, *rows) + for row in rows: + assert row.status == JobStatus.SUCCESS.value + assert row.object_store_url.startswith("s3://bucket/evaluations/tts/audio/") + assert row.metadata_ == {"duration_seconds": 1.0, "size_bytes": 48044} + assert run.status == "pending" + + def test_only_this_models_pending_rows_are_synthesized( + self, db: Session, run: EvaluationRun, boundaries: SimpleNamespace + ) -> None: + done = _row(db, run, status=JobStatus.SUCCESS) + done.object_store_url = "s3://bucket/original.wav" + db.add(done) + db.commit() + pending = _row(db, run) + other_model = _row(db, run, model=ELEVEN) + + _execute(run) + + # A redelivered task must not re-bill the provider for finished rows. + assert boundaries.sarvam.text_to_speech.convert.call_count == 1 + _refresh(db, done, pending, other_model) + assert done.object_store_url == "s3://bucket/original.wav" + assert pending.status == JobStatus.SUCCESS.value + assert other_model.status == JobStatus.PENDING.value + + def test_synthesis_failure_only_fails_that_row( + self, db: Session, run: EvaluationRun + ) -> None: + ok = _row(db, run) + rejected = _row(db, run, text=FAIL_TEXT) + + outcome = _execute(run) + + assert (outcome["processed"], outcome["failed"]) == (1, 1) + _refresh(db, ok, rejected) + assert ok.status == JobStatus.SUCCESS.value + assert (rejected.status, rejected.error_message) == ( + JobStatus.FAILED.value, + "sarvamai synthesis failed: bad request", + ) + + def test_upload_failure_fails_row( + self, db: Session, run: EvaluationRun, boundaries: SimpleNamespace + ) -> None: + boundaries.upload.side_effect = None + boundaries.upload.return_value = None + row = _row(db, run) + + _execute(run) + + _refresh(db, row) + assert (row.status, row.object_store_url, row.error_message) == ( + JobStatus.FAILED.value, + None, + "Audio upload to object store failed", + ) + + def test_unsupported_model_fails_its_rows( + self, db: Session, run: EvaluationRun + ) -> None: + row = _row(db, run, model="mystery-tts") + sarvam_row = _row(db, run) + + outcome = _execute(run, model="mystery-tts") + + assert outcome["success"] is False + _refresh(db, row, sarvam_row) + assert (row.status, row.error_message) == ( + JobStatus.FAILED.value, + "Unsupported sync TTS model: mystery-tts", + ) + assert sarvam_row.status == JobStatus.PENDING.value + + @pytest.mark.parametrize( + ("break_boundary", "error"), + [ + ( + lambda b: setattr(b.credentials, "return_value", None), + "sarvamai credentials are not configured", + ), + ( + lambda b: setattr(b.storage, "side_effect", RuntimeError("no bucket")), + "no bucket", + ), + ], + ids=["missing_credentials", "crash"], + ) + def test_setup_failure_fails_every_pending_row( + self, + db: Session, + run: EvaluationRun, + boundaries: SimpleNamespace, + break_boundary: Any, + error: str, + ) -> None: + break_boundary(boundaries) + rows = [_row(db, run), _row(db, run)] + + outcome = _execute(run) + + assert (outcome["success"], outcome["error"]) == (False, error) + _refresh(db, *rows) + for row in rows: + assert (row.status, row.error_message) == ( + JobStatus.FAILED.value, + f"sarvamai synthesis failed: {error}", + ) + + @pytest.mark.parametrize( + "timeout", [Timeout(), SoftTimeLimitExceeded()], ids=["gevent", "soft"] + ) + def test_timeout_fails_leftover_rows_and_reraises( + self, + db: Session, + run: EvaluationRun, + boundaries: SimpleNamespace, + timeout: BaseException, + ) -> None: + boundaries.storage.side_effect = timeout + row = _row(db, run) + + with pytest.raises(type(timeout)): + _execute(run) + + _refresh(db, row) + assert (row.status, row.error_message) == ( + JobStatus.FAILED.value, + "Synthesis timed out before this sample was processed", + ) diff --git a/backend/app/tests/utils/tts_evaluation.py b/backend/app/tests/utils/tts_evaluation.py new file mode 100644 index 000000000..81d79626f --- /dev/null +++ b/backend/app/tests/utils/tts_evaluation.py @@ -0,0 +1,185 @@ +"""Factories and session helpers for TTS evaluation tests.""" + +from collections.abc import Iterator +from contextlib import contextmanager +from typing import Any +from unittest.mock import MagicMock, patch + +from sqlmodel import Session + +from app.core.util import now +from app.crud.tts_evaluations import create_tts_run +from app.models import EvaluationDataset, EvaluationRun +from app.models.batch_job import BatchJob, BatchJobType +from app.models.job import JobStatus +from app.models.stt_evaluation import EvaluationType +from app.models.tts_evaluation import TTSResult +from app.services.tts_evaluations.constants import DEFAULT_TTS_MODEL +from app.tests.utils.utils import random_lower_string + +TEST_DATASET_URL = "s3://test-bucket/tts_datasets/test.csv" + + +class NonClosingSession: + """Lets a service's `with Session(engine)` reuse the transactional test session. + + Services open their own sessions; pointing them at the conftest `db` keeps + every write inside the per-test rollback and visible to the test's asserts. + """ + + def __init__(self, session: Session) -> None: + self._session = session + + def __enter__(self) -> Session: + return self._session + + def __exit__(self, *_exc: object) -> bool: + return False + + +@contextmanager +def use_test_session(module_path: str, db: Session) -> Iterator[Session]: + with patch( + f"{module_path}.Session", side_effect=lambda *_a, **_k: NonClosingSession(db) + ): + yield db + + +class FakeGeminiBatchProvider: + """Stands in for the Gemini Batch API; records each submitted JSONL per model path.""" + + def __init__(self) -> None: + self.submitted: dict[str, list[dict[str, Any]]] = {} + + def __call__( + self, *, client: Any, model: str | None = None + ) -> "FakeGeminiBatchProvider": + self._model = model or "" + return self + + def create_batch( + self, jsonl_data: list[dict[str, Any]], config: dict[str, Any] + ) -> dict[str, Any]: + self.submitted[self._model] = jsonl_data + return { + "provider_batch_id": f"batches/{self._model.split('/')[-1]}", + "provider_file_id": "files/input", + "provider_status": "JOB_STATE_PENDING", + "total_items": len(jsonl_data), + } + + +@contextmanager +def fake_gemini_batch(module_path: str) -> Iterator[FakeGeminiBatchProvider]: + fake = FakeGeminiBatchProvider() + with ( + patch(f"{module_path}.GeminiClient") as gemini_client, + patch(f"{module_path}.GeminiBatchProvider", side_effect=fake), + ): + gemini_client.from_credentials.return_value = MagicMock() + yield fake + + +def create_test_tts_dataset( + db: Session, + *, + organization_id: int, + project_id: int, + language_id: int | None = None, + object_store_url: str | None = TEST_DATASET_URL, + sample_count: int = 1, +) -> EvaluationDataset: + dataset = EvaluationDataset( + name=f"tts_dataset_{random_lower_string()}", + type=EvaluationType.TTS.value, + language_id=language_id, + object_store_url=object_store_url, + dataset_metadata={"sample_count": sample_count}, + organization_id=organization_id, + project_id=project_id, + inserted_at=now(), + updated_at=now(), + ) + db.add(dataset) + db.commit() + db.refresh(dataset) + return dataset + + +def create_test_tts_run_with_dataset( + db: Session, + *, + organization_id: int, + project_id: int, + models: list[str] | None = None, + language_id: int | None = None, + object_store_url: str | None = TEST_DATASET_URL, +) -> EvaluationRun: + dataset = create_test_tts_dataset( + db, + organization_id=organization_id, + project_id=project_id, + language_id=language_id, + object_store_url=object_store_url, + ) + return create_tts_run( + session=db, + run_name=f"tts_run_{random_lower_string()}", + dataset_id=dataset.id, + dataset_name=dataset.name, + org_id=organization_id, + project_id=project_id, + models=models or [DEFAULT_TTS_MODEL], + ) + + +def create_test_tts_result_row( + db: Session, + *, + run: EvaluationRun, + provider: str = DEFAULT_TTS_MODEL, + status: JobStatus = JobStatus.PENDING, + sample_text: str | None = None, + object_store_url: str | None = None, +) -> TTSResult: + result = TTSResult( + sample_text=sample_text or f"text {random_lower_string()}", + object_store_url=object_store_url, + provider=provider, + status=status.value, + evaluation_run_id=run.id, + organization_id=run.organization_id, + project_id=run.project_id, + inserted_at=now(), + updated_at=now(), + ) + db.add(result) + db.commit() + db.refresh(result) + return result + + +def create_test_tts_batch_job( + db: Session, + *, + run: EvaluationRun, + model: str = DEFAULT_TTS_MODEL, + provider_status: str, + provider_batch_id: str | None = "batches/test", + error_message: str | None = None, +) -> BatchJob: + batch_job = BatchJob( + provider="google-aistudio", + job_type=BatchJobType.TTS_EVALUATION.value, + config={"evaluation_run_id": run.id, "tts_provider": model}, + provider_batch_id=provider_batch_id, + provider_status=provider_status, + error_message=error_message, + total_items=1, + organization_id=run.organization_id, + project_id=run.project_id, + ) + db.add(batch_job) + db.commit() + db.refresh(batch_job) + return batch_job diff --git a/docs/wiki/modules/evaluations.md b/docs/wiki/modules/evaluations.md index 00669c4b8..23da35bda 100644 --- a/docs/wiki/modules/evaluations.md +++ b/docs/wiki/modules/evaluations.md @@ -1,6 +1,6 @@ # Module: Evaluations -Measures LLM config quality on golden datasets. Three families sharing one dataset/run schema: `text` (cosine + LLM-judge, OpenAI), `stt` (WER/CER, Gemini), `tts` (human annotation, Gemini). Batch-first (provider Batch APIs, not `/llm/call`), plus a synchronous `fast` text mode capped at `EVAL_FAST_MAX_UNIQUE_ROWS`. +Measures LLM config quality on golden datasets. Three families sharing one dataset/run schema: `text` (cosine + LLM-judge, OpenAI), `stt` (WER/CER, Gemini), `tts` (human annotation; Gemini batch + Sarvam/ElevenLabs sync). Batch-first (provider Batch APIs, not `/llm/call`), plus a synchronous `fast` text mode capped at `EVAL_FAST_MAX_UNIQUE_ROWS`. Deep dive: `docs/architecture/kaapi-evaluations-ARCHITECTURE.md` (§3 data model, §5 text pipeline, §6 STT, §7 TTS, §8 batch infra, §9 cron, §12 design decisions). All paths relative to `backend/app/`. @@ -32,11 +32,15 @@ v2 judge field on `EvaluationRun`: `is_judge_run` (bool marker gating native jud ## Services / CRUD - `services/evaluations/` — `evaluation.py`, `dataset.py` (`upload_dataset`; `use_langfuse=False` is the v2 Langfuse-free upload), `fast.py` (`validate_fast_evaluation_inputs` extracted for reuse by both the direct eval-start path and the iteration loop), `batch_job.py`, `validators.py`, `prompt_improvement.py`, `iteration.py` (`validate_and_start_evaluation_iteration`, `compute_round_scores`), `iteration_graph.py` (the LangGraph `StateGraph`: nodes + graph wiring only), `iteration_state.py` (`EvaluationIterationState`, `advance_round_state` — the pure per-round bookkeeping/stop arithmetic, `build_iteration_report`), `iteration_checkpointer.py` (`get_evaluation_iteration_checkpointer`, psycopg pool) - `services/stt_evaluations/`, `services/tts_evaluations/` +- TTS (`services/tts_evaluations/`): `constants.py` is flat module-level constants: `GEMINI_TTS_MODELS` (`gemini-2.5-pro-preview-tts`, `gemini-3.1-flash-tts-preview`; Gemini Batch), `SARVAM_TTS_MODELS` (`bulbul:v3`) and `ELEVENLABS_TTS_MODELS` (`eleven_v4`, `eleven_v3`) — the latter two synthesized directly; `SUPPORTED_TTS_MODELS` (the request validator's allow-list) is their concatenation. Language: submission maps the dataset's `language_id` → `Language.locale` → BCP-47 via `LOCALE_TO_BCP47` (`hi` → `hi-IN`, `or` → `od-IN`), defaulting to `DEFAULT_TTS_LANGUAGE` (`en-IN`); Sarvam gets the BCP-47 tag, ElevenLabs the ISO code via `BCP47_TO_ELEVENLABS_LANG` (unmapped → auto-detect). `sync_generation.py` holds `synthesize_sarvam` / `synthesize_elevenlabs` (raw 24 kHz/16-bit/mono PCM; Sarvam's multi-chunk WAVs are header-stripped per chunk) — retries are the SDKs' built-in ones; a failed sample is marked FAILED with `" synthesis failed: "`. `crud/tts_evaluations/run.py::finalize_tts_run_status` is the single run-status rule (locks the run row `FOR UPDATE`, counts results: `processing` while any row is PENDING, `failed` when no row succeeded, else `completed` with `"N synthesis(es) failed"` if some failed) and is called only by the TTS cron; workers only write `tts_result` rows, so one model failing never fails the run and the final status lags by up to one `CRON_INTERVAL_MINUTES`. - `crud/evaluations/` — `core.py`, `batch.py`, `fast.py` (orchestration + S3/OpenAI IO only; the pieces below carry the rest), `fast_results.py` (pure per-item S3 unit shapes, usage sums, failure threshold), `fast_chunks.py` (`batch_job` bookkeeping + chunk queries/cleanup), `fast_cosine.py` (pure v1 cosine scoring), `fast_traces.py` (pure `TraceData` building), `judge_stage.py` (run-level v2 judge: `resolve_config_prompt`, `judge_rows` pool, per-metric summary rollup), `judge_prompts.py` (rubric text, moved out of `score.py`), `retry.py` (shared OpenAI retry policy), `judge.py` (`METRIC_REGISTRY` + combined judge call; `ground_truth`, `prompt`, and `knowledge_base` metrics, applied per-row by which required inputs the row carries, each spec carrying a `weight` for the overall rollup), `score.py` (`VerdictEnum`/`verdict_from_score`, `OverallSummary`/`compute_overall_summary`), `summary.py` (`generate_run_ai_summary` — one-shot Anthropic `messages.create` call via `ClaudeProvider`, plain-text output — deliberately NOT a `json_schema` `output_config`, since a quotation mark in the prose reads as the JSON string terminator and silently truncates the note; `stop_reason != STOP_REASON_COMPLETE` (the shared `end_turn` constant in `services/llm/providers/claude.py`) is logged as a truncation warning; prompt carries every trace's raw per-question scores + judge rationale, golden/generated answers, and the evaluated config, each trace keyed on `question_id` (the 1-based dataset row number from `merge.py`, cited back as "Question N") rather than the Langfuse `trace_id`, and returns a severity-ranked diagnostic note, not just a qualitative band summary), `embeddings.py`, `cost.py`, `langfuse.py`, `merge.py`, `processing.py`, `cron.py`, `iteration.py` (thin-row CRUD for the iteration loop) - `core/batch/` — shared provider batch clients: `openai.py`, `gemini.py`, `anthropic.py`, `polling.py`, `operations.py` ## Async - Provider batches polled by cron (`crud/evaluations/cron.py`); no long-lived Celery task per run. +- TTS submission (`run_tts_batch_submission` → `services/tts_evaluations/batch_job.py::execute_batch_submission`) creates one `tts_result` per text×model and sets the run `processing` before any model work; Gemini rows go to `start_tts_evaluation_batch` (Gemini Batch), and each Sarvam/ElevenLabs model gets one `run_tts_sync_generation` task (`celery/utils.py::start_tts_sync_generation`, priority 2, global Celery time limits like every other task) → `services/tts_evaluations/sync_generation.py::execute_tts_sync_generation`, which synthesizes that model's PENDING rows `TTS_SYNC_MAX_WORKERS` (2) at a time in a thread pool (idempotent on redelivery), with clients from `SarvamAIProvider.create_client` / `ElevenlabsAIProvider.create_client`, and writes each outcome in its own short session; on timeout/crash it marks its leftover PENDING rows FAILED. A Gemini submission failure or a sync task that can't be queued marks only that model's rows FAILED; the other models carry on. +- TTS cron (`crud/tts_evaluations/cron.py::poll_all_pending_tts_evaluations`) has its own loop (not `cron_utils.poll_all_pending_evaluations_by_type`): it picks up every `processing` TTS run, including sync-only runs with NULL `batch_job_id`, and per run (`poll_tts_run`) builds the Gemini client only when the run has `batch_job` rows (a client failure fails only that run's Gemini rows), polls batches and dispatches `run_tts_result_processing` on SUCCEEDED, fails a model's PENDING rows when its batch ends FAILED/CANCELLED/EXPIRED, and — once the run is older than `CELERY_TASK_TIME_LIMIT` + `TTS_SYNC_STALE_GRACE_SECONDS` — fails leftover PENDING rows of sync models (`"Synthesis did not complete"`) and of Gemini models with no batch job. It then calls `finalize_tts_run_status`. A poll exception leaves the run `processing` for the next tick. Gemini result processing does not finalize; on failure it fails only its model's PENDING rows. +- TTS DB-connection discipline: submission, sync generation and Gemini result processing (`batch_result_processing.py`) never hold a session across network I/O — short sessions for prefetch, session-free CSV read/download/convert/upload (Gemini rows are modified detached and written back in one short session), fresh session for failure marking. `core/batch/operations.py::start_batch_job` commits (ending the read transaction opened by `create_batch_job`'s refresh) before `provider.create_batch` for every caller. The cron's `poll_batch_status` loop still runs inside the cron session's open transaction. - Fast text runs (v1 cosine + v2 judge) fan out `ceil(total_items / EVAL_FAST_CHUNK_SIZE)` `run_evaluation_fast_chunk` tasks (responses only), then a cron barrier (`dispatch_fast_evaluation_barriers`) enqueues one `run_evaluation_fast_aggregate` once every chunk has a `raw_output_url`. The aggregate merges chunks, then (v2) judges **every** row in that single task — so the judge pool is sized by its own `EVAL_JUDGE_CONCURRENCY` (not the response stage's `EVAL_FAST_API_CONCURRENCY`) to clear the max dataset (`EVAL_FAST_MAX_UNIQUE_ROWS` × `duplication_factor`) under the aggregate's `CELERY_TASK_SOFT_TIME_LIMIT`. No judge fan-out / second barrier. The aggregate persists the score unit (`save_score` → traces to S3 behind `score_trace_url`, summary + overall to the DB) **before** flipping the run to `completed`: on v2 that unit is the only copy of the per-row judge scores, so a crash between the two must leave the run `processing`, never `completed` with its scores gone. Chunk-artifact cleanup and the v1 Langfuse score sync run after completion as a best-effort tail that may never fail the run. - Prompt improvement is job-based with callback delivery: `POST /evaluations/{id}/improve-prompt` validates preconditions, enqueues a `Job` (`JobType.PROMPT_IMPROVEMENT`, `models/job.py`) run by Celery task `run_prompt_improvement` (`celery/tasks/job_execution.py`), and returns `202` with an `LLMJobImmediatePublic` handle. On finish the worker POSTs a single best-effort callback to the caller-supplied `callback_url` (SSRF-guarded via `validate_callback_url`): an `APIResponse[PromptImprovementJobPublic]` (`models/evaluation.py`) carrying the new `ConfigVersion` on success or `error_message` on failure (for an Anthropic fault `error_message` carries the provider's own response body, appended by `_anthropic_error_detail` in `services/evaluations/prompt_improvement.py`). The `ConfigVersion` is persisted regardless of callback outcome. `_call_prompt_drafting_llm` keeps `json_schema` structured output (two fields, `rationale` length-capped) but raises `prompt_generation_failed` when `stop_reason != STOP_REASON_COMPLETE`, so a `max_tokens`-truncated draft fails the job instead of minting a half-written system prompt. Celery redelivery of a `SUCCESS` job re-sends the callback without re-running the LLM. - v2 prompt iteration reuses the same job/Celery/config-version machinery (`services/evaluations/prompt_improvement.py`) and branches on `run.is_judge_run`: the v2 route calls `start_prompt_improvement_job(..., require_judge_run=True)` (a non-judge run → `422 not_a_judge_run`), and the worker drafts from the three-metric judge trace (`_draft_improved_prompt(is_judge_run=True)`, using each metric's score + reasoning) and delivers a `PromptRecommendationJobPublic` callback carrying `recommendation_type` (`Literal["prompt"]`, `models/evaluation.py`; widens to a union when knowledge-base / model recommendations land). Non-judge (v1) runs keep the default `_draft_improved_prompt` brief + `PromptImprovementJobPublic` path unchanged.