feat(generation): add support for custom structured output schema

- Introduce `generation.output_schema` and `output.format` parameters to configure direct structured output based on a user-defined JSON schema.
- When `output_schema` is provided, the generator now bypasses the standard conversation format and uses constrained decoding to generate output that adheres directly to the specified schema.
- Implement a dynamic model creation utility (`make_dynamic_model`) to convert a raw JSON schema dictionary into a Pydantic-compatible class, enabling its use with constrained decoding.
- Update the `outlines` and `transformers` dependencies to their latest versions to support these new capabilities.

Signed-off-by: Luke Hinds <lukehinds@gmail.com>

Author: lukehinds

deepfabric/config.py

@@ -263,6 +263,15 @@ class GenerationConfig(BaseModel):
     )
     save_as: str | None = Field(default=None, description="Where to save the generated samples")
 
+    # Custom structured output schema (bypasses conversation format)
+    output_schema: dict | None = Field(
+        default=None,
+        description=(
+            "JSON Schema for custom structured output. When set, bypasses the conversation "
+            "format and generates directly into this schema via constrained decoding."
+        ),
+    )
+
     # Optional LLM overrides
     llm: LLMConfig | None = Field(
         default=None, description="Optional LLM configuration overrides for generation"
@@ -314,6 +323,14 @@ class OutputConfig(BaseModel):
     )
     save_as: str = Field(..., min_length=1, description="Where to save the final dataset")
 
