@lobehub/lobehub 2.0.0-next.137 → 2.0.0-next.139

package/packages/database/src/models/__tests__/embedding.test.ts

@@ -0,0 +1,294 @@
+import { eq } from 'drizzle-orm';
+import { afterEach, beforeEach, describe, expect, it } from 'vitest';
+
+import { chunks, embeddings, users } from '../../schemas';
+import { LobeChatDatabase } from '../../type';
+import { EmbeddingModel } from '../embedding';
+import { getTestDB } from './_util';
+import { designThinkingQuery } from './fixtures/embedding';
+
+const userId = 'embedding-user-test';
+const otherUserId = 'other-user-test';
+
+const serverDB: LobeChatDatabase = await getTestDB();
+const embeddingModel = new EmbeddingModel(serverDB, userId);
+
+describe('EmbeddingModel', () => {
+  beforeEach(async () => {
+    await serverDB.delete(users);
+    await serverDB.insert(users).values([{ id: userId }, { id: otherUserId }]);
+  });
+
+  afterEach(async () => {
+    await serverDB.delete(users);
+  });
+
+  describe('create', () => {
+    it('should create a new embedding', async () => {
+      // Create a chunk first
+      const [chunk] = await serverDB
+        .insert(chunks)
+        .values({ text: 'Test chunk', userId })
+        .returning();
+
+      const id = await embeddingModel.create({
+        chunkId: chunk.id,
+        embeddings: designThinkingQuery,
+        model: 'text-embedding-ada-002',
+      });
+
+      expect(id).toBeDefined();
+
+      const created = await serverDB.query.embeddings.findFirst({
+        where: eq(embeddings.id, id),
+      });
+
+      expect(created).toBeDefined();
+      expect(created?.chunkId).toBe(chunk.id);
+      expect(created?.model).toBe('text-embedding-ada-002');
+      expect(created?.userId).toBe(userId);
+    });
+  });
+
+  describe('bulkCreate', () => {
+    it('should create multiple embeddings', async () => {
+      // Create chunks first
+      const [chunk1, chunk2] = await serverDB
+        .insert(chunks)
+        .values([
+          { text: 'Test chunk 1', userId },
+          { text: 'Test chunk 2', userId },
+        ])
+        .returning();
+
+      await embeddingModel.bulkCreate([
+        { chunkId: chunk1.id, embeddings: designThinkingQuery, model: 'text-embedding-ada-002' },
+        { chunkId: chunk2.id, embeddings: designThinkingQuery, model: 'text-embedding-ada-002' },
+      ]);
+
+      const created = await serverDB.query.embeddings.findMany({
+        where: eq(embeddings.userId, userId),
+      });
+
+      expect(created).toHaveLength(2);
+    });
+
+    it('should handle duplicate chunkId with onConflictDoNothing', async () => {
+      // Create a chunk
+      const [chunk] = await serverDB
+        .insert(chunks)
+        .values({ text: 'Test chunk', userId })
+        .returning();
+
+      // Create first embedding
+      await embeddingModel.create({
+        chunkId: chunk.id,
+        embeddings: designThinkingQuery,
+        model: 'text-embedding-ada-002',
+      });
+
+      // Try to create duplicate
+      await embeddingModel.bulkCreate([
+        { chunkId: chunk.id, embeddings: designThinkingQuery, model: 'text-embedding-3-small' },
+      ]);
+
+      const created = await serverDB.query.embeddings.findMany({
+        where: eq(embeddings.chunkId, chunk.id),
+      });
+
+      // Should still have only 1 embedding due to unique constraint
+      expect(created).toHaveLength(1);
+      expect(created[0].model).toBe('text-embedding-ada-002');
+    });
+  });
+
+  describe('delete', () => {
+    it('should delete an embedding', async () => {
+      const [chunk] = await serverDB
+        .insert(chunks)
+        .values({ text: 'Test chunk', userId })
+        .returning();
+
+      const id = await embeddingModel.create({
+        chunkId: chunk.id,
+        embeddings: designThinkingQuery,
+        model: 'text-embedding-ada-002',
+      });
+
+      await embeddingModel.delete(id);
+
+      const deleted = await serverDB.query.embeddings.findFirst({
+        where: eq(embeddings.id, id),
+      });
+
+      expect(deleted).toBeUndefined();
+    });
+
+    it('should not delete embedding belonging to another user', async () => {
+      // Create chunk and embedding for other user
+      const [chunk] = await serverDB
+        .insert(chunks)
+        .values({ text: 'Other user chunk', userId: otherUserId })
+        .returning();
+
+      const [otherEmbedding] = await serverDB
+        .insert(embeddings)
+        .values({
+          chunkId: chunk.id,
+          embeddings: designThinkingQuery,
+          model: 'text-embedding-ada-002',
+          userId: otherUserId,
+        })
+        .returning();
+
+      await embeddingModel.delete(otherEmbedding.id);
+
+      const stillExists = await serverDB.query.embeddings.findFirst({
+        where: eq(embeddings.id, otherEmbedding.id),
+      });
+
+      expect(stillExists).toBeDefined();
+    });
+  });
+
+  describe('query', () => {
+    it('should return all embeddings for the user', async () => {
+      const [chunk1, chunk2] = await serverDB
+        .insert(chunks)
+        .values([
+          { text: 'Test chunk 1', userId },
+          { text: 'Test chunk 2', userId },
+        ])
+        .returning();
+
+      await serverDB.insert(embeddings).values([
+        { chunkId: chunk1.id, embeddings: designThinkingQuery, userId },
+        { chunkId: chunk2.id, embeddings: designThinkingQuery, userId },
+      ]);
+
+      const result = await embeddingModel.query();
+
+      expect(result).toHaveLength(2);
+    });
+
+    it('should only return embeddings for the current user', async () => {
+      const [chunk1] = await serverDB
+        .insert(chunks)
+        .values([{ text: 'Test chunk 1', userId }])
+        .returning();
+
+      const [chunk2] = await serverDB
+        .insert(chunks)
+        .values([{ text: 'Other user chunk', userId: otherUserId }])
+        .returning();
+
+      await serverDB.insert(embeddings).values([
+        { chunkId: chunk1.id, embeddings: designThinkingQuery, userId },
+        { chunkId: chunk2.id, embeddings: designThinkingQuery, userId: otherUserId },
+      ]);
+
+      const result = await embeddingModel.query();
+
+      expect(result).toHaveLength(1);
+      expect(result[0].userId).toBe(userId);
+    });
+  });
+
+  describe('findById', () => {
+    it('should return an embedding by id', async () => {
+      const [chunk] = await serverDB
+        .insert(chunks)
+        .values({ text: 'Test chunk', userId })
+        .returning();
+
+      const id = await embeddingModel.create({
+        chunkId: chunk.id,
+        embeddings: designThinkingQuery,
+        model: 'text-embedding-ada-002',
+      });
+
+      const result = await embeddingModel.findById(id);
+
+      expect(result).toBeDefined();
+      expect(result?.id).toBe(id);
+      expect(result?.chunkId).toBe(chunk.id);
+    });
+
+    it('should return undefined for non-existent embedding', async () => {
+      // Use a valid UUID format that doesn't exist
+      const result = await embeddingModel.findById('00000000-0000-0000-0000-000000000000');
+
+      expect(result).toBeUndefined();
+    });
+
+    it('should not return embedding belonging to another user', async () => {
+      const [chunk] = await serverDB
+        .insert(chunks)
+        .values({ text: 'Other user chunk', userId: otherUserId })
+        .returning();
+
+      const [otherEmbedding] = await serverDB
+        .insert(embeddings)
+        .values({
+          chunkId: chunk.id,
+          embeddings: designThinkingQuery,
+          userId: otherUserId,
+        })
+        .returning();
+
+      const result = await embeddingModel.findById(otherEmbedding.id);
+
+      expect(result).toBeUndefined();
+    });
+  });
+
+  describe('countUsage', () => {
+    it('should return the count of embeddings for the user', async () => {
+      const [chunk1, chunk2, chunk3] = await serverDB
+        .insert(chunks)
+        .values([
+          { text: 'Test chunk 1', userId },
+          { text: 'Test chunk 2', userId },
+          { text: 'Test chunk 3', userId },
+        ])
+        .returning();
+
+      await serverDB.insert(embeddings).values([
+        { chunkId: chunk1.id, embeddings: designThinkingQuery, userId },
+        { chunkId: chunk2.id, embeddings: designThinkingQuery, userId },
+        { chunkId: chunk3.id, embeddings: designThinkingQuery, userId },
+      ]);
+
+      const count = await embeddingModel.countUsage();
+
+      expect(count).toBe(3);
+    });
+
+    it('should return 0 when user has no embeddings', async () => {
+      const count = await embeddingModel.countUsage();
+
+      expect(count).toBe(0);
+    });
+
+    it('should only count embeddings for the current user', async () => {
+      const [chunk1] = await serverDB
+        .insert(chunks)
+        .values([{ text: 'Test chunk 1', userId }])
+        .returning();
+
+      const [chunk2] = await serverDB
+        .insert(chunks)
+        .values([{ text: 'Other user chunk', userId: otherUserId }])
+        .returning();
+
+      await serverDB.insert(embeddings).values([
+        { chunkId: chunk1.id, embeddings: designThinkingQuery, userId },
+        { chunkId: chunk2.id, embeddings: designThinkingQuery, userId: otherUserId },
+      ]);
+
+      const count = await embeddingModel.countUsage();
+
+      expect(count).toBe(1);
+    });
+  });
+});

