AstrBot 3.5.6__py3-none-any.whl → 4.7.0__py3-none-any.whl

astrbot/core/agent/runners/coze/coze_agent_runner.py

@@ -0,0 +1,367 @@
+import base64
+import json
+import sys
+import typing as T
+
+import astrbot.core.message.components as Comp
+from astrbot import logger
+from astrbot.core import sp
+from astrbot.core.message.message_event_result import MessageChain
+from astrbot.core.provider.entities import (
+    LLMResponse,
+    ProviderRequest,
+)
+
+from ...hooks import BaseAgentRunHooks
+from ...response import AgentResponseData
+from ...run_context import ContextWrapper, TContext
+from ..base import AgentResponse, AgentState, BaseAgentRunner
+from .coze_api_client import CozeAPIClient
+
+if sys.version_info >= (3, 12):
+    from typing import override
+else:
+    from typing_extensions import override
+
+
+class CozeAgentRunner(BaseAgentRunner[TContext]):
+    """Coze Agent Runner"""
+
+    @override
+    async def reset(
+        self,
+        request: ProviderRequest,
+        run_context: ContextWrapper[TContext],
+        agent_hooks: BaseAgentRunHooks[TContext],
+        provider_config: dict,
+        **kwargs: T.Any,
+    ) -> None:
+        self.req = request
+        self.streaming = kwargs.get("streaming", False)
+        self.final_llm_resp = None
+        self._state = AgentState.IDLE
+        self.agent_hooks = agent_hooks
+        self.run_context = run_context
+
+        self.api_key = provider_config.get("coze_api_key", "")
+        if not self.api_key:
+            raise Exception("Coze API Key 不能为空。")
+        self.bot_id = provider_config.get("bot_id", "")
+        if not self.bot_id:
+            raise Exception("Coze Bot ID 不能为空。")
+        self.api_base: str = provider_config.get("coze_api_base", "https://api.coze.cn")
+
+        if not isinstance(self.api_base, str) or not self.api_base.startswith(
+            ("http://", "https://"),
+        ):
+            raise Exception(
+                "Coze API Base URL 格式不正确，必须以 http:// 或 https:// 开头。",
+            )
+
+        self.timeout = provider_config.get("timeout", 120)
+        if isinstance(self.timeout, str):
+            self.timeout = int(self.timeout)
+        self.auto_save_history = provider_config.get("auto_save_history", True)
+
+        # 创建 API 客户端
+        self.api_client = CozeAPIClient(api_key=self.api_key, api_base=self.api_base)
+
+        # 会话相关缓存
+        self.file_id_cache: dict[str, dict[str, str]] = {}
+
+    @override
+    async def step(self):
+        """
+        执行 Coze Agent 的一个步骤
+        """
+        if not self.req:
+            raise ValueError("Request is not set. Please call reset() first.")
+
+        if self._state == AgentState.IDLE:
+            try:
+                await self.agent_hooks.on_agent_begin(self.run_context)
+            except Exception as e:
+                logger.error(f"Error in on_agent_begin hook: {e}", exc_info=True)
+
+        # 开始处理，转换到运行状态
+        self._transition_state(AgentState.RUNNING)
+
+        try:
+            # 执行 Coze 请求并处理结果
+            async for response in self._execute_coze_request():
+                yield response
+        except Exception as e:
+            logger.error(f"Coze 请求失败：{str(e)}")
+            self._transition_state(AgentState.ERROR)
+            self.final_llm_resp = LLMResponse(
+                role="err", completion_text=f"Coze 请求失败：{str(e)}"
+            )
+            yield AgentResponse(
+                type="err",
+                data=AgentResponseData(
+                    chain=MessageChain().message(f"Coze 请求失败：{str(e)}")
+                ),
+            )
+        finally:
+            await self.api_client.close()
+
+    @override
+    async def step_until_done(
+        self, max_step: int = 30
+    ) -> T.AsyncGenerator[AgentResponse, None]:
+        while not self.done():
+            async for resp in self.step():
+                yield resp
+
+    async def _execute_coze_request(self):
+        """执行 Coze 请求的核心逻辑"""
+        prompt = self.req.prompt or ""
+        session_id = self.req.session_id or "unknown"
+        image_urls = self.req.image_urls or []
+        contexts = self.req.contexts or []
+        system_prompt = self.req.system_prompt
+
+        # 用户ID参数
+        user_id = session_id
+
+        # 获取或创建会话ID
+        conversation_id = await sp.get_async(
+            scope="umo",
+            scope_id=user_id,
+            key="coze_conversation_id",
+            default="",
+        )
+
+        # 构建消息
+        additional_messages = []
+
+        if system_prompt:
+            if not self.auto_save_history or not conversation_id:
+                additional_messages.append(
+                    {
+                        "role": "system",
+                        "content": system_prompt,
+                        "content_type": "text",
+                    },
+                )
+
+        # 处理历史上下文
+        if not self.auto_save_history and contexts:
+            for ctx in contexts:
+                if isinstance(ctx, dict) and "role" in ctx and "content" in ctx:
+                    # 处理上下文中的图片
+                    content = ctx["content"]
+                    if isinstance(content, list):
+                        # 多模态内容，需要处理图片
+                        processed_content = []
+                        for item in content:
+                            if isinstance(item, dict):
+                                if item.get("type") == "text":
+                                    processed_content.append(item)
+                                elif item.get("type") == "image_url":
+                                    # 处理图片上传
+                                    try:
+                                        image_data = item.get("image_url", {})
+                                        url = image_data.get("url", "")
+                                        if url:
+                                            file_id = (
+                                                await self._download_and_upload_image(
+                                                    url, session_id
+                                                )
+                                            )
+                                            processed_content.append(
+                                                {
+                                                    "type": "file",
+                                                    "file_id": file_id,
+                                                    "file_url": url,
+                                                }
+                                            )
+                                    except Exception as e:
+                                        logger.warning(f"处理上下文图片失败: {e}")
+                                        continue
+
+                        if processed_content:
+                            additional_messages.append(
+                                {
+                                    "role": ctx["role"],
+                                    "content": processed_content,
+                                    "content_type": "object_string",
+                                }
+                            )
+                    else:
+                        # 纯文本内容
+                        additional_messages.append(
+                            {
+                                "role": ctx["role"],
+                                "content": content,
+                                "content_type": "text",
+                            }
+                        )
+
+        # 构建当前消息
+        if prompt or image_urls:
+            if image_urls:
+                # 多模态
+                object_string_content = []
+                if prompt:
+                    object_string_content.append({"type": "text", "text": prompt})
+
+                for url in image_urls:
+                    # the url is a base64 string
+                    try:
+                        image_data = base64.b64decode(url)
+                        file_id = await self.api_client.upload_file(image_data)
+                        object_string_content.append(
+                            {
+                                "type": "image",
+                                "file_id": file_id,
+                            }
+                        )
+                    except Exception as e:
+                        logger.warning(f"处理图片失败 {url}: {e}")
+                        continue
+
+                if object_string_content:
+                    content = json.dumps(object_string_content, ensure_ascii=False)
+                    additional_messages.append(
+                        {
+                            "role": "user",
+                            "content": content,
+                            "content_type": "object_string",
+                        }
+                    )
+            elif prompt:
+                # 纯文本
+                additional_messages.append(
+                    {
+                        "role": "user",
+                        "content": prompt,
+                        "content_type": "text",
+                    },
+                )
+
+        # 执行 Coze API 请求
+        accumulated_content = ""
+        message_started = False
+
+        async for chunk in self.api_client.chat_messages(
+            bot_id=self.bot_id,
+            user_id=user_id,
+            additional_messages=additional_messages,
+            conversation_id=conversation_id,
+            auto_save_history=self.auto_save_history,
+            stream=True,
+            timeout=self.timeout,
+        ):
+            event_type = chunk.get("event")
+            data = chunk.get("data", {})
+
+            if event_type == "conversation.chat.created":
+                if isinstance(data, dict) and "conversation_id" in data:
+                    await sp.put_async(
+                        scope="umo",
+                        scope_id=user_id,
+                        key="coze_conversation_id",
+                        value=data["conversation_id"],
+                    )
+
+            if event_type == "conversation.message.delta":
+                # 增量消息
+                content = data.get("content", "")
+                if not content and "delta" in data:
+                    content = data["delta"].get("content", "")
+                if not content and "text" in data:
+                    content = data.get("text", "")
+
+                if content:
+                    accumulated_content += content
+                    message_started = True
+
+                    # 如果是流式响应，发送增量数据
+                    if self.streaming:
+                        yield AgentResponse(
+                            type="streaming_delta",
+                            data=AgentResponseData(
+                                chain=MessageChain().message(content)
+                            ),
+                        )
+
+            elif event_type == "conversation.message.completed":
+                # 消息完成
+                logger.debug("Coze message completed")
+                message_started = True
+
+            elif event_type == "conversation.chat.completed":
+                # 对话完成
+                logger.debug("Coze chat completed")
+                break
+
+            elif event_type == "error":
+                # 错误处理
+                error_msg = data.get("msg", "未知错误")
+                error_code = data.get("code", "UNKNOWN")
+                logger.error(f"Coze 出现错误: {error_code} - {error_msg}")
+                raise Exception(f"Coze 出现错误: {error_code} - {error_msg}")
+
+        if not message_started and not accumulated_content:
+            logger.warning("Coze 未返回任何内容")
+            accumulated_content = ""
+
+        # 创建最终响应
+        chain = MessageChain(chain=[Comp.Plain(accumulated_content)])
+        self.final_llm_resp = LLMResponse(role="assistant", result_chain=chain)
+        self._transition_state(AgentState.DONE)
+
+        try:
+            await self.agent_hooks.on_agent_done(self.run_context, self.final_llm_resp)
+        except Exception as e:
+            logger.error(f"Error in on_agent_done hook: {e}", exc_info=True)
+
+        # 返回最终结果
+        yield AgentResponse(
+            type="llm_result",
+            data=AgentResponseData(chain=chain),
+        )
+
+    async def _download_and_upload_image(
+        self,
+        image_url: str,
+        session_id: str | None = None,
+    ) -> str:
+        """下载图片并上传到 Coze，返回 file_id"""
+        import hashlib
+
+        # 计算哈希实现缓存
+        cache_key = hashlib.md5(image_url.encode("utf-8")).hexdigest()
+
+        if session_id:
+            if session_id not in self.file_id_cache:
+                self.file_id_cache[session_id] = {}
+
+            if cache_key in self.file_id_cache[session_id]:
+                file_id = self.file_id_cache[session_id][cache_key]
+                logger.debug(f"[Coze] 使用缓存的 file_id: {file_id}")
+                return file_id
+
+        try:
+            image_data = await self.api_client.download_image(image_url)
+            file_id = await self.api_client.upload_file(image_data)
+
+            if session_id:
+                self.file_id_cache[session_id][cache_key] = file_id
+                logger.debug(f"[Coze] 图片上传成功并缓存，file_id: {file_id}")
+
+            return file_id
+
+        except Exception as e:
+            logger.error(f"处理图片失败 {image_url}: {e!s}")
+            raise Exception(f"处理图片失败: {e!s}")
+
+    @override
+    def done(self) -> bool:
+        """检查 Agent 是否已完成工作"""
+        return self._state in (AgentState.DONE, AgentState.ERROR)
+
+    @override
+    def get_final_llm_resp(self) -> LLMResponse | None:
+        return self.final_llm_resp

