agents/ten_packages/extension/coze_python_async/extension.py

#
# This file is part of TEN Framework, an open source project.
# Licensed under the Apache License, Version 2.0.
# See the LICENSE file for more information.
#
import asyncio
import traceback
import aiohttp
import json
import copy

from typing import List, Any, AsyncGenerator
from dataclasses import dataclass

from cozepy import ChatEventType, Message, TokenAuth, AsyncCoze, ChatEvent, Chat

from ten import (
    AudioFrame,
    VideoFrame,
    AsyncTenEnv,
    Cmd,
    StatusCode,
    CmdResult,
    Data,
)

from ten_ai_base.config import BaseConfig
from ten_ai_base.chat_memory import ChatMemory
from ten_ai_base.types import (
    LLMChatCompletionUserMessageParam,
    LLMCallCompletionArgs,
    LLMDataCompletionArgs,
    LLMToolMetadata,
)
from ten_ai_base.llm import (
    AsyncLLMBaseExtension,
)

CMD_IN_FLUSH = "flush"
CMD_IN_ON_USER_JOINED = "on_user_joined"
CMD_IN_ON_USER_LEFT = "on_user_left"
CMD_OUT_FLUSH = "flush"
CMD_OUT_TOOL_CALL = "tool_call"

DATA_IN_TEXT_DATA_PROPERTY_IS_FINAL = "is_final"
DATA_IN_TEXT_DATA_PROPERTY_TEXT = "text"

DATA_OUT_TEXT_DATA_PROPERTY_TEXT = "text"
DATA_OUT_TEXT_DATA_PROPERTY_END_OF_SEGMENT = "end_of_segment"

CMD_PROPERTY_RESULT = "tool_result"


def is_punctuation(char):
    if char in [",", "，", ".", "。", "?", "？", "!", "！"]:
        return True
    return False


def parse_sentences(sentence_fragment, content):
    sentences = []
    current_sentence = sentence_fragment
    for char in content:
        current_sentence += char
        if is_punctuation(char):
            stripped_sentence = current_sentence
            if any(c.isalnum() for c in stripped_sentence):
                sentences.append(stripped_sentence)
            current_sentence = ""

    remain = current_sentence
    return sentences, remain


@dataclass
class CozeConfig(BaseConfig):
    base_url: str = "https://api.acoze.com"
    bot_id: str = ""
    token: str = ""
    user_id: str = "TenAgent"
    greeting: str = ""
    max_history: int = 32


