Knowledge_library_score/llms/baidullm.py

# coding=utf-8
from llms.llm import Llm
import json
import requests


class Baidullm(Llm):
    def __init__(self, **param):
        self.access_url = param['access_url']
        self.api_url = param['api_url']
        self.api_key = param['api_key']
        self.api_secret = param['api_secret']

    def get_access_token(self):
        # 获取accesstoken
        url = self.access_url + self.api_key
        url = url + '&client_secret=' + self.api_secret
        payload = json.dumps("")
        headers = {
            'Content-Type': 'application/json',
            'Accept': 'application/json'
        }
        response = requests.request("POST", url, headers=headers, data=payload)
        return response.json().get("access_token")

    def link(self, question):
        url = self.api_url + self.get_access_token()
        payload = json.dumps({
            "messages": [
                {
                    "role": "user",
                    "content": question
                }
            ]
        })
        headers = {
            'Content-Type': 'application/json'
            }

        response = requests.request("POST", url, headers=headers, data=payload)
        data = json.loads(response.text)
        answer = data['result']
        tokens = data['usage']['total_tokens']
        return answer, tokens