sdg-hub 0.2.1__py3-none-any.whl → 0.3.0__py3-none-any.whl

sdg_hub/core/blocks/llm/client_manager.py

@@ -107,9 +107,18 @@ class LLMClientManager:
                 f"Could not validate setup for model '{self.config.model}': {e}"
             )
 
+    def _message_to_dict(self, message: Any) -> dict[str, Any]:
+        """Convert a message to a dict."""
+        if hasattr(message, "to_dict"):
+            return message.to_dict()
+        elif hasattr(message, "__dict__"):
+            return message.__dict__
+        else:
+            return dict(message)
+
     def create_completion(
         self, messages: list[dict[str, Any]], **overrides: Any
-    ) -> Union[str, list[str]]:
+    ) -> Union[dict, list[dict]]:
         """Create a completion using LiteLLM.
 
         Parameters
@@ -121,9 +130,9 @@ class LLMClientManager:
 
         Returns
         -------
-        Union[str, List[str]]
-            The completion text(s). Returns a single string when n=1 or n is None,
-            returns a list of strings when n>1.
+        Union[dict, List[dict]]
+            The completion response(s). Returns a single response when n=1 or n is None,
+            returns a list of responses when n>1. Response dicts contain 'content' and may contain 'reasoning_content'.
 
         Raises
         ------
@@ -151,18 +160,80 @@ class LLMClientManager:
         # Make the completion call
         response = completion_func(kwargs)
 
-        # Extract content from response
+        # Extract message objects from response
         # Check if n > 1 to determine return type
         n_value = final_config.n or 1
         if n_value > 1:
-            return [choice.message.content for choice in response.choices]
+            return [
+                self._message_to_dict(choice.message) for choice in response.choices
+            ]
         else:
-            return response.choices[0].message.content
+            return self._message_to_dict(response.choices[0].message)
 
     async def acreate_completion(
+        self,
+        messages: Union[list[dict[str, Any]], list[list[dict[str, Any]]]],
+        max_concurrency: Optional[int] = None,
+        **overrides: Any,
+    ) -> Union[dict, list[dict]] | list[Union[dict, list[dict]]]:
+        """Create async completion(s) using LiteLLM with optional concurrency control.
+
+        Parameters
+        ----------
+        messages : Union[List[Dict[str, Any]], List[List[Dict[str, Any]]]]
+            Single message list or list of message lists.
+            - For single: List[Dict[str, Any]] - returns Union[Any, List[Any]]
+            - For multiple: List[List[Dict[str, Any]]] - returns List[Union[Any, List[Any]]]
+        max_concurrency : Optional[int], optional
+            Maximum number of concurrent requests when processing multiple messages.
+            If None, all requests run concurrently.
+        **overrides : Any
+            Runtime parameter overrides.
+
+        Returns
+        -------
+        Union[dict, List[dict], List[Union[dict, List[dict]]]]
+            For single message: completion response (dict when n=1, List[dict] when n>1)
+            For multiple messages: list of completion responses (each element can be dict or List[dict])
+
+        Raises
+        ------
+        Exception
+            If the completion fails after all retries.
+        """
+        # Detect if we have single message or multiple messages
+        if not messages:
+            raise ValueError("messages cannot be empty")
+
+        # Check if first element is a dict (single message) or list (multiple messages)
+        if isinstance(messages[0], dict):
+            # Single message case
+            return await self._acreate_single(messages, **overrides)
+        else:
+            # Multiple messages case
+            messages_list = messages
+
+            if max_concurrency is not None:
+                # Use semaphore for concurrency control
+                semaphore = asyncio.Semaphore(max_concurrency)
+
+                async def _create_with_semaphore(msgs):
+                    async with semaphore:
+                        return await self._acreate_single(msgs, **overrides)
+
+                tasks = [_create_with_semaphore(msgs) for msgs in messages_list]
+                return await asyncio.gather(*tasks)
+            else:
+                # No concurrency limit - process all at once
+                tasks = [
+                    self._acreate_single(msgs, **overrides) for msgs in messages_list
+                ]
+                return await asyncio.gather(*tasks)
+
+    async def _acreate_single(
         self, messages: list[dict[str, Any]], **overrides: Any
-    ) -> Union[str, list[str]]:
-        """Create an async completion using LiteLLM.
+    ) -> Union[dict, list[dict]]:
+        """Create a single async completion using LiteLLM.
 
         Parameters
         ----------
@@ -173,10 +244,9 @@ class LLMClientManager:
 
         Returns
         -------
-        Union[str, List[str]]
-            The completion text(s). Returns a single string when n=1 or n is None,
-            returns a list of strings when n>1.
-
+        Union[dict, List[dict]]
+            List of completion message objects. Each element is a dict when n=1 or n is None,
+            or a list of dicts when n>1. Message dicts contain 'content' and may contain 'reasoning_content'.
         Raises
         ------
         Exception
@@ -203,17 +273,19 @@ class LLMClientManager:
         # Make the async completion call
         response = await completion_func(kwargs)
 
-        # Extract content from response
+        # Extract message objects from response
         # Check if n > 1 to determine return type
         n_value = final_config.n or 1
         if n_value > 1:
-            return [choice.message.content for choice in response.choices]
+            return [
+                self._message_to_dict(choice.message) for choice in response.choices
+            ]
         else:
-            return response.choices[0].message.content
+            return self._message_to_dict(response.choices[0].message)
 
     def create_completions_batch(
         self, messages_list: list[list[dict[str, Any]]], **overrides: Any
-    ) -> list[Union[str, list[str]]]:
+    ) -> list[Union[dict, list[dict]]]:
         """Create multiple completions in batch.
 
         Parameters
