sdg-hub 0.1.1__py3-none-any.whl → 0.1.2__py3-none-any.whl

sdg_hub/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.1.1'
-__version_tuple__ = version_tuple = (0, 1, 1)
+__version__ = version = '0.1.2'
+__version_tuple__ = version_tuple = (0, 1, 2)

sdg_hub/blocks/__init__.py

@@ -6,6 +6,10 @@ This package provides various block implementations for data generation, process
 # Local
 from .block import Block
 from .llmblock import LLMBlock, ConditionalLLMBlock
+from .openaichatblock import (
+    OpenAIChatBlock,
+    OpenAIAsyncChatBlock
+)
 from .utilblocks import (
     SamplePopulatorBlock,
     SelectorBlock,
@@ -33,4 +37,6 @@ __all__ = [
     "RenameColumns",
     "SetToMajorityValue",
     "BlockRegistry",
+    "OpenAIChatBlock",
+    "OpenAIAsyncChatBlock"
 ]

sdg_hub/blocks/openaichatblock.py

@@ -0,0 +1,556 @@
+# SPDX-License-Identifier: Apache-2.0
+"""OpenAI-specific blocks for text generation.
+
+This module provides blocks for interacting with OpenAI's Chat Completions API.
+"""
+
+# Standard
+from typing import Any, Dict, List, Optional, Union
+import asyncio
+
+# Third Party
+from datasets import Dataset
+from tenacity import (
+    retry,
+    retry_if_exception_type,
+    stop_after_attempt,
+    wait_random_exponential,
+)
+import openai
+
+# Local
+from ..logger_config import setup_logger
+from ..registry import BlockRegistry
+from .block import Block
+
+logger = setup_logger(__name__)
+
+
+@BlockRegistry.register("OpenAIChatBlock")
+class OpenAIChatBlock(Block):
+    """Block for generating text using OpenAI Chat Completions API.
+
+    This block takes a column containing OpenAI message format and makes
+    direct calls to the chat completions endpoint.
+
+    Parameters
+    ----------
+    block_name : str
+        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 column name(s) for the response.
+    client : openai.OpenAI
+        OpenAI client instance.
+    model_id : str
+        Model ID to use.
+
+    ### Text-relevant OpenAI Chat Completions API parameters ###
+
+    frequency_penalty : Optional[float], optional
+        Penalize frequent tokens (-2.0 to 2.0).
+    logit_bias : Optional[Dict[str, int]], optional
+        Modify likelihood of specified tokens.
+    logprobs : Optional[bool], optional
+        Whether to return log probabilities.
+    max_completion_tokens : Optional[int], optional
+        Maximum tokens in completion.
+    max_tokens : Optional[int], optional
+        Maximum tokens in completion (legacy).
+    n : Optional[int], optional
+        Number of completions to generate.
+    presence_penalty : Optional[float], optional
+        Penalize repeated tokens (-2.0 to 2.0).
+    response_format : Optional[Dict[str, Any]], optional
+        Response format specification (e.g., JSON mode).
+    seed : Optional[int], optional
+        Seed for deterministic outputs.
+    stop : Optional[Union[str, List[str]]], optional
+        Stop sequences.
+    stream : Optional[bool], optional
+        Whether to stream responses.
+    temperature : Optional[float], optional
+        Sampling temperature (0.0 to 2.0).
+    tool_choice : Optional[Union[str, Dict[str, Any]]], optional
+        Tool selection strategy.
+    tools : Optional[List[Dict[str, Any]]], optional
+        Available tools for function calling.
+    top_logprobs : Optional[int], optional
+        Number of top log probabilities to return.
+    top_p : Optional[float], optional
+        Nucleus sampling parameter (0.0 to 1.0).
+    user : Optional[str], optional
+        End-user identifier.
+    extra_body : Optional[dict], optional
+        Dictionary of additional parameters if supported by inference backend
+    """
+
+    def __init__(
+        self,
+        block_name: str,
+        input_cols: Union[str, List[str]],
+        output_cols: Union[str, List[str]],
+        client: openai.OpenAI,
+        model_id: str,
+        # Text-relevant OpenAI Chat Completions API parameters
+        frequency_penalty: Optional[float] = None,
+        logit_bias: Optional[Dict[str, int]] = None,
+        logprobs: Optional[bool] = None,
+        max_completion_tokens: Optional[int] = None,
+        max_tokens: Optional[int] = None,
+        n: Optional[int] = None,
+        presence_penalty: Optional[float] = None,
+        response_format: Optional[Dict[str, Any]] = None,
+        seed: Optional[int] = None,
+        stop: Optional[Union[str, List[str]]] = None,
+        stream: Optional[bool] = None,
+        temperature: Optional[float] = None,
+        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        top_logprobs: Optional[int] = None,
+        top_p: Optional[float] = None,
+        user: Optional[str] = None,
+        extra_body: Optional[dict] = None,
+    ) -> None:
+        super().__init__(block_name)
+        self.input_cols = [input_cols] if isinstance(input_cols, str) else input_cols
+        self.output_cols = (
+            [output_cols] if isinstance(output_cols, str) else output_cols
+        )
+        self.client = client
+        self.model_id = model_id
+
+        # For this block, we expect exactly one input column (messages) and one output column
+        if len(self.input_cols) != 1:
+            raise ValueError("OpenAIChatBlock expects exactly one input column")
+        if len(self.output_cols) != 1:
+            raise ValueError("OpenAIChatBlock expects exactly one output column")
+
+        self.messages_column = self.input_cols[0]
+        self.output_column = self.output_cols[0]
+
+        # Store all generation parameters (only non-None values)
+        self.gen_kwargs = {}
+        params = {
+            "frequency_penalty": frequency_penalty,
+            "logit_bias": logit_bias,
+            "logprobs": logprobs,
+            "max_completion_tokens": max_completion_tokens,
+            "max_tokens": max_tokens,
+            "n": n,
+            "presence_penalty": presence_penalty,
+            "response_format": response_format,
+            "seed": seed,
+            "stop": stop,
+            "stream": stream,
+            "temperature": temperature,
+            "tool_choice": tool_choice,
+            "tools": tools,
+            "top_logprobs": top_logprobs,
+            "top_p": top_p,
+            "user": user,
+            "extra_body": extra_body,
+        }
+
+        # Only include non-None parameters
+        for key, value in params.items():
+            if value is not None:
+                self.gen_kwargs[key] = value
+
+        # Log initialization with model and parameters
+        logger.info(
+            f"Initialized OpenAIChatBlock '{block_name}' with model '{model_id}'",
+            extra={
+                "block_name": block_name,
+                "model_id": model_id,
+                "generation_params": self.gen_kwargs,
+            },
+        )
+
+    @retry(
+        wait=wait_random_exponential(min=1, max=60),
+        stop=stop_after_attempt(6),
+        retry=retry_if_exception_type(
+            (openai.RateLimitError, openai.APITimeoutError, openai.APIConnectionError)
+        ),
+    )
+    def _create_completion_with_retry(self, **kwargs):
+        """Create completion with retry logic."""
+        return self.client.chat.completions.create(**kwargs)
+
+    def generate(self, samples: Dataset, **override_kwargs: Dict[str, Any]) -> Dataset:
+        """Generate the output from the block.
+
+        Parameters set at runtime override those set during initialization.
+        Supports all text-relevant OpenAI Chat Completions API parameters.
+
+        Parameters
+        ----------
+        samples : Dataset
+            Input dataset containing the messages column.
+        **override_kwargs : Dict[str, Any]
+            Runtime parameters that override initialization defaults.
+            Valid parameters: frequency_penalty, logit_bias, logprobs,
+            max_completion_tokens, max_tokens, n, presence_penalty,
+            response_format, seed, stop, stream, temperature, tool_choice,
+            tools, top_logprobs, top_p, user.
+
+        Returns
+        -------
+        Dataset
+            Dataset with the response added to the output column.
+        """
+        # Define valid parameters for validation
+        valid_params = {
+            "frequency_penalty",
+            "logit_bias",
+            "logprobs",
+            "max_completion_tokens",
+            "max_tokens",
+            "n",
+            "presence_penalty",
+            "response_format",
+            "seed",
+            "stop",
+            "stream",
+            "temperature",
+            "tool_choice",
+            "tools",
+            "top_logprobs",
+            "top_p",
+            "user",
+            "extra_body",
+        }
+
+        # Filter and validate override parameters
+        filtered_kwargs = {
+            k: v for k, v in override_kwargs.items() if k in valid_params
+        }
+
+        # Warn about invalid parameters
+        invalid_params = set(override_kwargs.keys()) - valid_params
+        if invalid_params:
+            logger.warning(f"Ignoring invalid parameters: {invalid_params}")
+
+        # Merge kwargs with priority: runtime > init > defaults
+        final_kwargs = {**self.gen_kwargs, **filtered_kwargs}
+        final_kwargs["model"] = self.model_id
+
+        # Extract all messages
+        messages_list = samples[self.messages_column]
+
+        # Log generation start with model and effective parameters
+        logger.info(
+            f"Starting generation for {len(messages_list)} samples",
+            extra={
+                "block_name": self.block_name,
+                "model_id": self.model_id,
+                "batch_size": len(messages_list),
+                "effective_params": {
+                    k: v
+                    for k, v in final_kwargs.items()
+                    if k
+                    in [
+                        "temperature",
+                        "max_tokens",
+                        "max_completion_tokens",
+                        "top_p",
+                        "n",
+                        "seed",
+                    ]
+                },
+            },
+        )
+
+        # Get all responses
+        responses = []
+        for messages in messages_list:
+            response = self._create_completion_with_retry(
+                messages=messages, **final_kwargs
+            )
+            responses.append(response.choices[0].message.content)
+
+        # Log completion
+        logger.info(
+            f"Generation completed successfully for {len(responses)} samples",
+            extra={
+                "block_name": self.block_name,
+                "model_id": self.model_id,
+                "batch_size": len(responses),
+            },
+        )
+
+        # Add responses as new column
+        return samples.add_column(self.output_column, responses)
+
+
+@BlockRegistry.register("OpenAIAsyncChatBlock")
+class OpenAIAsyncChatBlock(Block):
+    """Async block for generating text using OpenAI Chat Completions API.
+
+    This block takes a column containing OpenAI message format and makes
+    asynchronous calls to the chat completions endpoint for better performance.
+
+    Parameters
+    ----------
+    block_name : str
+        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 column name(s) for the response.
+    async_client : openai.AsyncOpenAI
+        Async OpenAI client instance.
+    model_id : str
+        Model ID to use.
+
+    ### Text-relevant OpenAI Chat Completions API parameters ###
+
+    frequency_penalty : Optional[float], optional
+        Penalize frequent tokens (-2.0 to 2.0).
+    logit_bias : Optional[Dict[str, int]], optional
+        Modify likelihood of specified tokens.
+    logprobs : Optional[bool], optional
+        Whether to return log probabilities.
+    max_completion_tokens : Optional[int], optional
+        Maximum tokens in completion.
+    max_tokens : Optional[int], optional
+        Maximum tokens in completion (legacy).
+    n : Optional[int], optional
+        Number of completions to generate.
+    presence_penalty : Optional[float], optional
+        Penalize repeated tokens (-2.0 to 2.0).
+    response_format : Optional[Dict[str, Any]], optional
+        Response format specification (e.g., JSON mode).
+    seed : Optional[int], optional
+        Seed for deterministic outputs.
+    stop : Optional[Union[str, List[str]]], optional
+        Stop sequences.
+    stream : Optional[bool], optional
+        Whether to stream responses.
+    temperature : Optional[float], optional
+        Sampling temperature (0.0 to 2.0).
+    tool_choice : Optional[Union[str, Dict[str, Any]]], optional
+        Tool selection strategy.
+    tools : Optional[List[Dict[str, Any]]], optional
+        Available tools for function calling.
+    top_logprobs : Optional[int], optional
+        Number of top log probabilities to return.
+    top_p : Optional[float], optional
+        Nucleus sampling parameter (0.0 to 1.0).
+    user : Optional[str], optional
+        End-user identifier.
+    extra_body : Optional[dict], optional
+        Dictionary of additional parameters if supported by inference backend
+    """
+
+    def __init__(
+        self,
+        block_name: str,
+        input_cols: Union[str, List[str]],
+        output_cols: Union[str, List[str]],
+        async_client: openai.AsyncOpenAI,
+        model_id: str,
+        # Text-relevant OpenAI Chat Completions API parameters
+        frequency_penalty: Optional[float] = None,
+        logit_bias: Optional[Dict[str, int]] = None,
+        logprobs: Optional[bool] = None,
+        max_completion_tokens: Optional[int] = None,
+        max_tokens: Optional[int] = None,
+        n: Optional[int] = None,
+        presence_penalty: Optional[float] = None,
+        response_format: Optional[Dict[str, Any]] = None,
+        seed: Optional[int] = None,
+        stop: Optional[Union[str, List[str]]] = None,
+        stream: Optional[bool] = None,
+        temperature: Optional[float] = None,
+        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        top_logprobs: Optional[int] = None,
+        top_p: Optional[float] = None,
+        user: Optional[str] = None,
+        extra_body: Optional[dict] = None,
+    ) -> None:
+        super().__init__(block_name)
+        self.input_cols = [input_cols] if isinstance(input_cols, str) else input_cols
+        self.output_cols = (
+            [output_cols] if isinstance(output_cols, str) else output_cols
+        )
+        self.async_client = async_client
+        self.model_id = model_id
+
+        # For this block, we expect exactly one input column (messages) and one output column
+        if len(self.input_cols) != 1:
+            raise ValueError("OpenAIAsyncChatBlock expects exactly one input column")
+        if len(self.output_cols) != 1:
+            raise ValueError("OpenAIAsyncChatBlock expects exactly one output column")
+
+        self.messages_column = self.input_cols[0]
+        self.output_column = self.output_cols[0]
+
+        # Store all generation parameters (only non-None values)
+        self.gen_kwargs = {}
+        params = {
+            "frequency_penalty": frequency_penalty,
+            "logit_bias": logit_bias,
+            "logprobs": logprobs,
+            "max_completion_tokens": max_completion_tokens,
+            "max_tokens": max_tokens,
+            "n": n,
+            "presence_penalty": presence_penalty,
+            "response_format": response_format,
+            "seed": seed,
+            "stop": stop,
+            "stream": stream,
+            "temperature": temperature,
+            "tool_choice": tool_choice,
+            "tools": tools,
+            "top_logprobs": top_logprobs,
+            "top_p": top_p,
+            "user": user,
+            "extra_body": extra_body,
+        }
+
+        # Only include non-None parameters
+        for key, value in params.items():
+            if value is not None:
+                self.gen_kwargs[key] = value
+
+        # Log initialization with model and parameters
+        logger.info(
+            f"Initialized OpenAIAsyncChatBlock '{block_name}' with model '{model_id}'",
+            extra={
+                "block_name": block_name,
+                "model_id": model_id,
+                "generation_params": self.gen_kwargs,
+            },
+        )
+
+    @retry(
+        wait=wait_random_exponential(min=1, max=60),
+        stop=stop_after_attempt(6),
+        retry=retry_if_exception_type(
+            (openai.RateLimitError, openai.APITimeoutError, openai.APIConnectionError)
+        ),
+    )
+    async def _generate_single(
+        self, messages: List[Dict[str, Any]], **final_kwargs: Dict[str, Any]
+    ) -> str:
+        """Generate a single response asynchronously."""
+        response = await self.async_client.chat.completions.create(
+            messages=messages, **final_kwargs
+        )
+        return response.choices[0].message.content
+
+    def generate(self, samples: Dataset, **override_kwargs: Dict[str, Any]) -> Dataset:
+        """Generate the output from the block using async calls.
+
+        Parameters set at runtime override those set during initialization.
+        Supports all text-relevant OpenAI Chat Completions API parameters.
+
+        Parameters
+        ----------
+        samples : Dataset
+            Input dataset containing the messages column.
+        **override_kwargs : Dict[str, Any]
+            Runtime parameters that override initialization defaults.
+            Valid parameters: frequency_penalty, logit_bias, logprobs,
+            max_completion_tokens, max_tokens, n, presence_penalty,
+            response_format, seed, stop, stream, temperature, tool_choice,
+            tools, top_logprobs, top_p, user.
+
+        Returns
+        -------
+        Dataset
+            Dataset with the response added to the output column.
+        """
+        # Define valid parameters for validation
+        valid_params = {
+            "frequency_penalty",
+            "logit_bias",
+            "logprobs",
+            "max_completion_tokens",
+            "max_tokens",
+            "n",
+            "presence_penalty",
+            "response_format",
+            "seed",
+            "stop",
+            "stream",
+            "temperature",
+            "tool_choice",
+            "tools",
+            "top_logprobs",
+            "top_p",
+            "user",
+        }
+
+        # Filter and validate override parameters
+        filtered_kwargs = {
+            k: v for k, v in override_kwargs.items() if k in valid_params
+        }
+
+        # Warn about invalid parameters
+        invalid_params = set(override_kwargs.keys()) - valid_params
+        if invalid_params:
+            logger.warning(f"Ignoring invalid parameters: {invalid_params}")
+
+        # Merge kwargs with priority: runtime > init > defaults
+        final_kwargs = {**self.gen_kwargs, **filtered_kwargs}
+        final_kwargs["model"] = self.model_id
+
+        # Log generation start with model and effective parameters
+        logger.info(
+            f"Starting async generation for {len(samples)} samples",
+            extra={
+                "block_name": self.block_name,
+                "model_id": self.model_id,
+                "batch_size": len(samples),
+                "effective_params": {
+                    k: v
+                    for k, v in final_kwargs.items()
+                    if k
+                    in [
+                        "temperature",
+                        "max_tokens",
+                        "max_completion_tokens",
+                        "top_p",
+                        "n",
+                        "seed",
+                    ]
+                },
+            },
+        )
+
+        # Run async generation
+        return asyncio.run(self._generate_async(samples, final_kwargs))
+
+    async def _generate_async(
+        self, samples: Dataset, final_kwargs: Dict[str, Any]
+    ) -> Dataset:
+        """Internal async method to generate all responses concurrently."""
+        # Extract all messages
+        messages_list = samples[self.messages_column]
+
+        # Create all tasks
+        tasks = [
+            self._generate_single(messages, **final_kwargs)
+            for messages in messages_list
+        ]
+
+        # Execute all tasks concurrently and collect responses
+        responses = await asyncio.gather(*tasks)
+
+        # Log completion
+        logger.info(
+            f"Async generation completed successfully for {len(responses)} samples",
+            extra={
+                "block_name": self.block_name,
+                "model_id": final_kwargs["model"],
+                "batch_size": len(responses),
+            },
+        )
+
+        # Add responses as new column
+        return samples.add_column(self.output_column, responses)

sdg_hub/flow.py

@@ -38,6 +38,7 @@ from .logger_config import setup_logger
 from .prompts import *  # needed to register prompts
 from .registry import BlockRegistry, PromptRegistry
 from .utils.config_validation import validate_prompt_config_schema
+from .utils.path_resolution import resolve_path
 from .utils.validation_result import ValidationResult
 
 logger = setup_logger(__name__)
@@ -141,15 +142,7 @@ class Flow(ABC):
         str
             Selected file path.
         """
-        if os.path.isabs(filename):
-            return filename
-        for d in dirs:
-            full_file_path = os.path.join(d, filename)
-            if os.path.isfile(full_file_path):
-                return full_file_path
-        # If not found above then return the path unchanged i.e.
-        # assume the path is relative to the current directory
-        return filename
+        return resolve_path(filename, dirs)
 
     def _drop_duplicates(self, dataset: Dataset, cols: List[str]) -> Dataset:
         """Drop duplicates from the dataset based on the columns provided.
@@ -273,7 +266,9 @@ class Flow(ABC):
                 try:
                     with open(path, "r", encoding="utf-8") as f:
                         config_data = yaml.safe_load(f)
-                        _, validation_errors = validate_prompt_config_schema(config_data, path)
+                        _, validation_errors = validate_prompt_config_schema(
+                            config_data, path
+                        )
 
                         if validation_errors:
                             errors.extend(validation_errors)
@@ -320,9 +315,7 @@ class Flow(ABC):
         KeyError
             If a required block or prompt is not found in the registry.
         """
-        yaml_path_relative_to_base = os.path.join(self.base_path, yaml_path)
-        if os.path.isfile(yaml_path_relative_to_base):
-            yaml_path = yaml_path_relative_to_base
+        yaml_path = resolve_path(yaml_path, self.base_path)
         yaml_dir = os.path.dirname(yaml_path)
 
         try:
@@ -433,7 +426,11 @@ class Flow(ABC):
             config = block["block_config"]
 
             # LLM Block: parse Jinja vars
-            cls_name = block_type.__name__ if isinstance(block_type, type) else block_type.__class__.__name__
+            cls_name = (
+                block_type.__name__
+                if isinstance(block_type, type)
+                else block_type.__class__.__name__
+            )
             logger.info(f"Validating block: {name} ({cls_name})")
             if "LLM" in cls_name:
                 config_path = config.get("config_path")
@@ -445,7 +442,9 @@ class Flow(ABC):
                         vars_found = meta.find_undeclared_variables(ast)
                         for var in vars_found:
                             if var not in all_columns:
-                                errors.append(f"[{name}] Missing column for prompt var: '{var}'")
+                                errors.append(
+                                    f"[{name}] Missing column for prompt var: '{var}'"
+                                )
 
             # FilterByValueBlock
             if "FilterByValueBlock" in str(block_type):
@@ -462,13 +461,17 @@ class Flow(ABC):
                 choice_map = config.get("choice_map", {})
                 for col in choice_map.values():
                     if col not in all_columns:
-                        errors.append(f"[{name}] choice_map references missing column: '{col}'")
+                        errors.append(
+                            f"[{name}] choice_map references missing column: '{col}'"
+                        )
 
             # CombineColumnsBlock
             if "CombineColumnsBlock" in str(block_type):
                 cols = config.get("columns", [])
                 for col in cols:
                     if col not in all_columns:
-                        errors.append(f"[{name}] CombineColumnsBlock requires column: '{col}'")
+                        errors.append(
+                            f"[{name}] CombineColumnsBlock requires column: '{col}'"
+                        )
 
-        return ValidationResult(valid=(len(errors) == 0), errors=errors)
+        return ValidationResult(valid=(len(errors) == 0), errors=errors)