backend for streming works

CKennte11 · CKennte11 · commit 37d90e96a6c4 · 2025-02-05T16:58:00.000+01:00
diff --git a/backend/app/api/routes/learn.py b/backend/app/api/routes/learn.py
@@ -1,5 +1,6 @@
 from typing import Literal
 from fastapi import APIRouter, Depends, HTTPException
+from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 from app.api import deps
 from app.core.ai_client import ChatManager, AnthropicClient, OpenAIClient
@@ -66,6 +67,29 @@ async def chat_general(
     
     return ChatResponse(message=response)
 
+@router.post("/learn/chat/stream")
+async def chat_stream(
+    request: ChatRequest,
+    current_user = Depends(deps.get_current_user),
+):
+    """Streaming chat endpoint."""
+    chat_key = f"{current_user.id}_general"
+    if chat_key not in active_chats:
+        active_chats[chat_key] = ChatManager(client=request.model)
+    
+    return StreamingResponse(
+        active_chats[chat_key].stream_message(
+            request.message,
+            system=request.system_prompt
+        ),
+        media_type='text/event-stream',
+        headers={
+            'Cache-Control': 'no-cache',
+            'Connection': 'keep-alive',
+            'X-Accel-Buffering': 'no'  # Disable buffering in nginx
+        }
+    )
+
 @router.post("/learn/{path_id}", response_model=ChatResponse)
 async def chat(
     path_id: str,
diff --git a/backend/app/core/ai_client.py b/backend/app/core/ai_client.py
@@ -1,20 +1,27 @@
 from typing import AsyncGenerator, Dict, Protocol, Literal
 from typing import Literal, TypedDict
-from anthropic import Anthropic
+from anthropic import AsyncAnthropic
 from openai import AsyncOpenAI
 from app.core.config import settings
+import json
+import logging
+
+logger = logging.getLogger(__name__)
 
 class Message(TypedDict):
     role: Literal["user", "assistant", "system"]
     content: str
 
-class BaseAIClient:
+class BaseAIClient(Protocol):
     async def chat(self, messages: list[Message], system: str | None = None) -> str:
         raise NotImplementedError
+        
+    async def chat_stream(self, messages: list[Message], system: str | None = None) -> AsyncGenerator[str, None]:
+        raise NotImplementedError
 
 class AnthropicClient(BaseAIClient):
     def __init__(self):
-        self.client = Anthropic(api_key=settings.ANTHROPIC_API_KEY)
+        self.client = AsyncAnthropic(api_key=settings.ANTHROPIC_API_KEY)
         self.model = settings.ANTHROPIC_MODEL
     
     async def chat(self, messages: list[Message], system: str | None = None) -> str:
@@ -27,6 +34,26 @@ async def chat(self, messages: list[Message], system: str | None = None) -> str:
         )
         return response.content[0].text
 
+    async def chat_stream(self, messages: list[Message], system: str | None = None) -> AsyncGenerator[str, None]:
+        """Stream chat responses."""
+        request_params = {
+            'messages': [{"role": m["role"], "content": m["content"]} for m in messages],
+            'model': self.model,
+            'max_tokens': settings.MAX_TOKENS,
+            'temperature': settings.TEMPERATURE,
+        }
+        if system:
+            request_params['system'] = system
+        
+        async with self.client.messages.stream(**request_params) as stream:
+            async for text in stream.text_stream:
+                yield f"data: {json.dumps({'type': 'content', 'content': text})}\n\n"
+            
+            # Get the final message for history
+            message = await stream.get_final_message()
+            self.add_message("assistant", message.content)  # Add to history
+            yield f"data: {json.dumps({'type': 'done', 'content': ''})}\n\n"
+
 class OpenAIClient(BaseAIClient):
     def __init__(self):
         self.client = AsyncOpenAI(api_key=settings.OPENAI_API_KEY)
@@ -44,6 +71,22 @@ async def chat(self, messages: list[Message], system: str | None = None) -> str:
         )
         return response.choices[0].message.content
 
+    async def chat_stream(self, messages: list[Message], system: str | None = None) -> AsyncGenerator[str, None]:
+        """Stream chat responses."""
+        if system:
+            messages = [{"role": "system", "content": system}, *messages]
+            
+        response = await self.client.chat.completions.create(
+            model=self.model,
+            temperature=settings.TEMPERATURE,
+            max_tokens=settings.MAX_TOKENS,
+            messages=[{"role": m["role"], "content": m["content"]} for m in messages],
+            stream=True
+        )
+        async for chunk in response:
+            if chunk.choices:
+                yield chunk.choices[0].delta.content
+
 class ChatManager:
     def __init__(self, client: Literal["anthropic", "openai"] = "anthropic"):
         self.history: list[Message] = []
@@ -57,3 +100,15 @@ async def send_message(self, content: str, system: str | None = None) -> str:
         response = await self.client.chat(self.history, system)
         self.add_message("assistant", response)
         return response
