Browse Source

增加测试scokt请求

div_测试
majiahui@haimaqingfan.com 1 month ago
parent
commit
a83c265f4e
  1. 85
      ceshi_scokt.py

85
ceshi_scokt.py

@ -0,0 +1,85 @@
from openai import OpenAI
openai_api_key = "token-abc123"
openai_api_base = "http://127.0.0.1:12011/v1"
client = OpenAI(
api_key=openai_api_key,
base_url=openai_api_base,
)
models = client.models.list()
model = models.data[0].id
def model_generate_stream(prompt):
messages = [
{"role": "user", "content": prompt}
]
stream = client.chat.completions.create(model=model,
messages=messages,
stream=True)
printed_reasoning_content = False
printed_content = False
for chunk in stream:
reasoning_content = None
content = None
# Check the content is reasoning_content or content
if hasattr(chunk.choices[0].delta, "reasoning_content"):
reasoning_content = chunk.choices[0].delta.reasoning_content
elif hasattr(chunk.choices[0].delta, "content"):
content = chunk.choices[0].delta.content
if reasoning_content is not None:
if not printed_reasoning_content:
printed_reasoning_content = True
print("reasoning_content:", end="", flush=True)
print(reasoning_content, end="", flush=True)
elif content is not None:
if not printed_content:
printed_content = True
print("\ncontent:", end="", flush=True)
# Extract and print the content
# print(content, end="", flush=True)
print(content)
yield content
# if __name__ == '__main__':
# for i in model_generate_stream("你好"):
# print(i)
import asyncio
import websockets
import json
async def handle_websocket(websocket):
print("客户端已连接")
try:
while True:
message = await websocket.recv()
print("收到消息:", message)
data = json.loads(message)
texts = data.get("texts")
title = data.get("title")
top = data.get("top")
response = model_generate_stream(texts)
# response = message + "111"
for char in response:
await websocket.send(char)
# await asyncio.sleep(0.3)
await websocket.send("[DONE]")
except websockets.exceptions.ConnectionClosed:
print("客户端断开连接")
async def main():
async with websockets.serve(handle_websocket, "0.0.0.0", 5500):
print("WebSocket 服务器已启动,监听端口 5500")
await asyncio.Future() # 永久运行
if __name__ == "__main__":
asyncio.run(main()) # 正确启动事件循环
Loading…
Cancel
Save