Openai

zenml.integrations.openai special

Initialization of the OpenAI integration.

OpenAIIntegration (Integration)

Definition of OpenAI integration for ZenML.

Source code in zenml/integrations/openai/__init__.py

class OpenAIIntegration(Integration):
    """Definition of OpenAI integration for ZenML."""

    NAME = OPEN_AI
    REQUIREMENTS = ["openai>=0.27.0"]

hooks special

Initialization of the OpenAI hooks module.

open_ai_failure_hook

Functionality for OpenAI standard hooks.

openai_alerter_failure_hook_helper(context, params, exception, model_name)

Standard failure hook that sends a message to an Alerter.

Your OpenAI API key must be stored in the secret store under the name "openai" and with the key "api_key".

Parameters:

Name	Type	Description	Default
context	StepContext	The context of the step.	required
params	BaseParameters	The parameters of the step.	required
exception	BaseException	The exception that was raised.	required
model_name	str	The OpenAI model to use for the chatbot.	required

Source code in zenml/integrations/openai/hooks/open_ai_failure_hook.py

def openai_alerter_failure_hook_helper(
    context: StepContext,
    params: BaseParameters,
    exception: BaseException,
    model_name: str,
) -> None:
    """Standard failure hook that sends a message to an Alerter.

    Your OpenAI API key must be stored in the secret store under the name
    "openai" and with the key "api_key".

    Args:
        context: The context of the step.
        params: The parameters of the step.
        exception: The exception that was raised.
        model_name: The OpenAI model to use for the chatbot.
    """
    # get the api_key from the secret store
    try:
        c = Client()
        openai_secret = c.get_secret("openai", allow_partial_name_match=False)
        openai_api_key = openai_secret.secret_values.get("api_key")
    except (KeyError, NotImplementedError):
        openai_api_key = None

    if context.stack and context.stack.alerter and openai_api_key:
        output_captured = io.StringIO()
        original_stdout = sys.stdout
        sys.stdout = output_captured
        console = Console()
        console.print_exception(show_locals=False)

        sys.stdout = original_stdout
        rich_traceback = output_captured.getvalue()

        response = openai.ChatCompletion.create(  # type: ignore[no-untyped-call]
            model=model_name,
            messages=[
                {
                    "role": "user",
                    "content": f"This is an error message (following an exception of type '{type(exception)}') I encountered while executing a ZenML step. Please suggest ways I might fix the problem. Feel free to give code snippets as examples, and note that your response will be piped to a Slack bot so make sure the formatting is appropriate: {exception} -- {rich_traceback}. Thank you!",
                }
            ],
        )
        suggestion = response["choices"][0]["message"]["content"]
        message = "*Failure Hook Notification! Step failed!*" + "\n\n"
        message += f"Pipeline name: `{context.pipeline_name}`" + "\n"
        message += f"Run name: `{context.run_name}`" + "\n"
        message += f"Step name: `{context.step_name}`" + "\n"
        message += f"Parameters: `{params}`" + "\n"
        message += f"Exception: `({type(exception)}) {exception}`" + "\n\n"
        message += (
            f"Step Cache Enabled: `{'True' if context.cache_enabled else 'False'}`"
            + "\n\n"
        )
        message += (
            f"*OpenAI ChatGPT's suggestion (model = `{model_name}`) on how to fix it:*\n `{suggestion}`"
            + "\n"
        )
        context.stack.alerter.post(message)
    elif not openai_api_key:
        logger.warning(
            "Specified OpenAI failure hook but no OpenAI API key found. Skipping..."
        )
    else:
        logger.warning(
            "Specified OpenAI failure hook but no alerter configured in the stack. Skipping..."
        )

openai_chatgpt_alerter_failure_hook(context, params, exception)

Alerter hook that uses the OpenAI ChatGPT model.

Parameters:

Name	Type	Description	Default
context	StepContext	The context of the step.	required
params	BaseParameters	The parameters of the step.	required
exception	BaseException	The exception that was raised.	required

Source code in zenml/integrations/openai/hooks/open_ai_failure_hook.py

def openai_chatgpt_alerter_failure_hook(
    context: StepContext,
    params: BaseParameters,
    exception: BaseException,
) -> None:
    """Alerter hook that uses the OpenAI ChatGPT model.

    Args:
        context: The context of the step.
        params: The parameters of the step.
        exception: The exception that was raised.
    """
    openai_alerter_failure_hook_helper(
        context, params, exception, "gpt-3.5-turbo"
    )

openai_gpt4_alerter_failure_hook(context, params, exception)

Alerter hook that uses the OpenAI GPT-4 model.

Parameters:

Name	Type	Description	Default
context	StepContext	The context of the step.	required
params	BaseParameters	The parameters of the step.	required
exception	BaseException	The exception that was raised.	required

Source code in zenml/integrations/openai/hooks/open_ai_failure_hook.py

def openai_gpt4_alerter_failure_hook(
    context: StepContext,
    params: BaseParameters,
    exception: BaseException,
) -> None:
    """Alerter hook that uses the OpenAI GPT-4 model.

    Args:
        context: The context of the step.
        params: The parameters of the step.
        exception: The exception that was raised.
    """
    openai_alerter_failure_hook_helper(context, params, exception, "gpt-4")