Remote model inference streaming

[jngz-es]

Description

Support remote model streaming predict API. For now support openai v1 completion and bedrock converse stream APIs.

example

1. Create an openai connector

curl -X POST "http://localhost:9200/_plugins/_ml/connectors/_create" -H "Content-Type: application/json" -d'
{
    "name": "OpenAI Chat Connector",
    "description": "The connector to public OpenAI model service for GPT 3.5",
    "version": 1,
    "protocol": "http",
    "parameters": {
        "endpoint": "api.openai.com",
        "model": "gpt-3.5-turbo"
    },
    "credential": {
        "openAI_key": "<your_token_id>"
    },
    "actions": [
        {
            "action_type": "predict",
            "method": "POST",
            "url": "https://${parameters.endpoint}/v1/chat/completions",
            "headers": {
                "Authorization": "Bearer ${credential.openAI_key}"
            },
            "request_body": "{ \"model\": \"${parameters.model}\", \"messages\": ${parameters.messages} }"
        }
    ]
}
'

2. Create a model with the above connector.

curl -X POST "http://localhost:9200/_plugins/_ml/models/_register?deploy=true" -H "Content-Type: application/json" -d'
{
    "name": "openai gpt 3.5 turbo",
    "function_name": "remote",
    "description": "openai model",
    "connector_id": "<your_connector_id>"
}
'

3. Predict with stream api.

curl -X POST "http://localhost:9200/_plugins/_ml/models/AlI3QZcB-tBcnGSI1JlH/_predict/stream" -H "Transfer-Encoding: chunked" -H "Content-Type: application/json" -d'
{
  "parameters": {
    "messages": [
      {
        "role": "system",
        "content": "You are a helpful assistant."
      },
      {
        "role": "user",
        "content": "Can you summarize Prince Hamlet of William Shakespeare in around 1000 words?"
      }
    ],
    "_llm_interface": "openai/v1/chat/completions"
  }
}
'

4. You can run the non-stream predict too.

curl -X POST "http://localhost:9200/_plugins/_ml/models/AlI3QZcB-tBcnGSI1JlH/_predict" -H "Content-Type: application/json" -d'
{
  "parameters": {
    "messages": [
      {
        "role": "system",
        "content": "You are a helpful assistant."
      },
      {
        "role": "user",
        "content": "Can you summarize Prince Hamlet of William Shakespeare in around 1000 words?"
      }
    ]
  }
}
'

Related Issues

#3630

Check List

• New functionality includes testing.
• New functionality has been documented.
• API changes companion pull request created.
• Commits are signed per the DCO using --signoff.
• Public documentation issue/PR created.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

[jngz-es]

Hi team, I am publishing this pr for your early review. Meanwhile I am adding feature flag, fixing the existing tests and adding new tests.

[rithin-pullela-aws]

I see you mentioned Bedrock converse support too in the description. Are we only going with Open AI for now?

[jngz-es]

The neededStreamParameterInPayload is to check if we need to add "stream": true in http request body as openai does. But Bedrock converseStream doesn't use that parameter, but use a specific url like POST /model/modelId/converse-stream, so no converse here.

[rithin-pullela-aws]

Just to understand, how did we come up with these numbers?

[jngz-es]

I used remote configuration as a reference when configuring stream pool.

[rithin-pullela-aws]

Can we avoid thread.sleep by using a blocking queue?