astrbot/core/agent/runners/coze/coze_api_client.py

@@ -0,0 +1,324 @@
+import asyncio
+import io
+import json
+from collections.abc import AsyncGenerator
+from typing import Any
+
+import aiohttp
+
+from astrbot.core import logger
+
+
+class CozeAPIClient:
+    def __init__(self, api_key: str, api_base: str = "https://api.coze.cn"):
+        self.api_key = api_key
+        self.api_base = api_base
+        self.session = None
+
+    async def _ensure_session(self):
+        """确保HTTP session存在"""
+        if self.session is None:
+            connector = aiohttp.TCPConnector(
+                ssl=False if self.api_base.startswith("http://") else True,
+                limit=100,
+                limit_per_host=30,
+                keepalive_timeout=30,
+                enable_cleanup_closed=True,
+            )
+            timeout = aiohttp.ClientTimeout(
+                total=120,  # 默认超时时间
+                connect=30,
+                sock_read=120,
+            )
+            headers = {
+                "Authorization": f"Bearer {self.api_key}",
+                "Accept": "text/event-stream",
+            }
+            self.session = aiohttp.ClientSession(
+                headers=headers,
+                timeout=timeout,
+                connector=connector,
+            )
+        return self.session
+
+    async def upload_file(
+        self,
+        file_data: bytes,
+    ) -> str:
+        """上传文件到 Coze 并返回 file_id
+
+        Args:
+            file_data (bytes): 文件的二进制数据
+        Returns:
+            str: 上传成功后返回的 file_id
+
+        """
+        session = await self._ensure_session()
+        url = f"{self.api_base}/v1/files/upload"
+
+        try:
+            file_io = io.BytesIO(file_data)
+            async with session.post(
+                url,
+                data={
+                    "file": file_io,
+                },
+                timeout=aiohttp.ClientTimeout(total=60),
+            ) as response:
+                if response.status == 401:
+                    raise Exception("Coze API 认证失败，请检查 API Key 是否正确")
+
+                response_text = await response.text()
+                logger.debug(
+                    f"文件上传响应状态: {response.status}, 内容: {response_text}",
+                )
+
+                if response.status != 200:
+                    raise Exception(
+                        f"文件上传失败，状态码: {response.status}, 响应: {response_text}",
+                    )
+
+                try:
+                    result = await response.json()
+                except json.JSONDecodeError:
+                    raise Exception(f"文件上传响应解析失败: {response_text}")
+
+                if result.get("code") != 0:
+                    raise Exception(f"文件上传失败: {result.get('msg', '未知错误')}")
+
+                file_id = result["data"]["id"]
+                logger.debug(f"[Coze] 图片上传成功，file_id: {file_id}")
+                return file_id
+
+        except asyncio.TimeoutError:
+            logger.error("文件上传超时")
+            raise Exception("文件上传超时")
+        except Exception as e:
+            logger.error(f"文件上传失败: {e!s}")
+            raise Exception(f"文件上传失败: {e!s}")
+
+    async def download_image(self, image_url: str) -> bytes:
+        """下载图片并返回字节数据
+
+        Args:
+            image_url (str): 图片的URL
+        Returns:
+            bytes: 图片的二进制数据
+
+        """
+        session = await self._ensure_session()
+
+        try:
+            async with session.get(image_url) as response:
+                if response.status != 200:
+                    raise Exception(f"下载图片失败，状态码: {response.status}")
+
+                image_data = await response.read()
+                return image_data
+
+        except Exception as e:
+            logger.error(f"下载图片失败 {image_url}: {e!s}")
+            raise Exception(f"下载图片失败: {e!s}")
+
+    async def chat_messages(
+        self,
+        bot_id: str,
+        user_id: str,
+        additional_messages: list[dict] | None = None,
+        conversation_id: str | None = None,
+        auto_save_history: bool = True,
+        stream: bool = True,
+        timeout: float = 120,
+    ) -> AsyncGenerator[dict[str, Any], None]:
+        """发送聊天消息并返回流式响应
+
+        Args:
+            bot_id: Bot ID
+            user_id: 用户ID
+            additional_messages: 额外消息列表
+            conversation_id: 会话ID
+            auto_save_history: 是否自动保存历史
+            stream: 是否流式响应
+            timeout: 超时时间
+
+        """
+        session = await self._ensure_session()
+        url = f"{self.api_base}/v3/chat"
+
+        payload = {
+            "bot_id": bot_id,
+            "user_id": user_id,
+            "stream": stream,
+            "auto_save_history": auto_save_history,
+        }
+
+        if additional_messages:
+            payload["additional_messages"] = additional_messages
+
+        params = {}
+        if conversation_id:
+            params["conversation_id"] = conversation_id
+
+        logger.debug(f"Coze chat_messages payload: {payload}, params: {params}")
+
+        try:
+            async with session.post(
+                url,
+                json=payload,
+                params=params,
+                timeout=aiohttp.ClientTimeout(total=timeout),
+            ) as response:
+                if response.status == 401:
+                    raise Exception("Coze API 认证失败，请检查 API Key 是否正确")
+
+                if response.status != 200:
+                    raise Exception(f"Coze API 流式请求失败，状态码: {response.status}")
+
+                # SSE
+                buffer = ""
+                event_type = None
+                event_data = None
+
+                async for chunk in response.content:
+                    if chunk:
+                        buffer += chunk.decode("utf-8", errors="ignore")
+                        lines = buffer.split("\n")
+                        buffer = lines[-1]
+
+                        for line in lines[:-1]:
+                            line = line.strip()
+
+                            if not line:
+                                if event_type and event_data:
+                                    yield {"event": event_type, "data": event_data}
+                                    event_type = None
+                                    event_data = None
+                            elif line.startswith("event:"):
+                                event_type = line[6:].strip()
+                            elif line.startswith("data:"):
+                                data_str = line[5:].strip()
+                                if data_str and data_str != "[DONE]":
+                                    try:
+                                        event_data = json.loads(data_str)
+                                    except json.JSONDecodeError:
+                                        event_data = {"content": data_str}
+
+        except asyncio.TimeoutError:
+            raise Exception(f"Coze API 流式请求超时 ({timeout}秒)")
+        except Exception as e:
+            raise Exception(f"Coze API 流式请求失败: {e!s}")
+
+    async def clear_context(self, conversation_id: str):
+        """清空会话上下文
+
+        Args:
+            conversation_id: 会话ID
+        Returns:
+            dict: API响应结果
+
+        """
+        session = await self._ensure_session()
+        url = f"{self.api_base}/v3/conversation/message/clear_context"
+        payload = {"conversation_id": conversation_id}
+
+        try:
+            async with session.post(url, json=payload) as response:
+                response_text = await response.text()
+
+                if response.status == 401:
+                    raise Exception("Coze API 认证失败，请检查 API Key 是否正确")
+
+                if response.status != 200:
+                    raise Exception(f"Coze API 请求失败，状态码: {response.status}")
+
+                try:
+                    return json.loads(response_text)
+                except json.JSONDecodeError:
+                    raise Exception("Coze API 返回非JSON格式")
+
+        except asyncio.TimeoutError:
+            raise Exception("Coze API 请求超时")
+        except aiohttp.ClientError as e:
+            raise Exception(f"Coze API 请求失败: {e!s}")
+
+    async def get_message_list(
+        self,
+        conversation_id: str,
+        order: str = "desc",
+        limit: int = 10,
+        offset: int = 0,
+    ):
+        """获取消息列表
+
+        Args:
+            conversation_id: 会话ID
+            order: 排序方式 (asc/desc)
+            limit: 限制数量
+            offset: 偏移量
+        Returns:
+            dict: API响应结果
+
+        """
+        session = await self._ensure_session()
+        url = f"{self.api_base}/v3/conversation/message/list"
+        params = {
+            "conversation_id": conversation_id,
+            "order": order,
+            "limit": limit,
+            "offset": offset,
+        }
+
+        try:
+            async with session.get(url, params=params) as response:
+                response.raise_for_status()
+                return await response.json()
+
+        except Exception as e:
+            logger.error(f"获取Coze消息列表失败: {e!s}")
+            raise Exception(f"获取Coze消息列表失败: {e!s}")
+
+    async def close(self):
+        """关闭会话"""
+        if self.session:
+            await self.session.close()
+            self.session = None
+
+
+if __name__ == "__main__":
+    import asyncio
+    import os
+
+    async def test_coze_api_client():
+        api_key = os.getenv("COZE_API_KEY", "")
+        bot_id = os.getenv("COZE_BOT_ID", "")
+        client = CozeAPIClient(api_key=api_key)
+
+        try:
+            with open("README.md", "rb") as f:
+                file_data = f.read()
+            file_id = await client.upload_file(file_data)
+            print(f"Uploaded file_id: {file_id}")
+            async for event in client.chat_messages(
+                bot_id=bot_id,
+                user_id="test_user",
+                additional_messages=[
+                    {
+                        "role": "user",
+                        "content": json.dumps(
+                            [
+                                {"type": "text", "text": "这是什么"},
+                                {"type": "file", "file_id": file_id},
+                            ],
+                            ensure_ascii=False,
+                        ),
+                        "content_type": "object_string",
+                    },
+                ],
+                stream=True,
+            ):
+                print(f"Event: {event}")
+
+        finally:
+            await client.close()
+
+    asyncio.run(test_coze_api_client())