bedrock-auth.ts 3.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103
  1. import { AwsV4Signer } from "aws4fetch"
  2. import { Effect, Option, Schema } from "effect"
  3. import { Headers } from "effect/unstable/http"
  4. import { Auth, type AuthInput } from "../../route/auth"
  5. import type { LLMRequest } from "../../schema"
  6. import { ProviderShared } from "../shared"
  7. /**
  8. * AWS credentials for SigV4 signing. Bedrock also supports Bearer API key auth
  9. * via `model.apiKey`, which bypasses SigV4 signing. STS-vended credentials
  10. * should be refreshed by the consumer (rebuild the model) before they expire;
  11. * the route does not refresh.
  12. */
  13. export interface Credentials {
  14. readonly region: string
  15. readonly accessKeyId: string
  16. readonly secretAccessKey: string
  17. readonly sessionToken?: string
  18. }
  19. const NativeCredentials = Schema.Struct({
  20. accessKeyId: Schema.String,
  21. secretAccessKey: Schema.String,
  22. region: Schema.optional(Schema.String),
  23. sessionToken: Schema.optional(Schema.String),
  24. })
  25. const decodeNativeCredentials = Schema.decodeUnknownOption(NativeCredentials)
  26. export const region = (request: LLMRequest) => {
  27. const fromNative = request.model.native?.aws_region
  28. if (typeof fromNative === "string" && fromNative !== "") return fromNative
  29. return (
  30. decodeNativeCredentials(request.model.native?.aws_credentials).pipe(
  31. Option.map((credentials) => credentials.region),
  32. Option.getOrUndefined,
  33. ) ?? "us-east-1"
  34. )
  35. }
  36. const credentialsFromInput = (request: LLMRequest): Credentials | undefined =>
  37. decodeNativeCredentials(request.model.native?.aws_credentials).pipe(
  38. Option.map((creds) => ({ ...creds, region: creds.region ?? region(request) })),
  39. Option.getOrUndefined,
  40. )
  41. const signRequest = (input: {
  42. readonly url: string
  43. readonly body: string
  44. readonly headers: Headers.Headers
  45. readonly credentials: Credentials
  46. }) =>
  47. Effect.tryPromise({
  48. try: async () => {
  49. const signed = await new AwsV4Signer({
  50. url: input.url,
  51. method: "POST",
  52. headers: Object.entries(input.headers),
  53. body: input.body,
  54. region: input.credentials.region,
  55. accessKeyId: input.credentials.accessKeyId,
  56. secretAccessKey: input.credentials.secretAccessKey,
  57. sessionToken: input.credentials.sessionToken,
  58. service: "bedrock",
  59. }).sign()
  60. return Object.fromEntries(signed.headers.entries())
  61. },
  62. catch: (error) =>
  63. ProviderShared.invalidRequest(
  64. `Bedrock Converse SigV4 signing failed: ${error instanceof Error ? error.message : String(error)}`,
  65. ),
  66. })
  67. /**
  68. * Bedrock auth. `model.apiKey` (Bedrock's newer Bearer API key auth) wins if
  69. * set; otherwise sign the exact JSON bytes with SigV4 using credentials from
  70. * `model.native.aws_credentials`.
  71. */
  72. export const auth = Auth.custom((input: AuthInput) => {
  73. if (input.request.model.apiKey) return Auth.toEffect(Auth.bearer())(input)
  74. return Effect.gen(function* () {
  75. const credentials = credentialsFromInput(input.request)
  76. if (!credentials) {
  77. return yield* ProviderShared.invalidRequest(
  78. "Bedrock Converse requires either model.apiKey or AWS credentials in model.native.aws_credentials",
  79. )
  80. }
  81. const headersForSigning = Headers.set(input.headers, "content-type", "application/json")
  82. const signed = yield* signRequest({ url: input.url, body: input.body, headers: headersForSigning, credentials })
  83. return Headers.setAll(headersForSigning, signed)
  84. })
  85. })
  86. export const nativeCredentials = (native: Record<string, unknown> | undefined, credentials: Credentials | undefined) =>
  87. credentials
  88. ? {
  89. ...native,
  90. aws_credentials: credentials,
  91. aws_region: credentials.region,
  92. }
  93. : native
  94. export * as BedrockAuth from "./bedrock-auth"