easy-rag/tests/test_api.py

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