[AI Connector] Change completion subAction schema to be OpenAI compatible

x-pack/plugins/stack_connectors/common/inference/constants.ts

@@ -33,10 +33,12 @@ export enum ServiceProviderKeys {
 export const INFERENCE_CONNECTOR_ID = '.inference';
 export enum SUB_ACTION {
   COMPLETION = 'completion',
+  UNIFIED_COMPLETION = 'unified_completion',
   RERANK = 'rerank',
   TEXT_EMBEDDING = 'text_embedding',
   SPARSE_EMBEDDING = 'sparse_embedding',
   COMPLETION_STREAM = 'completion_stream',
+  COMPLETION_ASYNC_ITERATOR = 'completion_async_iterator',
 }
 
 export const DEFAULT_PROVIDER = 'openai';

x-pack/plugins/stack_connectors/common/inference/schema.ts

@@ -23,6 +23,162 @@ export const ChatCompleteParamsSchema = schema.object({
   input: schema.string(),
 });
 
+// subset of OpenAI.ChatCompletionMessageParam https://github.com/openai/openai-node/blob/master/src/resources/chat/completions.ts
+const AIMessage = schema.object({
+  role: schema.string(),
+  content: schema.maybe(schema.string()),
+  name: schema.maybe(schema.string()),
+  tool_calls: schema.maybe(
+    schema.arrayOf(
+      schema.object({
+        id: schema.string(),
+        function: schema.object({
+          arguments: schema.maybe(schema.string()),
+          name: schema.maybe(schema.string()),
+        }),
+        type: schema.string(),
+      })
+    )
+  ),
+  tool_call_id: schema.maybe(schema.string()),
+});
+
+const AITool = schema.object({
+  type: schema.string(),
+  function: schema.object({
+    name: schema.string(),
+    description: schema.maybe(schema.string()),
+    parameters: schema.maybe(schema.recordOf(schema.string(), schema.any())),
+  }),
+});
+
+// subset of OpenAI.ChatCompletionCreateParamsBase https://github.com/openai/openai-node/blob/master/src/resources/chat/completions.ts
+export const UnifiedChatCompleteParamsSchema = schema.object({
+  body: schema.object({
+    messages: schema.arrayOf(AIMessage),
+    model: schema.maybe(schema.string()),
+    /**
+     * The maximum number of [tokens](/tokenizer) that can be generated in the chat
+     * completion. This value can be used to control
+     * [costs](https://openai.com/api/pricing/) for text generated via API.
+     *
+     * This value is now deprecated in favor of `max_completion_tokens`, and is not
+     * compatible with
+     * [o1 series models](https://platform.openai.com/docs/guides/reasoning).
+     */
+    max_tokens: schema.maybe(schema.number()),
+    /**
+     * Developer-defined tags and values used for filtering completions in the
+     * [dashboard](https://platform.openai.com/chat-completions).
+     */
+    metadata: schema.maybe(schema.recordOf(schema.string(), schema.string())),
+    /**
+     * How many chat completion choices to generate for each input message. Note that
+     * you will be charged based on the number of generated tokens across all of the
+     * choices. Keep `n` as `1` to minimize costs.
+     */
+    n: schema.maybe(schema.number()),
+    /**
+     * Up to 4 sequences where the API will stop generating further tokens.
+     */
+    stop: schema.maybe(
+      schema.nullable(schema.oneOf([schema.string(), schema.arrayOf(schema.string())]))
+    ),
+    /**
+     * What sampling temperature to use, between 0 and 2. Higher values like 0.8 will
+     * make the output more random, while lower values like 0.2 will make it more
+     * focused and deterministic.
+     *
+     * We generally recommend altering this or `top_p` but not both.
+     */
+    temperature: schema.maybe(schema.number()),
+    /**
+     * Controls which (if any) tool is called by the model. `none` means the model will
+     * not call any tool and instead generates a message. `auto` means the model can
+     * pick between generating a message or calling one or more tools. `required` means
+     * the model must call one or more tools. Specifying a particular tool via
+     * `{"type": "function", "function": {"name": "my_function"}}` forces the model to
+     * call that tool.
+     *
+     * `none` is the default when no tools are present. `auto` is the default if tools
+     * are present.
+     */
+    tool_choice: schema.maybe(
+      schema.oneOf([
+        schema.string(),
+        schema.object({
+          type: schema.string(),
+          function: schema.object({
+            name: schema.string(),
+          }),
+        }),
+      ])
+    ),
+    /**
+     * A list of tools the model may call. Currently, only functions are supported as a
+     * tool. Use this to provide a list of functions the model may generate JSON inputs
+     * for. A max of 128 functions are supported.
+     */
+    tools: schema.maybe(schema.arrayOf(AITool)),
+    /**
+     * An alternative to sampling with temperature, called nucleus sampling, where the
+     * model considers the results of the tokens with top_p probability mass. So 0.1
+     * means only the tokens comprising the top 10% probability mass are considered.
+     *
+     * We generally recommend altering this or `temperature` but not both.
+     */
+    top_p: schema.maybe(schema.number()),
+    /**
+     * A unique identifier representing your end-user, which can help OpenAI to monitor
+     * and detect abuse.
+     * [Learn more](https://platform.openai.com/docs/guides/safety-best-practices/end-user-ids).
+     */
+    user: schema.maybe(schema.string()),
+  }),
+});
+
+export const UnifiedChatCompleteResponseSchema = schema.object({
+  id: schema.string(),
+  choices: schema.arrayOf(
+    schema.object({
+      finish_reason: schema.oneOf([
+        schema.literal('stop'),
+        schema.literal('length'),
+        schema.literal('tool_calls'),
+        schema.literal('content_filter'),
+        schema.literal('function_call'),
+      ]),
+      index: schema.number(),
+      message: schema.object({
+        content: schema.nullable(schema.string()),
+        refusal: schema.nullable(schema.string()),
+        role: schema.string(),
+        tool_calls: schema.arrayOf(
+          schema.object({
+            id: schema.string(),
+            function: schema.object({
+              arguments: schema.maybe(schema.string()),
+              name: schema.maybe(schema.string()),
+            }),
+            type: schema.string(),
+          })
+        ),
+      }),
+    }),
+    { defaultValue: [] }
+  ),
+  created: schema.maybe(schema.number()),
+  model: schema.maybe(schema.string()),
+  object: schema.maybe(schema.string()),
+  usage: schema.maybe(
+    schema.object({
+      completion_tokens: schema.maybe(schema.number()),
+      prompt_tokens: schema.maybe(schema.number()),
+      total_tokens: schema.maybe(schema.number()),
+    })
+  ),
+});
+
 export const ChatCompleteResponseSchema = schema.arrayOf(
   schema.object({
     result: schema.string(),
@@ -66,3 +222,12 @@ export const TextEmbeddingResponseSchema = schema.arrayOf(
 );
 
 export const StreamingResponseSchema = schema.stream();
+
+// Run action schema
+export const DashboardActionParamsSchema = schema.object({
+  dashboardId: schema.string(),
+});
+
+export const DashboardActionResponseSchema = schema.object({
+  available: schema.boolean(),
+});

x-pack/plugins/stack_connectors/common/inference/types.ts

@@ -18,12 +18,19 @@ import {
   SparseEmbeddingResponseSchema,
   TextEmbeddingParamsSchema,
   TextEmbeddingResponseSchema,
+  UnifiedChatCompleteParamsSchema,
+  UnifiedChatCompleteResponseSchema,
+  DashboardActionParamsSchema,
+  DashboardActionResponseSchema,
 } from './schema';
 import { ConfigProperties } from '../dynamic_config/types';
 
 export type Config = TypeOf<typeof ConfigSchema>;
 export type Secrets = TypeOf<typeof SecretsSchema>;
 
+export type UnifiedChatCompleteParams = TypeOf<typeof UnifiedChatCompleteParamsSchema>;
+export type UnifiedChatCompleteResponse = TypeOf<typeof UnifiedChatCompleteResponseSchema>;
+
 export type ChatCompleteParams = TypeOf<typeof ChatCompleteParamsSchema>;
 export type ChatCompleteResponse = TypeOf<typeof ChatCompleteResponseSchema>;
 
@@ -38,6 +45,9 @@ export type TextEmbeddingResponse = TypeOf<typeof TextEmbeddingResponseSchema>;
 
 export type StreamingResponse = TypeOf<typeof StreamingResponseSchema>;
 
+export type DashboardActionParams = TypeOf<typeof DashboardActionParamsSchema>;
+export type DashboardActionResponse = TypeOf<typeof DashboardActionResponseSchema>;
+
 export type FieldsConfiguration = Record<string, ConfigProperties>;
 
 export interface InferenceTaskType {

x-pack/plugins/stack_connectors/public/connector_types/inference/constants.tsx

@@ -25,13 +25,25 @@ export const DEFAULT_TEXT_EMBEDDING_BODY = {
   inputType: 'ingest',
 };
 
+export const DEFAULT_UNIFIED_CHAT_COMPLETE_BODY = {
+  body: {
+    messages: [
+      {
+        role: 'user',
+        content: 'Hello world',
+      },
+    ],
+  },
+};
+
 export const DEFAULTS_BY_TASK_TYPE: Record<string, unknown> = {
   [SUB_ACTION.COMPLETION]: DEFAULT_CHAT_COMPLETE_BODY,
+  [SUB_ACTION.UNIFIED_COMPLETION]: DEFAULT_UNIFIED_CHAT_COMPLETE_BODY,
   [SUB_ACTION.RERANK]: DEFAULT_RERANK_BODY,
   [SUB_ACTION.SPARSE_EMBEDDING]: DEFAULT_SPARSE_EMBEDDING_BODY,
   [SUB_ACTION.TEXT_EMBEDDING]: DEFAULT_TEXT_EMBEDDING_BODY,
 };
 
-export const DEFAULT_TASK_TYPE = 'completion';
+export const DEFAULT_TASK_TYPE = 'unified_completion';
 
 export const DEFAULT_PROVIDER = 'elasticsearch';

x-pack/plugins/stack_connectors/public/connector_types/inference/inference.test.tsx

@@ -55,6 +55,10 @@ describe('OpenAI action params validation', () => {
       subAction: SUB_ACTION.SPARSE_EMBEDDING,
       subActionParams: { input: 'message test' },
     },
+    {
+      subAction: SUB_ACTION.UNIFIED_COMPLETION,
+      subActionParams: { input: 'message test' },
+    },
   ])(
     'validation succeeds when params are valid for subAction $subAction',
     async ({ subAction, subActionParams }) => {

x-pack/plugins/stack_connectors/public/connector_types/inference/inference.tsx

@@ -19,6 +19,7 @@ import { InferenceActionParams, InferenceConnector } from './types';
 interface ValidationErrors {
   subAction: string[];
   input: string[];
+  body: string[];
   // rerank only
   query: string[];
   // text_embedding only
@@ -40,11 +41,18 @@ export function getConnectorType(): InferenceConnector {
       const translations = await import('./translations');
       const errors: ValidationErrors = {
         input: [],
+        body: [],
         subAction: [],
         inputType: [],
         query: [],
       };
 
+      if (subAction === SUB_ACTION.UNIFIED_COMPLETION) {
+        if (!subActionParams.body.messages.length) {
+          errors.body.push(translations.getRequiredMessage('Messages'));
+        }
+      }
+
       if (
         subAction === SUB_ACTION.RERANK ||
         subAction === SUB_ACTION.COMPLETION ||

x-pack/plugins/stack_connectors/public/connector_types/inference/params.tsx

@@ -17,6 +17,7 @@ import {
   RerankParams,
   SparseEmbeddingParams,
   TextEmbeddingParams,
+  UnifiedChatCompleteParams,
 } from '../../../common/inference/types';
 import { DEFAULTS_BY_TASK_TYPE } from './constants';
 import * as i18n from './translations';
@@ -28,25 +29,35 @@ const InferenceServiceParamsFields: React.FunctionComponent<
 > = ({ actionParams, editAction, index, errors, actionConnector }) => {
   const { subAction, subActionParams } = actionParams;
 
-  const { taskType } = (actionConnector as unknown as InferenceActionConnector).config;
+  const { taskType, provider } = (actionConnector as unknown as InferenceActionConnector).config;
 
   useEffect(() => {
     if (!subAction) {
-      editAction('subAction', taskType, index);
+      editAction(
+        'subAction',
+        provider === 'openai' && taskType === SUB_ACTION.COMPLETION
+          ? SUB_ACTION.UNIFIED_COMPLETION
+          : taskType,
+        index
+      );
     }
-  }, [editAction, index, subAction, taskType]);
+  }, [editAction, index, provider, subAction, taskType]);
 
   useEffect(() => {
     if (!subActionParams) {
       editAction(
         'subActionParams',
         {
-          ...(DEFAULTS_BY_TASK_TYPE[taskType] ?? {}),
+          ...(DEFAULTS_BY_TASK_TYPE[
+            provider === 'openai' && taskType === SUB_ACTION.COMPLETION
+              ? SUB_ACTION.UNIFIED_COMPLETION
+              : taskType
+          ] ?? {}),
         },
         index
       );
     }
-  }, [editAction, index, subActionParams, taskType]);
+  }, [editAction, index, provider, subActionParams, taskType]);
 
   const editSubActionParams = useCallback(
     (params: Partial<InferenceActionParams['subActionParams']>) => {
@@ -65,6 +76,16 @@ const InferenceServiceParamsFields: React.FunctionComponent<
     );
   }
 
+  if (subAction === SUB_ACTION.UNIFIED_COMPLETION) {
+    return (
+      <UnifiedCompletionParamsFields
+        errors={errors}
+        editSubActionParams={editSubActionParams}
+        subActionParams={subActionParams as UnifiedChatCompleteParams}
+      />
+    );
+  }
+
   if (subAction === SUB_ACTION.RERANK) {
     return (
       <RerankParamsFields
@@ -119,6 +140,32 @@ const InferenceInput: React.FunctionComponent<{
   );
 };
 
+const UnifiedCompletionParamsFields: React.FunctionComponent<{
+  subActionParams: UnifiedChatCompleteParams;
+  errors: RuleFormParamsErrors;
+  editSubActionParams: (params: Partial<InferenceActionParams['subActionParams']>) => void;
+}> = ({ subActionParams, editSubActionParams, errors }) => {
+  return (
+    <>
+      <JsonEditorWithMessageVariables
+        paramsProperty={'body'}
+        inputTargetValue={JSON.stringify(subActionParams.body)}
+        label={i18n.BODY}
+        errors={errors.body as string[]}
+        onDocumentsChange={(json: string) => {
+          editSubActionParams({ body: { ...JSON.parse(json) } });
+        }}
+        onBlur={() => {
+          if (!subActionParams.body) {
+            editSubActionParams({ body: { messages: [] } });
+          }
+        }}
+        dataTestSubj="inference-bodyJsonEditor"
+      />
+    </>
+  );
+};
+
 const CompletionParamsFields: React.FunctionComponent<{
   subActionParams: ChatCompleteParams;
   errors: RuleFormParamsErrors;

x-pack/plugins/stack_connectors/public/connector_types/inference/types.ts

@@ -13,10 +13,12 @@ import {
   RerankParams,
   SparseEmbeddingParams,
   TextEmbeddingParams,
+  UnifiedChatCompleteParams,
 } from '../../../common/inference/types';
 
 export type InferenceActionParams =
   | { subAction: SUB_ACTION.COMPLETION; subActionParams: ChatCompleteParams }
+  | { subAction: SUB_ACTION.UNIFIED_COMPLETION; subActionParams: UnifiedChatCompleteParams }
   | { subAction: SUB_ACTION.RERANK; subActionParams: RerankParams }
   | { subAction: SUB_ACTION.SPARSE_EMBEDDING; subActionParams: SparseEmbeddingParams }
   | { subAction: SUB_ACTION.TEXT_EMBEDDING; subActionParams: TextEmbeddingParams };