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

astrbot/core/agent/runners/dify/dify_agent_runner.py

@@ -0,0 +1,336 @@
+import base64
+import os
+import sys
+import typing as T
+
+import astrbot.core.message.components as Comp
+from astrbot.core import logger, sp
+from astrbot.core.message.message_event_result import MessageChain
+from astrbot.core.provider.entities import (
+    LLMResponse,
+    ProviderRequest,
+)
+from astrbot.core.utils.astrbot_path import get_astrbot_data_path
+from astrbot.core.utils.io import download_file
+
+from ...hooks import BaseAgentRunHooks
+from ...response import AgentResponseData
+from ...run_context import ContextWrapper, TContext
+from ..base import AgentResponse, AgentState, BaseAgentRunner
+from .dify_api_client import DifyAPIClient
+
+if sys.version_info >= (3, 12):
+    from typing import override
+else:
+    from typing_extensions import override
+
+
+class DifyAgentRunner(BaseAgentRunner[TContext]):
+    """Dify 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("dify_api_key", "")
+        self.api_base = provider_config.get("dify_api_base", "https://api.dify.ai/v1")
+        self.api_type = provider_config.get("dify_api_type", "chat")
+        self.workflow_output_key = provider_config.get(
+            "dify_workflow_output_key",
+            "astrbot_wf_output",
+        )
+        self.dify_query_input_key = provider_config.get(
+            "dify_query_input_key",
+            "astrbot_text_query",
+        )
+        self.variables: dict = provider_config.get("variables", {}) or {}
+        self.timeout = provider_config.get("timeout", 60)
+        if isinstance(self.timeout, str):
+            self.timeout = int(self.timeout)
+
+        self.api_client = DifyAPIClient(self.api_key, self.api_base)
+
+    @override
+    async def step(self):
+        """
+        执行 Dify 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:
+            # 执行 Dify 请求并处理结果
+            async for response in self._execute_dify_request():
+                yield response
+        except Exception as e:
+            logger.error(f"Dify 请求失败：{str(e)}")
+            self._transition_state(AgentState.ERROR)
+            self.final_llm_resp = LLMResponse(
+                role="err", completion_text=f"Dify 请求失败：{str(e)}"
+            )
+            yield AgentResponse(
+                type="err",
+                data=AgentResponseData(
+                    chain=MessageChain().message(f"Dify 请求失败：{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_dify_request(self):
+        """执行 Dify 请求的核心逻辑"""
+        prompt = self.req.prompt or ""
+        session_id = self.req.session_id or "unknown"
+        image_urls = self.req.image_urls or []
+        system_prompt = self.req.system_prompt
+
+        conversation_id = await sp.get_async(
+            scope="umo",
+            scope_id=session_id,
+            key="dify_conversation_id",
+            default="",
+        )
+        result = ""
+
+        # 处理图片上传
+        files_payload = []
+        for image_url in image_urls:
+            # image_url is a base64 string
+            try:
+                image_data = base64.b64decode(image_url)
+                file_response = await self.api_client.file_upload(
+                    file_data=image_data,
+                    user=session_id,
+                    mime_type="image/png",
+                    file_name="image.png",
+                )
+                logger.debug(f"Dify 上传图片响应：{file_response}")
+                if "id" not in file_response:
+                    logger.warning(
+                        f"上传图片后得到未知的 Dify 响应：{file_response}，图片将忽略。"
+                    )
+                    continue
+                files_payload.append(
+                    {
+                        "type": "image",
+                        "transfer_method": "local_file",
+                        "upload_file_id": file_response["id"],
+                    }
+                )
+            except Exception as e:
+                logger.warning(f"上传图片失败：{e}")
+                continue
+
+        # 获得会话变量
+        payload_vars = self.variables.copy()
+        # 动态变量
+        session_var = await sp.get_async(
+            scope="umo",
+            scope_id=session_id,
+            key="session_variables",
+            default={},
+        )
+        payload_vars.update(session_var)
+        payload_vars["system_prompt"] = system_prompt
+
+        # 处理不同的 API 类型
+        match self.api_type:
+            case "chat" | "agent" | "chatflow":
+                if not prompt:
+                    prompt = "请描述这张图片。"
+
+                async for chunk in self.api_client.chat_messages(
+                    inputs={
+                        **payload_vars,
+                    },
+                    query=prompt,
+                    user=session_id,
+                    conversation_id=conversation_id,
+                    files=files_payload,
+                    timeout=self.timeout,
+                ):
+                    logger.debug(f"dify resp chunk: {chunk}")
+                    if chunk["event"] == "message" or chunk["event"] == "agent_message":
+                        result += chunk["answer"]
+                        if not conversation_id:
+                            await sp.put_async(
+                                scope="umo",
+                                scope_id=session_id,
+                                key="dify_conversation_id",
+                                value=chunk["conversation_id"],
+                            )
+                            conversation_id = chunk["conversation_id"]
+
+                        # 如果是流式响应，发送增量数据
+                        if self.streaming and chunk["answer"]:
+                            yield AgentResponse(
+                                type="streaming_delta",
+                                data=AgentResponseData(
+                                    chain=MessageChain().message(chunk["answer"])
+                                ),
+                            )
+                    elif chunk["event"] == "message_end":
+                        logger.debug("Dify message end")
+                        break
+                    elif chunk["event"] == "error":
+                        logger.error(f"Dify 出现错误：{chunk}")
+                        raise Exception(
+                            f"Dify 出现错误 status: {chunk['status']} message: {chunk['message']}"
+                        )
+
+            case "workflow":
+                async for chunk in self.api_client.workflow_run(
+                    inputs={
+                        self.dify_query_input_key: prompt,
+                        "astrbot_session_id": session_id,
+                        **payload_vars,
+                    },
+                    user=session_id,
+                    files=files_payload,
+                    timeout=self.timeout,
+                ):
+                    logger.debug(f"dify workflow resp chunk: {chunk}")
+                    match chunk["event"]:
+                        case "workflow_started":
+                            logger.info(
+                                f"Dify 工作流(ID: {chunk['workflow_run_id']})开始运行。"
+                            )
+                        case "node_finished":
+                            logger.debug(
+                                f"Dify 工作流节点(ID: {chunk['data']['node_id']} Title: {chunk['data'].get('title', '')})运行结束。"
+                            )
+                        case "text_chunk":
+                            if self.streaming and chunk["data"]["text"]:
+                                yield AgentResponse(
+                                    type="streaming_delta",
+                                    data=AgentResponseData(
+                                        chain=MessageChain().message(
+                                            chunk["data"]["text"]
+                                        )
+                                    ),
+                                )
+                        case "workflow_finished":
+                            logger.info(
+                                f"Dify 工作流(ID: {chunk['workflow_run_id']})运行结束"
+                            )
+                            logger.debug(f"Dify 工作流结果：{chunk}")
+                            if chunk["data"]["error"]:
+                                logger.error(
+                                    f"Dify 工作流出现错误：{chunk['data']['error']}"
+                                )
+                                raise Exception(
+                                    f"Dify 工作流出现错误：{chunk['data']['error']}"
+                                )
+                            if self.workflow_output_key not in chunk["data"]["outputs"]:
+                                raise Exception(
+                                    f"Dify 工作流的输出不包含指定的键名：{self.workflow_output_key}"
+                                )
+                            result = chunk
+            case _:
+                raise Exception(f"未知的 Dify API 类型：{self.api_type}")
+
+        if not result:
+            logger.warning("Dify 请求结果为空，请查看 Debug 日志。")
+
+        # 解析结果
+        chain = await self.parse_dify_result(result)
+
+        # 创建最终响应
+        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 parse_dify_result(self, chunk: dict | str) -> MessageChain:
+        """解析 Dify 的响应结果"""
+        if isinstance(chunk, str):
+            # Chat
+            return MessageChain(chain=[Comp.Plain(chunk)])
+
+        async def parse_file(item: dict):
+            match item["type"]:
+                case "image":
+                    return Comp.Image(file=item["url"], url=item["url"])
+                case "audio":
+                    # 仅支持 wav
+                    temp_dir = os.path.join(get_astrbot_data_path(), "temp")
+                    path = os.path.join(temp_dir, f"{item['filename']}.wav")
+                    await download_file(item["url"], path)
+                    return Comp.Image(file=item["url"], url=item["url"])
+                case "video":
+                    return Comp.Video(file=item["url"])
+                case _:
+                    return Comp.File(name=item["filename"], file=item["url"])
+
+        output = chunk["data"]["outputs"][self.workflow_output_key]
+        chains = []
+        if isinstance(output, str):
+            # 纯文本输出
+            chains.append(Comp.Plain(output))
+        elif isinstance(output, list):
+            # 主要适配 Dify 的 HTTP 请求结点的多模态输出
+            for item in output:
+                # handle Array[File]
+                if (
+                    not isinstance(item, dict)
+                    or item.get("dify_model_identity", "") != "__dify__file__"
+                ):
+                    chains.append(Comp.Plain(str(output)))
+                    break
+        else:
+            chains.append(Comp.Plain(str(output)))
+
+        # scan file
+        files = chunk["data"].get("files", [])
+        for item in files:
+            comp = await parse_file(item)
+            chains.append(comp)
+
+        return MessageChain(chain=chains)
+
+    @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/dify/dify_api_client.py

@@ -0,0 +1,195 @@
+import codecs
+import json
+from collections.abc import AsyncGenerator
+from typing import Any
+
+from aiohttp import ClientResponse, ClientSession, FormData
+
+from astrbot.core import logger
+
+
+async def _stream_sse(resp: ClientResponse) -> AsyncGenerator[dict, None]:
+    decoder = codecs.getincrementaldecoder("utf-8")()
+    buffer = ""
+    async for chunk in resp.content.iter_chunked(8192):
+        buffer += decoder.decode(chunk)
+        while "\n\n" in buffer:
+            block, buffer = buffer.split("\n\n", 1)
+            if block.strip().startswith("data:"):
+                try:
+                    yield json.loads(block[5:])
+                except json.JSONDecodeError:
+                    logger.warning(f"Drop invalid dify json data: {block[5:]}")
+                    continue
+    # flush any remaining text
+    buffer += decoder.decode(b"", final=True)
+    if buffer.strip().startswith("data:"):
+        try:
+            yield json.loads(buffer[5:])
+        except json.JSONDecodeError:
+            logger.warning(f"Drop invalid dify json data: {buffer[5:]}")
+
+
+class DifyAPIClient:
+    def __init__(self, api_key: str, api_base: str = "https://api.dify.ai/v1"):
+        self.api_key = api_key
+        self.api_base = api_base
+        self.session = ClientSession(trust_env=True)
+        self.headers = {
+            "Authorization": f"Bearer {self.api_key}",
+        }
+
+    async def chat_messages(
+        self,
+        inputs: dict,
+        query: str,
+        user: str,
+        response_mode: str = "streaming",
+        conversation_id: str = "",
+        files: list[dict[str, Any]] | None = None,
+        timeout: float = 60,
+    ) -> AsyncGenerator[dict[str, Any], None]:
+        if files is None:
+            files = []
+        url = f"{self.api_base}/chat-messages"
+        payload = locals()
+        payload.pop("self")
+        payload.pop("timeout")
+        logger.info(f"chat_messages payload: {payload}")
+        async with self.session.post(
+            url,
+            json=payload,
+            headers=self.headers,
+            timeout=timeout,
+        ) as resp:
+            if resp.status != 200:
+                text = await resp.text()
+                raise Exception(
+                    f"Dify /chat-messages 接口请求失败：{resp.status}. {text}",
+                )
+            async for event in _stream_sse(resp):
+                yield event
+
+    async def workflow_run(
+        self,
+        inputs: dict,
+        user: str,
+        response_mode: str = "streaming",
+        files: list[dict[str, Any]] | None = None,
+        timeout: float = 60,
+    ):
+        if files is None:
+            files = []
+        url = f"{self.api_base}/workflows/run"
+        payload = locals()
+        payload.pop("self")
+        payload.pop("timeout")
+        logger.info(f"workflow_run payload: {payload}")
+        async with self.session.post(
+            url,
+            json=payload,
+            headers=self.headers,
+            timeout=timeout,
+        ) as resp:
+            if resp.status != 200:
+                text = await resp.text()
+                raise Exception(
+                    f"Dify /workflows/run 接口请求失败：{resp.status}. {text}",
+                )
+            async for event in _stream_sse(resp):
+                yield event
+
+    async def file_upload(
+        self,
+        user: str,
+        file_path: str | None = None,
+        file_data: bytes | None = None,
+        file_name: str | None = None,
+        mime_type: str | None = None,
+    ) -> dict[str, Any]:
+        """Upload a file to Dify. Must provide either file_path or file_data.
+
+        Args:
+            user: The user ID.
+            file_path: The path to the file to upload.
+            file_data: The file data in bytes.
+            file_name: Optional file name when using file_data.
+        Returns:
+            A dictionary containing the uploaded file information.
+        """
+        url = f"{self.api_base}/files/upload"
+
+        form = FormData()
+        form.add_field("user", user)
+
+        if file_data is not None:
+            # 使用 bytes 数据
+            form.add_field(
+                "file",
+                file_data,
+                filename=file_name or "uploaded_file",
+                content_type=mime_type or "application/octet-stream",
+            )
+        elif file_path is not None:
+            # 使用文件路径
+            import os
+
+            with open(file_path, "rb") as f:
+                file_content = f.read()
+                form.add_field(
+                    "file",
+                    file_content,
+                    filename=os.path.basename(file_path),
+                    content_type=mime_type or "application/octet-stream",
+                )
+        else:
+            raise ValueError("file_path 和 file_data 不能同时为 None")
+
+        async with self.session.post(
+            url,
+            data=form,
+            headers=self.headers,  # 不包含 Content-Type，让 aiohttp 自动设置
+        ) as resp:
+            if resp.status != 200 and resp.status != 201:
+                text = await resp.text()
+                raise Exception(f"Dify 文件上传失败：{resp.status}. {text}")
+            return await resp.json()  # {"id": "xxx", ...}
+
+    async def close(self):
+        await self.session.close()
+
+    async def get_chat_convs(self, user: str, limit: int = 20):
+        # conversations. GET
+        url = f"{self.api_base}/conversations"
+        payload = {
+            "user": user,
+            "limit": limit,
+        }
+        async with self.session.get(url, params=payload, headers=self.headers) as resp:
+            return await resp.json()
+
+    async def delete_chat_conv(self, user: str, conversation_id: str):
+        # conversation. DELETE
+        url = f"{self.api_base}/conversations/{conversation_id}"
+        payload = {
+            "user": user,
+        }
+        async with self.session.delete(url, json=payload, headers=self.headers) as resp:
+            return await resp.json()
+
+    async def rename(
+        self,
+        conversation_id: str,
+        name: str,
+        user: str,
+        auto_generate: bool = False,
+    ):
+        # /conversations/:conversation_id/name
+        url = f"{self.api_base}/conversations/{conversation_id}/name"
+        payload = {
+            "user": user,
+            "name": name,
+            "auto_generate": auto_generate,
+        }
+        async with self.session.post(url, json=payload, headers=self.headers) as resp:
+            return await resp.json()