Langchain-Chatchat

1
import sys
2
from pathlib import Path
3
root_path = Path(__file__).parent.parent
4
sys.path.append(str(root_path))
5

6
from configs import ONLINE_LLM_MODEL
7
from server.model_workers.base import *
8
from server.utils import get_model_worker_config, list_config_llm_models
9
from pprint import pprint
10
import pytest
11

12

13
workers = []
14
for x in list_config_llm_models()["online"]:
15
    if x in ONLINE_LLM_MODEL and x not in workers:
16
        workers.append(x)
17
print(f"all workers to test: {workers}")
18

19
# workers = ["fangzhou-api"]
20

21

22
@pytest.mark.parametrize("worker", workers)
23
def test_chat(worker):
24
    params = ApiChatParams(
25
        messages = [
26
            {"role": "user", "content": "你是谁"},
27
        ],
28
    )
29
    print(f"\nchat with {worker} \n")
30

31
    if worker_class := get_model_worker_config(worker).get("worker_class"):
32
        for x in worker_class().do_chat(params):
33
            pprint(x)
34
            assert isinstance(x, dict)
35
            assert x["error_code"] == 0
36

37

38
@pytest.mark.parametrize("worker", workers)
39
def test_embeddings(worker):
40
    params = ApiEmbeddingsParams(
41
        texts = [
42
            "LangChain-Chatchat (原 Langchain-ChatGLM): 基于 Langchain 与 ChatGLM 等大语言模型的本地知识库问答应用实现。",
43
            "一种利用 langchain 思想实现的基于本地知识库的问答应用，目标期望建立一套对中文场景与开源模型支持友好、可离线运行的知识库问答解决方案。",
44
        ]
45
    )
46

47
    if worker_class := get_model_worker_config(worker).get("worker_class"):
48
        if worker_class.can_embedding():
49
            print(f"\embeddings with {worker} \n")
50
            resp = worker_class().do_embeddings(params)
51

52
            pprint(resp, depth=2)
53
            assert resp["code"] == 200
54
            assert "data" in resp
55
            embeddings = resp["data"]
56
            assert isinstance(embeddings, list) and len(embeddings) > 0
57
            assert isinstance(embeddings[0], list) and len(embeddings[0]) > 0
58
            assert isinstance(embeddings[0][0], float)
59
            print("向量长度：", len(embeddings[0]))
60

61

62
# @pytest.mark.parametrize("worker", workers)
63
# def test_completion(worker):
64
#     params = ApiCompletionParams(prompt="五十六个民族")
65
    
66
#     print(f"\completion with {worker} \n")
67

68
#     worker_class = get_model_worker_config(worker)["worker_class"]
69
#     resp = worker_class().do_completion(params)
70
#     pprint(resp)
71