ws链接支持上下文

main.py

@@ -36,6 +36,32 @@ def chat(item: Item):
     return res
 
 
+# 存储每个连接的对话历史
+connection_history: Dict[str, List[Dict[str, str]]] = {}
+
+
+def get_ai_response(messages):
+    '''获取ChatGPT答复，使用流式返回'''
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {api_key}",
+    }
+
+    payload = json.dumps({
+        "model": "gpt-3.5-turbo",
+        "messages": messages,
+        "stream": True,
+    })
+
+    response = requests.post(
+        API_URL,
+        headers=headers,
+        data=payload,
+        stream=True
+    )
+    return response
+
+
 class Message(BaseModel):
     msg: str
 
@@ -48,31 +74,22 @@ class ResponseMessage(BaseModel):
 @app.websocket("/api/ws")
 async def websocket_endpoint(websocket: WebSocket):
     await websocket.accept()
+    connection_id = str(id(websocket))
+    connection_history[connection_id] = [
+        {"role": "system", "content": "You are a helpful assistant."}]
+
     try:
         while True:
             data = await websocket.receive_text()
             message = Message(**json.loads(data))
+            print(message)
 
-            headers = {
-                "Content-Type": "application/json",
-                "Authorization": f"Bearer {api_key}",
-            }
-
-            payload = json.dumps({
-                "model": "gpt-3.5-turbo",
-                "messages": [{"role": "system", "content": "start"}, {"role": "user", "content": message.msg}],
-                "stream": True,
-            })
-
-            response = requests.post(
-                API_URL,
-                headers=headers,
-                data=payload,
-                stream=True
-            )
+            user_message = {"role": "user", "content": message.msg}
+            connection_history[connection_id].append(user_message)
 
-            print(message)
+            response = get_ai_response(connection_history[connection_id])
 
+            whole_message = ""
             finished = False
             for line in response.iter_lines():
                 if line:
@@ -83,15 +100,22 @@ async def websocket_endpoint(websocket: WebSocket):
 
                     if "content" in content['choices'][0]["delta"]:
                         response_msg = content["choices"][0]["delta"]["content"]
+                        whole_message += response_msg
                     else:
                         response_msg = ""
 
                     if finish_reason == "stop":
                         finished = True
+                        print(whole_message)
+                        assistant_message = {
+                            "role": "assistant", "content": whole_message}
+                        connection_history[connection_id].append(
+                            assistant_message)
 
                     response_message = ResponseMessage(
                         msg=response_msg, finished=finished)
                     await websocket.send_text(response_message.json())
+                    # sleep 1ms 给发送协程让出执行窗口
                     await asyncio.sleep(0.001)
 
                     if finished: