kiln-ai 0.14.0__py3-none-any.whl → 0.15.0__py3-none-any.whl

kiln_ai/adapters/fine_tune/vertex_finetune.py

@@ -0,0 +1,217 @@
+import logging
+import time
+
+import vertexai
+from google.cloud import storage
+from google.cloud.aiplatform_v1beta1 import types as gca_types
+from vertexai.tuning import sft
+
+from kiln_ai.adapters.fine_tune.base_finetune import (
+    BaseFinetuneAdapter,
+    FineTuneParameter,
+    FineTuneStatus,
+    FineTuneStatusType,
+)
+from kiln_ai.adapters.fine_tune.dataset_formatter import DatasetFormat, DatasetFormatter
+from kiln_ai.datamodel import DatasetSplit, StructuredOutputMode, Task
+from kiln_ai.utils.config import Config
+
+logger = logging.getLogger(__name__)
+
+
+class VertexFinetune(BaseFinetuneAdapter):
+    """
+    A fine-tuning adapter for Vertex AI.
+    """
+
+    async def status(self) -> FineTuneStatus:
+        """
+        Get the status of the fine-tune.
+        """
+
+        # Update the datamodel with the latest status if it has changed
+        status = await self._status()
+        if status.status != self.datamodel.latest_status:
+            self.datamodel.latest_status = status.status
+            if self.datamodel.path:
+                self.datamodel.save_to_file()
+        return status
+
+    async def _status(self) -> FineTuneStatus:
+        if not self.datamodel or not self.datamodel.provider_id:
+            return FineTuneStatus(
+                status=FineTuneStatusType.pending,
+                message="This fine-tune has not been started or has not been assigned a provider ID.",
+            )
+
+        response = sft.SupervisedTuningJob(self.datamodel.provider_id)
+        # If the fine-tuned model ID has been updated, update the datamodel
+        try:
+            if self.datamodel.fine_tune_model_id != response.tuned_model_endpoint_name:
+                self.datamodel.fine_tune_model_id = response.tuned_model_endpoint_name
+                if self.datamodel.path:
+                    self.datamodel.save_to_file()
+        except Exception as e:
+            # Don't let this error crash the status call
+            logger.warning(f"Error updating fine-tune model ID: {e}")
+            pass
+
+        error = response.error
+        if error and error.code != 0:
+            return FineTuneStatus(
+                status=FineTuneStatusType.failed,
+                message=f"Fine Tune Job Error: {error.message} [{error.code}]",
+            )
+        state = response.state
+        if state in [
+            gca_types.JobState.JOB_STATE_FAILED,
+            gca_types.JobState.JOB_STATE_EXPIRED,
+        ]:
+            return FineTuneStatus(
+                status=FineTuneStatusType.failed,
+                message="Fine Tune Job Failed",
+            )
+        if state in [
+            gca_types.JobState.JOB_STATE_CANCELLED,
+            gca_types.JobState.JOB_STATE_CANCELLING,
+        ]:
+            return FineTuneStatus(
+                status=FineTuneStatusType.failed, message="Fine Tune Job Cancelled"
+            )
+        if state in [
+            gca_types.JobState.JOB_STATE_PENDING,
+            gca_types.JobState.JOB_STATE_QUEUED,
+        ]:
+            return FineTuneStatus(
+                status=FineTuneStatusType.pending, message="Fine Tune Job Pending"
+            )
+        if state in [
+            gca_types.JobState.JOB_STATE_RUNNING,
+        ]:
+            return FineTuneStatus(
+                status=FineTuneStatusType.running,
+                message="Fine Tune Job Running",
+            )
+        if state in [
+            gca_types.JobState.JOB_STATE_SUCCEEDED,
+            gca_types.JobState.JOB_STATE_PARTIALLY_SUCCEEDED,
+        ]:
+            return FineTuneStatus(
+                status=FineTuneStatusType.completed, message="Fine Tune Job Completed"
+            )
+
+        if state not in [
+            gca_types.JobState.JOB_STATE_UPDATING,
+            gca_types.JobState.JOB_STATE_UNSPECIFIED,
+            gca_types.JobState.JOB_STATE_PAUSED,
+        ]:
+            # While the above states map to "unknown", they are expected unknowns. Log if some new state appears we aren't expecting
+            logger.warning(f"Unknown Vertex AI Fine Tune Status: [{state}]")
+
+        return FineTuneStatus(
+            status=FineTuneStatusType.unknown, message=f"Unknown state: [{state}]"
+        )
+
+    async def _start(self, dataset: DatasetSplit) -> None:
+        task = self.datamodel.parent_task()
+        if not task:
+            raise ValueError("Task is required to start a fine-tune")
+
+        # Use chat format for unstructured output, and JSON for formatted output
+        format = DatasetFormat.VERTEX_GEMINI
+        if task.output_json_schema:
+            self.datamodel.structured_output_mode = StructuredOutputMode.json_mode
+        train_file_id = await self.generate_and_upload_jsonl(
+            dataset, self.datamodel.train_split_name, task, format
+        )
+        validation_file_id = None
+        if self.datamodel.validation_split_name:
+            validation_file_id = await self.generate_and_upload_jsonl(
+                dataset, self.datamodel.validation_split_name, task, format
+            )
+
+        hyperparameters = self.datamodel.parameters
+
+        project, location = self.get_vertex_provider_location()
+        vertexai.init(project=project, location=location)
+
+        sft_tuning_job = sft.train(
+            source_model=self.datamodel.base_model_id,
+            train_dataset=train_file_id,
+            validation_dataset=validation_file_id,
+            tuned_model_display_name=f"kiln_finetune_{self.datamodel.id}",
+            # It is recommended to use auto-selection and leave them unset
+            epochs=hyperparameters.get("epochs", None),  # type: ignore
+            adapter_size=hyperparameters.get("adapter_size", None),  # type: ignore
+            learning_rate_multiplier=hyperparameters.get(
+                "learning_rate_multiplier", None
+            ),  # type: ignore
+            labels={
+                "source": "kiln",
+                "kiln_finetune_id": str(self.datamodel.id),
+                "kiln_task_id": str(task.id),
+            },
+        )
+        self.datamodel.provider_id = sft_tuning_job.resource_name
+
+        return None
+
+    async def generate_and_upload_jsonl(
+        self, dataset: DatasetSplit, split_name: str, task: Task, format: DatasetFormat
+    ) -> str:
+        formatter = DatasetFormatter(
+            dataset, self.datamodel.system_message, self.datamodel.thinking_instructions
+        )
+        path = formatter.dump_to_file(split_name, format, self.datamodel.data_strategy)
+
+        project, location = self.get_vertex_provider_location()
+        storage_client = storage.Client(project=project)
+
+        bucket_name = "kiln-ai-data"
+
+        # Check if bucket exists and create it if it doesn't
+        if not storage_client.lookup_bucket(bucket_name):
+            bucket = storage_client.create_bucket(bucket_name, location=location)
+        else:
+            bucket = storage_client.bucket(bucket_name)
+
+        # Create a blob and upload
+        epoch_timestamp = int(time.time())
+        blob_name = f"{epoch_timestamp}/{path.name}"
+        blob = bucket.blob(blob_name)
+        blob.upload_from_filename(path)
+
+        return f"gs://{bucket.name}/{blob.name}"
+
+    @classmethod
+    def available_parameters(cls) -> list[FineTuneParameter]:
+        return [
+            FineTuneParameter(
+                name="learning_rate_multiplier",
+                type="float",
+                description="Scaling factor for the learning rate. A smaller learning rate may be useful to avoid overfitting. Defaults to 1.0 (don't scale vertex's learning rate).",
+                optional=True,
+            ),
+            FineTuneParameter(
+                name="epochs",
+                type="int",
+                description="The number of epochs to train the model for. An epoch refers to one full cycle through the training dataset. Defaults to 'auto'",
+                optional=True,
+            ),
+            FineTuneParameter(
+                name="adapter_size",
+                type="int",
+                description="The size of the adapter to use for the fine-tune. One of 1, 4, 8, or 16. By default Vertex will auto-select a size.",
+                optional=True,
+            ),
+        ]
+
+    @classmethod
+    def get_vertex_provider_location(cls) -> tuple[str, str]:
+        project = Config.shared().vertex_project_id
+        location = Config.shared().vertex_location
+        if not project or not location:
+            raise ValueError(
+                "Google Vertex project and location must be set in Kiln settings to fine tune."
+            )
+        return project, location

