Create Fine-tune with traces & tools

app/desktop/studio_server/eval_api.py

@@ -322,6 +322,27 @@ async def get_task_run_configs(
         task = task_from_id(project_id, task_id)
         return task.run_configs()
 
+    @app.get("/api/projects/{project_id}/tasks/{task_id}/run_configs/")
+    async def get_run_configs(project_id: str, task_id: str) -> list[TaskRunConfig]:
+        # Returns all run configs of a given task.
+        task = task_from_id(project_id, task_id)
+        configs = task.run_configs()
+
+        # Get run configs from finetunes
+        finetunes = task.finetunes()
+        for finetune in finetunes:
+            if finetune.run_config is not None:
+                configs.append(
+                    TaskRunConfig(
+                        id=f"finetune_run_config::{project_id}::{task_id}::{finetune.id}",
+                        name=finetune.name,
+                        description=finetune.description,
+                        run_config_properties=finetune.run_config,
+                    )
+                )
+
+        return configs
+
     @app.get("/api/projects/{project_id}/tasks/{task_id}/eval/{eval_id}")
     async def get_eval(project_id: str, task_id: str, eval_id: str) -> Eval:
         return eval_from_id(project_id, task_id, eval_id)

app/desktop/studio_server/finetune_api.py

@@ -1,9 +1,9 @@
 import logging
 from enum import Enum
-from typing import Dict
+from typing import Annotated, Dict
 
 import httpx
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Query
 from fastapi.responses import StreamingResponse
 from kiln_ai.adapters.fine_tune.base_finetune import FineTuneParameter, FineTuneStatus
 from kiln_ai.adapters.fine_tune.dataset_formatter import (
@@ -42,6 +42,7 @@
     Train80Test20SplitDefinition,
     Train80Val20SplitDefinition,
 )
+from kiln_ai.datamodel.run_config import RunConfigProperties
 from kiln_ai.utils.config import Config
 from kiln_ai.utils.name_generator import generate_memorable_name
 from kiln_server.task_api import task_from_id
@@ -61,6 +62,7 @@ class FinetuneProviderModel(BaseModel):
             ChatStrategy.two_message_cot,
         ]
     )
+    supports_function_calling: bool = True
 
 
 class FinetuneProvider(BaseModel):
@@ -89,6 +91,10 @@ class FinetuneDatasetInfo(BaseModel):
     existing_finetunes: list[Finetune]
     finetune_tags: list[FinetuneDatasetTagInfo]
 
+    # eligible finetune data based on tool selection
+    eligible_datasets: list[DatasetSplit]
+    eligible_finetune_tags: list[FinetuneDatasetTagInfo]
+
 
 class DatasetSplitType(Enum):
     """Dataset split types used in the API. Any split type can be created in code."""
@@ -133,6 +139,7 @@ class CreateFinetuneRequest(BaseModel):
     custom_system_message: str | None = None
     custom_thinking_instructions: str | None = None
     data_strategy: ChatStrategy
+    run_config_properties: RunConfigProperties | None = None
 
     @model_validator(mode="after")
     def validate_data_strategy(self):
@@ -170,6 +177,72 @@ def finetune_from_id(project_id: str, task_id: str, finetune_id: str) -> Finetun
     return finetune
 
 
+def compute_finetune_tag_info(
+    task: Task, tool_filter: list[str] | None = None
+) -> list[FinetuneDatasetTagInfo]:
+    """
+    Compute finetune tag counts with quality filters, optionally filtering by tools.
+
+    Args:
+        task: The task to analyze
+        tool_filter: Optional list of tool IDs to filter by. Only runs with exactly these tools will be included.
+
+    Returns:
+        List of FinetuneDatasetTagInfo with counts for each fine_tune tag
+    """
+    finetune_tag_counts: Dict[str, int] = {}
+    reasoning_count: Dict[str, int] = {}
+    high_quality_count: Dict[str, int] = {}
+    reasoning_and_high_quality_count: Dict[str, int] = {}
+
+    required_tools_set = set(tool_filter) if tool_filter else None
+
+    for sample in task.runs(readonly=True):
+        # filter by tools if provided
+        if required_tools_set is not None:
+            if sample.output is None:
+                # A few places in code guards against this though the field isn't optional. Just be defensive.
+                continue
+
+            run_tools_set: set[str] = set()
+            if (
+                sample.output.source
+                and sample.output.source.run_config
+                and sample.output.source.run_config.tools_config
+            ):
+                run_tools_set = set(sample.output.source.run_config.tools_config.tools)
+
+            if run_tools_set != required_tools_set:
+                continue
+
+        for tag in sample.tags:
+            if tag.startswith("fine_tune"):
+                finetune_tag_counts[tag] = finetune_tag_counts.get(tag, 0) + 1
+                is_reasoning = ThinkingModelDatasetFilter(sample)
+                is_high_quality = HighRatingDatasetFilter(sample)
+                if is_reasoning:
+                    reasoning_count[tag] = reasoning_count.get(tag, 0) + 1
+                if is_high_quality:
+                    high_quality_count[tag] = high_quality_count.get(tag, 0) + 1
+                if is_reasoning and is_high_quality:
+                    reasoning_and_high_quality_count[tag] = (
+                        reasoning_and_high_quality_count.get(tag, 0) + 1
+                    )
+
+    return [
+        FinetuneDatasetTagInfo(
+            tag=tag,
+            count=count,
+            reasoning_count=reasoning_count.get(tag, 0),
+            high_quality_count=high_quality_count.get(tag, 0),
+            reasoning_and_high_quality_count=reasoning_and_high_quality_count.get(
+                tag, 0
+            ),
+        )
+        for tag, count in finetune_tag_counts.items()
+    ]
+
+
 def connect_fine_tune_api(app: FastAPI):
     @app.get("/api/projects/{project_id}/tasks/{task_id}/dataset_splits")
     async def dataset_splits(project_id: str, task_id: str) -> list[DatasetSplit]:
