Langchain-Chatchat

1
import requests
2
import json
3
import sys
4
from pathlib import Path
5

6
root_path = Path(__file__).parent.parent.parent
7
sys.path.append(str(root_path))
8
from server.utils import api_address
9
from configs import VECTOR_SEARCH_TOP_K
10
from server.knowledge_base.utils import get_kb_path, get_file_path
11

12
from pprint import pprint
13

14

15
api_base_url = api_address()
16

17

18
kb = "kb_for_api_test"
19
test_files = {
20
    "wiki/Home.MD": get_file_path("samples", "wiki/Home.md"),
21
    "wiki/开发环境部署.MD": get_file_path("samples", "wiki/开发环境部署.md"),
22
    "test_files/test.txt": get_file_path("samples", "test_files/test.txt"),
23
}
24

25
print("\n\n直接url访问\n")
26

27

28
def test_delete_kb_before(api="/knowledge_base/delete_knowledge_base"):
29
    if not Path(get_kb_path(kb)).exists():
30
        return
31

32
    url = api_base_url + api
33
    print("\n测试知识库存在，需要删除")
34
    r = requests.post(url, json=kb)
35
    data = r.json()
36
    pprint(data)
37

38
    # check kb not exists anymore
39
    url = api_base_url + "/knowledge_base/list_knowledge_bases"
40
    print("\n获取知识库列表：")
41
    r = requests.get(url)
42
    data = r.json()
43
    pprint(data)
44
    assert data["code"] == 200
45
    assert isinstance(data["data"], list) and len(data["data"]) > 0
46
    assert kb not in data["data"]
47

48

49
def test_create_kb(api="/knowledge_base/create_knowledge_base"):
50
    url = api_base_url + api
51

52
    print(f"\n尝试用空名称创建知识库：")
53
    r = requests.post(url, json={"knowledge_base_name": " "})
54
    data = r.json()
55
    pprint(data)
56
    assert data["code"] == 404
57
    assert data["msg"] == "知识库名称不能为空，请重新填写知识库名称"
58

59
    print(f"\n创建新知识库： {kb}")
60
    r = requests.post(url, json={"knowledge_base_name": kb})
61
    data = r.json()
62
    pprint(data)
63
    assert data["code"] == 200
64
    assert data["msg"] == f"已新增知识库 {kb}"
65

66
    print(f"\n尝试创建同名知识库： {kb}")
67
    r = requests.post(url, json={"knowledge_base_name": kb})
68
    data = r.json()
69
    pprint(data)
70
    assert data["code"] == 404
71
    assert data["msg"] == f"已存在同名知识库 {kb}"
72

73

74
def test_list_kbs(api="/knowledge_base/list_knowledge_bases"):
75
    url = api_base_url + api
76
    print("\n获取知识库列表：")
77
    r = requests.get(url)
78
    data = r.json()
79
    pprint(data)
80
    assert data["code"] == 200
81
    assert isinstance(data["data"], list) and len(data["data"]) > 0
82
    assert kb in data["data"]
83

84

85
def test_upload_docs(api="/knowledge_base/upload_docs"):
86
    url = api_base_url + api
87
    files = [("files", (name, open(path, "rb"))) for name, path in test_files.items()]
88

89
    print(f"\n上传知识文件")
90
    data = {"knowledge_base_name": kb, "override": True}
91
    r = requests.post(url, data=data, files=files)
92
    data = r.json()
93
    pprint(data)
94
    assert data["code"] == 200
95
    assert len(data["data"]["failed_files"]) == 0
96

97
    print(f"\n尝试重新上传知识文件， 不覆盖")
98
    data = {"knowledge_base_name": kb, "override": False}
99
    files = [("files", (name, open(path, "rb"))) for name, path in test_files.items()]
100
    r = requests.post(url, data=data, files=files)
101
    data = r.json()
102
    pprint(data)
103
    assert data["code"] == 200
104
    assert len(data["data"]["failed_files"]) == len(test_files)
105

106
    print(f"\n尝试重新上传知识文件， 覆盖，自定义docs")
107
    docs = {"FAQ.MD": [{"page_content": "custom docs", "metadata": {}}]}
108
    data = {"knowledge_base_name": kb, "override": True, "docs": json.dumps(docs)}