kiln_ai/adapters/ml_model_list.py

@@ -70,9 +70,18 @@ class ModelName(str, Enum):
     llama_3_3_70b = "llama_3_3_70b"
     gpt_4o_mini = "gpt_4o_mini"
     gpt_4o = "gpt_4o"
+    gpt_4_1 = "gpt_4_1"
+    gpt_4_1_mini = "gpt_4_1_mini"
+    gpt_4_1_nano = "gpt_4_1_nano"
+    gpt_o3_low = "gpt_o3_low"
+    gpt_o3_medium = "gpt_o3_medium"
+    gpt_o3_high = "gpt_o3_high"
     gpt_o1_low = "gpt_o1_low"
     gpt_o1_medium = "gpt_o1_medium"
     gpt_o1_high = "gpt_o1_high"
+    gpt_o4_mini_low = "gpt_o4_mini_low"
+    gpt_o4_mini_medium = "gpt_o4_mini_medium"
+    gpt_o4_mini_high = "gpt_o4_mini_high"
     gpt_o3_mini_low = "gpt_o3_mini_low"
     gpt_o3_mini_medium = "gpt_o3_mini_medium"
     gpt_o3_mini_high = "gpt_o3_mini_high"
@@ -97,6 +106,9 @@ class ModelName(str, Enum):
     gemini_1_5_flash_8b = "gemini_1_5_flash_8b"
     gemini_1_5_pro = "gemini_1_5_pro"
     gemini_2_0_flash = "gemini_2_0_flash"
