@lobehub/chat 1.43.6 → 1.44.0

package/src/database/migrations/meta/_journal.json

@@ -91,6 +91,13 @@
       "when": 1731858381716,
       "tag": "0012_add_thread",
       "breakpoints": true
+    },
+    {
+      "idx": 13,
+      "version": "7",
+      "when": 1735834653361,
+      "tag": "0013_add_ai_infra",
+      "breakpoints": true
     }
   ],
   "version": "6"

package/src/database/repositories/aiInfra/index.ts

@@ -0,0 +1,115 @@
+import pMap from 'p-map';
+
+import { DEFAULT_MODEL_PROVIDER_LIST } from '@/config/modelProviders';
+import { AiModelModel } from '@/database/server/models/aiModel';
+import { AiProviderModel } from '@/database/server/models/aiProvider';
+import { LobeChatDatabase } from '@/database/type';
+import { AIChatModelCard, AiModelSourceEnum, AiProviderModelListItem } from '@/types/aiModel';
+import { AiProviderListItem, EnabledAiModel } from '@/types/aiProvider';
+import { ProviderConfig } from '@/types/user/settings';
+import { mergeArrayById } from '@/utils/merge';
+
+export class AiInfraRepos {
+  private userId: string;
+  private db: LobeChatDatabase;
+  aiProviderModel: AiProviderModel;
+  private providerConfigs: Record<string, ProviderConfig>;
+  private aiModelModel: AiModelModel;
+
+  constructor(
+    db: LobeChatDatabase,
+    userId: string,
+    providerConfigs: Record<string, ProviderConfig>,
+  ) {
+    this.userId = userId;
+    this.db = db;
+    this.aiProviderModel = new AiProviderModel(db, userId);
+    this.aiModelModel = new AiModelModel(db, userId);
+    this.providerConfigs = providerConfigs;
+  }
+
+  /**
+   * Calculate the final providerList based on the known providerConfig
+   */
+  getAiProviderList = async () => {
+    const userProviders = await this.aiProviderModel.getAiProviderList();
+
+    // 1. 先创建一个基于 DEFAULT_MODEL_PROVIDER_LIST id 顺序的映射
+    const orderMap = new Map(DEFAULT_MODEL_PROVIDER_LIST.map((item, index) => [item.id, index]));
+
+    const builtinProviders = DEFAULT_MODEL_PROVIDER_LIST.map((item) => ({
+      description: item.description,
+      enabled:
+        userProviders.some((provider) => provider.id === item.id && provider.enabled) ||
+        this.providerConfigs[item.id]?.enabled,
+      id: item.id,
+      name: item.name,
+      source: 'builtin',
+    })) as AiProviderListItem[];
+
+    const mergedProviders = mergeArrayById(builtinProviders, userProviders);
+
+    // 3. 根据 orderMap 排序
+    return mergedProviders.sort((a, b) => {
+      const orderA = orderMap.get(a.id) ?? Number.MAX_SAFE_INTEGER;
+      const orderB = orderMap.get(b.id) ?? Number.MAX_SAFE_INTEGER;
+      return orderA - orderB;
+    });
+  };
+
+  getUserEnabledProviderList = async () => {
+    const list = await this.getAiProviderList();
+    return list
+      .filter((item) => item.enabled)
+      .sort((a, b) => a.sort! - b.sort!)
+      .map((item) => ({ id: item.id, name: item.name, source: item.source }));
+  };
+
+  getEnabledModels = async () => {
+    const providers = await this.getAiProviderList();
+    const enabledProviders = providers.filter((item) => item.enabled);
+
+    const userEnabledModels = await this.aiModelModel.getEnabledModels();
+    const modelList = await pMap(
+      enabledProviders,
+      async (provider) => {
+        const aiModels = await this.fetchBuiltinModels(provider.id);
+
+        return (aiModels || [])
+          .filter((i) => i.enabled)
+          .map<EnabledAiModel>((item) => ({
+            ...item,
+            abilities: item.abilities || {},
+            providerId: provider.id,
+          }));
+      },
+      { concurrency: 10 },
+    );
+
+    return [...modelList.flat(), ...userEnabledModels] as EnabledAiModel[];
+  };
+
+  getAiProviderModelList = async (providerId: string) => {
+    const aiModels = await this.aiModelModel.getModelListByProviderId(providerId);
+
+    const defaultModels: AiProviderModelListItem[] =
+      (await this.fetchBuiltinModels(providerId)) || [];
+
+    return mergeArrayById(defaultModels, aiModels) as AiProviderModelListItem[];
+  };
+
+  private fetchBuiltinModels = async (
+    providerId: string,
+  ): Promise<AiProviderModelListItem[] | undefined> => {
+    try {
+      const { default: providerModels } = await import(`@/config/aiModels/${providerId}`);
+      return (providerModels as AIChatModelCard[]).map<AiProviderModelListItem>((m) => ({
+        ...m,
+        enabled: m.enabled || false,
+        source: AiModelSourceEnum.Builtin,
+      }));
+    } catch {
+      // maybe provider id not exist
+    }
+  };
+}

