easy-rag/tests/test_api.py

87 lines
2.8 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

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}")