+    gemini_2_0_flash_lite = "gemini_2_0_flash_lite"
+    gemini_2_5_pro = "gemini_2_5_pro"
+    gemini_2_5_flash = "gemini_2_5_flash"
     nemotron_70b = "nemotron_70b"
     mixtral_8x7b = "mixtral_8x7b"
     qwen_2p5_7b = "qwen_2p5_7b"
@@ -181,29 +193,77 @@ class KilnModel(BaseModel):
 
 
 built_in_models: List[KilnModel] = [
-    # GPT 4o Mini
+    # GPT 4.1
     KilnModel(
         family=ModelFamily.gpt,
-        name=ModelName.gpt_4o_mini,
-        friendly_name="GPT 4o Mini",
+        name=ModelName.gpt_4_1,
+        friendly_name="GPT 4.1",
         providers=[
             KilnModelProvider(
                 name=ModelProviderName.openai,
-                model_id="gpt-4o-mini",
-                provider_finetune_id="gpt-4o-mini-2024-07-18",
+                model_id="gpt-4.1",
+                provider_finetune_id="gpt-4.1-2025-04-14",
                 structured_output_mode=StructuredOutputMode.json_schema,
                 supports_logprobs=True,
             ),
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
-                model_id="openai/gpt-4o-mini",
+                model_id="openai/gpt-4.1",
                 structured_output_mode=StructuredOutputMode.json_schema,
                 supports_logprobs=True,
-                logprobs_openrouter_options=True,
             ),
             KilnModelProvider(
                 name=ModelProviderName.azure_openai,
-                model_id="gpt-4o-mini",
+                model_id="gpt-4.1",
+            ),
+        ],
+    ),
+    # GPT 4.1 Mini
+    KilnModel(
+        family=ModelFamily.gpt,
+        name=ModelName.gpt_4_1_mini,
+        friendly_name="GPT 4.1 Mini",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openai,
+                model_id="gpt-4.1-mini",
+                provider_finetune_id="gpt-4.1-mini-2025-04-14",
+                structured_output_mode=StructuredOutputMode.json_schema,
+                supports_logprobs=True,
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                model_id="openai/gpt-4.1-mini",
+                structured_output_mode=StructuredOutputMode.json_schema,
+                supports_logprobs=True,
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.azure_openai,
+                model_id="gpt-4.1-mini",
+            ),
+        ],
+    ),
+    # GPT 4.1 Nano
+    KilnModel(
+        family=ModelFamily.gpt,
+        name=ModelName.gpt_4_1_nano,
+        friendly_name="GPT 4.1 Nano",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openai,
+                model_id="gpt-4.1-nano",
+                structured_output_mode=StructuredOutputMode.json_schema,
+                supports_logprobs=True,
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                model_id="openai/gpt-4.1-nano",
+                structured_output_mode=StructuredOutputMode.json_schema,
+                supports_logprobs=True,
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.azure_openai,
+                model_id="gpt-4.1-nano",
             ),
         ],
     ),