class AsyncCozeExtension(AsyncLLMBaseExtension):
    config: CozeConfig = None
    sentence_fragment: str = ""
    ten_env: AsyncTenEnv = None
    loop: asyncio.AbstractEventLoop = None
    stopped: bool = False
    users_count = 0
    memory: ChatMemory = None

    acoze: AsyncCoze = None
    # conversation: str = ""

    async def on_init(self, ten_env: AsyncTenEnv) -> None:
        await super().on_init(ten_env)
        ten_env.log_debug("on_init")

    async def on_start(self, ten_env: AsyncTenEnv) -> None:
        await super().on_start(ten_env)
        ten_env.log_debug("on_start")

        self.loop = asyncio.get_event_loop()

        self.config = await CozeConfig.create_async(ten_env=ten_env)
        ten_env.log_info(f"config: {self.config}")

        if not self.config.bot_id or not self.config.token:
            ten_env.log_error("Missing required configuration")
            return

        self.memory = ChatMemory(self.config.max_history)
        try:
            self.acoze = AsyncCoze(
                auth=TokenAuth(token=self.config.token), base_url=self.config.base_url
            )

            # self.conversation = await self.acoze.conversations.create(messages = [
            #        Message.build_user_question_text(self.config.prompt)
            #    ] if self.config.prompt else [])

        except Exception as e:
            ten_env.log_error(f"Failed to create conversation {e}")

        self.ten_env = ten_env

    async def on_stop(self, ten_env: AsyncTenEnv) -> None:
        await super().on_stop(ten_env)
        ten_env.log_debug("on_stop")

        self.stopped = True

    async def on_deinit(self, ten_env: AsyncTenEnv) -> None:
        await super().on_deinit(ten_env)
        ten_env.log_debug("on_deinit")

    async def on_cmd(self, ten_env: AsyncTenEnv, cmd: Cmd) -> None:
        cmd_name = cmd.get_name()
        ten_env.log_debug("on_cmd name {}".format(cmd_name))

        status = StatusCode.OK
        detail = "success"

        if cmd_name == CMD_IN_FLUSH:
            await self.flush_input_items(ten_env)
            await ten_env.send_cmd(Cmd.create(CMD_OUT_FLUSH))
            ten_env.log_info("on flush")
        elif cmd_name == CMD_IN_ON_USER_JOINED:
            self.users_count += 1
            # Send greeting when first user joined
            if self.config.greeting and self.users_count == 1:
                self.send_text_output(ten_env, self.config.greeting, True)
        elif cmd_name == CMD_IN_ON_USER_LEFT:
            self.users_count -= 1
        else:
            await super().on_cmd(ten_env, cmd)
            return

        cmd_result = CmdResult.create(status)
        cmd_result.set_property_string("detail", detail)
        await ten_env.return_result(cmd_result, cmd)

    async def on_call_chat_completion(
        self, ten_env: AsyncTenEnv, **kargs: LLMCallCompletionArgs
    ) -> any:
        raise RuntimeError("Not implemented")

    async def on_data_chat_completion(
        self, ten_env: AsyncTenEnv, **kargs: LLMDataCompletionArgs
    ) -> None:
        if not self.acoze:
            await self._send_text(
                "Coze is not connected. Please check your configuration.", True
            )
            return

        input_messages: LLMChatCompletionUserMessageParam = kargs.get("messages", [])
        messages = copy.copy(self.memory.get())
        if not input_messages:
            ten_env.log_warn("No message in data")
        else:
            messages.extend(input_messages)
            for i in input_messages:
                self.memory.put(i)

        total_output = ""
        sentence_fragment = ""
        calls = {}

        sentences = []
        self.ten_env.log_info(f"messages: {messages}")
        response = self._stream_chat(messages=messages)
        async for message in response:
            self.ten_env.log_info(f"content: {message}")
            try:
                if message.event == ChatEventType.CONVERSATION_MESSAGE_DELTA:
                    total_output += message.message.content
                    sentences, sentence_fragment = parse_sentences(
                        sentence_fragment, message.message.content
                    )
                    for s in sentences:
                        await self._send_text(s, False)
                elif message.event == ChatEventType.CONVERSATION_MESSAGE_COMPLETED:
                    if sentence_fragment:
                        await self._send_text(sentence_fragment, True)
                    else:
                        await self._send_text("", True)
                elif message.event == ChatEventType.CONVERSATION_CHAT_FAILED:
                    last_error = message.chat.last_error
                    if last_error and last_error.code == 4011:
                        await self._send_text(
                            "The Coze token has been depleted. Please check your token usage.",
                            True,
                        )
                    else:
                        await self._send_text(last_error.msg, True)
            except Exception as e:
                self.ten_env.log_error(f"Failed to parse response: {message} {e}")
                traceback.print_exc()

        self.memory.put({"role": "assistant", "content": total_output})
        self.ten_env.log_info(f"total_output: {total_output} {calls}")

    async def on_tools_update(
        self, ten_env: AsyncTenEnv, tool: LLMToolMetadata
    ) -> None:
        # Implement the logic for tool updates
        return await super().on_tools_update(ten_env, tool)

    async def on_data(self, ten_env: AsyncTenEnv, data: Data) -> None:
        data_name = data.get_name()
        ten_env.log_info("on_data name {}".format(data_name))

        is_final = False
        input_text = ""
        try:
            is_final = data.get_property_bool(DATA_IN_TEXT_DATA_PROPERTY_IS_FINAL)
        except Exception as err:
            ten_env.log_info(
                f"GetProperty optional {DATA_IN_TEXT_DATA_PROPERTY_IS_FINAL} failed, err: {err}"
            )

        try:
            input_text = data.get_property_string(DATA_IN_TEXT_DATA_PROPERTY_TEXT)
        except Exception as err:
            ten_env.log_info(
                f"GetProperty optional {DATA_IN_TEXT_DATA_PROPERTY_TEXT} failed, err: {err}"
            )

        if not is_final:
            ten_env.log_info("ignore non-final input")
            return
        if not input_text:
            ten_env.log_info("ignore empty text")
            return

        ten_env.log_info(f"OnData input text: [{input_text}]")

        # Start an asynchronous task for handling chat completion
        message = LLMChatCompletionUserMessageParam(role="user", content=input_text)
        await self.queue_input_item(False, messages=[message])

    async def on_audio_frame(
        self, ten_env: AsyncTenEnv, audio_frame: AudioFrame
    ) -> None:
        pass

    async def on_video_frame(
        self, ten_env: AsyncTenEnv, video_frame: VideoFrame
    ) -> None:
        pass

    async def _send_text(self, text: str, end_of_segment: bool) -> None:
        data = Data.create("text_data")
        data.set_property_string(DATA_OUT_TEXT_DATA_PROPERTY_TEXT, text)
        data.set_property_bool(
            DATA_OUT_TEXT_DATA_PROPERTY_END_OF_SEGMENT, end_of_segment
        )
        asyncio.create_task(self.ten_env.send_data(data))

    async def _stream_chat(
        self, messages: List[Any]
    ) -> AsyncGenerator[ChatEvent, None]:
        additionals = []
        for m in messages:
            if m["role"] == "user":
                additionals.append(
                    Message.build_user_question_text(m["content"]).model_dump()
                )
            elif m["role"] == "assistant":
                additionals.append(
                    Message.build_assistant_answer(m["content"]).model_dump()
                )

        def chat_stream_handler(event: str, event_data: Any) -> ChatEvent:
            if event == ChatEventType.DONE:
                raise StopAsyncIteration
            elif event == ChatEventType.ERROR:
                raise RuntimeError(f"error event: {event_data}")
            elif event in [
                ChatEventType.CONVERSATION_MESSAGE_DELTA,
                ChatEventType.CONVERSATION_MESSAGE_COMPLETED,
            ]:
                return ChatEvent(
                    event=event, message=Message.model_validate_json(event_data)
                )
            elif event in [
                ChatEventType.CONVERSATION_CHAT_CREATED,
                ChatEventType.CONVERSATION_CHAT_IN_PROGRESS,
                ChatEventType.CONVERSATION_CHAT_COMPLETED,
                ChatEventType.CONVERSATION_CHAT_FAILED,
                ChatEventType.CONVERSATION_CHAT_REQUIRES_ACTION,
            ]:
                return ChatEvent(event=event, chat=Chat.model_validate_json(event_data))
            else:
                raise ValueError(f"invalid chat.event: {event}, {event_data}")

        async with aiohttp.ClientSession() as session:
            try:
                url = f"{self.config.base_url}/v3/chat"
                headers = {
                    "Authorization": f"Bearer {self.config.token}",
                }
                params = {
                    "bot_id": self.config.bot_id,
                    "user_id": self.config.user_id,
                    "additional_messages": additionals,
                    "stream": True,
                    "auto_save_history": True,
                    # "conversation_id": self.conversation.id
                }
                event = ""
                async with session.post(url, json=params, headers=headers) as response:
                    async for line in response.content:
                        if line:
                            try:
                                self.ten_env.log_info(f"line: {line}")
                                decoded_line = line.decode("utf-8").strip()
                                if decoded_line:
                                    if decoded_line.startswith("data:"):
                                        data = decoded_line[5:].strip()
                                        yield chat_stream_handler(
                                            event=event, event_data=data.strip()
                                        )
                                    elif decoded_line.startswith("event:"):
                                        event = decoded_line[6:]
                                        self.ten_env.log_info(f"event: {event}")
                                        if event == "done":
                                            break
                                    else:
                                        result = json.loads(decoded_line)
                                        code = result.get("code", 0)
                                        if code == 4000:
                                            await self._send_text(
                                                "Coze bot is not published.", True
                                            )
                                        else:
                                            self.ten_env.log_error(
                                                f"Failed to stream chat: {result['code']}"
                                            )
                                            await self._send_text(
                                                "Coze bot is not connected. Please check your configuration.",
                                                True,
                                            )
                            except Exception as e:
                                self.ten_env.log_error(f"Failed to stream chat: {e}")
            except Exception as e:
                traceback.print_exc()
                self.ten_env.log_error(f"Failed to stream chat: {e}")
            finally:
                await session.close()