feat: Implement streaming response for chat

main.py

@@ -1,6 +1,6 @@
 import requests
 from fastapi import FastAPI, Request, HTTPException
-from fastapi.responses import FileResponse, RedirectResponse
+from fastapi.responses import FileResponse, RedirectResponse, StreamingResponse
 from fastapi.staticfiles import StaticFiles
 
 OLLAMA_API_URL = "http://localhost:11434"
@@ -19,16 +19,20 @@ async def chat_endpoint(request: Request):
         raise HTTPException(status_code=400, detail="Prompt is required")
 
     url = f"{OLLAMA_API_URL}/api/generate"
-    
-    try:
-        response = requests.post(
-            url=url,
-            json={"model": model, "prompt": prompt, "stream": False}
-        )
-        
-        return response.json()
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+
+    async def stream_response():
+        try:
+            response = requests.post(
+                url=url,
+                json={"model": model, "prompt": prompt, "stream": True},
+                stream=True
+            )
+            for chunk in response.iter_content(chunk_size=None):
+                yield chunk
+        except Exception as e:
+            print(f"Error during streaming: {e}")
+
+    return StreamingResponse(stream_response(), media_type="application/x-ndjson")
 
 @app.get("/")
 async def root():
@@ -36,4 +40,4 @@ async def root():
 
 @app.get("/chat")
 async def chat_page():
-    return FileResponse('chat.html')
+    return FileResponse('chat.html')

static/script.js

@@ -18,6 +18,8 @@ function sendMessage() {
     appendMessage(message, 'user-message');
     messageInput.value = '';
 
+    appendMessage('', 'bot-message'); // Create an empty message element for the bot's response
+
     fetch('/chat_api', {
         method: 'POST',
         headers: {
@@ -28,13 +30,42 @@ function sendMessage() {
             prompt: message
         })
     })
-    .then(response => response.json())
-    .then(data => {
-        appendMessage(data.response, 'bot-message');
+    .then(response => {
+        const reader = response.body.getReader();
+        const decoder = new TextDecoder();
+        let partialResponse = '';
+
+        function read() {
+            reader.read().then(({ done, value }) => {
+                if (done) {
+                    return;
+                }
+                partialResponse += decoder.decode(value, { stream: true });
+                try {
+                    const jsonObjects = partialResponse.split('\n');
+                    partialResponse = jsonObjects.pop(); // Keep the last partial object
+                    for (const jsonObjStr of jsonObjects) {
+                        if (jsonObjStr) {
+                            const jsonObj = JSON.parse(jsonObjStr);
+                            if (jsonObj.response) {
+                                appendToLastMessage(jsonObj.response);
+                            }
+                        }
+                    }
+                } catch (e) {
+                    // JSON parsing error, wait for more chunks
+                }
+                read();
+            });
+        }
+        read();
     })
     .catch(error => {
         console.error('Error:', error);
-        appendMessage('Sorry, something went wrong.', 'bot-message');
+        const lastMessage = messageList.lastChild;
+        if (lastMessage && lastMessage.classList.contains('bot-message')) {
+            lastMessage.textContent = 'Sorry, something went wrong.';
+        }
     });
 }
 
@@ -45,3 +76,11 @@ function appendMessage(message, className) {
     messageList.appendChild(li);
     messageList.scrollTop = messageList.scrollHeight;
 }
+
+function appendToLastMessage(text) {
+    const lastMessage = messageList.lastChild;
+    if (lastMessage && lastMessage.classList.contains('bot-message')) {
+        lastMessage.textContent += text;
+        messageList.scrollTop = messageList.scrollHeight;
+    }
+}