Support for Azure OpenAI with AI Assistant

[matulef]

Description

Currently, the AI assistant works with any OpenAI API-compatible endpoint. However, when OpenAI models are deployed on Azure, there is a slightly different initialization strategy. See here:

https://github.com/openai/openai-python?tab=readme-ov-file#microsoft-azure-openai

The user needs to provide an API version number, then the "AzureOpenAI" client is created instead of a regular OpenAI client (the URL ultimately looks like https://[XXXX].openai.azure.com/openai/deployments/gpt-4o/chat/completions?api-version=2024-10-01-preview). Once the client is initialized, the API calls are the same, so I believe the rest should just work.

Personally, at my company we use an Azure deployment instead of OpenAI's own API for privacy/compliance reasons, so it would be nice to support this.

Suggested solution

We can use the marimo.toml file to specify which flavor of the API to use, say using an api_type variable, e.g.

[ai.open_ai]
api_type = "azure" # 'azure' or 'openai', default is 'openai' if unspecified
api_key = "*********" 
api_version = "2024-10-01-preview"
model = "gpt-4o" # for azure this is the deployment_name, which is often the model name but can differ 
base_url = "https://example-endpoint.openai.azure.com"

Then when the client is initialized (looks like the llm.py file here) we can check the api_type in the config, and use the either the AzureOpenAI class or regular OpenAI class accordingly.

Alternative

No response

Additional context

No response

[mscolnick]

Is this still true? I thought you could just pass https://[XXXX].openai.azure.com/openai/deployments/gpt-4o/chat/completions?api-version=2024-10-01-preview?

If not, i'd prefer to not update the config if possible (since the keys don't make sense in the non-azure case), and instead parse the URL https://[XXXX].openai.azure.com/openai/deployments/gpt-4o/chat/completions?api-version=2024-10-01-preview assuming that is quite stable

[matulef]

Unfortunately I tried passing that URL directly to the openAI client, but it just doesn't work (I also tried variants leaving off the chat/completions at the end and passing the api-version through an OPENAI_API_VERSION env variable, but the result is the same).

# this yields Error code: 404 - {'error': {'code': '404', 'message': 'Resource not found'}}
client = OpenAI(
    api_key=os.getenv("OPENAI_API_KEY"),
    base_url="https://my-deployment.openai.azure.com/openai/deployments/gpt-4o/chat/completions?api-version=2024-10-01-preview",
)

# But this works:
client = AzureOpenAI(
    api_key=os.getenv("OPENAI_API_KEY"),
    api_version="2024-10-01-preview",
    azure_endpoint="https://my-deployment.openai.azure.com",
)

Parsing the Azure params from the URL is definitely an option, and makes the config cleaner, though I have no idea how stable the URL format is. I suppose we can assume it's stable for now? Then we can just infer that any base_url of the form [XXXX].openai.azure.com is an Azure endpoint, and extract the model and api_version from https://[base_url]/deployments/[model]/chat/completions?api-version=[api_version]

[mscolnick]

That's approach is preferable with me (parsing the url)

[matulef]

Made a pull request that implements the url parsing approach. It works for me. Feel free to review and modify as you see fit. And thanks for building marimo!