214 lines
6.7 KiB
Python
214 lines
6.7 KiB
Python
import requests
|
||
import json
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
from server.knowledge_base.utils import get_file_path
|
||
|
||
root_path = Path(__file__).parent.parent.parent
|
||
sys.path.append(str(root_path))
|
||
from server.utils import api_address
|
||
from configs import VECTOR_SEARCH_TOP_K
|
||
|
||
from pprint import pprint
|
||
|
||
|
||
api_base_url = api_address()
|
||
|
||
|
||
kb = "kb_for_api_test"
|
||
test_files = {
|
||
"test.txt": get_file_path("samples", "test.txt"),
|
||
}
|
||
|
||
print("\n\n直接url访问\n")
|
||
|
||
|
||
def test_delete_kb_before(api="/knowledge_base/delete_knowledge_base"):
|
||
# if not Path(get_kb_path(kb)).exists():
|
||
# return
|
||
|
||
url = api_base_url + api
|
||
print("\n测试知识库存在，需要删除")
|
||
r = requests.post(url, json=kb)
|
||
data = r.json()
|
||
pprint(data)
|
||
|
||
# check kb not exists anymore
|
||
url = api_base_url + "/knowledge_base/list_knowledge_bases"
|
||
print("\n获取知识库列表：")
|
||
r = requests.get(url)
|
||
data = r.json()
|
||
pprint(data)
|
||
assert data["code"] == 200
|
||
assert isinstance(data["data"], list) and len(data["data"]) > 0
|
||
assert kb not in data["data"]
|
||
|
||
|
||
def test_create_kb(api="/knowledge_base/create_knowledge_base"):
|
||
url = api_base_url + api
|
||
|
||
print(f"\n尝试用空名称创建知识库：")
|
||
r = requests.post(url, json={"knowledge_base_name": " "})
|
||
data = r.json()
|
||
pprint(data)
|
||
# assert data["code"] == 404
|
||
# assert data["msg"] == "知识库名称不能为空，请重新填写知识库名称"
|
||
|
||
print(f"\n创建新知识库： {kb}")
|
||
r = requests.post(url, json={"knowledge_base_name": kb})
|
||
data = r.json()
|
||
pprint(data)
|
||
# assert data["code"] == 200
|
||
# assert data["msg"] == f"已新增知识库 {kb}"
|
||
|
||
print(f"\n尝试创建同名知识库： {kb}")
|
||
r = requests.post(url, json={"knowledge_base_name": kb})
|
||
data = r.json()
|
||
pprint(data)
|
||
# assert data["code"] == 404
|
||
# assert data["msg"] == f"已存在同名知识库 {kb}"
|
||
|
||
|
||
def test_list_kbs(api="/knowledge_base/list_knowledge_bases"):
|
||
url = api_base_url + api
|
||
print("\n获取知识库列表：")
|
||
r = requests.get(url)
|
||
data = r.json()
|
||
pprint(data)
|
||
# assert data["code"] == 200
|
||
# assert isinstance(data["data"], list) and len(data["data"]) > 0
|
||
# assert kb in data["data"]
|
||
|
||
|
||
def test_upload_docs(api="/knowledge_base/upload_docs"):
|
||
url = api_base_url + api
|
||
files = [("files", (name, open(path, "rb"))) for name, path in test_files.items()]
|
||
|
||
print(f"\n上传知识文件")
|
||
data = {"knowledge_base_name": kb, "override": True}
|
||
r = requests.post(url, data=data, files=files)
|
||
data = r.json()
|
||
pprint(data)
|
||
assert data["code"] == 200
|
||
assert len(data["data"]["failed_files"]) == 0
|
||
|
||
print(f"\n尝试重新上传知识文件， 不覆盖")
|
||
data = {"knowledge_base_name": kb, "override": False}
|
||
files = [("files", (name, open(path, "rb"))) for name, path in test_files.items()]
|
||
r = requests.post(url, data=data, files=files)
|
||
data = r.json()
|
||
pprint(data)
|
||
assert data["code"] == 200
|
||
assert len(data["data"]["failed_files"]) == len(test_files)
|
||
|
||
print(f"\n尝试重新上传知识文件， 覆盖，自定义docs")
|
||
docs = {"FAQ.MD": [{"page_content": "custom docs", "metadata": {}}]}
|
||
data = {"knowledge_base_name": kb, "override": True, "docs": json.dumps(docs)}
|
||
files = [("files", (name, open(path, "rb"))) for name, path in test_files.items()]
|
||
r = requests.post(url, data=data, files=files)
|
||
data = r.json()
|
||
pprint(data)
|
||
assert data["code"] == 200
|
||
assert len(data["data"]["failed_files"]) == 0
|
||
|
||
|
||
def test_list_files(api="/knowledge_base/list_files"):
|
||
url = api_base_url + api
|
||
print("\n获取知识库中文件列表：")
|
||
r = requests.get(url, params={"knowledge_base_name": kb})
|
||
data = r.json()
|
||
pprint(data)
|
||
assert data["code"] == 200
|
||
assert isinstance(data["data"], list)
|
||
for name in test_files:
|
||
assert name in data["data"]
|
||
|
||
|
||
def test_search_docs(api="/knowledge_base/search_docs"):
|
||
url = api_base_url + api
|
||
query = "林枫是谁"
|
||
print("\n检索知识库：")
|
||
print(query)
|
||
r = requests.post(url, json={"knowledge_base_name": kb, "query": query})
|
||
data = r.json()
|
||
pprint(data)
|
||
assert isinstance(data, list) and len(data) == VECTOR_SEARCH_TOP_K
|
||
|
||
|
||
def test_update_info(api="/knowledge_base/update_info"):
|
||
url = api_base_url + api
|
||
print("\n更新知识库介绍")
|
||
r = requests.post(url, json={"knowledge_base_name": kb, "kb_info": "你好"})
|
||
data = r.json()
|
||
pprint(data)
|
||
assert data["code"] == 200
|
||
|
||
def test_update_docs(api="/knowledge_base/update_docs"):
|
||
url = api_base_url + api
|
||
|
||
print(f"\n更新知识文件")
|
||
r = requests.post(url, json={"knowledge_base_name": kb, "file_names": list(test_files)})
|
||
data = r.json()
|
||
pprint(data)
|
||
assert data["code"] == 200
|
||
assert len(data["data"]["failed_files"]) == 0
|
||
|
||
|
||
def test_delete_docs(api="/knowledge_base/delete_docs"):
|
||
url = api_base_url + api
|
||
|
||
print(f"\n删除知识文件")
|
||
r = requests.post(url, json={"knowledge_base_name": kb, "file_names": list(test_files)})
|
||
data = r.json()
|
||
pprint(data)
|
||
# assert data["code"] == 200
|
||
# assert len(data["data"]["failed_files"]) == 0
|
||
|
||
url = api_base_url + "/knowledge_base/search_docs"
|
||
query = "介绍一下langchain-chatchat项目"
|
||
print("\n尝试检索删除后的检索知识库：")
|
||
print(query)
|
||
r = requests.post(url, json={"knowledge_base_name": kb, "query": query})
|
||
data = r.json()
|
||
pprint(data)
|
||
# assert isinstance(data, list) and len(data) == 0
|
||
|
||
|
||
def test_recreate_vs(api="/knowledge_base/recreate_vector_store"):
|
||
url = api_base_url + api
|
||
print("\n重建知识库：")
|
||
r = requests.post(url, json={"knowledge_base_name": kb}, stream=True)
|
||
for chunk in r.iter_content(None):
|
||
data = json.loads(chunk)
|
||
assert isinstance(data, dict)
|
||
assert data["code"] == 200
|
||
print(data["msg"])
|
||
|
||
url = api_base_url + "/knowledge_base/search_docs"
|
||
query = "本项目支持哪些文件格式?"
|
||
print("\n尝试检索重建后的检索知识库：")
|
||
print(query)
|
||
r = requests.post(url, json={"knowledge_base_name": kb, "query": query})
|
||
data = r.json()
|
||
pprint(data)
|
||
assert isinstance(data, list) and len(data) == VECTOR_SEARCH_TOP_K
|
||
|
||
|
||
def test_delete_kb_after(api="/knowledge_base/delete_knowledge_base"):
|
||
url = api_base_url + api
|
||
print("\n删除知识库")
|
||
r = requests.post(url, json=kb)
|
||
data = r.json()
|
||
pprint(data)
|
||
|
||
# check kb not exists anymore
|
||
url = api_base_url + "/knowledge_base/list_knowledge_bases"
|
||
print("\n获取知识库列表：")
|
||
r = requests.get(url)
|
||
data = r.json()
|
||
pprint(data)
|
||
assert data["code"] == 200
|
||
assert isinstance(data["data"], list) and len(data["data"]) > 0
|
||
assert kb not in data["data"]
|