fix websocket

app.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, Form, Request,Response
+from fastapi import FastAPI, Form, Request,Response, WebSocket, WebSocketDisconnect
 from fastapi.responses import HTMLResponse
 from jinja2 import Template
 import markdown
@@ -6,9 +6,11 @@ import time
 from datetime import datetime, timedelta
 from apscheduler.schedulers.background import BackgroundScheduler
 from agents import DeepResearchAgent, get_llms
+import hashlib
 import threading
 import logging
 from queue import Queue
+import json
 import uuid
 
 
@@ -346,10 +348,27 @@ html_template = """
     <script>
        {{ script}}
     </script>
+    <script>
+    const socket = new WebSocket("ws://localhost:7860/ws");
+
+    socket.addEventListener('open', function (event) {
+        const userId = document.getElementById("user_id").value;
+        socket.send(JSON.stringify({ action: "connect", user_id: userId }));
+    });
+
+    window.addEventListener("beforeunload", function (event) {
+        const userId = document.getElementById("user_id").value;
+        socket.send(JSON.stringify({ action: "disconnect", user_id: userId }));
+    });
+</script>
 </body>
 </html>
 """
 
+def generate_user_id(ip: str) -> str:
+    # 使用哈希函数生成用户 ID
+    return hashlib.md5(ip.encode()).hexdigest()
+
 # 重置每日计数器
 def reset_counter():
     global reply_count
@@ -393,23 +412,50 @@ script_template = """
 
 queue = Queue()
 
-@app.get("/", response_class=HTMLResponse)
-def form_get():
+@app.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    await websocket.accept()
+    print("WebSocket connection established.")
+    try:
+        while True:
+            data = await websocket.receive_text()
+            message = json.loads(data)
+            user_id = message.get("user_id")
+            action = message.get("action")
+            print(action)
+            if action == "disconnect":
+                for item in list(queue.queue):
+                    if item[0] == user_id:
+                        queue.queue.remove(item)
+                        break
+                print(f"User {user_id} disconnected.")
+
+    except WebSocketDisconnect:
+        print("WebSocket connection closed.")
 
-    script = script_template.format(user_id=str(uuid.uuid4()), state="generate")
+@app.get("/", response_class=HTMLResponse)
+def form_get(request: Request):
+    client_ip = request.client.host
+    user_id = generate_user_id(client_ip)
+    script = script_template.format(user_id=user_id, state="generate")
+    print(client_ip,user_id)
     return Template(html_template).render(idea= "This is a example of the idea geneartion", error=None, reply_count=reply_count,button_text="Generate",loading_text="Generating content, Usually takes 3-4 minutes, please wait...",script=script)
 
 @app.post("/", response_class=HTMLResponse)
-def form_post(topic: str = Form(...), user_id: str = Form(...), state: str = Form(...)):
+def form_post(request: Request,topic: str = Form(...), user_id: str = Form(...), state: str = Form(...)):
     global reply_count
     start_time = time.time()
-
+    client_ip = request.client.host
     if user_id == "":
-        user_id = str(uuid.uuid4())
+        user_id = generate_user_id(client_ip)
     if state == "":
         state = "generate"
 
     script = script_template.format(user_id=user_id, state=state)
+    if user_id in queue.queue:
+        error_message = "Your request is being processed. Please wait for the result, or close the previous page and try again."
+        script = script_template.format(user_id=user_id, state=state)
+        return Template(html_template).render(idea="", error=error_message, reply_count=reply_count, button_text="Generate",loading_text="Generating content, Usually takes 3-4 minutes, please wait...",script=script)
 
     print(f"current0 user_id={user_id}, state={state}")
     loading_text = "Generating content, Usually takes 3-4 minutes, please wait..."
@@ -437,7 +483,7 @@ def form_post(topic: str = Form(...), user_id: str = Form(...), state: str = For
 
             return Template(html_template).render(idea="", error=error_message, reply_count=reply_count, button_text=new_button_text,loading_text=f"Generating content, Usually takes {(queue_len+1)*3}-{(queue_len+1)*4} minutes, please wait...",script=script)
 
-    queue.put([user_id,topic])
+    queue.put(user_id)
     new_state = "generate"
     new_button_text = "Generate"
     queue_len = queue.qsize()
@@ -446,13 +492,24 @@ def form_post(topic: str = Form(...), user_id: str = Form(...), state: str = For
     print(f"current2 user_id={user_id}, state={new_state}")
     # 判断当前是否轮到该用户,如果没轮到则一直等待到轮到为止
     print(queue.queue[0], [user_id,topic])
-    while queue.queue[0] != [user_id,topic]:
+    while queue.queue[0] != user_id:
+        # 检查用户是否还在队列中
+        if not any(user_id == item[0] for item in queue.queue):
+            return Template(html_template).render(
+                idea="", 
+                error="Request was cancelled.", 
+                reply_count=reply_count, 
+                button_text="Generate",
+                loading_text=loading_text,
+                script=script
+            )
         time.sleep(10)
         continue
 
     try:
         with lock:
             logging.info(f"Processing request for topic: {topic}")
+            time.sleep(1000)
             start_time = time.time()
             error_message = None
             idea = ""