Add native LLM core foundation (#24712)

2026-05-08 16:56:20 -04:00
parent dc7d665e94
commit 5bb7b23440
144 changed files with 17052 additions and 2 deletions
@@ -0,0 +1,103 @@
+import { AwsV4Signer } from "aws4fetch"
+import { Effect, Option, Schema } from "effect"
+import { Headers } from "effect/unstable/http"
+import { Auth, type AuthInput } from "../../route/auth"
+import type { LLMRequest } from "../../schema"
+import { ProviderShared } from "../shared"
+
+/**
+ * AWS credentials for SigV4 signing. Bedrock also supports Bearer API key auth
+ * via `model.apiKey`, which bypasses SigV4 signing. STS-vended credentials
+ * should be refreshed by the consumer (rebuild the model) before they expire;
+ * the route does not refresh.
+ */
+export interface Credentials {
+  readonly region: string
+  readonly accessKeyId: string
+  readonly secretAccessKey: string
+  readonly sessionToken?: string
+}
+
+const NativeCredentials = Schema.Struct({
+  accessKeyId: Schema.String,
+  secretAccessKey: Schema.String,
+  region: Schema.optional(Schema.String),
+  sessionToken: Schema.optional(Schema.String),
+})
+
+const decodeNativeCredentials = Schema.decodeUnknownOption(NativeCredentials)
+
+export const region = (request: LLMRequest) => {
+  const fromNative = request.model.native?.aws_region
+  if (typeof fromNative === "string" && fromNative !== "") return fromNative
+  return (
+    decodeNativeCredentials(request.model.native?.aws_credentials).pipe(
+      Option.map((credentials) => credentials.region),
+      Option.getOrUndefined,
+    ) ?? "us-east-1"
+  )
+}
+
+const credentialsFromInput = (request: LLMRequest): Credentials | undefined =>
+  decodeNativeCredentials(request.model.native?.aws_credentials).pipe(
+    Option.map((creds) => ({ ...creds, region: creds.region ?? region(request) })),
+    Option.getOrUndefined,
+  )
+
+const signRequest = (input: {
+  readonly url: string
+  readonly body: string
+  readonly headers: Headers.Headers
+  readonly credentials: Credentials
+}) =>
+  Effect.tryPromise({
+    try: async () => {
+      const signed = await new AwsV4Signer({
+        url: input.url,
+        method: "POST",
+        headers: Object.entries(input.headers),
+        body: input.body,
+        region: input.credentials.region,
+        accessKeyId: input.credentials.accessKeyId,
+        secretAccessKey: input.credentials.secretAccessKey,
+        sessionToken: input.credentials.sessionToken,
+        service: "bedrock",
+      }).sign()
+      return Object.fromEntries(signed.headers.entries())
+    },
+    catch: (error) =>
+      ProviderShared.invalidRequest(
+        `Bedrock Converse SigV4 signing failed: ${error instanceof Error ? error.message : String(error)}`,
+      ),
+  })
+
+/**
+ * Bedrock auth. `model.apiKey` (Bedrock's newer Bearer API key auth) wins if
+ * set; otherwise sign the exact JSON bytes with SigV4 using credentials from
+ * `model.native.aws_credentials`.
+ */
+export const auth = Auth.custom((input: AuthInput) => {
+  if (input.request.model.apiKey) return Auth.toEffect(Auth.bearer())(input)
+  return Effect.gen(function* () {
+    const credentials = credentialsFromInput(input.request)
+    if (!credentials) {
+      return yield* ProviderShared.invalidRequest(
+        "Bedrock Converse requires either model.apiKey or AWS credentials in model.native.aws_credentials",
+      )
+    }
+    const headersForSigning = Headers.set(input.headers, "content-type", "application/json")
+    const signed = yield* signRequest({ url: input.url, body: input.body, headers: headersForSigning, credentials })
+    return Headers.setAll(headersForSigning, signed)
+  })
+})
+
+export const nativeCredentials = (native: Record<string, unknown> | undefined, credentials: Credentials | undefined) =>
+  credentials
+    ? {
+        ...native,
+        aws_credentials: credentials,
+        aws_region: credentials.region,
+      }
+    : native
+
+export * as BedrockAuth from "./bedrock-auth"