package/src/database/schemas/aiInfra.ts

@@ -0,0 +1,69 @@
+/* eslint-disable sort-keys-fix/sort-keys-fix  */
+import { boolean, integer, jsonb, pgTable, primaryKey, text, varchar } from 'drizzle-orm/pg-core';
+
+import { timestamps } from '@/database/schemas/_helpers';
+import { users } from '@/database/schemas/user';
+import { AiProviderSettings } from '@/types/aiProvider';
+
+export const aiProviders = pgTable(
+  'ai_providers',
+  {
+    id: varchar('id', { length: 64 }).notNull(),
+    name: text('name'),
+
+    userId: text('user_id')
+      .references(() => users.id, { onDelete: 'cascade' })
+      .notNull(),
+
+    sort: integer('sort'),
+    enabled: boolean('enabled'),
+    fetchOnClient: boolean('fetch_on_client'),
+    checkModel: text('check_model'),
+    logo: text('logo'),
+    description: text('description'),
+
+    // need to be encrypted
+    keyVaults: text('key_vaults'),
+    source: varchar('source', { enum: ['builtin', 'custom'], length: 20 }),
+    settings: jsonb('settings')
+      .$defaultFn(() => ({}))
+      .$type<AiProviderSettings>(),
+
+    ...timestamps,
+  },
+  (table) => [primaryKey({ columns: [table.id, table.userId] })],
+);
+
+export type NewAiProviderItem = Omit<typeof aiProviders.$inferInsert, 'userId'>;
+export type AiProviderSelectItem = typeof aiProviders.$inferSelect;
+
+export const aiModels = pgTable(
+  'ai_models',
+  {
+    id: varchar('id', { length: 150 }).notNull(),
+    displayName: varchar('display_name', { length: 200 }),
+    description: text('description'),
+    organization: varchar('organization', { length: 100 }),
+    enabled: boolean('enabled'),
+    providerId: varchar('provider_id', { length: 64 }).notNull(),
+    type: varchar('type', { length: 20 }).default('chat').notNull(),
+    sort: integer('sort'),
+
+    userId: text('user_id')
+      .references(() => users.id, { onDelete: 'cascade' })
+      .notNull(),
+    pricing: jsonb('pricing'),
+    parameters: jsonb('parameters').default({}),
+    config: jsonb('config'),
+    abilities: jsonb('abilities').default({}),
+    contextWindowTokens: integer('context_window_tokens'),
+    source: varchar('source', { enum: ['remote', 'custom', 'builtin'], length: 20 }),
+    releasedAt: varchar('released_at', { length: 10 }),
+
+    ...timestamps,
+  },
+  (table) => [primaryKey({ columns: [table.id, table.providerId, table.userId] })],
+);
+
+export type NewAiModelItem = Omit<typeof aiModels.$inferInsert, 'userId'>;
+export type AiModelSelectItem = typeof aiModels.$inferSelect;

package/src/database/schemas/index.ts

@@ -1,4 +1,5 @@
 export * from './agent';
+export * from './aiInfra';
 export * from './asyncTask';
 export * from './file';
 export * from './message';

package/src/database/server/models/__tests__/aiModel.test.ts

