diff --git a/src/shared/key-management/anthropic/provider.ts b/src/shared/key-management/anthropic/provider.ts index a56a5f2..144bd30 100644 --- a/src/shared/key-management/anthropic/provider.ts +++ b/src/shared/key-management/anthropic/provider.ts @@ -15,6 +15,16 @@ export const ANTHROPIC_SUPPORTED_MODELS = [ ] as const; export type AnthropicModel = (typeof ANTHROPIC_SUPPORTED_MODELS)[number]; +type AnthropicKeyUsage = { + [K in AnthropicModelFamily as `${K}Tokens`]: number; +}; + +const SERIALIZABLE_FIELDS = ["key", "service", "hash", "claudeTokens"] as const; +type SerializableAnthropicKey = Partial< + Pick +> & + Pick; + export type AnthropicKeyUpdate = Omit< Partial, | "key" @@ -25,10 +35,6 @@ export type AnthropicKeyUpdate = Omit< | "rateLimitedUntil" >; -type AnthropicKeyUsage = { - [K in AnthropicModelFamily as `${K}Tokens`]: number; -}; - export interface AnthropicKey extends Key, AnthropicKeyUsage { readonly service: "anthropic"; readonly modelFamilies: AnthropicModelFamily[]; @@ -81,27 +87,7 @@ export class AnthropicKeyProvider implements KeyProvider { let bareKeys: string[]; bareKeys = [...new Set(keyConfig.split(",").map((k) => k.trim()))]; for (const key of bareKeys) { - const newKey: AnthropicKey = { - key, - service: this.service, - modelFamilies: ["claude"], - isDisabled: false, - isRevoked: false, - isPozzed: false, - promptCount: 0, - lastUsed: 0, - rateLimitedAt: 0, - rateLimitedUntil: 0, - requiresPreamble: false, - hash: `ant-${crypto - .createHash("sha256") - .update(key) - .digest("hex") - .slice(0, 8)}`, - lastChecked: 0, - claudeTokens: 0, - }; - this.keys.push(newKey); + this.keys.push(AnthropicKeyProvider.deserialize({ key })); } this.log.info({ keyCount: this.keys.length }, "Loaded Anthropic keys."); } @@ -226,4 +212,28 @@ export class AnthropicKeyProvider implements KeyProvider { }); this.checker?.scheduleNextCheck(); } + + static deserialize({ key, ...rest }: SerializableAnthropicKey): AnthropicKey { + return { + key, + service: "anthropic" as const, + modelFamilies: ["claude" as const], + isTrial: false, + isDisabled: false, + isPozzed: false, + promptCount: 0, + lastUsed: 0, + rateLimitedAt: 0, + rateLimitedUntil: 0, + requiresPreamble: false, + hash: `ant-${crypto + .createHash("sha256") + .update(key) + .digest("hex") + .slice(0, 8)}`, + lastChecked: 0, + claudeTokens: 0, + ...rest, + }; + } } diff --git a/src/shared/key-management/openai/provider.ts b/src/shared/key-management/openai/provider.ts index e55bbbf..72cb1cb 100644 --- a/src/shared/key-management/openai/provider.ts +++ b/src/shared/key-management/openai/provider.ts @@ -109,8 +109,7 @@ export class OpenAIKeyProvider implements KeyProvider { bareKeys = keyString.split(",").map((k) => k.trim()); bareKeys = [...new Set(bareKeys)]; for (const k of bareKeys) { - const newKey = OpenAIKeyProvider.deserialize({ key: k }); - this.keys.push(newKey); + this.keys.push(OpenAIKeyProvider.deserialize({ key: k })); } this.log.info({ keyCount: this.keys.length }, "Loaded OpenAI keys."); }