lobe-chat

1
// @vitest-environment node
2
import { beforeEach, describe, expect, it, vi } from 'vitest';
3

4
/**
5
 * This file contains the root router of your tRPC-backend
6
 */
7
import { createCallerFactory } from '@/libs/trpc';
8
import { AuthContext, createContextInner } from '@/server/context';
9
import { GlobalServerConfig } from '@/types/serverConfig';
10

11
import { configRouter } from './index';
12

13
const createCaller = createCallerFactory(configRouter);
14
let ctx: AuthContext;
15
let router: ReturnType<typeof createCaller>;
16

17
vi.mock('@/libs/next-auth/edge', () => {
18
  return {
19
    auth: vi.fn().mockResolvedValue(undefined),
20
  };
21
});
22

23
beforeEach(async () => {
24
  vi.resetAllMocks();
25
  ctx = await createContextInner();
26
  router = createCaller(ctx);
27
});
28

29
describe('configRouter', () => {
30
  describe('getGlobalConfig', () => {
31
    describe('Model Provider env', () => {
32
      describe('OPENAI_MODEL_LIST', () => {
33
        it('custom deletion, addition, and renaming of models', async () => {
34
          process.env.OPENAI_MODEL_LIST =
35
            '-all,+llama,+claude-2，-gpt-3.5-turbo,gpt-4-0125-preview=gpt-4-turbo,gpt-4-0125-preview=gpt-4-32k';
36

37
          const response = await router.getGlobalConfig();
38

39
          // Assert
40
          const result = response.languageModel?.openai;
41

42
          expect(result).toMatchSnapshot();
43
          process.env.OPENAI_MODEL_LIST = '';
44
        });
45

46
        it('should work correct with gpt-4', async () => {
47
          process.env.OPENAI_MODEL_LIST =
48
            '-all,+gpt-3.5-turbo-1106,+gpt-3.5-turbo,+gpt-4,+gpt-4-32k,+gpt-4-1106-preview,+gpt-4-vision-preview';
49

50
          const response = await router.getGlobalConfig();
51

52
          const result = response.languageModel?.openai?.serverModelCards;
53

54
          expect(result).toMatchSnapshot();
55

56
          process.env.OPENAI_MODEL_LIST = '';
57
        });
58

59
        it('duplicate naming model', async () => {
60
          process.env.OPENAI_MODEL_LIST =
61
            'gpt-4-0125-preview=gpt-4-turbo，gpt-4-0125-preview=gpt-4-32k';
62

63
          const response = await router.getGlobalConfig();
64

65
          const result = response.languageModel?.openai?.serverModelCards;
66

67
          expect(result?.find((s) => s.id === 'gpt-4-0125-preview')?.displayName).toEqual(
68
            'gpt-4-32k',
69
          );
70

71
          process.env.OPENAI_MODEL_LIST = '';
72
        });
73

74
        it('should delete model', async () => {
75
          process.env.OPENAI_MODEL_LIST = '-gpt-4';
76

77
          const response = await router.getGlobalConfig();
78

79
          const result = response.languageModel?.openai?.serverModelCards;
80

81
          expect(result?.find((r) => r.id === 'gpt-4')).toBeUndefined();
82

83
          process.env.OPENAI_MODEL_LIST = '';
84
        });
85

86
        it('show the hidden model', async () => {
87
          process.env.OPENAI_MODEL_LIST = '+gpt-4-1106-preview';
88

89
          const response = await router.getGlobalConfig();
90

91
          const result = response.languageModel?.openai?.serverModelCards;
92

93
          const model = result?.find((o) => o.id === 'gpt-4-1106-preview');
94

95
          expect(model).toMatchSnapshot();
96

97
          process.env.OPENAI_MODEL_LIST = '';
98
        });
99

100
        it('only add the model', async () => {
101
          process.env.OPENAI_MODEL_LIST = 'model1,model2,model3，model4';
102

103
          const response = await router.getGlobalConfig();
104

105
          const result = response.languageModel?.openai?.serverModelCards;
106

107
          expect(result).toContainEqual({
108
            displayName: 'model1',
109
            id: 'model1',
110
            enabled: true,
111
          });
112
          expect(result).toContainEqual({
113
            displayName: 'model2',
114
            enabled: true,
115
            id: 'model2',
116
          });
117
          expect(result).toContainEqual({
118
            displayName: 'model3',
119
            enabled: true,
120
            id: 'model3',
121
          });
122
          expect(result).toContainEqual({
123
            displayName: 'model4',
124
            enabled: true,
125
            id: 'model4',
126
          });
127

128
          process.env.OPENAI_MODEL_LIST = '';
129
        });
130
      });
131

132
      describe('OPENROUTER_MODEL_LIST', () => {
133
        it('custom deletion, addition, and renaming of models', async () => {
134
          process.env.OPENROUTER_MODEL_LIST =
135
            '-all,+meta-llama/llama-3.1-8b-instruct:free,+google/gemma-2-9b-it:free';
136

137
          const response = await router.getGlobalConfig();
138

139
          // Assert
140
          const result = response.languageModel?.openrouter;
141

142
          expect(result).toMatchSnapshot();
143

144
          process.env.OPENROUTER_MODEL_LIST = '';
145
        });
146
      });
147
    });
148
  });
149

150
  describe('getDefaultAgentConfig', () => {
151
    it('should return the default agent config', async () => {
152
      process.env.DEFAULT_AGENT_CONFIG =
153
        'plugins=search-engine,lobe-image-designer;enableAutoCreateTopic=true;model=gemini-pro;provider=google;';
154

155
      const response = await router.getDefaultAgentConfig();
156

157
      expect(response).toEqual({
158
        enableAutoCreateTopic: true,
159
        model: 'gemini-pro',
160
        plugins: ['search-engine', 'lobe-image-designer'],
161
        provider: 'google',
162
      });
163

164
      process.env.DEFAULT_AGENT_CONFIG = '';
165
    });
166

167
    it('should return another config', async () => {
168
      process.env.DEFAULT_AGENT_CONFIG =
169
        'model=meta-11ama/11ama-3-70b-instruct:nitro;provider=openrouter;enableAutoCreateTopic=true;params.max_tokens=700';
170

171
      const response = await router.getDefaultAgentConfig();
172

173
      expect(response).toEqual({
174
        enableAutoCreateTopic: true,
175
        model: 'meta-11ama/11ama-3-70b-instruct:nitro',
176
        params: { max_tokens: 700 },
177
        provider: 'openrouter',
178
      });
179

180
      process.env.DEFAULT_AGENT_CONFIG = '';
181
    });
182
  });
183
});
184