+    format: Literal["messages", "custom"] = Field(
+        default="messages",
+        description=(
+            "'messages' (default): OpenAI chat format. "
+            "'custom': emit records matching generation.output_schema directly."
+        ),
+    )
+
     # Optional checkpoint configuration (nested inside output)
     checkpoint: CheckpointConfig | None = Field(
         None, description="Checkpoint configuration for resumable generation"
@@ -656,6 +673,10 @@ def get_generation_params(self, **overrides) -> dict:
             params["scenario_seed"] = self.generation.tools.scenario_seed
             params["max_agent_steps"] = self.generation.tools.max_agent_steps
 
+        if self.generation.output_schema:
+            params["output_schema"] = self.generation.output_schema
+            params["output_format"] = self.output.format
+
         # Handle overrides
         override_provider = overrides.pop("provider", None)
         override_model = overrides.pop("model", None)

deepfabric/generator.py

@@ -39,6 +39,7 @@
 from .prompts import (
     AGENT_COT_TOOLS_PROMPT,
     CONVERSATION_GENERATION_PROMPT,
+    CUSTOM_SCHEMA_PROMPT,
     FREETEXT_COT_PROMPT,
     STRUCTURED_COT_PROMPT,
     AgentPromptBuilder,
@@ -132,6 +133,23 @@ class DataSetGeneratorConfig(BaseModel):
         description="Rate limiting and retry configuration (uses provider defaults if not specified)",
     )
 
+    # Custom structured output schema
+    output_schema: dict | None = Field(
+        default=None,
+        description=(
+            "JSON Schema for custom structured output. When set, bypasses the conversation "
+            "format and generates directly into this schema via constrained decoding. "
+            "Records in the JSONL output will match the schema instead of the OpenAI messages format."
+        ),
+    )
+    output_format: Literal["messages", "custom"] = Field(
+        default="messages",
+        description=(
+            "'messages' (default): OpenAI chat format. "
+            "'custom': emit records matching output_schema directly."
+        ),
+    )
+
     # Modular conversation configuration
     conversation_type: Literal["basic", "cot"] = Field(
         default="basic",
@@ -955,6 +973,17 @@ def _get_minimal_schema(self) -> type:
         """Get the conversation schema for the current config."""
         return get_conversation_schema(self.config.conversation_type)
 
+    def _get_custom_schema_model(self) -> type:
+        """Build a dynamic model class from the configured output_schema."""
+        from .llm.client import make_dynamic_model
+        return make_dynamic_model(self.config.output_schema)
+
+    def _get_prompt_template(self) -> str:
+        """Return the prompt template, using CUSTOM_SCHEMA_PROMPT when output_schema is set."""
+        if self.config.output_schema:
+            return CUSTOM_SCHEMA_PROMPT
+        return self._get_cot_prompt_template()
+
     def _emit_retry(
         self,
         sample_idx: int,
@@ -1014,6 +1043,30 @@ async def _generate_with_retry(
             Each parallel task gets its own builder instance to avoid Spin session
             conflicts when running samples concurrently (batch_size > 1).
             """
+            last_error: Exception | None = None
+            max_attempts = self.config.sample_retries + 1
+
+            # Custom schema mode: bypass conversation builder entirely and
+            # generate directly into the user-defined schema via constrained decoding.
+            if config.output_schema:
+                schema_model = self._get_custom_schema_model()
+                for attempt in range(max_attempts):
+                    try:
+                        result = await self.llm_client.generate_async(
+                            prompt,
+                            schema_model,
+                            max_tokens=config.max_tokens,
+                        )
+                        return True, result
+                    except Exception as e:  # noqa: BLE001
+                        last_error = e
+                        if is_validation_error(e) and attempt < self.config.sample_retries:
+                            self._emit_retry(sample_idx, attempt, max_attempts, e)
+                            continue
+                        return False, last_error
+                return False, last_error or Exception("Custom schema generation failed")
+
+            # Normal conversation builder mode
             # Create a fresh builder for this sample to avoid session conflicts
             # when running in parallel batches
             builder = ConversationBuilderFactory.create(
@@ -1023,9 +1076,7 @@ async def _generate_with_retry(
                 progress_reporter=self.progress_reporter,
             )
 
-            last_error: Exception | None = None
             error_feedback: str | None = None
-            max_attempts = self.config.sample_retries + 1
             logger.debug(
                 "Sample %d: max_attempts=%d (sample_retries=%d)",
                 sample_idx + 1,
@@ -1295,7 +1346,7 @@ async def create_data_async(
 
         # Calculate total samples requested
         total_samples = num_steps * batch_size
-        data_creation_prompt = self._get_cot_prompt_template()
+        data_creation_prompt = self._get_prompt_template()
 
         # Ensure checkpoint_interval is at least as large as concurrency/batch_size
         # so checkpoints align with batch boundaries
@@ -1388,7 +1439,7 @@ async def create_data_with_events_async(
 
         # Calculate total samples requested
         total_samples = num_steps * batch_size
-        data_creation_prompt = self._get_cot_prompt_template()
+        data_creation_prompt = self._get_prompt_template()
 
         # Ensure checkpoint_interval is at least as large as concurrency/batch_size
         # so checkpoints align with batch boundaries

deepfabric/llm/client.py

@@ -1,4 +1,5 @@
 import asyncio
+import json
 import logging
 import os
 import sys
@@ -1218,3 +1219,40 @@ def make_async_outlines_model(provider: str, model_name: str, **kwargs) -> Any |
     # Outlines does not currently expose async structured generation wrappers
     # for the remaining providers. Fallback to synchronous execution later.
     return None
+
+
+def make_dynamic_model(schema_dict: dict) -> type:
+    """Create a schema-compatible model class from a raw JSON schema dict.
+
+    The returned class satisfies the interface expected by LLMClient.generate_async:
+    - model_json_schema() — returns the schema dict (used by provider clients)
+    - model_validate_json(json_str) — parses JSON and returns an instance
+    - instance.model_dump() — returns the parsed dict
+
+    This lets user-defined JSON schemas flow through the existing structured-output
+    pipeline without requiring a hand-written Pydantic model.
+    """
+    _schema = dict(schema_dict)
+
+    class DynamicModel:
+        def __init__(self, data: dict) -> None:
+            self._data = data
+
+        @classmethod
+        def model_json_schema(cls, **_kwargs: Any) -> dict:
+            return _schema
+
+        @classmethod
+        def model_validate_json(cls, json_data: str) -> "DynamicModel":
+            try:
+                data = json.loads(json_data)
+            except json.JSONDecodeError as exc:
+                raise ValueError(f"Invalid JSON from model: {exc}") from exc
+            return cls(data)
+
+        def model_dump(self, exclude_none: bool = False, **_kwargs: Any) -> dict:
+            if exclude_none:
+                return {k: v for k, v in self._data.items() if v is not None}
+            return dict(self._data)
+
+    return DynamicModel

deepfabric/prompts.py

@@ -195,6 +195,12 @@ def build_tool_context_prompt(tool_registry, max_tools_per_query: int = 3) -> st
 {{{{examples}}}}
 {{{{subtopics}}}}"""
 
+CUSTOM_SCHEMA_PROMPT = """{{{{system_prompt}}}}
+{{{{instructions}}}}
+{{{{subtopics}}}}
+
+Generate a single JSON object for the topic above. Fill in all required fields accurately and completely."""
+
 CONVERSATION_GENERATION_PROMPT = """Generate a training conversation for a language model with this system prompt:
 
 <system_prompt>