File size: 568 Bytes
4231e0d
 
 
 
 
 
 
 
387ca34
4231e0d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
import zhipuai
from .base import register_llm


def ask_chatglm(message: str, api_key: str):
    zhipuai.api_key = api_key

    response = zhipuai.model_api.invoke(
        model="chatglm_turbo",
        prompt=[{
            "role": "user",
            "content": message
        }],
        top_p=0.7,
        temperature=0.9,
    )

    response_msg = response['data']['choices'][0]['content']
    # strip the front and end '"'
    if len(response_msg) >= 2:
        response_msg = response_msg[1:-1]

    return response_msg


register_llm('chatglm', ask_chatglm)