human/nlp/nlp_doubao.py

#encoding = utf8

import logging
import time

from nlp.nlp_base import NLPBase
from volcenginesdkarkruntime import AsyncArk

logger = logging.getLogger(__name__)


class DouBao(NLPBase):
    def __init__(self, split, callback=None):
        super().__init__(split, callback)
        logger.info("DouBao init")
        # 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()
        logger.info(f'_request:{question}')
        logger.info(f'-------dou_bao ask:{question}')
        try:
            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:
                sec = sec + completion.choices[0].delta.content
                sec, message = self._split_handle.handle(sec)
                if len(message) > 0:
                    self._on_callback(message)
            self._on_callback(sec)
            await stream.close()
        except Exception as e:
            print(e)
        logger.info(f'_request:{question}, time:{time.time() - t:.4f}s')
        logger.info(f'-------dou_bao nlp time:{time.time() - t:.4f}s')

    async def _on_close(self):
        logger.info('AsyncArk close')
        if self.__client is not None and not self.__client.is_closed():
            await self.__client.close()