Both uploading and chat work

quzhi1 · quzhi1 · commit cae8ba947069 · 2025-04-12T15:33:40.000-07:00
diff --git a/app/api/upload/route.ts b/app/api/upload/route.ts
@@ -34,71 +34,117 @@ export async function POST(req: Request) {
     }
 
     console.log('JSON parsed successfully, items:', jsonData.length);
-    const results = [];
+    interface ProcessedItem {
+      id: string;
+      source: string;
+      chapter: string;
+    }
+    const results: ProcessedItem[] = [];
     const totalItems = jsonData.length;
 
-    for (let i = 0; i < jsonData.length; i++) {
-      const item = jsonData[i];
-      if (!item.source || !item.chapter || !item.text) {
-        console.log('Skipping invalid item:', item);
-        continue;
-      }
+    // Create a TransformStream for streaming responses
+    const stream = new TransformStream();
+    const writer = stream.writable.getWriter();
+    const encoder = new TextEncoder();
 
-      console.log('Processing item', i + 1, 'of', totalItems);
-      
+    // Start processing in the background
+    (async () => {
       try {
-        // Insert into resources table
-        console.log('Inserting into resources table...');
-        const [resource] = await db
-          .insert(resources)
-          .values({
-            content: item.text,
-            source: item.source,
-            chapter: item.chapter,
-          })
-          .returning();
+        // Process items in batches to show progress
+        const batchSize = 5;
+        for (let i = 0; i < jsonData.length; i += batchSize) {
+          const batch = jsonData.slice(i, i + batchSize);
+          
+          for (const item of batch) {
+            if (!item.source || !item.chapter || !item.text) {
+              console.log('Skipping invalid item:', item);
+              continue;
+            }
 
-        console.log('Resource inserted:', resource.id);
+            try {
+              // Insert into resources table
+              const [resource] = await db
+                .insert(resources)
+                .values({
+                  content: item.text,
+                  source: item.source,
+                  chapter: item.chapter,
+                })
+                .returning();
 
-        // Generate embeddings
-        console.log('Generating embeddings...');
-        const embeddingResults = await generateEmbeddings(item.text);
-        console.log('Embeddings generated:', embeddingResults.length);
-        
-        // Insert embeddings
-        console.log('Inserting embeddings...');
-        for (const embeddingResult of embeddingResults) {
-          await db.insert(embeddings).values({
-            resourceId: resource.id,
-            content: embeddingResult.content,
-            embedding: embeddingResult.embedding,
-          });
-        }
+              // Generate embeddings
+              const embeddingResults = await generateEmbeddings(item.text);
+              
+              // Insert embeddings
+              for (const embeddingResult of embeddingResults) {
+                await db.insert(embeddings).values({
+                  resourceId: resource.id,
+                  content: embeddingResult.content,
+                  embedding: embeddingResult.embedding,
+                });
+              }
+
+              results.push({
+                id: resource.id,
+                source: item.source,
+                chapter: item.chapter,
+              });
 
-        results.push({
-          id: resource.id,
-          source: item.source,
-          chapter: item.chapter,
-        });
+            } catch (dbError) {
+              console.error('Database error processing item:', dbError);
+              continue;
+            }
+          }
 
-      } catch (dbError) {
-        console.error('Database error processing item:', dbError);
-        // Continue with next item even if this one fails
-        continue;
+          // Send progress update after each batch
+          const processed = Math.min(i + batchSize, jsonData.length);
+          console.log(`Processed ${processed} of ${totalItems} items`);
+          
+          await writer.write(encoder.encode(
+            `data: ${JSON.stringify({
+              type: 'progress',
+              processed,
+              total: totalItems,
+            })}\n\n`
+          ));
+        }
+
+        // Send final completion message
+        await writer.write(encoder.encode(
+          `data: ${JSON.stringify({
+            type: 'complete',
+            success: true,
+            processed: results.length,
+            total: totalItems,
+            results,
+          })}\n\n`
+        ));
+      } catch (error) {
+        console.error('Error during processing:', error);
+        await writer.write(encoder.encode(
+          `data: ${JSON.stringify({
+            type: 'error',
+            error: 'Error processing file',
+            details: error instanceof Error ? error.message : 'Unknown error',
+          })}\n\n`
+        ));
+      } finally {
+        await writer.close();
       }
-    }
+    })();
 
