fixes AWS Sonnet 3.5 key assignment bug

This commit is contained in:
nai-degen 2024-06-20 12:00:08 -05:00
parent 5871025245
commit 84b917f726
2 changed files with 40 additions and 22 deletions

View File

@ -64,9 +64,15 @@ export class AwsKeyChecker extends KeyCheckerBase<AwsBedrockKey> {
// a ResourceNotFoundException when trying to invoke it which will fail the
// entire key As a temporary measure we will trap thrown errors for this
// particular check and ignore them.
checks.unshift(
checks.push(
this.invokeModel("anthropic.claude-3-5-sonnet-20240620-v1:0", key).catch(
() => false
({ response }) => {
this.log.debug(
{ key: key.hash, error: response.data, status: response.status },
"AWS Sonnet 3.5 model is not accessible."
);
return false;
}
)
);
@ -75,9 +81,14 @@ export class AwsKeyChecker extends KeyCheckerBase<AwsBedrockKey> {
const [_logging, claudeV2, sonnet, haiku, opus, sonnet35] =
await Promise.all(checks);
this.log.debug(
{ key: key.hash, _logging, claudeV2, sonnet, haiku, opus, sonnet35 },
"AWS model tests complete."
);
if (isInitialCheck) {
const families: AwsBedrockModelFamily[] = [];
if (claudeV2 || sonnet || haiku) families.push("aws-claude");
if (claudeV2 || sonnet || sonnet35 || haiku) families.push("aws-claude");
if (opus) families.push("aws-claude-opus");
if (families.length === 0) {
@ -214,13 +225,6 @@ export class AwsKeyChecker extends KeyCheckerBase<AwsBedrockKey> {
const correctErrorMessage = errorMessage?.match(/max_tokens/);
if (!correctErrorType || !correctErrorMessage) {
return false;
// throw new AxiosError(
// `Unexpected error when invoking model ${model}: ${errorMessage}`,
// "AWS_ERROR",
// response.config,
// response.request,
// response
// );
}
this.log.debug(

View File

@ -99,8 +99,6 @@ export class AwsBedrockKeyProvider implements KeyProvider<AwsBedrockKey> {
}
public get(model: string) {
const availableKeys = this.keys.filter((k) => {
const isNotLogged = k.awsLoggingStatus !== "enabled";
const neededFamily = getAwsBedrockModelFamily(model);
// this is a horrible mess
@ -109,10 +107,12 @@ export class AwsBedrockKeyProvider implements KeyProvider<AwsBedrockKey> {
// model variants are added.
const needsSonnet =
model.includes("sonnet") && neededFamily === "aws-claude";
const needsHaiku =
model.includes("haiku") && neededFamily === "aws-claude";
const needsHaiku = model.includes("haiku") && neededFamily === "aws-claude";
const needsSonnet35 =
model.includes("claude-3-5-sonnet") && neededFamily === "aws-claude";
const availableKeys = this.keys.filter((k) => {
const isNotLogged = k.awsLoggingStatus !== "enabled";
return (
!k.isDisabled &&
(isNotLogged || config.allowAwsLogging) &&
@ -122,6 +122,20 @@ export class AwsBedrockKeyProvider implements KeyProvider<AwsBedrockKey> {
k.modelFamilies.includes(neededFamily)
);
});
this.log.debug(
{
model,
neededFamily,
needsSonnet,
needsHaiku,
needsSonnet35,
availableKeys: availableKeys.length,
totalKeys: this.keys.length,
},
"Selecting AWS key"
);
if (availableKeys.length === 0) {
throw new PaymentRequiredError(
`No AWS Bedrock keys available for model ${model}`