OpenAI SDK wrapper for Azure AD authentications

openai-azure-ad-wrapper.py

from azure.identity import DefaultAzureCredential
from typing import Awaitable
import asyncio
import openai

# Set up the Azure credential
credential = DefaultAzureCredential()

def openai_refresh() -> None:
  """
  Refreshes the OpenAI API key using the Azure credential.
  """
  openai.api_key = credential.get_token().token

# First refresh
openai_refresh()

async def completion_wrapper(prompt: str) -> Awaitable[str]:
  """
  Execute the OpenAI completion API.
  """
  return openai.Completion.acreate(
    engine="davinci",
    max_tokens=5,
    prompt=prompt,
  )

async def completion(prompt: str):
  """
  Wrapper around the OpenAI completion API that handles authentication errors and rotates the API key if necessary.
  """
  try:
    return await completion_wrapper(prompt)
  except openai.error.AuthenticationError:
    openai_refresh()
    return await completion_wrapper(prompt)

async def main():
  """
  Main function that runs the bot.
  """
  completion("Write me a poem about a cat sitting on a chair.")

if __name__ == "__main__":
  asyncio.run(main())