Fix InvokeAIRaw bedrock subaction's schema (elastic#192015)

## Summary Fix regression on `InvokeAIRaw` bedrock subaction introduced by elastic#191434
nreese · Sep 3, 2024 · 84c6815 · 84c6815
1 parent 448c9e1
commit 84c6815
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 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
diff --git 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
@@ -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
@@ -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;