From 4e4c6385f3ece91d4c74b4f5824a2f3c6a86b627 Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Tue, 6 Oct 2026 20:59:36 +0530 Subject: [PATCH 1/4] feat(assessment): copy assessment_run columns onto assessment Additive expand step toward folding the 1:1 assessment_run child into its parent. Adds the seven run-only columns to assessment (nullable, no table rewrite), backfills them from assessment_run, and leaves assessment_run untouched so deployed code keeps working. The code switch and the child-table drop follow in a separate PR. --- ...py_assessment_run_columns_to_assessment.py | 141 ++++++++++++++++++ 1 file changed, 141 insertions(+) create mode 100644 backend/app/alembic/versions/086_copy_assessment_run_columns_to_assessment.py diff --git a/backend/app/alembic/versions/086_copy_assessment_run_columns_to_assessment.py b/backend/app/alembic/versions/086_copy_assessment_run_columns_to_assessment.py new file mode 100644 index 000000000..e42594e56 --- /dev/null +++ b/backend/app/alembic/versions/086_copy_assessment_run_columns_to_assessment.py @@ -0,0 +1,141 @@ +"""Copy assessment_run's per-config columns onto assessment (expand step) + +Revision ID: 086 +Revises: 085 +Create Date: 2026-10-06 00:00:00.000000 + +`assessment_run` was designed as a 1:N child so one submission could fan out to up to +four configs. The console dropped multi-config submission and the API-client path +always creates exactly one execution, so in every environment the two tables are 1:1 +and the child is redundant. This is the additive half of folding it into the parent: +the seven run-only columns are added to `assessment` and backfilled from +`assessment_run`, with `assessment_run` left untouched so the deployed code keeps +working. The code switch and the `assessment_run` drop ship in a later migration once +the application reads from these columns. + +Every new column is nullable with no server default so the ADD COLUMN is a catalog-only +change (no table rewrite, no long lock on a live table). RESPONSE assessments have no +run and stay NULL in all seven. The backfill is a single statement guarded by +`config_id IS NULL`, so re-running it after a partial failure copies only what is +missing. Downgrade drops the columns; nothing is lost because `assessment_run` still +holds the source rows. +""" + +import sqlalchemy as sa +from alembic import op +from sqlalchemy.dialects import postgresql + +revision = "086" +down_revision = "085" +branch_labels = None +depends_on = None + +CONFIG_INDEX = "idx_assessment_config" +BATCH_JOB_INDEX = "idx_assessment_batch_job" +CONFIG_FK = "fk_assessment_config_id_config" +BATCH_JOB_FK = "fk_assessment_batch_job_id_batch_job" + + +def upgrade() -> None: + op.add_column( + "assessment", + sa.Column( + "config_id", + postgresql.UUID(as_uuid=True), + nullable=True, + comment="ASSESSMENT-tagged config this assessment runs; NULL for RESPONSE", + ), + ) + op.add_column( + "assessment", + sa.Column("config_version", sa.Integer(), nullable=True), + ) + op.add_column( + "assessment", + sa.Column( + "total_items", + sa.Integer(), + nullable=True, + comment="Submitted row count; drives result sizing without reading rows", + ), + ) + op.add_column( + "assessment", + sa.Column( + "batch_job_id", + sa.Integer(), + nullable=True, + comment="Provider batch job; staged runs track per-stage jobs in execution", + ), + ) + op.add_column( + "assessment", + sa.Column( + "execution", + postgresql.JSONB(astext_type=sa.Text()), + nullable=True, + comment="Staged-batch runtime bag (BatchRunState for BATCH, RunExecution for RUN)", + ), + ) + op.add_column( + "assessment", + sa.Column( + "post_processing_config", + postgresql.JSONB(astext_type=sa.Text()), + nullable=True, + ), + ) + op.add_column( + "assessment", + sa.Column("error_message", sa.Text(), nullable=True), + ) + + op.create_foreign_key(CONFIG_FK, "assessment", "config", ["config_id"], ["id"]) + op.create_foreign_key( + BATCH_JOB_FK, + "assessment", + "batch_job", + ["batch_job_id"], + ["id"], + ondelete="SET NULL", + ) + op.create_index( + CONFIG_INDEX, "assessment", ["config_id", "config_version"], unique=False + ) + op.create_index( + BATCH_JOB_INDEX, + "assessment", + ["batch_job_id"], + unique=False, + postgresql_where=sa.text("batch_job_id IS NOT NULL"), + ) + + op.execute( + """ + UPDATE assessment AS a + SET config_id = r.config_id, + config_version = r.config_version, + total_items = r.total_items, + batch_job_id = r.batch_job_id, + execution = r.execution, + post_processing_config = r.post_processing_config, + error_message = r.error_message + FROM assessment_run AS r + WHERE r.assessment_id = a.id + AND a.config_id IS NULL + """ + ) + + +def downgrade() -> None: + op.drop_index(BATCH_JOB_INDEX, table_name="assessment") + op.drop_index(CONFIG_INDEX, table_name="assessment") + op.drop_constraint(BATCH_JOB_FK, "assessment", type_="foreignkey") + op.drop_constraint(CONFIG_FK, "assessment", type_="foreignkey") + op.drop_column("assessment", "error_message") + op.drop_column("assessment", "post_processing_config") + op.drop_column("assessment", "execution") + op.drop_column("assessment", "batch_job_id") + op.drop_column("assessment", "total_items") + op.drop_column("assessment", "config_version") + op.drop_column("assessment", "config_id") From 7df8645c258a7496d025036706a336662701ef3d Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Wed, 7 Oct 2026 00:21:06 +0530 Subject: [PATCH 2/4] Refactor assessment module to consolidate execution handling and remove legacy RUN pipeline - Merged AssessmentRun into Assessment, simplifying the data model and removing unnecessary complexity. - Updated AssessmentExecution to encapsulate execution state and related data, replacing BatchRunState. - Refactored error handling and result file management to streamline error reporting and file uploads. - Adjusted API routes and documentation to reflect the removal of the RUN method and the focus on BATCH and RESPONSE methods. - Enhanced type annotations and improved code readability across various service modules. - Updated submission handling to ensure consistent execution state management and error handling. --- ...py_assessment_run_columns_to_assessment.py | 167 ++++---- backend/app/api/routes/assessment/__init__.py | 5 +- backend/app/api/routes/assessment/api.py | 6 +- backend/app/celery/tasks/job_execution.py | 7 +- backend/app/crud/assessment/api.py | 154 +++----- backend/app/models/assessment/__init__.py | 14 +- backend/app/models/assessment/assessment.py | 18 +- .../app/models/assessment/assessment_api.py | 99 +++-- backend/app/services/assessment/api/batch.py | 372 ++++++++---------- .../services/assessment/api/result_files.py | 82 ++-- .../app/services/assessment/api/results.py | 107 +++-- .../app/services/assessment/api/submission.py | 67 ++-- docs/wiki/domain-map.md | 2 +- docs/wiki/modules/assessment.md | 21 +- 14 files changed, 501 insertions(+), 620 deletions(-) diff --git a/backend/app/alembic/versions/086_copy_assessment_run_columns_to_assessment.py b/backend/app/alembic/versions/086_copy_assessment_run_columns_to_assessment.py index e42594e56..8b16453ab 100644 --- a/backend/app/alembic/versions/086_copy_assessment_run_columns_to_assessment.py +++ b/backend/app/alembic/versions/086_copy_assessment_run_columns_to_assessment.py @@ -1,24 +1,15 @@ -"""Copy assessment_run's per-config columns onto assessment (expand step) +"""Fold assessment_run into assessment Revision ID: 086 Revises: 085 Create Date: 2026-10-06 00:00:00.000000 -`assessment_run` was designed as a 1:N child so one submission could fan out to up to -four configs. The console dropped multi-config submission and the API-client path -always creates exactly one execution, so in every environment the two tables are 1:1 -and the child is redundant. This is the additive half of folding it into the parent: -the seven run-only columns are added to `assessment` and backfilled from -`assessment_run`, with `assessment_run` left untouched so the deployed code keeps -working. The code switch and the `assessment_run` drop ship in a later migration once -the application reads from these columns. - -Every new column is nullable with no server default so the ADD COLUMN is a catalog-only -change (no table rewrite, no long lock on a live table). RESPONSE assessments have no -run and stay NULL in all seven. The backfill is a single statement guarded by -`config_id IS NULL`, so re-running it after a partial failure copies only what is -missing. Downgrade drops the columns; nothing is lost because `assessment_run` still -holds the source rows. +The child was 1:1 with its parent everywhere, so its config + runtime columns move +onto `assessment` and the table is dropped. `error_message` lands in +`execution.error`; keys the code no longer reads are stripped. `input` goes: nothing +writes it since the console pipeline was retired. Downgrade recreates `assessment_run` +empty: take a `pg_dump -t assessment_run` before upgrading, the rows come back from it. +An assessment with no run fails the `SET NOT NULL` on purpose rather than migrating half. """ import sqlalchemy as sa @@ -31,24 +22,19 @@ depends_on = None CONFIG_INDEX = "idx_assessment_config" -BATCH_JOB_INDEX = "idx_assessment_batch_job" CONFIG_FK = "fk_assessment_config_id_config" -BATCH_JOB_FK = "fk_assessment_batch_job_id_batch_job" +RUN_TABLE = "assessment_run" +# Derivable from the pinned config version, or duplicated in result_files. +DROPPED_EXECUTION_KEYS = ("stage_output_urls", "provider", "model", "input_schema") def upgrade() -> None: op.add_column( "assessment", - sa.Column( - "config_id", - postgresql.UUID(as_uuid=True), - nullable=True, - comment="ASSESSMENT-tagged config this assessment runs; NULL for RESPONSE", - ), + sa.Column("config_id", postgresql.UUID(as_uuid=True), nullable=True), ) op.add_column( - "assessment", - sa.Column("config_version", sa.Integer(), nullable=True), + "assessment", sa.Column("config_version", sa.Integer(), nullable=True) ) op.add_column( "assessment", @@ -56,16 +42,7 @@ def upgrade() -> None: "total_items", sa.Integer(), nullable=True, - comment="Submitted row count; drives result sizing without reading rows", - ), - ) - op.add_column( - "assessment", - sa.Column( - "batch_job_id", - sa.Integer(), - nullable=True, - comment="Provider batch job; staged runs track per-stage jobs in execution", + comment="Submitted row count; sizes the result without reading the rows", ), ) op.add_column( @@ -74,68 +51,98 @@ def upgrade() -> None: "execution", postgresql.JSONB(astext_type=sa.Text()), nullable=True, - comment="Staged-batch runtime bag (BatchRunState for BATCH, RunExecution for RUN)", + comment="Staged-batch runtime (AssessmentExecution); NULL for RESPONSE", ), ) - op.add_column( - "assessment", - sa.Column( - "post_processing_config", - postgresql.JSONB(astext_type=sa.Text()), - nullable=True, - ), - ) - op.add_column( - "assessment", - sa.Column("error_message", sa.Text(), nullable=True), - ) - - op.create_foreign_key(CONFIG_FK, "assessment", "config", ["config_id"], ["id"]) - op.create_foreign_key( - BATCH_JOB_FK, - "assessment", - "batch_job", - ["batch_job_id"], - ["id"], - ondelete="SET NULL", - ) - op.create_index( - CONFIG_INDEX, "assessment", ["config_id", "config_version"], unique=False - ) - op.create_index( - BATCH_JOB_INDEX, - "assessment", - ["batch_job_id"], - unique=False, - postgresql_where=sa.text("batch_job_id IS NOT NULL"), - ) + strip = " ".join(f"- '{key}'" for key in DROPPED_EXECUTION_KEYS) op.execute( - """ + f""" UPDATE assessment AS a SET config_id = r.config_id, config_version = r.config_version, total_items = r.total_items, - batch_job_id = r.batch_job_id, - execution = r.execution, - post_processing_config = r.post_processing_config, - error_message = r.error_message - FROM assessment_run AS r + execution = CASE + WHEN r.error_message IS NULL THEN r.execution {strip} + ELSE (COALESCE(r.execution, '{{}}'::jsonb) {strip}) + || jsonb_build_object('error', r.error_message) + END + FROM {RUN_TABLE} AS r WHERE r.assessment_id = a.id AND a.config_id IS NULL """ ) + op.alter_column("assessment", "config_id", nullable=False) + op.alter_column("assessment", "config_version", nullable=False) + op.alter_column("assessment", "total_items", nullable=False) + op.create_foreign_key(CONFIG_FK, "assessment", "config", ["config_id"], ["id"]) + op.create_index(CONFIG_INDEX, "assessment", ["config_id", "config_version"]) + + op.drop_column("assessment", "input") + op.drop_table(RUN_TABLE) + def downgrade() -> None: - op.drop_index(BATCH_JOB_INDEX, table_name="assessment") + op.create_table( + RUN_TABLE, + sa.Column("id", sa.Integer(), primary_key=True), + sa.Column( + "assessment_id", + postgresql.UUID(as_uuid=True), + sa.ForeignKey( + "assessment.id", + name="fk_assessment_run_assessment_id", + ondelete="CASCADE", + ), + nullable=False, + ), + sa.Column( + "config_id", + postgresql.UUID(as_uuid=True), + sa.ForeignKey("config.id", name="fk_assessment_run_config_id"), + nullable=False, + ), + sa.Column("config_version", sa.Integer(), nullable=False), + sa.Column( + "status", + postgresql.ENUM(name="assessment_status", create_type=False), + server_default="PENDING", + nullable=False, + ), + sa.Column( + "batch_job_id", + sa.Integer(), + sa.ForeignKey( + "batch_job.id", + name="fk_assessment_run_batch_job_id", + ondelete="SET NULL", + ), + nullable=True, + ), + sa.Column("total_items", sa.Integer(), server_default="0", nullable=False), + sa.Column("execution", postgresql.JSONB(astext_type=sa.Text()), nullable=True), + sa.Column( + "post_processing_config", + postgresql.JSONB(astext_type=sa.Text()), + nullable=True, + ), + sa.Column("error_message", sa.Text(), nullable=True), + sa.Column("inserted_at", sa.DateTime(), nullable=False), + sa.Column("updated_at", sa.DateTime(), nullable=False), + ) + op.create_index("idx_assessment_run_assessment_id", RUN_TABLE, ["assessment_id"]) + op.create_index( + "idx_assessment_run_config", RUN_TABLE, ["config_id", "config_version"] + ) + op.create_index("idx_assessment_run_status", RUN_TABLE, ["status"]) + op.add_column( + "assessment", + sa.Column("input", postgresql.JSONB(astext_type=sa.Text()), nullable=True), + ) op.drop_index(CONFIG_INDEX, table_name="assessment") - op.drop_constraint(BATCH_JOB_FK, "assessment", type_="foreignkey") op.drop_constraint(CONFIG_FK, "assessment", type_="foreignkey") - op.drop_column("assessment", "error_message") - op.drop_column("assessment", "post_processing_config") op.drop_column("assessment", "execution") - op.drop_column("assessment", "batch_job_id") op.drop_column("assessment", "total_items") op.drop_column("assessment", "config_version") op.drop_column("assessment", "config_id") diff --git a/backend/app/api/routes/assessment/__init__.py b/backend/app/api/routes/assessment/__init__.py index 266e55c4b..53651f90b 100644 --- a/backend/app/api/routes/assessment/__init__.py +++ b/backend/app/api/routes/assessment/__init__.py @@ -1,7 +1,7 @@ from fastapi import APIRouter, Depends from app.api.permissions import require_feature -from app.api.routes.assessment import assessments, datasets, runs +from app.api.routes.assessment import datasets from app.core.feature_flags import FeatureFlag router = APIRouter( @@ -10,8 +10,7 @@ dependencies=[Depends(require_feature(FeatureFlag.ASSESSMENT))], ) +# RUN routers (assessments, runs) are retired; their modules go in the follow-up PR. router.include_router(datasets.router) -router.include_router(assessments.router) -router.include_router(runs.router) __all__ = ["router"] diff --git a/backend/app/api/routes/assessment/api.py b/backend/app/api/routes/assessment/api.py index c8246a4e8..52444bb29 100644 --- a/backend/app/api/routes/assessment/api.py +++ b/backend/app/api/routes/assessment/api.py @@ -85,7 +85,7 @@ def list_assessments( offset: Annotated[int, Query(ge=0)] = 0, ) -> APIResponse[list[AssessmentSummary]]: """BATCH assessments newest-first, optionally narrowed to one config or version.""" - rows = api_crud.list_assessments_with_execution( + rows = api_crud.list_assessments( session=session, organization_id=auth_context.organization_.id, project_id=auth_context.project_.id, @@ -96,8 +96,8 @@ def list_assessments( ) return APIResponse.success_response( data=[ - results.build_summary(assessment, execution, submission_name) - for assessment, execution, submission_name in rows + results.build_summary(assessment, submission_name) + for assessment, submission_name in rows ] ) diff --git a/backend/app/celery/tasks/job_execution.py b/backend/app/celery/tasks/job_execution.py index 483384364..baa9722c6 100644 --- a/backend/app/celery/tasks/job_execution.py +++ b/backend/app/celery/tasks/job_execution.py @@ -15,6 +15,7 @@ import logging from collections.abc import Callable from typing import TYPE_CHECKING, TypeVar +from uuid import UUID from asgi_correlation_id import correlation_id from celery import Task, current_task @@ -370,7 +371,7 @@ def run_assessment_pipeline( @gevent_timeout(settings.CELERY_TASK_SOFT_TIME_LIMIT, "run_assessment_api_batch") def run_assessment_api_batch( self, - execution_id: int, + assessment_id: str, organization_id: int, project_id: int, trace_id: str, @@ -389,7 +390,7 @@ def run_assessment_api_batch( result = _run_with_otel_parent( self, lambda: run_batch_stage( - execution_id=execution_id, + assessment_id=UUID(assessment_id), organization_id=organization_id, project_id=project_id, ), @@ -397,7 +398,7 @@ def run_assessment_api_batch( if result and result.get("requeue"): self.apply_async( kwargs={ - "execution_id": execution_id, + "assessment_id": assessment_id, "organization_id": organization_id, "project_id": project_id, "trace_id": trace_id, diff --git a/backend/app/crud/assessment/api.py b/backend/app/crud/assessment/api.py index f6f7cd360..7b1dca1b9 100644 --- a/backend/app/crud/assessment/api.py +++ b/backend/app/crud/assessment/api.py @@ -1,12 +1,10 @@ -"""Assessment API-client CRUD — method-based Assessment / AssessmentRun writes. +"""Assessment API-client CRUD — writes for the BATCH pipeline on the `assessment` row. -Kept separate from the UI-only crud (core/cron/processing/batch): -writes only the new method-based columns and leaves the RUN-only `execution` -and `submission_id` fields NULL. +Kept separate from the legacy RUN crud (core/cron/processing/batch), which is retired. """ import logging -from typing import Any, TypeVar, cast +from typing import Any from uuid import UUID, uuid4 from sqlalchemy import cast as sa_cast @@ -18,36 +16,38 @@ from app.core.util import now from app.models.assessment import ( Assessment, + AssessmentExecution, AssessmentMethod, - AssessmentRun, AssessmentStatus, AssessmentSubmission, - BatchRunState, ) logger = logging.getLogger(__name__) -# update_status works on either status-bearing row; the TypeVar preserves which. -StatusModel = TypeVar("StatusModel", Assessment, AssessmentRun) - def create_assessment( *, session: Session, method: AssessmentMethod, - input: dict[str, Any] | None, + config_id: UUID, + config_version: int, + total_items: int, organization_id: int, project_id: int, assessment_id: UUID | None = None, submission_id: UUID | None = None, submission_input: str | None = None, experiment_name: str | None = None, + execution: AssessmentExecution | None = None, ) -> Assessment: - """Insert the parent row; pass ``assessment_id`` when the object key already used it.""" + """Insert the row; pass ``assessment_id`` when the object key already used it.""" assessment = Assessment( id=assessment_id or uuid4(), method=method, - input=input, + config_id=config_id, + config_version=config_version, + total_items=total_items, + execution=execution.model_dump(mode="json") if execution else None, submission_id=submission_id, submission_input=submission_input, experiment_name=experiment_name, @@ -60,7 +60,8 @@ def create_assessment( session.refresh(assessment) logger.info( f"[create_assessment] Created | assessment_id: {assessment.id} | " - f"method: {method} | org: {organization_id} | project: {project_id}" + f"method: {method} | config: {config_id} v{config_version} | " + f"org: {organization_id} | project: {project_id}" ) return assessment @@ -79,67 +80,29 @@ def set_assessment_job( return assessment -def create_execution( - *, - session: Session, - assessment_id: UUID, - config_id: UUID, - config_version: int, - total_items: int, -) -> AssessmentRun: - execution = AssessmentRun( - assessment_id=assessment_id, - config_id=config_id, - config_version=config_version, - status=AssessmentStatus.PENDING, - total_items=total_items, - ) - session.add(execution) - session.commit() - session.refresh(execution) - logger.info( - f"[create_execution] Created | execution_id: {execution.id} | " - f"assessment_id: {assessment_id} | config_id: {config_id} v{config_version}" - ) - return execution - - -def set_execution_batch_job( - *, session: Session, execution: AssessmentRun, batch_job_id: int -) -> AssessmentRun: - execution.batch_job_id = batch_job_id - execution.updated_at = now() - session.add(execution) +def save_execution_state( + *, session: Session, assessment: Assessment, state: AssessmentExecution +) -> Assessment: + """Persist the whole runtime bag. Reassigned + flagged: JSONB in-place edits are invisible to SQLAlchemy.""" + assessment.execution = state.model_dump(mode="json") + flag_modified(assessment, "execution") + assessment.updated_at = now() + session.add(assessment) session.commit() - session.refresh(execution) + session.refresh(assessment) logger.info( - f"[set_execution_batch_job] Linked batch job | execution_id: {execution.id} | " - f"batch_job_id: {batch_job_id}" + f"[save_execution_state] Saved | assessment_id: {assessment.id} | " + f"stage: {state.stage} | stage_status: {state.stage_status}" ) - return execution - + return assessment -def save_execution_state( - *, session: Session, execution: AssessmentRun, state: BatchRunState -) -> AssessmentRun: - """Persist the whole staged-batch runtime bag onto ``execution.execution``. - JSONB in-place mutation is invisible to SQLAlchemy, so we reassign the column - and flag it modified rather than mutating the existing dict. - """ - # mypy treats a TypedDict as incompatible with the column's plain dict[str, Any]; - # the cast is erased at runtime (a TypedDict already is a dict). - execution.execution = cast(dict[str, Any], state) - flag_modified(execution, "execution") - execution.updated_at = now() - session.add(execution) - session.commit() - session.refresh(execution) - logger.info( - f"[save_execution_state] Saved | execution_id: {execution.id} | " - f"stage: {state.get('stage')} | stage_status: {state.get('stage_status')}" +def load_execution_state(assessment: Assessment) -> AssessmentExecution | None: + return ( + AssessmentExecution.model_validate(assessment.execution) + if assessment.execution + else None ) - return execution def set_result_files( @@ -147,9 +110,8 @@ def set_result_files( ) -> Assessment: """Shallow-merge ``files`` into ``assessment.result_files``, one record per file kind. - The merge is server-side (``||``, right-hand side wins per key) because two drivers - can touch this row within the same second; a read-modify-write would drop the loser's - kinds instead of keeping both. + Server-side ``||`` (right-hand wins per key): two drivers can touch this row within + the same second, and a read-modify-write would drop the loser's kinds. """ statement = ( update(Assessment) @@ -174,21 +136,20 @@ def set_result_files( def update_status( - *, session: Session, obj: StatusModel, status: AssessmentStatus -) -> StatusModel: - """Set status on an Assessment or AssessmentRun; both carry `status`/`updated_at`.""" - obj.status = status - obj.updated_at = now() - session.add(obj) + *, session: Session, assessment: Assessment, status: AssessmentStatus +) -> Assessment: + assessment.status = status + assessment.updated_at = now() + session.add(assessment) session.commit() - session.refresh(obj) + session.refresh(assessment) logger.info( - f"[update_status] Updated | {type(obj).__name__}: {obj.id} | status: {status}" + f"[update_status] Updated | assessment_id: {assessment.id} | status: {status}" ) - return obj + return assessment -def list_assessments_with_execution( +def list_assessments( *, session: Session, organization_id: int, @@ -197,20 +158,10 @@ def list_assessments_with_execution( config_version: int | None = None, limit: int = 50, offset: int = 0, -) -> list[tuple[Assessment, AssessmentRun | None, str | None]]: - """BATCH assessments newest-first, each with its execution and submission name. - - BATCH only: it has exactly one execution, so the join yields one row per assessment; - a RUN parent has one run per config and would repeat. Outer joins so an inline BATCH - (no submission) and a failed execution insert still list. - """ +) -> list[tuple[Assessment, str | None]]: + """BATCH assessments newest-first with their submission name (outer join: inline BATCH has none).""" statement = ( - select(Assessment, AssessmentRun, AssessmentSubmission.name) - .join( - AssessmentRun, - col(AssessmentRun.assessment_id) == col(Assessment.id), - isouter=True, - ) + select(Assessment, AssessmentSubmission.name) .join( AssessmentSubmission, col(AssessmentSubmission.id) == col(Assessment.submission_id), @@ -221,9 +172,9 @@ def list_assessments_with_execution( .where(Assessment.project_id == project_id) ) if config_id is not None: - statement = statement.where(AssessmentRun.config_id == config_id) + statement = statement.where(Assessment.config_id == config_id) if config_version is not None: - statement = statement.where(AssessmentRun.config_version == config_version) + statement = statement.where(Assessment.config_version == config_version) statement = ( statement.order_by( @@ -233,12 +184,3 @@ def list_assessments_with_execution( .offset(offset) ) return list(session.exec(statement).all()) - - -def list_executions(*, session: Session, assessment_id: UUID) -> list[AssessmentRun]: - statement = ( - select(AssessmentRun) - .where(AssessmentRun.assessment_id == assessment_id) - .order_by(AssessmentRun.id.asc()) - ) - return list(session.exec(statement).all()) diff --git a/backend/app/models/assessment/__init__.py b/backend/app/models/assessment/__init__.py index 0e52cdc37..f7f608156 100644 --- a/backend/app/models/assessment/__init__.py +++ b/backend/app/models/assessment/__init__.py @@ -29,11 +29,13 @@ StageStatus, ) from app.models.assessment.assessment_api import ( + ApiStage, AssessmentBatchResult, AssessmentCallback, AssessmentCounts, AssessmentCreate, AssessmentDetailResponse, + AssessmentExecution, AssessmentInput, AssessmentOutput, AssessmentResult, @@ -44,13 +46,14 @@ AssessmentSummary, Attachment, BatchInput, - BatchRunState, ParsedResult, + PipelineStep, PreFilter, PreFilterVerdict, ResponseInput, + StageCounters, + StageKind, Submission, - Verdict, derive_method, ) from app.models.assessment.submission import ( @@ -91,10 +94,13 @@ "Attachment", "ResponseInput", "BatchInput", - "BatchRunState", "ParsedResult", - "Verdict", "Submission", + "ApiStage", + "StageKind", + "PipelineStep", + "StageCounters", + "AssessmentExecution", "AssessmentInput", "derive_method", "AssessmentCreate", diff --git a/backend/app/models/assessment/assessment.py b/backend/app/models/assessment/assessment.py index 259f975ad..c59877e38 100644 --- a/backend/app/models/assessment/assessment.py +++ b/backend/app/models/assessment/assessment.py @@ -138,7 +138,7 @@ class RunExecution(BaseModel): class Assessment(SQLModel, table=True): - """Parent — one submission: method + shared data source.""" + """One submission assessed against one pinned config version.""" __tablename__ = "assessment" __table_args__ = ( @@ -150,6 +150,7 @@ class Assessment(SQLModel, table=True): Index( "idx_assessment_job", "job_id", postgresql_where=text("job_id IS NOT NULL") ), + Index("idx_assessment_config", "config_id", "config_version"), ) id: UUID = SQLField( @@ -185,12 +186,21 @@ class Assessment(SQLModel, table=True): }, ) - input: dict[str, Any] | None = SQLField( + config_id: UUID = SQLField(foreign_key="config.id", nullable=False) + config_version: int = SQLField(nullable=False) + total_items: int = SQLField( + default=0, + nullable=False, + sa_column_kwargs={ + "comment": "Submitted row count; sizes the result without reading the rows" + }, + ) + execution: dict[str, Any] | None = SQLField( default=None, sa_column=Column( JSONB, nullable=True, - comment="Method-shaped: ResponseInput (RESPONSE) / InputBinding (RUN); NULL for API-client BATCH, which uses submission_input", + comment="Staged-batch runtime (AssessmentExecution); NULL for RESPONSE", ), ) submission_input: str | None = SQLField( @@ -245,7 +255,7 @@ class Assessment(SQLModel, table=True): class AssessmentRun(SQLModel, table=True): - """Child — one config execution (BATCH + RUN; RESPONSE has no run).""" + """Retired: the table is dropped in 086. The class leaves with the RUN code.""" __tablename__ = "assessment_run" __table_args__ = ( diff --git a/backend/app/models/assessment/assessment_api.py b/backend/app/models/assessment/assessment_api.py index 38229d4f9..d680e20ce 100644 --- a/backend/app/models/assessment/assessment_api.py +++ b/backend/app/models/assessment/assessment_api.py @@ -5,7 +5,8 @@ """ from datetime import datetime -from typing import Annotated, Any, NotRequired, TypedDict +from enum import StrEnum +from typing import Annotated, Any, TypedDict from uuid import UUID from pydantic import BaseModel, ConfigDict, Field, HttpUrl, JsonValue, model_validator @@ -15,6 +16,7 @@ AssessmentConfigRef, AssessmentMethod, AssessmentStatus, + StageStatus, ) from app.models.llm.request import ImageInput, PDFInput @@ -62,11 +64,68 @@ def _exactly_one_row_source(self) -> "BatchInput": return self -class Verdict(TypedDict): - """A pre-filter's parsed judgement for one row.""" +class ApiStage(StrEnum): + """Pipeline stage identifiers; names match the pre-filter config + result fields.""" + + TOPIC_RELEVANCE = "topic_relevance" + ASSESSMENT = "assessment" + + +class StageKind(StrEnum): + GATE = "GATE" + PASS_THROUGH = "PASS_THROUGH" + ASSESSMENT = "ASSESSMENT" + + +class PreFilterVerdict(BaseModel): + """Structured pre-filter result per item.""" verdict: bool - reasoning: str + reasoning: str = "" + + +class PipelineStep(BaseModel): + stage: ApiStage + kind: StageKind + + +class StageCounters(BaseModel): + total: int = 0 + passed: int = 0 + rejected: int = 0 + + +class AssessmentExecution(BaseModel): + """Runtime state of the staged BATCH pipeline, stored on ``assessment.execution``. + + Advanced one Celery task at a time by ``run_batch_stage``; ``stage_status`` keys the + idempotent redelivery. Row-indexed maps use int keys (JSON strings coerce on load). + """ + + model_config = ConfigDict(extra="forbid") + + pipeline: list[PipelineStep] + stage: ApiStage + stage_status: StageStatus + stage_batches: dict[ApiStage, int] = Field(default_factory=dict) + stage_errors: dict[ApiStage, dict[int, str]] = Field(default_factory=dict) + verdicts: dict[ApiStage, dict[int, PreFilterVerdict]] = Field(default_factory=dict) + counters: dict[ApiStage, StageCounters] = Field(default_factory=dict) + gate_passed: list[bool] + callback_url: str | None = None + request_metadata: dict[str, JsonValue] | None = None + error: str | None = None + + def stage_kind(self, stage: ApiStage) -> StageKind: + for step in self.pipeline: + if step.stage == stage: + return step.kind + raise ValueError(f"[stage_kind] Stage {stage} not in pipeline") + + def next_stage(self, current: ApiStage) -> ApiStage | None: + stages = [step.stage for step in self.pipeline] + idx = stages.index(current) + return stages[idx + 1] if idx + 1 < len(stages) else None class ParsedResult(TypedDict): @@ -81,31 +140,6 @@ class ParsedResult(TypedDict): response_id: str | None -class BatchRunState(TypedDict): - """Persisted runtime state of the staged BATCH pipeline, stored on - ``AssessmentRun.execution`` (JSONB). Advanced one Celery tick at a time by - ``run_batch_stage``; keyed off ``stage_status`` for idempotent redelivery.""" - - pipeline: list[dict[str, str]] # ordered [{"stage","kind"}] - stage: str # current stage - stage_status: str # PENDING | PROCESSING | COMPLETED | FAILED - # Values are None-typed at the write sites: batch_job.id is an ORM-optional PK and - # raw_output_url is Optional, so the map value types must admit None. - stage_batches: dict[str, int | None] # stage -> provider batch_job id - stage_output_urls: dict[str, str | None] # stage -> raw result url - # stage -> {row_index -> error}, captured at parse time (raw dumps are too big here) - stage_errors: NotRequired[dict[str, dict[str, str]]] - verdicts: dict[str, dict[str, Verdict]] # stage -> {item_idx -> verdict} - counters: dict[str, dict[str, int]] # stage -> {total,passed,rejected} - gate_passed: list[bool] # per-item still-eligible flag - provider: str - model: str - input_schema: dict[str, Any] | None - callback_url: str | None # None when the client polls instead of receiving a push - request_metadata: dict[str, Any] | None - error: NotRequired[str] # only set on failure (_fail) - - # Strict, tagless discrimination via extra=forbid: an input carrying `data` is a # BatchInput, one carrying only `attachments` is a ResponseInput. The two cannot be # confused — neither accepts the other's distinguishing field. @@ -156,13 +190,6 @@ class AssessmentSubmitResponse(BaseModel): updated_at: datetime -class PreFilterVerdict(BaseModel): - """Structured pre-filter result per item.""" - - verdict: bool - reasoning: str = "" - - class PreFilter(BaseModel): """Grouped pre-filter verdicts for one item; each is null if not configured.""" diff --git a/backend/app/services/assessment/api/batch.py b/backend/app/services/assessment/api/batch.py index 3f296e552..363ecad72 100644 --- a/backend/app/services/assessment/api/batch.py +++ b/backend/app/services/assessment/api/batch.py @@ -1,9 +1,8 @@ """BATCH API-client staged pipeline: stage build/submit, verdict parsing, advance. -One config -> one execution -> an ordered series of provider batches: -all GATE pre-filters first (config order), then all PASS-THROUGH pre-filters, -then the assessment stage. Runtime state lives entirely in the execution bag -(``AssessmentRun.execution``); no dedicated columns. +One assessment -> an ordered series of provider batches: all GATE pre-filters first +(config order), then all PASS-THROUGH pre-filters, then the assessment stage. Runtime +state lives entirely in ``assessment.execution`` (``AssessmentExecution``). Conditional forwarding: - a GATE stage runs on every row; rows whose verdict fails are marked @@ -16,7 +15,6 @@ import json import logging -from enum import StrEnum from typing import Any, cast from uuid import UUID @@ -49,15 +47,18 @@ from app.crud.credentials import get_provider_credential from app.crud.job import get_batch_job from app.models.assessment import ( + ApiStage, Assessment, AssessmentAttachment, - AssessmentRun, + AssessmentExecution, AssessmentStatus, BatchInput, - BatchRunState, ParsedResult, + PipelineStep, + PreFilterVerdict, + StageCounters, + StageKind, StageStatus, - Verdict, ) from app.models.batch_job import BatchJob, BatchJobType from app.models.config.assessment_blob import ( @@ -105,19 +106,6 @@ def _google_gcp_credential( POLL_COUNTDOWN_SECONDS = settings.CRON_INTERVAL_MINUTES * 60 -class ApiStage(StrEnum): - """Pipeline stage identifiers; names match the pre-filter config + result fields.""" - - TOPIC_RELEVANCE = "topic_relevance" - ASSESSMENT = "assessment" - - -class StageKind(StrEnum): - GATE = "GATE" - PASS_THROUGH = "PASS_THROUGH" - ASSESSMENT = "ASSESSMENT" - - # Structured-output schema injected into every pre-filter stage's completion; the # verdict parser reads ``verdict``. Both keys are required because provider strict # JSON mode (OpenAI) rejects optional properties — ``reasoning`` may be empty. @@ -163,9 +151,7 @@ def is_supported_provider(provider_name: str) -> bool: return provider_name in _SUPPORTED_PROVIDERS -def build_pipeline( - pre_filters: AssessmentPreFilters | None, -) -> list[dict[str, str]]: +def build_pipeline(pre_filters: AssessmentPreFilters | None) -> list[PipelineStep]: """Ordered stages: GATE pre-filters, then PASS-THROUGH pre-filters, then assessment.""" present: list[tuple[ApiStage, TopicRelevanceFilter]] = [] if pre_filters is not None: @@ -173,35 +159,22 @@ def build_pipeline( present.append((ApiStage.TOPIC_RELEVANCE, pre_filters.topic_relevance)) gates = [ - {"stage": stage.value, "kind": StageKind.GATE.value} + PipelineStep(stage=stage, kind=StageKind.GATE) for stage, flt in present if flt.stop_on_fail ] passthrough = [ - {"stage": stage.value, "kind": StageKind.PASS_THROUGH.value} + PipelineStep(stage=stage, kind=StageKind.PASS_THROUGH) for stage, flt in present if not flt.stop_on_fail ] return [ *gates, *passthrough, - {"stage": ApiStage.ASSESSMENT.value, "kind": StageKind.ASSESSMENT.value}, + PipelineStep(stage=ApiStage.ASSESSMENT, kind=StageKind.ASSESSMENT), ] -def next_stage(pipeline: list[dict[str, str]], current: str) -> str | None: - stages = [s["stage"] for s in pipeline] - idx = stages.index(current) - return stages[idx + 1] if idx + 1 < len(stages) else None - - -def _stage_kind(pipeline: list[dict[str, str]], stage: str) -> str: - for entry in pipeline: - if entry["stage"] == stage: - return entry["kind"] - raise ValueError(f"[_stage_kind] Stage {stage} not in pipeline") - - def column_kinds( columns: list[str], input_columns: dict[str, Any] ) -> tuple[list[str], list[AssessmentAttachment]]: @@ -247,7 +220,7 @@ def build_rows( return rows, text_columns, attachments -def _stage_prompt(blob: AssessmentConfigBlob, stage: str) -> str | None: +def _stage_prompt(blob: AssessmentConfigBlob, stage: ApiStage) -> str | None: """Per-row prompt template for a stage, sourced from the config's ``submission``. The assessment carries a mandatory ``submission``; a pre-filter may carry its own @@ -262,7 +235,7 @@ def _stage_prompt(blob: AssessmentConfigBlob, stage: str) -> str | None: def _prefilter_for_stage( - blob: AssessmentConfigBlob, stage: str + blob: AssessmentConfigBlob, stage: ApiStage ) -> TopicRelevanceFilter: """The pre-filter config object for a pre-filter stage (topic_relevance).""" pre = blob.pre_filters @@ -274,7 +247,7 @@ def _prefilter_for_stage( ) -def _stage_params(blob: AssessmentConfigBlob, stage: str) -> dict[str, Any]: +def _stage_params(blob: AssessmentConfigBlob, stage: ApiStage) -> dict[str, Any]: """Kaapi params for a stage: the assessment uses its own params + output schema; each pre-filter uses ITS own params (criteria in params.instructions), with the verdict schema + gate directive layered on. @@ -301,7 +274,9 @@ def _stage_params(blob: AssessmentConfigBlob, stage: str) -> dict[str, Any]: return params -def _stage_provider_model(blob: AssessmentConfigBlob, stage: str) -> tuple[str, str]: +def stage_provider_model( + blob: AssessmentConfigBlob, stage: ApiStage +) -> tuple[str, str]: """Provider + model for a stage: each pre-filter uses its own; assessment uses the config's.""" if stage == ApiStage.TOPIC_RELEVANCE: flt = _prefilter_for_stage(blob, stage) @@ -309,6 +284,21 @@ def _stage_provider_model(blob: AssessmentConfigBlob, stage: str) -> tuple[str, return blob.assessment.provider, blob.assessment.params["model"] +def resolve_blob(session: Session, assessment: Assessment) -> AssessmentConfigBlob: + """The pinned config version's blob; raises on a deleted version or an old shape.""" + from app.crud.config.version import ConfigVersionCrud + from app.models.config.config import ConfigTag + + version_crud = ConfigVersionCrud( + session=session, + config_id=assessment.config_id, + project_id=assessment.project_id, + tag=ConfigTag.ASSESSMENT, + ) + version = version_crud.exists_or_raise(version_number=assessment.config_version) + return AssessmentConfigBlob.model_validate(version.config_blob) + + def _submit_provider_batch( *, session: Session, @@ -559,63 +549,63 @@ def parse_batch_results( return parsed -def _parse_verdict(output: str | None) -> Verdict: +def _parse_verdict(output: str | None) -> PreFilterVerdict: """Read a pre-filter verdict. Unparseable output fails open (verdict=True).""" if not output: - return {"verdict": True, "reasoning": ""} + return PreFilterVerdict(verdict=True) try: data = json.loads(output) - return { - "verdict": bool(data.get("verdict", True)), - "reasoning": str(data.get("reasoning", "")), - } + return PreFilterVerdict( + verdict=bool(data.get("verdict", True)), + reasoning=str(data.get("reasoning", "")), + ) except (json.JSONDecodeError, TypeError, AttributeError): logger.warning( "[_parse_verdict] Unparseable verdict, failing open | output=%s", output[:200], ) - return {"verdict": True, "reasoning": ""} + return PreFilterVerdict(verdict=True) def _record_stage( - bag: BatchRunState, stage: str, kind: str, parsed: dict[int, ParsedResult] + bag: AssessmentExecution, + stage: ApiStage, + kind: StageKind, + parsed: dict[int, ParsedResult], ) -> None: """Fold a completed stage's parsed results into the bag per its kind.""" if kind == StageKind.ASSESSMENT: return # assessment outputs are read from object store at result time - verdicts: dict[str, Verdict] = {} + verdicts: dict[int, PreFilterVerdict] = {} passed_count = 0 for idx, out in parsed.items(): verdict = _parse_verdict(out.get("output")) - verdicts[str(idx)] = verdict - if verdict["verdict"]: + verdicts[idx] = verdict + if verdict.verdict: passed_count += 1 elif kind == StageKind.GATE: - bag["gate_passed"][idx] = False + bag.gate_passed[idx] = False - bag.setdefault("verdicts", {})[stage] = verdicts - bag.setdefault("counters", {})[stage] = { - "total": len(parsed), - "passed": passed_count, - "rejected": len(parsed) - passed_count, - } + bag.verdicts[stage] = verdicts + bag.counters[stage] = StageCounters( + total=len(parsed), passed=passed_count, rejected=len(parsed) - passed_count + ) -def _row_subset(bag: BatchRunState, stage: str, kind: str, total: int) -> list[int]: +def _row_subset(bag: AssessmentExecution, kind: StageKind, total: int) -> list[int]: if kind == StageKind.ASSESSMENT: - return [i for i in range(total) if bag["gate_passed"][i]] + return [i for i in range(total) if bag.gate_passed[i]] return list(range(total)) def _submit_stage( *, session: Session, - execution: AssessmentRun, assessment: Assessment, blob: AssessmentConfigBlob, - bag: BatchRunState, - stage: str, + bag: AssessmentExecution, + stage: ApiStage, organization_id: int, project_id: int, ) -> bool: @@ -625,18 +615,19 @@ def _submit_stage( in-memory bag can predate a concurrent task's write, so the row is re-read here. Rows are streamed from storage once per stage submit; only this stage's subset is held. """ - session.refresh(execution) - persisted = cast(BatchRunState, execution.execution or {}) - in_flight_batch_id = (persisted.get("stage_batches") or {}).get(stage) + session.refresh(assessment) + persisted = api.load_execution_state(assessment) + in_flight_batch_id = persisted.stage_batches.get(stage) if persisted else None if ( - in_flight_batch_id is not None - and persisted.get("stage_status") == StageStatus.PROCESSING.value + persisted is not None + and in_flight_batch_id is not None + and persisted.stage_status == StageStatus.PROCESSING ): # Deliberately leaves `bag` unpersisted: the stored state is the fresher one. logger.warning( "[_submit_stage] Stage already submitted, skipping duplicate | " - "execution_id=%s | stage=%s | batch_job=%s", - execution.id, + "assessment_id=%s | stage=%s | batch_job=%s", + assessment.id, stage, in_flight_batch_id, ) @@ -644,32 +635,32 @@ def _submit_stage( from app.services.assessment.api.submission_store import open_submission_rows - kind = _stage_kind(bag["pipeline"], stage) + kind = bag.stage_kind(stage) input_columns = { name: col.model_dump(exclude_none=True) for name, col in blob.input_schema.items() } text_columns, attachments = column_kinds(list(input_columns), input_columns) - subset = _row_subset(bag, stage, kind, execution.total_items) + subset = _row_subset(bag, kind, assessment.total_items) if not subset: # No rows left for this stage (everything gated out upstream). Persist the # empty counters and return False so the caller advances/finalizes instead of # re-submitting a PENDING stage forever. logger.info( - "[_submit_stage] Empty subset, skipping | execution_id=%s | stage=%s", - execution.id, + "[_submit_stage] Empty subset, skipping | assessment_id=%s | stage=%s", + assessment.id, stage, ) - bag.setdefault("counters", {})[stage] = {"total": 0, "passed": 0, "rejected": 0} - api.save_execution_state(session=session, execution=execution, state=bag) + bag.counters[stage] = StageCounters() + api.save_execution_state(session=session, assessment=assessment, state=bag) return False wanted = set(subset) with open_submission_rows(session=session, assessment=assessment) as stream: rows = [row for idx, row in enumerate(stream) if idx in wanted] - provider_name, model = _stage_provider_model(blob, stage) + provider_name, model = stage_provider_model(blob, stage) batch_job = _submit_provider_batch( session=session, provider_name=provider_name, @@ -682,19 +673,18 @@ def _submit_stage( row_indices=subset, organization_id=organization_id, project_id=project_id, - description=f"assessment-{execution.id}-{stage}", + description=f"assessment-{assessment.id}-{stage}", ) + if batch_job.id is None: + raise ValueError("[_submit_stage] Batch job was not persisted") - bag["stage"] = stage - bag["stage_status"] = StageStatus.PROCESSING.value - bag.setdefault("stage_batches", {})[stage] = batch_job.id - api.set_execution_batch_job( - session=session, execution=execution, batch_job_id=batch_job.id - ) - api.save_execution_state(session=session, execution=execution, state=bag) + bag.stage = stage + bag.stage_status = StageStatus.PROCESSING + bag.stage_batches[stage] = batch_job.id + api.save_execution_state(session=session, assessment=assessment, state=bag) logger.info( - "[_submit_stage] Submitted | execution_id=%s | stage=%s | batch_job=%s | rows=%s", - execution.id, + "[_submit_stage] Submitted | assessment_id=%s | stage=%s | batch_job=%s | rows=%s", + assessment.id, stage, batch_job.id, len(subset), @@ -767,20 +757,17 @@ def _poll_outcome( def _finalize( - session: Session, - execution: AssessmentRun, - assessment: Assessment, - bag: BatchRunState, + session: Session, assessment: Assessment, bag: AssessmentExecution ) -> None: from app.services.assessment.api.callbacks import deliver from app.services.assessment.api.result_files import finalize_result_files from app.services.assessment.api.results import build_result - bag["stage"] = ApiStage.ASSESSMENT.value - bag["stage_status"] = StageStatus.COMPLETED.value - api.save_execution_state(session=session, execution=execution, state=bag) + bag.stage = ApiStage.ASSESSMENT + bag.stage_status = StageStatus.COMPLETED + api.save_execution_state(session=session, assessment=assessment, state=bag) - result = build_result(session=session, assessment=assessment) + result = build_result(session=session, assessment=assessment, bag=bag) errors = sum(1 for item in result.items if item.error) if result.items and errors == len(result.items): status = AssessmentStatus.FAILED @@ -789,72 +776,64 @@ def _finalize( else: status = AssessmentStatus.COMPLETED - api.update_status(session=session, obj=execution, status=status) - api.update_status(session=session, obj=assessment, status=status) + api.update_status(session=session, assessment=assessment, status=status) logger.info( - "[_finalize] Completed | execution_id=%s | status=%s | items=%s | errors=%s", - execution.id, + "[_finalize] Completed | assessment_id=%s | status=%s | items=%s | errors=%s", + assessment.id, status, len(result.items), errors, ) # Before the callback check: durability must not depend on a callback existing. - finalize_result_files( - session=session, execution=execution, assessment=assessment, bag=bag - ) + finalize_result_files(session=session, assessment=assessment, bag=bag) - callback_url = bag.get("callback_url") - if callback_url: + if bag.callback_url: deliver( session=session, assessment=assessment, result=result, - callback_url=callback_url, - request_metadata=bag.get("request_metadata"), + callback_url=bag.callback_url, + request_metadata=bag.request_metadata, failure_message=None, ) def _fail( session: Session, - execution: AssessmentRun, assessment: Assessment, - bag: BatchRunState, + bag: AssessmentExecution, message: str, ) -> None: from app.services.assessment.api.callbacks import deliver from app.services.assessment.api.result_files import finalize_result_files from app.services.assessment.api.results import build_result - bag["stage_status"] = StageStatus.FAILED.value - bag["error"] = message - api.save_execution_state(session=session, execution=execution, state=bag) - execution.error_message = message - api.update_status(session=session, obj=execution, status=AssessmentStatus.FAILED) - api.update_status(session=session, obj=assessment, status=AssessmentStatus.FAILED) + bag.stage_status = StageStatus.FAILED + bag.error = message + api.save_execution_state(session=session, assessment=assessment, state=bag) + api.update_status( + session=session, assessment=assessment, status=AssessmentStatus.FAILED + ) logger.error( - "[_fail] Execution failed | execution_id=%s | message=%s", execution.id, message + "[_fail] Assessment failed | assessment_id=%s | message=%s", + assessment.id, + message, ) # Before the callback check: durability must not depend on a callback existing. finalize_result_files( - session=session, - execution=execution, - assessment=assessment, - bag=bag, - failure_message=message, + session=session, assessment=assessment, bag=bag, failure_message=message ) - callback_url = bag.get("callback_url") - if callback_url: - result = build_result(session=session, assessment=assessment) + if bag.callback_url: + result = build_result(session=session, assessment=assessment, bag=bag) deliver( session=session, assessment=assessment, result=result, - callback_url=callback_url, - request_metadata=bag.get("request_metadata"), + callback_url=bag.callback_url, + request_metadata=bag.request_metadata, failure_message=message, ) @@ -862,11 +841,10 @@ def _fail( def _advance_or_finalize( *, session: Session, - execution: AssessmentRun, assessment: Assessment, blob: AssessmentConfigBlob, - bag: BatchRunState, - stage: str, + bag: AssessmentExecution, + stage: ApiStage, organization_id: int, project_id: int, ) -> dict[str, bool]: @@ -878,16 +856,15 @@ def _advance_or_finalize( """ from app.services.assessment.api.submission_store import SubmissionUnavailableError - nxt = next_stage(bag["pipeline"], stage) + nxt = bag.next_stage(stage) if nxt is None: - _finalize(session, execution, assessment, bag) + _finalize(session, assessment, bag) return {"requeue": False} try: - bag["stage"] = nxt - bag["stage_status"] = StageStatus.PENDING.value + bag.stage = nxt + bag.stage_status = StageStatus.PENDING submitted = _submit_stage( session=session, - execution=execution, assessment=assessment, blob=blob, bag=bag, @@ -899,21 +876,20 @@ def _advance_or_finalize( # Storage blip, not a bad run: the stage is still PENDING, so retry the task. logger.warning( "[_advance_or_finalize] Submission unreadable, will retry | " - "execution_id=%s | stage=%s | %s", - execution.id, + "assessment_id=%s | stage=%s | %s", + assessment.id, nxt, exc, ) return {"requeue": True} except Exception as exc: - _fail(session, execution, assessment, bag, str(exc)) + _fail(session, assessment, bag, str(exc)) return {"requeue": False} if not submitted: - bag["stage_status"] = StageStatus.COMPLETED.value - api.save_execution_state(session=session, execution=execution, state=bag) + bag.stage_status = StageStatus.COMPLETED + api.save_execution_state(session=session, assessment=assessment, state=bag) return _advance_or_finalize( session=session, - execution=execution, assessment=assessment, blob=blob, bag=bag, @@ -925,56 +901,45 @@ def _advance_or_finalize( def run_batch_stage( - *, execution_id: int, organization_id: int, project_id: int + *, assessment_id: UUID, organization_id: int, project_id: int ) -> dict[str, bool]: """Run one step of the staged pipeline. Returns ``{"requeue": bool}``. Idempotent: keyed off ``stage_status`` in the bag — a redelivery either re-polls the in-flight batch or re-submits a stage that was never dispatched. """ - # Lazy (like callbacks/results below): result_files imports ApiStage from this module. from app.services.assessment.api.result_files import record_stage_dump from app.services.assessment.api.submission_store import SubmissionUnavailableError with Session(engine) as session: - execution = session.exec( - select(AssessmentRun) - .where(col(AssessmentRun.id) == execution_id) + assessment = session.exec( + select(Assessment) + .where(col(Assessment.id) == assessment_id) .with_for_update(skip_locked=True) ).first() - if execution is None: - logger.warning( - "[run_batch_stage] execution_id=%s missing or held by another task", - execution_id, - ) - return {"requeue": False} - assessment = session.get(Assessment, execution.assessment_id) if assessment is None: - logger.error( - "[run_batch_stage] parent assessment missing | execution_id=%s", - execution_id, + logger.warning( + "[run_batch_stage] assessment_id=%s missing or held by another task", + assessment_id, ) return {"requeue": False} - if execution.status in ( - AssessmentStatus.COMPLETED, - AssessmentStatus.COMPLETED_WITH_ERRORS, - AssessmentStatus.FAILED, - ): + if assessment.status in AssessmentStatus.terminal(): return {"requeue": False} - bag = cast(BatchRunState, dict(execution.execution or {})) - stage = bag.get("stage") - stage_status = bag.get("stage_status") - if not stage or not bag.get("pipeline"): + bag = api.load_execution_state(assessment) + if bag is None: logger.error( - "[run_batch_stage] uninitialised bag | execution_id=%s", execution_id + "[run_batch_stage] uninitialised execution | assessment_id=%s", + assessment_id, ) return {"requeue": False} + stage = bag.stage + stage_status = bag.stage_status # Entry log: distinguishes "the task never ran" from "it ran and did nothing". logger.info( - "[run_batch_stage] Step | execution_id=%s | stage=%s | stage_status=%s", - execution_id, + "[run_batch_stage] Step | assessment_id=%s | stage=%s | stage_status=%s", + assessment_id, stage, stage_status, ) @@ -983,18 +948,15 @@ def run_batch_stage( # invalid/old-shape blob). Route these to _fail so the client gets a terminal # callback instead of the run stranding in PROCESSING. try: - blob = AssessmentConfigBlob.model_validate( - _resolve_blob(session, execution, project_id) - ) + blob = resolve_blob(session, assessment) except Exception as exc: - _fail(session, execution, assessment, bag, str(exc)) + _fail(session, assessment, bag, str(exc)) return {"requeue": False} if stage_status == StageStatus.PENDING: try: submitted = _submit_stage( session=session, - execution=execution, assessment=assessment, blob=blob, bag=bag, @@ -1006,8 +968,8 @@ def run_batch_stage( # Storage blip, not a bad run: the stage stays PENDING, retry the task. logger.warning( "[run_batch_stage] Submission unreadable, will retry | " - "execution_id=%s | stage=%s | %s", - execution_id, + "assessment_id=%s | stage=%s | %s", + assessment_id, stage, exc, ) @@ -1015,18 +977,17 @@ def run_batch_stage( except Exception as exc: # Credential/provider/network errors from _submit_provider_batch, not # just ValueError — all are terminal for this run. - _fail(session, execution, assessment, bag, str(exc)) + _fail(session, assessment, bag, str(exc)) return {"requeue": False} if not submitted: # Empty subset (all rows gated out): the stage is done — advance/finalize # rather than re-submitting a PENDING stage forever. - bag["stage_status"] = StageStatus.COMPLETED.value + bag.stage_status = StageStatus.COMPLETED api.save_execution_state( - session=session, execution=execution, state=bag + session=session, assessment=assessment, state=bag ) return _advance_or_finalize( session=session, - execution=execution, assessment=assessment, blob=blob, bag=bag, @@ -1037,18 +998,19 @@ def run_batch_stage( return {"requeue": True} # stage_status == PROCESSING: poll the in-flight batch. - batch_id = (bag.get("stage_batches") or {}).get(stage) + batch_id = bag.stage_batches.get(stage) batch_job = ( get_batch_job(session=session, batch_job_id=batch_id) if batch_id else None ) if batch_job is None: - _fail(session, execution, assessment, bag, f"Stage {stage} batch not found") + _fail(session, assessment, bag, f"Stage {stage} batch not found") return {"requeue": False} + provider_name, _ = stage_provider_model(blob, stage) try: provider = _build_batch_provider( session=session, - provider_name=bag["provider"], + provider_name=provider_name, organization_id=organization_id, project_id=project_id, ) @@ -1058,8 +1020,8 @@ def run_batch_stage( except Exception as exc: # Transient (network/provider hiccup) — the batch is still running; retry. logger.warning( - "[run_batch_stage] poll error, will retry | execution_id=%s | stage=%s | %s", - execution_id, + "[run_batch_stage] poll error, will retry | assessment_id=%s | stage=%s | %s", + assessment_id, stage, exc, ) @@ -1070,7 +1032,6 @@ def run_batch_stage( if outcome == "failed": _fail( session, - execution, assessment, bag, batch_job.error_message or f"Stage {stage} batch failed", @@ -1078,19 +1039,14 @@ def run_batch_stage( return {"requeue": False} # outcome == "completed" - kind = _stage_kind(bag["pipeline"], stage) - parsed = parse_batch_results(results or [], bag["provider"]) + kind = bag.stage_kind(stage) + parsed = parse_batch_results(results or [], provider_name) _record_stage(bag, stage, kind, parsed) - bag["stage_status"] = StageStatus.COMPLETED.value - - stage_errors: dict[str, str] = {} - for idx, out in parsed.items(): - error = out.get("error") - if error: - stage_errors[str(idx)] = error - bag.setdefault("stage_errors", {})[stage] = stage_errors + bag.stage_status = StageStatus.COMPLETED + bag.stage_errors[stage] = { + idx: error for idx, out in parsed.items() if (error := out.get("error")) + } - bag.setdefault("stage_output_urls", {})[stage] = batch_job.raw_output_url # Per stage, not only at terminal time: a run that never terminates still has dumps. record_stage_dump( session=session, @@ -1098,11 +1054,10 @@ def run_batch_stage( stage=stage, url=batch_job.raw_output_url, ) - api.save_execution_state(session=session, execution=execution, state=bag) + api.save_execution_state(session=session, assessment=assessment, state=bag) return _advance_or_finalize( session=session, - execution=execution, assessment=assessment, blob=blob, bag=bag, @@ -1110,20 +1065,3 @@ def run_batch_stage( organization_id=organization_id, project_id=project_id, ) - - -def _resolve_blob( - session: Session, execution: AssessmentRun, project_id: int -) -> dict[str, Any]: - """Fetch the execution's stored ASSESSMENT config_blob dict for re-parsing.""" - from app.crud.config.version import ConfigVersionCrud - from app.models.config.config import ConfigTag - - version_crud = ConfigVersionCrud( - session=session, - config_id=execution.config_id, - project_id=project_id, - tag=ConfigTag.ASSESSMENT, - ) - version = version_crud.exists_or_raise(version_number=execution.config_version) - return version.config_blob diff --git a/backend/app/services/assessment/api/result_files.py b/backend/app/services/assessment/api/result_files.py index 7a08ae56d..2de74232e 100644 --- a/backend/app/services/assessment/api/result_files.py +++ b/backend/app/services/assessment/api/result_files.py @@ -17,10 +17,10 @@ from app.crud.assessment import api from app.crud.job import get_batch_job from app.models.assessment import ( + ApiStage, Assessment, + AssessmentExecution, AssessmentResultFiles, - AssessmentRun, - BatchRunState, ) from app.models.batch_job import BatchJob from app.services.assessment.api.batch import ( @@ -75,20 +75,20 @@ def record_stage_dump( ) -def _row_error_rows(stage_errors: dict[str, dict[str, str]]) -> list[dict[str, Any]]: +def _row_error_rows( + stage_errors: dict[ApiStage, dict[int, str]], +) -> list[dict[str, Any]]: """Per-row errors captured at parse time, flattened across stages.""" - rows: list[dict[str, Any]] = [] - for stage, errors in stage_errors.items(): - for row_index, error in errors.items(): - rows.append( - { - "type": ErrorRecordEnum.ROW_ERROR.value, - "stage": stage, - "row_index": int(row_index) if row_index.isdigit() else row_index, - "error": error, - } - ) - return rows + return [ + { + "type": ErrorRecordEnum.ROW_ERROR.value, + "stage": stage.value, + "row_index": row_index, + "error": error, + } + for stage, errors in stage_errors.items() + for row_index, error in errors.items() + ] def _error_file_rows( @@ -147,9 +147,8 @@ def _error_file_rows( def build_and_upload_errors( *, session: Session, - execution: AssessmentRun, assessment: Assessment, - bag: BatchRunState, + bag: AssessmentExecution, failure_message: str | None, ) -> str | None: """Assemble and upload the run's ``errors.jsonl``; ``None`` when there is nothing @@ -159,23 +158,19 @@ def build_and_upload_errors( rows.append( { "type": ErrorRecordEnum.EXECUTION_ERROR.value, - "stage": bag.get("stage"), + "stage": bag.stage.value, "error": failure_message, } ) - rows.extend(_row_error_rows(bag.get("stage_errors") or {})) - for stage, batch_job_id in (bag.get("stage_batches") or {}).items(): - batch_job = ( - get_batch_job(session=session, batch_job_id=batch_job_id) - if batch_job_id - else None - ) + rows.extend(_row_error_rows(bag.stage_errors)) + for stage, batch_job_id in bag.stage_batches.items(): + batch_job = get_batch_job(session=session, batch_job_id=batch_job_id) if batch_job is not None: rows.extend( _error_file_rows( session=session, assessment=assessment, - stage=stage, + stage=stage.value, batch_job=batch_job, ) ) @@ -188,8 +183,8 @@ def build_and_upload_errors( except Exception: logger.error( "[build_and_upload_errors] Storage unavailable, no errors dump | " - "execution_id=%s | rows=%s", - execution.id, + "assessment_id=%s | rows=%s", + assessment.id, len(rows), exc_info=True, ) @@ -202,9 +197,9 @@ def build_and_upload_errors( subdirectory=assessment_prefix(assessment.id), ) logger.info( - "[build_and_upload_errors] Errors dump %s | execution_id=%s | rows=%s | url=%s", + "[build_and_upload_errors] Errors dump %s | assessment_id=%s | rows=%s | url=%s", "uploaded" if url else "upload failed", - execution.id, + assessment.id, len(rows), url, ) @@ -214,41 +209,32 @@ def build_and_upload_errors( def finalize_result_files( *, session: Session, - execution: AssessmentRun, assessment: Assessment, - bag: BatchRunState, + bag: AssessmentExecution, failure_message: str | None = None, ) -> None: - """Persist every stage dump plus the run's errors.jsonl at terminal time. + """Persist the run's errors.jsonl at terminal time; stage dumps were recorded per stage. - Idempotent (a re-merge replaces only its own kind), so a redelivered tick is + Idempotent (a re-merge replaces only its own kind), so a redelivered task is harmless. Never raises: durability is best-effort, terminating the run is not. """ try: - files: dict[str, dict[str, Any]] = {} - for stage, url in (bag.get("stage_output_urls") or {}).items(): - if not url: - continue - files[stage] = {"object_store_url": url} - errors_url = build_and_upload_errors( session=session, - execution=execution, assessment=assessment, bag=bag, failure_message=failure_message, ) if errors_url: - files[ERRORS_FILE_KIND] = {"object_store_url": errors_url} - - if files: - api.set_result_files(session=session, assessment=assessment, files=files) + api.set_result_files( + session=session, + assessment=assessment, + files={ERRORS_FILE_KIND: {"object_store_url": errors_url}}, + ) except Exception: logger.error( - "[finalize_result_files] Could not persist result files | " - "assessment_id=%s | execution_id=%s", + "[finalize_result_files] Could not persist result files | assessment_id=%s", assessment.id, - execution.id, exc_info=True, ) diff --git a/backend/app/services/assessment/api/results.py b/backend/app/services/assessment/api/results.py index 32463b238..134be7a2f 100644 --- a/backend/app/services/assessment/api/results.py +++ b/backend/app/services/assessment/api/results.py @@ -8,32 +8,34 @@ import json import logging -from typing import Any, cast +from typing import Any from sqlmodel import Session from app.core.cloud import get_cloud_storage from app.crud.assessment import api from app.models.assessment import ( + ApiStage, Assessment, AssessmentBatchResult, AssessmentConfigRef, AssessmentCounts, AssessmentDetailResponse, + AssessmentExecution, AssessmentOutput, AssessmentResult, AssessmentResultRow, - AssessmentRun, AssessmentSubmission, AssessmentSummary, - BatchRunState, ParsedResult, PreFilter, - PreFilterVerdict, Submission, - Verdict, ) -from app.services.assessment.api.batch import ApiStage, parse_batch_results +from app.services.assessment.api.batch import ( + parse_batch_results, + resolve_blob, + stage_provider_model, +) from app.services.assessment.utils.parsing import parse_stored_results logger = logging.getLogger(__name__) @@ -54,31 +56,25 @@ def _parse_assessment(out: ParsedResult) -> dict[str, Any] | str | None: return parsed if isinstance(parsed, dict) else text -def _verdict_obj(verdict: Verdict | None) -> PreFilterVerdict | None: - if not verdict: - return None - return PreFilterVerdict( - verdict=bool(verdict.get("verdict")), - reasoning=str(verdict.get("reasoning") or ""), - ) - - def _load_assessment_outputs( - session: Session, bag: BatchRunState, project_id: int + session: Session, assessment: Assessment ) -> tuple[dict[int, ParsedResult], str | None]: - """Stream + parse the assessment stage output. Returns ``(outputs, load_error)``. + """Stream + parse the assessment stage dump. Returns ``(outputs, load_error)``. ``load_error`` is set only on a failed read of a present URL, so the caller - can flag it per-row instead of mistaking it for a clean run. Missing URL + can flag it per-row instead of mistaking it for a clean run. A missing URL (all rows gated) is a legit empty, not an error. """ - url = (bag.get("stage_output_urls") or {}).get(ApiStage.ASSESSMENT.value) + record = assessment.result_files.get(ApiStage.ASSESSMENT.value) or {} + url = record.get("object_store_url") if not url: return {}, None try: - storage = get_cloud_storage(session=session, project_id=project_id) + blob = resolve_blob(session, assessment) + provider, _ = stage_provider_model(blob, ApiStage.ASSESSMENT) + storage = get_cloud_storage(session=session, project_id=assessment.project_id) raw = parse_stored_results(storage.stream(url).read().decode("utf-8")) - return parse_batch_results(raw, bag.get("provider")), None + return parse_batch_results(raw, provider), None except Exception as exc: logger.warning( "[_load_assessment_outputs] Could not read assessment output | url=%s | %s", @@ -92,29 +88,23 @@ def build_result( *, session: Session, assessment: Assessment, - execution: AssessmentRun | None = None, + bag: AssessmentExecution | None = None, ) -> AssessmentBatchResult: """Build the per-row result from stored verdicts (bag) + assessment output (store). - Status lives on the response envelope (sourced from the parent assessment), so this - result body carries only the rows and their tallies. Pass ``execution`` when the - caller already holds it, so a poll does not re-query it. + Status lives on the response envelope, so this body carries only rows and tallies. + Pass ``bag`` when the caller already holds it, so a poll does not re-parse the column. """ - if execution is None: - executions = api.list_executions(session=session, assessment_id=assessment.id) - execution = executions[0] if executions else None - bag = cast(BatchRunState, (execution.execution or {}) if execution else {}) + if bag is None: + bag = api.load_execution_state(assessment) - # Off the execution, not re-derived: the terminal path must not fetch rows to count them. - total_items = execution.total_items if execution else 0 - - gate_passed = bag.get("gate_passed") or [True] * total_items - verdicts = bag.get("verdicts") or {} - outputs, load_error = _load_assessment_outputs(session, bag, assessment.project_id) + # Off the row, not re-derived: the terminal path must not fetch rows to count them. + total_items = assessment.total_items + gate_passed = bag.gate_passed if bag else [True] * total_items + outputs, load_error = _load_assessment_outputs(session, assessment) # Rows the provider rejected are absent from the output dump; the bag kept their error. - row_errors = (bag.get("stage_errors") or {}).get(ApiStage.ASSESSMENT.value, {}) - - tr_verdicts = verdicts.get(ApiStage.TOPIC_RELEVANCE.value, {}) + row_errors = bag.stage_errors.get(ApiStage.ASSESSMENT, {}) if bag else {} + tr_verdicts = bag.verdicts.get(ApiStage.TOPIC_RELEVANCE, {}) if bag else {} items: list[AssessmentResult] = [] counts = AssessmentCounts() @@ -128,12 +118,12 @@ def build_result( out = outputs[idx] assessment_output = _parse_assessment(out) error = out.get("error") - elif str(idx) in row_errors: - error = row_errors[str(idx)] + elif idx in row_errors: + error = row_errors[idx] elif load_error: error = load_error - topic_relevance = _verdict_obj(tr_verdicts.get(str(idx))) + topic_relevance = tr_verdicts.get(idx) pre_filter = ( PreFilter(topic_relevance=topic_relevance) if topic_relevance else None ) @@ -185,16 +175,12 @@ def _submission_rows(session: Session, assessment: Assessment) -> list[Submissio def build_summary( assessment: Assessment, - execution: AssessmentRun | None, submission_name: str | None = None, + bag: AssessmentExecution | None = None, ) -> AssessmentSummary: """List row for one assessment. Reads only columns and the exec bag, never storage.""" - bag = cast(BatchRunState, (execution.execution or {}) if execution else {}) - config = ( - AssessmentConfigRef(id=execution.config_id, version=execution.config_version) - if execution - else None - ) + if bag is None: + bag = api.load_execution_state(assessment) return AssessmentSummary( assessment_id=assessment.id, method=assessment.method, @@ -202,14 +188,14 @@ def build_summary( experiment_name=assessment.experiment_name, submission_id=assessment.submission_id, submission_name=submission_name, - config=config, - total_items=execution.total_items if execution else 0, - stages=[ - str(step["stage"]) for step in bag.get("pipeline", []) if step.get("stage") - ], - stage=bag.get("stage"), - stage_status=bag.get("stage_status"), - error=execution.error_message if execution else None, + config=AssessmentConfigRef( + id=assessment.config_id, version=assessment.config_version + ), + total_items=assessment.total_items, + stages=[step.stage.value for step in bag.pipeline] if bag else [], + stage=bag.stage.value if bag else None, + stage_status=bag.stage_status.value if bag else None, + error=bag.error if bag else None, inserted_at=assessment.inserted_at, updated_at=assessment.updated_at, ) @@ -218,20 +204,19 @@ def build_summary( def build_detail( *, session: Session, assessment: Assessment, include_input: bool = False ) -> AssessmentDetailResponse: - """Poll payload: run status plus every row produced so far. + """Poll payload: status plus every row produced so far. ``include_input`` echoes the submitted columns; it costs a storage read, so a tight poll loop leaves it off. """ - executions = api.list_executions(session=session, assessment_id=assessment.id) - execution = executions[0] if executions else None + bag = api.load_execution_state(assessment) submission = ( session.get(AssessmentSubmission, assessment.submission_id) if assessment.submission_id else None ) - result = build_result(session=session, assessment=assessment, execution=execution) + result = build_result(session=session, assessment=assessment, bag=bag) rows = _submission_rows(session, assessment) if include_input else [] items = [ @@ -245,7 +230,7 @@ def build_detail( ] summary = build_summary( - assessment, execution, submission.name if submission else None + assessment, submission.name if submission else None, bag=bag ) return AssessmentDetailResponse( **summary.model_dump(), diff --git a/backend/app/services/assessment/api/submission.py b/backend/app/services/assessment/api/submission.py index c8a3a934a..e6672c1f9 100644 --- a/backend/app/services/assessment/api/submission.py +++ b/backend/app/services/assessment/api/submission.py @@ -1,7 +1,7 @@ """BATCH API-client submit entrypoint. -Resolves the ASSESSMENT config, persists the parent assessment + one execution, -seeds the staged-batch runtime bag, and hands the pipeline off to Celery. Only the +Resolves the ASSESSMENT config, persists the assessment with its seeded runtime bag, +and hands the pipeline off to Celery. Only the BATCH method is wired here; RESPONSE stays a route-level 501 stub (deferred). """ @@ -19,11 +19,12 @@ from app.crud.config import ConfigCrud, ConfigVersionCrud from app.models.assessment import ( AssessmentCreate, + AssessmentExecution, AssessmentMethod, AssessmentStatus, AssessmentSubmitResponse, BatchInput, - BatchRunState, + StageStatus, derive_method, ) from app.models.config.assessment_blob import ( @@ -284,79 +285,59 @@ def submit( status_code=503, detail="Failed to store the assessment submission. Please retry.", ) + pipeline = batch_service.build_pipeline(blob.pre_filters) + bag = AssessmentExecution( + pipeline=pipeline, + stage=pipeline[0].stage, + stage_status=StageStatus.PENDING, + gate_passed=[True] * total_items, + callback_url=str(request.callback_url) if request.callback_url else None, + request_metadata=request.request_metadata, + ) assessment = api.create_assessment( session=session, assessment_id=assessment_id, method=AssessmentMethod.BATCH, - input=None, + config_id=request.config.id, + config_version=request.config.version, + total_items=total_items, + execution=bag, submission_id=submission_id, submission_input=submission_url, experiment_name=request.experiment_name, organization_id=organization_id, project_id=project_id, ) - - execution = api.create_execution( - session=session, - assessment_id=assessment.id, - config_id=request.config.id, - config_version=request.config.version, - total_items=total_items, - ) - - pipeline = batch_service.build_pipeline(blob.pre_filters) - bag: BatchRunState = { - "pipeline": pipeline, - "stage": pipeline[0]["stage"], - "stage_status": AssessmentStatus.PENDING.value, - "stage_batches": {}, - "stage_output_urls": {}, - "verdicts": {}, - "counters": {}, - "gate_passed": [True] * total_items, - "provider": provider, - "model": model, - "input_schema": input_columns or None, - "callback_url": str(request.callback_url) if request.callback_url else None, - "request_metadata": request.request_metadata, - } - api.save_execution_state(session=session, execution=execution, state=bag) api.update_status( - session=session, obj=assessment, status=AssessmentStatus.PROCESSING + session=session, assessment=assessment, status=AssessmentStatus.PROCESSING ) trace_id = correlation_id.get() or "" try: dispatched = run_assessment_api_batch.delay( - execution_id=execution.id, + assessment_id=str(assessment.id), organization_id=organization_id, project_id=project_id, trace_id=trace_id, ) except Exception as exc: logger.exception( - "[submit] Failed to enqueue BATCH task | assessment_id=%s | execution_id=%s", - assessment.id, - execution.id, - ) - api.update_status( - session=session, obj=execution, status=AssessmentStatus.FAILED + "[submit] Failed to enqueue BATCH task | assessment_id=%s", assessment.id ) api.update_status( - session=session, obj=assessment, status=AssessmentStatus.FAILED + session=session, assessment=assessment, status=AssessmentStatus.FAILED ) raise HTTPException( status_code=503, detail="Failed to dispatch the assessment for processing. Please retry.", ) from exc logger.info( - "[submit] Dispatched BATCH assessment | assessment_id=%s | execution_id=%s | " - "task_id=%s | provider=%s | stages=%s | rows=%s", + "[submit] Dispatched BATCH assessment | assessment_id=%s | task_id=%s | " + "provider=%s | stages=%s | rows=%s", assessment.id, - execution.id, dispatched.id, provider, - [s["stage"] for s in pipeline], + [step.stage.value for step in pipeline], total_items, ) diff --git a/docs/wiki/domain-map.md b/docs/wiki/domain-map.md index 073d65540..f45ecfcdd 100644 --- a/docs/wiki/domain-map.md +++ b/docs/wiki/domain-map.md @@ -35,7 +35,7 @@ APIKey → Organization, Project, User # programmatic access | EvaluationIterationRun | evaluation_iteration.py | Dataset, Config, Org, Project | EvaluationRun, Job (referenced only inside the LangGraph checkpoint state, not FK columns on this table, logical); callback_url caller (logical) | | STTSample / STTResult | stt_evaluation.py | Dataset, Run, File, Language | human annotation UI (logical) | | TTSResult | tts_evaluation.py | Run, Org, Project | human annotation UI (logical) | -| Assessment / AssessmentRun | assessment.py | Config, AssessmentSubmission, BatchJob, Job, Org, Project | console UI (logical); webhook consumers (logical) | +| Assessment | assessment/assessment.py | Config (`config_id`+`config_version`), AssessmentSubmission, BatchJob (via `execution.stage_batches`), Job, Org, Project | console UI (logical); webhook consumers (logical) | | Document | document.py | Project, Document (parent) | DocumentCollection, DocTransformationJob, FineTuning, ModelEvaluation | | Collection | collection.py | Project | DocumentCollection, CollectionJob; provider vector stores (logical) | | DocumentCollection | document_collection.py | Document, Collection | RAG lookups (logical) | diff --git a/docs/wiki/modules/assessment.md b/docs/wiki/modules/assessment.md index 7abfeddac..12f610ac2 100644 --- a/docs/wiki/modules/assessment.md +++ b/docs/wiki/modules/assessment.md @@ -1,11 +1,11 @@ # Module: Assessment -Assessments of datasets/inputs against ASSESSMENT-tagged configs. Three methods, chosen by input shape: **RESPONSE** (object input → `job` → llm_call/llm_chain), **BATCH** (list input → assessment_run → batch_job), **RUN** (legacy UI pipeline). No deep-dive doc yet. +Assessments of datasets/inputs against ASSESSMENT-tagged configs. Two methods, chosen by input shape: **RESPONSE** (object input → `job` → llm_call/llm_chain; not wired yet) and **BATCH** (list input → staged provider batches). The legacy **RUN** console pipeline is retired: its routers are unregistered and its modules go in the follow-up PR. One assessment = one pinned config version; the former `assessment_run` child is folded into `assessment` and dropped (migration 086). No deep-dive doc yet. All paths relative to `backend/app/`. ## Routes -- `api/routes/assessment/` — legacy UI (assessments, runs, datasets); RUN only, and the results export 422s a BATCH assessment (its rows come from `GET /assessments/{id}`) +- `api/routes/assessment/` — only `datasets.py` is mounted (submission upload/preview/list/delete). `assessments.py` and `runs.py` are RUN-only, unregistered, pending deletion. - `api/routes/assessment/api.py` — API-client routes, mounted top-level at `/assessments`: `POST` (submit; method inferred from input shape, BATCH wired, RESPONSE returns 501), `GET` (list, **BATCH only** so the execution join yields one row per assessment; `config_id` + optional `version` filters), `GET /{assessment_id}` (detail/poll; `include_input=true` echoes the submitted rows at the cost of a storage read, off by default). `callback_url` is optional and adds a webhook push on top of polling. ## Tables (SQLModel) @@ -13,13 +13,12 @@ All paths relative to `backend/app/`. | Table | Model | |---|---| -| `assessment` (Assessment; parent — `method`, data source, `submission_input`, `result_files`; FK → job (RESPONSE), assessment_submission (RUN), org, project) | `models/assessment/assessment.py` | -| `assessment_run` (AssessmentRun; child — one config execution, BATCH/RUN; FK → assessment, config, batch_job) | `models/assessment/assessment.py` | +| `assessment` (Assessment — `method`, `config_id` + `config_version`, `total_items`, `execution` (`AssessmentExecution` JSONB), data source (`submission_id` / `submission_input`), `result_files`; FK → config, job (RESPONSE), assessment_submission, org, project) | `models/assessment/assessment.py` | | `assessment_submission` (AssessmentSubmission; an uploaded CSV/XLSX a run can read its rows from; FK → org, project) | `models/assessment/submission.py` | -Config version (tag=ASSESSMENT, `models/config/assessment_blob.py`) owns system / pre-filters / params / schemas: `input_schema` is a **top-level** field on the blob (sibling of `pre_filters`/`assessment`, **mandatory, non-empty** per-column spec `{type, format}` for the BATCH `data` rows — `type` is required per column; `strict: true` means the column must be present and non-blank in every submission row; it is **off by default** because the console declares every sheet column (the backend rejects undeclared ones) and real sheets have blanks — a blank on a non-strict column passes and `build_rows` fills omitted keys with `""` so the placeholder still resolves; note the config store persists the validated blob *with defaults*, so a saved version pins the default in force at save time; `InputColumn` is `extra=forbid`, so a misspelt flag fails at config save; attachment columns (`image`/`pdf`/`video`) are url-format only, and `video` is Google-only — the OpenAI and Anthropic part builders log and skip a video cell) — it describes the shared input rows once, so both the pre-filter and assessment consumers read the same schema. `submission` (the per-row prompt template with `{column}` placeholders — **mandatory** on `assessment.params`, optional on each pre-filter's `params`; every `{placeholder}` is validated at config save against the top-level `input_schema` keys, and an unknown placeholder rejects the save). `json_output_schema` stays in `assessment.params` (assessment-specific, object-typed structured-output schema, omit for free text). The only API-client pre-filter is `topic_relevance` (duplicate_detection was removed from the API-client pipeline; it survives only in the legacy RUN pipeline); it carries its own `provider` (default `openai`) + `params` (TextLLMParams: model, temperature, ...) and runs its own llm call; its criteria live in `params.instructions` (a **mandatory** field, same shape as the assessment call — pre-filters no longer have a top-level `prompt`/`content`) and it may carry its own `params.submission` template. The prompt template lives on the config, not the request. Strict input types `ResponseInput` (RESPONSE, `{attachments}`) / `BatchInput` (BATCH) no longer carry `query`; they discriminate structurally (a `data` or `submission_doc_id` key ⇒ BATCH, else RESPONSE) with `extra=forbid` keeping them disjoint and rejecting a stray `query` — no `mode` tag (`models/assessment/assessment_api.py`). `BatchInput` takes the rows **either** inline as `data` (a list of submission rows, each a flat column→string map, an attachment column's value being a url string) **or** by reference as `submission_doc_id` (an `assessment_submission` id); a `model_validator` requires exactly one, so both-or-neither is a 422. A `submission_doc_id` is resolved at submit: its rows are read, projected onto the schema's columns (a sheet column `input_schema` does not name is dropped, so the schema is the selection and the console need only declare the `@`-referenced columns), validated against `input_schema` like inline ones, and copied into that run's own `submission.jsonl`, so downstream stays one code path and the run keeps an immutable snapshot if the submission file is later replaced. `assessment.submission_id` records which submission it came from, so the provenance survives the copy and `delete_submission` refuses while a BATCH run still points at it. That column is set by RUN and by a `submission_doc_id` BATCH alike; it stays NULL only when BATCH sent rows inline. Legacy RUN runtime lives in `assessment_run.execution` (`RunExecution`). +Config version (tag=ASSESSMENT, `models/config/assessment_blob.py`) owns system / pre-filters / params / schemas: `input_schema` is a **top-level** field on the blob (sibling of `pre_filters`/`assessment`, **mandatory, non-empty** per-column spec `{type, format}` for the BATCH `data` rows — `type` is required per column; `strict: true` means the column must be present and non-blank in every submission row; it is **off by default** because the console declares every sheet column (the backend rejects undeclared ones) and real sheets have blanks — a blank on a non-strict column passes and `build_rows` fills omitted keys with `""` so the placeholder still resolves; note the config store persists the validated blob *with defaults*, so a saved version pins the default in force at save time; `InputColumn` is `extra=forbid`, so a misspelt flag fails at config save; attachment columns (`image`/`pdf`/`video`) are url-format only, and `video` is Google-only — the OpenAI and Anthropic part builders log and skip a video cell) — it describes the shared input rows once, so both the pre-filter and assessment consumers read the same schema. `submission` (the per-row prompt template with `{column}` placeholders — **mandatory** on `assessment.params`, optional on each pre-filter's `params`; every `{placeholder}` is validated at config save against the top-level `input_schema` keys, and an unknown placeholder rejects the save). `json_output_schema` stays in `assessment.params` (assessment-specific, object-typed structured-output schema, omit for free text). The only API-client pre-filter is `topic_relevance` (duplicate_detection was removed from the API-client pipeline; it survives only in the legacy RUN pipeline); it carries its own `provider` (default `openai`) + `params` (TextLLMParams: model, temperature, ...) and runs its own llm call; its criteria live in `params.instructions` (a **mandatory** field, same shape as the assessment call — pre-filters no longer have a top-level `prompt`/`content`) and it may carry its own `params.submission` template. The prompt template lives on the config, not the request. Strict input types `ResponseInput` (RESPONSE, `{attachments}`) / `BatchInput` (BATCH) no longer carry `query`; they discriminate structurally (a `data` or `submission_doc_id` key ⇒ BATCH, else RESPONSE) with `extra=forbid` keeping them disjoint and rejecting a stray `query` — no `mode` tag (`models/assessment/assessment_api.py`). `BatchInput` takes the rows **either** inline as `data` (a list of submission rows, each a flat column→string map, an attachment column's value being a url string) **or** by reference as `submission_doc_id` (an `assessment_submission` id); a `model_validator` requires exactly one, so both-or-neither is a 422. A `submission_doc_id` is resolved at submit: its rows are read, projected onto the schema's columns (a sheet column `input_schema` does not name is dropped, so the schema is the selection and the console need only declare the `@`-referenced columns), validated against `input_schema` like inline ones, and copied into that run's own `submission.jsonl`, so downstream stays one code path and the run keeps an immutable snapshot if the submission file is later replaced. `assessment.submission_id` records which submission it came from, so the provenance survives the copy and `delete_submission` refuses while a BATCH run still points at it. That column is set by RUN and by a `submission_doc_id` BATCH alike; it stays NULL only when BATCH sent rows inline. -`assessment.id` is a **UUID** (like config/job/llm_call). Per-item result = `AssessmentResult {output: {assessment, pre_filter}, error}` (no `metadata` — the provider/model/usage block was removed from the API-client output) where `output.assessment` = the LLM output parsed to an object when the config has a `json_output_schema`, else string (null for gated/failed rows), and `output.pre_filter` holds the `{topic_relevance}` verdict (`{verdict, reasoning}` or null) and is itself null when no pre-filter ran. The API-client BATCH rows do **not** live in `assessment.input` (that column is now RESPONSE/RUN only, NULL for BATCH). They are uploaded to `submission.jsonl` at submit and `assessment.submission_input` holds its `s3://` url: 3-6MB of JSONB was dragged along by every full-row `SELECT` of the assessment. `services/assessment/api/submission_store.py` owns the round trip: `upload_submission_rows` at submit, and the `open_submission_rows` context manager on read, which streams the JSONL line by line and closes the body on exit. Rows are read only inside `_submit_stage`, never on a poll, and only the rows the stage's subset needs are held (the count comes from `execution.total_items`, so an empty subset never opens the file). A storage read failure raises `SubmissionUnavailableError`, which requeues the task instead of failing the execution; a corrupt line is a `ValueError` and terminal. An upload failure at submit is a 503 (an assessment without its rows is unrunnable). `build_result` reads `execution.total_items` rather than re-deriving the count from the rows, so the terminal path never fetches them. +`assessment.id` is a **UUID** (like config/job/llm_call). Per-item result = `AssessmentResult {output: {assessment, pre_filter}, error}` (no `metadata` — the provider/model/usage block was removed from the API-client output) where `output.assessment` = the LLM output parsed to an object when the config has a `json_output_schema`, else string (null for gated/failed rows), and `output.pre_filter` holds the `{topic_relevance}` verdict (`{verdict, reasoning}` or null) and is itself null when no pre-filter ran. The API-client BATCH rows do **not** live in `assessment.input` (that column is now RESPONSE/RUN only, NULL for BATCH). They are uploaded to `submission.jsonl` at submit and `assessment.submission_input` holds its `s3://` url: 3-6MB of JSONB was dragged along by every full-row `SELECT` of the assessment. `services/assessment/api/submission_store.py` owns the round trip: `upload_submission_rows` at submit, and the `open_submission_rows` context manager on read, which streams the JSONL line by line and closes the body on exit. Rows are read only inside `_submit_stage`, never on a poll, and only the rows the stage's subset needs are held (the count comes from `assessment.total_items`, so an empty subset never opens the file). A storage read failure raises `SubmissionUnavailableError`, which requeues the task instead of failing the execution; a corrupt line is a `ValueError` and terminal. An upload failure at submit is a 503 (an assessment without its rows is unrunnable). `build_result` reads `assessment.total_items` rather than re-deriving the count from the rows, so the terminal path never fetches them. `assessment.result_files` is the durable record of every provider batch dump held: JSONB, **NOT NULL default `{}`**, CHECK-constrained to a JSON object (`ck_assessment_result_files_is_object`), keyed by stage name (`assessment` / `topic_relevance`, the `ApiStage` values) plus `errors`, with `{object_store_url}` per key. It stores the raw `s3://` url (presigning is per-delivery, so the column never holds an expiring link); writes go through `crud/assessment/api.py::set_result_files`, which merges **server-side** (`result_files || :files::jsonb`) rather than read-modify-write, because two drivers can touch the row in the same second. `batch_job.provider_error_file_id` persists OpenAI's `error_file_id` (OpenAI only; Anthropic and Gemini report per-item errors inline) so the error dump is still fetchable at terminal time, in a later task than the poll that surfaced it. @@ -27,7 +26,7 @@ Delivery is poll-or-push: the `POST /assessments` ack is the flat `AssessmentSub ## Services / CRUD - `services/assessment/utils/attachments.py` — cell→provider attachment conversion (Drive URL normalization; OpenAI/Anthropic/Gemini part builders). `rewrite_gcs_attachment_urls` bulk-resolves `gs://` cells to provider-reachable URLs via `services/buckets/` (Path A native passthrough for google-gcp / Path B signed HTTPS otherwise) **before** JSONL build. Called in: `services/assessment/api/batch.py::_submit_provider_batch` (API pipeline, all stages), `crud/assessment/batch.py::submit_assessment_batch` (legacy L2), `services/assessment/tasks.py` (legacy prefilter). Submit-time validation (`services/assessment/api/submission.py`) allows `gs://` alongside `http(s)://`. -- `services/assessment/` — legacy RUN pipeline (service, stages, processing, batch, cron, tasks) +- `services/assessment/{service,stages,tasks}.py` — legacy RUN pipeline, unreachable (routers unregistered), pending deletion - `services/assessment/submission.py` — submission-file upload + preview (`upload_submission`, `preview_submission`, `file_extension_of`). The name is sanitized then checked via `crud/assessment/submission.py::get_submission_by_name` **before** the object-store write (early 409 on a duplicate); the id is pre-generated and the key is `assessment/submissions//.` (`SUBMISSIONS_SUBDIRECTORY`), so two concurrent same-name uploads never share a key and the unique constraint is the only arbiter. No `file_extension` column exists; the format is read off the key's suffix. Upload parses the sheet once and writes the rows beside the file as `submission.jsonl` (`_store_parsed_rows`, best effort — a failure only costs a later reader one re-parse); `services/assessment/api/submission_store.py::open_uploaded_rows` streams that copy and falls back to `crud/assessment/batch.py::load_submission_file_rows` only when the copy cannot be opened, never mid-stream. XLSX goes through `clean_sheet` in all three readers (upload count, preview, run loader): blank rows drop first, then columns with no header or no values, and cells come back stripped — so a 1000×17 sheet with 100×10 filled is 100×10 everywhere. CSV keeps only the blank-row / blank-header filter. - `services/assessment/validators.py` — the module's one leaf of shared helpers, in three sections: upload validation (`MAX_FILE_SIZE`, `ALLOWED_EXTENSIONS`, `ALLOWED_MIME_TYPES`, `validate_dataset_file`), storage-key constants and builders (`SUBMISSION_FILENAME`, `ERRORS_FILENAME`, `assessment_prefix`, `stage_batch_prefix`, `submission_prefix`, `submission_rows_url`, `file_extension_of`), and sheet parsing (`clean_sheet`, `parse_csv_rows`, `parse_excel_rows`, `parse_rows`). Keep new shared helpers here rather than adding another utils module. - **Provider param mapping lives in `services/llm/mappers.py`, not in the assessment tree.** The assessment fork (`services/assessment/mappers.py`) is deleted; `crud/assessment/batch.py`, `services/assessment/api/batch.py` and `prefilter/request_builder.py` all import the shared mappers. The structured-output key is `output_schema` on every provider (the wire name stays `json_output_schema` on `AssessmentTextParams`; `api/batch.py::_stage_params` renames it). `normalize_llm_text` is exported from the same module but no mapper calls it — assessment normalises `instructions` at its own call sites, so prod LLM prompts are untouched. @@ -35,13 +34,13 @@ Delivery is poll-or-push: the `POST /assessments` ack is the flat `AssessmentSub - `services/assessment/api/` — API-client pipeline: `submission.py` (submit), `submission_store.py` (submission rows in object storage), `batch.py` (staged provider batches — gate pre-filters → pass-through → assessment, over `core/batch`; `PREFILTER_VERDICT_SCHEMA`), `results.py` (builds `AssessmentBatchResult`), `result_files.py` (dump bookkeeping + `errors.jsonl` + presigned callback `metadata`), `callbacks.py` (webhook) - `result_files.py` — `record_stage_dump` (called per completed stage from `batch.py`), `build_and_upload_errors` (one self-describing JSONL of `execution_error` / `row_error` / `provider_error_file` records, **not uploaded at all when there are no rows**, so a clean run leaves no empty object and the callback reports `errors: null`; per-row errors come from the exec bag's `stage_errors`; OpenAI's error-file rows are merged into the stage's parsed results at poll time via `batch.error_file_entries`, so they land in `stage_errors` and `build_result` reports them per row instead of as silent "no output, no error" gaps; the same helper feeds the errors.jsonl dump), `finalize_result_files` (called from `_finalize`/`_fail` *before* the callback_url check, so durability never depends on a callback), `presign_result_files` (presigns each key at `expires_in=86400` into `AssessmentResultFiles`; never raises). Storage keys land under one per-assessment prefix, `/assessment//` (`assessment_subdirectory`), holding `errors.jsonl` plus `batch-/results.jsonl` per stage; the API path passes that prefix into `process_completed_batch`'s optional `subdirectory` (legacy RUN and evaluations keep the default `/batch-`). Uploaded submission files sit at `assessment/submissions//.` beside their parsed `submission.jsonl`. Only an **inline** BATCH gets a per-assessment copy at `/submission.jsonl` (`assessment.submission_input`); a run submitted by `submission_doc_id` streams the uploaded submission's own rows, because a submission is immutable (create/read/delete only, no update crud) and the copy would never diverge. - `crud/assessment/submission.py` — `create_submission`, `get_submission_by_name`, `get_submission_by_id`, `list_submissions`, `delete_submission` (refuses while an assessment still references it) -- `crud/assessment/api.py` — new API-client crud (method-based Assessment/AssessmentRun writes): `create_assessment` (takes a pre-generated `assessment_id` + `submission_input`, so submit uploads `submission.jsonl` first and then does one insert — no FAILED orphan when the upload fails), `set_assessment_job`, `create_execution`, `set_execution_batch_job`, `save_execution_state`, `set_result_files`, `update_status`, `list_executions`, `list_assessments_with_execution` (BATCH only; the detail route reads through the legacy `core.get_assessment_by_id`). Namespaced under `api` (`from app.crud.assessment import api`) to avoid colliding with the legacy `create_assessment`. -- `crud/assessment/{core,cron,processing,batch}.py` — legacy RUN pipeline crud. RUN runtime for the dropped columns (`stage`/`stage_status`/`pipeline`/`stage_batches`/`prefilter_total_*`/`object_store_url`) now lives in the `assessment_run.execution` bag via `core._read_exec`/`_write_exec`; `status` is the `AssessmentStatus` enum; the RUN input binding lives on the parent `assessment.input`. `batch.py` builds per-row prompts (`{column}` substitution from the parent `InputBinding.prompt`). +- `crud/assessment/api.py` — API-client crud on the `assessment` row: `create_assessment` (takes the pre-generated `assessment_id`, the pinned config, `total_items` and the seeded `AssessmentExecution`, so submit uploads `submission.jsonl` first and then does one insert — no FAILED orphan when the upload fails), `set_assessment_job`, `save_execution_state` / `load_execution_state` (Pydantic ⇄ JSONB), `set_result_files`, `update_status`, `list_assessments` (BATCH only, optional `config_id`/`config_version` filter; the detail route reads through the legacy `core.get_assessment_by_id`). Namespaced under `api` (`from app.crud.assessment import api`) to avoid colliding with the legacy `create_assessment`. +- `crud/assessment/{core,cron,processing}.py` — legacy RUN pipeline crud, unreachable, pending deletion. `batch.py` keeps the shared JSONL builders + `load_submission_file_rows` the BATCH path uses. ## Async - Rides shared `core/batch/` provider batch infra + cron polling (same as evaluations). -- **Two drivers, strictly separated:** the cron poller (`crud/assessment/cron.py::poll_all_pending_assessment_evaluations`) selects `method == RUN` only and owns the legacy pipeline; the API-client BATCH path is driven solely by the Celery self-re-enqueue. The poller used to pick up BATCH runs too and crash on their list-shaped `pipeline`, leaving the run un-finalised and the callback unfired. Its deterministic-error branch (`DETERMINISTIC_ERRORS`) marks a run FAILED on a programming error instead of logging "will retry" forever; bare `Exception` stays the transient path. -- API-client BATCH: `celery/tasks/job_execution.py::run_assessment_api_batch` drives the staged pipeline (poll → parse verdict/gate → advance/finalize → callback), self-re-enqueuing between stages. The task carries `autoretry_for=(Exception, Timeout)` with backoff (app-default `max_retries`), so a task that raises is re-run instead of silently ending the chain, and `run_batch_stage` takes its execution row `FOR UPDATE SKIP LOCKED`, so an overlapping task no-ops and a duplicate chain collapses. It also sets `reject_on_worker_lost=False` (per task; the global flag stays `True` for every other task), so a task whose worker is lost is acked instead of re-queued in a loop. Known gap: the execution then stays PROCESSING (no callback) until retriggered; marking it FAILED is a follow-up. Staged state + `callback_url`/`request_metadata` live in the `assessment_run.execution` bag (`BatchRunState`), which also holds `stage_errors` (`stage -> {row_index -> error}`, captured at parse time; the raw dumps go to object storage, not the bag). +- The RUN cron poller (`crud/assessment/cron.py`) still runs from `api/routes/cron.py` but selects `method == RUN`, which no longer exists in data; it goes with the RUN code. +- API-client BATCH: `celery/tasks/job_execution.py::run_assessment_api_batch` drives the staged pipeline (poll → parse verdict/gate → advance/finalize → callback), self-re-enqueuing between stages. The task carries `autoretry_for=(Exception, Timeout)` with backoff (app-default `max_retries`), so a task that raises is re-run instead of silently ending the chain, and `run_batch_stage` takes the `assessment` row `FOR UPDATE SKIP LOCKED`, so an overlapping task no-ops and a duplicate chain collapses. It also sets `reject_on_worker_lost=False` (per task; the global flag stays `True` for every other task), so a task whose worker is lost is acked instead of re-queued in a loop. Known gap: the execution then stays PROCESSING (no callback) until retriggered; marking it FAILED is a follow-up. Staged state + `callback_url`/`request_metadata` live in `assessment.execution` as `AssessmentExecution` (Pydantic, `extra="forbid"`, int row keys): `pipeline`, `stage`, `stage_status`, `stage_batches`, `stage_errors`, `verdicts`, `counters`, `gate_passed`, `error`. Provider/model come from the pinned config version (`stage_provider_model`), stage dump urls from `result_files` — neither is duplicated in the bag. The task takes `assessment_id`; drain the Celery queue before deploying 086, a queued pre-fold `execution_id` task has no row to resolve. - `core/batch/polling.py::poll_batch_status` diffs **per field**, not only on a `provider_status` flip, so `error_file_id` and `provider_output_file_id` persist even when the status is unchanged. Shared with the evaluations path; return value unchanged. ## External From d3d34cfa9dd58aab3b3d74b58b0d17e5fc6e9658 Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Wed, 7 Oct 2026 00:55:02 +0530 Subject: [PATCH 3/4] feat(assessment): enhance execution model by adding verdicts and gate_passed fields, and update related logic for stage processing --- ...py_assessment_run_columns_to_assessment.py | 11 +- backend/app/api/routes/assessment/__init__.py | 1 - .../app/models/assessment/assessment_api.py | 8 +- backend/app/services/assessment/api/batch.py | 103 +++++++++++++----- .../app/services/assessment/api/results.py | 81 ++++++++------ .../app/services/assessment/api/submission.py | 1 - docs/wiki/modules/assessment.md | 2 +- 7 files changed, 135 insertions(+), 72 deletions(-) diff --git a/backend/app/alembic/versions/086_copy_assessment_run_columns_to_assessment.py b/backend/app/alembic/versions/086_copy_assessment_run_columns_to_assessment.py index 8b16453ab..287a0ab6d 100644 --- a/backend/app/alembic/versions/086_copy_assessment_run_columns_to_assessment.py +++ b/backend/app/alembic/versions/086_copy_assessment_run_columns_to_assessment.py @@ -24,8 +24,15 @@ CONFIG_INDEX = "idx_assessment_config" CONFIG_FK = "fk_assessment_config_id_config" RUN_TABLE = "assessment_run" -# Derivable from the pinned config version, or duplicated in result_files. -DROPPED_EXECUTION_KEYS = ("stage_output_urls", "provider", "model", "input_schema") +# Derivable from config_version or already in the stage dumps. +DROPPED_EXECUTION_KEYS = ( + "stage_output_urls", + "provider", + "model", + "input_schema", + "verdicts", + "gate_passed", +) def upgrade() -> None: diff --git a/backend/app/api/routes/assessment/__init__.py b/backend/app/api/routes/assessment/__init__.py index 53651f90b..89fb233dd 100644 --- a/backend/app/api/routes/assessment/__init__.py +++ b/backend/app/api/routes/assessment/__init__.py @@ -10,7 +10,6 @@ dependencies=[Depends(require_feature(FeatureFlag.ASSESSMENT))], ) -# RUN routers (assessments, runs) are retired; their modules go in the follow-up PR. router.include_router(datasets.router) __all__ = ["router"] diff --git a/backend/app/models/assessment/assessment_api.py b/backend/app/models/assessment/assessment_api.py index d680e20ce..0a440ff46 100644 --- a/backend/app/models/assessment/assessment_api.py +++ b/backend/app/models/assessment/assessment_api.py @@ -99,7 +99,8 @@ class AssessmentExecution(BaseModel): """Runtime state of the staged BATCH pipeline, stored on ``assessment.execution``. Advanced one Celery task at a time by ``run_batch_stage``; ``stage_status`` keys the - idempotent redelivery. Row-indexed maps use int keys (JSON strings coerce on load). + idempotent redelivery. Deliberately O(stages), never O(rows): per-row verdicts and + outputs live in each stage's dump (``assessment.result_files``), not here. """ model_config = ConfigDict(extra="forbid") @@ -109,9 +110,7 @@ class AssessmentExecution(BaseModel): stage_status: StageStatus stage_batches: dict[ApiStage, int] = Field(default_factory=dict) stage_errors: dict[ApiStage, dict[int, str]] = Field(default_factory=dict) - verdicts: dict[ApiStage, dict[int, PreFilterVerdict]] = Field(default_factory=dict) counters: dict[ApiStage, StageCounters] = Field(default_factory=dict) - gate_passed: list[bool] callback_url: str | None = None request_metadata: dict[str, JsonValue] | None = None error: str | None = None @@ -127,6 +126,9 @@ def next_stage(self, current: ApiStage) -> ApiStage | None: idx = stages.index(current) return stages[idx + 1] if idx + 1 < len(stages) else None + def gate_stages(self) -> list[ApiStage]: + return [step.stage for step in self.pipeline if step.kind == StageKind.GATE] + class ParsedResult(TypedDict): """One provider batch response row, normalised across providers. diff --git a/backend/app/services/assessment/api/batch.py b/backend/app/services/assessment/api/batch.py index 363ecad72..d1c610118 100644 --- a/backend/app/services/assessment/api/batch.py +++ b/backend/app/services/assessment/api/batch.py @@ -5,12 +5,14 @@ state lives entirely in ``assessment.execution`` (``AssessmentExecution``). Conditional forwarding: - - a GATE stage runs on every row; rows whose verdict fails are marked - ``gate_passed=False`` but still flow through the remaining pass-through stages - so their metadata is complete. - - a PASS-THROUGH stage runs on every row and never changes ``gate_passed``. + - a GATE stage runs on every row; a row whose verdict fails is excluded from the + assessment stage but still flows through the remaining pass-through stages. + - a PASS-THROUGH stage runs on every row and never excludes one. - the assessment stage batches ONLY gate-passed rows; gate-failed rows get - ``response=null`` in the result, carrying their pre-filter verdicts. + ``assessment=null`` in the result, carrying their pre-filter verdicts. + +Per-row state is never held in the bag: a stage's verdicts are re-read from its dump +(``assessment.result_files[stage]``) when the next stage needs them. """ import json @@ -35,6 +37,7 @@ start_batch_job, ) from app.core.batch.base import BatchProvider +from app.core.cloud import get_cloud_storage from app.core.batch.client import GeminiClient from app.core.config import settings from app.core.db import engine @@ -74,6 +77,7 @@ map_kaapi_to_openai_params, ) from app.services.assessment.utils.attachments import rewrite_gcs_attachment_urls +from app.services.assessment.utils.parsing import parse_stored_results from app.services.assessment.validators import normalize_llm_text, stage_batch_prefix from app.services.llm.providers.registry import LLMProvider from app.utils import ( @@ -84,6 +88,10 @@ logger = logging.getLogger(__name__) +class StageDumpUnavailableError(Exception): + """A completed stage's dump could not be read; the caller should retry the task.""" + + def _google_gcp_credential( *, session: Session, organization_id: int, project_id: int ) -> dict[str, Any]: @@ -573,30 +581,67 @@ def _record_stage( kind: StageKind, parsed: dict[int, ParsedResult], ) -> None: - """Fold a completed stage's parsed results into the bag per its kind.""" + """Fold a completed stage's tallies into the bag. Verdicts stay in the stage dump.""" if kind == StageKind.ASSESSMENT: - return # assessment outputs are read from object store at result time - - verdicts: dict[int, PreFilterVerdict] = {} - passed_count = 0 - for idx, out in parsed.items(): - verdict = _parse_verdict(out.get("output")) - verdicts[idx] = verdict - if verdict.verdict: - passed_count += 1 - elif kind == StageKind.GATE: - bag.gate_passed[idx] = False - - bag.verdicts[stage] = verdicts + return + passed = sum( + 1 for out in parsed.values() if _parse_verdict(out.get("output")).verdict + ) bag.counters[stage] = StageCounters( - total=len(parsed), passed=passed_count, rejected=len(parsed) - passed_count + total=len(parsed), passed=passed, rejected=len(parsed) - passed ) -def _row_subset(bag: AssessmentExecution, kind: StageKind, total: int) -> list[int]: - if kind == StageKind.ASSESSMENT: - return [i for i in range(total) if bag.gate_passed[i]] - return list(range(total)) +def load_stage_outputs( + session: Session, + assessment: Assessment, + blob: AssessmentConfigBlob, + stage: ApiStage, +) -> dict[int, ParsedResult] | None: + """Stream + parse one stage's dump. ``None`` when the stage left no dump.""" + record = assessment.result_files.get(stage.value) or {} + url = record.get("object_store_url") + if not url: + return None + provider, _ = stage_provider_model(blob, stage) + storage = get_cloud_storage(session=session, project_id=assessment.project_id) + raw = parse_stored_results(storage.stream(url).read().decode("utf-8")) + return parse_batch_results(raw, provider) + + +def gate_rejected_rows( + session: Session, + assessment: Assessment, + blob: AssessmentConfigBlob, + bag: AssessmentExecution, +) -> set[int]: + """Row indices some GATE stage rejected, read from those stages' dumps.""" + rejected: set[int] = set() + for stage in bag.gate_stages(): + try: + outputs = load_stage_outputs(session, assessment, blob, stage) + except Exception as exc: + raise StageDumpUnavailableError( + f"Gate stage {stage} dump could not be read: {exc}" + ) from exc + for idx, out in (outputs or {}).items(): + if not _parse_verdict(out.get("output")).verdict: + rejected.add(idx) + return rejected + + +def _row_subset( + session: Session, + assessment: Assessment, + blob: AssessmentConfigBlob, + bag: AssessmentExecution, + kind: StageKind, +) -> list[int]: + total = assessment.total_items + if kind != StageKind.ASSESSMENT: + return list(range(total)) + rejected = gate_rejected_rows(session, assessment, blob, bag) + return [i for i in range(total) if i not in rejected] def _submit_stage( @@ -641,7 +686,7 @@ def _submit_stage( for name, col in blob.input_schema.items() } text_columns, attachments = column_kinds(list(input_columns), input_columns) - subset = _row_subset(bag, kind, assessment.total_items) + subset = _row_subset(session, assessment, blob, bag, kind) if not subset: # No rows left for this stage (everything gated out upstream). Persist the @@ -872,10 +917,10 @@ def _advance_or_finalize( organization_id=organization_id, project_id=project_id, ) - except SubmissionUnavailableError as exc: + except (SubmissionUnavailableError, StageDumpUnavailableError) as exc: # Storage blip, not a bad run: the stage is still PENDING, so retry the task. logger.warning( - "[_advance_or_finalize] Submission unreadable, will retry | " + "[_advance_or_finalize] Storage unreadable, will retry | " "assessment_id=%s | stage=%s | %s", assessment.id, nxt, @@ -964,10 +1009,10 @@ def run_batch_stage( organization_id=organization_id, project_id=project_id, ) - except SubmissionUnavailableError as exc: + except (SubmissionUnavailableError, StageDumpUnavailableError) as exc: # Storage blip, not a bad run: the stage stays PENDING, retry the task. logger.warning( - "[run_batch_stage] Submission unreadable, will retry | " + "[run_batch_stage] Storage unreadable, will retry | " "assessment_id=%s | stage=%s | %s", assessment_id, stage, diff --git a/backend/app/services/assessment/api/results.py b/backend/app/services/assessment/api/results.py index 134be7a2f..06b149a93 100644 --- a/backend/app/services/assessment/api/results.py +++ b/backend/app/services/assessment/api/results.py @@ -2,8 +2,8 @@ One result unit per input row. Gate-failed rows carry ``assessment=null`` plus their pre-filter verdicts; gate-passed rows carry the assessment call's parsed -output. Pre-filter verdicts live in the execution bag; assessment outputs are -streamed back from object storage. +output. Both verdicts and outputs are streamed back from the stage dumps in object +storage; the execution bag holds nothing per row. """ import json @@ -12,8 +12,8 @@ from sqlmodel import Session -from app.core.cloud import get_cloud_storage from app.crud.assessment import api +from app.models.config.assessment_blob import AssessmentConfigBlob from app.models.assessment import ( ApiStage, Assessment, @@ -32,11 +32,10 @@ Submission, ) from app.services.assessment.api.batch import ( - parse_batch_results, + _parse_verdict, + load_stage_outputs, resolve_blob, - stage_provider_model, ) -from app.services.assessment.utils.parsing import parse_stored_results logger = logging.getLogger(__name__) @@ -56,32 +55,28 @@ def _parse_assessment(out: ParsedResult) -> dict[str, Any] | str | None: return parsed if isinstance(parsed, dict) else text -def _load_assessment_outputs( - session: Session, assessment: Assessment +def _load_stage( + session: Session, + assessment: Assessment, + blob: AssessmentConfigBlob, + stage: ApiStage, ) -> tuple[dict[int, ParsedResult], str | None]: - """Stream + parse the assessment stage dump. Returns ``(outputs, load_error)``. + """One stage's dump as ``(outputs, load_error)``. - ``load_error`` is set only on a failed read of a present URL, so the caller - can flag it per-row instead of mistaking it for a clean run. A missing URL - (all rows gated) is a legit empty, not an error. + ``load_error`` is set only on a failed read of a present dump, so the caller can + flag it per-row instead of mistaking it for a clean run. No dump (stage not in the + pipeline, or every row gated out) is a legit empty, not an error. """ - record = assessment.result_files.get(ApiStage.ASSESSMENT.value) or {} - url = record.get("object_store_url") - if not url: - return {}, None try: - blob = resolve_blob(session, assessment) - provider, _ = stage_provider_model(blob, ApiStage.ASSESSMENT) - storage = get_cloud_storage(session=session, project_id=assessment.project_id) - raw = parse_stored_results(storage.stream(url).read().decode("utf-8")) - return parse_batch_results(raw, provider), None + return load_stage_outputs(session, assessment, blob, stage) or {}, None except Exception as exc: logger.warning( - "[_load_assessment_outputs] Could not read assessment output | url=%s | %s", - url, + "[_load_stage] Could not read stage output | assessment_id=%s | stage=%s | %s", + assessment.id, + stage, exc, ) - return {}, "Assessment output could not be read from storage." + return {}, f"{stage.value} output could not be read from storage." def build_result( @@ -90,7 +85,7 @@ def build_result( assessment: Assessment, bag: AssessmentExecution | None = None, ) -> AssessmentBatchResult: - """Build the per-row result from stored verdicts (bag) + assessment output (store). + """Build the per-row result from the stage dumps: gate verdicts + assessment output. Status lives on the response envelope, so this body carries only rows and tallies. Pass ``bag`` when the caller already holds it, so a poll does not re-parse the column. @@ -98,34 +93,50 @@ def build_result( if bag is None: bag = api.load_execution_state(assessment) - # Off the row, not re-derived: the terminal path must not fetch rows to count them. total_items = assessment.total_items - gate_passed = bag.gate_passed if bag else [True] * total_items - outputs, load_error = _load_assessment_outputs(session, assessment) + blob = resolve_blob(session, assessment) + + tr_outputs: dict[int, ParsedResult] = {} + tr_error: str | None = None + rejected: set[int] = set() + if bag and ApiStage.TOPIC_RELEVANCE in {step.stage for step in bag.pipeline}: + tr_outputs, tr_error = _load_stage( + session, assessment, blob, ApiStage.TOPIC_RELEVANCE + ) + if ApiStage.TOPIC_RELEVANCE in bag.gate_stages(): + rejected = { + idx + for idx, out in tr_outputs.items() + if not _parse_verdict(out.get("output")).verdict + } + + outputs, load_error = _load_stage(session, assessment, blob, ApiStage.ASSESSMENT) # Rows the provider rejected are absent from the output dump; the bag kept their error. row_errors = bag.stage_errors.get(ApiStage.ASSESSMENT, {}) if bag else {} - tr_verdicts = bag.verdicts.get(ApiStage.TOPIC_RELEVANCE, {}) if bag else {} items: list[AssessmentResult] = [] counts = AssessmentCounts() for idx in range(total_items): + gate_passed = idx not in rejected # dict shape is the config's own json_output_schema (runtime-defined, no fixed # model); str for free-text output, None when gated/failed. assessment_output: dict[str, Any] | str | None = None error: str | None = None - if gate_passed[idx]: + if gate_passed: if idx in outputs: out = outputs[idx] assessment_output = _parse_assessment(out) error = out.get("error") elif idx in row_errors: error = row_errors[idx] - elif load_error: - error = load_error + elif load_error or tr_error: + error = load_error or tr_error - topic_relevance = tr_verdicts.get(idx) + tr_out = tr_outputs.get(idx) pre_filter = ( - PreFilter(topic_relevance=topic_relevance) if topic_relevance else None + PreFilter(topic_relevance=_parse_verdict(tr_out.get("output"))) + if tr_out + else None ) items.append( AssessmentResult( @@ -139,7 +150,7 @@ def build_result( if assessment_output is not None: counts.assessed += 1 - if not gate_passed[idx]: + if not gate_passed: counts.filtered += 1 if error: counts.errors += 1 diff --git a/backend/app/services/assessment/api/submission.py b/backend/app/services/assessment/api/submission.py index e6672c1f9..8a91474bb 100644 --- a/backend/app/services/assessment/api/submission.py +++ b/backend/app/services/assessment/api/submission.py @@ -290,7 +290,6 @@ def submit( pipeline=pipeline, stage=pipeline[0].stage, stage_status=StageStatus.PENDING, - gate_passed=[True] * total_items, callback_url=str(request.callback_url) if request.callback_url else None, request_metadata=request.request_metadata, ) diff --git a/docs/wiki/modules/assessment.md b/docs/wiki/modules/assessment.md index 12f610ac2..64974ecb8 100644 --- a/docs/wiki/modules/assessment.md +++ b/docs/wiki/modules/assessment.md @@ -40,7 +40,7 @@ Delivery is poll-or-push: the `POST /assessments` ack is the flat `AssessmentSub ## Async - Rides shared `core/batch/` provider batch infra + cron polling (same as evaluations). - The RUN cron poller (`crud/assessment/cron.py`) still runs from `api/routes/cron.py` but selects `method == RUN`, which no longer exists in data; it goes with the RUN code. -- API-client BATCH: `celery/tasks/job_execution.py::run_assessment_api_batch` drives the staged pipeline (poll → parse verdict/gate → advance/finalize → callback), self-re-enqueuing between stages. The task carries `autoretry_for=(Exception, Timeout)` with backoff (app-default `max_retries`), so a task that raises is re-run instead of silently ending the chain, and `run_batch_stage` takes the `assessment` row `FOR UPDATE SKIP LOCKED`, so an overlapping task no-ops and a duplicate chain collapses. It also sets `reject_on_worker_lost=False` (per task; the global flag stays `True` for every other task), so a task whose worker is lost is acked instead of re-queued in a loop. Known gap: the execution then stays PROCESSING (no callback) until retriggered; marking it FAILED is a follow-up. Staged state + `callback_url`/`request_metadata` live in `assessment.execution` as `AssessmentExecution` (Pydantic, `extra="forbid"`, int row keys): `pipeline`, `stage`, `stage_status`, `stage_batches`, `stage_errors`, `verdicts`, `counters`, `gate_passed`, `error`. Provider/model come from the pinned config version (`stage_provider_model`), stage dump urls from `result_files` — neither is duplicated in the bag. The task takes `assessment_id`; drain the Celery queue before deploying 086, a queued pre-fold `execution_id` task has no row to resolve. +- API-client BATCH: `celery/tasks/job_execution.py::run_assessment_api_batch` drives the staged pipeline (poll → parse verdict/gate → advance/finalize → callback), self-re-enqueuing between stages. The task carries `autoretry_for=(Exception, Timeout)` with backoff (app-default `max_retries`), so a task that raises is re-run instead of silently ending the chain, and `run_batch_stage` takes the `assessment` row `FOR UPDATE SKIP LOCKED`, so an overlapping task no-ops and a duplicate chain collapses. It also sets `reject_on_worker_lost=False` (per task; the global flag stays `True` for every other task), so a task whose worker is lost is acked instead of re-queued in a loop. Known gap: the execution then stays PROCESSING (no callback) until retriggered; marking it FAILED is a follow-up. Staged state + `callback_url`/`request_metadata` live in `assessment.execution` as `AssessmentExecution` (Pydantic, `extra="forbid"`): `pipeline`, `stage`, `stage_status`, `stage_batches`, `stage_errors` (sparse), `counters`, `error`. The bag is O(stages), never O(rows): gate verdicts are re-read from the pre-filter stage dump (`result_files[stage]`, via `batch.load_stage_outputs` / `gate_rejected_rows`) when the assessment stage is submitted and when the result is built; an unreadable dump at submit raises `StageDumpUnavailableError`, which requeues the task like a submission-rows blip. Provider/model come from the pinned config version (`stage_provider_model`). The task takes `assessment_id`; drain the Celery queue before deploying 086, a queued pre-fold `execution_id` task has no row to resolve. - `core/batch/polling.py::poll_batch_status` diffs **per field**, not only on a `provider_status` flip, so `error_file_id` and `provider_output_file_id` persist even when the status is unchanged. Shared with the evaluations path; return value unchanged. ## External From 4aa75a4696e8788fa9669cd24365cca3986893ed Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Wed, 7 Oct 2026 00:59:01 +0530 Subject: [PATCH 4/4] feat(assessment): fold assessment_run into assessment and update schema --- ...o_assessment.py => 086_fold_assessment_run_into_assessment.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename backend/app/alembic/versions/{086_copy_assessment_run_columns_to_assessment.py => 086_fold_assessment_run_into_assessment.py} (100%) diff --git a/backend/app/alembic/versions/086_copy_assessment_run_columns_to_assessment.py b/backend/app/alembic/versions/086_fold_assessment_run_into_assessment.py similarity index 100% rename from backend/app/alembic/versions/086_copy_assessment_run_columns_to_assessment.py rename to backend/app/alembic/versions/086_fold_assessment_run_into_assessment.py