Add LiteLLM+instructor (for structured output) backend for curator

examples/litellm_recipe_prompting.py

@@ -0,0 +1,45 @@
+from typing import List
+from pydantic import BaseModel, Field
+from bespokelabs import curator
+from datasets import Dataset
+
+
+def main():
+    # List of cuisines to generate recipes for
+    cuisines = [
+        {"cuisine": cuisine}
+        for cuisine in [
+            "Chinese",
+            "Italian",
+            "Mexican",
+            "French",
+            "Japanese",
+            "Indian",
+            "Thai",
+            "Korean",
+            "Vietnamese",
+            "Brazilian",
+        ]
+    ]
+    cuisines = Dataset.from_list(cuisines)
+
+    # Create prompter using LiteLLM backend
+    recipe_prompter = curator.Prompter(
+        model_name="gpt-4o-mini",
+        prompt_func=lambda row: f"Generate a random {row['cuisine']} recipe. Be creative but keep it realistic.",
+        parse_func=lambda row, response: {
+            "recipe": response,
+            "cuisine": row["cuisine"],
+        },
+        backend="litellm",
+    )
+
+    # Generate recipes for all cuisines
+    recipes = recipe_prompter(cuisines)
+
+    # Print results
+    print(recipes.to_pandas())
+
+
+if __name__ == "__main__":
+    main()

examples/litellm_recipe_structured_output.py

@@ -0,0 +1,60 @@
+from typing import List
+from pydantic import BaseModel, Field
+from bespokelabs import curator
+from datasets import Dataset
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+# Define response format using Pydantic
+class Recipe(BaseModel):
+    title: str = Field(description="Title of the recipe")
+    ingredients: List[str] = Field(description="List of ingredients needed")
+    instructions: List[str] = Field(description="Step by step cooking instructions")
+    prep_time: int = Field(description="Preparation time in minutes")
+    cook_time: int = Field(description="Cooking time in minutes")
+    servings: int = Field(description="Number of servings")
+
+
+class Cuisines(BaseModel):
+    cuisines_list: List[str] = Field(description="A list of cuisines.")
+
+
+def main():
+    # We define a prompter that generates cuisines
+    cuisines_generator = curator.Prompter(
+        prompt_func=lambda: f"Generate 10 diverse cuisines.",
+        model_name="gpt-4o-mini",
+        response_format=Cuisines,
+        parse_func=lambda _, cuisines: [{"cuisine": t} for t in cuisines.cuisines_list],
+        backend="litellm",
+    )
+    cuisines = cuisines_generator()
+    print(cuisines.to_pandas())
+
+    recipe_prompter = curator.Prompter(
+        model_name="claude-3-5-haiku-20241022",
+        prompt_func=lambda row: f"Generate a random {row['cuisine']} recipe. Be creative but keep it realistic.",
+        parse_func=lambda row, response: {
+            "title": response.title,
+            "ingredients": response.ingredients,
+            "instructions": response.instructions,
+            "prep_time": response.prep_time,
+            "cook_time": response.cook_time,
+            "servings": response.servings,
+            "cuisine": row["cuisine"],
+        },
+        response_format=Recipe,
+        backend="litellm",
+    )
+
+    # Generate recipes for all cuisines
+    recipes = recipe_prompter(cuisines)
+
+    # Print results
+    print(recipes.to_pandas())
+
+
+if __name__ == "__main__":
+    main()

src/bespokelabs/curator/prompter/prompter.py

@@ -21,6 +21,9 @@
 from bespokelabs.curator.request_processor.openai_online_request_processor import (
     OpenAIOnlineRequestProcessor,
 )
+from bespokelabs.curator.request_processor.litellm_online_request_processor import (
+    LiteLLMOnlineRequestProcessor,
+)
 
 _CURATOR_DEFAULT_CACHE_DIR = "~/.cache/curator"
 T = TypeVar("T")
@@ -45,6 +48,7 @@ def __init__(
             ]
         ] = None,
         response_format: Optional[Type[BaseModel]] = None,
+        backend: str = "openai",
         batch: bool = False,
         batch_size: Optional[int] = None,
         temperature: Optional[float] = None,
