Browse Source

first demo

Signed-off-by: ryjiang <jiangruiyi@gmail.com>
ryjiang 1 month ago
parent
commit
dd0c4f31b6
3 changed files with 241 additions and 37 deletions
  1. 167 23
      client/src/pages/ai/AIChat.tsx
  2. 32 3
      server/src/ai/ai.controller.ts
  3. 42 11
      server/src/ai/ai.service.ts

+ 167 - 23
client/src/pages/ai/AIChat.tsx

@@ -1,22 +1,151 @@
-import { FC } from 'react';
+import { FC, useState, useEffect, useRef } from 'react';
 import { useTranslation } from 'react-i18next';
 import { Box, Typography, Chip, Stack } from '@mui/material';
-import { useChat } from '@ai-sdk/react';
+
+interface Message {
+  id: string;
+  role: 'user' | 'assistant';
+  content: string;
+  createdAt: string;
+  parts?: Array<{
+    type: string;
+    text: string;
+  }>;
+}
 
 const AIChat: FC = () => {
   const { t } = useTranslation('ai');
   const apiKey = localStorage.getItem('attu.ui.openai_api_key');
+  const [messages, setMessages] = useState<Message[]>([]);
+  const [input, setInput] = useState('');
+  const [isLoading, setIsLoading] = useState(false);
+  const eventSourceRef = useRef<EventSource | null>(null);
+  const currentMessageRef = useRef<Message | null>(null);
+  const accumulatedContentRef = useRef('');
+
+  const handleSubmit = async (e: React.FormEvent) => {
+    e.preventDefault();
+    if (!input.trim() || isLoading) return;
+
+    // Add user message
+    const userMessage: Message = {
+      id: `msg_${Date.now()}`,
+      role: 'user',
+      content: input,
+      createdAt: new Date().toISOString(),
+    };
+    setMessages(prev => [...prev, userMessage]);
+    setInput('');
+    setIsLoading(true);
+    accumulatedContentRef.current = ''; // Reset accumulated content
+
+    // Close existing EventSource if any
+    if (eventSourceRef.current) {
+      eventSourceRef.current.close();
+    }
+
+    try {
+      // First, send the initial request with headers
+      const response = await fetch('/api/v1/ai/chat', {
+        method: 'POST',
+        headers: {
+          'Content-Type': 'application/json',
+          'x-openai-api-key': apiKey || '',
+        },
+        body: JSON.stringify({
+          messages: [...messages, userMessage],
+        }),
+      });
+
+      if (!response.ok) {
+        throw new Error(`HTTP error! status: ${response.status}`);
+      }
+
+      // Then create EventSource for streaming
+      const eventSource = new EventSource(
+        `/api/v1/ai/chat?messages=${encodeURIComponent(JSON.stringify([...messages, userMessage]))}&x-openai-api-key=${encodeURIComponent(apiKey || '')}`
+      );
+      eventSourceRef.current = eventSource;
+
+      // Initialize current message
+      const messageId = `msg_${Date.now()}`;
+      currentMessageRef.current = {
+        id: messageId,
+        role: 'assistant',
+        content: '',
+        createdAt: new Date().toISOString(),
+        parts: [
+          {
+            type: 'text',
+            text: '',
+          },
+        ],
+      };
+
+      // Add initial empty message to the list
+      setMessages(prev => [...prev, { ...currentMessageRef.current! }]);
 
-  const { messages, input, handleInputChange, handleSubmit, isLoading } =
-    useChat({
-      api: '/api/v1/ai/chat',
-      headers: {
-        'x-openai-api-key': apiKey || '',
-      },
-      onError: (error: Error) => {
-        console.error('Chat error:', error);
-      },
-    });
+      // Handle messages
+      eventSource.onmessage = event => {
+        if (event.data === '[DONE]') {
+          // console.log('Stream completed');
+          eventSource.close();
+          setIsLoading(false);
+          return;
+        }
+
+        try {
+          const data = JSON.parse(event.data);
+          // console.log('Received message:', data);
+          if (data.value) {
+            const newContent = data.value.parts[0].text;
+            // console.log('New content:', newContent);
+
+            // Accumulate content
+            accumulatedContentRef.current += newContent;
+
+            setMessages(prev => {
+              const newMessages = [...prev];
+              const lastMessage = newMessages[newMessages.length - 1];
+              if (lastMessage && lastMessage.id === messageId) {
+                return newMessages.map(msg =>
+                  msg.id === messageId
+                    ? {
+                        ...msg,
+                        content: accumulatedContentRef.current,
+                        parts: [
+                          {
+                            type: 'text',
+                            text: accumulatedContentRef.current,
+                          },
+                        ],
+                      }
+                    : msg
+                );
+              }
+              return newMessages;
+            });
+          }
+        } catch (error) {
+          console.error('Error parsing message:', error);
+        }
+      };
+
+      // Handle errors
+      eventSource.onerror = error => {
+        console.error('EventSource error:', error);
+        eventSource.close();
+        setIsLoading(false);
+      };
+    } catch (error) {
+      console.error('Error:', error);
+      setIsLoading(false);
+    }
+  };
+
+  const handleInputChange = (e: React.ChangeEvent<HTMLInputElement>) => {
+    setInput(e.target.value);
+  };
 
   const suggestedCommands = [
     t('suggestions.search'),
@@ -25,6 +154,15 @@ const AIChat: FC = () => {
     t('suggestions.insert'),
   ];
 
+  // Cleanup EventSource on unmount
+  useEffect(() => {
+    return () => {
+      if (eventSourceRef.current) {
+        eventSourceRef.current.close();
+      }
+    };
+  }, []);
+
   return (
     <Box
       sx={{ height: '100%', display: 'flex', flexDirection: 'column', p: 2 }}
@@ -46,6 +184,7 @@ const AIChat: FC = () => {
           flexDirection: 'column',
           gap: 2,
           mb: 2,
+          overflowY: 'auto',
         }}
       >
         {messages.map(message => (
@@ -71,14 +210,21 @@ const AIChat: FC = () => {
                 boxShadow: 1,
               }}
             >
-              <Typography variant="body1">
-                {message.parts.map((part, index) => {
-                  if (part.type === 'text') {
-                    return <span key={index}>{part.text}</span>;
-                  }
-                  return null;
-                })}
-              </Typography>
+              {message.role === 'assistant' &&
+              message.id === currentMessageRef.current?.id ? (
+                <Typography variant="body1">{message.content || ''}</Typography>
+              ) : (
+                <Typography variant="body1">
+                  {message.parts?.map((part, index) => {
+                    if (part.type === 'text') {
+                      return <span key={index}>{part.text || ''}</span>;
+                    }
+                    return null;
+                  }) ||
+                    message.content ||
+                    ''}
+                </Typography>
+              )}
             </Box>
           </Box>
         ))}
@@ -101,9 +247,7 @@ const AIChat: FC = () => {
             <Chip
               key={index}
               label={command}
-              onClick={() =>
-                handleInputChange({ target: { value: command } } as any)
-              }
+              onClick={() => setInput(command)}
               sx={{ mb: 1 }}
             />
           ))}

