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(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
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(
    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:
        exception: The exception that was raised.
        model_name: The OpenAI model to use for the chatbot.
    """
    client = Client()
    context = get_step_context()

    # get the api_key from the secret store
    try:
        openai_secret = client.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

    alerter = client.active_stack.alerter
    if 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"Run name: `{context.pipeline_run.name}`" + "\n"
        message += f"Step name: `{context.step_run.name}`" + "\n"
        message += f"Parameters: `{context.step_run.config.parameters}`" + "\n"
        message += f"Exception: `({type(exception)}) {exception}`" + "\n\n"
        message += (
            f"*OpenAI ChatGPT's suggestion (model = `{model_name}`) on how to fix it:*\n `{suggestion}`"
            + "\n"
        )
        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(exception)

Alerter hook that uses the OpenAI ChatGPT model.

Parameters:

Name	Type	Description	Default
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(
    exception: BaseException,
) -> None:
    """Alerter hook that uses the OpenAI ChatGPT model.

    Args:
        exception: The exception that was raised.
    """
    openai_alerter_failure_hook_helper(exception, "gpt-3.5-turbo")

openai_gpt4_alerter_failure_hook(exception)

Alerter hook that uses the OpenAI GPT-4 model.

Parameters:

Name	Type	Description	Default
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(
    exception: BaseException,
) -> None:
    """Alerter hook that uses the OpenAI GPT-4 model.

    Args:
        exception: The exception that was raised.
    """
    openai_alerter_failure_hook_helper(exception, "gpt-4")