lobe-chat

1
import { JWTPayload, LOBE_CHAT_AUTH_HEADER } from '@/const/auth';
2
import { ModelProvider } from '@/libs/agent-runtime';
3
import { useUserStore } from '@/store/user';
4
import { keyVaultsConfigSelectors, userProfileSelectors } from '@/store/user/selectors';
5
import { GlobalLLMProviderKey } from '@/types/user/settings';
6
import { createJWT } from '@/utils/jwt';
7

8
export const getProviderAuthPayload = (provider: string) => {
9
  switch (provider) {
10
    case ModelProvider.Bedrock: {
11
      const { accessKeyId, region, secretAccessKey, sessionToken } =
12
        keyVaultsConfigSelectors.bedrockConfig(useUserStore.getState());
13

14
      const awsSecretAccessKey = secretAccessKey;
15
      const awsAccessKeyId = accessKeyId;
16

17
      const apiKey = (awsSecretAccessKey || '') + (awsAccessKeyId || '');
18

19
      return {
20
        apiKey,
21
        awsAccessKeyId,
22
        awsRegion: region,
23
        awsSecretAccessKey,
24
        awsSessionToken: sessionToken,
25
      };
26
    }
27

28
    case ModelProvider.Wenxin: {
29
      const { secretKey, accessKey } = keyVaultsConfigSelectors.wenxinConfig(
30
        useUserStore.getState(),
31
      );
32

33
      const apiKey = (accessKey || '') + (secretKey || '');
34

35
      return {
36
        apiKey,
37
        wenxinAccessKey: accessKey,
38
        wenxinSecretKey: secretKey,
39
      };
40
    }
41

42
    case ModelProvider.Azure: {
43
      const azure = keyVaultsConfigSelectors.azureConfig(useUserStore.getState());
44

45
      return {
46
        apiKey: azure.apiKey,
47
        azureApiVersion: azure.apiVersion,
48
        endpoint: azure.endpoint,
49
      };
50
    }
51

52
    case ModelProvider.Ollama: {
53
      const config = keyVaultsConfigSelectors.ollamaConfig(useUserStore.getState());
54

55
      return { endpoint: config?.baseURL };
56
    }
57

58
    default: {
59
      const config = keyVaultsConfigSelectors.getVaultByProvider(provider as GlobalLLMProviderKey)(
60
        useUserStore.getState(),
61
      );
62

63
      return { apiKey: config?.apiKey, endpoint: config?.baseURL };
64
    }
65
  }
66
};
67

68
const createAuthTokenWithPayload = async (payload = {}) => {
69
  const accessCode = keyVaultsConfigSelectors.password(useUserStore.getState());
70
  const userId = userProfileSelectors.userId(useUserStore.getState());
71

72
  return await createJWT<JWTPayload>({ accessCode, userId, ...payload });
73
};
74

75
interface AuthParams {
76
  // eslint-disable-next-line no-undef
77
  headers?: HeadersInit;
78
  payload?: Record<string, any>;
79
  provider?: string;
80
}
81

82
// eslint-disable-next-line no-undef
83
export const createHeaderWithAuth = async (params?: AuthParams): Promise<HeadersInit> => {
84
  let payload = params?.payload || {};
85

86
  if (params?.provider) {
87
    payload = { ...payload, ...getProviderAuthPayload(params?.provider) };
88
  }
89

90
  const token = await createAuthTokenWithPayload(payload);
91

92
  // eslint-disable-next-line no-undef
93
  return { ...params?.headers, [LOBE_CHAT_AUTH_HEADER]: token };
94
};
95