+ 32 - 3
server/src/ai/ai.controller.ts

@@ -12,7 +12,19 @@ export class AIController {
   }
 
   generateRoutes() {
+    // Handle preflight requests
+    this.router.options('/chat', (req, res) => {
+      res.header('Access-Control-Allow-Origin', '*');
+      res.header('Access-Control-Allow-Methods', 'GET, POST, OPTIONS');
+      res.header(
+        'Access-Control-Allow-Headers',
+        'Content-Type, x-openai-api-key'
+      );
+      res.status(204).end();
+    });
+
     this.router.post('/chat', this.handleChat.bind(this));
+    this.router.get('/chat', this.handleChat.bind(this));
     return this.router;
   }
 
@@ -22,16 +34,33 @@ export class AIController {
     next: NextFunction
   ) {
     try {
-      const apiKey = req.headers['x-openai-api-key'] as string;
+      // Get API key from either headers or query parameters
+      const apiKey =
+        (req.headers['x-openai-api-key'] as string) ||
+        (req.query['x-openai-api-key'] as string);
       if (!apiKey) {
         return res.status(400).json({
           error: 'API key is required',
         });
       }
 
-      await this.aiService.chat(req.body, res, apiKey);
+      // Set CORS headers for SSE
+      res.header('Access-Control-Allow-Origin', '*');
+      res.header(
+        'Access-Control-Allow-Headers',
+        'Content-Type, x-openai-api-key'
+      );
+      res.header('Access-Control-Allow-Methods', 'GET, POST, OPTIONS');
+
+      // Get messages from either body (POST) or query (GET)
+      const messages =
+        req.method === 'POST'
+          ? req.body.messages
+          : JSON.parse(req.query.messages as string);
+
+      await this.aiService.chat({ messages }, res, apiKey);
     } catch (error) {
       next(error);
     }
   }
-}
+}

+ 42 - 11
server/src/ai/ai.service.ts

@@ -11,39 +11,70 @@ export class AIService {
         throw new Error('API key is required');
       }
 
-      // Initialize OpenAI client only when needed
+      // Initialize OpenAI client with DeepSeek configuration
       this.openai = new OpenAI({
-        apiKey,
+        baseURL: 'https://api.deepseek.com',
+        apiKey: apiKey,
       });
 
       const stream = await this.openai.chat.completions.create({
-        model: 'gpt-3.5-turbo',
+        model: 'deepseek-chat',
         messages: chatRequest.messages,
         stream: true,
+        max_tokens: 4096, // Default max output tokens for deepseek-chat
       });
 
-      // Set headers for streaming response
-      res.setHeader('Content-Type', 'text/event-stream');
-      res.setHeader('Cache-Control', 'no-cache');
-      res.setHeader('Connection', 'keep-alive');
+      // Set headers for SSE
+      res.writeHead(200, {
+        'Content-Type': 'text/event-stream',
+        'Cache-Control': 'no-cache',
+        Connection: 'keep-alive',
+        'Transfer-Encoding': 'chunked',
+      });
+
+      let messageId = `msg_${Date.now()}`;
 
       // Handle the stream
       for await (const chunk of stream) {
         const content = chunk.choices[0]?.delta?.content || '';
         if (content) {
-          res.write(`data: ${JSON.stringify({ content })}\n\n`);
+          // Send each chunk immediately without accumulation
+          const message = {
+            id: messageId,
+            role: 'assistant',
+            content: content,
+            createdAt: new Date().toISOString(),
+            parts: [
+              {
+                type: 'text',
+                text: content,
+              },
+            ],
+          };
+
+          // Flush the response immediately
+          const data = JSON.stringify({
+            type: 'text',
+            value: message,
+          });
+          res.write(`data: ${data}\n\n`);
+          res.flush?.(); // Flush if available
         }
       }
 
+      // Send the final message
       res.write('data: [DONE]\n\n');
       res.end();
     } catch (error) {
-      console.error('OpenAI API error:', error);
+      console.error('DeepSeek API error:', error);
       if (!res.headersSent) {
         res.status(500).json({
-          error: error instanceof Error ? error.message : 'Failed to get response from OpenAI',
+          error:
+            error instanceof Error
+              ? error.message
+              : 'Failed to get response from DeepSeek',
         });
       }
     }
   }
-}
+}