@@ -225,9 +297,9 @@ class LLMClientManager:
 
         Returns
         -------
-        List[Union[str, List[str]]]
-            List of completion texts. Each element is a single string when n=1 or n is None,
-            or a list of strings when n>1.
+        List[dict] | List[List[dict]]
+            List of completion responses. Each element is a dict when n=1 or n is None,
+            or a list of dicts when n>1. Response dicts contain 'content' and may contain 'reasoning_content'.
         """
         results = []
         for messages in messages_list:
@@ -235,29 +307,6 @@ class LLMClientManager:
             results.append(result)
         return results
 
-    async def acreate_completions_batch(
-        self, messages_list: list[list[dict[str, Any]]], **overrides: Any
-    ) -> list[Union[str, list[str]]]:
-        """Create multiple completions in batch asynchronously.
-
-        Parameters
-        ----------
-        messages_list : List[List[Dict[str, Any]]]
-            List of message lists to process.
-        **overrides : Any
-            Runtime parameter overrides.
-
-        Returns
-        -------
-        List[Union[str, List[str]]]
-            List of completion texts. Each element is a single string when n=1 or n is None,
-            or a list of strings when n>1.
-        """
-        tasks = [
-            self.acreate_completion(messages, **overrides) for messages in messages_list
-        ]
-        return await asyncio.gather(*tasks)
-
     def _build_completion_kwargs(
         self, messages: list[dict[str, Any]], config: LLMConfig
     ) -> dict[str, Any]:

sdg_hub/core/blocks/llm/config.py

@@ -240,6 +240,7 @@ class LLMConfig:
             "logprobs",
             "top_logprobs",
             "user",
+            "timeout",
         ]:
             value = getattr(self, param)
             if value is not None:

sdg_hub/core/blocks/llm/llm_chat_block.py

@@ -42,9 +42,10 @@ class LLMChatBlock(BaseBlock):
         Name of the block.
     input_cols : Union[str, List[str]]
         Input column name(s). Should contain the messages list.
-    output_cols : Union[str, List[str]]
+    output_cols : Union[dict, List[dict]]
         Output column name(s) for the response. When n > 1, the column will contain
-        a list of responses instead of a single string.
+        a list of responses instead of a single response. Responses contain 'content',
+        may contain 'reasoning_content' and other fields if any.
     model : str
         Model identifier in LiteLLM format. Examples:
         - "openai/gpt-4"
@@ -131,7 +132,7 @@ class LLMChatBlock(BaseBlock):
     >>> block = LLMChatBlock(
     ...     block_name="gpt4_multiple",
     ...     input_cols="messages",
-    ...     output_cols="responses",  # Will contain lists of strings
+    ...     output_cols="responses",  # Will contain lists of responses
     ...     model="openai/gpt-4",
     ...     n=3,  # Generate 3 responses per input
     ...     temperature=0.8
@@ -297,6 +298,10 @@ class LLMChatBlock(BaseBlock):
             temperature, max_tokens, top_p, frequency_penalty, presence_penalty,
             stop, seed, response_format, stream, n, and provider-specific params.
 
+            Special flow-level parameters:
+            _flow_max_concurrency : int, optional
+                Maximum concurrency for async requests (passed by Flow).
+
         Returns
         -------
         Dataset
@@ -314,27 +319,73 @@ class LLMChatBlock(BaseBlock):
                 f"Call flow.set_model_config() before generating."
             )
 
+        # Extract max_concurrency if provided by flow
+        flow_max_concurrency = override_kwargs.pop("_flow_max_concurrency", None)
+
         # Extract messages
         messages_list = samples[self.input_cols[0]]
 
         # Log generation start
         logger.info(
-            f"Starting {'async' if self.async_mode else 'sync'} generation for {len(messages_list)} samples",
+            f"Starting {'async' if self.async_mode else 'sync'} generation for {len(messages_list)} samples"
+            + (
+                f" (max_concurrency={flow_max_concurrency})"
+                if flow_max_concurrency
+                else ""
+            ),
             extra={
                 "block_name": self.block_name,
                 "model": self.model,
                 "provider": self.client_manager.config.get_provider(),
                 "batch_size": len(messages_list),
                 "async_mode": self.async_mode,
-                "override_params": override_kwargs,
+                "flow_max_concurrency": flow_max_concurrency,
+                "override_params": {
+                    k: (
+                        "***"
+                        if any(
+                            s in k.lower()
+                            for s in ["key", "token", "secret", "authorization"]
+                        )
+                        else v
+                    )
+                    for k, v in override_kwargs.items()
+                },
             },
         )
 
         # Generate responses
         if self.async_mode:
-            responses = asyncio.run(
-                self._generate_async(messages_list, **override_kwargs)
-            )
+            try:
+                # Check if there's already a running event loop
+                loop = asyncio.get_running_loop()
+                # Check if nest_asyncio is applied (allows nested asyncio.run)
+                # Use multiple detection methods for robustness
+                nest_asyncio_applied = (
+                    hasattr(loop, "_nest_patched")
+                    or getattr(asyncio.run, "__module__", "") == "nest_asyncio"
+                )
+
+                if nest_asyncio_applied:
+                    # nest_asyncio is applied, safe to use asyncio.run
+                    responses = asyncio.run(
+                        self._generate_async(
+                            messages_list, flow_max_concurrency, **override_kwargs
+                        )
+                    )
+                else:
+                    # Running inside an event loop without nest_asyncio
+                    raise BlockValidationError(
+                        f"async_mode=True cannot be used from within a running event loop for '{self.block_name}'. "
+                        "Use an async entrypoint, set async_mode=False, or apply nest_asyncio.apply() in notebook environments."
+                    )
+            except RuntimeError:
+                # No running loop; safe to create one
+                responses = asyncio.run(
+                    self._generate_async(
+                        messages_list, flow_max_concurrency, **override_kwargs
+                    )
+                )
         else:
             responses = self._generate_sync(messages_list, **override_kwargs)
 
@@ -356,7 +407,7 @@ class LLMChatBlock(BaseBlock):
         self,
         messages_list: list[list[dict[str, Any]]],
         **override_kwargs: dict[str, Any],
-    ) -> list[Union[str, list[str]]]:
+    ) -> list[Union[dict, list[dict]]]:
         """Generate responses synchronously.
 
         Parameters
