feat: 添加对 chat 和 workflow 的支持

2026-06-28 00:14:21 +00:00 · 2024-12-14 17:51:11 +08:00
parent 2ea3ff0b5c
commit dbf9f2398e
18 changed files with 1301 additions and 187 deletions
@@ -0,0 +1,125 @@
+from __future__ import annotations
+
+import httpx
+import typing
+import json
+
+from .errors import DifyAPIError
+
+
+class AsyncDifyServiceClient:
+    """Dify Service API 客户端"""
+
+    api_key: str
+    base_url: str
+    
+    def __init__(
+        self,
+        api_key: str,
+        base_url: str = "https://api.dify.ai/v1",
+    ) -> None:
+        self.api_key = api_key
+        self.base_url = base_url
+
+    async def chat_messages(
+        self,
+        inputs: dict[str, typing.Any],
+        query: str,
+        user: str,
+        response_mode: str = "blocking",  # 当前不支持 streaming
+        conversation_id: str = "",
+        files: list[dict[str, typing.Any]] = [],
+        timeout: float = 30.0,
+    ) -> dict[str, typing.Any]:
+        """发送消息"""
+        if response_mode != "blocking":
+            raise DifyAPIError("当前仅支持 blocking 模式")
+        
+        async with httpx.AsyncClient(
+            base_url=self.base_url,
+            trust_env=True,
+            timeout=timeout,
+        ) as client:
+            response = await client.post(
+                "/chat-messages",
+                headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"},
+                json={
+                    "inputs": inputs,
+                    "query": query,
+                    "user": user,
+                    "response_mode": response_mode,
+                    "conversation_id": conversation_id,
+                    "files": files,
+                },
+            )
+
+            if response.status_code != 200:
+                raise DifyAPIError(f"{response.status_code} {response.text}")
+
+            return response.json()
+        
+    async def workflow_run(
+        self,
+        inputs: dict[str, typing.Any],
+        user: str,
+        response_mode: str = "streaming",  # 当前不支持 blocking
+        files: list[dict[str, typing.Any]] = [],
+        timeout: float = 30.0,
+    ) -> typing.AsyncGenerator[dict[str, typing.Any], None]:
+        """运行工作流"""
+        if response_mode != "streaming":
+            raise DifyAPIError("当前仅支持 streaming 模式")
+        
+        async with httpx.AsyncClient(
+            base_url=self.base_url,
+            trust_env=True,
+            timeout=timeout,
+        ) as client:
+
+            async with client.stream(
+                "POST",
+                "/workflows/run",
+                headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"},
+                json={
+                    "inputs": inputs,
+                    "user": user,
+                    "response_mode": response_mode,
+                    "files": files,
+                },
+            ) as r:
+                async for chunk in r.aiter_lines():
+                    if chunk.strip() == "":
+                        continue
+                    if chunk.startswith("data:"):
+                        yield json.loads(chunk[5:])
+
+    async def upload_file(
+        self,
+        file: httpx._types.FileTypes,
+        user: str,
+        timeout: float = 30.0,
+    ) -> str:
+        """上传文件"""
+        # curl -X POST 'http://dify.rockchin.top/v1/files/upload' \
+        # --header 'Authorization: Bearer {api_key}' \
+        # --form 'file=@localfile;type=image/[png|jpeg|jpg|webp|gif] \
+        # --form 'user=abc-123'
+        async with httpx.AsyncClient(
+            base_url=self.base_url,
+            trust_env=True,
+            timeout=timeout,
+        ) as client:
+            # multipart/form-data
+            response = await client.post(
+                "/files/upload",
+                headers={"Authorization": f"Bearer {self.api_key}"},
+                files={
+                    "file": file,
+                    "user": (None, user),
+                },
+            )
+
+            if response.status_code != 201:
+                raise DifyAPIError(f"{response.status_code} {response.text}")
+
+            return response.json()