87 lines
2.8 KiB
Python
87 lines
2.8 KiB
Python
import requests
|
||
import json
|
||
from datetime import datetime
|
||
|
||
def test_upload_and_chat():
|
||
"""测试文档上传和聊天功能"""
|
||
base_url = "http://localhost:8000"
|
||
|
||
# 测试健康检查
|
||
print("1. 测试健康检查...")
|
||
response = requests.get(f"{base_url}/health")
|
||
print(f"状态码: {response.status_code}")
|
||
print(f"响应: {response.json()}")
|
||
print()
|
||
|
||
# 测试文档上传
|
||
print("2. 测试文档上传...")
|
||
test_content = "这是一个测试文档。它包含了关于人工智能的基本信息。人工智能是计算机科学的一个分支。"
|
||
|
||
# 创建临时文件
|
||
with open("test_doc.txt", "w", encoding="utf-8") as f:
|
||
f.write(test_content)
|
||
|
||
with open("test_doc.txt", "rb") as f:
|
||
files = {"file": ("test_doc.txt", f, "text/plain")}
|
||
response = requests.post(f"{base_url}/upload", files=files)
|
||
|
||
print(f"状态码: {response.status_code}")
|
||
if response.status_code == 200:
|
||
upload_result = response.json()
|
||
print(f"上传成功: {upload_result}")
|
||
doc_id = upload_result["data"]["document_id"]
|
||
else:
|
||
print(f"上传失败: {response.text}")
|
||
return
|
||
print()
|
||
|
||
# 测试文档列表
|
||
print("3. 测试文档列表...")
|
||
response = requests.get(f"{base_url}/documents")
|
||
print(f"状态码: {response.status_code}")
|
||
print(f"文档列表: {response.json()}")
|
||
print()
|
||
|
||
# 测试聊天
|
||
print("4. 测试聊天...")
|
||
chat_data = {"question": "什么是人工智能?", "top_k": 3, "temperature": 0.7}
|
||
start_time = datetime.now()
|
||
response = requests.post(
|
||
f"{base_url}/chat/stream",
|
||
json=chat_data,
|
||
headers={"Content-Type": "application/json"},
|
||
stream=True,
|
||
)
|
||
|
||
print(f"状态码: {response.status_code}")
|
||
if response.status_code == 200:
|
||
# 遍历响应体,逐行处理流式数据(适用于text/event-stream 或 chunked json)
|
||
last_line = None
|
||
for line in response.iter_lines(decode_unicode=True):
|
||
if line:
|
||
last_line = line
|
||
print(f"回答: {line}")
|
||
end_time = datetime.now()
|
||
processing_time = (end_time - start_time).total_seconds()
|
||
print(f"处理时间: {processing_time:.2f}秒")
|
||
print(f"来源数量: {len(json.loads(last_line.replace('data: ', ''))['sources'])}")
|
||
else:
|
||
print(f"聊天失败: {response.text}")
|
||
print()
|
||
|
||
# 清理测试文件
|
||
import os
|
||
|
||
if os.path.exists("test_doc.txt"):
|
||
os.remove("test_doc.txt")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
try:
|
||
test_upload_and_chat()
|
||
except requests.exceptions.ConnectionError:
|
||
print("错误: 无法连接到服务器")
|
||
print("请确保服务器正在运行: python main.py")
|
||
except Exception as e:
|
||
print(f"测试失败: {e}")
|