@@ -239,10 +312,17 @@ async def finetune_providers() -> list[FinetuneProvider]:
                 if provider.provider_finetune_id:
                     if provider.name not in provider_models:
                         provider_models[provider.name] = []
+
+                    # special case, Together AI doesn't support fine-tuning with tools
+                    supports_tools = (
+                        provider.supports_function_calling
+                        and provider.name != ModelProviderName.together_ai
+                    )
                     provider_models[provider.name].append(
                         FinetuneProviderModel(
                             name=model.friendly_name,
                             id=provider.provider_finetune_id,
+                            supports_function_calling=supports_tools,
                         )
                     )
 
@@ -285,46 +365,32 @@ async def finetune_hyperparameters(
 
     @app.get("/api/projects/{project_id}/tasks/{task_id}/finetune_dataset_info")
     async def finetune_dataset_info(
-        project_id: str, task_id: str
+        project_id: str,
+        task_id: str,
+        tool_ids: Annotated[list[str] | None, Query()] = None,
     ) -> FinetuneDatasetInfo:
         task = task_from_id(project_id, task_id)
         existing_datasets = task.dataset_splits()
         existing_finetunes = task.finetunes()
 
-        finetune_tag_counts: Dict[str, int] = {}
-        reasoning_count: Dict[str, int] = {}
-        high_quality_count: Dict[str, int] = {}
-        reasoning_and_high_quality_count: Dict[str, int] = {}
-        for sample in task.runs(readonly=True):
-            for tag in sample.tags:
-                if tag.startswith("fine_tune"):
-                    finetune_tag_counts[tag] = finetune_tag_counts.get(tag, 0) + 1
-                    is_reasoning = ThinkingModelDatasetFilter(sample)
-                    is_high_quality = HighRatingDatasetFilter(sample)
-                    if is_reasoning:
-                        reasoning_count[tag] = reasoning_count.get(tag, 0) + 1
-                    if is_high_quality:
-                        high_quality_count[tag] = high_quality_count.get(tag, 0) + 1
-                    if is_reasoning and is_high_quality:
-                        reasoning_and_high_quality_count[tag] = (
-                            reasoning_and_high_quality_count.get(tag, 0) + 1
-                        )
+        finetune_tags = compute_finetune_tag_info(task, tool_filter=None)
+        eligible_finetune_tags = compute_finetune_tag_info(task, tool_filter=tool_ids)
+
+        eligible_datasets = existing_datasets
+        if tool_ids:
+            required_tools_set = set(tool_ids)
+            eligible_datasets = [
+                dataset
+                for dataset in existing_datasets
+                if set(dataset.tool_info().tools) == required_tools_set
+            ]
 
         return FinetuneDatasetInfo(
             existing_datasets=existing_datasets,
             existing_finetunes=existing_finetunes,
-            finetune_tags=[
-                FinetuneDatasetTagInfo(
-                    tag=tag,
-                    count=count,
-                    reasoning_count=reasoning_count.get(tag, 0),
-                    high_quality_count=high_quality_count.get(tag, 0),
-                    reasoning_and_high_quality_count=reasoning_and_high_quality_count.get(
-                        tag, 0
-                    ),
-                )
-                for tag, count in finetune_tag_counts.items()
-            ],
+            finetune_tags=finetune_tags,
+            eligible_datasets=eligible_datasets,
+            eligible_finetune_tags=eligible_finetune_tags,
         )
 
     @app.post("/api/projects/{project_id}/tasks/{task_id}/dataset_splits")
@@ -392,6 +458,7 @@ async def create_finetune(
             description=request.description,
             validation_split_name=request.validation_split_name,
             data_strategy=request.data_strategy,
+            run_config=request.run_config_properties,
         )
 
         return finetune_model
@@ -447,7 +514,7 @@ async def download_dataset_jsonl(
             system_message=system_message,
             thinking_instructions=thinking_instructions,
         )
-        path = dataset_formatter.dump_to_file(
+        path = await dataset_formatter.dump_to_file(
             split_name,
             format_type_typed,
             data_strategy_typed,
@@ -564,10 +631,14 @@ async def fetch_fireworks_finetune_models() -> list[FinetuneProviderModel]:
             else:
                 name = display_name + " (" + id_tail + ")"
 
+            # Check if the model supports tools via 'supportTools'
+            supports_tools = model.get("supportsTools", False)
+
             tuneable_models.append(
                 FinetuneProviderModel(
                     name=name,
                     id=id,
+                    supports_function_calling=supports_tools,
                 )
             )
 

app/desktop/studio_server/provider_api.py

@@ -33,6 +33,7 @@
 )
 from kiln_ai.adapters.provider_tools import provider_name_from_id, provider_warnings
 from kiln_ai.adapters.reranker_list import built_in_rerankers
+from kiln_ai.datamodel.finetune import Finetune
 from kiln_ai.datamodel.registry import all_projects
 from kiln_ai.utils.config import Config
 from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
@@ -155,6 +156,8 @@ class ModelDetails(BaseModel):
     # True if this is a untested model (typically user added). We don't know if these support structured output, data gen, etc. They should appear in their own section in the UI.
     untested_model: bool = Field(default=False)
     task_filter: List[str] | None = Field(default=None)
+    # if the model has a model-specific run config which should be used when running the model (like a fine-tune model's baked in run config)
+    model_specific_run_config: str | None = Field(default=None)
 
 
 class AvailableModels(BaseModel):
@@ -1438,6 +1441,22 @@ def custom_models() -> AvailableModels | None:
     )
 
 
+def fine_tune_model_structured_output_mode(
+    fine_tune: Finetune,
+) -> StructuredOutputMode:
+    # Current field
+    if fine_tune.run_config and fine_tune.run_config.structured_output_mode is not None:
+        return fine_tune.run_config.structured_output_mode
+    # Legacy field
+    legacy_structured_output_mode = fine_tune.structured_output_mode
+    if legacy_structured_output_mode is not None and isinstance(
+        legacy_structured_output_mode, StructuredOutputMode
+    ):
+        return legacy_structured_output_mode
+    # Fallback
+    return StructuredOutputMode.json_instructions
+
+
 def all_fine_tuned_models() -> AvailableModels | None:
     # Add any fine tuned models
     models: List[ModelDetails] = []
@@ -1447,9 +1466,14 @@ def all_fine_tuned_models() -> AvailableModels | None:
             for fine_tune in task.finetunes():
                 # check if the fine tune is completed
                 if fine_tune.fine_tune_model_id:
+                    model_specific_run_config = (
+                        f"finetune_run_config::{project.id}::{task.id}::{fine_tune.id}"
+                        if fine_tune.run_config is not None
+                        else None
+                    )
                     models.append(
                         ModelDetails(
-                            id=f"{project.id}::{task.id}::{fine_tune.id}",
+                            id=fine_tune.model_id(),
                             name=fine_tune.name
                             + f" ({provider_name_from_id(fine_tune.provider)})",
                             # YMMV, but we'll assume all fine tuned models support structured output, data gen, and tools as they may have been trained with them
@@ -1462,16 +1486,10 @@ def all_fine_tuned_models() -> AvailableModels | None:
                             suggested_for_evals=False,
                             uncensored=False,
                             suggested_for_uncensored_data_gen=False,
-                            structured_output_mode=(
-                                fine_tune_mode
-                                if (
-                                    fine_tune_mode := getattr(
-                                        fine_tune, "structured_output_mode", None
-                                    )
-                                )
-                                and isinstance(fine_tune_mode, StructuredOutputMode)
-                                else StructuredOutputMode.json_instructions
+                            structured_output_mode=fine_tune_model_structured_output_mode(
+                                fine_tune
                             ),
+                            model_specific_run_config=model_specific_run_config,
                             supports_vision=False,
                             supports_doc_extraction=False,
                             suggested_for_doc_extraction=False,