feat(assessment): Add audio input column (#1317)

Author: vprashrex

backend/app/crud/assessment/batch.py

@@ -25,15 +25,12 @@
     AssessmentSubmission,
 )
 from app.models.batch_job import BatchJob, BatchJobType
-from app.models.config.assessment_blob import VIDEO_COLUMN_TYPE
 from app.models.llm.constants import DEFAULT_ASSESSMENT_BATCH_MAX_TOKENS
 from app.models.llm.request import ConfigBlob
 from app.services.assessment.utils.attachments import (
-    attachment_type_for_row,
     build_anthropic_attachment_parts,
     build_gemini_attachment_parts,
     resolve_attachment_values,
-    resolve_item_type,
     rewrite_gcs_attachment_urls,
 )
 from app.services.assessment.validators import (
@@ -130,13 +127,7 @@ def build_openai_jsonl(
         # Attachments
         for att in attachments:
             cell_value = row.get(att.column, "")
-            input_parts.extend(
-                resolve_attachment_values(
-                    cell_value,
-                    att,
-                    type_override=attachment_type_for_row(att, row),
-                )
-            )
+            input_parts.extend(resolve_attachment_values(cell_value, att))
 
         if not input_parts:
             logger.warning("[build_openai_jsonl] Skipping empty row | idx=%s", idx)
@@ -195,17 +186,11 @@ def build_google_jsonl(
         has_video = False
         for att in attachments:
             cell_value = row.get(att.column, "")
-            type_override = attachment_type_for_row(att, row)
-            if cell_value.strip() and (
-                resolve_item_type(att.type, type_override) == VIDEO_COLUMN_TYPE
-            ):
+            if cell_value.strip() and att.type == "video":
                 has_video = True
             parts.extend(
                 build_gemini_attachment_parts(
-                    cell_value,
-                    att,
-                    type_override=type_override,
-                    video_part_config=video_part_config,
+                    cell_value, att, video_part_config=video_part_config
                 )
             )
 
@@ -282,13 +267,7 @@ def build_anthropic_jsonl(
         # Attachments (Anthropic uses url-source image/document blocks)
         for att in attachments:
             cell_value = row.get(att.column, "")
-            content.extend(
-                build_anthropic_attachment_parts(
-                    cell_value,
-                    att,
-                    type_override=attachment_type_for_row(att, row),
-                )
-            )
+            content.extend(build_anthropic_attachment_parts(cell_value, att))
 
         if not content:
             logger.warning("[build_anthropic_jsonl] Skipping empty row | idx=%s", idx)

backend/app/models/assessment/assessment.py

@@ -9,7 +9,7 @@
 from typing import TYPE_CHECKING, Any, Literal, Optional
 from uuid import UUID, uuid4
 
-from pydantic import BaseModel, Field, model_validator
+from pydantic import BaseModel, Field
 from sqlalchemy import Column, Index, Text, text
 from sqlalchemy.dialects import postgresql
 from sqlalchemy.dialects.postgresql import JSONB
@@ -84,30 +84,10 @@ class AssessmentAttachment(BaseModel):
     """External-dataset attachment column config (RUN / BATCH-by-ref)."""
 
     column: str = Field(..., description="Dataset column holding the attachment")
-    type: AttachmentColumnType | Literal["mixed"] = Field(
-        ...,
-        description=(
-            "'image'/'pdf'/'video' fix the type; 'mixed' resolves per-row via type_column"
-        ),
+    type: AttachmentColumnType = Field(
+        ..., description="'image', 'pdf', 'video' or 'audio'"
     )
     format: Literal["url", "base64"] = Field(..., description="Data format")
-    type_column: str | None = Field(
-        None, description="'mixed' only: column whose value decides each row's type"
-    )
-    type_value_map: dict[str, AttachmentColumnType] | None = Field(
-        None,
-        description="'mixed' only: maps a type_column value to 'image', 'pdf' or 'video'",
-    )
-
-    @model_validator(mode="after")
-    def _validate_mixed(self) -> "AssessmentAttachment":
-        if self.type == "mixed" and (
-            self.type_column is None or not self.type_value_map
-        ):
-            raise ValueError(
-                "type='mixed' requires 'type_column' and a non-empty 'type_value_map'."
-            )
-        return self
 
 
 # NOTE: Only used for UI. The API-client Assessment pipeline does not use this.

backend/app/models/config/assessment_blob.py

@@ -1,6 +1,6 @@
 import logging
 import re
-from typing import Literal
+from typing import Literal, get_args
 
 from pydantic import JsonValue, model_validator
 from sqlmodel import Field, SQLModel
@@ -18,14 +18,15 @@
 # object-typed dict. Provider strict-mode normalisation is a run-mode concern.
 JSON_SCHEMA_OBJECT_TYPE = "object"
 
-IMAGE_COLUMN_TYPE = "image"
-PDF_COLUMN_TYPE = "pdf"
-VIDEO_COLUMN_TYPE = "video"
+AttachmentColumnType = Literal["image", "pdf", "video", "audio"]
+ColumnType = Literal["text"] | AttachmentColumnType
 # Input-column types carrying a media reference rather than prompt text.
-ATTACHMENT_COLUMN_TYPES = (IMAGE_COLUMN_TYPE, PDF_COLUMN_TYPE, VIDEO_COLUMN_TYPE)
+ATTACHMENT_COLUMN_TYPES = get_args(AttachmentColumnType)
 
-AttachmentColumnType = Literal["image", "pdf", "video"]
-ColumnType = Literal["text"] | AttachmentColumnType
+# Media column types a provider cannot take, so a config pairing them is rejected at save.
+UNSUPPORTED_COLUMN_TYPES_BY_PROVIDER: dict[str, tuple[str, ...]] = dict.fromkeys(
+    (Provider.OPENAI, Provider.ANTHROPIC), ("video", "audio")
+)
 
 # {column} placeholders in a submission template; the capture group is the column name.
 PLACEHOLDER_RE = re.compile(r"\{(\w+)\}")
@@ -51,8 +52,8 @@ class InputColumn(SQLModel):
 
     @model_validator(mode="after")
     def _validate_video_format(self):
-        if self.type == VIDEO_COLUMN_TYPE and self.format == "base64":
-            raise ValueError("A 'video' input column must be url-format.")
+        if self.type in ("video", "audio") and self.format == "base64":
+            raise ValueError(f"A '{self.type}' input column must be url-format.")
         return self
 
 
@@ -225,3 +226,26 @@ def _validate_submission_placeholders(self):
                     f"{', '.join(unknown)}"
                 )
         return self
+
+    @model_validator(mode="after")
+    def _validate_column_types_supported(self):
+        """Every stage's provider must accept each media column type in input_schema."""
+        stage_providers = {"assessment": self.assessment.provider}
+        if self.pre_filters and self.pre_filters.topic_relevance:
+            stage_providers[
+                "pre_filter topic_relevance"
+            ] = self.pre_filters.topic_relevance.provider
+
+        for stage, provider in stage_providers.items():
+            unsupported = UNSUPPORTED_COLUMN_TYPES_BY_PROVIDER.get(provider, ())
+            offending = sorted(
+                column
+                for column, spec in self.input_schema.items()
+                if spec.type in unsupported
+            )
+            if offending:
+                raise ValueError(
+                    f"{stage} provider '{provider}' does not support "
+                    f"{'/'.join(unsupported)} input columns: {', '.join(offending)}"
+                )
+        return self

backend/app/services/assessment/api/batch.py

@@ -207,7 +207,7 @@ def column_kinds(
 ) -> tuple[list[str], list[AssessmentAttachment]]:
     """Split columns into text names and attachment specs per the config's ``input_schema``.
 
-    A column typed image/pdf/video is an attachment (url-format only); anything else is text.
+    A column typed image/pdf/video/audio is an attachment (url-format only); anything else is text.
     """
     text_columns: list[str] = []
     attachments: list[AssessmentAttachment] = []

backend/app/services/assessment/api/submission.py

@@ -6,7 +6,6 @@
 """
 
 import logging
-from urllib.parse import urlparse
 from uuid import UUID, uuid4
 
 from asgi_correlation_id import correlation_id
@@ -36,40 +35,16 @@
     open_uploaded_rows,
     upload_submission_rows,
 )
+from app.services.assessment.utils.attachments import (
+    attachment_type_from_url,
+    split_attachment_urls,
+)
 from app.utils import validate_callback_url
 
 logger = logging.getLogger(__name__)
 
 # Attachment cell values are provided as URLs (base64 is unsupported for batch).
 _URL_PREFIXES = ("http://", "https://", "gs://")
-_EXTENSION_TYPES = {
-    ".pdf": "pdf",
-    ".png": "image",
-    ".jpg": "image",
-    ".jpeg": "image",
-    ".gif": "image",
-    ".webp": "image",
-    ".mp4": "video",
-    ".mov": "video",
-    ".mpeg": "video",
-    ".mpg": "video",
-    ".avi": "video",
-    ".wmv": "video",
-    ".flv": "video",
-    ".webm": "video",
-    ".3gp": "video",
-}
-
-
-def _url_attachment_type(url: str) -> str | None:
-    """The kind a URL's own extension implies, or None when it names no extension.
-
-    Presigned URLs carry a query string, so only the path is read.
-    """
-    path = urlparse(url).path.lower()
-    return next(
-        (kind for ext, kind in _EXTENSION_TYPES.items() if path.endswith(ext)), None
-    )
 
 
 def _validate_rows_against_schema(
@@ -104,8 +79,11 @@ def _validate_rows_against_schema(
         for column, spec in input_schema.items():
             column_type = (spec or {}).get("type")
             value = (row.get(column) or "").strip()
-            if value and column_type in ATTACHMENT_COLUMN_TYPES:
-                if not value.startswith(_URL_PREFIXES):
+            if not value or column_type not in ATTACHMENT_COLUMN_TYPES:
+                continue
+            # A cell may hold several URLs; the builder sends each one separately.
+            for url in split_attachment_urls(value):
+                if not url.startswith(_URL_PREFIXES):
                     raise HTTPException(
                         status_code=422,
                         detail=(
@@ -114,7 +92,7 @@ def _validate_rows_against_schema(
                         ),
                     )
                 # The provider only reports a mistyped column later, as an opaque format error.
-                actual = _url_attachment_type(value)
+                actual = attachment_type_from_url(url)
                 if actual is not None and actual != column_type:
                     raise HTTPException(
                         status_code=422,

backend/app/services/assessment/prefilter/topic_relevance.py

@@ -12,7 +12,6 @@
 from app.services.assessment.prefilter import constants
 from app.services.assessment.prefilter.request_builder import build_request_line
 from app.services.assessment.utils.attachments import (
-    attachment_type_for_row,
     build_gemini_attachment_parts,
     resolve_attachment_values,
 )
@@ -65,11 +64,10 @@ def build_topic_relevance_requests(
             cell = row.get(att.column, "")
             if not cell.strip():
                 continue
-            override = attachment_type_for_row(att, row)
             attachment_parts.extend(
-                resolve_attachment_values(cell, att, type_override=override)
+                resolve_attachment_values(cell, att)
                 if is_openai
-                else build_gemini_attachment_parts(cell, att, type_override=override)
+                else build_gemini_attachment_parts(cell, att)
             )
         lines.append(
             build_request_line(

backend/app/services/assessment/service.py

@@ -30,6 +30,7 @@
     InputBinding,
     StageStatus,
 )
+from app.models.config.assessment_blob import UNSUPPORTED_COLUMN_TYPES_BY_PROVIDER
 from app.models.config.config import ConfigTag
 from app.services.llm.providers.registry import LLMProvider
 
@@ -173,6 +174,23 @@ def start_assessment(
                     f"Supported providers: {sorted(_SUPPORTED_BATCH_PROVIDERS)}"
                 ),
             )
+        unsupported_types = UNSUPPORTED_COLUMN_TYPES_BY_PROVIDER.get(
+            str(provider).removesuffix("-native"), ()
+        )
+        unsupported_columns = sorted(
+            att.column
+            for att in request.input_binding.attachments
+            if att.type in unsupported_types
+        )
+        if unsupported_columns:
+            raise HTTPException(
+                status_code=422,
+                detail=(
+                    f"Config {cfg.id} v{cfg.version} uses provider '{provider}', "
+                    f"which does not support {'/'.join(unsupported_types)} "
+                    f"attachments: {', '.join(unsupported_columns)}"
+                ),
+            )
         resolved_configs.append(cfg)
 
     assessment = create_assessment(