[Inference API] Fix Azure AI Studio Integration for Completions and Embeddings

[brendan-jugan-elastic]

This draft PR fixes the Inference API integration with Azure AI Foundry (previously Azure AI Studio). The previous integration was broken for both completions and embeddings models due to API changes from Microsoft.

Core Changes:

• no longer referencing an approved list of providers
• introducing a required AzureAiStudioDeploymentType to the service settings
  • either azure_ai_model_inference_service or serverless_api
• modifying auth configuration for each deployment type
• slight request format modifications after API changes
• testing and rebranding changes are in-progress, wanted to get some eyes on the implementation while I complete them

Once testing is complete, I will add more detailed docs describing the deployment types, their configurations, and describe how to use this integration with screenshots from the Azure console.

Local Testing:

Embeddings:

PUT http://localhost:9200/_inference/text_embedding/cohere_serverless_embed

curl --location --request PUT 'http://localhost:9200/_inference/text_embedding/cohere_serverless_embed' \
--header 'Content-Type: application/json' \
--header 'Authorization: Basic *****' \
--data '{
  "service": "azureaistudio",
  "service_settings": {
    "api_key": "*****",
    "target": "https://example-target.eastus.models.ai.azure.com/embeddings",
    "deployment_type": "serverless_api",
    "deployment_name": "Cohere-embed-v3-english-hmcek"
  }
}

POST http://localhost:9200/_inference/text_embedding/cohere_serverless_embed

curl --location 'http://localhost:9200/_inference/text_embedding/cohere_serverless_embed' \
--header 'Content-Type: application/json' \
--header 'Authorization: Basic *****' \
--data '{
  "input": "What is Elastic?"
}'



PUT http://localhost:9200/_inference/text_embedding/cohere_amlis_embed

curl --location --request PUT 'http://localhost:9200/_inference/text_embedding/cohere_amlis_embed' \
--header 'Content-Type: application/json' \
--header 'Authorization: Basic *****' \
--data '{
  "service": "azureaistudio",
  "service_settings": {
    "api_key": "*****",
    "target": "https://example-target/models/embeddings",
    "deployment_type": "azure_ai_model_inference_service",
    "deployment_name": "Cohere-embed-v3-english"
  }
}'

curl --location 'http://localhost:9200/_inference/text_embedding/cohere_amlis_embed' \
--header 'Content-Type: application/json' \
--header 'Authorization: Basic *****' \
--data '{
  "input": "What is Elastic?"
}'

Completions:

PUT http://localhost:9200/_inference/completion/cohere_serverless_completion

curl --location --request PUT 'http://localhost:9200/_inference/completion/cohere_serverless_completion' \
--header 'Content-Type: application/json' \
--header 'Authorization: Basic *****' \
--data '{
  "service": "azureaistudio",
  "service_settings": {
    "api_key": "*****",
    "target": "https://example-target.eastus.models.ai.azure.com/chat/completions",
    "deployment_type": "serverless_api",
    "deployment_name": "Cohere-command-r"
  }
}'

POST http://localhost:9200/_inference/completion/cohere_serverless_completion

curl --location 'http://localhost:9200/_inference/completion/cohere_serverless_completion' \
--header 'Content-Type: application/json' \
--header 'Authorization: Basic *****' \
--data '{
  "input": "What is Elastic?"
}'



PUT http://localhost:9200/_inference/completion/cohere_amlis_completion

curl --location --request PUT 'http://localhost:9200/_inference/completion/cohere_amlis_completion' \
--header 'Content-Type: application/json' \
--header 'Authorization: Basic *****' \
--data '{
  "service": "azureaistudio",
  "service_settings": {
    "api_key": "*****",
    "target": "https://example-target.services.ai.azure.com/models/chat/completions?api-version=2024-05-01-preview",
    "deployment_type": "azure_ai_model_inference_service",
    "deployment_name": "Cohere-command-r"
  }
}'

POST http://localhost:9200/_inference/completion/cohere_amlis_completion

curl --location 'http://localhost:9200/_inference/completion/cohere_amlis_completion' \
--header 'Content-Type: application/json' \
--header 'Authorization: Basic *****' \
--data '{
  "input": "What is Elastic?"
}'

Related Issues:

• https://github.com/elastic/search-team/issues/8893
• https://github.com/elastic/search-team/issues/8895
• https://github.com/elastic/search-team/issues/8894

Helpful Links:

• https://learn.microsoft.com/en-us/azure/ai-studio/reference/reference-model-inference-chat-completions
• https://learn.microsoft.com/en-us/azure/ai-studio/reference/reference-model-inference-embeddings

[timgrein]

Nice, good stuff! 👏 Gave it a first pass and left some comments, already looking good

[brendan-jugan-elastic]

Note: I'm waiting to complete the Azure AI Studio -> Azure AI Foundry renaming until tomorrow. The above commits contain all of the functional/test changes for this Inference API fix.

[timgrein]

Mainly left comments around changes we need to address with regards to the transport level changes - we can also sync on that, it can be a bit confusing in the beginning

[davidkyle]

The old node will does not know about the enum value AzureAiFoundryProvider.NONE as that is added in this PR. Enums are written by their ordinal values, the old node will read the ordinal value for AzureAiFoundryProvider.NONE (let's say it is 2) but then throw an error because it does not know of any AzureAiFoundryProvider enum with ordinal value 2.

If there isn't a logical mapping from the old fields (provider, endpointtype) to the new (deploymenttype, model) then it is a question of how we want this to fail.

Inference endpoint creation is a master node action. The request will be serialised from whichever node it lands on to the master node and the response in turn will be serialised back to the originating node. PutInferenceModelAction.Response contains the ServiceSettings (we return the new endpoint configuration), if the originating node is an old node that doesn't know about the new options it won't return the proper config as some fields have been lost in serialisation. We can solve for that by not allowing new AzureAiFoundry inference endpoints to be created in a mixed cluster where some nodes do not know about this change.

[davidkyle]

NAME has changed from azure_ai_studio_chat_completion_service_settings to azure_ai_foundry_chat_completion_service_settings. When an old node writes a named writable with the old name azure_ai_studio_chat_completion_service_settings this node will not know about it.

Because there is logic in the AzureAiFoundryChatCompletionServiceSettings serialisation code to handle backwards compatibility it is better not to change the name. Just add a comment explaining that NAME hasn't changed to maintain BWC