207 lines
7.5 KiB
Python
207 lines
7.5 KiB
Python
"""
|
||
ParamHub LLM 智能解析模块
|
||
"""
|
||
import json
|
||
import base64
|
||
import requests
|
||
|
||
from config import CATEGORIES_FILE, IMAGES_DIR
|
||
from utils import load_data, get_llm_config
|
||
|
||
|
||
def get_parse_prompt_template(category_type, category_id=None, subcategory_id=None):
|
||
"""获取解析 prompt 模板（供前端显示和编辑）"""
|
||
categories = load_data(CATEGORIES_FILE)
|
||
|
||
if category_id:
|
||
cat = next((c for c in categories if c['id'] == category_id), None)
|
||
else:
|
||
type_to_cat_id = {'model': 'ai-models', 'gpu': 'gpus', 'cpu': 'cpus', 'dynamic': None}
|
||
cat_id = type_to_cat_id.get(category_type)
|
||
cat = next((c for c in categories if c['id'] == cat_id), None)
|
||
|
||
fields = _build_fields(cat, subcategory_id)
|
||
fields_json = json.dumps(fields, ensure_ascii=False, indent=2)
|
||
|
||
image_prompt = f"""请分析图片中的产品参数信息，提取结构化数据。
|
||
|
||
需要提取的字段：
|
||
{fields_json}
|
||
|
||
重要要求：
|
||
1. 图片中可能包含1个或多个产品，请识别所有产品
|
||
2. 如果是多张图片，请综合分析所有图片内容
|
||
3. **提取数据时保留原始单位**：字段标签中如有单位标注（如($)、(GB)、(MHz)等），提取时请带上对应单位，保持数据完整性
|
||
4. 如果某字段没有提及，返回null
|
||
5. 返回格式：如果识别到多个产品，返回数组 [对象列表]; 如果只有一个产品，返回单个对象
|
||
6. 只返回JSON数据，不要其他内容"""
|
||
|
||
return {
|
||
'fields': fields,
|
||
'fields_json': fields_json,
|
||
'image_prompt': image_prompt,
|
||
'category_name': cat.get('name', '') if cat else ''
|
||
}
|
||
|
||
|
||
def parse_with_llm(text, category_type, images=None, category_id=None,
|
||
subcategory_id=None, custom_prompt=None):
|
||
"""使用大模型解析文本/图片，提取结构化数据"""
|
||
categories = load_data(CATEGORIES_FILE)
|
||
|
||
if category_id:
|
||
cat = next((c for c in categories if c['id'] == category_id), None)
|
||
else:
|
||
type_to_cat_id = {'model': 'ai-models', 'gpu': 'gpus', 'cpu': 'cpus'}
|
||
cat_id = type_to_cat_id.get(category_type)
|
||
cat = next((c for c in categories if c['id'] == cat_id), None)
|
||
|
||
fields = _build_fields(cat, subcategory_id)
|
||
fields_json = json.dumps(fields, ensure_ascii=False, indent=2)
|
||
|
||
content_parts = []
|
||
|
||
if images and len(images) > 0:
|
||
if custom_prompt and custom_prompt.strip():
|
||
prompt_text = custom_prompt
|
||
else:
|
||
prompt_text = f"""请分析图片中的产品参数信息，提取结构化数据。
|
||
|
||
需要提取的字段：
|
||
{fields_json}
|
||
|
||
重要要求：
|
||
1. 图片中可能包含1个或多个产品，请识别所有产品
|
||
2. 如果是多张图片，请综合分析所有图片内容
|
||
3. **提取数据时保留原始单位**：字段标签中如有单位标注，提取时请带上对应单位
|
||
4. 如果某字段没有提及，返回null
|
||
5. 返回格式：如果识别到多个产品，返回数组; 如果只有一个产品，返回单个对象
|
||
6. 只返回JSON数据，不要其他内容"""
|
||
|
||
content_parts.append({"type": "text", "text": prompt_text})
|
||
|
||
for img in images:
|
||
if isinstance(img, str):
|
||
if img.startswith('http'):
|
||
content_parts.append({"type": "image_url", "image_url": {"url": img}})
|
||
elif img.startswith('data:'):
|
||
content_parts.append({"type": "image_url", "image_url": {"url": img}})
|
||
else:
|
||
b64 = _load_local_image(img)
|
||
if b64:
|
||
content_parts.append({"type": "image_url", "image_url": {"url": b64}})
|
||
else:
|
||
prompt_text = f"""请解析以下文本，提取结构化数据。
|
||
|
||
文本内容：
|
||
{text}
|
||
|
||
需要提取的字段：
|
||
{fields_json}
|
||
|
||
要求：
|
||
1. 根据文本内容智能提取各个字段的值
|
||
2. **提取数据时保留原始单位**
|
||
3. 如果某字段在文本中没有提及，返回null
|
||
4. 返回JSON格式，不要包含任何其他内容
|
||
|
||
请直接返回JSON数据："""
|
||
content_parts.append({"type": "text", "text": prompt_text})
|
||
|
||
try:
|
||
llm_config = get_llm_config()
|
||
model = llm_config.get('vision_model', 'gpt-4-vision-preview') if images else llm_config['model']
|
||
|
||
response = requests.post(
|
||
f"{llm_config['base_url']}/chat/completions",
|
||
headers={
|
||
"Content-Type": "application/json",
|
||
"Authorization": f"Bearer {llm_config['api_key']}"
|
||
},
|
||
json={
|
||
"model": model,
|
||
"messages": [
|
||
{"role": "system",
|
||
"content": "你是一个产品参数提取助手。只返回JSON，不要其他内容。"},
|
||
{"role": "user", "content": content_parts}
|
||
],
|
||
"max_tokens": 2000,
|
||
"temperature": 0.1
|
||
},
|
||
timeout=60
|
||
)
|
||
|
||
if response.status_code == 200:
|
||
data = response.json()
|
||
content = data['choices'][0]['message']['content'].strip()
|
||
|
||
if content.startswith('```'):
|
||
content = content.split('\n', 1)[1] if '\n' in content else content[3:]
|
||
content = content.rsplit('```', 1)[0] if '```' in content else content
|
||
|
||
parsed = json.loads(content)
|
||
results = parsed if isinstance(parsed, list) else [parsed]
|
||
|
||
return [_clean_result(item) for item in results]
|
||
except Exception as e:
|
||
print(f"LLM解析失败: {e}")
|
||
|
||
return [{'name': (text or '未命名产品')[:50], 'description': text}]
|
||
|
||
|
||
# ---- 内部函数 ----
|
||
|
||
def _build_fields(cat, subcategory_id):
|
||
if not cat or 'fields' not in cat:
|
||
return {
|
||
'name': '名称', 'brand': '品牌', 'price': '价格(数字)',
|
||
'year': '年份(数字)', 'specs': '规格参数(JSON对象)',
|
||
'description': '简介描述',
|
||
}
|
||
fields = {}
|
||
for field in cat['fields']:
|
||
desc = field['label']
|
||
desc += '(长文本)' if field.get('input_style') == 'long' else '(文本)'
|
||
if field.get('description'):
|
||
desc += f" - {field['description']}"
|
||
fields[field['key']] = desc
|
||
|
||
if subcategory_id:
|
||
subcat = next((s for s in cat.get('subcategories', []) if s['id'] == subcategory_id), None)
|
||
if subcat and 'extra_fields' in subcat:
|
||
for field in subcat['extra_fields']:
|
||
desc = field['label']
|
||
desc += '(长文本)' if field.get('input_style') == 'long' else '(文本)'
|
||
if field.get('description'):
|
||
desc += f" - {field['description']}"
|
||
fields[field['key']] = desc
|
||
return fields
|
||
|
||
|
||
def _load_local_image(img_src: str):
|
||
try:
|
||
img_path = IMAGES_DIR / img_src.replace('/static/uploads/', '')
|
||
if img_path.exists():
|
||
with open(img_path, 'rb') as f:
|
||
img_data = base64.b64encode(f.read()).decode()
|
||
ext = img_path.suffix.lower().lstrip('.')
|
||
mime = f'image/{"jpeg" if ext == "jpg" else ext}'
|
||
return f'data:{mime};base64,{img_data}'
|
||
except Exception:
|
||
pass
|
||
return None
|
||
|
||
|
||
def _clean_result(item: dict) -> dict:
|
||
cleaned = {}
|
||
for k, v in item.items():
|
||
if v is not None and v != '' and v != 'null':
|
||
if isinstance(v, str):
|
||
try:
|
||
cleaned[k] = float(v) if '.' in v else int(v)
|
||
except (ValueError, TypeError):
|
||
cleaned[k] = v
|
||
else:
|
||
cleaned[k] = v
|
||
return cleaned
|