@@ -233,6 +293,102 @@ built_in_models: List[KilnModel] = [
             ),
         ],
     ),
+    # GPT 4o Mini
+    KilnModel(
+        family=ModelFamily.gpt,
+        name=ModelName.gpt_4o_mini,
+        friendly_name="GPT 4o Mini",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openai,
+                model_id="gpt-4o-mini",
+                provider_finetune_id="gpt-4o-mini-2024-07-18",
+                structured_output_mode=StructuredOutputMode.json_schema,
+                supports_logprobs=True,
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                model_id="openai/gpt-4o-mini",
+                structured_output_mode=StructuredOutputMode.json_schema,
+                supports_logprobs=True,
+                logprobs_openrouter_options=True,
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.azure_openai,
+                model_id="gpt-4o-mini",
+            ),
+        ],
+    ),
+    # GPT o4 Mini Low
+    KilnModel(
+        family=ModelFamily.gpt,
+        name=ModelName.gpt_o4_mini_low,
+        friendly_name="GPT o4 Mini - Low",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openai,
+                model_id="o4-mini",
+                thinking_level="low",
+                structured_output_mode=StructuredOutputMode.json_schema,
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.azure_openai,
+                model_id="o4-mini",
+                structured_output_mode=StructuredOutputMode.json_schema,
+                thinking_level="low",
+            ),
+        ],
+    ),
+    # GPT o4 Mini Medium
+    KilnModel(
+        family=ModelFamily.gpt,
+        name=ModelName.gpt_o4_mini_medium,
+        friendly_name="GPT o4 Mini - Medium",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openai,
+                model_id="o4-mini",
+                thinking_level="medium",
+                structured_output_mode=StructuredOutputMode.json_schema,
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.azure_openai,
+                model_id="o4-mini",
+                structured_output_mode=StructuredOutputMode.json_schema,
+                thinking_level="medium",
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                model_id="openai/o4-mini",
+                structured_output_mode=StructuredOutputMode.json_schema,
+            ),
+        ],
+    ),
+    # GPT o4 Mini High
+    KilnModel(
+        family=ModelFamily.gpt,
+        name=ModelName.gpt_o4_mini_high,
+        friendly_name="GPT o4 Mini - High",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openai,
+                model_id="o4-mini",
+                thinking_level="high",
+                structured_output_mode=StructuredOutputMode.json_schema,
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.azure_openai,
+                model_id="o4-mini",
+                structured_output_mode=StructuredOutputMode.json_schema,
+                thinking_level="high",
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                model_id="openai/o4-mini-high",
+                structured_output_mode=StructuredOutputMode.json_schema,
+            ),
+        ],
+    ),
     # GPT o3 Mini Low
     KilnModel(
         family=ModelFamily.gpt,
@@ -293,6 +449,66 @@ built_in_models: List[KilnModel] = [
             ),
         ],
     ),
