
1 changed files with 85 additions and 0 deletions
@ -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…
Reference in new issue