From 84c6815a1152049c74db00da48522d0624603856 Mon Sep 17 00:00:00 2001
From: Pierre Gayvallet <pierre.gayvallet@elastic.co>
Date: Tue, 3 Sep 2024 22:29:56 +0200
Subject: [PATCH] Fix `InvokeAIRaw` bedrock subaction's schema (#192015)

## Summary

Fix regression on `InvokeAIRaw` bedrock subaction introduced by
https://github.com/elastic/kibana/pull/191434
---
 .../connector_types.test.ts.snap              | 49 +------------------
 .../stack_connectors/common/bedrock/schema.ts |  7 ++-
 .../stack_connectors/common/bedrock/types.ts  |  2 +-
 .../server/connector_types/bedrock/bedrock.ts | 14 +++---
 4 files changed, 16 insertions(+), 56 deletions(-)

diff --git a/x-pack/plugins/actions/server/integration_tests/__snapshots__/connector_types.test.ts.snap b/x-pack/plugins/actions/server/integration_tests/__snapshots__/connector_types.test.ts.snap
index 6e4392e3f737e..c52ef2f2f6478 100644
--- a/x-pack/plugins/actions/server/integration_tests/__snapshots__/connector_types.test.ts.snap
+++ b/x-pack/plugins/actions/server/integration_tests/__snapshots__/connector_types.test.ts.snap
@@ -1083,51 +1083,14 @@ Object {
           "keys": Object {
             "content": Object {
               "flags": Object {
-                "default": [Function],
                 "error": [Function],
-                "presence": "optional",
               },
               "metas": Array [
                 Object {
-                  "x-oas-optional": true,
-                },
-              ],
-              "rules": Array [
-                Object {
-                  "args": Object {
-                    "method": [Function],
-                  },
-                  "name": "custom",
-                },
-              ],
-              "type": "string",
-            },
-            "rawContent": Object {
-              "flags": Object {
-                "default": [Function],
-                "error": [Function],
-                "presence": "optional",
-              },
-              "items": Array [
-                Object {
-                  "flags": Object {
-                    "error": [Function],
-                    "presence": "optional",
-                  },
-                  "metas": Array [
-                    Object {
-                      "x-oas-any-type": true,
-                    },
-                  ],
-                  "type": "any",
-                },
-              ],
-              "metas": Array [
-                Object {
-                  "x-oas-optional": true,
+                  "x-oas-any-type": true,
                 },
               ],
-              "type": "array",
+              "type": "any",
             },
             "role": Object {
               "flags": Object {
@@ -1144,14 +1107,6 @@ Object {
               "type": "string",
             },
           },
-          "rules": Array [
-            Object {
-              "args": Object {
-                "method": [Function],
-              },
-              "name": "custom",
-            },
-          ],
           "type": "object",
         },
       ],
diff --git a/x-pack/plugins/stack_connectors/common/bedrock/schema.ts b/x-pack/plugins/stack_connectors/common/bedrock/schema.ts
index 03f4f5cc01735..15ac45c0cf597 100644
--- a/x-pack/plugins/stack_connectors/common/bedrock/schema.ts
+++ b/x-pack/plugins/stack_connectors/common/bedrock/schema.ts
@@ -78,7 +78,12 @@ export const InvokeAIActionResponseSchema = schema.object({
 });
 
 export const InvokeAIRawActionParamsSchema = schema.object({
-  messages: schema.arrayOf(BedrockMessageSchema),
+  messages: schema.arrayOf(
+    schema.object({
+      role: schema.string(),
+      content: schema.any(),
+    })
+  ),
   model: schema.maybe(schema.string()),
   temperature: schema.maybe(schema.number()),
   stopSequences: schema.maybe(schema.arrayOf(schema.string())),
diff --git a/x-pack/plugins/stack_connectors/common/bedrock/types.ts b/x-pack/plugins/stack_connectors/common/bedrock/types.ts
index 3b02f40d2de62..9d742e5f892a8 100644
--- a/x-pack/plugins/stack_connectors/common/bedrock/types.ts
+++ b/x-pack/plugins/stack_connectors/common/bedrock/types.ts
@@ -35,5 +35,5 @@ export type RunActionResponse = TypeOf<typeof RunActionResponseSchema>;
 export type StreamingResponse = TypeOf<typeof StreamingResponseSchema>;
 export type DashboardActionParams = TypeOf<typeof DashboardActionParamsSchema>;
 export type DashboardActionResponse = TypeOf<typeof DashboardActionResponseSchema>;
-export type BedRockMessage = TypeOf<typeof BedrockMessageSchema>;
+export type BedrockMessage = TypeOf<typeof BedrockMessageSchema>;
 export type BedrockToolChoice = TypeOf<typeof BedrockToolChoiceSchema>;
diff --git a/x-pack/plugins/stack_connectors/server/connector_types/bedrock/bedrock.ts b/x-pack/plugins/stack_connectors/server/connector_types/bedrock/bedrock.ts
index c2c773bdeaf87..9bd5c64404f64 100644
--- a/x-pack/plugins/stack_connectors/server/connector_types/bedrock/bedrock.ts
+++ b/x-pack/plugins/stack_connectors/server/connector_types/bedrock/bedrock.ts
@@ -32,7 +32,7 @@ import type {
   InvokeAIRawActionParams,
   InvokeAIRawActionResponse,
   RunApiLatestResponse,
-  BedRockMessage,
+  BedrockMessage,
   BedrockToolChoice,
 } from '../../../common/bedrock/types';
 import {
@@ -422,7 +422,7 @@ const formatBedrockBody = ({
   tools,
   toolChoice,
 }: {
-  messages: BedRockMessage[];
+  messages: BedrockMessage[];
   stopSequences?: string[];
   temperature?: number;
   maxTokens?: number;
@@ -440,9 +440,9 @@ const formatBedrockBody = ({
   tool_choice: toolChoice,
 });
 
-interface FormattedBedRockMessage {
+interface FormattedBedrockMessage {
   role: string;
-  content: string | BedRockMessage['rawContent'];
+  content: string | BedrockMessage['rawContent'];
 }
 
 /**
@@ -452,15 +452,15 @@ interface FormattedBedRockMessage {
  * @param messages
  */
 const ensureMessageFormat = (
-  messages: BedRockMessage[],
+  messages: BedrockMessage[],
   systemPrompt?: string
 ): {
-  messages: FormattedBedRockMessage[];
+  messages: FormattedBedrockMessage[];
   system?: string;
 } => {
   let system = systemPrompt ? systemPrompt : '';
 
-  const newMessages = messages.reduce<FormattedBedRockMessage[]>((acc, m) => {
+  const newMessages = messages.reduce<FormattedBedrockMessage[]>((acc, m) => {
     if (m.role === 'system') {
       system = `${system.length ? `${system}\n` : ''}${m.content}`;
       return acc;