Improve the R1 distill support:

- Fail if openrouter fails to return reasoning for a reasoning model - Don't use json_object mode for structured
Kiln-AI · Feb 5, 2025 · 6fa75dd · 6fa75dd
1 parent 331dd7f
commit 6fa75dd
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 16 deletions.
diff --git a/app/desktop/studio_server/finetune_api.py b/app/desktop/studio_server/finetune_api.py
@@ -28,7 +28,6 @@
     Train60Test20Val20SplitDefinition,
     Train80Test10Val10SplitDefinition,
     Train80Test20SplitDefinition,
-    dataset_filters,
 )
 from kiln_ai.utils.name_generator import generate_memorable_name
 from kiln_server.task_api import task_from_id

diff --git a/libs/core/kiln_ai/adapters/ml_model_list.py b/libs/core/kiln_ai/adapters/ml_model_list.py
@@ -818,14 +818,14 @@ class KilnModel(BaseModel):
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
                 reasoning_capable=True,
-                structured_output_mode=StructuredOutputMode.json_instruction_and_object,
+                structured_output_mode=StructuredOutputMode.json_instructions,
                 provider_options={"model": "deepseek/deepseek-r1-distill-qwen-32b"},
             ),
             KilnModelProvider(
                 name=ModelProviderName.ollama,
                 parser=ModelParserID.r1_thinking,
                 reasoning_capable=True,
-                structured_output_mode=StructuredOutputMode.json_schema,
+                structured_output_mode=StructuredOutputMode.json_instructions,
                 provider_options={"model": "deepseek-r1:32b"},
             ),
         ],
@@ -839,15 +839,15 @@ class KilnModel(BaseModel):
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
                 reasoning_capable=True,
-                structured_output_mode=StructuredOutputMode.json_instruction_and_object,
+                structured_output_mode=StructuredOutputMode.json_instructions,
                 provider_options={"model": "deepseek/deepseek-r1-distill-llama-70b"},
             ),
             KilnModelProvider(
                 name=ModelProviderName.ollama,
                 supports_data_gen=False,
                 parser=ModelParserID.r1_thinking,
                 reasoning_capable=True,
-                structured_output_mode=StructuredOutputMode.json_schema,
+                structured_output_mode=StructuredOutputMode.json_instructions,
                 provider_options={"model": "deepseek-r1:70b"},
             ),
         ],
@@ -862,15 +862,15 @@ class KilnModel(BaseModel):
                 name=ModelProviderName.openrouter,
                 supports_data_gen=False,
                 reasoning_capable=True,
-                structured_output_mode=StructuredOutputMode.json_instruction_and_object,
+                structured_output_mode=StructuredOutputMode.json_instructions,
                 provider_options={"model": "deepseek/deepseek-r1-distill-qwen-14b"},
             ),
             KilnModelProvider(
                 name=ModelProviderName.ollama,
                 supports_data_gen=False,
                 parser=ModelParserID.r1_thinking,
                 reasoning_capable=True,
-                structured_output_mode=StructuredOutputMode.json_schema,
+                structured_output_mode=StructuredOutputMode.json_instructions,
                 provider_options={"model": "deepseek-r1:14b"},
             ),
         ],

diff --git a/libs/core/kiln_ai/adapters/model_adapters/openai_model_adapter.py b/libs/core/kiln_ai/adapters/model_adapters/openai_model_adapter.py
@@ -93,11 +93,17 @@ async def _run(self, input: Dict | str) -> RunOutput:
                 ]
             )
 
+        # OpenRouter specific options for reasoning models
         extra_body = {}
-        if self.config.openrouter_style_reasoning and provider.reasoning_capable:
+        require_or_reasoning = (
+            self.config.openrouter_style_reasoning and provider.reasoning_capable
+        )
+        if require_or_reasoning:
             extra_body["include_reasoning"] = True
             # Filter to providers that support the reasoning parameter
-            extra_body["provider"] = {"require_parameters": True}
+            extra_body["provider"] = {
+                "require_parameters": True,
+            }
 
         # Main completion call
         response_format_options = await self.response_format_options()
@@ -124,13 +130,16 @@ async def _run(self, input: Dict | str) -> RunOutput:
 
         message = response.choices[0].message
 
-        # Save reasoning if it exists
-        if (
-            self.config.openrouter_style_reasoning
-            and hasattr(message, "reasoning")
-            and message.reasoning  # pyright: ignore
-        ):
-            intermediate_outputs["reasoning"] = message.reasoning  # pyright: ignore
+        # Save reasoning if it exists (OpenRouter specific format)
+        if require_or_reasoning:
+            if (
+                hasattr(message, "reasoning") and message.reasoning  # pyright: ignore
+            ):
+                intermediate_outputs["reasoning"] = message.reasoning  # pyright: ignore
+            else:
+                raise RuntimeError(
+                    "Reasoning is required for this model, but no reasoning was returned from OpenRouter."
+                )
 
         # the string content of the response
         response_content = message.content