Feature/structured output

coolprompt/assistant.py

@@ -17,6 +17,7 @@
 from coolprompt.utils.enums import PD_Method, Task
 from coolprompt.utils.correction.corrector import correct
 from coolprompt.utils.correction.rule import LanguageRule
+from coolprompt.utils.utils import get_dataset_split
 
 from coolprompt.optimizer.autoprompting_method import AutoPromptingMethod
 
@@ -91,33 +92,6 @@ def reset_stats(self):
         if hasattr(self._target_model, "reset_stats"):
             self._target_model.reset_stats()
 
-    def _get_dataset_split(
-        self,
-        dataset: Iterable[str],
-        target: Iterable[str],
-        validation_size: float,
-        train_as_test: bool,
-    ) -> Tuple[Iterable[str], Iterable[str], Iterable[str], Iterable[str]]:
-        """Split the dataset into training and validation sets.
-
-        Args:
-            dataset (Iterable[str]): Input texts.
-            target (Iterable[str]): Corresponding labels/targets.
-            validation_size (float): Fraction of data to use for validation.
-            train_as_test (bool): If True, use the full dataset as both
-                train and validation (ignoring `validation_size`).
-
-        Returns:
-            Tuple[Iterable[str], Iterable[str], Iterable[str], Iterable[str]]:
-                A tuple (train_data, val_data, train_targets, val_targets).
-        """
-        if train_as_test:
-            return (dataset, dataset, target, target)
-        train_data, val_data, train_targets, val_targets = train_test_split(
-            dataset, target, test_size=validation_size
-        )
-        return (train_data, val_data, train_targets, val_targets)
-
     def run(
         self,
         start_prompt: str,
@@ -144,6 +118,7 @@ def run(
         return_final_prompt: bool = True,
         hyper_meta_info: dict = None,
         system_model_as_optimizer: bool = False,
+        use_structured_output: bool = True,
         **kwargs,
     ) -> Optional[str]:
         """Run prompt optimization using the selected method.
@@ -205,6 +180,8 @@ def run(
                 merged into the meta-info block for ``hyper`` and ``hyper_light``.
             system_model_as_optimizer (bool): If True, use the system model for
                 optimizing processes, while target model will be used for inference.
+            use_structured_output (bool): Either to use structured output or not.
+                Defaults to True.
             **kwargs: Additional arguments passed to the optimization method.
 
         Returns:
@@ -220,7 +197,10 @@ def run(
             validate_verbose(verbose)
             set_verbose(verbose)
 
-        task_detector = TaskDetector(self._system_model)
+        task_detector = TaskDetector(
+            model=self._system_model,
+            use_structured_output=use_structured_output
+        )
         if task is None:
             task = task_detector.generate(start_prompt)
 
@@ -255,10 +235,17 @@ def run(
         )
         metric_name = base_metric._get_name()
         evaluator = Evaluator(
-            self._target_model, task_value, base_metric, batch_size=batch_size
+            model=self._target_model,
+            task=task_value,
+            metric=base_metric,
+            batch_size=batch_size,
+            use_structured_output=use_structured_output
         )
         final_prompt = ""
-        generator = SyntheticDataGenerator(self._system_model)
+        generator = SyntheticDataGenerator(
+            model=self._system_model,
+            use_structured_output=use_structured_output
+        )
 
         if dataset is None:
             dataset, target, problem_description = generator.generate(
@@ -271,7 +258,7 @@ def run(
             self.synthetic_dataset = dataset
             self.synthetic_target = target
 
-        dataset_split = self._get_dataset_split(
+        dataset_split = get_dataset_split(
             dataset=dataset,
             target=target,
             validation_size=validation_size,
@@ -320,13 +307,14 @@ def run(
             dataset_split=dataset_split,
             evaluator=evaluator,
             problem_description=problem_description,
+            use_structured_output=use_structured_output,
             **kwargs,
         )
 
         logger.info("Running the prompt format checking...")
         final_prompt = correct(
             prompt=final_prompt,
-            rule=LanguageRule(self._system_model),
+            rule=LanguageRule(self._system_model, use_structured_output=use_structured_output),
             start_prompt=start_prompt,
         )
 
@@ -368,6 +356,7 @@ def test(
         metric: Optional[str] = None,
         batch_size: int = 25,
         return_raw_outputs: bool = True,
+        use_structured_output: bool = True
     ) -> List[str] | Tuple[List[str], float]:
         """
         Generate model predictions for a test dataset and optionally compute a metric.
@@ -390,6 +379,8 @@ def test(
             batch_size (int, default=25): Number of samples per inference batch.
             return_raw_outputs (bool, default=True): If True, return raw model outputs;
                 if False, return parsed outputs via metric.parse_output().
+            use_structured_output: a boolean variable.
+                Either to use structured output or nor.
 
         Returns:
             If targets is None: List[str] of raw/parsed outputs.
@@ -404,44 +395,45 @@ def test(
                 "No prompt provided and self.final_prompt is not set. "
                 "Either call .run() first or pass prompt explicitly."
             )
-        
+
         if task is None:
             task_detector = TaskDetector(self._system_model)
             task = task_detector.generate(use_prompt)
-        
+
         task_str = task.lower()
         if task_str not in ("classification", "generation"):
             raise ValueError("task must be 'classification' or 'generation'.")
-        
+
         task_enum = Task.CLASSIFICATION if task_str == "classification" else Task.GENERATION
-        
+
         if metric is None:
             metric = "accuracy" if task_enum == Task.CLASSIFICATION else "meteor"
-        
+
         metric_impl = validate_and_create_metric(task_enum, metric)
-        
+
         evaluator = Evaluator(
             model=self._target_model,
             task=task_enum,
             metric=metric_impl,
             batch_size=batch_size,
+            use_structured_output=use_structured_output
         )
-        
+
         dataset_list = list(dataset)
         use_targets = list(targets) if targets is not None else [""] * len(dataset_list)
-        
+
         result = evaluator.evaluate(
             prompt=use_prompt,
             dataset=dataset_list,
             targets=use_targets,
             template=None,
             return_detailed=True,
         )
-        
+
         outputs = result.raw_outputs if return_raw_outputs else [
             metric_impl.parse_output(a) for a in result.raw_outputs
         ]
-        
+
         if targets is not None:
             return outputs, result.aggregate_score
         return outputs

coolprompt/data_generator/generator.py

@@ -1,15 +1,14 @@
 from typing import Optional, List, Tuple, Any
 from langchain_core.language_models.base import BaseLanguageModel
-from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.messages.ai import AIMessage
 from pydantic import BaseModel
 
-from coolprompt.data_generator.pydantic_formatters import (
-    ProblemDescriptionStructuredOutputSchema,
-    ClassificationTaskStructuredOutputSchema,
+from coolprompt.utils.structured_schemas.data_generator import (
     ClassificationTaskExample,
+    ClassificationTaskResponse,
     GenerationTaskExample,
-    GenerationTaskStructuredOutputSchema,
+    GenerationTaskResponse,
+    ProblemDescriptionResponse,
 )
 from coolprompt.utils.prompt_templates.data_generator_templates import (
     PROBLEM_DESCRIPTION_TEMPLATE,
@@ -25,39 +24,49 @@
 
 
 class SyntheticDataGenerator:
-    """Synthetic Data Generator
-    Generates synthetic dataset for prompt optimization
-    based on given initial prompt and optional problem description
+    """Synthetic Data Generator.
+
+    Generates synthetic datasets for prompt optimization based on a
+    given initial prompt and an optional problem description.
 
     Attributes:
-        model: langchain.BaseLanguageModel class of model to use.
+        model: ``langchain.BaseLanguageModel`` instance to use for LLM calls.
+        use_structured_output: When ``True``, every LLM call routes through
+            ``model.with_structured_output(schema, method="json_schema")``
+            using the Pydantic schemas defined in
+            :mod:`coolprompt.utils.structured_schemas.data_generator`.
+            When ``False`` (default), the generator falls back to a plain
+            ``model.invoke`` call followed by JSON extraction from the
+            raw model response — the same convention used by
+            :mod:`coolprompt.optimizer` submodules.
     """
 
-    def __init__(self, model: BaseLanguageModel) -> None:
+    def __init__(
+        self,
+        model: BaseLanguageModel,
+        use_structured_output: bool = False,
+    ) -> None:
         self.model = model
+        self.use_structured_output = use_structured_output
 
     def _generate(
-            self, request: str, schema: BaseModel, field_name: str
+            self, request: str, schema: type[BaseModel], field_name: str
     ) -> Any:
-        """Generates model output
-        either using structured output from langchain
-        or just strict json output format for LLM
+        """Generates model output using either structured-output via
+        LangChain or a raw JSON parsing fallback.
 
         Args:
-            request (str): request to LLM
-                when langchain structured output is used
-            schema (BaseModel): Pydantic output format
-            field_name (str): field name to select from output
+            request (str): request to send to the LLM.
+            schema (type[BaseModel]): Pydantic schema describing the
+                expected structured output. Only used when
+                ``self.use_structured_output`` is ``True``.
+            field_name (str): top-level field name to extract from the
+                model output.
 
         Returns:
-            Any: generated data
+            Any: extracted value of ``field_name`` from the model output.
         """
-        if hasattr(self.model, "model"):
-            wrapped_model = self.model.model
-        else:
-            wrapped_model = self.model
-
-        if not isinstance(wrapped_model, BaseChatModel):
+        if not self.use_structured_output:
             output = self.model.invoke(request)
             if isinstance(output, AIMessage):
                 output = output.content
@@ -92,13 +101,16 @@ def _examples_to_str(self, examples: List[Tuple[str, str]]) -> str:
     def _generate_problem_description(
             self, prompt: str, examples: Optional[List[Tuple[str, str]]] = None
     ) -> str:
-        """Generates problem description based on given user prompt
+        """Generates problem description based on given user prompt.
 
         Args:
-            prompt (str): initial user prompt
+            prompt (str): initial user prompt.
+            examples (Optional[List[Tuple[str, str]]]): optional list of
+                ``(input, output)`` examples drawn from the task dataset
+                to ground the description.
 
         Returns:
-            str: generated problem description
+            str: generated problem description.
         """
         if examples:
             request = PROBLEM_DESCRIPTION_BASED_ON_EXAMPLES_TEMPLATE.format(
@@ -109,7 +121,7 @@ def _generate_problem_description(
 
         return self._generate(
             request,
-            ProblemDescriptionStructuredOutputSchema,
+            ProblemDescriptionResponse,
             "problem_description",
         )
 
@@ -119,7 +131,7 @@ def _convert_dataset(
                 dict | ClassificationTaskExample | GenerationTaskExample
                 ],
     ) -> Tuple[List[str], List[str]]:
-        """Converts outputs to the dataset format
+        """Converts outputs to the dataset format.
 
         Args:
             examples (
@@ -128,11 +140,11 @@ def _convert_dataset(
                     ClassificationTaskExample |
                     GenerationTaskExample
                 ]
-            ): outputs of the model
+            ): outputs of the model.
 
         Returns:
             Tuple[List[str], List[str]]:
-                converted dataset and target
+                converted dataset and target.
         """
         dataset = []
         targets = []
@@ -156,12 +168,11 @@ def generate(
             num_samples: int = 8,
             corner_ratio: float = 0.4,
     ) -> Tuple[List[str], List[str], str]:
-        """Generates synthetic dataset
-        based on given user prompt, optimization task
-        and optionally provided problem description
+        """Generates synthetic dataset based on the given user prompt,
+        optimization task and optionally provided problem description.
 
-        If problem description isn't provided -
-            it will be generated automatically
+        If problem description isn't provided it will be generated
+        automatically.
 
         Args:
             prompt (str): initial user prompt
@@ -184,7 +195,7 @@ def generate(
 
         Returns:
             Tuple[List[str], List[str], str]:
-                generated dataset, target and problem description
+                generated dataset, target and problem description.
         """
 
         if not 1 <= num_samples <= 100:
@@ -204,11 +215,11 @@ def generate(
         if task == Task.CLASSIFICATION:
             regular_template = CLASSIFICATION_DATA_GENERATING_TEMPLATE
             corner_template = CLASSIFICATION_CORNER_CASE_GENERATING_TEMPLATE
-            schema = ClassificationTaskStructuredOutputSchema
+            schema = ClassificationTaskResponse
         else:
             regular_template = GENERATION_DATA_GENERATING_TEMPLATE
             corner_template = GENERATION_CORNER_CASE_GENERATING_TEMPLATE
-            schema = GenerationTaskStructuredOutputSchema
+            schema = GenerationTaskResponse
 
         n_corner = int(num_samples * corner_ratio)
         n_regular = num_samples - n_corner