Skip to content

Conversation

@jonathan-buttner
Copy link
Contributor

@jonathan-buttner jonathan-buttner commented Sep 10, 2025

This PR adds support for including custom headers in the chat completion and completion openai request. We'll need the support for other openai task types but I'll do that in a follow up PR to keep this one small.

The headers are stored in the task_settings.

This PR also introduce a new type for the _inference/_services configuration results called map to represent the headers field.

Example create inference API request

PUT _inference/chat_completion/openai
{
    "service": "openai",
    "service_settings": {
        "api_key": "<api_key>",
        "model_id": "o4-mini"
    },
    "task_settings": {
        "headers": {
            "OpenAI-Organization": "org id",
            "x-request-id": "jon-test"
        }
    }
}
OpenAI configuration
{
    "service": "openai",
    "name": "OpenAI",
    "task_types": [
        "text_embedding",
        "completion",
        "chat_completion"
    ],
    "configurations": {
        "headers": { <---------- New field
            "description": "Custom headers to include in the requests to OpenAI.",
            "label": "Custom Headers",
            "required": false,
            "sensitive": false,
            "updatable": true,
            "type": "map", <-------- New type
            "supported_task_types": [
                "completion",
                "chat_completion"
            ]
        },
        "api_key": {
            "description": "The OpenAI API authentication key. For more details about generating OpenAI API keys, refer to the https://platform.openai.com/account/api-keys.",
            "label": "API Key",
            "required": true,
            "sensitive": true,
            "updatable": true,
            "type": "str",
            "supported_task_types": [
                "text_embedding",
                "completion",
                "chat_completion"
            ]
        },
        "organization_id": {
            "description": "The unique identifier of your organization.",
            "label": "Organization ID",
            "required": false,
            "sensitive": false,
            "updatable": false,
            "type": "str",
            "supported_task_types": [
                "text_embedding",
                "completion",
                "chat_completion"
            ]
        },
        "rate_limit.requests_per_minute": {
            "description": "Default number of requests allowed per minute. For text_embedding is 3000. For completion is 500.",
            "label": "Rate Limit",
            "required": false,
            "sensitive": false,
            "updatable": false,
            "type": "int",
            "supported_task_types": [
                "text_embedding",
                "completion",
                "chat_completion"
            ]
        },
        "model_id": {
            "description": "The name of the model to use for the inference task.",
            "label": "Model ID",
            "required": true,
            "sensitive": false,
            "updatable": false,
            "type": "str",
            "supported_task_types": [
                "text_embedding",
                "completion",
                "chat_completion"
            ]
        },
        "url": {
            "description": "The absolute URL of the external service to send requests to.",
            "label": "URL",
            "required": false,
            "sensitive": false,
            "updatable": false,
            "type": "str",
            "supported_task_types": [
                "text_embedding",
                "completion",
                "chat_completion"
            ]
        },
        "dimensions": {
            "description": "The number of dimensions the resulting embeddings should have. For more information refer to https://platform.openai.com/docs/api-reference/embeddings/create#embeddings-create-dimensions.",
            "label": "Dimensions",
            "required": false,
            "sensitive": false,
            "updatable": false,
            "type": "int",
            "supported_task_types": [
                "text_embedding"
            ]
        }
    }
}

@jonathan-buttner jonathan-buttner added >enhancement :ml Machine learning Team:ML Meta label for the ML team v9.2.0 labels Sep 10, 2025
@elasticsearchmachine
Copy link
Collaborator

Hi @jonathan-buttner, I've created a changelog YAML for you.


public void testIsEmpty() {
var randomSettings = new OpenAiChatCompletionTaskSettings(randomBoolean() ? null : "username");
var randomSettings = new OpenAiEmbeddingsTaskSettings(randomBoolean() ? null : "username");
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Randomly found that this was testing the wrong class so switch it to the appropriate one.

ValidationException validationException = new ValidationException();

String user = extractOptionalString(map, USER, ModelConfigurations.TASK_SETTINGS, validationException);
Map<String, Object> headers = extractOptionalMapRemoveNulls(map, HEADERS, validationException);
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since we need to add the headers to each task type I'll move this functionality to an abstract class in a later PR.

@jonathan-buttner jonathan-buttner marked this pull request as ready for review September 11, 2025 12:55
@jonathan-buttner jonathan-buttner requested a review from a team as a code owner September 11, 2025 12:55
@elasticsearchmachine
Copy link
Collaborator

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

@jonathan-buttner jonathan-buttner changed the title [ML] Adding headers support for openai [ML] Adding headers support for OpenAI chat completion Sep 11, 2025
}

/**
* Read an optional {@link Map} using the given key and value readers. The return Map is immutable.
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor typo, should be "the returned Map is immutable"


public void testFromMap_ReturnsUser() {
var settings = OpenAiChatCompletionRequestTaskSettings.fromMap(new HashMap<>(Map.of(OpenAiServiceFields.USER, "user")));
public void testFromMap_ParsesCorrectly() {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It would be good to also have a test for the case where the headers map is null, since that's new functionality in this PR.

@jonathan-buttner jonathan-buttner added the cloud-deploy Publish cloud docker image for Cloud-First-Testing label Sep 11, 2025
@jonathan-buttner jonathan-buttner merged commit 4a39b4c into elastic:main Sep 16, 2025
36 checks passed
@jonathan-buttner jonathan-buttner deleted the ml-openai-headers branch September 16, 2025 12:52
mridula-s109 pushed a commit to mridula-s109/elasticsearch that referenced this pull request Sep 17, 2025
* Adding headers support for openai chat completion and completion

* Update docs/changelog/134504.yaml

* [CI] Auto commit changes from spotless

* Adjusting test name

* Updating the changelog

* Adding headers to the configuration

* [CI] Auto commit changes from spotless

* [CI] Update transport version definitions

* Fixing transport versions

---------

Co-authored-by: elasticsearchmachine <infra-root+elasticsearchmachine@elastic.co>
gmjehovich pushed a commit to gmjehovich/elasticsearch that referenced this pull request Sep 18, 2025
* Adding headers support for openai chat completion and completion

* Update docs/changelog/134504.yaml

* [CI] Auto commit changes from spotless

* Adjusting test name

* Updating the changelog

* Adding headers to the configuration

* [CI] Auto commit changes from spotless

* [CI] Update transport version definitions

* Fixing transport versions

---------

Co-authored-by: elasticsearchmachine <infra-root+elasticsearchmachine@elastic.co>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cloud-deploy Publish cloud docker image for Cloud-First-Testing >enhancement :ml Machine learning Team:ML Meta label for the ML team v9.2.0

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants