@lobehub/chat 1.43.5 → 1.44.0

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

@@ -0,0 +1,373 @@
+// @vitest-environment node
+import { eq } from 'drizzle-orm/expressions';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { getTestDBInstance } from '@/database/server/core/dbForTest';
+import { ModelProvider } from '@/libs/agent-runtime';
+
+import { aiProviders, users } from '../../../schemas';
+import { AiProviderModel } from '../aiProvider';
+
+let serverDB = await getTestDBInstance();
+
+const userId = 'session-group-model-test-user-id';
+const aiProviderModel = new AiProviderModel(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(aiProviders).where(eq(aiProviders.userId, userId));
+});
+
+describe('AiProviderModel', () => {
+  describe('create', () => {
+    it('should create a new ai provider', async () => {
+      const params = {
+        name: 'AiHubMix',
+        id: 'aihubmix',
+        source: 'custom',
+      } as const;
+
+      const result = await aiProviderModel.create(params);
+      expect(result.id).toBeDefined();
+      expect(result).toMatchObject({ ...params, userId });
+
+      const group = await serverDB.query.aiProviders.findFirst({
+        where: eq(aiProviders.id, result.id),
+      });
+      expect(group).toMatchObject({ ...params, userId });
+    });
+  });
+  describe('delete', () => {
+    it('should delete a ai provider by id', async () => {
+      const { id } = await aiProviderModel.create({
+        name: 'AiHubMix',
+        id: 'aihubmix',
+        source: 'custom',
+      });
+
+      await aiProviderModel.delete(id);
+
+      const group = await serverDB.query.aiProviders.findFirst({
+        where: eq(aiProviders.id, id),
+      });
+      expect(group).toBeUndefined();
+    });
+  });
+  describe('deleteAll', () => {
+    it('should delete all ai providers for the user', async () => {
+      await aiProviderModel.create({ name: 'AiHubMix', source: 'custom', id: 'aihubmix' });
+      await aiProviderModel.create({ name: 'AiHubMix', source: 'custom', id: 'aihubmix-2' });
+
+      await aiProviderModel.deleteAll();
+
+      const userGroups = await serverDB.query.aiProviders.findMany({
+        where: eq(aiProviders.userId, userId),
+      });
+      expect(userGroups).toHaveLength(0);
+    });
+    it('should only delete ai providers for the user, not others', async () => {
+      await aiProviderModel.create({ name: 'AiHubMix', source: 'custom', id: 'aihubmix' });
+      await aiProviderModel.create({ name: 'AiHubMix', source: 'custom', id: 'aihubmix-2' });
+
+      const anotherAiProviderModel = new AiProviderModel(serverDB, 'user2');
+      await anotherAiProviderModel.create({
+        id: 'aihubmix',
+        source: 'custom',
+        name: 'Another provider',
+      });
+
+      await aiProviderModel.deleteAll();
+
+      const userGroups = await serverDB.query.aiProviders.findMany({
+        where: eq(aiProviders.userId, userId),
+      });
+      const total = await serverDB.query.aiProviders.findMany();
+      expect(userGroups).toHaveLength(0);
+      expect(total).toHaveLength(1);
+    });
+  });
+
+  describe('query', () => {
+    it('should query ai providers for the user', async () => {
+      await aiProviderModel.create({ name: 'AiHubMix', source: 'custom', id: 'aihubmix' });
+      await aiProviderModel.create({ name: 'AiHubMix', source: 'custom', id: 'aihubmix-2' });
+
+      const userGroups = await aiProviderModel.query();
+      expect(userGroups).toHaveLength(2);
+      expect(userGroups[0].id).toBe('aihubmix-2');
+      expect(userGroups[1].id).toBe('aihubmix');
+    });
+  });
+
+  describe('findById', () => {
+    it('should find a ai provider by id', async () => {
+      const { id } = await aiProviderModel.create({
+        name: 'AiHubMix',
+        source: 'custom',
+        id: 'aihubmix',
+      });
+
+      const group = await aiProviderModel.findById(id);
+      expect(group).toMatchObject({
+        id,
+        name: 'AiHubMix',
+        userId,
+      });
+    });
+  });
+
+  describe('update', () => {
+    it('should update a ai provider', async () => {
+      const { id } = await aiProviderModel.create({
+        name: 'AiHubMix',
+        source: 'custom',
+        id: 'aihubmix',
+      });
+
+      await aiProviderModel.update(id, { name: 'Updated Test Group', sort: 3 });
+
+      const updatedGroup = await serverDB.query.aiProviders.findFirst({
+        where: eq(aiProviders.id, id),
+      });
+      expect(updatedGroup).toMatchObject({
+        id,
+        name: 'Updated Test Group',
+        sort: 3,
+        userId,
+      });
+    });
+  });
+
+  describe('updateOrder', () => {
+    it('should update order of ai providers', async () => {
+      const group1 = await aiProviderModel.create({
+        name: 'AiHubMix',
+        source: 'custom',
+        id: 'aihubmix',
+      });
+      const group2 = await aiProviderModel.create({
+        name: 'AiHubMix',
+        source: 'custom',
+        id: 'aihubmix-2',
+      });
+
+      await aiProviderModel.updateOrder([
+        { id: group1.id, sort: 3 },
+        { id: group2.id, sort: 4 },
+      ]);
+
+      const updatedGroup1 = await serverDB.query.aiProviders.findFirst({
+        where: eq(aiProviders.id, group1.id),
+      });
+      const updatedGroup2 = await serverDB.query.aiProviders.findFirst({
+        where: eq(aiProviders.id, group2.id),
+      });
+
+      expect(updatedGroup1?.sort).toBe(3);
+      expect(updatedGroup2?.sort).toBe(4);
+    });
+  });
+
+  describe('getAiProviderList', () => {
+    it('should return a list of ai providers with selected fields', async () => {
+      await serverDB.insert(aiProviders).values({
+        description: 'Test description',
+        enabled: true,
+        id: 'aihubmix',
+        logo: 'test-logo',
+        name: 'AiHubMix',
+        sort: 1,
+        source: 'custom',
+        userId,
+      });
+
+      const list = await aiProviderModel.getAiProviderList();
+      expect(list).toHaveLength(1);
+      expect(list[0]).toMatchObject({
+        description: 'Test description',
+        enabled: true,
+        id: 'aihubmix',
+        logo: 'test-logo',
+        name: 'AiHubMix',
+        sort: 1,
+        source: 'custom',
+      });
+    });
+  });
+
+  describe('updateConfig', () => {
+    it('should update provider config with encryption', async () => {
+      const providerId = 'aihubmix';
+      await serverDB.insert(aiProviders).values({
+        id: providerId,
+        keyVaults: JSON.stringify({ key: 'value' }),
+        name: 'AiHubMix',
+        source: 'custom',
+        userId,
+      });
+
+      const mockEncryptor = vi.fn().mockResolvedValue('encrypted-data');
+      await aiProviderModel.updateConfig(
+        providerId,
+        {
+          keyVaults: { newKey: 'newValue' },
+          fetchOnClient: true,
+        },
+        mockEncryptor,
+      );
+
+      const updated = await serverDB.query.aiProviders.findFirst({
+        where: eq(aiProviders.id, providerId),
+      });
+
+      expect(mockEncryptor).toHaveBeenCalledWith(JSON.stringify({ newKey: 'newValue' }));
+      expect(updated?.keyVaults).toBe('encrypted-data');
+      expect(updated?.fetchOnClient).toBeTruthy();
+    });
+
+    it('should update provider config without encryption', async () => {
+      const providerId = 'aihubmix';
+      await serverDB.insert(aiProviders).values({
+        id: providerId,
+        keyVaults: JSON.stringify({ key: 'value' }),
+        name: 'AiHubMix',
+        source: 'custom',
+        userId,
+      });
+
+      await aiProviderModel.updateConfig(providerId, {
+        keyVaults: { newKey: 'newValue' },
+      });
+
+      const updated = await serverDB.query.aiProviders.findFirst({
+        where: eq(aiProviders.id, providerId),
+      });
+
+      expect(updated?.keyVaults).toBe(JSON.stringify({ newKey: 'newValue' }));
+    });
+  });
+
+  describe('toggleProviderEnabled', () => {
+    it('should toggle builtin provider enabled status', async () => {
+      const builtinId = ModelProvider.OpenAI;
+      await aiProviderModel.toggleProviderEnabled(builtinId, true);
+
+      const provider = await serverDB.query.aiProviders.findFirst({
+        where: eq(aiProviders.id, builtinId),
+      });
+
+      expect(provider?.enabled).toBe(true);
+      expect(provider?.source).toBe('builtin');
+    });
+
+    it('should toggle custom provider enabled status', async () => {
+      const customId = 'custom-provider';
+      await aiProviderModel.toggleProviderEnabled(customId, false);
+
+      const provider = await serverDB.query.aiProviders.findFirst({
+        where: eq(aiProviders.id, customId),
+      });
+
+      expect(provider?.enabled).toBe(false);
+      expect(provider?.source).toBe('custom');
+    });
+  });
+
+  describe('getAiProviderById', () => {
+    it('should get provider details with decryption', async () => {
+      const providerId = 'aihubmix';
+      const mockDecryptor = vi.fn().mockResolvedValue({ decryptedKey: 'value' });
+
+      await serverDB.insert(aiProviders).values({
+        id: providerId,
+        keyVaults: JSON.stringify({ key: 'value' }),
+        name: 'AiHubMix',
+        settings: { setting1: true } as any,
+        source: 'custom',
+        userId,
+      });
+
+      const provider = await aiProviderModel.getAiProviderById(providerId, mockDecryptor);
+
+      expect(provider).toBeDefined();
+      expect(provider?.keyVaults).toEqual({ decryptedKey: 'value' });
+    });
+
+    it('should handle non-existent provider for builtin provider', async () => {
+      const builtinId = ModelProvider.OpenAI;
+      const provider = await aiProviderModel.getAiProviderById(builtinId, (text) =>
+        JSON.parse(text as string),
+      );
+
+      expect(provider).toBeDefined();
+      expect(provider?.source).toBe('builtin');
+    });
+
+    it('should return undefined for non-existent custom provider', async () => {
+      const provider = await aiProviderModel.getAiProviderById('non-existent', (text) =>
+        JSON.parse(text as string),
+      );
+
+      expect(provider).toBeUndefined();
+    });
+
+    it('should handle null keyVaults', async () => {
+      const providerId = 'aihubmix';
+      await serverDB.insert(aiProviders).values({
+        id: providerId,
+        name: 'AiHubMix',
+        source: 'custom',
+        userId,
+      });
+
+      const provider = await aiProviderModel.getAiProviderById(providerId, (text) =>
+        JSON.parse(text as string),
+      );
+
+      expect(provider?.keyVaults).toEqual({});
+    });
+  });
+
+  describe('getAiProviderRuntimeConfig', () => {
+    it('should get runtime config for all providers', async () => {
+      const mockDecryptor = vi.fn().mockResolvedValue({ decryptedKey: 'value' });
+
+      await serverDB.insert(aiProviders).values([
+        {
+          fetchOnClient: true,
+          id: 'provider1',
+          keyVaults: JSON.stringify({ key: 'value' }),
+          name: 'Provider 1',
+          settings: { setting1: true } as any,
+          source: 'custom',
+          userId,
+        },
+        {
+          id: 'provider2',
+          name: 'Provider 2',
+          source: 'custom',
+          userId,
+        },
+      ]);
+
+      const config = await aiProviderModel.getAiProviderRuntimeConfig(mockDecryptor);
+
+      expect(config.provider1).toEqual({
+        fetchOnClient: true,
+        keyVaults: { decryptedKey: 'value' },
+        settings: { setting1: true },
+      });
+
+      expect(config.provider2).toEqual({
+        fetchOnClient: undefined,
+        keyVaults: {},
+        settings: {},
+      });
+    });
+  });
+});

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

@@ -0,0 +1,250 @@
+import { and, asc, desc, eq, inArray } from 'drizzle-orm/expressions';
+import pMap from 'p-map';
+
+import { LobeChatDatabase } from '@/database/type';
+import {
+  AiModelSortMap,
+  AiModelSourceEnum,
+  AiProviderModelListItem,
+  ToggleAiModelEnableParams,
+} from '@/types/aiModel';
+
+import { AiModelSelectItem, NewAiModelItem, aiModels } from '../../schemas';
+
+export class AiModelModel {
+  private userId: string;
+  private db: LobeChatDatabase;
+
+  constructor(db: LobeChatDatabase, userId: string) {
+    this.userId = userId;
+    this.db = db;
+  }
+
+  create = async (params: NewAiModelItem) => {
+    const [result] = await this.db
+      .insert(aiModels)
+      .values({
+        ...params,
+        enabled: true, // enabled by default
+        source: AiModelSourceEnum.Custom,
+        userId: this.userId,
+      })
+      .returning();
+
+    return result;
+  };
+
+  delete = async (id: string, providerId: string) => {
+    return this.db
+      .delete(aiModels)
+      .where(
+        and(
+          eq(aiModels.id, id),
+          eq(aiModels.providerId, providerId),
+          eq(aiModels.userId, this.userId),
+        ),
+      );
+  };
+
+  deleteAll = async () => {
+    return this.db.delete(aiModels).where(eq(aiModels.userId, this.userId));
+  };
+
+  query = async () => {
+    return this.db.query.aiModels.findMany({
+      orderBy: [desc(aiModels.updatedAt)],
+      where: eq(aiModels.userId, this.userId),
+    });
+  };
+
+  getModelListByProviderId = async (providerId: string) => {
+    const result = await this.db
+      .select({
+        abilities: aiModels.abilities,
+        config: aiModels.config,
+        contextWindowTokens: aiModels.contextWindowTokens,
+        description: aiModels.description,
+        displayName: aiModels.displayName,
+        enabled: aiModels.enabled,
+        id: aiModels.id,
+        pricing: aiModels.pricing,
+        source: aiModels.source,
+        type: aiModels.type,
+      })
+      .from(aiModels)
+      .where(and(eq(aiModels.providerId, providerId), eq(aiModels.userId, this.userId)))
+      .orderBy(
+        asc(aiModels.sort),
+        desc(aiModels.enabled),
+        desc(aiModels.releasedAt),
+        desc(aiModels.updatedAt),
+      );
+
+    return result as AiProviderModelListItem[];
+  };
+
+  getEnabledModels = async () => {
+    return this.db
+      .select({
+        abilities: aiModels.abilities,
+        config: aiModels.config,
+        contextWindowTokens: aiModels.contextWindowTokens,
+        displayName: aiModels.displayName,
+        enabled: aiModels.enabled,
+        id: aiModels.id,
+        providerId: aiModels.providerId,
+        source: aiModels.source,
+        type: aiModels.type,
+      })
+      .from(aiModels)
+      .where(and(eq(aiModels.userId, this.userId), eq(aiModels.enabled, true)));
+  };
+
+  findById = async (id: string) => {
+    return this.db.query.aiModels.findFirst({
+      where: and(eq(aiModels.id, id), eq(aiModels.userId, this.userId)),
+    });
+  };
+
+  update = async (id: string, providerId: string, value: Partial<AiModelSelectItem>) => {
+    return this.db
+      .insert(aiModels)
+      .values({ ...value, id, providerId, updatedAt: new Date(), userId: this.userId })
+      .onConflictDoUpdate({
+        set: value,
+        target: [aiModels.id, aiModels.providerId, aiModels.userId],
+      });
+  };
+
+  toggleModelEnabled = async (value: ToggleAiModelEnableParams) => {
+    return this.db
+      .insert(aiModels)
+      .values({ ...value, updatedAt: new Date(), userId: this.userId })
+      .onConflictDoUpdate({
+        set: { enabled: value.enabled, updatedAt: new Date() },
+        target: [aiModels.id, aiModels.providerId, aiModels.userId],
+      });
+  };
+
+  batchUpdateAiModels = async (providerId: string, models: AiProviderModelListItem[]) => {
+    return this.db.transaction(async (trx) => {
+      const records = models.map(({ id, ...model }) => ({
+        ...model,
+        id,
+        providerId,
+        updatedAt: new Date(),
+        userId: this.userId,
+      }));
+
+      // 第一步：尝试插入所有记录，忽略冲突
+      const insertedRecords = await trx
+        .insert(aiModels)
+        .values(records)
+        .onConflictDoNothing({
+          target: [aiModels.id, aiModels.userId, aiModels.providerId],
+        })
+        .returning();
+      // 第二步：找出需要更新的记录（即插入时发生冲突的记录）
+      // 找出未能插入的记录（需要更新的记录）
+      const insertedIds = new Set(insertedRecords.map((r) => r.id));
+      const recordsToUpdate = records.filter((r) => !insertedIds.has(r.id));
+
+      // 第三步：更新已存在的记录
+      if (recordsToUpdate.length > 0) {
+        await pMap(
+          recordsToUpdate,
+          async (record) => {
+            await trx
+              .update(aiModels)
+              .set({
+                ...record,
+                updatedAt: new Date(),
+              })
+              .where(
+                and(
+                  eq(aiModels.id, record.id),
+                  eq(aiModels.userId, this.userId),
+                  eq(aiModels.providerId, providerId),
+                ),
+              );
+          },
+          { concurrency: 10 }, // 限制并发数为 10
+        );
+      }
+    });
+  };
+
+  batchToggleAiModels = async (providerId: string, models: string[], enabled: boolean) => {
+    return this.db.transaction(async (trx) => {
+      // 1. insert models that are not in the db
+      const insertedRecords = await trx
+        .insert(aiModels)
+        .values(
+          models.map((i) => ({
+            enabled,
+            id: i,
+            providerId,
+            // if the model is not in the db, it's a builtin model
+            source: AiModelSourceEnum.Builtin,
+            updatedAt: new Date(),
+            userId: this.userId,
+          })),
+        )
+        .onConflictDoNothing({
+          target: [aiModels.id, aiModels.userId, aiModels.providerId],
+        })
+        .returning();
+
+      // 2. update models that are in the db
+      const insertedIds = new Set(insertedRecords.map((r) => r.id));
+      const recordsToUpdate = models.filter((r) => !insertedIds.has(r));
+
+      await trx
+        .update(aiModels)
+        .set({ enabled })
+        .where(
+          and(
+            eq(aiModels.providerId, providerId),
+            inArray(aiModels.id, recordsToUpdate),
+            eq(aiModels.userId, this.userId),
+          ),
+        );
+    });
+  };
+
+  clearRemoteModels(providerId: string) {
+    return this.db
+      .delete(aiModels)
+      .where(
+        and(
+          eq(aiModels.providerId, providerId),
+          eq(aiModels.source, AiModelSourceEnum.Remote),
+          eq(aiModels.userId, this.userId),
+        ),
+      );
+  }
+
+  updateModelsOrder = async (providerId: string, sortMap: AiModelSortMap[]) => {
+    await this.db.transaction(async (tx) => {
+      const updates = sortMap.map(({ id, sort }) => {
+        return tx
+          .insert(aiModels)
+          .values({
+            enabled: true,
+            id,
+            providerId,
+            sort,
+            // source: isBuiltin ? 'builtin' : 'custom',
+            updatedAt: new Date(),
+            userId: this.userId,
+          })
+          .onConflictDoUpdate({
+            set: { sort, updatedAt: new Date() },
+            target: [aiModels.id, aiModels.userId, aiModels.providerId],
+          });
+      });
+
+      await Promise.all(updates);
+    });
+  };
+}