+    # GPT o3 Low
+    KilnModel(
+        family=ModelFamily.gpt,
+        name=ModelName.gpt_o3_low,
+        friendly_name="GPT o3 - Low",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openai,
+                model_id="o3",
+                thinking_level="low",
+                structured_output_mode=StructuredOutputMode.json_schema,
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.azure_openai,
+                model_id="o3",
+                structured_output_mode=StructuredOutputMode.json_schema,
+                thinking_level="low",
+            ),
+        ],
+    ),
+    # GPT o3 Medium
+    KilnModel(
+        family=ModelFamily.gpt,
+        name=ModelName.gpt_o3_medium,
+        friendly_name="GPT o3 - Medium",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openai,
+                model_id="o3",
+                thinking_level="medium",
+                structured_output_mode=StructuredOutputMode.json_schema,
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.azure_openai,
+                model_id="o3",
+                structured_output_mode=StructuredOutputMode.json_schema,
+                thinking_level="medium",
+            ),
+        ],
+    ),
+    # GPT o3 High
+    KilnModel(
+        family=ModelFamily.gpt,
+        name=ModelName.gpt_o3_high,
+        friendly_name="GPT o3 - High",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openai,
+                model_id="o3",
+                thinking_level="high",
+                structured_output_mode=StructuredOutputMode.json_schema,
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.azure_openai,
+                model_id="o3",
+                structured_output_mode=StructuredOutputMode.json_schema,
+                thinking_level="high",
+            ),
+        ],
+    ),
     # GPT o1 Low
     KilnModel(
         family=ModelFamily.gpt,
@@ -440,6 +656,100 @@ built_in_models: List[KilnModel] = [
             ),
         ],
     ),
+    # Gemini 2.5 Pro
+    KilnModel(
+        family=ModelFamily.gemini,
+        name=ModelName.gemini_2_5_pro,
+        friendly_name="Gemini 2.5 Pro",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                model_id="google/gemini-2.5-pro-preview-03-25",
+                structured_output_mode=StructuredOutputMode.json_schema,
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.gemini_api,
+                model_id="gemini-2.5-pro-preview-03-25",
+                structured_output_mode=StructuredOutputMode.json_schema,
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.vertex,
+                model_id="gemini-2.5-pro-preview-03-25",
+                structured_output_mode=StructuredOutputMode.json_schema,
+            ),
+        ],
+    ),
+    # Gemini 2.5 Flash
+    KilnModel(
+        family=ModelFamily.gemini,
+        name=ModelName.gemini_2_5_flash,
+        friendly_name="Gemini 2.5 Flash",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                model_id="google/gemini-2.5-flash-preview",
+                structured_output_mode=StructuredOutputMode.json_schema,
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.gemini_api,
+                model_id="gemini-2.5-flash-preview-04-17",
+                structured_output_mode=StructuredOutputMode.json_schema,
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.vertex,
+                model_id="gemini-2.5-flash-preview-04-17",
+                structured_output_mode=StructuredOutputMode.json_schema,
+            ),
+        ],
+    ),
+    # Gemini 2.0 Flash
+    KilnModel(
+        family=ModelFamily.gemini,
+        name=ModelName.gemini_2_0_flash,
+        friendly_name="Gemini 2.0 Flash",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                model_id="google/gemini-2.0-flash-001",
+                structured_output_mode=StructuredOutputMode.json_instruction_and_object,
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.gemini_api,
+                model_id="gemini-2.0-flash",
+                structured_output_mode=StructuredOutputMode.json_instruction_and_object,
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.vertex,
+                model_id="gemini-2.0-flash",
+                structured_output_mode=StructuredOutputMode.json_instruction_and_object,
+                provider_finetune_id="gemini-2.0-flash-001",
+            ),
+        ],
+    ),
+    # Gemini 2.0 Flash Lite
+    KilnModel(
+        family=ModelFamily.gemini,
+        name=ModelName.gemini_2_0_flash_lite,
+        friendly_name="Gemini 2.0 Flash Lite",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                model_id="google/gemini-2.0-flash-lite-001",
+                structured_output_mode=StructuredOutputMode.json_instruction_and_object,
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.gemini_api,
+                model_id="gemini-2.0-flash-lite",
+                structured_output_mode=StructuredOutputMode.json_instruction_and_object,
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.vertex,
+                model_id="gemini-2.0-flash-lite",
+                structured_output_mode=StructuredOutputMode.json_instruction_and_object,
+                provider_finetune_id="gemini-2.0-flash-lite-001",
+            ),
+        ],
+    ),
     # Gemini 1.5 Pro
     KilnModel(
         family=ModelFamily.gemini,
@@ -506,29 +816,6 @@ built_in_models: List[KilnModel] = [
             ),
         ],
     ),