-    console.log('Upload process completed');
-    return NextResponse.json({ 
-      success: true, 
-      processed: results.length,
-      total: totalItems,
-      results 
+    return new Response(stream.readable, {
+      headers: {
+        'Content-Type': 'text/event-stream',
+        'Cache-Control': 'no-cache',
+        'Connection': 'keep-alive',
+      },
     });
 
   } catch (error) {
     console.error('Error processing file:', error);
     return NextResponse.json({ 
+      type: 'error',
       error: 'Error processing file',
       details: error instanceof Error ? error.message : 'Unknown error'
     }, { status: 500 });
diff --git a/components/FileUpload.tsx b/components/FileUpload.tsx
@@ -12,6 +12,14 @@ export function FileUpload() {
   const [processedItems, setProcessedItems] = useState(0);
   const abortControllerRef = useRef<AbortController | null>(null);
 
+  const resetState = () => {
+    setIsUploading(false);
+    setProgress(0);
+    setTotalItems(0);
+    setProcessedItems(0);
+    abortControllerRef.current = null;
+  };
+
   const handleFileUpload = async (event: React.ChangeEvent<HTMLInputElement>) => {
     const file = event.target.files?.[0];
     if (!file) return;
@@ -20,6 +28,7 @@ export function FileUpload() {
     setIsUploading(true);
     setProgress(0);
     setProcessedItems(0);
+    setTotalItems(0);
     abortControllerRef.current = new AbortController();
 
     const formData = new FormData();
@@ -38,30 +47,48 @@ export function FileUpload() {
         throw new Error('Upload failed');
       }
 
-      const data = await response.json();
-      console.log('Upload completed:', data);
-
-      if (data.success) {
-        setProcessedItems(data.processed);
-        setTotalItems(data.total);
-        setProgress(100);
-        toast.success(`Successfully processed ${data.processed} items`);
-      } else {
-        toast.error(data.error || 'Error uploading file');
+      const reader = response.body?.getReader();
+      if (!reader) {
+        throw new Error('No reader available');
       }
 
+      while (true) {
+        const { done, value } = await reader.read();
+        if (done) break;
+
+        const text = new TextDecoder().decode(value);
+        const lines = text.split('\n');
+        
+        for (const line of lines) {
+          if (line.startsWith('data: ')) {
+            const data = JSON.parse(line.slice(6));
+            console.log('Received data:', data);
+
+            if (data.type === 'progress') {
+              setProcessedItems(data.processed);
+              setTotalItems(data.total);
+              setProgress((data.processed / data.total) * 100);
+            } else if (data.type === 'complete') {
+              setProcessedItems(data.processed);
+              setTotalItems(data.total);
+              setProgress(100);
+              toast.success(`Successfully processed ${data.processed} items`);
+              setTimeout(resetState, 2000);
+            } else if (data.type === 'error') {
+              toast.error(data.error || 'Error uploading file');
+              resetState();
+            }
+          }
+        }
+      }
     } catch (error) {
       console.error('Upload error:', error);
       if (error instanceof Error && error.name === 'AbortError') {
         toast.info('Upload cancelled');
       } else {
         toast.error('Error uploading file');
       }
-    } finally {
-      setIsUploading(false);
-      setProgress(0);
-      setProcessedItems(0);
-      abortControllerRef.current = null;
+      resetState();
     }
   };
 
@@ -98,7 +125,7 @@ export function FileUpload() {
           </Button>
         )}
       </div>
-      {isUploading && (
+      {(isUploading || progress > 0) && (
         <div className="w-full max-w-md space-y-2">
           <Progress value={progress} className="w-full" />
           <p className="text-sm text-center text-gray-500">
diff --git a/install.sh b/install.sh
@@ -41,4 +41,8 @@ pnpm db:push
 echo "🤖 Installing AI dependencies..."
 pnpm add ai @ai-sdk/react @ai-sdk/openai @radix-ui/react-progress
 
+# Install test dependencies
+echo "🧪 Installing test dependencies..."
+pnpm add -D @types/jest jest ts-jest @jest/globals babel-jest @babel/core @babel/preset-env @babel/plugin-transform-runtime crypto-browserify
+
 echo "✨ Installation complete! You can now start the development server with 'pnpm dev'" 
diff --git a/jest.config.js b/jest.config.js
@@ -0,0 +1,19 @@
+/** @type {import('ts-jest').JestConfigWithTsJest} */
+module.exports = {
+  preset: 'ts-jest',
+  testEnvironment: 'node',
+  moduleFileExtensions: ['ts', 'tsx', 'js', 'jsx', 'json', 'mjs'],
+  transform: {
+    '^.+\\.tsx?$': ['ts-jest', { useESM: true }],
+    '^.+\\.m?js$': 'babel-jest'
+  },
+  moduleNameMapper: {
+    '^@/(.*)$': '<rootDir>/$1',
+    '^nanoid$': '<rootDir>/node_modules/nanoid/nanoid.js',
+    '^node:crypto$': '<rootDir>/node_modules/crypto-browserify/index.js'
+  },
+  transformIgnorePatterns: [
+    'node_modules/(?!(nanoid|clsx|tailwind-merge|@ai-sdk)/)'
+  ],
+  extensionsToTreatAsEsm: ['.ts'],
+}; 
diff --git a/lib/ai/embedding.test.ts b/lib/ai/embedding.test.ts
@@ -0,0 +1,71 @@
+import { describe, expect, test, jest } from '@jest/globals';
+import { generateChunks } from './embedding';
+
+// Mock the dependencies
+jest.mock('nanoid', () => ({
+  customAlphabet: () => () => 'test-id'
+}));
+
+jest.mock('clsx', () => ({
+  clsx: () => ''
+}));
+
+jest.mock('tailwind-merge', () => ({
+  twMerge: () => ''
+}));
+
+jest.mock('@/lib/db', () => ({}));
+jest.mock('@/lib/env.mjs', () => ({}));
+jest.mock('ai', () => ({}));
+
+describe('Text Chunking', () => {
+  test('should split Chinese text into sentences correctly', () => {
+    const input = '高祖，沛丰邑中阳里人，姓刘氏，字季。父曰太公，母曰刘媪。其先刘媪尝息大泽之陂，梦与神遇。是时雷电晦冥，太公往视，则见蛟龙於其上。已而有身，遂产高祖。';
+    const expected = [
+      '高祖，沛丰邑中阳里人，姓刘氏，字季',
+      '父曰太公，母曰刘媪',
+      '其先刘媪尝息大泽之陂，梦与神遇',
+      '是时雷电晦冥，太公往视，则见蛟龙於其上',
+      '已而有身，遂产高祖'
+    ];
+    
+    const result = generateChunks(input);
+    expect(result).toEqual(expected);
+  });
+
+  test('should handle mixed Chinese and English punctuation', () => {
+    const input = '这是第一句。This is the second sentence! 这是第三句？';
+    const expected = [
+      '这是第一句',
+      'This is the second sentence',
+      '这是第三句'
+    ];
+    
+    const result = generateChunks(input);
+    expect(result).toEqual(expected);
+  });
+
+  test('should handle empty input', () => {
+    const input = '';
+    const expected: string[] = [];
+    
+    const result = generateChunks(input);
+    expect(result).toEqual(expected);
+  });
+
+  test('should handle input with only whitespace', () => {
+    const input = '   \n  \t  ';
+    const expected: string[] = [];
+    
+    const result = generateChunks(input);
+    expect(result).toEqual(expected);
+  });
+
+  test('should trim whitespace from chunks', () => {
+    const input = '  第一句。  第二句  。  第三句  ';
+    const expected = ['第一句', '第二句', '第三句'];
+    
+    const result = generateChunks(input);
+    expect(result).toEqual(expected);
+  });
+}); 
diff --git a/lib/ai/embedding.ts b/lib/ai/embedding.ts
@@ -9,10 +9,13 @@ const embeddingModel = openai.embedding('text-embedding-ada-002');
 const generateChunks = (input: string): string[] => {
   return input
     .trim()
-    .split(/[。.]/)
-    .filter(i => i !== '');
+    .split(/[。！？!?]/)
+    .filter(i => i.trim().length > 0)
+    .map(i => i.trim());
 };
 
+export { generateChunks };
+
 export const generateEmbeddings = async (
   value: string,
 ): Promise<Array<{ embedding: number[]; content: string }>> => {
diff --git a/package.json b/package.json
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
