mirror of
https://github.com/DrHo1y/ezrknn-llm.git
synced 2026-03-24 09:36:44 +07:00
82 lines
3.3 KiB
Python
82 lines
3.3 KiB
Python
import sys
|
||
import requests
|
||
import json
|
||
|
||
# 设置 Server 服务器的地址
|
||
server_url = 'http://172.16.10.102:8080/rkllm_chat'
|
||
# 设置是否开启流式对话
|
||
is_streaming = True
|
||
|
||
# 创建一个会话对象
|
||
session = requests.Session()
|
||
session.keep_alive = False # 关闭连接池，保持长连接
|
||
adapter = requests.adapters.HTTPAdapter(max_retries=5)
|
||
session.mount('https://', adapter)
|
||
session.mount('http://', adapter)
|
||
|
||
if __name__ == '__main__':
|
||
print("============================")
|
||
print("在终端中输入您的问题，即可与 RKLLM 模型进行对话....")
|
||
print("============================")
|
||
# 进入循环，持续获取用户输入，并与RKLLM模型进行对话
|
||
while True:
|
||
try:
|
||
user_message = input("请输入您的问题：")
|
||
if user_message == "exit":
|
||
print("============================")
|
||
print("程序正在退出......")
|
||
print("============================")
|
||
break
|
||
else:
|
||
# 设置请求头，此处的请求头实际并无作用，仅为模拟OpenAI接口设计
|
||
headers = {
|
||
'Content-Type': 'application/json',
|
||
'Authorization': 'not_required'
|
||
}
|
||
|
||
# 准备要发送的数据
|
||
# model: 为用户在设置RKLLM-Server时定义的模型，此处并无作用
|
||
# messages: 用户输入的问题，RKLLM-Server将会把它作为输入，并返回模型的回复；支持在 messags 加入多个问题
|
||
# stream: 是否开启流式对话，与OpenAI接口相同
|
||
data = {
|
||
"model": 'your_model_deploy_with_RKLLM_Server',
|
||
"messages": [{"role": "user", "content": user_message}],
|
||
"stream": is_streaming
|
||
}
|
||
|
||
# 发送 POST 请求
|
||
responses = session.post(server_url, json=data, headers=headers, stream=is_streaming, verify=False)
|
||
|
||
if not is_streaming:
|
||
# 解析响应
|
||
if responses.status_code == 200:
|
||
print("Q:", data["messages"][-1]["content"])
|
||
print("A:", json.loads(responses.text)["choices"][-1]["message"]["content"])
|
||
else:
|
||
print("Error:", responses.text)
|
||
else:
|
||
if responses.status_code == 200:
|
||
print("Q:", data["messages"][-1]["content"])
|
||
print("A:", end="")
|
||
for line in responses.iter_lines():
|
||
if line:
|
||
line = json.loads(line.decode('utf-8'))
|
||
if line["choices"][-1]["finish_reason"] != "stop":
|
||
print(line["choices"][-1]["delta"]["content"], end="")
|
||
sys.stdout.flush()
|
||
else:
|
||
print('Error:', responses.text)
|
||
|
||
|
||
|
||
|
||
except KeyboardInterrupt:
|
||
# 捕获 Ctrl-C 信号，关闭会话
|
||
session.close()
|
||
|
||
print("\n")
|
||
print("============================")
|
||
print("程序正在退出......")
|
||
print("============================")
|
||
break
|