diff --git a/src/proxy/aws.ts b/src/proxy/aws.ts index 094eec5..d63bd46 100644 --- a/src/proxy/aws.ts +++ b/src/proxy/aws.ts @@ -37,6 +37,7 @@ const getModelsResponse = () => { "anthropic.claude-v2:1", "anthropic.claude-3-haiku-20240307-v1:0", "anthropic.claude-3-sonnet-20240229-v1:0", + "anthropic.claude-3-opus-20240229-v1:0" ]; const models = variants.map((id) => ({ diff --git a/src/service-info.ts b/src/service-info.ts index 7adf0fd..57a60a2 100644 --- a/src/service-info.ts +++ b/src/service-info.ts @@ -141,8 +141,6 @@ const SERVICE_ENDPOINTS: { [s in LLMService]: Record } = { }, anthropic: { anthropic: `%BASE%/anthropic`, - "anthropic-sonnet (⚠️Temporary: for Claude 3 Sonnet)": `%BASE%/anthropic/sonnet`, - "anthropic-opus (⚠️Temporary: for Claude 3 Opus)": `%BASE%/anthropic/opus`, }, "google-ai": { "google-ai": `%BASE%/google-ai`, @@ -152,7 +150,6 @@ const SERVICE_ENDPOINTS: { [s in LLMService]: Record } = { }, aws: { aws: `%BASE%/aws/claude`, - "aws-sonnet (⚠️Temporary: for AWS Claude 3 Sonnet)": `%BASE%/aws/claude/sonnet`, }, azure: { azure: `%BASE%/azure/openai`, @@ -212,7 +209,8 @@ export function buildInfo(baseUrl: string, forAdmin = false): ServiceInfo { } function getStatus() { - if (!config.checkKeys) return "Key checking is disabled. The data displayed are not reliable."; + if (!config.checkKeys) + return "Key checking is disabled. The data displayed are not reliable."; let unchecked = 0; for (const service of LLM_SERVICES) { @@ -444,13 +442,15 @@ function getInfoForFamily(family: ModelFamily): BaseFamilyInfo { info.prefilledKeys = modelStats.get(`${family}__pozzed`) || 0; break; case "aws": - info.sonnetKeys = modelStats.get(`${family}__awsSonnet`) || 0; - info.haikuKeys = modelStats.get(`${family}__awsHaiku`) || 0; - const logged = modelStats.get(`${family}__awsLogged`) || 0; - if (logged > 0) { - info.privacy = config.allowAwsLogging - ? `${logged} active keys are potentially logged.` - : `${logged} active keys are potentially logged and can't be used. Set ALLOW_AWS_LOGGING=true to override.`; + if (family === "aws-claude") { + info.sonnetKeys = modelStats.get(`${family}__awsSonnet`) || 0; + info.haikuKeys = modelStats.get(`${family}__awsHaiku`) || 0; + const logged = modelStats.get(`${family}__awsLogged`) || 0; + if (logged > 0) { + info.privacy = config.allowAwsLogging + ? `${logged} active keys are potentially logged.` + : `${logged} active keys are potentially logged and can't be used. Set ALLOW_AWS_LOGGING=true to override.`; + } } break; } diff --git a/src/shared/key-management/aws/checker.ts b/src/shared/key-management/aws/checker.ts index 9c11aa3..164886d 100644 --- a/src/shared/key-management/aws/checker.ts +++ b/src/shared/key-management/aws/checker.ts @@ -5,6 +5,7 @@ import axios, { AxiosError, AxiosRequestConfig, AxiosHeaders } from "axios"; import { URL } from "url"; import { KeyCheckerBase } from "../key-checker-base"; import type { AwsBedrockKey, AwsBedrockKeyProvider } from "./provider"; +import { AwsBedrockModelFamily } from "../../models"; const MIN_CHECK_INTERVAL = 3 * 1000; // 3 seconds const KEY_CHECK_PERIOD = 30 * 60 * 1000; // 30 minutes @@ -54,18 +55,32 @@ export class AwsKeyChecker extends KeyCheckerBase { this.invokeModel("anthropic.claude-v2", key), this.invokeModel("anthropic.claude-3-sonnet-20240229-v1:0", key), this.invokeModel("anthropic.claude-3-haiku-20240307-v1:0", key), + this.invokeModel("anthropic.claude-3-opus-20240229-v1:0", key), ]; } checks.unshift(this.checkLoggingConfiguration(key)); - const [_logging, _claudeV2, sonnet, haiku] = await Promise.all(checks); + const [_logging, claudeV2, sonnet, haiku, opus] = await Promise.all(checks); if (isInitialCheck) { - this.updateKey(key.hash, { sonnetEnabled: sonnet, haikuEnabled: haiku }); + const families: AwsBedrockModelFamily[] = []; + if (claudeV2 || sonnet || haiku) families.push("aws-claude"); + if (opus) families.push("aws-claude-opus"); + this.updateKey(key.hash, { + sonnetEnabled: sonnet, + haikuEnabled: haiku, + modelFamilies: families, + }); } this.log.info( - { key: key.hash, sonnet, haiku, logged: key.awsLoggingStatus }, + { + key: key.hash, + sonnet, + haiku, + families: key.modelFamilies, + logged: key.awsLoggingStatus, + }, "Checked key." ); } diff --git a/src/shared/key-management/aws/provider.ts b/src/shared/key-management/aws/provider.ts index 98bbd9c..57c07d1 100644 --- a/src/shared/key-management/aws/provider.ts +++ b/src/shared/key-management/aws/provider.ts @@ -61,7 +61,7 @@ export class AwsBedrockKeyProvider implements KeyProvider { const newKey: AwsBedrockKey = { key, service: this.service, - modelFamilies: ["aws-claude", "aws-claude-opus"], + modelFamilies: ["aws-claude"], isDisabled: false, isRevoked: false, promptCount: 0, @@ -99,13 +99,17 @@ export class AwsBedrockKeyProvider implements KeyProvider { public get(model: string) { const availableKeys = this.keys.filter((k) => { const isNotLogged = k.awsLoggingStatus === "disabled"; - const needsSonnet = model.includes("sonnet"); - const needsHaiku = model.includes("haiku"); + const neededFamily = getAwsBedrockModelFamily(model); + const needsSonnet = + model.includes("sonnet") && neededFamily === "aws-claude"; + const needsHaiku = + model.includes("haiku") && neededFamily === "aws-claude"; return ( !k.isDisabled && (isNotLogged || config.allowAwsLogging) && - (k.sonnetEnabled || !needsSonnet) && - (k.haikuEnabled || !needsHaiku) + (k.sonnetEnabled || !needsSonnet) && // sonnet and haiku are both under aws-claude, while opus is not + (k.haikuEnabled || !needsHaiku) && + k.modelFamilies.includes(neededFamily) ); }); if (availableKeys.length === 0) {