[Inference API] Add unified api for chat completions

[maxhniebergall]

Unified API communicating with OpenAI

Testing

Running ES

The _unified route is behind a feature flag, so to enable it run es like this:

./gradlew :run -Drun.license_type=trial -Des.inference_unified_feature_flag_enabled=true

Creating endpoint and sending requestions

Creating a completion endpoint

PUT http://localhost:9200/_inference/completion/test
{
    "service": "openai",
    "service_settings": {
        "api_key": "<api key>",
        "model_id": "gpt-4o"
    }
}

Completion request

POST http://localhost:9200/_inference/completion/test/_unified
{
    "model": "gpt-4o",
    "messages": [
        {
            "role": "user",
            "content": "What is the weather like in Boston today?"
        }
    ],
    "stop": "none",
    "tools": [
        {
            "type": "function",
            "function": {
                "name": "get_current_weather",
                "description": "Get the current weather in a given location",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "location": {
                            "type": "string",
                            "description": "The city and state, e.g. San Francisco, CA"
                        },
                        "unit": {
                            "type": "string",
                            "enum": [
                                "celsius",
                                "fahrenheit"
                            ]
                        }
                    },
                    "required": [
                        "location"
                    ]
                }
            }
        }
    ],
    "tool_choice": "auto"
}

Response format

A sequence of partial responses:

{
    "id": "chatcmpl-AamGq3bqfB5bhmDaBLYi5Z5kTUz5R",
    "choices": [
        {
            "delta": {
                "content": " ol"
            },
            "index": 0
        }
    ],
    "model": "gpt-3.5-turbo-0125",
    "object": "chat.completion.chunk"
}

{
    "id": "chatcmpl-AamGq3bqfB5bhmDaBLYi5Z5kTUz5R",
    "choices": [
        {
            "delta": {
                "content": "iv"
            },
            "index": 0
        }
    ],
    "model": "gpt-3.5-turbo-0125",
    "object": "chat.completion.chunk"
}

{
    "id": "chatcmpl-AamGq3bqfB5bhmDaBLYi5Z5kTUz5R",
    "choices": [
        {
            "delta": {
                "content": "ine"
            },
            "index": 0
        }
    ],
    "model": "gpt-3.5-turbo-0125",
    "object": "chat.completion.chunk"
}

{
    "id": "chatcmpl-AamGq3bqfB5bhmDaBLYi5Z5kTUz5R",
    "choices": [
        {
            "delta": {
                "content": "."
            },
            "index": 0
        }
    ],
    "model": "gpt-3.5-turbo-0125",
    "object": "chat.completion.chunk"
}

{
    "id": "chatcmpl-AamGq3bqfB5bhmDaBLYi5Z5kTUz5R",
    "choices": [
        {
            "delta": {},
            "finish_reason": "stop",
            "index": 0
        }
    ],
    "model": "gpt-3.5-turbo-0125",
    "object": "chat.completion.chunk"
}

[DONE
]

[elasticsearchmachine]

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

[elasticsearchmachine]

Hi @maxhniebergall, I've created a changelog YAML for you.

[maxhniebergall]

Tests to add:

• UnifiedCompletionAction (@jonathan-buttner) ✅
• TestStreamingCompletionServiceExtension (@jonathan-buttner) ✅
• Rolling update tests (We don't think we need these?)
• TransportInferenceActionTests (@jonathan-buttner) ✅
• InferenceInputs (@jonathan-buttner) ✅
  • we should double check the castTo method and add it to the other subclasses (Let's do this later because the PR is large enough)
• UnifiedChatInput (@jonathan-buttner) ✅
  • for the conversions
• OpenAiUnifiedCompletionRequestEntity (@maxhniebergall)✅
• BaseInferenceAction (@jonathan-buttner) ✅
• RestUnifiedCompletionInferenceAction (@jonathan-buttner) ✅
• OpenAiService (@jonathan-buttner) ✅
• OpenAiChatCompletionModel (@jonathan-buttner) ✅
• OpenAiUnifiedStreamingProcessor & StreamingUnifiedChatCompletionResults (@maxhniebergall )✅

TODO

• Create a list of all new named writables and add them to the registry ✅
  Address outstanding TODOs ✅

[maxhniebergall]

LGTM

[jonathan-buttner]

I just realized that the OpenAiChatCompletionServiceSettings has a similar field that isn't used (even previous to this PR). I wonder if we should sync those fields up like we do the modelId 🤔 . I think we've typically used the max tokens to do the truncation for text embedding so I don't think we really used it for completions.

@maxhniebergall what do you think?

[maxhniebergall]

hmm, yea it sounds like we should use the value from the service settings if its available

[davidkyle]

LGTM

[davidkyle]

Is function required? The class does not declare that is implements toXContent

[jonathan-buttner]

Oops, yeah we can remove that. Thanks.

[davidkyle]

Suggested change

	if (taskType != TaskType.COMPLETION) {
	if (taskType.isAnyOrSame(TaskType.COMPLETION)) {

For the case where tasktype is not set in the URL and defaulted to ANY

[jonathan-buttner]

Ah good catch!

[davidkyle]

Why size 2 and not 1?

[maxhniebergall]

I'm not sure, it was in the code Pat sent me for this. I also thought it was odd prwhelan@4c573ba

[davidkyle]

Nice! TIL

[davidkyle]

It took me a while to grok that onRequest is not part of the Flow interface. Maybe call this upstreamRequest not to confuse it with the various Flow on* methods

[jonathan-buttner]

Yep sounds good. I still don't quite understand what all that stuff is doing haha. I'll have Pat give an overview when he gets back maybe.

[maxhniebergall]

💚 All backports created successfully

Status	Branch	Result
✅	8.x

Questions ?

Please refer to the Backport tool documentation

[maxhniebergall]

💚 All backports created successfully

Status	Branch	Result
✅	8.x

Questions ?

Please refer to the Backport tool documentation