增加ws接口

main.py

@@ -1,14 +1,20 @@
+import json
+from typing import Dict, List, Union
 import os
+from fastapi import FastAPI, WebSocket
 import openai
-from fastapi import FastAPI
 from pydantic import BaseModel
-from typing import Dict, List
+import requests
+import uvicorn
 
-openai.api_key = os.environ.get('api_key')
+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():
@@ -16,18 +22,81 @@ def read_root():
 
 
 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
 
 
+class Message(BaseModel):
+    msg: str
+
+
+class ResponseMessage(BaseModel):
+    msg: str
+    finished: bool
+
+
+@app.websocket("/api/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    await websocket.accept()
+    while True:
+        data = await websocket.receive_text()
+        message = Message(**json.loads(data))
+
+        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
+        )
+
+        print(message)
+
+        finished = False
+        for line in response.iter_lines():
+            if line:
+                decoded_line = line.decode("utf-8")
+                # print(decoded_line)
+                content = json.loads(decoded_line[6:])
+                finish_reason = content["choices"][0]["finish_reason"]
+
+                if "content" in content['choices'][0]["delta"]:
+                    response_msg = content["choices"][0]["delta"]["content"]
+                else:
+                    response_msg = ""
+
+                if finish_reason == "stop":
+                    finished = True
+
+                response_message = ResponseMessage(
+                    msg=response_msg, finished=finished)
+                await websocket.send_text(response_message.json())
+
+                if finished:
+                    break
+
+
 def get_response(system_prompt, history):
+
     history = [construct_system(system_prompt), *history]
 
     payload = {
@@ -55,3 +124,7 @@ def construct_assistant(text):
 
 def construct_text(role, text):
     return {"role": role, "content": text}
+
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)