@@ -0,0 +1,318 @@
+// @vitest-environment node
+import { eq } from 'drizzle-orm/expressions';
+import { afterEach, beforeEach, describe, expect, it } from 'vitest';
+
+import { getTestDBInstance } from '@/database/server/core/dbForTest';
+import { AiProviderModelListItem } from '@/types/aiModel';
+
+import { AiModelSelectItem, NewAiModelItem, aiModels, users } from '../../../schemas';
+import { AiModelModel } from '../aiModel';
+
+let serverDB = await getTestDBInstance();
+
+const userId = 'ai-model-test-user-id';
+const aiProviderModel = new AiModelModel(serverDB, userId);
+
+beforeEach(async () => {
+  await serverDB.delete(users);
+  await serverDB.insert(users).values([{ id: userId }, { id: 'user2' }]);
+});
+
+afterEach(async () => {
+  await serverDB.delete(users).where(eq(users.id, userId));
+  await serverDB.delete(aiModels).where(eq(aiModels.userId, userId));
+});
+
+describe('AiModelModel', () => {
+  describe('create', () => {
+    it('should create a new ai provider', async () => {
+      const params: NewAiModelItem = {
+        organization: 'Qwen',
+        id: 'qvq',
+        providerId: 'openai',
+      };
+
+      const result = await aiProviderModel.create(params);
+      expect(result.id).toBeDefined();
+      expect(result).toMatchObject({ ...params, userId });
+
+      const group = await serverDB.query.aiModels.findFirst({
+        where: eq(aiModels.id, result.id),
+      });
+      expect(group).toMatchObject({ ...params, userId });
+    });
+  });
+  describe('delete', () => {
+    it('should delete a ai provider by id', async () => {
+      const { id } = await aiProviderModel.create({
+        organization: 'Qwen',
+        providerId: 'openai',
+        id: 'qvq',
+      });
+
+      await aiProviderModel.delete(id, 'openai');
+
+      const group = await serverDB.query.aiModels.findFirst({
+        where: eq(aiModels.id, id),
+      });
+      expect(group).toBeUndefined();
+    });
+  });
+  describe('deleteAll', () => {
+    it('should delete all ai providers for the user', async () => {
+      await aiProviderModel.create({ organization: 'Qwen', providerId: 'openai', id: 'qvq' });
+      await aiProviderModel.create({
+        organization: 'Qwen',
+        providerId: 'openai',
+        id: 'aihubmix-2',
+      });
+
+      await aiProviderModel.deleteAll();
+
+      const userGroups = await serverDB.query.aiModels.findMany({
+        where: eq(aiModels.userId, userId),
+      });
+      expect(userGroups).toHaveLength(0);
+    });
+    it('should only delete ai providers for the user, not others', async () => {
+      await aiProviderModel.create({ organization: 'Qwen', providerId: 'openai', id: 'qvq' });
+      await aiProviderModel.create({
+        organization: 'Qwen',
+        providerId: 'openai',
+        id: 'aihubmix-2',
+      });
+
+      const anotherAiModelModel = new AiModelModel(serverDB, 'user2');
+      await anotherAiModelModel.create({ id: 'qvq', providerId: 'openai' });
+
+      await aiProviderModel.deleteAll();
+
+      const userGroups = await serverDB.query.aiModels.findMany({
+        where: eq(aiModels.userId, userId),
+      });
+      const total = await serverDB.query.aiModels.findMany();
+      expect(userGroups).toHaveLength(0);
+      expect(total).toHaveLength(1);
+    });
+  });
+
+  describe('query', () => {
+    it('should query ai providers for the user', async () => {
+      await aiProviderModel.create({ organization: 'Qwen', providerId: 'openai', id: 'qvq' });
+      await aiProviderModel.create({
+        organization: 'Qwen',
+        providerId: 'openai',
+        id: 'aihubmix-2',
+      });
+
+      const userGroups = await aiProviderModel.query();
+      expect(userGroups).toHaveLength(2);
+      expect(userGroups[0].id).toBe('aihubmix-2');
+      expect(userGroups[1].id).toBe('qvq');
+    });
+  });
+
+  describe('findById', () => {
+    it('should find a ai provider by id', async () => {
+      const { id } = await aiProviderModel.create({
+        organization: 'Qwen',
+        providerId: 'openai',
+        id: 'qvq',
+      });
+
+      const group = await aiProviderModel.findById(id);
+      expect(group).toMatchObject({
+        id,
+        organization: 'Qwen',
+        providerId: 'openai',
+
+        userId,
+      });
+    });
+  });
+
+  describe('update', () => {
+    it('should update a ai provider', async () => {
+      const { id } = await aiProviderModel.create({
+        organization: 'Qwen',
+        providerId: 'openai',
+        id: 'qvq',
+      });
+
+      await aiProviderModel.update(id, 'openai', {
+        displayName: 'Updated Test Group',
+        contextWindowTokens: 3000,
+      });
+
+      const updatedGroup = await serverDB.query.aiModels.findFirst({
+        where: eq(aiModels.id, id),
+      });
+      expect(updatedGroup).toMatchObject({
+        id,
+        displayName: 'Updated Test Group',
+        contextWindowTokens: 3000,
+        userId,
+      });
+    });
+  });
+
+  describe('getModelListByProviderId', () => {
+    it('should get model list by provider id', async () => {
+      await aiProviderModel.create({
+        id: 'model1',
+        providerId: 'openai',
+        sort: 1,
+        enabled: true,
+      });
+      await aiProviderModel.create({
+        id: 'model2',
+        providerId: 'openai',
+        sort: 2,
+        enabled: false,
+      });
+
+      const models = await aiProviderModel.getModelListByProviderId('openai');
+      expect(models).toHaveLength(2);
+      expect(models[0].id).toBe('model1');
+      expect(models[1].id).toBe('model2');
+    });
+
+    it('should only return models for specified provider', async () => {
+      await aiProviderModel.create({
+        id: 'model1',
+        providerId: 'openai',
+      });
+      await aiProviderModel.create({
+        id: 'model2',
+        providerId: 'anthropic',
+      });
+
+      const models = await aiProviderModel.getModelListByProviderId('openai');
+      expect(models).toHaveLength(1);
+      expect(models[0].id).toBe('model1');
+    });
+  });
+
+  describe('getEnabledModels', () => {
+    it('should only return enabled models', async () => {
+      await serverDB.insert(aiModels).values([
+        { id: 'model1', providerId: 'openai', enabled: true, source: 'custom', userId },
+        { id: 'model2', providerId: 'openai', enabled: false, source: 'custom', userId },
+      ]);
+
+      const models = await aiProviderModel.getEnabledModels();
+      expect(models).toHaveLength(1);
+      expect(models[0].id).toBe('model1');
+      expect(models[0].enabled).toBe(true);
+    });
+  });
+
+  describe('toggleModelEnabled', () => {
+    it('should toggle model enabled status', async () => {
+      const model = await aiProviderModel.create({
+        id: 'model1',
+        providerId: 'openai',
+        enabled: true,
+      });
+
+      await aiProviderModel.toggleModelEnabled({
+        id: model.id,
+        providerId: 'openai',
+        enabled: false,
+      });
+
+      const updatedModel = await aiProviderModel.findById(model.id);
+      expect(updatedModel?.enabled).toBe(false);
+    });
+  });
+
+  describe('batchUpdateAiModels', () => {
+    it('should insert new models and update existing ones', async () => {
+      // Create an initial model
+      await aiProviderModel.create({
+        id: 'existing-model',
+        providerId: 'openai',
+        displayName: 'Old Name',
+      });
+
+      const models = [
+        {
+          id: 'existing-model',
+          displayName: 'Updated Name',
+        },
+        {
+          id: 'new-model',
+          displayName: 'New Model',
+        },
+      ] as AiProviderModelListItem[];
+
+      await aiProviderModel.batchUpdateAiModels('openai', models);
+
+      const allModels = await aiProviderModel.query();
+      expect(allModels).toHaveLength(2);
+      expect(allModels.find((m) => m.id === 'existing-model')?.displayName).toBe('Updated Name');
+      expect(allModels.find((m) => m.id === 'new-model')?.displayName).toBe('New Model');
+    });
+  });
+
+  describe('batchToggleAiModels', () => {
+    it('should toggle multiple models enabled status', async () => {
+      await aiProviderModel.create({
+        id: 'model1',
+        providerId: 'openai',
+        enabled: false,
+      });
+      await aiProviderModel.create({
+        id: 'model2',
+        providerId: 'openai',
+        enabled: false,
+      });
+
+      await aiProviderModel.batchToggleAiModels('openai', ['model1', 'model2'], true);
+
+      const models = await aiProviderModel.query();
+      expect(models.every((m) => m.enabled)).toBe(true);
+    });
+  });
+
+  describe('clearRemoteModels', () => {
+    it('should delete all remote models for a provider', async () => {
+      await serverDB.insert(aiModels).values([
+        { id: 'remote1', providerId: 'openai', source: 'remote', userId },
+        { id: 'custom1', providerId: 'openai', source: 'custom', userId },
+      ]);
+
+      await aiProviderModel.clearRemoteModels('openai');
+
+      const remainingModels = await aiProviderModel.query();
+      expect(remainingModels).toHaveLength(1);
+      expect(remainingModels[0].id).toBe('custom1');
+    });
+  });
+
+  describe('updateModelsOrder', () => {
+    it('should update the sort order of models', async () => {
+      await aiProviderModel.create({
+        id: 'model1',
+        providerId: 'openai',
+        sort: 1,
+      });
+      await aiProviderModel.create({
+        id: 'model2',
+        providerId: 'openai',
+        sort: 2,
+      });
+
+      const sortMap = [
+        { id: 'model1', sort: 2 },
+        { id: 'model2', sort: 1 },
+      ];
+
+      await aiProviderModel.updateModelsOrder('openai', sortMap);
+
+      const models = await aiProviderModel.getModelListByProviderId('openai');
+      expect(models[0].id).toBe('model2');
+      expect(models[1].id).toBe('model1');
+    });
+  });
+});