Mirrors
/
oai-reverse-proxy
mirror of https://gitgud.io/khanon/oai-reverse-proxy.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

allows KeyPool.available to be scoped per-service

This commit is contained in:
nai-degen 2023-06-01 08:31:41 -05:00
parent b7fa65d9f9
commit a6c6e4c694
2 changed files with 7 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
src/key-management/key-pool.ts
View File

@ -13,7 +13,7 @@ export class KeyPool {
public init() {
this.keyProviders.forEach((provider) => provider.init());
const availableKeys = this.available();
const availableKeys = this.available("all");
if (availableKeys === 0) {
throw new Error(
"No keys loaded. Ensure either OPENAI_KEY or ANTHROPIC_KEY is set."
@ -35,14 +35,11 @@ export class KeyPool {
service.disable(key);
}
// TODO: this probably needs to be scoped to a specific provider. I think the
// only code calling this is the error handler which needs to know how many
// more keys are available for the provider the user tried to use.
public available(): number {
return this.keyProviders.reduce(
(sum, provider) => sum + provider.available(),
0
);
public available(service: AIService | "all" = "all"): number {
return this.keyProviders.reduce((sum, provider) => {
const includeProvider = service === "all" || service === provider.service;
return sum + (includeProvider ? provider.available() : 0);
}, 0);
}
public anyUnchecked(): boolean {

2
src/proxy/middleware/response/index.ts
View File

@ -223,7 +223,7 @@ const handleUpstreamErrors: ProxyResHandlerWithBody = async (
let errorPayload: Record<string, any>;
// Subtract 1 from available keys because if this message is being shown,
// it's because the key is about to be disabled.
const availableKeys = keyPool.available() - 1;
const availableKeys = keyPool.available(req.outboundApi) - 1;
const tryAgainMessage = Boolean(availableKeys)
? `There are ${availableKeys} more keys available; try your request again.`
: "There are no more keys available.";