Langchain-Chatchat

1
import json
2
import time
3
import hashlib
4

5
from fastchat.conversation import Conversation
6
from server.model_workers.base import *
7
from server.utils import get_httpx_client
8
from fastchat import conversation as conv
9
import json
10
from typing import List, Literal, Dict
11
import requests
12

13

14
class TianGongWorker(ApiModelWorker):
15
    def __init__(
16
            self,
17
            *,
18
            controller_addr: str = None,
19
            worker_addr: str = None,
20
            model_names: List[str] = ["tiangong-api"],
21
            version: Literal["SkyChat-MegaVerse"] = "SkyChat-MegaVerse",
22
            **kwargs,
23
    ):
24
        kwargs.update(model_names=model_names, controller_addr=controller_addr, worker_addr=worker_addr)
25
        kwargs.setdefault("context_len", 32768)
26
        super().__init__(**kwargs)
27
        self.version = version
28

29
    def do_chat(self, params: ApiChatParams) -> Dict:
30
        params.load_config(self.model_names[0])
31

32
        url = 'https://sky-api.singularity-ai.com/saas/api/v4/generate'
33
        data = {
34
            "messages": params.messages,
35
            "model": "SkyChat-MegaVerse"
36
        }
37
        timestamp = str(int(time.time()))
38
        sign_content = params.api_key + params.secret_key + timestamp
39
        sign_result = hashlib.md5(sign_content.encode('utf-8')).hexdigest()
40
        headers = {
41
            "app_key": params.api_key,
42
            "timestamp": timestamp,
43
            "sign": sign_result,
44
            "Content-Type": "application/json",
45
            "stream": "true"  # or change to "false" 不处理流式返回内容
46
        }
47

48
        # 发起请求并获取响应
49
        response = requests.post(url, headers=headers, json=data, stream=True)
50

51
        text = ""
52
        # 处理响应流
53
        for line in response.iter_lines(chunk_size=None, decode_unicode=True):
54
            if line:
55
                # 处理接收到的数据
56
                # print(line.decode('utf-8'))
57
                resp = json.loads(line)
58
                if resp["code"] == 200:
59
                    text += resp['resp_data']['reply']
60
                    yield {
61
                        "error_code": 0,
62
                        "text": text
63
                    }
64
                else:
65
                    data = {
66
                        "error_code": resp["code"],
67
                        "text": resp["code_msg"]
68
                    }
69
                    self.logger.error(f"请求天工 API 时出错：{data}")
70
                    yield data
71

72
    def get_embeddings(self, params):
73
        print("embedding")
74
        print(params)
75

76
    def make_conv_template(self, conv_template: str = None, model_path: str = None) -> Conversation:
77
        return conv.Conversation(
78
            name=self.model_names[0],
79
            system_message="",
80
            messages=[],
81
            roles=["user", "system"],
82
            sep="\n### ",
83
            stop_str="###",
84
        )
85