changes from review

Author: narengogi

src/providers/bedrock/index.ts

@@ -203,6 +203,9 @@ const BedrockConfig: ProviderConfigs = {
       if (!config.chatComplete) {
         config.chatComplete = BedrockConverseChatCompleteConfig;
       }
+      if (!config.messages) {
+        config.messages = BedrockConverseMessagesConfig;
+      }
       if (!config.responseTransforms?.['stream-chatComplete']) {
         config.responseTransforms = {
           ...(config.responseTransforms ?? {}),

src/providers/bedrock/messages.ts

@@ -11,7 +11,7 @@ import {
 import {
   ContentBlock,
   MessagesResponse,
-  STOP_REASON,
+  ANTHROPIC_STOP_REASON,
 } from '../../types/messagesResponse';
 import { RawContentBlockDeltaEvent } from '../../types/MessagesStreamResponse';
 import {
@@ -22,7 +22,10 @@ import {
   ANTHROPIC_MESSAGE_STOP_EVENT,
 } from '../anthropic-base/constants';
 import { ErrorResponse, ProviderConfig } from '../types';
-import { generateInvalidProviderResponseError } from '../utils';
+import {
+  generateInvalidProviderResponseError,
+  transformToAnthropicStopReason,
+} from '../utils';
 import { BedrockErrorResponseTransform } from './chatComplete';
 import { BedrockErrorResponse } from './embed';
 import {
@@ -185,7 +188,6 @@ const appendToolResultBlock = (
           text: item.text,
         });
       } else if (item.type === 'image') {
-        // TODO: test this
         appendImageBlock(transformedToolResultContent, item);
       }
     }
@@ -419,8 +421,7 @@ export const BedrockMessagesResponseTransform = (
       type: 'message',
       role: 'assistant',
       content: transformedContent,
-      // TODO: pull changes from stop reason transformation PR
-      stop_reason: response.stopReason as STOP_REASON,
+      stop_reason: transformToAnthropicStopReason(response.stopReason),
       usage: {
         cache_read_input_tokens: response.usage.cacheReadInputTokens,
         cache_creation_input_tokens: response.usage.cacheWriteInputTokens,
@@ -535,15 +536,16 @@ export const BedrockConverseMessagesStreamChunkTransform = (
       parsedChunk.usage.cacheReadInputTokens;
     messageDeltaEvent.usage.cache_creation_input_tokens =
       parsedChunk.usage.cacheWriteInputTokens;
-    messageDeltaEvent.delta.stop_reason = streamState.stopReason || '';
+    messageDeltaEvent.delta.stop_reason = transformToAnthropicStopReason(
+      streamState.stopReason
+    );
     const contentBlockStopEvent = { ...ANTHROPIC_CONTENT_BLOCK_STOP_EVENT };
     contentBlockStopEvent.index = streamState.currentContentBlockIndex;
     let returnChunk = `event: content_block_stop\ndata: ${JSON.stringify(contentBlockStopEvent)}\n\n`;
     returnChunk += `event: message_delta\ndata: ${JSON.stringify(messageDeltaEvent)}\n\n`;
     returnChunk += `event: message_stop\ndata: ${JSON.stringify(ANTHROPIC_MESSAGE_STOP_EVENT)}\n\n`;
     return returnChunk;
   }
-  // console.log(JSON.stringify(parsedChunk, null, 2));
 };
 
 function getMessageStartEvent(fallbackId: string, gatewayRequest: Params<any>) {

src/providers/bedrock/types.ts

@@ -108,7 +108,7 @@ export interface BedrockChatCompletionResponse {
       content: BedrockContentItem[];
     };
   };
-  stopReason: string;
+  stopReason: BEDROCK_STOP_REASON;
   usage: {
     inputTokens: number;
     outputTokens: number;
@@ -154,8 +154,9 @@ export type BedrockContentItem = {
     type: string;
   };
 };
+
 export interface BedrockStreamState {
-  stopReason?: string;
+  stopReason?: BEDROCK_STOP_REASON;
   currentToolCallIndex?: number;
   currentContentBlockIndex?: number;
 }
@@ -185,7 +186,7 @@ export interface BedrockChatCompleteStreamChunk {
       input?: object;
     };
   };
-  stopReason?: string;
+  stopReason?: BEDROCK_STOP_REASON;
   metrics?: {
     latencyMs: number;
   };

src/providers/utils.ts

@@ -1,5 +1,9 @@
+import { ANTHROPIC_STOP_REASON } from './anthropic/types';
 import { FINISH_REASON, ErrorResponse, PROVIDER_FINISH_REASON } from './types';
-import { finishReasonMap } from './utils/finishReasonMap';
+import {
+  AnthropicFinishReasonMap,
+  finishReasonMap,
+} from './utils/finishReasonMap';
 
 export const generateInvalidProviderResponseError: (
   response: Record<string, any>,
@@ -78,3 +82,19 @@ export const transformFinishReason = (
   }
   return transformedFinishReason;
 };
+
+/*
+  Transforms the finish reason from the provider to the finish reason used by the Anthropic API.
+  If the finish reason is not found in the map, it will return the stop reason.
+  NOTE: this function always returns a finish reason
+*/
+export const transformToAnthropicStopReason = (
+  finishReason?: PROVIDER_FINISH_REASON
+): ANTHROPIC_STOP_REASON => {
+  if (!finishReason) return ANTHROPIC_STOP_REASON.end_turn;
+  const transformedFinishReason = AnthropicFinishReasonMap.get(finishReason);
+  if (!transformedFinishReason) {
+    return ANTHROPIC_STOP_REASON.end_turn;
+  }
+  return transformedFinishReason;
+};

src/providers/utils/finishReasonMap.ts

@@ -2,6 +2,7 @@ import { ANTHROPIC_STOP_REASON } from '../anthropic/types';
 import { FINISH_REASON, PROVIDER_FINISH_REASON } from '../types';
 import { BEDROCK_STOP_REASON } from '../bedrock/types';
 
+// TODO: rename this to OpenAIFinishReasonMap
 export const finishReasonMap = new Map<PROVIDER_FINISH_REASON, FINISH_REASON>([
   // https://docs.anthropic.com/en/api/messages#response-stop-reason
   [ANTHROPIC_STOP_REASON.stop_sequence, FINISH_REASON.stop],
@@ -17,3 +18,16 @@ export const finishReasonMap = new Map<PROVIDER_FINISH_REASON, FINISH_REASON>([
   [BEDROCK_STOP_REASON.guardrail_intervened, FINISH_REASON.content_filter],
   [BEDROCK_STOP_REASON.content_filtered, FINISH_REASON.content_filter],
 ]);
+
+export const AnthropicFinishReasonMap = new Map<
+  PROVIDER_FINISH_REASON,
+  ANTHROPIC_STOP_REASON
+>([
+  // https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Converse.html#API_runtime_Converse_ResponseSyntax
+  [BEDROCK_STOP_REASON.end_turn, ANTHROPIC_STOP_REASON.end_turn],
+  [BEDROCK_STOP_REASON.tool_use, ANTHROPIC_STOP_REASON.tool_use],
+  [BEDROCK_STOP_REASON.max_tokens, ANTHROPIC_STOP_REASON.max_tokens],
+  [BEDROCK_STOP_REASON.stop_sequence, ANTHROPIC_STOP_REASON.stop_sequence],
+  [BEDROCK_STOP_REASON.guardrail_intervened, ANTHROPIC_STOP_REASON.end_turn],
+  [BEDROCK_STOP_REASON.content_filtered, ANTHROPIC_STOP_REASON.end_turn],
+]);

src/types/MessagesStreamResponse.ts

@@ -6,7 +6,7 @@ import {
   MessagesResponse,
   RedactedThinkingBlock,
   ServerToolUseBlock,
-  STOP_REASON,
+  ANTHROPIC_STOP_REASON,
   TextBlock,
   ThinkingBlock,
   ToolUseBlock,
@@ -21,7 +21,7 @@ export interface RawMessageStartEvent {
 }
 
 export interface RawMessageDelta {
-  stop_reason: STOP_REASON | null;
+  stop_reason: ANTHROPIC_STOP_REASON | null;
 
   stop_sequence: string | null;
 }

src/types/messagesResponse.ts

@@ -150,7 +150,7 @@ export type ContentBlock =
   | ThinkingBlock
   | RedactedThinkingBlock;
 
-export enum STOP_REASON {
+export enum ANTHROPIC_STOP_REASON {
   end_turn = 'end_turn',
   max_tokens = 'max_tokens',
   stop_sequence = 'stop_sequence',
@@ -234,7 +234,7 @@ export interface MessagesResponse {
    * In non-streaming mode this value is always non-null. In streaming mode, it is
    * null in the `message_start` event and non-null otherwise.
    */
-  stop_reason: STOP_REASON | null;
+  stop_reason: ANTHROPIC_STOP_REASON | null;
 
   /**
    * Which custom stop sequence was generated, if any.