97 lines
2.9 KiB
Python
97 lines
2.9 KiB
Python
import requests
|
|
import json
|
|
import sseclient
|
|
|
|
def test_non_stream_chat():
|
|
"""测试非流式调用 /api/chat 接口"""
|
|
url = "http://localhost:9621/api/chat"
|
|
|
|
# 构造请求数据
|
|
data = {
|
|
"model": "lightrag:latest",
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": "孙悟空"
|
|
}
|
|
],
|
|
"stream": False
|
|
}
|
|
|
|
# 发送请求
|
|
response = requests.post(url, json=data)
|
|
|
|
# 打印响应
|
|
print("\n=== 非流式调用响应 ===")
|
|
response_json = response.json()
|
|
|
|
# 打印消息内容
|
|
print("=== 响应内容 ===")
|
|
print(json.dumps({
|
|
"model": response_json["model"],
|
|
"message": response_json["message"]
|
|
}, ensure_ascii=False, indent=2))
|
|
|
|
# 打印性能统计
|
|
print("\n=== 性能统计 ===")
|
|
stats = {
|
|
"total_duration": response_json["total_duration"],
|
|
"load_duration": response_json["load_duration"],
|
|
"prompt_eval_count": response_json["prompt_eval_count"],
|
|
"prompt_eval_duration": response_json["prompt_eval_duration"],
|
|
"eval_count": response_json["eval_count"],
|
|
"eval_duration": response_json["eval_duration"]
|
|
}
|
|
print(json.dumps(stats, ensure_ascii=False, indent=2))
|
|
|
|
def test_stream_chat():
|
|
"""测试流式调用 /api/chat 接口"""
|
|
url = "http://localhost:9621/api/chat"
|
|
|
|
# 构造请求数据
|
|
data = {
|
|
"model": "lightrag:latest",
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": "孙悟空有什么法力，性格特征是什么"
|
|
}
|
|
],
|
|
"stream": True
|
|
}
|
|
|
|
# 发送请求并获取 SSE 流
|
|
response = requests.post(url, json=data, stream=True)
|
|
client = sseclient.SSEClient(response)
|
|
|
|
print("\n=== 流式调用响应 ===")
|
|
output_buffer = []
|
|
try:
|
|
for event in client.events():
|
|
try:
|
|
data = json.loads(event.data)
|
|
if data.get("done", False): # 如果是完成标记
|
|
if "total_duration" in data: # 最终的性能统计消息
|
|
print("\n=== 性能统计 ===")
|
|
print(json.dumps(data, ensure_ascii=False, indent=2))
|
|
break
|
|
else: # 正常的内容消息
|
|
message = data.get("message", {})
|
|
content = message.get("content", "")
|
|
if content: # 只收集非空内容
|
|
output_buffer.append(content)
|
|
except json.JSONDecodeError:
|
|
print("Error decoding JSON from SSE event")
|
|
finally:
|
|
response.close() # 确保关闭响应连接
|
|
|
|
# 一次性打印所有收集到的内容
|
|
print("".join(output_buffer))
|
|
|
|
if __name__ == "__main__":
|
|
# 先测试非流式调用
|
|
test_non_stream_chat()
|
|
|
|
# 再测试流式调用
|
|
test_stream_chat()
|