ai_generation/b.py

253 lines
8.3 KiB
Python
Raw Normal View History

2025-01-13 16:34:36 +08:00
from flask import Flask, request, jsonify
from llama_index.llms.ollama import Ollama
from llama_index.core.llms import MockLLM
import logging
import time
# 初始化日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
app = Flask(__name__)
# 定义任务映射
TASK_MAPPING = {
"1": ["文章内容", 400], # 表示文章
"2": ["文章标题", 20], # 表示标题
"3": ["文章摘要", 200], # 表示摘要
"4": ["文章关键词", 10], # 表示关键词
"5": ["文章描述", 50] # 表示问题描述
}
# 自定义 MockLLM 实现
class CustomMockLLM(MockLLM):
def complete(self, prompt):
# 模拟生成过程
time.sleep(1) # 模拟处理时间
return f"这是MockLLM生成的测试响应。\n原始输入: {prompt}"
# 统一的模型接口
class ModelInterface:
def __init__(self, model):
self.model = model
def generate_text(self, prompt):
if isinstance(self.model, Ollama):
response = self.model.complete(prompt)
return response.text
elif isinstance(self.model, CustomMockLLM):
return self.model.complete(prompt)
else:
raise ValueError("Unsupported model type")
# 初始化LLM
try:
logging.info("尝试初始化 Ollama 模型...")
raw_llm = Ollama(
model="qwen2.5",
temperature=0,
request_timeout=120.0,
base_url="http://localhost:11434"
)
logging.info("Ollama 模型初始化成功!")
except Exception as e:
logging.warning(f"无法连接到 Ollama 服务器,使用 MockLLM。错误: {e}")
raw_llm = CustomMockLLM()
# 创建统一的模型接口
llm = ModelInterface(raw_llm)
# 获取映射后的任务类型
def get_mapped_task(task_id):
return TASK_MAPPING.get(task_id, ["未知任务类型", 200]) # 默认返回未知任务类型和200字限制
# 调用模型生成内容
def generate_with_ollama(x, y, text):
source_type, _ = get_mapped_task(x)
target_type, word_limit = get_mapped_task(y)
prompt = (
f"你是一名专业的SEO内容策划专家。\n"
f"任务:根据以下{source_type}生成符合要求的高质量SEO{target_type}\n"
f"- 内容要求:清晰、连贯,准确概括核心内容。包含专业但易懂的科学解释,自然地融入关键词,不要出现特殊符号和表情符号,地名要在[]之中\n"
f"- 写作风格:专业但通俗易懂,适合一般读者阅读,保持科学准确性\n"
f"- 输出限制:仅返回生成的{target_type},不需要任何解释或补充内容。\n"
f"- 字数限制:{word_limit}字以内。\n"
f"- 语言要求:输出内容符合中文语言习惯。\n"
f"- 输出格式:请严格按照以下格式返回:\n"
f"{target_type}\n\n"
f"输入:\n{text}\n"
)
logging.info(f"生成 Prompt{prompt}")
try:
response = llm.generate_text(prompt)
logging.info(f"生成结果:{response}")
return {"result": response}
except Exception as e:
logging.error(f"生成失败:{e}")
return {"error": f"模型生成失败: {e}"}
# 调用模型优化内容
def optimize_with_ollama(x, text):
task_type, word_limit = get_mapped_task(x)
prompt = (
f"任务:对以下{task_type}进行优化。\n"
f"要求:\n"
f"1. 保持原文信息完整性,优化内容更清晰、逻辑更连贯。\n"
f"2. 适当简化句式,使内容更简洁流畅。\n"
f"3. 控制字数在{word_limit}字以内。\n"
f"4. 请直接输出优化后的{task_type}内容,不需要额外解释。\n"
f"\n输入:\n{text}\n"
)
logging.info(f"优化 Prompt{prompt}")
try:
response = llm.generate_text(prompt)
logging.info(f"优化结果:{response}")
return {"result": response}
except Exception as e:
logging.error(f"优化失败:{e}")
return {"error": f"模型优化失败: {e}"}
# 调用模型根据 x 和 y 进行优化
def optimize_with_target(x, y, text):
source_type, _ = get_mapped_task(x)
target_type, word_limit = get_mapped_task(y)
prompt = (
f"任务:将以下{source_type}优化为{target_type}\n"
f"要求:\n"
f"1. 确保内容符合{target_type}的表达特点,逻辑清晰,语言简洁。\n"
f"2. 控制内容在{word_limit}字以内。\n"
f"3. 请直接输出优化后的{target_type}内容,不要附加多余解释。\n"
f"\n输入:\n{text}\n"
)
logging.info(f"优化带目标类型Prompt{prompt}")
try:
response = llm.generate_text(prompt)
logging.info(f"优化结果:{response}")
return {"result": response}
except Exception as e:
logging.error(f"优化失败:{e}")
return {"error": f"模型优化失败: {e}"}
# 调用模型生成关键词
def generate_keywords(x, y, text):
task_type, word_limit = get_mapped_task(x)
# 构建关键词生成的 Prompt
prompt = (
f"任务:根据以下{task_type}提取关键词。\n"
f"要求:\n"
f"1. 提取{y}个关键词,每个关键词应准确概括内容核心思想。\n"
f"2. 关键词之间用逗号分隔。\n"
f"3. 请仅输出关键词列表,不要添加其他内容。\n"
f"\n输入:\n{text}\n"
)
logging.info(f"生成关键词 Prompt{prompt}")
try:
response = llm.generate_text(prompt)
logging.info(f"生成关键词结果:{response}")
return {"result": response}
except Exception as e:
logging.error(f"生成关键词失败:{e}")
return {"error": f"模型生成关键词失败: {e}"}
# 定义生成接口
@app.route('/generate', methods=['POST'])
def generate():
params = request.json
if not params:
return jsonify({"error": "缺少参数"}), 400
x = params.get("x")
y = params.get("y")
text = params.get("text")
if not x or not y or not text:
return jsonify({"error": "参数不完整"}), 400
if x not in TASK_MAPPING:
return jsonify({"error": f"无效的源任务类型: {x}"}), 400
if y not in TASK_MAPPING:
return jsonify({"error": f"无效的目标任务类型: {y}"}), 400
result = generate_with_ollama(x, y, text)
return jsonify(result)
# 定义优化接口
@app.route('/optimize', methods=['POST'])
def optimize():
params = request.json
if not params:
return jsonify({"error": "缺少参数"}), 400
x = params.get("x")
text = params.get("text")
if not x or not text:
return jsonify({"error": "参数不完整"}), 400
if x not in TASK_MAPPING:
return jsonify({"error": f"无效的任务类型: {x}"}), 400
result = optimize_with_ollama(x, text)
return jsonify(result)
# 定义优化带目标接口
@app.route('/optimize_with_target', methods=['POST'])
def optimize_with_target_api():
params = request.json
if not params:
return jsonify({"error": "缺少参数"}), 400
x = params.get("x")
y = params.get("y")
text = params.get("text")
if not x or not y or not text:
return jsonify({"error": "参数不完整"}), 400
if x not in TASK_MAPPING:
return jsonify({"error": f"无效的源任务类型: {x}"}), 400
if y not in TASK_MAPPING:
return jsonify({"error": f"无效的目标任务类型: {y}"}), 400
result = optimize_with_target(x, y, text)
return jsonify(result)
# 定义关键词生成接口
@app.route('/generate_keywords', methods=['POST'])
def generate_keywords_api():
params = request.json
if not params:
return jsonify({"error": "缺少参数"}), 400
x = params.get("x")
y = params.get("y")
text = params.get("text")
if not x or not y or not text:
return jsonify({"error": "参数不完整"}), 400
# 验证任务类型是否有效
if x not in TASK_MAPPING:
return jsonify({"error": f"无效的任务类型: {x}"}), 400
try:
y = int(y)
if y <= 0:
return jsonify({"error": "关键词数量必须为正整数"}), 400
except ValueError:
return jsonify({"error": "关键词数量必须为整数"}), 400
result = generate_keywords(x, y, text)
return jsonify(result)
if __name__ == '__main__':
app.run(host='0.0.0.0', port=5000)