alexion
4def357d09
Imported as a linear vendor commit so forge rebases do not replay the upstream root tree. git-subtree-dir: third_party/pi git-subtree-split: d981de1229ef899957bbe968bc8dcda02a21f477
91 lines
3.3 KiB
TypeScript
91 lines
3.3 KiB
TypeScript
import { bedrockConverseStreamApi } from "../api/bedrock-converse-stream.lazy.ts";
|
|
import type { ApiKeyAuth } from "../auth/types.ts";
|
|
import { createProvider, type Provider } from "../models.ts";
|
|
import { AMAZON_BEDROCK_MODELS } from "./amazon-bedrock.models.ts";
|
|
|
|
/**
|
|
* Bedrock accepts a bearer token or the AWS SDK's default credential chain.
|
|
* The login flow can store a token/profile choice; resolve also detects ambient
|
|
* AWS credentials without copying them into pi's credential store.
|
|
*/
|
|
const bedrockAuth: ApiKeyAuth = {
|
|
name: "AWS credentials or bearer token",
|
|
login: async (interaction) => {
|
|
interaction.signal.throwIfAborted();
|
|
const method = await interaction.prompt({
|
|
type: "select",
|
|
message: "Select Amazon Bedrock authentication method:",
|
|
options: [
|
|
{ id: "bearer-token", label: "Bearer token" },
|
|
{ id: "aws-profile", label: "AWS profile" },
|
|
{ id: "credential-chain", label: "Existing AWS credential chain" },
|
|
],
|
|
});
|
|
interaction.signal.throwIfAborted();
|
|
if (method === "bearer-token") {
|
|
return {
|
|
type: "api_key",
|
|
key: await interaction.prompt({ type: "secret", message: "Enter Amazon Bedrock bearer token" }),
|
|
};
|
|
}
|
|
interaction.notify({
|
|
type: "info",
|
|
message: "Amazon Bedrock supports AWS profiles, IAM credentials, and role-based credentials.",
|
|
links: [
|
|
{
|
|
label: "AWS credential provider chain",
|
|
url: "https://docs.aws.amazon.com/sdkref/latest/guide/standardized-credentials.html",
|
|
},
|
|
],
|
|
});
|
|
if (method === "aws-profile") {
|
|
return {
|
|
type: "api_key",
|
|
env: { AWS_PROFILE: await interaction.prompt({ type: "text", message: "Enter AWS profile name" }) },
|
|
};
|
|
}
|
|
if (method !== "credential-chain") throw new Error(`Unknown Amazon Bedrock auth method: ${method}`);
|
|
await interaction.prompt({
|
|
type: "text",
|
|
message: "Configure AWS credentials, then press Enter to continue",
|
|
});
|
|
return { type: "api_key" };
|
|
},
|
|
resolve: async ({ ctx, credential, signal }) => {
|
|
const env = async (name: string) => {
|
|
signal.throwIfAborted();
|
|
const value = await ctx.env(name);
|
|
signal.throwIfAborted();
|
|
return value;
|
|
};
|
|
if (credential?.key) {
|
|
return { auth: { apiKey: credential.key }, env: credential.env, source: "stored credential" };
|
|
}
|
|
if (await env("AWS_BEARER_TOKEN_BEDROCK")) return { auth: {}, source: "AWS_BEARER_TOKEN_BEDROCK" };
|
|
if (credential?.env?.AWS_PROFILE ?? (await env("AWS_PROFILE"))) {
|
|
return {
|
|
auth: {},
|
|
env: credential?.env,
|
|
source: credential?.env?.AWS_PROFILE ? "stored credential" : "AWS_PROFILE",
|
|
};
|
|
}
|
|
if ((await env("AWS_ACCESS_KEY_ID")) && (await env("AWS_SECRET_ACCESS_KEY"))) {
|
|
return { auth: {}, source: "AWS access keys" };
|
|
}
|
|
if (await env("AWS_CONTAINER_CREDENTIALS_RELATIVE_URI")) return { auth: {}, source: "ECS task role" };
|
|
if (await env("AWS_CONTAINER_CREDENTIALS_FULL_URI")) return { auth: {}, source: "ECS task role" };
|
|
if (await env("AWS_WEB_IDENTITY_TOKEN_FILE")) return { auth: {}, source: "web identity token" };
|
|
return undefined;
|
|
},
|
|
};
|
|
|
|
export function amazonBedrockProvider(): Provider<"bedrock-converse-stream"> {
|
|
return createProvider({
|
|
id: "amazon-bedrock",
|
|
name: "Amazon Bedrock",
|
|
auth: { apiKey: bedrockAuth },
|
|
models: Object.values(AMAZON_BEDROCK_MODELS),
|
|
api: bedrockConverseStreamApi(),
|
|
});
|
|
}
|