+
+    async def stream_message(self, content: str, system: str | None = None) -> AsyncGenerator[str, None]:
+        """Stream a message response."""
+        self.add_message("user", content)
+        async for chunk in self.client.chat_stream(self.history, system):
+            yield chunk
+        # Add the complete message to history after streaming
+        if self.history[-1]["role"] == "user":
+            last_chunk = None
+            async for chunk in self.client.chat_stream(self.history, system):
+                last_chunk = chunk
+            self.add_message("assistant", last_chunk)  # Add the last chunk as the complete response
diff --git a/frontend/src/routes/learn/_layout.tsx b/frontend/src/routes/learn/_layout.tsx
@@ -0,0 +1,35 @@
+import { Flex, Spinner } from "@chakra-ui/react"
+import { Outlet, createFileRoute, redirect } from "@tanstack/react-router"
+
+import Sidebar from "../../components/Common/Sidebar"
+import UserMenu from "../../components/Common/UserMenu"
+import useAuth, { isLoggedIn } from "../../hooks/useAuth"
+
+export const Route = createFileRoute("/learn/_layout")({
+  component: Layout,
+  beforeLoad: async () => {
+    if (!isLoggedIn()) {
+      throw redirect({
+        to: "/login",
+      })
+    }
+  },
+})
+
+function Layout() {
+  const { isLoading } = useAuth()
+
+  return (
+    <Flex maxW="large" h="auto" position="relative">
+      <Sidebar />
+      {isLoading ? (
+        <Flex justify="center" align="center" height="100vh" width="full">
+          <Spinner size="xl" color="ui.main" />
+        </Flex>
+      ) : (
+        <Outlet />
+      )}
+      <UserMenu />
+    </Flex>
+  )
+}
diff --git a/frontend/src/routes/learn/chat/index.tsx b/frontend/src/routes/learn/chat/index.tsx
@@ -40,7 +40,7 @@ const ChatMessage = ({ message }: { message: ChatMessage }) => {
   )
 }
 
-export const Route = createFileRoute('/learn/chat')({
+export const Route = createFileRoute('/learn/_layout/chat')({
   component: ChatRoute
 })
 
@@ -74,30 +74,77 @@ function ChatRoute() {
     setCurrentMessage("")
     setIsLoading(true)
 
-    // Simulate streaming response
-    const response = "This is a simulated streaming response..."
-    let streamedContent = ""
-    
-    const newMessage: ChatMessage = {
+    const assistantMessage: ChatMessage = {
       id: window.crypto.randomUUID(),
       content: "",
       isUser: false,
     }
-    setMessages(prev => [...prev, newMessage])
+    setMessages(prev => [...prev, assistantMessage])
 
-    for (const char of response) {
-      await new Promise(resolve => setTimeout(resolve, 50))
-      streamedContent += char
+    try {
+      const response = await fetch('/api/v1/learn/chat/stream', {
+        method: 'POST',
+        headers: {
+          'Content-Type': 'application/json',
+          'Accept': 'text/event-stream',
+          'Authorization': `Bearer ${localStorage.getItem('access_token')}`,
+        },
+        body: JSON.stringify({
+          message: currentMessage,
+          model: "anthropic"
+        })
+      })
+
+      if (!response.ok) {
+        const errorData = await response.json()
+        throw new Error(`HTTP error! status: ${response.status}, message: ${errorData.detail || 'Unknown error'}`)
+      }
+
+      const reader = response.body?.getReader()
+      if (!reader) throw new Error('No reader available')
+
+      let streamedContent = ""
+      const decoder = new TextDecoder()
+      
+      while (true) {
+        const {done, value} = await reader.read()
+        if (done) break
+        
+        const chunk = decoder.decode(value)
+        const lines = chunk.split('\n')
+        
+        for (const line of lines) {
+          if (line.startsWith('data: ')) {
+            try {
+              const data = JSON.parse(line.slice(6))
+              if (data.type === 'content' && data.content) {
+                streamedContent += data.content
+                setMessages(prev => 
+                  prev.map(msg => 
+                    msg.id === assistantMessage.id 
+                      ? { ...msg, content: streamedContent }
+                      : msg
+                  )
+                )
+              }
+            } catch (e) {
+              console.error('Error parsing SSE data:', e)
+            }
+          }
+        }
+      }
+    } catch (error) {
+      console.error('Error:', error)
       setMessages(prev => 
         prev.map(msg => 
-          msg.id === newMessage.id 
-            ? { ...msg, content: streamedContent }
+          msg.id === assistantMessage.id 
+            ? { ...msg, content: "Sorry, there was an error processing your request." }
             : msg
         )
       )
+    } finally {
+      setIsLoading(false)
     }
-
-    setIsLoading(false)
   }
 
   return (
