Skip to content

Commit 2564d6d

Browse files
[ML] Fix OpenAI connector does not use the action proxy configuration for all subactions (#219617)
## Summary This PR fixes #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>
1 parent aa488b4 commit 2564d6d

4 files changed

Lines changed: 155 additions & 2 deletions

File tree

‎x-pack/platform/plugins/shared/actions/server/lib/get_custom_agents.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,9 @@ import type { ActionsConfigurationUtilities } from '../actions_config';
1515
import { getNodeSSLOptions, getSSLSettingsFromConfig } from './get_node_ssl_options';
1616
import type { SSLSettings } from '../types';
1717

18+
/**
19+
* Create http and https proxy agents with custom proxy /hosts/SSL settings from configurationUtilities
20+
*/
1821
interface GetCustomAgentsResponse {
1922
httpAgent: HttpAgent | undefined;
2023
httpsAgent: HttpsAgent | undefined;

‎x-pack/platform/plugins/shared/actions/server/sub_action_framework/sub_action_connector.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ export abstract class SubActionConnector<Config, Secrets> {
4444
[k: string]: ((params: unknown) => unknown) | unknown;
4545
private axiosInstance: AxiosInstance;
4646
private subActions: Map<string, SubAction> = new Map();
47-
private configurationUtilities: ActionsConfigurationUtilities;
47+
protected configurationUtilities: ActionsConfigurationUtilities;
4848
protected readonly kibanaRequest?: KibanaRequest;
4949
protected logger: Logger;
5050
protected esClient: ElasticsearchClient;
Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
import { DEFAULT_TIMEOUT_MS, OPENAI_CONNECTOR_ID } from '../../../common/openai/constants';
9+
import { actionsMock } from '@kbn/actions-plugin/server/mocks';
10+
import { DEFAULT_OPENAI_MODEL } from '../../../common/openai/constants';
11+
import { actionsConfigMock } from '@kbn/actions-plugin/server/actions_config.mock';
12+
import { OpenAIConnector } from './openai';
13+
import { OpenAiProviderType } from '../../../common/openai/constants';
14+
import { loggingSystemMock } from '@kbn/core/server/mocks';
15+
import { ConnectorUsageCollector } from '@kbn/actions-plugin/server/types';
16+
import { RunActionResponseSchema } from '../../../common/openai/schema';
17+
18+
const logger = loggingSystemMock.createLogger();
19+
20+
// Mock an instance of the OpenAI class
21+
// with overridden flag for purpose of jest test
22+
jest.mock('openai', () => {
23+
const UnmodifiedOpenAIClient = jest.requireActual('openai').default;
24+
25+
return {
26+
__esModule: true,
27+
default: jest.fn().mockImplementation((config) => {
28+
return new UnmodifiedOpenAIClient({
29+
...config,
30+
dangerouslyAllowBrowser: true,
31+
});
32+
}),
33+
};
34+
});
35+
describe('OpenAI with proxy config', () => {
36+
let mockProxiedRequest: jest.Mock;
37+
let connectorUsageCollector: ConnectorUsageCollector;
38+
const mockDefaults = {
39+
timeout: DEFAULT_TIMEOUT_MS,
40+
url: 'https://api.openai.com/v1/chat/completions',
41+
method: 'post',
42+
responseSchema: RunActionResponseSchema,
43+
};
44+
45+
const mockResponse = {
46+
headers: {},
47+
data: {},
48+
};
49+
50+
const configurationUtilities = actionsConfigMock.create();
51+
const PROXY_HOST = 'proxy.custom.elastic.co';
52+
const PROXY_URL = `http://${PROXY_HOST}`;
53+
54+
configurationUtilities.getProxySettings.mockReturnValue({
55+
proxyUrl: PROXY_URL,
56+
proxySSLSettings: {
57+
verificationMode: 'none',
58+
},
59+
proxyBypassHosts: undefined,
60+
proxyOnlyHosts: undefined,
61+
});
62+
63+
const connector = new OpenAIConnector({
64+
configurationUtilities,
65+
connector: { id: '1', type: OPENAI_CONNECTOR_ID },
66+
config: {
67+
apiUrl: 'https://api.openai.com/v1/chat/completions',
68+
apiProvider: OpenAiProviderType.OpenAi,
69+
defaultModel: DEFAULT_OPENAI_MODEL,
70+
organizationId: 'org-id',
71+
projectId: 'proj-id',
72+
headers: {
73+
'X-My-Custom-Header': 'foo',
74+
Authorization: 'override',
75+
},
76+
},
77+
secrets: { apiKey: '123' },
78+
logger,
79+
services: actionsMock.createServices(),
80+
});
81+
82+
const sampleOpenAiBody = {
83+
messages: [
84+
{
85+
role: 'user',
86+
content: 'Hello world',
87+
},
88+
],
89+
};
90+
91+
beforeEach(() => {
92+
connectorUsageCollector = new ConnectorUsageCollector({
93+
logger,
94+
connectorId: 'test-connector-id',
95+
});
96+
mockProxiedRequest = jest.fn().mockResolvedValue(mockResponse);
97+
// @ts-ignore
98+
connector.request = mockProxiedRequest;
99+
jest.clearAllMocks();
100+
});
101+
102+
it('verifies that the OpenAI client is initialized with the custom proxy HTTP agent', () => {
103+
// @ts-ignore .openAI is private
104+
const openAIClient = connector.openAI;
105+
106+
// Verify the client was initialized with the custom agent configuration
107+
expect(openAIClient).toBeDefined();
108+
expect(openAIClient.httpAgent).toBeDefined();
109+
expect(openAIClient.httpAgent.proxy).toBeDefined();
110+
expect(openAIClient.httpAgent.proxy.host).toBe(PROXY_HOST);
111+
expect(openAIClient.httpAgent.proxy.port).toBe(80);
112+
});
113+
114+
it('verifies that requests use the configured HTTP agent', async () => {
115+
// Make a test request
116+
const response = await connector.runApi(
117+
{ body: JSON.stringify(sampleOpenAiBody) },
118+
connectorUsageCollector
119+
);
120+
expect(mockProxiedRequest).toBeCalledTimes(1);
121+
expect(mockProxiedRequest).toHaveBeenCalledWith(
122+
{
123+
...mockDefaults,
124+
signal: undefined,
125+
data: JSON.stringify({
126+
...sampleOpenAiBody,
127+
stream: false,
128+
model: DEFAULT_OPENAI_MODEL,
129+
}),
130+
headers: {
131+
Authorization: 'Bearer 123',
132+
'X-My-Custom-Header': 'foo',
133+
'content-type': 'application/json',
134+
'OpenAI-Organization': 'org-id',
135+
'OpenAI-Project': 'proj-id',
136+
},
137+
},
138+
connectorUsageCollector
139+
);
140+
expect(response).toEqual(mockResponse.data);
141+
});
142+
});

‎x-pack/platform/plugins/shared/stack_connectors/server/connector_types/openai/openai.ts‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ import type {
1818
} from 'openai/resources/chat/completions';
1919
import type { Stream } from 'openai/streaming';
2020
import type { ConnectorUsageCollector } from '@kbn/actions-plugin/server/types';
21+
import { getCustomAgents } from '@kbn/actions-plugin/server/lib/get_custom_agents';
2122
import { removeEndpointFromUrl } from './lib/openai_utils';
2223
import {
2324
RunActionParamsSchema,
@@ -64,7 +65,6 @@ export class OpenAIConnector extends SubActionConnector<Config, Secrets> {
6465

6566
constructor(params: ServiceParams<Config, Secrets>) {
6667
super(params);
67-
6868
this.url = this.config.apiUrl;
6969
this.provider = this.config.apiProvider;
7070
// 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> {
7777
...('projectId' in this.config ? { 'OpenAI-Project': this.config.projectId } : {}),
7878
};
7979

80+
const { httpAgent, httpsAgent } = getCustomAgents(
81+
this.configurationUtilities,
82+
this.logger,
83+
this.url
84+
);
85+
8086
this.openAI =
8187
this.config.apiProvider === OpenAiProviderType.AzureAi
8288
? new OpenAI({
@@ -87,13 +93,15 @@ export class OpenAIConnector extends SubActionConnector<Config, Secrets> {
8793
...this.headers,
8894
'api-key': this.key,
8995
},
96+
httpAgent: httpsAgent ?? httpAgent,
9097
})
9198
: new OpenAI({
9299
baseURL: removeEndpointFromUrl(this.config.apiUrl),
93100
apiKey: this.key,
94101
defaultHeaders: {
95102
...this.headers,
96103
},
104+
httpAgent: httpsAgent ?? httpAgent,
97105
});
98106

99107
this.registerSubActions();

0 commit comments

Comments
 (0)