diff --git a/xiaogpt/bot/doubao_bot.py b/xiaogpt/bot/doubao_bot.py index 4a268cc9..581fe7f0 100644 --- a/xiaogpt/bot/doubao_bot.py +++ b/xiaogpt/bot/doubao_bot.py @@ -1,4 +1,4 @@ -"""ChatGLM bot""" +"""Doubao (Volcengine Ark) bot""" from __future__ import annotations @@ -8,22 +8,23 @@ from xiaogpt.bot.base_bot import BaseBot, ChatHistoryMixin from xiaogpt.config import Config +from xiaogpt.utils import split_sentences class DoubaoBot(ChatHistoryMixin, BaseBot): name = "豆包" - default_options = {"model": "skylark-chat"} # 根据官方示例修改默认模型 + default_options = {"model": "doubao-pro-32k"} def __init__(self, api_key: str) -> None: - from volcenginesdkarkruntime import Ark # 引入官方 SDK + from volcenginesdkarkruntime import AsyncArk self.api_key = api_key self.history = [] - self.client = Ark(api_key=api_key) # 初始化客户端 + self.client = AsyncArk(api_key=api_key) @classmethod def from_config(cls, config: Config): - return cls(api_key=config.volc_api_key) # 假设配置文件中有 volc_api_key 字段 + return cls(api_key=config.volc_api_key) def _get_data(self, query: str, **options: Any): options = {**self.default_options, **options} @@ -35,29 +36,40 @@ def _get_data(self, query: str, **options: Any): async def ask(self, query, **options): data = self._get_data(query, **options) try: - completion = self.client.chat.completions.create(**data) + completion = await self.client.chat.completions.create(**data) message = completion.choices[0].message.content self.add_message(query, message) print(message) return message except Exception as e: print(str(e)) - return + return "" async def ask_stream(self, query: str, **options: Any): data = self._get_data(query, **options) data["stream"] = True try: - full_content = "" - for chunk in self.client.chat.completions.create(**data): - content = chunk.choices[0].delta.content - if content: - full_content += content - print(content, end="", flush=True) - yield content - print() - self.add_message(query, full_content) + completion = await self.client.chat.completions.create(**data) except Exception as e: print(str(e)) return + + async def text_gen(): + async for chunk in completion: + if not chunk.choices: + continue + content = chunk.choices[0].delta.content + if content is None: + continue + print(content, end="", flush=True) + yield content + + message = "" + try: + async for sentence in split_sentences(text_gen()): + message += sentence + yield sentence + finally: + print() + self.add_message(query, message)