[ML] Fix OpenAI connector does not use the action proxy configuration for all subactions (#219617)

## Summary

This PR fixes https://github.com/elastic/kibana/issues/214057 by adding
the httpsAgent/httpAgent to the OpenAI client.

### Checklist

Check the PR satisfies following conditions. 

Reviewers should verify this PR satisfies this list as well.

- [ ] Any text added follows [EUI's writing
guidelines](https://elastic.github.io/eui/#/guidelines/writing), uses
sentence case text and includes [i18n
support](https://github.com/elastic/kibana/blob/main/src/platform/packages/shared/kbn-i18n/README.md)
- [ ]
[Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html)
was added for features that require explanation or tutorials
- [x] [Unit or functional
tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html)
were updated or added to match the most common scenarios
- [ ] If a plugin configuration key changed, check if it needs to be
allowlisted in the cloud and added to the [docker
list](https://github.com/elastic/kibana/blob/main/src/dev/build/tasks/os_packages/docker_generator/resources/base/bin/kibana-docker)
- [ ] This was checked for breaking HTTP API changes, and any breaking
changes have been approved by the breaking-change committee. The
`release_note:breaking` label should be applied in these situations.
- [ ] [Flaky Test
Runner](https://ci-stats.kibana.dev/trigger_flaky_test_runner/1) was
used on any tests changed
- [ ] The PR description includes the appropriate Release Notes section,
and the correct `release_note:*` label is applied per the
[guidelines](https://www.elastic.co/guide/en/kibana/master/contributing.html#kibana-release-notes-process)

### Identify risks

Does this PR introduce any risks? For example, consider risks like hard
to test bugs, performance regression, potential of data loss.

Describe the risk, its severity, and mitigation for each identified
risk. Invite stakeholders and evaluate how to proceed before merging.

- [ ] [See some risk
examples](https://github.com/elastic/kibana/blob/main/RISK_MATRIX.mdx)
- [ ] ...

---------

Co-authored-by: Elastic Machine <elasticmachine@users.noreply.github.com>
This commit is contained in:
Quynh Nguyen (Quinn) 2025-06-05 12:34:27 -05:00 committed by GitHub
parent aa488b41e7
commit 2564d6de38
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 155 additions and 2 deletions

View file

@ -15,6 +15,9 @@ import type { ActionsConfigurationUtilities } from '../actions_config';
import { getNodeSSLOptions, getSSLSettingsFromConfig } from './get_node_ssl_options'; import { getNodeSSLOptions, getSSLSettingsFromConfig } from './get_node_ssl_options';
import type { SSLSettings } from '../types'; import type { SSLSettings } from '../types';
/**
* Create http and https proxy agents with custom proxy /hosts/SSL settings from configurationUtilities
*/
interface GetCustomAgentsResponse { interface GetCustomAgentsResponse {
httpAgent: HttpAgent | undefined; httpAgent: HttpAgent | undefined;
httpsAgent: HttpsAgent | undefined; httpsAgent: HttpsAgent | undefined;

View file

@ -44,7 +44,7 @@ export abstract class SubActionConnector<Config, Secrets> {
[k: string]: ((params: unknown) => unknown) | unknown; [k: string]: ((params: unknown) => unknown) | unknown;
private axiosInstance: AxiosInstance; private axiosInstance: AxiosInstance;
private subActions: Map<string, SubAction> = new Map(); private subActions: Map<string, SubAction> = new Map();
private configurationUtilities: ActionsConfigurationUtilities; protected configurationUtilities: ActionsConfigurationUtilities;
protected readonly kibanaRequest?: KibanaRequest; protected readonly kibanaRequest?: KibanaRequest;
protected logger: Logger; protected logger: Logger;
protected esClient: ElasticsearchClient; protected esClient: ElasticsearchClient;

View file

@ -0,0 +1,142 @@
/*
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
* or more contributor license agreements. Licensed under the Elastic License
* 2.0; you may not use this file except in compliance with the Elastic License
* 2.0.
*/
import { DEFAULT_TIMEOUT_MS, OPENAI_CONNECTOR_ID } from '../../../common/openai/constants';
import { actionsMock } from '@kbn/actions-plugin/server/mocks';
import { DEFAULT_OPENAI_MODEL } from '../../../common/openai/constants';
import { actionsConfigMock } from '@kbn/actions-plugin/server/actions_config.mock';
import { OpenAIConnector } from './openai';
import { OpenAiProviderType } from '../../../common/openai/constants';
import { loggingSystemMock } from '@kbn/core/server/mocks';
import { ConnectorUsageCollector } from '@kbn/actions-plugin/server/types';
import { RunActionResponseSchema } from '../../../common/openai/schema';
const logger = loggingSystemMock.createLogger();
// Mock an instance of the OpenAI class
// with overridden flag for purpose of jest test
jest.mock('openai', () => {
const UnmodifiedOpenAIClient = jest.requireActual('openai').default;
return {
__esModule: true,
default: jest.fn().mockImplementation((config) => {
return new UnmodifiedOpenAIClient({
...config,
dangerouslyAllowBrowser: true,
});
}),
};
});
describe('OpenAI with proxy config', () => {
let mockProxiedRequest: jest.Mock;
let connectorUsageCollector: ConnectorUsageCollector;
const mockDefaults = {
timeout: DEFAULT_TIMEOUT_MS,
url: 'https://api.openai.com/v1/chat/completions',
method: 'post',
responseSchema: RunActionResponseSchema,
};
const mockResponse = {
headers: {},
data: {},
};
const configurationUtilities = actionsConfigMock.create();
const PROXY_HOST = 'proxy.custom.elastic.co';
const PROXY_URL = `http://${PROXY_HOST}`;
configurationUtilities.getProxySettings.mockReturnValue({
proxyUrl: PROXY_URL,
proxySSLSettings: {
verificationMode: 'none',
},
proxyBypassHosts: undefined,
proxyOnlyHosts: undefined,
});
const connector = new OpenAIConnector({
configurationUtilities,
connector: { id: '1', type: OPENAI_CONNECTOR_ID },
config: {
apiUrl: 'https://api.openai.com/v1/chat/completions',
apiProvider: OpenAiProviderType.OpenAi,
defaultModel: DEFAULT_OPENAI_MODEL,
organizationId: 'org-id',
projectId: 'proj-id',
headers: {
'X-My-Custom-Header': 'foo',
Authorization: 'override',
},
},
secrets: { apiKey: '123' },
logger,
services: actionsMock.createServices(),
});
const sampleOpenAiBody = {
messages: [
{
role: 'user',
content: 'Hello world',
},
],
};
beforeEach(() => {
connectorUsageCollector = new ConnectorUsageCollector({
logger,
connectorId: 'test-connector-id',
});
mockProxiedRequest = jest.fn().mockResolvedValue(mockResponse);
// @ts-ignore
connector.request = mockProxiedRequest;
jest.clearAllMocks();
});
it('verifies that the OpenAI client is initialized with the custom proxy HTTP agent', () => {
// @ts-ignore .openAI is private
const openAIClient = connector.openAI;
// Verify the client was initialized with the custom agent configuration
expect(openAIClient).toBeDefined();
expect(openAIClient.httpAgent).toBeDefined();
expect(openAIClient.httpAgent.proxy).toBeDefined();
expect(openAIClient.httpAgent.proxy.host).toBe(PROXY_HOST);
expect(openAIClient.httpAgent.proxy.port).toBe(80);
});
it('verifies that requests use the configured HTTP agent', async () => {
// Make a test request
const response = await connector.runApi(
{ body: JSON.stringify(sampleOpenAiBody) },
connectorUsageCollector
);
expect(mockProxiedRequest).toBeCalledTimes(1);
expect(mockProxiedRequest).toHaveBeenCalledWith(
{
...mockDefaults,
signal: undefined,
data: JSON.stringify({
...sampleOpenAiBody,
stream: false,
model: DEFAULT_OPENAI_MODEL,
}),
headers: {
Authorization: 'Bearer 123',
'X-My-Custom-Header': 'foo',
'content-type': 'application/json',
'OpenAI-Organization': 'org-id',
'OpenAI-Project': 'proj-id',
},
},
connectorUsageCollector
);
expect(response).toEqual(mockResponse.data);
});
});

View file

@ -18,6 +18,7 @@ import type {
} from 'openai/resources/chat/completions'; } from 'openai/resources/chat/completions';
import type { Stream } from 'openai/streaming'; import type { Stream } from 'openai/streaming';
import type { ConnectorUsageCollector } from '@kbn/actions-plugin/server/types'; import type { ConnectorUsageCollector } from '@kbn/actions-plugin/server/types';
import { getCustomAgents } from '@kbn/actions-plugin/server/lib/get_custom_agents';
import { removeEndpointFromUrl } from './lib/openai_utils'; import { removeEndpointFromUrl } from './lib/openai_utils';
import { import {
RunActionParamsSchema, RunActionParamsSchema,
@ -64,7 +65,6 @@ export class OpenAIConnector extends SubActionConnector<Config, Secrets> {
constructor(params: ServiceParams<Config, Secrets>) { constructor(params: ServiceParams<Config, Secrets>) {
super(params); super(params);
this.url = this.config.apiUrl; this.url = this.config.apiUrl;
this.provider = this.config.apiProvider; this.provider = this.config.apiProvider;
// apiKey could be undefined if PKI, use mock value when this is the case // apiKey could be undefined if PKI, use mock value when this is the case
@ -77,6 +77,12 @@ export class OpenAIConnector extends SubActionConnector<Config, Secrets> {
...('projectId' in this.config ? { 'OpenAI-Project': this.config.projectId } : {}), ...('projectId' in this.config ? { 'OpenAI-Project': this.config.projectId } : {}),
}; };
const { httpAgent, httpsAgent } = getCustomAgents(
this.configurationUtilities,
this.logger,
this.url
);
this.openAI = this.openAI =
this.config.apiProvider === OpenAiProviderType.AzureAi this.config.apiProvider === OpenAiProviderType.AzureAi
? new OpenAI({ ? new OpenAI({
@ -87,6 +93,7 @@ export class OpenAIConnector extends SubActionConnector<Config, Secrets> {
...this.headers, ...this.headers,
'api-key': this.key, 'api-key': this.key,
}, },
httpAgent: httpsAgent ?? httpAgent,
}) })
: new OpenAI({ : new OpenAI({
baseURL: removeEndpointFromUrl(this.config.apiUrl), baseURL: removeEndpointFromUrl(this.config.apiUrl),
@ -94,6 +101,7 @@ export class OpenAIConnector extends SubActionConnector<Config, Secrets> {
defaultHeaders: { defaultHeaders: {
...this.headers, ...this.headers,
}, },
httpAgent: httpsAgent ?? httpAgent,
}); });
this.registerSubActions(); this.registerSubActions();