package/packages/database/src/models/__tests__/oauthHandoff.test.ts

@@ -0,0 +1,261 @@
+import { eq } from 'drizzle-orm';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { oauthHandoffs } from '../../schemas';
+import { LobeChatDatabase } from '../../type';
+import { OAuthHandoffModel } from '../oauthHandoff';
+import { getTestDB } from './_util';
+
+const serverDB: LobeChatDatabase = await getTestDB();
+const oauthHandoffModel = new OAuthHandoffModel(serverDB);
+
+describe('OAuthHandoffModel', () => {
+  beforeEach(async () => {
+    await serverDB.delete(oauthHandoffs);
+  });
+
+  afterEach(async () => {
+    await serverDB.delete(oauthHandoffs);
+    vi.useRealTimers();
+  });
+
+  describe('create', () => {
+    it('should create a new OAuth handoff record', async () => {
+      const result = await oauthHandoffModel.create({
+        id: 'handoff-1',
+        client: 'desktop',
+        payload: { code: 'auth-code', state: 'state-value' },
+      });
+
+      expect(result).toBeDefined();
+      expect(result.id).toBe('handoff-1');
+      expect(result.client).toBe('desktop');
+      expect(result.payload).toEqual({ code: 'auth-code', state: 'state-value' });
+    });
+
+    it('should handle conflict by doing nothing', async () => {
+      // Create first record
+      await oauthHandoffModel.create({
+        id: 'handoff-1',
+        client: 'desktop',
+        payload: { code: 'original-code' },
+      });
+
+      // Try to create duplicate - should not throw
+      const result = await oauthHandoffModel.create({
+        id: 'handoff-1',
+        client: 'desktop',
+        payload: { code: 'new-code' },
+      });
+
+      // Result is undefined when conflict occurs with onConflictDoNothing
+      expect(result).toBeUndefined();
+
+      // Original should remain unchanged
+      const record = await serverDB.query.oauthHandoffs.findFirst({
+        where: eq(oauthHandoffs.id, 'handoff-1'),
+      });
+
+      expect(record?.payload).toEqual({ code: 'original-code' });
+    });
+  });
+
+  describe('fetchAndConsume', () => {
+    it('should fetch and delete valid credentials', async () => {
+      await oauthHandoffModel.create({
+        id: 'handoff-1',
+        client: 'desktop',
+        payload: { code: 'auth-code', state: 'state-value' },
+      });
+
+      const result = await oauthHandoffModel.fetchAndConsume('handoff-1', 'desktop');
+
+      expect(result).toBeDefined();
+      expect(result?.id).toBe('handoff-1');
+      expect(result?.payload).toEqual({ code: 'auth-code', state: 'state-value' });
+
+      // Verify it was deleted
+      const deleted = await serverDB.query.oauthHandoffs.findFirst({
+        where: eq(oauthHandoffs.id, 'handoff-1'),
+      });
+      expect(deleted).toBeUndefined();
+    });
+
+    it('should return null for non-existent credentials', async () => {
+      const result = await oauthHandoffModel.fetchAndConsume('non-existent', 'desktop');
+
+      expect(result).toBeNull();
+    });
+
+    it('should return null when client type does not match', async () => {
+      await oauthHandoffModel.create({
+        id: 'handoff-1',
+        client: 'desktop',
+        payload: { code: 'auth-code' },
+      });
+
+      const result = await oauthHandoffModel.fetchAndConsume('handoff-1', 'browser-extension');
+
+      expect(result).toBeNull();
+
+      // Record should still exist
+      const record = await serverDB.query.oauthHandoffs.findFirst({
+        where: eq(oauthHandoffs.id, 'handoff-1'),
+      });
+      expect(record).toBeDefined();
+    });
+
+    it('should return null for expired credentials (older than 5 minutes)', async () => {
+      // Create a record with old timestamp
+      const sixMinutesAgo = new Date(Date.now() - 6 * 60 * 1000);
+
+      await serverDB.insert(oauthHandoffs).values({
+        id: 'handoff-expired',
+        client: 'desktop',
+        payload: { code: 'auth-code' },
+        createdAt: sixMinutesAgo,
+        updatedAt: sixMinutesAgo,
+      });
+
+      const result = await oauthHandoffModel.fetchAndConsume('handoff-expired', 'desktop');
+
+      expect(result).toBeNull();
+    });
+
+    it('should return valid credentials within 5 minutes', async () => {
+      // Create a record with timestamp 4 minutes ago
+      const fourMinutesAgo = new Date(Date.now() - 4 * 60 * 1000);
+
+      await serverDB.insert(oauthHandoffs).values({
+        id: 'handoff-valid',
+        client: 'desktop',
+        payload: { code: 'auth-code' },
+        createdAt: fourMinutesAgo,
+        updatedAt: fourMinutesAgo,
+      });
+
+      const result = await oauthHandoffModel.fetchAndConsume('handoff-valid', 'desktop');
+
+      expect(result).toBeDefined();
+      expect(result?.id).toBe('handoff-valid');
+    });
+  });
+
+  describe('cleanupExpired', () => {
+    it('should delete expired records', async () => {
+      const sixMinutesAgo = new Date(Date.now() - 6 * 60 * 1000);
+      const now = new Date();
+
+      // Create expired record
+      await serverDB.insert(oauthHandoffs).values({
+        id: 'handoff-expired',
+        client: 'desktop',
+        payload: { code: 'expired-code' },
+        createdAt: sixMinutesAgo,
+        updatedAt: sixMinutesAgo,
+      });
+
+      // Create valid record
+      await serverDB.insert(oauthHandoffs).values({
+        id: 'handoff-valid',
+        client: 'desktop',
+        payload: { code: 'valid-code' },
+        createdAt: now,
+        updatedAt: now,
+      });
+
+      await oauthHandoffModel.cleanupExpired();
+
+      // Verify expired record is deleted
+      const expiredRecord = await serverDB.query.oauthHandoffs.findFirst({
+        where: eq(oauthHandoffs.id, 'handoff-expired'),
+      });
+      expect(expiredRecord).toBeUndefined();
+
+      // Verify valid record still exists
+      const validRecord = await serverDB.query.oauthHandoffs.findFirst({
+        where: eq(oauthHandoffs.id, 'handoff-valid'),
+      });
+      expect(validRecord).toBeDefined();
+    });
+
+    it('should not throw when no expired records exist', async () => {
+      await expect(oauthHandoffModel.cleanupExpired()).resolves.not.toThrow();
+    });
+
+    it('should delete multiple expired records', async () => {
+      const sixMinutesAgo = new Date(Date.now() - 6 * 60 * 1000);
+
+      await serverDB.insert(oauthHandoffs).values([
+        {
+          id: 'handoff-expired-1',
+          client: 'desktop',
+          payload: { code: 'code-1' },
+          createdAt: sixMinutesAgo,
+          updatedAt: sixMinutesAgo,
+        },
+        {
+          id: 'handoff-expired-2',
+          client: 'browser',
+          payload: { code: 'code-2' },
+          createdAt: sixMinutesAgo,
+          updatedAt: sixMinutesAgo,
+        },
+      ]);
+
+      await oauthHandoffModel.cleanupExpired();
+
+      // Verify all expired records are deleted
+      const remainingRecords = await serverDB.query.oauthHandoffs.findMany();
+      expect(remainingRecords).toHaveLength(0);
+    });
+  });
+
+  describe('exists', () => {
+    it('should return true for existing valid credentials', async () => {
+      await oauthHandoffModel.create({
+        id: 'handoff-1',
+        client: 'desktop',
+        payload: { code: 'auth-code' },
+      });
+
+      const result = await oauthHandoffModel.exists('handoff-1', 'desktop');
+
+      expect(result).toBe(true);
+    });
+
+    it('should return false for non-existent credentials', async () => {
+      const result = await oauthHandoffModel.exists('non-existent', 'desktop');
+
+      expect(result).toBe(false);
+    });
+
+    it('should return false when client type does not match', async () => {
+      await oauthHandoffModel.create({
+        id: 'handoff-1',
+        client: 'desktop',
+        payload: { code: 'auth-code' },
+      });
+
+      const result = await oauthHandoffModel.exists('handoff-1', 'browser-extension');
+
+      expect(result).toBe(false);
+    });
+
+    it('should return false for expired credentials', async () => {
+      const sixMinutesAgo = new Date(Date.now() - 6 * 60 * 1000);
+
+      await serverDB.insert(oauthHandoffs).values({
+        id: 'handoff-expired',
+        client: 'desktop',
+        payload: { code: 'auth-code' },
+        createdAt: sixMinutesAgo,
+        updatedAt: sixMinutesAgo,
+      });
+
+      const result = await oauthHandoffModel.exists('handoff-expired', 'desktop');
+
+      expect(result).toBe(false);
+    });
+  });
+});