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

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

@@ -54,6 +54,169 @@ const createTestDocument = async (model: DocumentModel, fModel: FileModel, conte
 };
 
 describe('DocumentModel', () => {
+  describe('create', () => {
+    it('should create a new document', async () => {
+      const { id: fileId } = await fileModel.create({
+        fileType: 'text/plain',
+        name: 'test.txt',
+        size: 100,
+        url: 'https://example.com/test.txt',
+      });
+
+      const file = await fileModel.findById(fileId);
+      if (!file) throw new Error('File not found');
+
+      const result = await documentModel.create({
+        content: 'Test content',
+        fileId: file.id,
+        fileType: 'text/plain',
+        source: file.url,
+        sourceType: 'file',
+        totalCharCount: 12,
+        totalLineCount: 1,
+      });
+
+      expect(result).toBeDefined();
+      expect(result.content).toBe('Test content');
+      expect(result.fileId).toBe(file.id);
+    });
+  });
+
+  describe('delete', () => {
+    it('should delete a document', async () => {
+      const { documentId } = await createTestDocument(documentModel, fileModel, 'Test content');
+
+      await documentModel.delete(documentId);
+
+      const deleted = await documentModel.findById(documentId);
+      expect(deleted).toBeUndefined();
+    });
+
+    it('should not delete document belonging to another user', async () => {
+      const { documentId } = await createTestDocument(documentModel, fileModel, 'Test content');
+
+      // Try to delete with another user's model
+      await documentModel2.delete(documentId);
+
+      // Document should still exist
+      const stillExists = await documentModel.findById(documentId);
+      expect(stillExists).toBeDefined();
+    });
+  });
+
+  describe('deleteAll', () => {
+    it('should delete all documents for the user', async () => {
+      await createTestDocument(documentModel, fileModel, 'First document');
+      await createTestDocument(documentModel, fileModel, 'Second document');
+      await createTestDocument(documentModel2, fileModel2, 'Other user document');
+
+      await documentModel.deleteAll();
+
+      const userDocs = await documentModel.query();
+      const otherUserDocs = await documentModel2.query();
+
+      expect(userDocs).toHaveLength(0);
+      expect(otherUserDocs).toHaveLength(1);
+    });
+  });
+
+  describe('query', () => {
+    it('should return all documents for the user', async () => {
+      await createTestDocument(documentModel, fileModel, 'First document');
+      await createTestDocument(documentModel, fileModel, 'Second document');
+
+      const result = await documentModel.query();
+
+      expect(result).toHaveLength(2);
+    });
+
+    it('should only return documents for the current user', async () => {
+      await createTestDocument(documentModel, fileModel, 'User 1 document');
+      await createTestDocument(documentModel2, fileModel2, 'User 2 document');
+
+      const result = await documentModel.query();
+
+      expect(result).toHaveLength(1);
+      expect(result[0].content).toBe('User 1 document');
+    });
+
+    it('should return documents ordered by updatedAt desc', async () => {
+      const { documentId: doc1Id } = await createTestDocument(
+        documentModel,
+        fileModel,
+        'First document',
+      );
+      const { documentId: doc2Id } = await createTestDocument(
+        documentModel,
+        fileModel,
+        'Second document',
+      );
+
+      // Wait a bit to ensure timestamp difference
+      await new Promise((resolve) => setTimeout(resolve, 50));
+
+      // Update first document to make it more recent
+      await documentModel.update(doc1Id, { content: 'Updated first document' });
+
+      const result = await documentModel.query();
+
+      expect(result[0].id).toBe(doc1Id);
+      expect(result[1].id).toBe(doc2Id);
+    });
+  });
+
+  describe('findById', () => {
+    it('should find document by id', async () => {
+      const { documentId } = await createTestDocument(documentModel, fileModel, 'Test content');
+
+      const found = await documentModel.findById(documentId);
+
+      expect(found).toBeDefined();
+      expect(found?.id).toBe(documentId);
+      expect(found?.content).toBe('Test content');
+    });
+
+    it('should return undefined for non-existent document', async () => {
+      const found = await documentModel.findById('non-existent-id');
+
+      expect(found).toBeUndefined();
+    });
+
+    it('should not find document belonging to another user', async () => {
+      const { documentId } = await createTestDocument(documentModel, fileModel, 'Test content');
+
+      const found = await documentModel2.findById(documentId);
+
+      expect(found).toBeUndefined();
+    });
+  });
+
+  describe('update', () => {
+    it('should update a document', async () => {
+      const { documentId } = await createTestDocument(documentModel, fileModel, 'Original content');
+
+      await documentModel.update(documentId, {
+        content: 'Updated content',
+        totalCharCount: 15,
+      });
+
+      const updated = await documentModel.findById(documentId);
+
+      expect(updated?.content).toBe('Updated content');
+      expect(updated?.totalCharCount).toBe(15);
+    });
+
+    it('should not update document belonging to another user', async () => {
+      const { documentId } = await createTestDocument(documentModel, fileModel, 'Original content');
+
+      await documentModel2.update(documentId, { content: 'Hacked content' });
+
+      const unchanged = await documentModel.findById(documentId);
+
+      expect(unchanged?.content).toBe('Original content');
+    });
+  });
+
   describe('findByFileId', () => {
     it('should find document by fileId', async () => {
       const { documentId, file } = await createTestDocument(

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);
+    });
+  });
+});