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,<1.0.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
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")