-    # Gemini 2.0 Flash
-    KilnModel(
-        family=ModelFamily.gemini,
-        name=ModelName.gemini_2_0_flash,
-        friendly_name="Gemini 2.0 Flash",
-        providers=[
-            KilnModelProvider(
-                name=ModelProviderName.openrouter,
-                model_id="google/gemini-2.0-flash-001",
-                structured_output_mode=StructuredOutputMode.json_instruction_and_object,
-            ),
-            KilnModelProvider(
-                name=ModelProviderName.gemini_api,
-                model_id="gemini-2.0-flash",
-                structured_output_mode=StructuredOutputMode.json_instruction_and_object,
-            ),
-            KilnModelProvider(
-                name=ModelProviderName.vertex,
-                model_id="gemini-2.0-flash",
-                structured_output_mode=StructuredOutputMode.json_instruction_and_object,
-            ),
-        ],
-    ),
     # Nemotron 70B
     KilnModel(
         family=ModelFamily.llama,
@@ -760,12 +1047,6 @@ built_in_models: List[KilnModel] = [
                 supports_data_gen=False,
                 model_id="llama3.2",
             ),
-            KilnModelProvider(
-                name=ModelProviderName.fireworks_ai,
-                supports_structured_output=False,
-                supports_data_gen=False,
-                model_id="accounts/fireworks/models/llama-v3p2-3b-instruct",
-            ),
             KilnModelProvider(
                 name=ModelProviderName.huggingface,
                 model_id="meta-llama/Llama-3.2-3B-Instruct",

kiln_ai/adapters/model_adapters/base_adapter.py

@@ -3,20 +3,16 @@ from abc import ABCMeta, abstractmethod
 from dataclasses import dataclass
 from typing import Dict, Literal, Tuple
 
+import jsonschema
+
 from kiln_ai.adapters.ml_model_list import KilnModelProvider, StructuredOutputMode
 from kiln_ai.adapters.parsers.json_parser import parse_json_string
 from kiln_ai.adapters.parsers.parser_registry import model_parser_from_id
 from kiln_ai.adapters.prompt_builders import prompt_builder_from_id
 from kiln_ai.adapters.provider_tools import kiln_model_provider_from
 from kiln_ai.adapters.run_output import RunOutput
-from kiln_ai.datamodel import (
-    DataSource,
-    DataSourceType,
-    Task,
-    TaskOutput,
-    TaskRun,
-)
-from kiln_ai.datamodel.json_schema import validate_schema
+from kiln_ai.datamodel import DataSource, DataSourceType, Task, TaskOutput, TaskRun
+from kiln_ai.datamodel.json_schema import validate_schema_with_value_error
 from kiln_ai.datamodel.task import RunConfig
 from kiln_ai.utils.config import Config
 
@@ -103,7 +99,12 @@ class BaseAdapter(metaclass=ABCMeta):
         if self.input_schema is not None:
             if not isinstance(input, dict):
                 raise ValueError(f"structured input is not a dict: {input}")
-            validate_schema(input, self.input_schema)
+
+            validate_schema_with_value_error(
+                input,
+                self.input_schema,
+                "This task requires a specific input schema. While the model produced JSON, that JSON didn't meet the schema. Search 'Troubleshooting Structured Data Issues' in our docs for more information.",
+            )
 
         # Run
         run_output = await self._run(input)
@@ -125,7 +126,11 @@ class BaseAdapter(metaclass=ABCMeta):
                 raise RuntimeError(
                     f"structured response is not a dict: {parsed_output.output}"
                 )
-            validate_schema(parsed_output.output, self.output_schema)
+            validate_schema_with_value_error(
+                parsed_output.output,
+                self.output_schema,
+                "This task requires a specific output schema. While the model produced JSON, that JSON didn't meet the schema. Search 'Troubleshooting Structured Data Issues' in our docs for more information.",
+            )
         else:
             if not isinstance(parsed_output.output, str):
                 raise RuntimeError(