File size: 2,032 Bytes
68d7816 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | import { errorStatusCode } from '#/llm/errors';
import type { LlmModel } from '#/llm/model';
import type { LlmRecovery } from '#/llm/requester/recovery';
import {
mergeRequestHeaders,
type LlmCredential,
type LlmCredentialProvider,
} from '#/llm/requester/requester';
export interface AccessTokenResolver {
(options?: { readonly force?: boolean }): Promise<string | undefined>;
}
export function createStaticCredentialProvider(apiKey?: string): LlmCredentialProvider {
return {
resolve: () =>
apiKey === undefined || apiKey.trim().length === 0 ? undefined : { apiKey },
};
}
export function createOAuthCredentialProvider(
getToken: AccessTokenResolver,
): LlmCredentialProvider {
let refreshed: Promise<string | undefined> | undefined;
return {
resolve: async () => {
const pending = refreshed;
refreshed = undefined;
const apiKey = pending === undefined ? await getToken() : await pending;
return apiKey === undefined ? undefined : { apiKey };
},
canRecover: (error) => errorStatusCode(error) === 401,
invalidate: () => {
refreshed ??= getToken({ force: true });
refreshed.catch(() => {});
},
};
}
export function applyCredential(
model: LlmModel,
credential: LlmCredential | undefined,
): LlmModel {
if (credential === undefined) {
return model;
}
return {
...model,
apiKey: credential.apiKey ?? model.apiKey,
defaultHeaders: mergeRequestHeaders(model.defaultHeaders, credential.headers),
};
}
const CREDENTIALS_RECOVERY_ID = 'credentials';
export const credentialsRecovery: LlmRecovery = {
propose: ({ error, appliedRecoveries, credentialProvider }) => {
if (
credentialProvider?.canRecover?.(error) !== true ||
appliedRecoveries.some((record) => record.strategy === CREDENTIALS_RECOVERY_ID)
) {
return undefined;
}
return {
strategy: CREDENTIALS_RECOVERY_ID,
action: 'refresh',
beforeNextAttempt: () => credentialProvider?.invalidate?.(),
};
},
};
|