Langchain-Chatchat

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

6
sys.path.append(str(Path(__file__).parent.parent.parent))
7
from configs import BING_SUBSCRIPTION_KEY
8
from server.utils import api_address
9

10
from pprint import pprint
11

12

13
api_base_url = api_address()
14

15

16
def dump_input(d, title):
17
    print("\n")
18
    print("=" * 30 + title + "  input " + "="*30)
19
    pprint(d)
20

21

22
def dump_output(r, title):
23
    print("\n")
24
    print("=" * 30 + title + "  output" + "="*30)
25
    for line in r.iter_content(None, decode_unicode=True):
26
        print(line, end="", flush=True)
27

28

29
headers = {
30
    'accept': 'application/json',
31
    'Content-Type': 'application/json',
32
}
33

34
data = {
35
    "query": "请用100字左右的文字介绍自己",
36
    "history": [
37
        {
38
            "role": "user",
39
            "content": "你好"
40
        },
41
        {
42
            "role": "assistant",
43
            "content": "你好，我是人工智能大模型"
44
        }
45
    ],
46
    "stream": True,
47
    "temperature": 0.7,
48
}
49

50

51
def test_chat_chat(api="/chat/chat"):
52
    url = f"{api_base_url}{api}"
53
    dump_input(data, api)
54
    response = requests.post(url, headers=headers, json=data, stream=True)
55
    dump_output(response, api)
56
    assert response.status_code == 200
57

58

59
def test_knowledge_chat(api="/chat/knowledge_base_chat"):
60
    url = f"{api_base_url}{api}"
61
    data = {
62
        "query": "如何提问以获得高质量答案",
63
        "knowledge_base_name": "samples",
64
        "history": [
65
            {
66
                "role": "user",
67
                "content": "你好"
68
            },
69
            {
70
                "role": "assistant",
71
                "content": "你好，我是 ChatGLM"
72
            }
73
        ],
74
        "stream": True
75
    }
76
    dump_input(data, api)
77
    response = requests.post(url, headers=headers, json=data, stream=True)
78
    print("\n")
79
    print("=" * 30 + api + "  output" + "="*30)
80
    for line in response.iter_content(None, decode_unicode=True):
81
        data = json.loads(line[6:])
82
        if "answer" in data:
83
            print(data["answer"], end="", flush=True)
84
    pprint(data)
85
    assert "docs" in data and len(data["docs"]) > 0
86
    assert response.status_code == 200
87

88

89
def test_search_engine_chat(api="/chat/search_engine_chat"):
90
    global data
91

92
    data["query"] = "室温超导最新进展是什么样？"
93

94
    url = f"{api_base_url}{api}"
95
    for se in ["bing", "duckduckgo"]:
96
        data["search_engine_name"] = se
97
        dump_input(data, api + f" by {se}")
98
        response = requests.post(url, json=data, stream=True)
99
        if se == "bing" and not BING_SUBSCRIPTION_KEY:
100
            data = response.json()
101
            assert data["code"] == 404
102
            assert data["msg"] == f"要使用Bing搜索引擎，需要设置 `BING_SUBSCRIPTION_KEY`"
103

104
        print("\n")
105
        print("=" * 30 + api + f" by {se}  output" + "="*30)
106
        for line in response.iter_content(None, decode_unicode=True):
107
            data = json.loads(line[6:])
108
            if "answer" in data:
109
                print(data["answer"], end="", flush=True)
110
        assert "docs" in data and len(data["docs"]) > 0
111
        pprint(data["docs"])
112
        assert response.status_code == 200
113

114