human/nlp/nlp_doubao.py

#encoding = utf8
import re
import time
from queue import Queue

import volcenginesdkark
import volcenginesdkcore
from volcenginesdkcore.rest import ApiException

from nlp.nlp_base import NLPBase
from volcenginesdkarkruntime import AsyncArk

nlp_queue = Queue()


class DouBao(NLPBase):
    def __init__(self, split):
        super().__init__(split)
        # Access Key ID
        # AKLTYTdmOTBmNWFjODkxNDE2Zjk3MjU0NjRhM2JhM2IyN2Y
        # AKLTNDZjNTdhNDlkZGE3NDZjMDlkMzk5YWQ3MDA4MTY1ZDc
        # Secret Access Key
        # WmpRelltRXhNbVkyWWpnNU5HRmpNamc0WTJZMFpUWmpOV1E1TTJFME1tTQ==
        # TkRJMk1tTTFZamt4TkRVNE5HRTNZMkUyTnpFeU5qQmxNMkUwWXpaak1HRQ==
        # endpoint_id
        # ep-20241008152048-fsgzf
        # api_key
        # c9635f9e-0f9e-4ca1-ac90-8af25a541b74
        # api_ky
        # eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJhcmstY29uc29sZSIsImV4cCI6MTczMDk2NTMxOSwiaWF0IjoxNzI4MzczMzE5LCJ0IjoidXNlciIsImt2IjoxLCJhaWQiOiIyMTAyMjc3NDc1IiwidWlkIjoiMCIsImlzX291dGVyX3VzZXIiOnRydWUsInJlc291cmNlX3R5cGUiOiJlbmRwb2ludCIsInJlc291cmNlX2lkcyI6WyJlcC0yMDI0MTAwODE1MjA0OC1mc2d6ZiJdfQ.BHgFj-UKeu7IGG5VL2e6iPQEMNMkQrgmM46zYmTpoNG_ySgSFJLWYzbrIABZmqVDB4Rt58j8kvoORs-RHJUz81rXUlh3BYl9-ZwbggtAU7Z1pm54_qZ00jF0jQ6r-fUSXZo2PVCLxb_clNuEh06NyaV7ullZwUCyLKx3vhCsxPAuEvQvLc_qDBx-IYNT-UApVADaqMs-OyewoxahqQ7RvaHFF14R6ihmg9H0uvl00_JiGThJveszKvy_T-Qk6iPOy-EDI2pwJxdHMZ7By0bWK5EfZoK2hOvOSRD0BNTYnvrTfI0l2JgS0nwCVEPR4KSTXxU_oVVtuUSZp1UHvvkhvA
        self.__token = 'c9635f9e-0f9e-4ca1-ac90-8af25a541b74'
        self.__client = AsyncArk(api_key=self.__token)

    async def _request(self, question):
        t = time.time()
        print(f'-------dou_bao ask:', question)
        stream = await self.__client.chat.completions.create(
            model="ep-20241008152048-fsgzf",
            messages=[
                {"role": "system", "content": "你是测试客服，是由字节跳动开发的 AI 人工智能助手"},
                {"role": "user", "content": question},
            ],
            stream=True
        )
        sec = ''
        async for completion in stream:
            # print(f'-------dou_bao nlp time:{time.time() - t:.4f}s')
            # nlp_queue.put(completion.choices[0].delta.content)
            # print(completion.choices[0].delta.content, end="")
            sec = sec + completion.choices[0].delta.content
            sec, message = self._split_handle.handle(sec)
            if len(message) > 0:
                print(message)
        print(sec)
        print(f'-------dou_bao nlp time:{time.time() - t:.4f}s')

'''
if __name__ == "__main__":
    # print(get_dou_bao_api())
    dou_bao = DouBao()
    dou_bao.ask('你好。')
    dou_bao.ask('你好，你是谁？')
    dou_bao.ask('你能做什么？')
    dou_bao.ask('介绍一下，我自己。')
    count = 1000
    sec = ''
    while count >= 0:
        count = count - 1
        if nlp_queue.empty():
            time.sleep(0.1)
            continue
        sec = sec + nlp_queue.get(block=True, timeout=0.01)

        pattern = r'[，。、；？！,.!?]'
        match = re.search(pattern, sec)
        if match:
            pos = match.start() + 1
            print(sec[:pos])
            sec = sec[pos:]
    print(sec)


    dou_bao.stop()
'''