Spaces:
Running
Running
import requests | |
import json | |
import sys | |
from pathlib import 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 server.knowledge_base.utils import get_kb_path, get_file_path | |
from pprint import pprint | |
api_base_url = api_address() | |
kb = "kb_for_api_test" | |
test_files = { | |
"wiki/Home.MD": get_file_path("samples", "wiki/Home.md"), | |
"wiki/开发环境部署.MD": get_file_path("samples", "wiki/开发环境部署.md"), | |
"test_files/test.txt": get_file_path("samples", "test_files/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 = "介绍一下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) == 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": "samples", "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[6:]) | |
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"] | |