database support

controllers.py

@@ -1,12 +1,13 @@
 import uuid
+import json
 from typing import Dict, List, Tuple
 from starlette.responses import JSONResponse
 from starlette.requests import Request
 from sse_starlette.sse import EventSourceResponse
 from chatgraph import get_messages, get_llm
+from models.Chat import Chat
 
 
-CHATS: Dict[str, List[dict]] = {}          # chat_id -> messages
 PENDING: Dict[str, Tuple[str, str]] = {}   # message_id -> (chat_id, provider)
 
 MODELS = {
@@ -32,16 +33,22 @@ async def create_chat(request: Request):
     provider = body.get("model","")
     if provider not in MODELS:
             return JSONResponse({"error": "Unknown model"}, status_code=400)
-    chat_id = str(uuid.uuid4())[:8]
-    CHATS[chat_id] = []
+    chat = Chat()
+    chat_id = str(uuid.uuid4())
+    chat.id = chat_id
+    chat.title = "New Chat"
+    chat.messages = json.dumps([])
+    chat.save()
     return JSONResponse({"id": chat_id, "model": provider})
 
 async def history(request : Request):
     """GET /chats/{chat_id}  -> previous messages"""
     chat_id = request.path_params["chat_id"]
-    if chat_id not in CHATS:
+    chat = Chat.find(chat_id)
+    if not chat:
         return JSONResponse({"error": "Not found"}, status_code=404)
-    return JSONResponse({"messages": CHATS[chat_id]})
+    messages = json.loads(chat.messages) if chat.messages else []
+    return JSONResponse({"messages": messages})
 
 async def post_message(request: Request):
     """POST /chats/{chat_id}/messages
@@ -49,7 +56,8 @@ async def post_message(request: Request):
        Returns: {"message_id": "<chat_id>"}
     """
     chat_id = request.path_params["chat_id"]
-    if chat_id not in CHATS:
+    chat = Chat.find(chat_id)
+    if not chat:
         return JSONResponse({"error": "Chat not found"}, status_code=404)
 
     body = await request.json()
@@ -58,9 +66,14 @@ async def post_message(request: Request):
     if provider not in MODELS:
         return JSONResponse({"error": "Unknown model"}, status_code=400)
 
+    # Load existing messages and add the new user message
+    messages = json.loads(chat.messages) if chat.messages else []
+    messages.append({"role": "human", "content": user_text})
+    chat.messages = json.dumps(messages)
+    chat.save()
+
     message_id = str(uuid.uuid4())
     PENDING[message_id] = (chat_id, provider)
-    CHATS[chat_id].append({"role": "human", "content": user_text})
 
     return JSONResponse({
         "status": "queued",
@@ -72,13 +85,18 @@ async def chat_stream(request):
     chat_id = request.path_params["chat_id"]
     message_id = request.query_params.get("message_id")
 
-    if chat_id not in CHATS or message_id not in PENDING:
+    if message_id not in PENDING:
         return JSONResponse({"error": "Not found"}, status_code=404)
 
     chat_id_from_map, provider = PENDING.pop(message_id)
     assert chat_id == chat_id_from_map
 
-    msgs = get_messages(CHATS, chat_id)
+    chat = Chat.find(chat_id)
+    if not chat:
+        return JSONResponse({"error": "Chat not found"}, status_code=404)
+
+    messages = json.loads(chat.messages) if chat.messages else []
+    msgs = get_messages( messages , chat_id)
     llm = get_llm(provider)
 
     async def event_generator():
@@ -88,7 +106,10 @@ async def chat_stream(request):
             buffer += token
             yield {"data": token}
         # Finished: store assistant reply
-        CHATS[chat_id].append({"role": "assistant", "content": buffer})
+        messages.append({"role": "assistant", "content": buffer})
+        chat.messages = json.dumps(messages)
+        chat.save()
         yield {"event": "done", "data": ""}
 
     return EventSourceResponse(event_generator())
+