testing-proxy/src/lib/providers/bedrock.js

import {
  BedrockRuntimeClient,
  ConverseCommand,
  ConverseStreamCommand
} from '@aws-sdk/client-bedrock-runtime';

/**
 * Amazon Bedrock provider — translates OpenAI chat completion requests
 * to AWS Bedrock Converse API calls.
 */
export function createProvider(config) {
  const { id, region, credentials, models } = config;

  // Singleton BedrockRuntimeClient per provider config
  const client = new BedrockRuntimeClient({
    region,
    credentials
  });

  /**
   * Translate OpenAI messages[] to Bedrock Converse messages[].
   * Extracts system messages into a top-level `system` array.
   */
  function translateMessages(messages) {
    const systemPrompts = [];
    const bedrockMessages = [];

    for (const msg of messages) {
      if (msg.role === 'system') {
        systemPrompts.push({ text: msg.content });
      } else {
        bedrockMessages.push({
          role:    msg.role,
          content: [{ text: msg.content }]
        });
      }
    }

    return {
      system:   systemPrompts.length > 0 ? systemPrompts : undefined,
      messages: bedrockMessages
    };
  }

  /**
   * Build inferenceConfig from options, only including defined values.
   */
  function buildInferenceConfig(options) {
    const inferenceConfig = {};
    if (options.temperature !== undefined) inferenceConfig.temperature = options.temperature;
    if (options.max_tokens   !== undefined) inferenceConfig.maxTokens   = options.max_tokens;
    if (options.top_p        !== undefined) inferenceConfig.topP        = options.top_p;
    return Object.keys(inferenceConfig).length > 0 ? inferenceConfig : undefined;
  }

  /**
   * Generate a unique OpenAI-style chat completion ID.
   */
  function generateId() {
    return `chatcmpl-${Date.now()}`;
  }

  /**
   * Build common OpenAI response metadata.
   */
  function buildMetadata(modelId) {
    return {
      id:      generateId(),
      object:  'chat.completion',
      created: Math.floor(Date.now() / 1000),
      model:   modelId
    };
  }

  /**
   * Non-streaming chat completion.
   */
  async function chat(messages, modelId, options = {}) {
    const providerModelId = models[modelId];
    const { system, messages: bedrockMessages } = translateMessages(messages);
    const inferenceConfig = buildInferenceConfig(options);

    const command = new ConverseCommand({
      modelId:         providerModelId,
      messages:        bedrockMessages,
      system,
      inferenceConfig
    });

    let response;
    try {
      response = await client.send(command);
    } catch (error) {
      throw new Error(`Bedrock provider request failed: ${error.message}`);
    }

    const text = response.output?.message?.content?.[0]?.text ?? '';
    const usage = response.usage || {};

    return {
      ...buildMetadata(modelId),
      choices: [{
        index: 0,
        message: {
          role:    'assistant',
          content: text
        },
        finish_reason: 'stop'
      }],
      usage: {
        prompt_tokens:     usage.inputTokens  || 0,
        completion_tokens: usage.outputTokens || 0,
        total_tokens:      usage.totalTokens  || 0
      }
    };
  }

  /**
   * Streaming chat completion — returns an async iterator of OpenAI chunks.
   */
  async function* chatStream(messages, modelId, options = {}) {
    const providerModelId = models[modelId];
    const { system, messages: bedrockMessages } = translateMessages(messages);
    const inferenceConfig = buildInferenceConfig(options);

    const command = new ConverseStreamCommand({
      modelId:         providerModelId,
      messages:        bedrockMessages,
      system,
      inferenceConfig
    });

    let streamResponse;
    try {
      streamResponse = await client.send(command);
    } catch (error) {
      throw new Error(`Bedrock provider stream request failed: ${error.message}`);
    }

    const metadata = {
      id:      generateId(),
      object:  'chat.completion.chunk',
      created: Math.floor(Date.now() / 1000),
      model:   modelId
    };

    for await (const event of streamResponse.stream) {
      if (event.contentBlockDelta?.delta?.text) {
        yield {
          ...metadata,
          choices: [{
            index: 0,
            delta: { content: event.contentBlockDelta.delta.text },
            finish_reason: null
          }]
        };
      }

      if (event.messageStop?.stopReason) {
        yield {
          ...metadata,
          choices: [{
            index: 0,
            delta: {},
            finish_reason: 'stop'
          }]
        };
      }
    }
  }

  return { id, chat, chatStream };
}