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

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

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

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

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

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

@@ -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);
+  });
+});

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

@@ -18,6 +18,7 @@ import type {
 } from 'openai/resources/chat/completions';
 import type { Stream } from 'openai/streaming';
 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 {
   RunActionParamsSchema,
@@ -64,7 +65,6 @@ export class OpenAIConnector extends SubActionConnector<Config, Secrets> {
 
   constructor(params: ServiceParams<Config, Secrets>) {
     super(params);
-
     this.url = this.config.apiUrl;
     this.provider = this.config.apiProvider;
     // 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 } : {}),
     };
 
+    const { httpAgent, httpsAgent } = getCustomAgents(
+      this.configurationUtilities,
+      this.logger,
+      this.url
+    );
+
     this.openAI =
       this.config.apiProvider === OpenAiProviderType.AzureAi
         ? new OpenAI({
@@ -87,13 +93,15 @@ export class OpenAIConnector extends SubActionConnector<Config, Secrets> {
               ...this.headers,
               'api-key': this.key,
             },
+            httpAgent: httpsAgent ?? httpAgent,
           })
         : new OpenAI({
             baseURL: removeEndpointFromUrl(this.config.apiUrl),
             apiKey: this.key,
             defaultHeaders: {
               ...this.headers,
             },
+            httpAgent: httpsAgent ?? httpAgent,
           });
 
     this.registerSubActions();