nolabs-ai
diff --git a/‎deepfabric/config.py‎
Lines changed: 21 additions & 0 deletions b/‎deepfabric/config.py‎
Lines changed: 21 additions & 0 deletions
diff --git a/‎deepfabric/generator.py‎
Lines changed: 57 additions & 5 deletions b/‎deepfabric/generator.py‎
Lines changed: 57 additions & 5 deletions
diff --git a/‎deepfabric/llm/client.py‎
Lines changed: 38 additions & 0 deletions b/‎deepfabric/llm/client.py‎
Lines changed: 38 additions & 0 deletions
diff --git a/‎deepfabric/prompts.py‎
Lines changed: 6 additions & 0 deletions b/‎deepfabric/prompts.py‎
Lines changed: 6 additions & 0 deletions
diff --git a/‎pyproject.toml‎
Lines changed: 1 addition & 1 deletion b/‎pyproject.toml‎
Lines changed: 1 addition & 1 deletion
@@ -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)
 
@@ -34,11 +34,13 @@
 from .error_codes import classify_error
 from .exceptions import DataSetGeneratorError
 from .llm import LLMClient
+from .llm.client import make_dynamic_model
 from .metrics import trace
 from .progress import ProgressReporter
 from .prompts import (
     AGENT_COT_TOOLS_PROMPT,
     CONVERSATION_GENERATION_PROMPT,
+    CUSTOM_SCHEMA_PROMPT,
     FREETEXT_COT_PROMPT,
     STRUCTURED_COT_PROMPT,
     AgentPromptBuilder,
@@ -132,6 +134,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 +974,16 @@ 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."""
+        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,
@@ -1006,14 +1035,39 @@ async def _generate_structured_samples_async(
             # Create a copy of config with sys_msg overridden
             config = self.config.model_copy(update={"sys_msg": include_sys_msg})
 
-        async def _generate_with_retry(
+        async def _generate_with_retry(  # noqa: PLR0911
             prompt: str, sample_idx: int, topic_path_info: TopicPath | None
         ) -> tuple[bool, Exception | Conversation]:
             """Generate a single sample with per-sample retry for validation errors.
 
             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,
+                        )
+                    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
+                    else:
+                        return True, result
+                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 +1077,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 +1347,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 +1440,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
 
@@ -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
@@ -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>
 
@@ -1,6 +1,6 @@
 [project]
 name = "DeepFabric"
-version = "4.12.0"
+version = "4.13.0"
 description = "Curate High Quality Datasets, Train, Evaluate and Ship"
 authors = [{name = "DeepFabric Team", email = "oss@alwaysfurther.ai"}]
 readme = "README.md"