@@ -368,8 +419,9 @@ class LLMChatBlock(BaseBlock):
 
         Returns
         -------
-        List[Union[str, List[str]]]
-            List of response strings or lists of response strings (when n > 1).
+        List[Union[dict, List[dict]]]
+            List of responses. Each element is a dict when n=1 or n is None,
+            or a list of dicts when n>1. Response dicts contain 'content', may contain 'reasoning_content' and other fields if any.
         """
         responses = []
 
@@ -409,26 +461,32 @@ class LLMChatBlock(BaseBlock):
     async def _generate_async(
         self,
         messages_list: list[list[dict[str, Any]]],
+        flow_max_concurrency: Optional[int] = None,
         **override_kwargs: dict[str, Any],
-    ) -> list[Union[str, list[str]]]:
+    ) -> list[Union[dict, list[dict]]]:
         """Generate responses asynchronously.
 
         Parameters
         ----------
         messages_list : List[List[Dict[str, Any]]]
             List of message lists to process.
+        flow_max_concurrency : Optional[int], optional
+            Maximum concurrency for async requests.
         **override_kwargs : Dict[str, Any]
             Runtime parameter overrides.
 
         Returns
         -------
-        List[Union[str, List[str]]]
-            List of response strings or lists of response strings (when n > 1).
+        List[Union[dict, List[dict]]]
+            List of responses. Each element is a dict when n=1 or n is None,
+            or a list of dicts when n>1. Response dicts contain 'content', may contain 'reasoning_content' and other fields if any.
         """
         try:
-            responses = await self.client_manager.acreate_completions_batch(
-                messages_list, **override_kwargs
+            # Use unified client manager method with optional concurrency control
+            responses = await self.client_manager.acreate_completion(
+                messages_list, max_concurrency=flow_max_concurrency, **override_kwargs
             )
+
             return responses
 
         except Exception as e: