250 lines
8.4 KiB
Python
250 lines
8.4 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
产品难度评分脚本
|
||
读取product_analysis表，增加难度评分字段，使用Ollama API进行智能评分
|
||
"""
|
||
|
||
import sqlite3
|
||
import os
|
||
import time
|
||
from typing import List, Tuple, Optional
|
||
from loguru import logger
|
||
import requests
|
||
import json
|
||
|
||
class DifficultyScorer:
|
||
"""产品难度评分器"""
|
||
|
||
def __init__(self, db_path: str = "products.db"):
|
||
"""
|
||
初始化评分器
|
||
|
||
Args:
|
||
db_path: 数据库文件路径
|
||
"""
|
||
self.db_path = db_path
|
||
self.api_url = "http://localhost:11434/api/generate"
|
||
|
||
# 检查数据库文件是否存在
|
||
if not os.path.exists(db_path):
|
||
current_dir_db = os.path.join(os.path.dirname(__file__), db_path)
|
||
if os.path.exists(current_dir_db):
|
||
self.db_path = current_dir_db
|
||
logger.info(f"使用当前目录下的数据库文件: {current_dir_db}")
|
||
else:
|
||
raise FileNotFoundError(f"数据库文件不存在: {db_path} 和 {current_dir_db}")
|
||
|
||
logger.info(f"初始化产品难度评分器，数据库: {self.db_path}")
|
||
|
||
def connect_to_database(self) -> sqlite3.Connection:
|
||
"""连接到SQLite数据库"""
|
||
try:
|
||
conn = sqlite3.connect(self.db_path)
|
||
logger.success(f"成功连接到数据库: {self.db_path}")
|
||
return conn
|
||
except Exception as e:
|
||
logger.error(f"连接数据库失败: {e}")
|
||
raise
|
||
|
||
def add_difficulty_score_column(self, conn: sqlite3.Connection):
|
||
"""添加难度评分字段"""
|
||
try:
|
||
cursor = conn.cursor()
|
||
|
||
# 检查字段是否已存在
|
||
cursor.execute("PRAGMA table_info(product_analysis)")
|
||
columns = [row[1] for row in cursor.fetchall()]
|
||
|
||
if 'difficulty_score' not in columns:
|
||
cursor.execute("ALTER TABLE product_analysis ADD COLUMN difficulty_score INTEGER")
|
||
conn.commit()
|
||
logger.success("成功添加difficulty_score字段")
|
||
else:
|
||
logger.info("difficulty_score字段已存在")
|
||
|
||
except Exception as e:
|
||
logger.error(f"添加难度评分字段失败: {e}")
|
||
raise
|
||
|
||
def get_unscored_products(self, conn: sqlite3.Connection) -> List[Tuple]:
|
||
"""
|
||
获取未评分的产品数据
|
||
|
||
Args:
|
||
conn: 数据库连接
|
||
|
||
Returns:
|
||
产品数据列表，每个元素为(id, ai_response)
|
||
"""
|
||
try:
|
||
cursor = conn.cursor()
|
||
|
||
# 查询未评分的产品
|
||
cursor.execute("""
|
||
SELECT id, ai_response
|
||
FROM product_analysis
|
||
WHERE difficulty_score IS NULL
|
||
AND ai_response IS NOT NULL
|
||
AND ai_response != ''
|
||
""")
|
||
|
||
products = cursor.fetchall()
|
||
logger.info(f"找到 {len(products)} 个未评分的产品")
|
||
|
||
return products
|
||
|
||
except Exception as e:
|
||
logger.error(f"获取未评分产品数据失败: {e}")
|
||
raise
|
||
|
||
def call_ollama_for_scoring(self, ai_response: str) -> Optional[int]:
|
||
"""
|
||
调用Ollama API进行难度评分
|
||
|
||
Args:
|
||
ai_response: AI响应内容
|
||
|
||
Returns:
|
||
评分（0-100），失败时返回None
|
||
"""
|
||
try:
|
||
# 构建评分提示
|
||
prompt = f"""
|
||
请根据以下产品开发难度描述，给出一个0-100分的难度评分：
|
||
|
||
难度描述：{ai_response}
|
||
|
||
评分标准：
|
||
- 90-100分：个人开发极其困难，需要大量专业知识和团队协作
|
||
- 70-89分：相对困难，需要较强的技术能力和较多时间
|
||
- 50-69分：中等难度，需要一定的技术基础
|
||
- 30-49分：相对简单，有基础即可开发
|
||
- 10-29分：非常简单，入门级别
|
||
- 0-9分：极其简单，几乎无难度
|
||
|
||
请只返回一个数字，不要有任何其他文字。
|
||
"""
|
||
|
||
data = {
|
||
"model": "qwen3:8b",
|
||
"prompt": prompt.strip(),
|
||
"stream": False
|
||
}
|
||
|
||
headers = {
|
||
"Content-Type": "application/json"
|
||
}
|
||
|
||
logger.info(f"调用Ollama API进行难度评分")
|
||
|
||
response = requests.post(
|
||
self.api_url,
|
||
headers=headers,
|
||
data=json.dumps(data, ensure_ascii=False),
|
||
timeout=60
|
||
)
|
||
|
||
if response.status_code == 200:
|
||
result = response.json()
|
||
score_text = result.get("response", "").strip()
|
||
|
||
# 尝试解析评分
|
||
try:
|
||
score = int(score_text)
|
||
# 确保评分在有效范围内
|
||
score = max(0, min(100, score))
|
||
logger.success(f"获得评分: {score}")
|
||
return score
|
||
except ValueError:
|
||
logger.error(f"无法解析评分: {score_text}")
|
||
return None
|
||
else:
|
||
logger.error(f"API调用失败: {response.status_code}, {response.text}")
|
||
return None
|
||
|
||
except Exception as e:
|
||
logger.error(f"调用Ollama API时出错: {e}")
|
||
return None
|
||
|
||
def update_difficulty_score(self, conn: sqlite3.Connection, product_id: int, score: int):
|
||
"""更新产品难度评分"""
|
||
try:
|
||
cursor = conn.cursor()
|
||
|
||
cursor.execute("""
|
||
UPDATE product_analysis
|
||
SET difficulty_score = ?
|
||
WHERE id = ?
|
||
""", (score, product_id))
|
||
|
||
conn.commit()
|
||
logger.success(f"更新产品ID {product_id} 的难度评分为: {score}")
|
||
|
||
except Exception as e:
|
||
logger.error(f"更新难度评分失败: {e}")
|
||
raise
|
||
|
||
def score_products(self):
|
||
"""评分所有未评分的产品"""
|
||
logger.info("开始产品难度评分")
|
||
|
||
conn = None
|
||
try:
|
||
# 连接数据库
|
||
conn = self.connect_to_database()
|
||
|
||
# 添加难度评分字段
|
||
self.add_difficulty_score_column(conn)
|
||
|
||
# 获取未评分的产品
|
||
products = self.get_unscored_products(conn)
|
||
|
||
if not products:
|
||
logger.info("没有需要评分的产品")
|
||
return
|
||
|
||
logger.info(f"准备评分 {len(products)} 个产品")
|
||
|
||
# 逐个评分
|
||
success_count = 0
|
||
for i, (product_id, ai_response) in enumerate(products, 1):
|
||
logger.info(f"\n评分进度: {i}/{len(products)} - 产品ID: {product_id}")
|
||
|
||
# 调用AI进行评分
|
||
score = self.call_ollama_for_scoring(ai_response)
|
||
|
||
if score is not None:
|
||
# 更新数据库
|
||
self.update_difficulty_score(conn, product_id, score)
|
||
success_count += 1
|
||
logger.success(f"评分完成: {score}分")
|
||
else:
|
||
logger.error(f"评分失败: 产品ID {product_id}")
|
||
|
||
# 延时避免API过载
|
||
logger.info("等待2秒后继续...")
|
||
time.sleep(2)
|
||
|
||
logger.success(f"评分完成! 成功评分 {success_count} 个产品")
|
||
|
||
except Exception as e:
|
||
logger.error(f"评分过程中出错: {e}")
|
||
finally:
|
||
if conn:
|
||
conn.close()
|
||
logger.info("数据库连接已关闭")
|
||
|
||
def main():
|
||
"""主函数"""
|
||
# 配置日志
|
||
logger.add("difficulty_scorer.log", rotation="10 MB", level="INFO")
|
||
|
||
# 创建评分器
|
||
scorer = DifficultyScorer()
|
||
|
||
# 开始评分
|
||
scorer.score_products()
|
||
|
||
if __name__ == "__main__":
|
||
main() |