bedrock.ts 2.0 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950
  1. import AnthropicBedrock from "@anthropic-ai/bedrock-sdk"
  2. import { Anthropic } from "@anthropic-ai/sdk"
  3. import { ApiHandler, ApiHandlerMessageResponse } from "."
  4. import { ApiHandlerOptions, bedrockDefaultModelId, BedrockModelId, bedrockModels, ModelInfo } from "../shared/api"
  5. // https://docs.anthropic.com/en/api/claude-on-amazon-bedrock
  6. export class AwsBedrockHandler implements ApiHandler {
  7. private options: ApiHandlerOptions
  8. private client: AnthropicBedrock
  9. constructor(options: ApiHandlerOptions) {
  10. this.options = options
  11. this.client = new AnthropicBedrock({
  12. // Authenticate by either providing the keys below or use the default AWS credential providers, such as
  13. // using ~/.aws/credentials or the "AWS_SECRET_ACCESS_KEY" and "AWS_ACCESS_KEY_ID" environment variables.
  14. ...(this.options.awsAccessKey ? { awsAccessKey: this.options.awsAccessKey } : {}),
  15. ...(this.options.awsSecretKey ? { awsSecretKey: this.options.awsSecretKey } : {}),
  16. ...(this.options.awsSessionToken ? { awsSessionToken: this.options.awsSessionToken } : {}),
  17. // awsRegion changes the aws region to which the request is made. By default, we read AWS_REGION,
  18. // and if that's not present, we default to us-east-1. Note that we do not read ~/.aws/config for the region.
  19. awsRegion: this.options.awsRegion,
  20. })
  21. }
  22. async createMessage(
  23. systemPrompt: string,
  24. messages: Anthropic.Messages.MessageParam[],
  25. tools: Anthropic.Messages.Tool[]
  26. ): Promise<ApiHandlerMessageResponse> {
  27. const message = await this.client.messages.create({
  28. model: this.getModel().id,
  29. max_tokens: this.getModel().info.maxTokens,
  30. system: systemPrompt,
  31. messages,
  32. tools,
  33. tool_choice: { type: "auto" },
  34. })
  35. return { message }
  36. }
  37. getModel(): { id: BedrockModelId; info: ModelInfo } {
  38. const modelId = this.options.apiModelId
  39. if (modelId && modelId in bedrockModels) {
  40. const id = modelId as BedrockModelId
  41. return { id, info: bedrockModels[id] }
  42. }
  43. return { id: bedrockDefaultModelId, info: bedrockModels[bedrockDefaultModelId] }
  44. }
  45. }