109
    files = [("files", (name, open(path, "rb"))) for name, path in test_files.items()]
110
    r = requests.post(url, data=data, files=files)
111
    data = r.json()
112
    pprint(data)
113
    assert data["code"] == 200
114
    assert len(data["data"]["failed_files"]) == 0
115

116

117
def test_list_files(api="/knowledge_base/list_files"):
118
    url = api_base_url + api
119
    print("\n获取知识库中文件列表：")
120
    r = requests.get(url, params={"knowledge_base_name": kb})
121
    data = r.json()
122
    pprint(data)
123
    assert data["code"] == 200
124
    assert isinstance(data["data"], list)
125
    for name in test_files:
126
        assert name in data["data"]
127

128

129
def test_search_docs(api="/knowledge_base/search_docs"):
130
    url = api_base_url + api
131
    query = "介绍一下langchain-chatchat项目"
132
    print("\n检索知识库：")
133
    print(query)
134
    r = requests.post(url, json={"knowledge_base_name": kb, "query": query})
135
    data = r.json()
136
    pprint(data)
137
    assert isinstance(data, list) and len(data) == VECTOR_SEARCH_TOP_K
138

139

140
def test_update_info(api="/knowledge_base/update_info"):
141
    url = api_base_url + api
142
    print("\n更新知识库介绍")
143
    r = requests.post(url, json={"knowledge_base_name": "samples", "kb_info": "你好"})
144
    data = r.json()
145
    pprint(data)
146
    assert data["code"] == 200
147

148
def test_update_docs(api="/knowledge_base/update_docs"):
149
    url = api_base_url + api
150

151
    print(f"\n更新知识文件")
152
    r = requests.post(url, json={"knowledge_base_name": kb, "file_names": list(test_files)})
153
    data = r.json()
154
    pprint(data)
155
    assert data["code"] == 200
156
    assert len(data["data"]["failed_files"]) == 0
157

158

159
def test_delete_docs(api="/knowledge_base/delete_docs"):
160
    url = api_base_url + api
161

162
    print(f"\n删除知识文件")
163
    r = requests.post(url, json={"knowledge_base_name": kb, "file_names": list(test_files)})
164
    data = r.json()
165
    pprint(data)
166
    assert data["code"] == 200
167
    assert len(data["data"]["failed_files"]) == 0
168

169
    url = api_base_url + "/knowledge_base/search_docs"
170
    query = "介绍一下langchain-chatchat项目"
171
    print("\n尝试检索删除后的检索知识库：")
172
    print(query)
173
    r = requests.post(url, json={"knowledge_base_name": kb, "query": query})
174
    data = r.json()
175
    pprint(data)
176
    assert isinstance(data, list) and len(data) == 0
177

178

179
def test_recreate_vs(api="/knowledge_base/recreate_vector_store"):
180
    url = api_base_url + api
181
    print("\n重建知识库：")
182
    r = requests.post(url, json={"knowledge_base_name": kb}, stream=True)
183
    for chunk in r.iter_content(None):
184
        data = json.loads(chunk[6:])
185
        assert isinstance(data, dict)
186
        assert data["code"] == 200
187
        print(data["msg"])
188

189
    url = api_base_url + "/knowledge_base/search_docs"
190
    query = "本项目支持哪些文件格式?"
191
    print("\n尝试检索重建后的检索知识库：")
192
    print(query)
193
    r = requests.post(url, json={"knowledge_base_name": kb, "query": query})
194
    data = r.json()
195
    pprint(data)
196
    assert isinstance(data, list) and len(data) == VECTOR_SEARCH_TOP_K
197

198

199
def test_delete_kb_after(api="/knowledge_base/delete_knowledge_base"):
200
    url = api_base_url + api
201
    print("\n删除知识库")
202
    r = requests.post(url, json=kb)
203
    data = r.json()
204
    pprint(data)
205

206
    # check kb not exists anymore
207
    url = api_base_url + "/knowledge_base/list_knowledge_bases"
208
    print("\n获取知识库列表：")
209
    r = requests.get(url)
210
    data = r.json()
211
    pprint(data)
212
    assert data["code"] == 200
213
    assert isinstance(data["data"], list) and len(data["data"]) > 0
214
    assert kb not in data["data"]
215