bisheng-langchain 0.0.1__py3-none-any.whl

bisheng_langchain/chat_models/interface/wenxin.py

@@ -0,0 +1,114 @@
+import json
+
+import requests
+
+from .types import ChatInput, ChatOutput, Choice, Message, Usage
+from .utils import get_ts
+
+
+def get_access_token(api_key, sec_key):
+    url = (f'https://aip.baidubce.com/oauth/2.0/token?'
+           f'grant_type=client_credentials'
+           f'&client_id={api_key}&client_secret={sec_key}')
+
+    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')
+
+
+class ChatCompletion(object):
+
+    def __init__(self, api_key, sec_key, **kwargs):
+        self.api_key = api_key
+        self.sec_key = sec_key
+        self.ep_url = ('https://aip.baidubce.com/rpc/2.0/ai_custom/v1/'
+                       'wenxinworkshop/chat/completions')
+        self.ep_url_turbo = ('https://aip.baidubce.com/rpc/2.0/ai_custom/v1/'
+                             'wenxinworkshop/chat/eb-instant')
+
+        # token = get_access_token(api_key, sec_key)
+        # self.endpoint = f"{self.ep_url}?access_token={token}"
+        self.headers = {'Content-Type': 'application/json'}
+
+    def __call__(self, inp: ChatInput, verbose=False):
+        messages = inp.messages
+        model = inp.model
+        top_p = 0.8 if inp.top_p is None else inp.top_p
+        temperature = 0.95 if inp.temperature is None else inp.temperature
+        stream = False if inp.stream is None else inp.stream
+        # max_tokens = 1024 if inp.max_tokens is None else inp.max_tokens
+
+        system_content = ''
+        new_messages = []
+        for m in messages:
+            role = m.role
+            if role == 'system':
+                system_content = m.content
+                continue
+            new_messages.append({'role': role, 'content': m.content})
+
+        if system_content:
+            new_messages[-1]['content'] = system_content + '\n' + new_messages[
+                -1]['content']
+
+        payload = {
+            'stream': stream,
+            'messages': new_messages,
+            'temperature': temperature,
+            'top_p': top_p
+        }
+
+        if verbose:
+            print('payload', payload)
+
+        token = get_access_token(self.api_key, self.sec_key)
+        endpoint = f'{self.ep_url}?access_token={token}'
+        if model == 'ernie-bot-turbo':
+            endpoint = f'{self.ep_url_turbo}?access_token={token}'
+
+        response = requests.post(endpoint, headers=self.headers, json=payload)
+
+        req_type = 'chat.completion'
+        status_message = 'success'
+        status_code = response.status_code
+        created = get_ts()
+        choices = []
+        usage = None
+        if status_code == 200:
+            try:
+                info = json.loads(response.text)
+                status_code = info.get('error_code', 200)
+                status_message = info.get('error_msg', status_message)
+                if status_code == 200:
+                    created = info['created']
+                    result = info['result']
+                    finish_reason = 'default'
+                    msg = Message(role='assistant', content=result)
+                    choices = [
+                        Choice(index=0,
+                               message=msg,
+                               finish_reason=finish_reason)
+                    ]
+                    usage = Usage(**info['usage'])
+            except Exception as e:
+                status_code = 401
+                status_message = str(e)
+        else:
+            status_code = 400
+            status_message = 'requests error'
+
+        if status_code != 200:
+            raise Exception(status_message)
+
+        return ChatOutput(status_code=status_code,
+                          status_message=status_message,
+                          model=model,
+                          object=req_type,
+                          created=created,
+                          choices=choices,
+                          usage=usage)

bisheng_langchain/chat_models/interface/xunfei.py

