增加重置对话功能

main.py

@@ -10,13 +10,10 @@ import requests
 import uvicorn
 
 api_key = os.environ.get('api_key')
-initial_prompt = "You are a helpful assistant."
 API_URL = "https://api.openai.com/v1/chat/completions"
 
 app = FastAPI()
 
-openai.api_key = api_key
-
 
 @app.get("/")
 def read_root():
@@ -30,10 +27,14 @@ connection_history: Dict[str, List[Dict[str, str]]] = {}
 def get_sys_prompt():
     return [{
         "role": "system",
-        "content": f"You are a helpful assistant. Current time is {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}. 默认用户在成都"
+        "content": f"You are a helpful assistant. Current time is {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}."
     }]
 
 
+def test_reset(msg):
+    return msg == "重置对话" or msg == "开启新对话"
+
+
 def get_ai_response(messages, stream=True):
     '''获取ChatGPT答复，使用流式返回'''
     headers = {
@@ -69,13 +70,6 @@ class ResponseMessage(BaseModel):
 
 @app.websocket("/api/ws")
 async def websocket_endpoint(websocket: WebSocket):
-    # auth_header = websocket.headers.get("Authorization")
-    # ws_key = "Bearer YOUR_API_KEY"  # 将 YOUR_API_KEY 替换为实际的 API 密钥
-
-    # if auth_header != ws_key:
-    #     await websocket.close(code=1008)  # 关闭连接，发送策略原因错误
-    #     return
-
     await websocket.accept()
     connection_id = str(id(websocket))
     connection_history[connection_id] = []
@@ -86,6 +80,12 @@ async def websocket_endpoint(websocket: WebSocket):
             message = Message(**json.loads(data))
             print(message)
 
+            # 判断是否重置对话
+            if test_reset(message.msg):
+                connection_history[connection_id] = []
+                print("OK，对话已重置")
+                continue
+
             user_message = {"role": "user", "content": message.msg}
             connection_history[connection_id].append(user_message)
 
@@ -125,6 +125,8 @@ async def websocket_endpoint(websocket: WebSocket):
     except WebSocketDisconnect as e:
         # 在这里处理断开连接的情况，例如记录日志、清理资源等
         print(f"WebSocket disconnected with code: {e.code}")
+    except:
+        print(f"WebSocket disconnected with unknown reason")
 
 
 class Item(BaseModel):
@@ -141,6 +143,10 @@ def chat(item: Item):
     return res
 
 
+openai.api_key = api_key
+initial_prompt = "You are a helpful assistant."
+
+
 def get_response(system_prompt, history):
 
     history = [construct_system(system_prompt), *history]