feat: Implement structured output for AI agent Tasks

products/tasks/backend/api.py

@@ -14,6 +14,7 @@
 from django.utils import timezone
 
 import requests as http_requests
+import jsonschema
 import posthoganalytics
 from drf_spectacular.utils import OpenApiResponse, extend_schema
 from rest_framework import status, viewsets
@@ -58,6 +59,7 @@
     TaskRunRelayMessageRequestSerializer,
     TaskRunRelayMessageResponseSerializer,
     TaskRunSessionLogsQuerySerializer,
+    TaskRunSetOutputRequestSerializer,
     TaskRunUpdateSerializer,
     TaskSerializer,
 )
@@ -542,7 +544,7 @@ def perform_create(self, serializer):
         serializer.save(team=self.team, task=task)
 
     @validated_request(
-        request_serializer=None,
+        request_serializer=TaskRunSetOutputRequestSerializer,
         responses={
             200: OpenApiResponse(response=TaskRunDetailSerializer, description="Run with updated output"),
             404: OpenApiResponse(description="Run not found"),
@@ -558,17 +560,22 @@ def perform_create(self, serializer):
     )
     def set_output(self, request, pk=None, **kwargs):
         task_run = cast(TaskRun, self.get_object())
+        task = cast(Task, task_run.task)
+        output_data = request.validated_data["output"]
 
-        output_data = request.data.get("output", {})
-        if not isinstance(output_data, dict):
-            return Response(
-                ErrorResponseSerializer({"error": "output must be a dictionary"}).data,
-                status=status.HTTP_400_BAD_REQUEST,
-            )
-
-        # TODO: Validate output data according to schema for the task type.
+        if task.json_schema:
+            try:
+                jsonschema.validate(instance=output_data, schema=task.json_schema)
+            except jsonschema.ValidationError as e:
+                return Response(
+                    ErrorResponseSerializer({"error": f"Output validation error: {e.message}"}).data,
+                    status=status.HTTP_400_BAD_REQUEST,
+                )
         task_run.output = output_data
         task_run.save(update_fields=["output", "updated_at"])
+        # We only really want to complete the task run if it's a structured output task.
+        if task.json_schema:
+            self._signal_workflow_completion(task_run, TaskRun.Status.COMPLETED, None)
         task_run.publish_stream_state_event()
         self._post_slack_update_for_pr(task_run)
 

products/tasks/backend/models.py

@@ -6,6 +6,11 @@
 import secrets
 from typing import TYPE_CHECKING, Any, Literal, Optional
 
+from django.db.models.signals import post_save
+from django.dispatch import receiver
+
+from pydantic import BaseModel
+
 if TYPE_CHECKING:
     from products.slack_app.backend.slack_thread import SlackThreadContext
 
@@ -35,6 +40,12 @@
 LogLevel = Literal["debug", "info", "warn", "error"]
 
 
+def resolve_schema(schema: type[BaseModel] | dict) -> dict:
+    if isinstance(schema, dict):
+        return schema
+    return schema.model_json_schema()
+
+
 class Task(DeletedMetaFields, models.Model):
     class OriginProduct(models.TextChoices):
         ERROR_TRACKING = "error_tracking", "Error Tracking"
@@ -244,6 +255,7 @@ def create_and_run(
         signal_report_id: str | None = None,
         sandbox_environment_id: str | None = None,
         internal: bool = False,
+        output_schema: type[BaseModel] | dict | None = None,
     ) -> "Task":
         from products.tasks.backend.temporal.client import execute_task_processing_workflow
 
@@ -270,6 +282,7 @@ def create_and_run(
             github_integration=github_integration,
             repository=repository,
             internal=internal,
+            json_schema=resolve_schema(output_schema) if output_schema else None,
             **({"signal_report_id": signal_report_id} if signal_report_id else {}),
         )
 
@@ -521,6 +534,20 @@ def mark_completed(self):
             {"duration_seconds": self._duration_seconds()},
         )
 
+    def track_structured_result(self):
+        """Track a structured result event with properties from the run output."""
+        if not self.output:
+            return
+
+        try:
+            self.capture_event("task_run_structured_result", {"result": self.output})
+        except Exception as e:
+            logger.warning(
+                "task_run.track_structured_result_failed",
+                task_run_id=str(self.id),
+                error=str(e),
+            )
+
     def mark_failed(self, error: str):
         """Mark the progress as failed with an error message."""
         self.status = self.Status.FAILED
@@ -857,3 +884,21 @@ class Meta:
 
     def __str__(self):
         return f"{self.user} redeemed {self.invite_code}"
+
+
+@receiver(post_save, sender=TaskRun)
+def track_task_run_completion(sender, instance: TaskRun, created: bool, **kwargs):
+    try:
+        if (
+            not created
+            and instance.status == TaskRun.Status.COMPLETED
+            and instance.output
+            and instance.task.json_schema
+        ):
+            instance.track_structured_result()
+    except Exception as e:
+        logger.warning(
+            "task_run.track_task_run_completion_failed",
+            task_run_id=str(instance.id),
+            error=str(e),
+        )

products/tasks/backend/serializers.py

@@ -218,6 +218,12 @@ def update(self, instance, validated_data):
         return super().update(instance, validated_data)
 
 
+class TaskRunSetOutputRequestSerializer(serializers.Serializer):
+    output = serializers.JSONField(
+        help_text="Output data from the run. Validated against the task's json_schema if one is set."
+    )
+
+
 class ErrorResponseSerializer(serializers.Serializer):
     error = serializers.CharField(help_text="Error message")
 

products/tasks/backend/temporal/process_task/activities/get_task_processing_context.py

@@ -38,6 +38,7 @@ class TaskProcessingContext:
     _branch: str | None = None
     sandbox_environment_name: str | None = None
     allowed_domains: list[str] | None = None
+    json_schema: dict | None = None
 
     @property
     def mode(self) -> str:
@@ -164,4 +165,5 @@ def get_task_processing_context(input: GetTaskProcessingContextInput) -> TaskPro
         _branch=task_run.branch,
         sandbox_environment_name=sandbox_environment_name,
         allowed_domains=allowed_domains,
+        json_schema=task.json_schema,
     )