(obsolete - do not merge) [Inference API] Add Azure AI Studio Embeddings and Chat Completion Support

[markjhoy]

This PR adds support for Azure AI Studio integration into the Inference API. Currently this supports text_embedding and completion task types.

Prerequisites to Model Creation

• You must have an Azure subscription with Azure AI Studio access
• You must have a deployed model either Chat Completion or Embeddings

Model Creation:

PUT _inference/{tasktype}/{model_id}
{
  "service": "azureaistudio",
  "service_settings": {
    "api_key": "{api_key}",
    "target": “{deployment_target}”,
    “provider”: “(model provider}”,
    “endpoint_type”: “(endpoint type)”
  }
}

• Valid {tasktype} types are: [text_embedding, completion]

Required Service Settings:

• api_key: The API key can be found on your Azure AI Studio deployment's overview page
• target: The target URL can be found on your Azure AI Studio deployment's overview page
• provider: Valid provider types are (case insensitive):
  • openai - available for embeddings and completion
  • mistral - available for completion only
  • meta - available for completion only
  • microsoft_phi - available for completion only
  • cohere - available for embeddings and completion
  • snowflake - available for completion only
  • databricks - available for completion only
• endpoint_type: Valid endpoint types are:
  • token - a "pay as you go" endpoint (charged by token)
    • Available for OpenAI, Meta and Cohere
  • realtime - a realtime endpoint VM deployment (charged by the hour)
    • Available for Mistral, Meta, Microsoft Phi, Snowflake and Databricks

Embeddings Service Settings

• dimensions: (optional) the number of dimensions the resulting output embeddings should have.

Embeddings Task Settings

(this is also overridable in the inference request)

• user: (optional) a string that is a unique identifier representing your end-user. This helps Azure AI Studio in the case of abuse or issues for debugging.

Completion Service Settings

(no additional service settings)

Completion Task Settings

(these are all optional and can be overridden in the inference request)

• temperature: What sampling temperature to use, between 0 and 2. Higher values mean the model takes more risks. Try 0.9 for more creative applications, and 0 (argmax sampling) for ones with a well-defined answer. Microsoft recommends altering this or top_p but not both.
• top_p: An alternative to sampling with temperature, called nucleus sampling, where the model considers the results of the tokens with top_p probability mass. So 0.1 means only the tokens comprising the top 10% probability mass are considered. Microsoft recommends altering this or temperature but not both.
• do_sample: request to perform the sampling or not
• max_new_tokens: the maximum number of new tokens the chat completion inference should produce in the output

Text Embedding Inference

POST _inference/text_embedding/{model_id}
{
    "input": "The answer to the universe is"
}

Chat Completion Inference

POST _inference/completion/{model_id}
{
    "input": "The answer to the universe is"
}

[elasticsearchmachine]

Pinging @elastic/ml-core (Team:ML)

[elasticsearchmachine]

Pinging @elastic/ent-search-eng (Team:Enterprise Search)

[davidkyle]

Looks great

It's a very large PR I'd like to do some testing then take another look

[markjhoy]

buildkite test this

[markjhoy]

Oh - I sometimes hate git and it's ability to not allow a clean merge... going to close this and re-open a new one that's clean.

[jonathan-buttner]

I'm going to add these here so I don't lose them haha, I'll copy them over to the new PR when it's up 👍

[jonathan-buttner]

Thanks for refactoring!

[jonathan-buttner]

I think we can get this field from AzureAiStudioRequest right?

[jonathan-buttner]

To avoid the cast, I think we can change the model to be AzureAiStudioEmbeddingsModel. Seems like each place we initialize this class we already know it's an embedding model.

[jonathan-buttner]

How about we make the default constructor private?

[jonathan-buttner]

hmm, since this class and EmbeddingResponseEntity are pretty similar for now, how do you feel about collapsing them into a single class?

I think we can achieve that by doing something like:

protected abstract InferenceServiceResults fromResponse(Request request, HttpResult response) throws IOException;

...

Or I suppose we could omit this class and have the child classes implement ResponseParser directly?

My reasoning for using InferenceServiceResults is because parsers like cohere can return different types (TextEmbeddingResults and TextEmbeddingByteResults).

[jonathan-buttner]

Same here, do we need this function?

[jonathan-buttner]

Do we need this since it's not used?

[jonathan-buttner]

Do we need this since it's not used?

[jonathan-buttner]

How about we keep with the visitor pattern and add an accept() method here similar to azure openai: https://github.com/elastic/elasticsearch/blob/main/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/azureopenai/AzureOpenAiModel.java#L59

If we add a new task type, hopefully it'd limit the number of changes to the azure studio service class.

[jonathan-buttner]

I mentioned this above, but the reason I chose the visitor pattern for the other services was to let the model dictate the action that is used. That way we don't need the instanceof checks here, except one to ensure that it's a AzureAiStudioModel model with an accept() method.

[markjhoy]

Sorry @jonathan-buttner ! :) -- new PR is up here: #108472