OpenDify/main.py
2025-04-10 18:07:24 +08:00

682 lines
28 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import json
import logging
import asyncio
from flask import Flask, request, Response, stream_with_context, jsonify
import httpx
import time
from dotenv import load_dotenv
import os
import ast
# 配置日志
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)
# 设置httpx的日志级别为WARNING减少不必要的输出
logging.getLogger("httpx").setLevel(logging.WARNING)
# 加载环境变量
load_dotenv()
# 从环境变量读取有效的API密钥逗号分隔
VALID_API_KEYS = [key.strip() for key in os.getenv("VALID_API_KEYS", "").split(",") if key]
# 获取会话记忆功能模式配置
# 0: 不开启会话记忆
# 1: 构造history_message附加到消息中的模式
# 2: 当前的零宽字符模式(默认)
CONVERSATION_MEMORY_MODE = int(os.getenv('CONVERSATION_MEMORY_MODE', '2'))
class DifyModelManager:
def __init__(self):
self.api_keys = []
self.name_to_api_key = {} # 应用名称到API Key的映射
self.api_key_to_name = {} # API Key到应用名称的映射
self.load_api_keys()
def load_api_keys(self):
"""从环境变量加载API Keys"""
api_keys_str = os.getenv('DIFY_API_KEYS', '')
if api_keys_str:
self.api_keys = [key.strip() for key in api_keys_str.split(',') if key.strip()]
logger.info(f"Loaded {len(self.api_keys)} API keys")
async def fetch_app_info(self, api_key):
"""获取Dify应用信息"""
try:
async with httpx.AsyncClient() as client:
headers = {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json"
}
response = await client.get(
f"{DIFY_API_BASE}/info",
headers=headers,
params={"user": "default_user"}
)
if response.status_code == 200:
app_info = response.json()
return app_info.get("name", "Unknown App")
else:
logger.error(f"Failed to fetch app info for API key: {api_key[:8]}...")
return None
except Exception as e:
logger.error(f"Error fetching app info: {str(e)}")
return None
async def refresh_model_info(self):
"""刷新所有应用信息"""
self.name_to_api_key.clear()
self.api_key_to_name.clear()
for api_key in self.api_keys:
app_name = await self.fetch_app_info(api_key)
if app_name:
self.name_to_api_key[app_name] = api_key
self.api_key_to_name[api_key] = app_name
logger.info(f"Mapped app '{app_name}' to API key: {api_key[:8]}...")
def get_api_key(self, model_name):
"""根据模型名称获取API Key"""
return self.name_to_api_key.get(model_name)
def get_available_models(self):
"""获取可用模型列表"""
return [
{
"id": name,
"object": "model",
"created": int(time.time()),
"owned_by": "dify"
}
for name in self.name_to_api_key.keys()
]
# 创建模型管理器实例
model_manager = DifyModelManager()
# 从环境变量获取API基础URL
DIFY_API_BASE = os.getenv("DIFY_API_BASE", "")
app = Flask(__name__)
def get_api_key(model_name):
"""根据模型名称获取对应的API密钥"""
api_key = model_manager.get_api_key(model_name)
if not api_key:
logger.warning(f"No API key found for model: {model_name}")
return api_key
def transform_openai_to_dify(openai_request, endpoint):
"""将OpenAI格式的请求转换为Dify格式"""
if endpoint == "/chat/completions":
messages = openai_request.get("messages", [])
stream = openai_request.get("stream", False)
# 尝试从历史消息中提取conversation_id
conversation_id = None
if CONVERSATION_MEMORY_MODE == 2: # 零宽字符模式
if len(messages) > 1:
# 遍历历史消息找到最近的assistant消息
for msg in reversed(messages[:-1]): # 除了最后一条消息
if msg.get("role") == "assistant":
content = msg.get("content", "")
# 尝试解码conversation_id
conversation_id = decode_conversation_id(content)
if conversation_id:
break
dify_request = {
"inputs": {},
"query": messages[-1]["content"] if messages else "",
"response_mode": "streaming" if stream else "blocking",
"conversation_id": conversation_id,
"user": openai_request.get("user", "default_user")
}
elif CONVERSATION_MEMORY_MODE == 1: # history_message模式
# 使用history_messages直接作为上下文
user_query = messages[-1]["content"] if messages else ""
if len(messages) > 1:
# 构造历史消息
history_messages = []
for msg in messages[:-1]: # 除了最后一条消息
role = msg.get("role", "")
content = msg.get("content", "")
if role and content:
history_messages.append(f"{role}: {content}")
# 将历史消息添加到查询中
if history_messages:
history_context = "\n\n".join(history_messages)
user_query = f"<history>\n{history_context}\n</history>\n\n用户当前问题: {user_query}"
dify_request = {
"inputs": {},
"query": user_query,
"response_mode": "streaming" if stream else "blocking",
"user": openai_request.get("user", "default_user")
}
else: # 不开启会话记忆
dify_request = {
"inputs": {},
"query": messages[-1]["content"] if messages else "",
"response_mode": "streaming" if stream else "blocking",
"user": openai_request.get("user", "default_user")
}
return dify_request
return None
def transform_dify_to_openai(dify_response, model="claude-3-5-sonnet-v2", stream=False):
"""将Dify格式的响应转换为OpenAI格式"""
if not stream:
answer = dify_response.get("answer", "")
# 只在零宽字符会话记忆模式时处理conversation_id
if CONVERSATION_MEMORY_MODE == 2:
conversation_id = dify_response.get("conversation_id", "")
history = dify_response.get("conversation_history", [])
# 检查历史消息中是否已经有会话ID
has_conversation_id = False
if history:
for msg in history:
if msg.get("role") == "assistant":
content = msg.get("content", "")
if decode_conversation_id(content) is not None:
has_conversation_id = True
break
# 只在新会话且历史消息中没有会话ID时插入
if conversation_id and not has_conversation_id:
logger.info(f"[Debug] Inserting conversation_id: {conversation_id}, history_length: {len(history)}")
encoded = encode_conversation_id(conversation_id)
answer = answer + encoded
logger.info(f"[Debug] Response content after insertion: {repr(answer)}")
return {
"id": dify_response.get("message_id", ""),
"object": "chat.completion",
"created": dify_response.get("created", int(time.time())),
"model": model,
"choices": [{
"index": 0,
"message": {
"role": "assistant",
"content": answer
},
"finish_reason": "stop"
}]
}
else:
# 流式响应的转换在stream_response函数中处理
return dify_response
def create_openai_stream_response(content, message_id, model="claude-3-5-sonnet-v2"):
"""创建OpenAI格式的流式响应"""
return {
"id": message_id,
"object": "chat.completion.chunk",
"created": int(time.time()),
"model": model,
"choices": [{
"index": 0,
"delta": {
"content": content
},
"finish_reason": None
}]
}
def encode_conversation_id(conversation_id):
"""将conversation_id编码为不可见的字符序列"""
if not conversation_id:
return ""
# 使用Base64编码减少长度
import base64
encoded = base64.b64encode(conversation_id.encode()).decode()
# 使用8种不同的零宽字符表示3位数字
# 这样可以将编码长度进一步减少
char_map = {
'0': '\u200b', # 零宽空格
'1': '\u200c', # 零宽非连接符
'2': '\u200d', # 零宽连接符
'3': '\ufeff', # 零宽非断空格
'4': '\u2060', # 词组连接符
'5': '\u180e', # 蒙古语元音分隔符
'6': '\u2061', # 函数应用
'7': '\u2062', # 不可见乘号
}
# 将Base64字符串转换为八进制数字
result = []
for c in encoded:
# 将每个字符转换为8进制数字0-7
if c.isalpha():
if c.isupper():
val = ord(c) - ord('A')
else:
val = ord(c) - ord('a') + 26
elif c.isdigit():
val = int(c) + 52
elif c == '+':
val = 62
elif c == '/':
val = 63
else: # '='
val = 0
# 每个Base64字符可以产生2个3位数字
first = (val >> 3) & 0x7
second = val & 0x7
result.append(char_map[str(first)])
if c != '=': # 不编码填充字符的后半部分
result.append(char_map[str(second)])
return ''.join(result)
def decode_conversation_id(content):
"""从消息内容中解码conversation_id"""
try:
# 零宽字符到3位数字的映射
char_to_val = {
'\u200b': '0', # 零宽空格
'\u200c': '1', # 零宽非连接符
'\u200d': '2', # 零宽连接符
'\ufeff': '3', # 零宽非断空格
'\u2060': '4', # 词组连接符
'\u180e': '5', # 蒙古语元音分隔符
'\u2061': '6', # 函数应用
'\u2062': '7', # 不可见乘号
}
# 提取最后一段零宽字符序列
space_chars = []
for c in reversed(content):
if c not in char_to_val:
break
space_chars.append(c)
if not space_chars:
return None
# 将零宽字符转换回Base64字符串
space_chars.reverse()
base64_chars = []
for i in range(0, len(space_chars), 2):
first = int(char_to_val[space_chars[i]], 8)
if i + 1 < len(space_chars):
second = int(char_to_val[space_chars[i + 1]], 8)
val = (first << 3) | second
else:
val = first << 3
# 转换回Base64字符
if val < 26:
base64_chars.append(chr(val + ord('A')))
elif val < 52:
base64_chars.append(chr(val - 26 + ord('a')))
elif val < 62:
base64_chars.append(str(val - 52))
elif val == 62:
base64_chars.append('+')
else:
base64_chars.append('/')
# 添加Base64填充
padding = len(base64_chars) % 4
if padding:
base64_chars.extend(['='] * (4 - padding))
# 解码Base64字符串
import base64
base64_str = ''.join(base64_chars)
return base64.b64decode(base64_str).decode()
except Exception as e:
logger.debug(f"Failed to decode conversation_id: {e}")
return None
@app.route('/v1/chat/completions', methods=['POST'])
def chat_completions():
try:
# 新增验证API密钥
auth_header = request.headers.get('Authorization')
if not auth_header:
return jsonify({
"error": {
"message": "Missing Authorization header",
"type": "invalid_request_error",
"param": None,
"code": "invalid_api_key"
}
}), 401
parts = auth_header.split()
if len(parts) != 2 or parts[0].lower() != 'bearer':
return jsonify({
"error": {
"message": "Invalid Authorization header format. Expected: Bearer <API_KEY>",
"type": "invalid_request_error",
"param": None,
"code": "invalid_api_key"
}
}), 401
provided_api_key = parts[1]
if provided_api_key not in VALID_API_KEYS:
return jsonify({
"error": {
"message": "Invalid API key",
"type": "invalid_request_error",
"param": None,
"code": "invalid_api_key"
}
}), 401
# 继续处理原始逻辑
openai_request = request.get_json()
logger.info(f"Received request: {json.dumps(openai_request, ensure_ascii=False)}")
model = openai_request.get("model", "claude-3-5-sonnet-v2")
# 验证模型是否支持
api_key = get_api_key(model)
if not api_key:
error_msg = f"Model {model} is not supported. Available models: {', '.join(model_manager.name_to_api_key.keys())}"
logger.error(error_msg)
return {
"error": {
"message": error_msg,
"type": "invalid_request_error",
"code": "model_not_found"
}
}, 404
dify_request = transform_openai_to_dify(openai_request, "/chat/completions")
if not dify_request:
logger.error("Failed to transform request")
return {
"error": {
"message": "Invalid request format",
"type": "invalid_request_error",
}
}, 400
headers = {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json"
}
stream = openai_request.get("stream", False)
dify_endpoint = f"{DIFY_API_BASE}/chat-messages"
logger.info(f"Sending request to Dify endpoint: {dify_endpoint}, stream={stream}")
if stream:
def generate():
client = httpx.Client(timeout=None)
def flush_chunk(chunk_data):
"""Helper function to flush chunks immediately"""
return chunk_data.encode('utf-8')
def calculate_delay(buffer_size):
"""
根据缓冲区大小动态计算延迟
buffer_size: 缓冲区中剩余的字符数量
"""
if buffer_size > 30: # 缓冲区内容较多,快速输出
return 0.001 # 5ms延迟
elif buffer_size > 20: # 中等数量,适中速度
return 0.002 # 10ms延迟
elif buffer_size > 10: # 较少内容,稍慢速度
return 0.01 # 20ms延迟
else: # 内容很少,使用较慢的速度
return 0.02 # 30ms延迟
def send_char(char, message_id):
"""Helper function to send single character"""
openai_chunk = {
"id": message_id,
"object": "chat.completion.chunk",
"created": int(time.time()),
"model": model,
"choices": [{
"index": 0,
"delta": {
"content": char
},
"finish_reason": None
}]
}
chunk_data = f"data: {json.dumps(openai_chunk)}\n\n"
return flush_chunk(chunk_data)
# 初始化缓冲区
output_buffer = []
try:
with client.stream(
'POST',
dify_endpoint,
json=dify_request,
headers={
**headers,
'Accept': 'text/event-stream',
'Cache-Control': 'no-cache',
'Connection': 'keep-alive'
}
) as response:
generate.message_id = None
buffer = ""
for raw_bytes in response.iter_raw():
if not raw_bytes:
continue
try:
buffer += raw_bytes.decode('utf-8')
while '\n' in buffer:
line, buffer = buffer.split('\n', 1)
line = line.strip()
if not line or not line.startswith('data: '):
continue
try:
json_str = line[6:]
dify_chunk = json.loads(json_str)
if dify_chunk.get("event") == "message" and "answer" in dify_chunk:
current_answer = dify_chunk["answer"]
if not current_answer:
continue
message_id = dify_chunk.get("message_id", "")
if not generate.message_id:
generate.message_id = message_id
# 将当前批次的字符添加到输出缓冲区
for char in current_answer:
output_buffer.append((char, generate.message_id))
# 根据缓冲区大小动态调整输出速度
while output_buffer:
char, msg_id = output_buffer.pop(0)
yield send_char(char, msg_id)
# 根据剩余缓冲区大小计算延迟
delay = calculate_delay(len(output_buffer))
time.sleep(delay)
# 立即继续处理下一个请求
continue
elif dify_chunk.get("event") == "message_end":
# 快速输出剩余内容
while output_buffer:
char, msg_id = output_buffer.pop(0)
yield send_char(char, msg_id)
time.sleep(0.001) # 固定使用最小延迟快速输出剩余内容
# 只在零宽字符会话记忆模式时处理conversation_id
if CONVERSATION_MEMORY_MODE == 2:
conversation_id = dify_chunk.get("conversation_id")
history = dify_chunk.get("conversation_history", [])
has_conversation_id = False
if history:
for msg in history:
if msg.get("role") == "assistant":
content = msg.get("content", "")
if decode_conversation_id(content) is not None:
has_conversation_id = True
break
# 只在新会话且历史消息中没有会话ID时插入
if conversation_id and not has_conversation_id:
logger.info(f"[Debug] Inserting conversation_id in stream: {conversation_id}")
encoded = encode_conversation_id(conversation_id)
logger.info(f"[Debug] Stream encoded content: {repr(encoded)}")
for char in encoded:
yield send_char(char, generate.message_id)
final_chunk = {
"id": generate.message_id,
"object": "chat.completion.chunk",
"created": int(time.time()),
"model": model,
"choices": [{
"index": 0,
"delta": {},
"finish_reason": "stop"
}]
}
yield flush_chunk(f"data: {json.dumps(final_chunk)}\n\n")
yield flush_chunk("data: [DONE]\n\n")
except json.JSONDecodeError as e:
logger.error(f"JSON decode error: {str(e)}")
continue
except Exception as e:
logger.error(f"Error processing chunk: {str(e)}")
continue
finally:
client.close()
return Response(
stream_with_context(generate()),
content_type='text/event-stream',
headers={
'Cache-Control': 'no-cache, no-transform',
'Connection': 'keep-alive',
'Transfer-Encoding': 'chunked',
'X-Accel-Buffering': 'no',
'Content-Encoding': 'none'
},
direct_passthrough=True
)
else:
async def sync_response():
try:
async with httpx.AsyncClient() as client:
response = await client.post(
dify_endpoint,
json=dify_request,
headers=headers
)
if response.status_code != 200:
error_msg = f"Dify API error: {response.text}"
logger.error(f"Request failed: {error_msg}")
return {
"error": {
"message": error_msg,
"type": "api_error",
"code": response.status_code
}
}, response.status_code
dify_response = response.json()
logger.info(f"Received response from Dify: {json.dumps(dify_response, ensure_ascii=False)}")
logger.info(f"[Debug] Response content: {repr(dify_response.get('answer', ''))}")
openai_response = transform_dify_to_openai(dify_response, model=model)
conversation_id = dify_response.get("conversation_id")
if conversation_id:
# 在响应头中传递conversation_id
return Response(
json.dumps(openai_response),
content_type='application/json',
headers={
'Conversation-Id': conversation_id
}
)
else:
return openai_response
except httpx.RequestError as e:
error_msg = f"Failed to connect to Dify: {str(e)}"
logger.error(error_msg)
return {
"error": {
"message": error_msg,
"type": "api_error",
"code": "connection_error"
}
}, 503
return asyncio.run(sync_response())
except Exception as e:
logger.exception("Unexpected error occurred")
return {
"error": {
"message": str(e),
"type": "internal_error",
}
}, 500
@app.route('/v1/models', methods=['GET'])
def list_models():
"""返回可用的模型列表"""
logger.info("Listing available models")
# 刷新模型信息
asyncio.run(model_manager.refresh_model_info())
# 获取可用模型列表
available_models = model_manager.get_available_models()
response = {
"object": "list",
"data": available_models
}
logger.info(f"Available models: {json.dumps(response, ensure_ascii=False)}")
return response
# 在main.py的最后初始化时添加环境变量检查
if __name__ == '__main__':
if not VALID_API_KEYS:
print("Warning: No API keys configured. Set the VALID_API_KEYS environment variable with comma-separated keys.")
# 启动时初始化模型信息
asyncio.run(model_manager.refresh_model_info())
host = os.getenv("SERVER_HOST", "127.0.0.1")
port = int(os.getenv("SERVER_PORT", 5000))
logger.info(f"Starting server on http://{host}:{port}")
app.run(debug=True, host=host, port=port)