feat: add support for n samples in generation params

src/bespokelabs/curator/llm/prompt_formatter.py

@@ -84,7 +84,6 @@ def create_generic_request(self, row: _DictOrBaseModel, idx: int) -> GenericRequ
         if isinstance(prompts, str):
             messages = [{"role": "user", "content": prompts}]
         elif isinstance(prompts, list):
-            multimodal_prompt = False
             _validate_messages(prompts)
             messages = prompts
         elif isinstance(prompts, tuple):
@@ -133,17 +132,24 @@ def response_to_response_format(self, response_message: str | dict) -> Optional[
 
         try:
             # First try to parse the response message as JSON
-            if isinstance(response_message, str):
+            def _load_response_message(response_message: str | dict) -> dict:
                 try:
                     response_dict = json.loads(response_message)
                 except json.JSONDecodeError as e:
                     logger.warning(f"Failed to parse response message as JSON: {response_message}. The model likely returned an invalid JSON format.")
                     raise e
+                return response_dict
+
+            if isinstance(response_message, str):
+                response_dict = _load_response_message(response_message)
+                response_message = self.response_format(**response_dict)
+            elif isinstance(response_message, list):
+                response_dicts = list(map(_load_response_message, response_message))
+                response_message = list(map(lambda r: self.response_format(**r), response_dicts))  # noqa: C417
             else:
                 response_dict = response_message
+                response_message = self.response_format(**response_dict)
 
-            # Then construct the Pydantic model from the parsed dict
-            response_message = self.response_format(**response_dict)
             return response_message
 
         except ValidationError as e:

src/bespokelabs/curator/request_processor/online/litellm_online_request_processor.py

@@ -346,6 +346,8 @@ async def call_single_request(
                     raise Exception("Response is empty")
                 if self.config.return_completions_object:
                     response_message = dict(completion_obj)
+                elif self.config.generation_params.get("n", 1) > 1:
+                    response_message = [choice["message"]["content"] for choice in completion_obj["choices"]]
                 else:
                     response_message = completion_obj["choices"][0]["message"]["content"]
         except litellm.RateLimitError as e:

src/bespokelabs/curator/request_processor/online/openai_online_request_processor.py

@@ -237,11 +237,14 @@ async def call_single_request(
 
             if response_obj.status != 200:
                 raise Exception(f"API request failed with status {response_obj.status}: {response}")
-
             if self.config.return_completions_object:
                 response_message = dict(response)
+            elif self.config.generation_params.get("n", 1) > 1:
+                response_message = [choice["message"]["content"] for choice in response["choices"]]
             else:
                 response_message = response["choices"][0]["message"]["content"]
+
+            # TODO: check if this is the correct way to get finish_reason with `n` > 1
             finish_reason = response["choices"][0].get("finish_reason", "unkown")
             usage = response["usage"]
             token_usage = TokenUsage(

src/bespokelabs/curator/types/generic_response.py

@@ -28,7 +28,7 @@
 class GenericResponse(BaseModel):
     """A generic response model for LLM API requests."""
 
-    response_message: Optional[Dict[str, Any]] | str = None
+    response_message: Optional[Dict[str, Any] | str | List] = None
     response_errors: Optional[List[str]] = None
     raw_response: Optional[Dict[str, Any]]
     raw_request: Optional[Dict[str, Any]] = None