百度千帆大模型API SDK

本文介绍百度千帆大模型API SDK

SDK如下:

import requests
import json
import datetime,time

class QIANFAN:

    _api_url = "https://aip.baidubce.com"
    
    def __init__(self, api_key, secret_key):
        self.API_KEY = api_key
        self.SECRET_KEY = secret_key

        url = self._api_url + "/oauth/2.0/token"
        params = {"grant_type": "client_credentials", "client_id": self.API_KEY, "client_secret": self.SECRET_KEY}
        result = self.http_request_v2(url, method="POST", params=params)
        if 'access_token' in result:
            self.access_token = result["access_token"]
        else:
            exit(print(result))

    def chat(self, model="ernie-lite-8k", message=None, **kwargs):
        url = f"{self._api_url}/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/{model}?access_token={self.access_token}"
        payload = {
            "messages": [{"role": "user", "content": message}],
            "temperature": 0.95,
            "penalty_score": 1
        }

        for key, value in kwargs.items():
            payload[key] = value

        response = self.http_request_v2(url, method="POST", params=json.dumps(payload))
        return response

    # 生成headers头
    def headers(self, params=None):
        headers = {}
        headers['Content-Type'] = 'application/json'
        return headers
                    
    def http_request_v2(self, url, method="GET", headers={}, params=None):
        headers['User-Agent'] = 'Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/39.0.2171.71 Safari/537.36'

        max_seconds, max_times = 10, 3
        for attempt in range(max_times):
            try:
                if method == "GET":
                    response = requests.get(url, headers=headers, timeout=max_seconds)
                elif method == "POST":
                    response = requests.post(url, data=params, headers=headers, timeout=max_seconds)
                elif method == "DELETE":
                    response = requests.delete(url, data=params, headers=headers, timeout=max_seconds)

                return response.json()
            
            except Exception as e:
                print([datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'), attempt, url, method, params, repr(e)])
                time.sleep(max_seconds * (attempt + 1))

调用方法如下：

    API_KEY = ""
    SECRET_KEY = ""

    chat_client = QIANFAN(API_KEY, SECRET_KEY)
    print(vars(chat_client))
    model = "ernie-lite-8k"
  
    personality = "你是数学专家"
   
    prompt1 = f"1+1等于多少"

    result = chat_client.chat(model=model, message=prompt1, system = personality)
    print(result)
    print(result['result'])

返回如下

{
    'id': 'as-t2wx1kt5wm',
    'object': 'chat.completion',
    'created': 1733380170,
    'result': '在数学运算中，1+1 等于 2。这是一 个基本的加法运算，是算术教学的基础之一。如果你有任何疑问，请随时向我提问。',
    'is_truncated': False,
    'need_clear_history': False,
    'finish_reason': 'normal',
    'usage': {
        'prompt_tokens': 10,
        'completion_tokens': 35,
        'total_tokens': 45
    }
}