607 lines
21 KiB
Python
607 lines
21 KiB
Python
#!/usr/bin/env python
|
||
# -*- coding: utf-8 -*-
|
||
|
||
import os
|
||
import json
|
||
import base64
|
||
import requests
|
||
from flask import Flask, render_template, request, jsonify, redirect, url_for, session, Response, stream_with_context
|
||
from flask_cors import CORS
|
||
from werkzeug.utils import secure_filename
|
||
from app.api.baidu_image_search import BaiduImageSearch
|
||
from app.api.image_utils import ImageUtils
|
||
from app.api.azure_openai import AzureOpenAI
|
||
from app.api.type_manager_mongo import TypeManagerMongo
|
||
from app.api.robot_manager import RobotManager
|
||
|
||
app = Flask(__name__, template_folder='app/templates', static_folder='app/static')
|
||
CORS(app) # 启用CORS支持跨域请求
|
||
|
||
# 设置会话密钥
|
||
app.secret_key = os.urandom(24)
|
||
|
||
# 配置上传文件夹
|
||
UPLOAD_FOLDER = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'uploads')
|
||
if not os.path.exists(UPLOAD_FOLDER):
|
||
os.makedirs(UPLOAD_FOLDER)
|
||
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
|
||
|
||
# 允许的文件扩展名
|
||
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif'}
|
||
|
||
# 初始化百度图像搜索API
|
||
image_search_api = BaiduImageSearch()
|
||
|
||
# 初始化Azure OpenAI API
|
||
try:
|
||
azure_openai_api = AzureOpenAI()
|
||
except ValueError as e:
|
||
print(f"警告: {str(e)}")
|
||
azure_openai_api = None
|
||
|
||
# 初始化类型管理器
|
||
type_manager = TypeManagerMongo()
|
||
|
||
# 初始化机器人角色管理器
|
||
robot_manager = RobotManager()
|
||
|
||
def allowed_file(filename):
|
||
"""检查文件扩展名是否允许"""
|
||
return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
|
||
|
||
@app.route('/imgsearcherApi/')
|
||
def index():
|
||
"""首页"""
|
||
return render_template('index.html')
|
||
|
||
@app.route('/imgsearcherApi/upload', methods=['POST'])
|
||
def upload_image():
|
||
"""上传图片到图库"""
|
||
if 'file' not in request.files:
|
||
return jsonify({'error': '没有文件上传'}), 400
|
||
|
||
file = request.files['file']
|
||
if file.filename == '':
|
||
return jsonify({'error': '没有选择文件'}), 400
|
||
|
||
if not allowed_file(file.filename):
|
||
return jsonify({'error': '不支持的文件类型'}), 400
|
||
|
||
# 获取表单数据
|
||
image_type = request.form.get('type', '')
|
||
description = request.form.get('description', '')
|
||
name = request.form.get('name', '')
|
||
tags = request.form.get('tags', '1') # 默认标签为1
|
||
|
||
# 将类型和描述信息保存到本地
|
||
type_manager.add_type(image_type, description)
|
||
|
||
# 创建brief信息（不包含描述信息）
|
||
brief = {
|
||
'type': image_type,
|
||
'name': name
|
||
}
|
||
|
||
# 保存文件
|
||
filename = secure_filename(file.filename)
|
||
file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
|
||
file.save(file_path)
|
||
|
||
try:
|
||
# 调用API添加图片
|
||
result = image_search_api.add_image(
|
||
image_path=file_path,
|
||
brief=brief,
|
||
tags=tags
|
||
)
|
||
|
||
# 添加本地存储路径到结果
|
||
result['file_path'] = file_path
|
||
|
||
return jsonify(result)
|
||
except Exception as e:
|
||
return jsonify({'error': str(e)}), 500
|
||
|
||
@app.route('/imgsearcherApi/search', methods=['POST'])
|
||
def search_image():
|
||
"""搜索相似图片"""
|
||
if 'file' not in request.files:
|
||
return jsonify({'error': '没有文件上传'}), 400
|
||
|
||
file = request.files['file']
|
||
if file.filename == '':
|
||
return jsonify({'error': '没有选择文件'}), 400
|
||
|
||
if not allowed_file(file.filename):
|
||
return jsonify({'error': '不支持的文件类型'}), 400
|
||
|
||
# 获取表单数据
|
||
image_type = request.form.get('type', '')
|
||
tags = request.form.get('tags', '1') # 默认标签为1
|
||
tag_logic = request.form.get('tag_logic', '0')
|
||
|
||
# 保存文件
|
||
filename = secure_filename(file.filename)
|
||
file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
|
||
file.save(file_path)
|
||
|
||
try:
|
||
# 调用API搜索图片
|
||
result = image_search_api.search_image(
|
||
image_path=file_path,
|
||
tags=tags,
|
||
tag_logic=tag_logic,
|
||
type_filter=image_type
|
||
)
|
||
|
||
# 从本地获取类型描述信息，并添加到结果中
|
||
if 'result' in result and result['result']:
|
||
for item in result['result']:
|
||
try:
|
||
brief = item.get('brief', '{}')
|
||
if isinstance(brief, str):
|
||
brief_dict = json.loads(brief)
|
||
item_type = brief_dict.get('type', '')
|
||
|
||
# 从本地获取描述信息
|
||
description = type_manager.get_description(item_type)
|
||
|
||
# 将描述信息添加到brief中
|
||
brief_dict['description'] = description
|
||
item['brief'] = json.dumps(brief_dict, ensure_ascii=False)
|
||
except Exception as e:
|
||
print(f"处理搜索结果项出错: {e}")
|
||
continue
|
||
|
||
return jsonify(result)
|
||
except Exception as e:
|
||
return jsonify({'error': str(e)}), 500
|
||
|
||
@app.route('/imgsearcherApi/delete', methods=['POST'])
|
||
def delete_image():
|
||
"""删除图库中的图片"""
|
||
data = request.get_json()
|
||
|
||
if not data or 'cont_sign' not in data:
|
||
return jsonify({'error': '缺少必要参数'}), 400
|
||
|
||
cont_sign = data['cont_sign']
|
||
|
||
try:
|
||
# 调用API删除图片
|
||
result = image_search_api.delete_image(cont_sign=cont_sign)
|
||
return jsonify(result)
|
||
except Exception as e:
|
||
return jsonify({'error': str(e)}), 500
|
||
|
||
@app.route('/imgsearcherApi/update', methods=['POST'])
|
||
def update_image():
|
||
"""更新图库中的图片信息"""
|
||
data = request.get_json()
|
||
|
||
if not data or 'cont_sign' not in data:
|
||
return jsonify({'error': '缺少必要参数'}), 400
|
||
|
||
cont_sign = data['cont_sign']
|
||
brief = data.get('brief')
|
||
tags = data.get('tags')
|
||
|
||
# 如果提供了brief信息，将类型和描述信息保存到本地
|
||
if brief and isinstance(brief, dict):
|
||
image_type = brief.get('type', '')
|
||
description = brief.get('description', '')
|
||
|
||
# 如果有描述信息，则更新到本地
|
||
if image_type and description:
|
||
type_manager.add_type(image_type, description)
|
||
|
||
# 从要上传的brief中移除描述信息
|
||
brief.pop('description', None)
|
||
|
||
try:
|
||
# 调用API更新图片
|
||
result = image_search_api.update_image(
|
||
cont_sign=cont_sign,
|
||
brief=brief,
|
||
tags=tags
|
||
)
|
||
return jsonify(result)
|
||
except Exception as e:
|
||
return jsonify({'error': str(e)}), 500
|
||
|
||
@app.route('/imgsearcherApi/api/token')
|
||
def get_token():
|
||
"""获取API访问令牌"""
|
||
try:
|
||
token = image_search_api.get_access_token()
|
||
return jsonify({'access_token': token})
|
||
except Exception as e:
|
||
return jsonify({'error': str(e)}), 500
|
||
|
||
@app.route('/imgsearcherApi/chat-with-image', methods=['GET'])
|
||
def chat_with_image_page():
|
||
"""图片对话页面"""
|
||
return render_template('chat.html')
|
||
|
||
@app.route('/imgsearcherApi/api/upload-chat-image', methods=['POST'])
|
||
def upload_chat_image():
|
||
"""上传图片用于对话"""
|
||
if 'file' not in request.files:
|
||
return jsonify({'error': '没有文件上传'}), 400
|
||
|
||
file = request.files['file']
|
||
if file.filename == '':
|
||
return jsonify({'error': '没有选择文件'}), 400
|
||
|
||
if not allowed_file(file.filename):
|
||
return jsonify({'error': '不支持的文件类型'}), 400
|
||
|
||
# 获取选择的机器人角色ID
|
||
robot_id = request.form.get('robot_id', '')
|
||
|
||
# 保存文件
|
||
filename = secure_filename(file.filename)
|
||
file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
|
||
file.save(file_path)
|
||
|
||
# 存储图片路径到会话
|
||
session['chat_image_path'] = file_path
|
||
|
||
# 存储机器人角色ID到会话
|
||
session['robot_id'] = robot_id
|
||
|
||
# 清空对话历史
|
||
session['conversation_history'] = []
|
||
|
||
try:
|
||
# 使用百度图像搜索API搜索图片
|
||
search_result = image_search_api.search_image(image_path=file_path)
|
||
|
||
# 使用方法三确定最可信的类型
|
||
reliable_types = image_search_api._determine_most_reliable_types(search_result.get('result', []))
|
||
image_type = reliable_types.get('method3', '')
|
||
|
||
# 从本地获取描述信息
|
||
description = type_manager.get_description(image_type)
|
||
|
||
# 如果本地没有描述信息，尝试从搜索结果中获取
|
||
if not description:
|
||
for item in search_result.get('result', []):
|
||
brief = item.get('brief', '{}')
|
||
if isinstance(brief, str):
|
||
try:
|
||
brief_dict = json.loads(brief)
|
||
if brief_dict.get('type') == image_type:
|
||
description = brief_dict.get('description', '')
|
||
if description:
|
||
# 将获取到的描述信息保存到本地
|
||
type_manager.add_type(image_type, description)
|
||
break
|
||
except:
|
||
continue
|
||
|
||
# 获取机器人角色信息
|
||
robot_info = None
|
||
if robot_id:
|
||
robot = robot_manager.get_robot(robot_id)
|
||
if robot:
|
||
robot_info = {
|
||
'name': robot.get('name', ''),
|
||
'background': robot.get('background', '')
|
||
}
|
||
|
||
return jsonify({
|
||
'success': True,
|
||
'image_path': file_path,
|
||
'image_type': image_type,
|
||
'description': description,
|
||
'robot': robot_info
|
||
})
|
||
except Exception as e:
|
||
return jsonify({'error': str(e)}), 500
|
||
|
||
@app.route('/imgsearcherApi/api/chat', methods=['POST'])
|
||
def chat():
|
||
"""与图片进行对话"""
|
||
if not azure_openai_api:
|
||
return jsonify({'error': 'Azure OpenAI API未正确配置'}), 500
|
||
|
||
data = request.get_json()
|
||
if not data or 'message' not in data:
|
||
return jsonify({'error': '缺少消息内容'}), 400
|
||
|
||
message = data['message']
|
||
image_path = session.get('chat_image_path')
|
||
robot_id = session.get('robot_id', '')
|
||
|
||
if not image_path or not os.path.exists(image_path):
|
||
return jsonify({'error': '没有上传图片或图片已失效'}), 400
|
||
|
||
# 获取对话历史
|
||
conversation_history = session.get('conversation_history', [])
|
||
|
||
# 获取机器人角色信息
|
||
robot_info = None
|
||
if robot_id:
|
||
robot = robot_manager.get_robot(robot_id)
|
||
if robot:
|
||
robot_info = {
|
||
'name': robot.get('name', ''),
|
||
'background': robot.get('background', '')
|
||
}
|
||
|
||
try:
|
||
# 调用Azure OpenAI API进行对话
|
||
response = azure_openai_api.chat_with_image(
|
||
image_path=image_path,
|
||
message=message,
|
||
conversation_history=conversation_history,
|
||
robot_info=robot_info
|
||
)
|
||
|
||
# 提取回复内容
|
||
reply = response['choices'][0]['message']['content']
|
||
|
||
# 更新对话历史
|
||
conversation_history.append({"role": "user", "content": message})
|
||
conversation_history.append({"role": "assistant", "content": reply})
|
||
session['conversation_history'] = conversation_history
|
||
|
||
return jsonify({
|
||
'reply': reply,
|
||
'conversation_history': conversation_history
|
||
})
|
||
except Exception as e:
|
||
return jsonify({'error': str(e)}), 500
|
||
|
||
@app.route('/imgsearcherApi/robots', methods=['GET'])
|
||
def robot_page():
|
||
"""机器人角色管理页面"""
|
||
return render_template('robots.html')
|
||
|
||
@app.route('/imgsearcherApi/api/robots', methods=['GET'])
|
||
def get_robots():
|
||
"""获取所有机器人角色"""
|
||
try:
|
||
robots = robot_manager.get_all_robots()
|
||
return jsonify({
|
||
'success': True,
|
||
'robots': robots
|
||
})
|
||
except Exception as e:
|
||
return jsonify({'error': str(e)}), 500
|
||
|
||
@app.route('/imgsearcherApi/api/robots', methods=['POST'])
|
||
def add_robot():
|
||
"""添加新的机器人角色"""
|
||
try:
|
||
name = request.form.get('name')
|
||
background = request.form.get('background')
|
||
|
||
if not name or not background:
|
||
return jsonify({'error': '机器人名称和背景故事不能为空'}), 400
|
||
|
||
avatar_file = None
|
||
if 'avatar' in request.files:
|
||
avatar_file = request.files['avatar']
|
||
if avatar_file.filename == '':
|
||
avatar_file = None
|
||
elif not allowed_file(avatar_file.filename):
|
||
return jsonify({'error': '不支持的文件类型'}), 400
|
||
|
||
result = robot_manager.add_robot(name, background, avatar_file)
|
||
|
||
if 'error' in result:
|
||
return jsonify(result), 400
|
||
|
||
return jsonify({
|
||
'success': True,
|
||
'robot': result
|
||
})
|
||
except Exception as e:
|
||
return jsonify({'error': str(e)}), 500
|
||
|
||
@app.route('/imgsearcherApi/api/robots/<robot_id>', methods=['GET'])
|
||
def get_robot(robot_id):
|
||
"""获取指定机器人角色"""
|
||
try:
|
||
robot = robot_manager.get_robot(robot_id)
|
||
if not robot:
|
||
return jsonify({'error': '找不到该机器人'}), 404
|
||
|
||
return jsonify({
|
||
'success': True,
|
||
'robot': robot
|
||
})
|
||
except Exception as e:
|
||
return jsonify({'error': str(e)}), 500
|
||
|
||
@app.route('/imgsearcherApi/api/robots/<robot_id>', methods=['PUT'])
|
||
def update_robot(robot_id):
|
||
"""更新机器人角色"""
|
||
try:
|
||
name = request.form.get('name')
|
||
background = request.form.get('background')
|
||
|
||
avatar_file = None
|
||
if 'avatar' in request.files:
|
||
avatar_file = request.files['avatar']
|
||
if avatar_file.filename == '':
|
||
avatar_file = None
|
||
elif not allowed_file(avatar_file.filename):
|
||
return jsonify({'error': '不支持的文件类型'}), 400
|
||
|
||
result = robot_manager.update_robot(robot_id, name, background, avatar_file)
|
||
|
||
if 'error' in result:
|
||
return jsonify(result), 400
|
||
|
||
return jsonify({
|
||
'success': True,
|
||
'robot': result
|
||
})
|
||
except Exception as e:
|
||
return jsonify({'error': str(e)}), 500
|
||
|
||
@app.route('/imgsearcherApi/api/robots/<robot_id>', methods=['DELETE'])
|
||
def delete_robot(robot_id):
|
||
"""删除机器人角色"""
|
||
try:
|
||
success = robot_manager.delete_robot(robot_id)
|
||
if not success:
|
||
return jsonify({'error': '找不到该机器人或删除失败'}), 404
|
||
|
||
return jsonify({
|
||
'success': True
|
||
})
|
||
except Exception as e:
|
||
return jsonify({'error': str(e)}), 500
|
||
|
||
# ==================== Chat Stream ====================
|
||
@app.route('/imgsearcherApi/api/chat-stream', methods=['POST'])
|
||
def chat_stream():
|
||
"""流式返回文本增量 + 最终语音 (SSE)"""
|
||
if not azure_openai_api:
|
||
return jsonify({'error': 'Azure OpenAI API未正确配置'}), 500
|
||
|
||
data = request.get_json()
|
||
if not data or 'message' not in data:
|
||
return jsonify({'error': '缺少消息内容'}), 400
|
||
|
||
message = data['message']
|
||
image_path = session.get('chat_image_path')
|
||
robot_id = session.get('robot_id', '')
|
||
if not image_path or not os.path.exists(image_path):
|
||
return jsonify({'error': '没有上传图片或图片已失效'}), 400
|
||
|
||
conversation_history = session.get('conversation_history', [])
|
||
robot_info = None
|
||
if robot_id:
|
||
robot = robot_manager.get_robot(robot_id)
|
||
if robot:
|
||
robot_info = {'name': robot.get('name', ''), 'background': robot.get('background', '')}
|
||
|
||
# 语音音色选择
|
||
voice_wav = data.get('voice') or 'zh-CN-XiaoXiao-Assistant-Audio.wav'
|
||
# 与Azure流式对话
|
||
openai_resp = azure_openai_api.chat_with_image_stream(
|
||
image_path=image_path,
|
||
message=message,
|
||
conversation_history=conversation_history,
|
||
robot_info=robot_info
|
||
)
|
||
|
||
def event_stream():
|
||
accumulated_text = ""
|
||
pending = ""
|
||
punctuation = "。！？.!?"
|
||
def send_tts(sentence):
|
||
if not sentence.strip():
|
||
return
|
||
tts_params = {
|
||
'tts_text': sentence,
|
||
'prompt_wav': voice_wav,
|
||
'text_split_method': 'cut5'
|
||
}
|
||
try:
|
||
tts_resp = requests.get(
|
||
"https://cloud.infini-ai.com/AIStudio/v1/inference/api/te-c7zfd4hrdoj2vqmw/tts/CosyVoice/v1/zero_shot",
|
||
params=tts_params,
|
||
headers={'Authorization': f'Bearer {os.getenv("INFINI_API_KEY", "sk-daooolufaf7ienn6")}'},
|
||
stream=True,
|
||
timeout=30
|
||
)
|
||
if tts_resp.status_code == 200:
|
||
audio_bytes = bytearray()
|
||
for chunk in tts_resp.iter_content(chunk_size=8192):
|
||
if chunk:
|
||
audio_bytes.extend(chunk)
|
||
if audio_bytes:
|
||
b64_audio = base64.b64encode(audio_bytes).decode()
|
||
yield f"data: {{\"type\": \"audio\", \"content\": \"{b64_audio}\" }}\n\n"
|
||
except Exception as e:
|
||
print("TTS 调用失败", e)
|
||
# 读取 GPT 流
|
||
for line in openai_resp.iter_lines():
|
||
if not line:
|
||
continue
|
||
try:
|
||
if line.strip() == b'data: [DONE]':
|
||
break
|
||
if line.startswith(b'data:'):
|
||
content_json = json.loads(line[5:].strip())
|
||
choices = content_json.get('choices', [])
|
||
if not choices:
|
||
continue
|
||
delta = choices[0].get('delta', {}).get('content', '')
|
||
if not delta:
|
||
continue
|
||
accumulated_text += delta
|
||
pending += delta
|
||
# 立即发送文本增量
|
||
yield f"data: {{\"type\": \"text\", \"content\": {json.dumps(delta, ensure_ascii=False)} }}\n\n"
|
||
# 检测标点
|
||
while True:
|
||
idx = -1
|
||
for i, ch in enumerate(pending):
|
||
if ch in punctuation:
|
||
idx = i
|
||
break
|
||
if idx == -1:
|
||
break
|
||
sentence = pending[:idx+1]
|
||
pending = pending[idx+1:]
|
||
# 调用 TTS 并流式发送
|
||
for audio_evt in send_tts(sentence):
|
||
yield audio_evt
|
||
except Exception as e:
|
||
print("解析OpenAI流错误", e)
|
||
continue
|
||
# 处理剩余 pending
|
||
if pending.strip():
|
||
for audio_evt in send_tts(pending):
|
||
yield audio_evt
|
||
pending = ""
|
||
# 更新会话历史
|
||
conversation_history.append({"role": "user", "content": message})
|
||
conversation_history.append({"role": "assistant", "content": accumulated_text})
|
||
session['conversation_history'] = conversation_history
|
||
yield "event: end\ndata: {}\n\n"
|
||
|
||
# 保持请求上下文，避免在生成器中操作 session 报错
|
||
return Response(stream_with_context(event_stream()), content_type='text/event-stream')
|
||
|
||
# ==================== TTS ====================
|
||
@app.route('/imgsearcherApi/api/tts', methods=['POST'])
|
||
def tts():
|
||
"""文本转语音接口，返回音频数据流"""
|
||
data = request.get_json()
|
||
if not data or 'text' not in data:
|
||
return jsonify({'error': '缺少text参数'}), 400
|
||
|
||
text = data['text']
|
||
prompt_text = data.get('prompt_text') or "我是威震天，我只代表月亮消灭你"
|
||
prompt_wav = data.get('voice') or data.get('prompt_wav') or "zh-CN-XiaoXiao-Assistant-Audio.wav"
|
||
|
||
tts_base_url = "https://cloud.infini-ai.com/AIStudio/v1/inference/api/te-c7zfd4hrdoj2vqmw/tts/CosyVoice/v1/zero_shot"
|
||
api_key = os.getenv('INFINI_API_KEY', 'sk-daooolufaf7ienn6')
|
||
headers = {'Authorization': f'Bearer {api_key}'}
|
||
params = {
|
||
'tts_text': text,
|
||
'prompt_wav': prompt_wav,
|
||
'text_split_method': 'cut5'
|
||
}
|
||
|
||
try:
|
||
r = requests.get(tts_base_url, params=params, headers=headers, timeout=20)
|
||
if r.status_code == 200:
|
||
return Response(r.content, content_type=r.headers.get('Content-Type', 'audio/wav'))
|
||
else:
|
||
print(f"TTS remote error {r.status_code}: {r.text}")
|
||
return jsonify({'error': f'TTS服务错误: {r.text}'}), 502
|
||
except Exception as e:
|
||
return jsonify({'error': f'TTS请求失败: {str(e)}'}), 500
|
||
|
||
if __name__ == '__main__':
|
||
app.run(debug=True, host='0.0.0.0', port=5001)
|