78 lines
2.4 KiB
Python
78 lines
2.4 KiB
Python
import requests
|
|
import json
|
|
|
|
|
|
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}
|
|
|
|
response = requests.post(
|
|
f"{base_url}/chat", json=chat_data, headers={"Content-Type": "application/json"}
|
|
)
|
|
|
|
print(f"状态码: {response.status_code}")
|
|
if response.status_code == 200:
|
|
chat_result = response.json()
|
|
print(f"回答: {chat_result['answer']}")
|
|
print(f"处理时间: {chat_result['processing_time']:.2f}秒")
|
|
print(f"来源数量: {len(chat_result['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}")
|