@@ -62,6 +66,7 @@ def __init__(
                 response object and returns the parsed output
             response_format (Optional[Type[BaseModel]]): A Pydantic model specifying the
                 response format from the LLM.
+            backend (str): The backend to use ("openai" or "litellm")
             batch (bool): Whether to use batch processing
             batch_size (Optional[int]): The size of the batch to use, only used if batch is True
             temperature (Optional[float]): The temperature to use for the LLM, only used if batch is False
@@ -86,32 +91,49 @@ def __init__(
             model_name, prompt_func, parse_func, response_format
         )
         self.batch_mode = batch
-        if batch:
-            if batch_size is None:
-                batch_size = 1_000
-                logger.info(
-                    f"batch=True but no batch_size provided, using default batch_size of {batch_size:,}"
+        self.backend = backend
+        # Select request processor based on backend
+        if backend == "openai":
+            if batch:
+                if batch_size is None:
+                    batch_size = 1_000
+                    logger.info(
+                        f"batch=True but no batch_size provided, using default batch_size of {batch_size:,}"
+                    )
+                self._request_processor = OpenAIBatchRequestProcessor(
+                    model=model_name,
+                    batch_size=batch_size,
+                    temperature=temperature,
+                    top_p=top_p,
+                    presence_penalty=presence_penalty,
+                    frequency_penalty=frequency_penalty,
                 )
-            self._request_processor = OpenAIBatchRequestProcessor(
-                model=model_name,
-                batch_size=batch_size,
-                temperature=temperature,
-                top_p=top_p,
-                presence_penalty=presence_penalty,
-                frequency_penalty=frequency_penalty,
-            )
-        else:
-            if batch_size is not None:
+            else:
+                if batch_size is not None:
+                    logger.warning(
+                        f"Prompter argument `batch_size` {batch_size} is ignored because `batch` is False"
+                    )
+                self._request_processor = OpenAIOnlineRequestProcessor(
+                    model=model_name,
+                    temperature=temperature,
+                    top_p=top_p,
+                    presence_penalty=presence_penalty,
+                    frequency_penalty=frequency_penalty,
+                )
+        elif backend == "litellm":
+            if batch:
                 logger.warning(
-                    f"Prompter argument `batch_size` {batch_size} is ignored because `batch` is False"
+                    "Batch mode is not supported with LiteLLM backend, ignoring batch=True"
                 )
-            self._request_processor = OpenAIOnlineRequestProcessor(
+            self._request_processor = LiteLLMOnlineRequestProcessor(
                 model=model_name,
                 temperature=temperature,
                 top_p=top_p,
                 presence_penalty=presence_penalty,
                 frequency_penalty=frequency_penalty,
             )
+        else:
+            raise ValueError(f"Unknown backend: {backend}")
 
     def __call__(self, dataset: Optional[Iterable] = None, working_dir: str = None) -> Dataset:
         """
@@ -176,6 +198,7 @@ def _completions(
                     else "text"
                 ),
                 str(self.batch_mode),
+                str(self.backend),
             ]
         )
 

src/bespokelabs/curator/request_processor/base_request_processor.py

@@ -239,16 +239,32 @@ def create_dataset_files(
                             # Response message is a string, which is converted to a dict
                             # The dict is then used to construct the response_format Pydantic model
                             try:
+                                # First try to parse the response message as JSON
+                                if isinstance(response.response_message, str):
+                                    try:
+                                        response_dict = json.loads(response.response_message)
+                                    except json.JSONDecodeError as e:
+                                        warning_msg = (
+                                            f"Failed to parse response message as JSON: {response.response_message}. "
+                                            f"The model likely returned an invalid JSON format. Will skip this response."
+                                        )
+                                        logger.warning(warning_msg)
+                                        failed_responses_count += 1
+                                        continue
+                                else:
+                                    response_dict = response.response_message
+
+                                # Then construct the Pydantic model from the parsed dict
                                 response.response_message = prompt_formatter.response_format(
-                                    **response.response_message
+                                    **response_dict
                                 )
                             except ValidationError as e:
                                 schema_str = json.dumps(
                                     prompt_formatter.response_format.model_json_schema(),
                                     indent=2,
                                 )
                                 warning_msg = (
-                                    f"Pydantic failed to parse response message {response.response_message} with `response_format` {schema_str}."
+                                    f"Pydantic failed to parse response message {response.response_message} with `response_format` {schema_str}. "
                                     f"The model likely returned a JSON that does not match the schema of the `response_format`. Will skip this response."
                                 )
                                 logger.warning(warning_msg)