diff --git a/src/shared/key-management/openai/provider.ts b/src/shared/key-management/openai/provider.ts index f5a1441..2e6dd53 100644 --- a/src/shared/key-management/openai/provider.ts +++ b/src/shared/key-management/openai/provider.ts @@ -149,32 +149,51 @@ export class OpenAIKeyProvider implements KeyProvider { return this.keys.map((key) => Object.freeze({ ...key, key: undefined })); } - public get(requestModel: string) { - let model = requestModel; - - const neededFamily = getOpenAIModelFamily(model); - const excludeTrials = model === "text-embedding-ada-002"; - - const availableKeys = this.keys.filter( - // Allow keys which - (key) => - !key.isDisabled && // are not disabled - key.modelFamilies.includes(neededFamily) && // have access to the model family we need - (!excludeTrials || !key.isTrial) && // and are not trials if we don't want them - (!config.checkKeys || key.modelIds.includes(model)) // and have the specific snapshot we need - ); - - if (availableKeys.length === 0) { - throw new PaymentRequiredError( - `No OpenAI keys available for model ${model}` - ); + public get(requestModel: string): OpenAIKey { + if (!this.keys.length) { + throw new PaymentRequiredError(`No OpenAI keys available for model ${requestModel}`); } - const keysByPriority = prioritizeKeys( - availableKeys, - (a, b) => +a.isTrial - +b.isTrial + // Get the model family, which determines key compatibility. + const family = getOpenAIModelFamily(requestModel, "turbo"); + + // For image models, we need extra validation to ensure the key has been properly validated + const isImageModel = family === "gpt-image"; + + // Only use keys that are enabled, not revoked, and can serve this model family. + const suitableKeys = this.keys.filter( + (k) => { + const basicRequirements = + !k.isDisabled && + !k.isRevoked && + !k.isOverQuota && + k.modelFamilies.includes(family) && + Date.now() >= k.rateLimitedUntil; + + // For image models, we need to ensure the key has been properly validated + // and has actually been tested for gpt-image-1 access + if (isImageModel && basicRequirements) { + // If the key hasn't been checked at all, it shouldn't be used for image models yet + if (!k.lastChecked) { + this.log.debug({ key: k.hash }, "Skipping unchecked key for image model request"); + return false; + } + return true; + } + + return basicRequirements; + } ); + if (suitableKeys.length === 0) { + throw new PaymentRequiredError(`No suitable OpenAI keys available for model ${requestModel}`); + } + + // Sort keys by usefulness and select the most useful key. + // Trial keys are prioritized because they can't be used for anything else. + // Otherwise, the key with the most remaining quota is chosen. + const keysByPriority = prioritizeKeys(suitableKeys, (a, b) => +a.isTrial - +b.isTrial); + const selectedKey = keysByPriority[0]; selectedKey.lastUsed = Date.now(); this.throttle(selectedKey.hash);