102 lines
3.0 KiB
Python
102 lines
3.0 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
AI 新闻分类与严重度判定
|
|
优先使用 Ollama 本地模型（免费），失败则回退到规则
|
|
设置 PARSER_AI_DISABLED=1 可只用规则（更快）
|
|
"""
|
|
import os
|
|
from typing import Literal, Optional, Tuple
|
|
|
|
Category = Literal["deployment", "alert", "intel", "diplomatic", "other"]
|
|
Severity = Literal["low", "medium", "high", "critical"]
|
|
|
|
PARSER_AI_DISABLED = os.environ.get("PARSER_AI_DISABLED", "0") == "1"
|
|
OLLAMA_MODEL = os.environ.get("OLLAMA_MODEL", "llama3.1") # 或 qwen2.5:7b
|
|
|
|
_CATEGORIES = ("deployment", "alert", "intel", "diplomatic", "other")
|
|
_SEVERITIES = ("low", "medium", "high", "critical")
|
|
|
|
|
|
def _parse_ai_response(text: str) -> Tuple[Category, Severity]:
|
|
"""从 AI 回复解析 category:severity"""
|
|
t = (text or "").strip().lower()
|
|
cat, sev = "other", "low"
|
|
for c in _CATEGORIES:
|
|
if c in t:
|
|
cat = c
|
|
break
|
|
for s in _SEVERITIES:
|
|
if s in t:
|
|
sev = s
|
|
break
|
|
return cat, sev # type: ignore
|
|
|
|
|
|
def _call_ollama(text: str, timeout: int = 5) -> Optional[Tuple[Category, Severity]]:
|
|
"""调用 Ollama 本地模型。需先运行 ollama run llama3.1 或 qwen2.5:7b"""
|
|
if PARSER_AI_DISABLED:
|
|
return None
|
|
try:
|
|
import requests
|
|
prompt = f"""Classify this news about US-Iran/middle east (one line only):
|
|
- category: deployment|alert|intel|diplomatic|other
|
|
- severity: low|medium|high|critical
|
|
|
|
News: {text[:300]}
|
|
|
|
Reply format: category:severity (e.g. alert:high)"""
|
|
r = requests.post(
|
|
"http://localhost:11434/api/chat",
|
|
json={
|
|
"model": OLLAMA_MODEL,
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
"stream": False,
|
|
"options": {"num_predict": 32},
|
|
},
|
|
timeout=timeout,
|
|
)
|
|
if r.status_code != 200:
|
|
return None
|
|
out = r.json().get("message", {}).get("content", "")
|
|
return _parse_ai_response(out)
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def _rule_classify(text: str) -> Category:
|
|
from parser import classify
|
|
return classify(text)
|
|
|
|
|
|
def _rule_severity(text: str, category: Category) -> Severity:
|
|
from parser import severity
|
|
return severity(text, category)
|
|
|
|
|
|
def classify(text: str) -> Category:
|
|
"""分类。AI 失败时回退规则"""
|
|
res = _call_ollama(text)
|
|
if res:
|
|
return res[0]
|
|
return _rule_classify(text)
|
|
|
|
|
|
def severity(text: str, category: Category) -> Severity:
|
|
"""严重度。AI 失败时回退规则"""
|
|
res = _call_ollama(text)
|
|
if res:
|
|
return res[1]
|
|
return _rule_severity(text, category)
|
|
|
|
|
|
def classify_and_severity(text: str) -> Tuple[Category, Severity]:
|
|
"""一次调用返回分类和严重度（减少 AI 调用）"""
|
|
if PARSER_AI_DISABLED:
|
|
from parser import classify, severity
|
|
c = classify(text)
|
|
return c, severity(text, c)
|
|
res = _call_ollama(text)
|
|
if res:
|
|
return res
|
|
return _rule_classify(text), _rule_severity(text, _rule_classify(text))
|