少许更新

main.py

@@ -1,4 +1,5 @@
 import asyncio
+from datetime import datetime
 import json
 from typing import Dict, List, Union
 import os
@@ -22,35 +23,30 @@ def read_root():
     return {"Hello": "World!"}
 
 
-class Item(BaseModel):
-    _msgid: Union[str, None] = None
-    input: str
-    history: List[Dict] = []
-
-
-@app.post("/api/chat")
-def chat(item: Item):
-    print(item)
-    history = [construct_user(item.input)]
-    res = get_response(initial_prompt, history)
-    return res
-
-
 # 存储每个连接的对话历史
 connection_history: Dict[str, List[Dict[str, str]]] = {}
 
 
-def get_ai_response(messages):
+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')}. 默认用户在成都"
+    }]
+
+
+def get_ai_response(messages, stream=True):
     '''获取ChatGPT答复，使用流式返回'''
     headers = {
         "Content-Type": "application/json",
         "Authorization": f"Bearer {api_key}",
     }
 
+    sys_prompt = get_sys_prompt()
+
     payload = json.dumps({
         "model": "gpt-3.5-turbo",
-        "messages": messages,
-        "stream": True,
+        "messages": sys_prompt + messages,
+        "stream": stream,
     })
 
     response = requests.post(
@@ -73,10 +69,16 @@ 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] = [
-        {"role": "system", "content": "You are a helpful assistant."}]
+    connection_history[connection_id] = []
 
     try:
         while True:
@@ -125,6 +127,20 @@ async def websocket_endpoint(websocket: WebSocket):
         print(f"WebSocket disconnected with code: {e.code}")
 
 
+class Item(BaseModel):
+    _msgid: Union[str, None] = None
+    input: str
+    history: List[Dict] = []
+
+
+@app.post("/api/chat")
+def chat(item: Item):
+    print(item)
+    history = [construct_user(item.input)]
+    res = get_response(initial_prompt, history)
+    return res
+
+
 def get_response(system_prompt, history):
 
     history = [construct_system(system_prompt), *history]