@@ -0,0 +1,233 @@
+import base64
+import hashlib
+import hmac
+import json
+from datetime import datetime
+from time import mktime
+from urllib.parse import urlencode, urlparse
+from wsgiref.handlers import format_date_time
+
+import websocket
+from websocket import create_connection
+
+import _thread as thread
+
+from .types import ChatInput, ChatOutput, Choice, Message, Usage
+from .utils import get_ts
+
+# import ssl
+# import threading
+
+
+class Ws_Param(object):
+    # 初始化
+    def __init__(self, APPID, APIKey, APISecret, gpt_url):
+        self.APPID = APPID
+        self.APIKey = APIKey
+        self.APISecret = APISecret
+        self.host = urlparse(gpt_url).netloc
+        self.path = urlparse(gpt_url).path
+        self.gpt_url = gpt_url
+
+    # 生成url
+    def create_url(self):
+        # 生成RFC1123格式的时间戳
+        now = datetime.now()
+        date = format_date_time(mktime(now.timetuple()))
+
+        # 拼接字符串
+        signature_origin = 'host: ' + self.host + '\n'
+        signature_origin += 'date: ' + date + '\n'
+        signature_origin += 'GET ' + self.path + ' HTTP/1.1'
+
+        # 进行hmac-sha256进行加密
+        signature_sha = hmac.new(self.APISecret.encode('utf-8'),
+                                 signature_origin.encode('utf-8'),
+                                 digestmod=hashlib.sha256).digest()
+
+        signature_sha_base64 = base64.b64encode(signature_sha).decode(
+            encoding='utf-8')
+
+        authorization_origin = (
+            f'api_key="{self.APIKey}", '
+            f'algorithm="hmac-sha256", headers="host date request-line",'
+            f' signature="{signature_sha_base64}"')
+
+        authorization = base64.b64encode(
+            authorization_origin.encode('utf-8')).decode(encoding='utf-8')
+
+        # 将请求的鉴权参数组合为字典
+        v = {'authorization': authorization, 'date': date, 'host': self.host}
+        # 拼接鉴权参数，生成url
+        url = self.gpt_url + '?' + urlencode(v)
+        # 此处打印出建立连接时候的url,参考本demo的时候可取消上方打印的注释，
+        # 比对相同参数时生成的url与自己代码生成的url是否一致
+        return url
+
+
+# 收到websocket错误的处理
+def on_error(ws, error):
+    print('### error:', error)
+
+
+# 收到websocket关闭的处理
+def on_close(ws):
+    print('### closed ###')
+
+
+# 收到websocket连接建立的处理
+def on_open(ws):
+    thread.start_new_thread(run, (ws, ))
+
+
+def run(ws, *args):
+    data = json.dumps(gen_params(appid=ws.appid, question=ws.question))
+    ws.send(data)
+
+
+# 收到websocket消息的处理
+def on_message(ws, message):
+    print(message)
+    data = json.loads(message)
+    code = data['header']['code']
+    if code != 0:
+        print(f'请求错误: {code}, {data}')
+        ws.close()
+    else:
+        choices = data['payload']['choices']
+        status = choices['status']
+        content = choices['text'][0]['content']
+        print(content, end='')
+        if status == 2:
+            ws.close()
+
+
+def gen_params(appid, question):
+    data = {
+        'header': {
+            'app_id': appid,
+            'uid': '1234'
+        },
+        'parameter': {
+            'chat': {
+                'domain': 'general',
+                'random_threshold': 0.5,
+                'max_tokens': 2048,
+                'auditing': 'default'
+            }
+        },
+        'payload': {
+            'message': {
+                'text': [{
+                    'role': 'user',
+                    'content': question
+                }]
+            }
+        }
+    }
+    return data
+
+
+class ChatCompletion(object):
+
+    def __init__(self, appid, api_key, api_secret, **kwargs):
+        gpt_url = 'ws://spark-api.xf-yun.com/v1.1/chat'
+        self.wsParam = Ws_Param(appid, api_key, api_secret, gpt_url)
+        websocket.enableTrace(False)
+        # wsUrl = wsParam.create_url()
+
+        # todo: modify to the ws pool
+        # self.mutex = threading.Lock()
+        # self.ws = websocket.WebSocket()
+        # self.ws.connect(wsUrl)
+
+        self.header = {'app_id': appid, 'uid': 'elem'}
+
+    def __call__(self, inp: ChatInput, verbose=False):
+        messages = inp.messages
+        model = inp.model
+        # top_p = 0.7 if inp.top_p is None else inp.top_p
+        temperature = 0.5 if inp.temperature is None else inp.temperature
+        # stream = False if inp.stream is None else inp.stream
+        max_tokens = 1024 if inp.max_tokens is None else inp.max_tokens
+        # stop = None
+        # if inp.stop is not None:
+        #     stop = inp.stop.split('||')
+
+        new_messages = []
+        for m in messages:
+            role = m.role
+            if role == 'system':
+                role = 'user'
+            new_messages.append({'role': role, 'content': m.content})
+
+        created = get_ts()
+        payload = {
+            'header': self.header,
+            'payload': {
+                'message': {
+                    'text': new_messages
+                }
+            },
+            'parameter': {
+                'chat': {
+                    'domain': 'general',
+                    'temperature': temperature,
+                    'max_tokens': max_tokens,
+                    'auditing': 'default'
+                }
+            }
+        }
+
+        if verbose:
+            print('payload', payload)
+
+        req_type = 'chat.completion'
+        status_code = 200
+        status_message = 'success'
+        choices = []
+        usage = None
+        texts = []
+        ws = None
+        try:
+            # self.mutex.acquire()
+            wsUrl = self.wsParam.create_url()
+            ws = create_connection(wsUrl)
+            ws.send(json.dumps(payload))
+            while True:
+                raw_data = ws.recv()
+                if not raw_data:
+                    break
+
+                resp = json.loads(raw_data)
+                if resp['header']['code'] == 0:
+                    texts.append(
+                        resp['payload']['choices']['text'][0]['content'])
+                if resp['header']['code'] == 0 and resp['header'][
+                        'status'] == 2:
+                    usage_dict = resp['payload']['usage']['text']
+                    usage_dict.pop('question_tokens')
+                    usage = Usage(**usage_dict)
+        except Exception as e:
+            status_code = 401
+            status_message = str(e)
+        finally:
+            if ws:
+                ws.close()
+
+        if texts:
+            finish_reason = 'default'
+            msg = Message(role='assistant', content=''.join(texts))
+            cho = Choice(index=0, message=msg, finish_reason=finish_reason)
+            choices.append(cho)
+
+        if status_code != 200:
+            raise Exception(status_message)
+
+        return ChatOutput(status_code=status_code,
+                          status_message=status_message,
+                          model=model,
+                          object=req_type,
+                          created=created,
+                          choices=choices,
+                          usage=usage)

