zhaolj
122 lines
4.5 KiB
Python
122 lines
4.5 KiB
Python
import json
|
|
from astrbot.core import logger
|
|
from aiohttp import ClientSession
|
|
from typing import Dict, List, Any, AsyncGenerator
|
|
|
|
|
|
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()
|
|
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]] = [],
|
|
timeout: float = 60,
|
|
) -> AsyncGenerator[Dict[str, Any], None]:
|
|
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"chat_messages 请求失败：{resp.status}. {text}")
|
|
|
|
buffer = ""
|
|
while True:
|
|
# 保持原有的8192字节限制，防止数据过大导致高水位报错
|
|
chunk = await resp.content.read(8192)
|
|
if not chunk:
|
|
break
|
|
|
|
buffer += chunk.decode('utf-8')
|
|
blocks = buffer.split('\n\n')
|
|
|
|
# 处理完整的数据块
|
|
for block in blocks[:-1]:
|
|
if block.strip() and block.startswith('data:'):
|
|
try:
|
|
json_str = block[5:] # 移除 "data:" 前缀
|
|
json_obj = json.loads(json_str)
|
|
yield json_obj
|
|
except json.JSONDecodeError as e:
|
|
logger.error(f"JSON解析错误: {str(e)}")
|
|
logger.error(f"原始数据块: {json_str}")
|
|
|
|
# 保留最后一个可能不完整的块
|
|
buffer = blocks[-1] if blocks else ""
|
|
|
|
async def workflow_run(
|
|
self,
|
|
inputs: Dict,
|
|
user: str,
|
|
response_mode: str = "streaming",
|
|
files: List[Dict[str, Any]] = [],
|
|
timeout: float = 60,
|
|
):
|
|
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"workflow_run 请求失败：{resp.status}. {text}")
|
|
|
|
buffer = ""
|
|
while True:
|
|
# 保持原有的8192字节限制，防止数据过大导致高水位报错
|
|
chunk = await resp.content.read(8192)
|
|
if not chunk:
|
|
break
|
|
|
|
buffer += chunk.decode('utf-8')
|
|
blocks = buffer.split('\n\n')
|
|
|
|
# 处理完整的数据块
|
|
for block in blocks[:-1]:
|
|
if block.strip() and block.startswith('data:'):
|
|
try:
|
|
json_str = block[5:] # 移除 "data:" 前缀
|
|
json_obj = json.loads(json_str)
|
|
yield json_obj
|
|
except json.JSONDecodeError as e:
|
|
logger.error(f"JSON解析错误: {str(e)}")
|
|
logger.error(f"原始数据块: {json_str}")
|
|
|
|
# 保留最后一个可能不完整的块
|
|
buffer = blocks[-1] if blocks else ""
|
|
|
|
async def file_upload(
|
|
self,
|
|
file_path: str,
|
|
user: str,
|
|
) -> Dict[str, Any]:
|
|
url = f"{self.api_base}/files/upload"
|
|
payload = {
|
|
"user": user,
|
|
"file": open(file_path, "rb"),
|
|
}
|
|
async with self.session.post(
|
|
url, data=payload, headers=self.headers
|
|
) as resp:
|
|
return await resp.json() # {"id": "xxx", ...}
|
|
|
|
async def close(self):
|
|
await self.session.close() |