adds aws opus maybe, idk cannot test

src/shared/key-management/aws/provider.ts

@@ -2,7 +2,7 @@ import crypto from "crypto";
 import { Key, KeyProvider } from "..";
 import { config } from "../../../config";
 import { logger } from "../../../logger";
-import type { AwsBedrockModelFamily } from "../../models";
+import { AwsBedrockModelFamily, getAwsBedrockModelFamily } from "../../models";
 import { AwsKeyChecker } from "./checker";
 import { PaymentRequiredError } from "../../errors";
 
@@ -61,7 +61,7 @@ export class AwsBedrockKeyProvider implements KeyProvider<AwsBedrockKey> {
       const newKey: AwsBedrockKey = {
         key,
         service: this.service,
-        modelFamilies: ["aws-claude"],
+        modelFamilies: ["aws-claude", "aws-claude-opus"],
         isDisabled: false,
         isRevoked: false,
         promptCount: 0,
@@ -78,6 +78,7 @@ export class AwsBedrockKeyProvider implements KeyProvider<AwsBedrockKey> {
         sonnetEnabled: true,
         haikuEnabled: false,
         ["aws-claudeTokens"]: 0,
+        ["aws-claude-opusTokens"]: 0,
       };
       this.keys.push(newKey);
     }
@@ -157,11 +158,11 @@ export class AwsBedrockKeyProvider implements KeyProvider<AwsBedrockKey> {
     return this.keys.filter((k) => !k.isDisabled).length;
   }
 
-  public incrementUsage(hash: string, _model: string, tokens: number) {
+  public incrementUsage(hash: string, model: string, tokens: number) {
     const key = this.keys.find((k) => k.hash === hash);
     if (!key) return;
     key.promptCount++;
-    key["aws-claudeTokens"] += tokens;
+    key[`${getAwsBedrockModelFamily(model)}Tokens`] += tokens;
   }
 
   public getLockoutPeriod() {

src/shared/models.ts

@@ -29,7 +29,7 @@ export type MistralAIModelFamily =
   | "mistral-small"
   | "mistral-medium"
   | "mistral-large";
-export type AwsBedrockModelFamily = "aws-claude";
+export type AwsBedrockModelFamily = "aws-claude" | "aws-claude-opus";
 export type AzureOpenAIModelFamily = `azure-${OpenAIModelFamily}`;
 export type ModelFamily =
   | OpenAIModelFamily
@@ -55,6 +55,7 @@ export const MODEL_FAMILIES = (<A extends readonly ModelFamily[]>(
   "mistral-medium",
   "mistral-large",
   "aws-claude",
+  "aws-claude-opus",
   "azure-turbo",
   "azure-gpt4",
   "azure-gpt4-32k",
@@ -98,6 +99,7 @@ export const MODEL_FAMILY_SERVICE: {
   claude: "anthropic",
   "claude-opus": "anthropic",
   "aws-claude": "aws",
+  "aws-claude-opus": "aws",
   "azure-turbo": "azure",
   "azure-gpt4": "azure",
   "azure-gpt4-32k": "azure",
@@ -150,8 +152,8 @@ export function getMistralAIModelFamily(model: string): MistralAIModelFamily {
   }
 }
 
-export function getAwsBedrockModelFamily(model: string): ModelFamily {
-  if (model.includes("opus")) return "claude-opus";
+export function getAwsBedrockModelFamily(model: string): AwsBedrockModelFamily {
+  if (model.includes("opus")) return "aws-claude-opus";
   return "aws-claude";
 }
 

src/shared/stats.ts

@@ -29,6 +29,7 @@ export function getTokenCostUsd(model: ModelFamily, tokens: number) {
     case "claude":
       cost = 0.000008;
       break;
+    case "aws-claude-opus":
     case "claude-opus":
       cost = 0.000015;
       break;