bisheng_langchain/chat_models/interface/zhipuai.py

@@ -0,0 +1,81 @@
+# import json
+
+import zhipuai
+
+from .types import ChatInput, ChatOutput, Choice, Message, Usage
+from .utils import get_ts
+
+
+class ChatCompletion(object):
+
+    def __init__(self, api_key, **kwargs):
+        zhipuai.api_key = api_key
+
+    def __call__(self, inp: ChatInput, verbose=False):
+        messages = inp.messages
+        model = inp.model
+        top_p = 0.7 if inp.top_p is None else inp.top_p
+        temperature = 0.95 if inp.temperature is None else inp.temperature
+        # stream = False if inp.stream is None else inp.stream
+        # max_tokens = 1024 if inp.max_tokens is None else inp.max_tokens
+
+        new_messages = []
+        system_content = ''
+        for m in messages:
+            content = m.content
+            role = m.role
+            if role == 'system':
+                system_content += content
+                continue
+            new_messages.append({'role': role, 'content': content})
+
+        if system_content:
+            new_messages[-1]['content'] = (system_content +
+                                           new_messages[-1]['content'])
+
+        created = get_ts()
+        payload = {
+            'model': model,
+            'prompt': new_messages,
+            'temperature': temperature,
+            'top_p': top_p,
+            'request_id': str(created),
+            'incremental': False
+        }
+
+        if verbose:
+            print('payload', payload)
+
+        req_type = 'chat.completion'
+        status_message = 'success'
+        choices = []
+        usage = None
+        try:
+            resp = zhipuai.model_api.invoke(**payload)
+            status_code = resp['code']
+            status_message = resp['msg']
+            if status_code == 200:
+                choices = []
+                for index, choice in enumerate(resp['data']['choices']):
+                    finish_reason = 'default'
+                    msg = Message(**choice)
+                    cho = Choice(index=index,
+                                 message=msg,
+                                 finish_reason=finish_reason)
+                    choices.append(cho)
+                usage = Usage(**resp['data']['usage'])
+
+        except Exception as e:
+            status_code = 400
+            status_message = str(e)
+
+        if status_code != 200:
+            raise Exception(status_message)
+
+        return ChatOutput(status_code=status_code,
+                          status_message=status_message,
+                          model=model,
+                          object=req_type,
+                          created=created,
+                          choices=choices,
+                          usage=usage)