359 lines
13 KiB
Python
359 lines
13 KiB
Python
import difflib
|
||
import logging
|
||
import re
|
||
import uuid
|
||
from datetime import datetime, timedelta
|
||
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
from sqlalchemy import select
|
||
|
||
|
||
def _sanitize_raw_response(text: str | None) -> str:
|
||
"""清理原始响应中的无效控制字符，避免 PostgreSQL UTF-8 插入失败"""
|
||
if not text:
|
||
return ""
|
||
# 移除 NULL 字节及其他非法控制字符，保留 \n \t \r
|
||
return re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f]", "", text)
|
||
|
||
from app.models.citation_record import CitationRecord
|
||
from app.models.query import Query
|
||
from app.models.query_task import QueryTask
|
||
from app.workers.platforms.kimi import KimiAdapter
|
||
from app.workers.platforms.wenxin import WenxinAdapter
|
||
from app.workers.platforms.tongyi import TongyiAdapter
|
||
from app.workers.platforms.doubao import DoubaoAdapter
|
||
from app.workers.platforms.qingyan import QingyanAdapter
|
||
from app.workers.platforms.tiangong import TiangongAdapter
|
||
from app.workers.platforms.xinghuo import XinghuoAdapter
|
||
from app.workers.citation_extractor import analyze_citations
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class BrandMatcher:
|
||
"""品牌匹配器：检测文本中是否引用了目标品牌"""
|
||
|
||
def __init__(self, target_brand: str, brand_aliases: list[str] | None = None):
|
||
self.target_brand = target_brand
|
||
self.brand_aliases = brand_aliases or []
|
||
|
||
def match(self, text: str) -> dict:
|
||
"""
|
||
返回: {
|
||
"cited": bool,
|
||
"confidence": float, # 0.0-1.0
|
||
"match_type": str, # "exact"/"alias"/"fuzzy"/None
|
||
"position": int|None, # 在文本段落中的位置（第几段提到，1-based）
|
||
"citation_text": str|None, # 被引用的上下文片段
|
||
}
|
||
"""
|
||
if not text:
|
||
return {
|
||
"cited": False,
|
||
"confidence": 0.0,
|
||
"match_type": None,
|
||
"position": None,
|
||
"citation_text": None,
|
||
}
|
||
|
||
# 1. 精确匹配
|
||
if self.target_brand in text:
|
||
position, citation_text = self._extract_position_and_context(text, self.target_brand)
|
||
return {
|
||
"cited": True,
|
||
"confidence": 1.0,
|
||
"match_type": "exact",
|
||
"position": position,
|
||
"citation_text": citation_text,
|
||
}
|
||
|
||
# 2. 别名匹配
|
||
for alias in self.brand_aliases:
|
||
if alias in text:
|
||
position, citation_text = self._extract_position_and_context(text, alias)
|
||
return {
|
||
"cited": True,
|
||
"confidence": 0.9,
|
||
"match_type": "alias",
|
||
"position": position,
|
||
"citation_text": citation_text,
|
||
}
|
||
|
||
# 3. 模糊匹配
|
||
best_ratio = 0.0
|
||
best_match = None
|
||
for word in self._extract_candidates(text):
|
||
ratio = difflib.SequenceMatcher(None, self.target_brand, word).ratio()
|
||
if ratio > best_ratio:
|
||
best_ratio = ratio
|
||
best_match = word
|
||
|
||
for alias in self.brand_aliases:
|
||
for word in self._extract_candidates(text):
|
||
ratio = difflib.SequenceMatcher(None, alias, word).ratio()
|
||
if ratio > best_ratio:
|
||
best_ratio = ratio
|
||
best_match = word
|
||
|
||
if best_ratio > 0.4 and best_match:
|
||
position, citation_text = self._extract_position_and_context(text, best_match)
|
||
return {
|
||
"cited": True,
|
||
"confidence": round(best_ratio, 2),
|
||
"match_type": "fuzzy",
|
||
"position": position,
|
||
"citation_text": citation_text,
|
||
}
|
||
|
||
return {
|
||
"cited": False,
|
||
"confidence": 0.0,
|
||
"match_type": None,
|
||
"position": None,
|
||
"citation_text": None,
|
||
}
|
||
|
||
def _extract_candidates(self, text: str) -> list[str]:
|
||
"""从文本中提取候选词（按非文字字符分割）"""
|
||
# 匹配中文词组、英文单词等
|
||
return [w for w in re.split(r'[^\w\u4e00-\u9fff]+', text) if len(w) >= 2]
|
||
|
||
def _extract_position_and_context(self, text: str, keyword: str) -> tuple[int | None, str | None]:
|
||
"""提取品牌首次出现的段落位置（1-based）和上下文片段"""
|
||
paragraphs = [p.strip() for p in text.split('\n') if p.strip()]
|
||
if not paragraphs:
|
||
paragraphs = [text]
|
||
|
||
for idx, paragraph in enumerate(paragraphs, start=1):
|
||
if keyword in paragraph:
|
||
# 截取前200字符
|
||
snippet = paragraph[:200]
|
||
return idx, snippet
|
||
|
||
return None, None
|
||
|
||
|
||
class CompetitorDetector:
|
||
"""竞争品牌检测器"""
|
||
|
||
# 预定义一些常见行业品牌列表
|
||
KNOWN_BRANDS = {
|
||
"保险": ["中国平安", "中国人寿", "太平洋保险", "新华保险", "泰康保险", "中国人保", "友邦保险"],
|
||
"金融": ["工商银行", "建设银行", "农业银行", "中国银行", "招商银行", "交通银行"],
|
||
"科技": ["华为", "腾讯", "阿里巴巴", "百度", "字节跳动", "小米", "京东"],
|
||
}
|
||
|
||
def detect(self, text: str, target_brand: str) -> list[str]:
|
||
"""检测文本中出现的其他品牌（排除 target_brand）"""
|
||
if not text:
|
||
return []
|
||
|
||
competitors = set()
|
||
for category, brands in self.KNOWN_BRANDS.items():
|
||
for brand in brands:
|
||
if brand == target_brand:
|
||
continue
|
||
if brand in text:
|
||
competitors.add(brand)
|
||
|
||
return sorted(list(competitors))
|
||
|
||
|
||
class CitationEngine:
|
||
"""引用检测引擎核心"""
|
||
|
||
def __init__(self):
|
||
self.platforms = {
|
||
"wenxin": WenxinAdapter(),
|
||
"kimi": KimiAdapter(),
|
||
"tongyi": TongyiAdapter(),
|
||
"doubao": DoubaoAdapter(),
|
||
"qingyan": QingyanAdapter(),
|
||
"tiangong": TiangongAdapter(),
|
||
"xinghuo": XinghuoAdapter(),
|
||
}
|
||
self.matcher = None
|
||
self.competitor_detector = CompetitorDetector()
|
||
|
||
async def execute_query(self, query: Query, db: AsyncSession) -> list[CitationRecord]:
|
||
"""
|
||
执行一个查询任务：
|
||
1. 创建 BrandMatcher
|
||
2. 遍历 query.platforms
|
||
3. 对每个 platform 执行查询和检测
|
||
4. 更新 query.last_queried_at 和 query.next_query_at
|
||
"""
|
||
self.matcher = BrandMatcher(
|
||
target_brand=query.target_brand,
|
||
brand_aliases=query.brand_aliases or [],
|
||
)
|
||
|
||
records: list[CitationRecord] = []
|
||
platforms = query.platforms or ["wenxin", "kimi"]
|
||
|
||
for platform_name in platforms:
|
||
# 查找或创建 QueryTask
|
||
task = await self._get_or_create_task(db, query.id, platform_name)
|
||
|
||
# 更新状态为 running
|
||
task.status = "running"
|
||
task.started_at = datetime.utcnow()
|
||
task.error_message = None
|
||
await db.commit()
|
||
|
||
try:
|
||
result = await self.execute_single_platform(
|
||
keyword=query.keyword,
|
||
platform=platform_name,
|
||
target_brand=query.target_brand,
|
||
brand_aliases=query.brand_aliases or [],
|
||
)
|
||
|
||
# 创建 CitationRecord
|
||
record = CitationRecord(
|
||
query_id=query.id,
|
||
platform=platform_name,
|
||
cited=result["cited"],
|
||
citation_position=result.get("position"),
|
||
citation_text=result.get("citation_text"),
|
||
competitor_brands=result.get("competitor_brands", []),
|
||
raw_response=_sanitize_raw_response(result.get("raw_response", "")),
|
||
confidence=result.get("confidence"),
|
||
match_type=result.get("match_type"),
|
||
# 引用源分析字段
|
||
data_source=result.get("data_source"),
|
||
source_urls=result.get("source_urls"),
|
||
source_titles=result.get("source_titles"),
|
||
citation_contexts=result.get("citation_contexts"),
|
||
ai_response_text=_sanitize_raw_response(result.get("ai_response_text", "")),
|
||
)
|
||
db.add(record)
|
||
records.append(record)
|
||
|
||
# 更新 QueryTask 状态为 success
|
||
task.status = "success"
|
||
task.completed_at = datetime.utcnow()
|
||
await db.commit()
|
||
|
||
except Exception as e:
|
||
logger.error(f"平台 {platform_name} 查询失败: {e}")
|
||
error_msg = str(e)
|
||
task.status = "failed"
|
||
task.error_message = error_msg
|
||
task.completed_at = datetime.utcnow()
|
||
|
||
# 创建一条 cited=False 的记录作为占位
|
||
record = CitationRecord(
|
||
query_id=query.id,
|
||
platform=platform_name,
|
||
cited=False,
|
||
raw_response=_sanitize_raw_response(error_msg),
|
||
)
|
||
db.add(record)
|
||
records.append(record)
|
||
await db.commit()
|
||
|
||
# 更新 Query 时间字段
|
||
query.last_queried_at = datetime.utcnow()
|
||
query.next_query_at = self._calculate_next_query_at(query.frequency)
|
||
await db.commit()
|
||
|
||
return records
|
||
|
||
async def execute_single_platform(
|
||
self,
|
||
keyword: str,
|
||
platform: str,
|
||
target_brand: str,
|
||
brand_aliases: list,
|
||
) -> dict:
|
||
"""执行单个平台的查询和检测"""
|
||
adapter = self.platforms.get(platform)
|
||
if not adapter:
|
||
raise ValueError(f"不支持的平台: {platform}")
|
||
|
||
# 获取平台内容（将关键词与目标品牌组合，确保结果包含品牌信息）
|
||
search_keyword = f"{keyword} {target_brand}"
|
||
raw_response = await adapter.query(search_keyword)
|
||
|
||
# 引用源分析
|
||
citation_analysis = analyze_citations(raw_response)
|
||
|
||
# 品牌匹配（使用清理后的纯文本进行匹配）
|
||
matcher = BrandMatcher(target_brand=target_brand, brand_aliases=brand_aliases)
|
||
match_result = matcher.match(citation_analysis.clean_response)
|
||
|
||
# 竞争品牌检测
|
||
competitor_brands = self.competitor_detector.detect(
|
||
citation_analysis.clean_response, target_brand
|
||
)
|
||
|
||
# 提取引用源信息
|
||
source_urls = [
|
||
c.source_url for c in citation_analysis.citations if c.source_url
|
||
]
|
||
source_titles = [
|
||
c.source_title for c in citation_analysis.citations if c.source_title
|
||
]
|
||
citation_contexts = [
|
||
c.citation_context for c in citation_analysis.citations if c.citation_context
|
||
]
|
||
|
||
return {
|
||
"cited": match_result["cited"],
|
||
"confidence": match_result["confidence"],
|
||
"match_type": match_result["match_type"],
|
||
"position": match_result["position"],
|
||
"citation_text": match_result["citation_text"],
|
||
"competitor_brands": competitor_brands,
|
||
"raw_response": raw_response,
|
||
# 引用源分析新增字段
|
||
"data_source": citation_analysis.data_source,
|
||
"source_urls": source_urls,
|
||
"source_titles": source_titles,
|
||
"citation_contexts": citation_contexts,
|
||
"ai_response_text": citation_analysis.clean_response,
|
||
}
|
||
|
||
async def _get_or_create_task(
|
||
self, db: AsyncSession, query_id: uuid.UUID, platform: str
|
||
) -> QueryTask:
|
||
"""获取或创建 QueryTask"""
|
||
stmt = select(QueryTask).where(
|
||
QueryTask.query_id == query_id,
|
||
QueryTask.platform == platform,
|
||
)
|
||
result = await db.execute(stmt)
|
||
task = result.scalar_one_or_none()
|
||
|
||
if not task:
|
||
task = QueryTask(
|
||
query_id=query_id,
|
||
platform=platform,
|
||
status="pending",
|
||
)
|
||
db.add(task)
|
||
await db.commit()
|
||
await db.refresh(task)
|
||
|
||
return task
|
||
|
||
def _calculate_next_query_at(self, frequency: str | None) -> datetime:
|
||
"""根据频率计算下次查询时间"""
|
||
now = datetime.utcnow()
|
||
freq_map = {
|
||
"daily": timedelta(days=1),
|
||
"weekly": timedelta(days=7),
|
||
"monthly": timedelta(days=30),
|
||
}
|
||
delta = freq_map.get(frequency or "weekly", timedelta(days=7))
|
||
return now + delta
|
||
|
||
async def close(self):
|
||
"""关闭所有平台适配器"""
|
||
for adapter in self.platforms.values():
|
||
try:
|
||
await adapter.close()
|
||
except Exception as e:
|
||
logger.warning(f"关闭适配器 {adapter.platform_name} 时出错: {e}")
|