wehub-resource-sync
116 lines
4.7 KiB
Python
116 lines
4.7 KiB
Python
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import json
|
|
|
|
import requests
|
|
|
|
|
|
def build_messages(query, history=None):
|
|
"""
|
|
根据传入的 query 和 history 构造符合 OpenAI 格式的消息列表。
|
|
如果 history 为 list 且每项为 dict，则直接使用；如果为 list 且每项为字符串，
|
|
则依次按用户（user）与助手（assistant）交替添加；否则直接只添加当前用户消息。
|
|
"""
|
|
messages = []
|
|
if history:
|
|
if isinstance(history, list):
|
|
if all(isinstance(item, dict) for item in history):
|
|
messages.extend(history)
|
|
else:
|
|
# 假设 history 按顺序依次为用户、助手、用户、助手……
|
|
for idx, item in enumerate(history):
|
|
role = "user" if idx % 2 == 0 else "assistant"
|
|
messages.append({"role": role, "content": str(item)})
|
|
else:
|
|
messages.append({"role": "user", "content": str(history)})
|
|
# 当前请求作为最新的用户消息
|
|
messages.append({"role": "user", "content": query})
|
|
return messages
|
|
|
|
|
|
def send_request(query, history=None, stream=True):
|
|
# 构造 OpenAI 格式的请求体
|
|
payload = {
|
|
"messages": build_messages(query, history),
|
|
# 以下生成参数可根据需要调整
|
|
# "top_k": 0,
|
|
# "top_p": 0.7,
|
|
# "temperature": 0.8,
|
|
# "repetition_penalty": 1.3,
|
|
"max_length": 1024,
|
|
"src_length": 1024,
|
|
"min_length": 1,
|
|
"stream": stream,
|
|
}
|
|
res = requests.post("http://localhost:8011/v1/chat/completions", json=payload, stream=True)
|
|
result_text = ""
|
|
printed_reasoning_content = False
|
|
printed_content = False
|
|
for line in res.iter_lines():
|
|
# https://github.com/vllm-project/vllm/blob/433c4a49230a470f13657f06e7612cde86e4fb40/examples/online_serving/openai_chat_completion_with_reasoning_streaming.py#L67-L69
|
|
if not line:
|
|
continue
|
|
|
|
decoded_line = line.decode("utf-8").strip()
|
|
# OpenAI 流返回每行以 "data:" 开头
|
|
if decoded_line.startswith("data:"):
|
|
data = decoded_line[5:].strip() # Remove "data:" prefix
|
|
if data == "[DONE]": # End of stream
|
|
print("\nclient: Stream completed.\n")
|
|
break
|
|
try:
|
|
# Parse the JSON data
|
|
chunk = json.loads(data)
|
|
reasoning_content = chunk["choices"][0]["delta"].get("reasoning_content", "")
|
|
content = chunk["choices"][0]["delta"].get("content", "")
|
|
|
|
if reasoning_content:
|
|
if not printed_reasoning_content:
|
|
printed_reasoning_content = True
|
|
print("reasoning_content:", end="", flush=True)
|
|
print(reasoning_content, end="", flush=True)
|
|
elif content:
|
|
if not printed_content:
|
|
printed_content = True
|
|
print("\ncontent:", end="", flush=True)
|
|
# Extract and print the content
|
|
print(content, end="", flush=True)
|
|
result_text += content
|
|
except Exception as e:
|
|
print("解析响应出错:", e)
|
|
continue
|
|
else:
|
|
try:
|
|
data = json.loads(decoded_line)
|
|
content = data["choices"][0]["message"].get("content", "")
|
|
print(content, end="", flush=True)
|
|
result_text += content
|
|
except Exception as e:
|
|
print("解析响应出错:", e)
|
|
continue
|
|
|
|
print()
|
|
return result_text
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# 示例调用：仅发送当前用户消息
|
|
send_request("你好啊")
|
|
send_request("你好啊", stream=False)
|
|
# 示例调用：使用 history 为字符串列表（交替为用户与助手的对话）
|
|
send_request("再加一等于多少", ["一加一等于多少", "一加一等于二"])
|
|
# 示例调用：history 为字典格式，明确指定对话角色
|
|
send_request("再加一等于多少", [{"role": "user", "content": "一加一等于多少"}, {"role": "assistant", "content": "一加一等于二"}])
|