218 lines
7.5 KiB
Python
218 lines
7.5 KiB
Python
"""CitationDetector Agent - 将现有 CitationEngine 封装为 Agent"""
|
||
|
||
import logging
|
||
import re
|
||
import time
|
||
from datetime import datetime, timezone
|
||
|
||
from app.agent_framework.base import BaseAgent
|
||
from app.agent_framework.protocol import (
|
||
AgentCapability,
|
||
AgentType,
|
||
TaskMessage,
|
||
TaskResult,
|
||
TaskStatus,
|
||
)
|
||
from app.database import AsyncSessionLocal
|
||
from app.models.citation_record import CitationRecord
|
||
from app.models.query import Query
|
||
from app.workers.citation_engine import CitationEngine
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class CitationDetectorAgent(BaseAgent):
|
||
"""
|
||
引用检测 Agent：将现有 CitationEngine 封装为 BaseAgent 实现。
|
||
|
||
支持的任务类型:
|
||
- citation_detect: 执行完整的引用检测（遍历 query 的所有平台）
|
||
- citation_detect_single: 执行单个平台的引用检测
|
||
"""
|
||
|
||
def __init__(self):
|
||
super().__init__(
|
||
name="citation_detector",
|
||
agent_type=AgentType.CITATION_DETECTOR,
|
||
version="1.0.0",
|
||
)
|
||
self._engine = CitationEngine()
|
||
|
||
def get_capabilities(self) -> AgentCapability:
|
||
return AgentCapability(
|
||
agent_name=self.name,
|
||
agent_type=self.agent_type,
|
||
version=self.version,
|
||
supported_tasks=["citation_detect", "citation_detect_single"],
|
||
max_concurrency=3,
|
||
description="AI平台引用检测Agent：检测目标品牌在各AI平台回答中的引用情况",
|
||
)
|
||
|
||
async def execute(self, task: TaskMessage) -> TaskResult:
|
||
"""执行引用检测任务"""
|
||
started_at = datetime.now(timezone.utc)
|
||
start_time = time.monotonic()
|
||
|
||
try:
|
||
if task.task_type == "citation_detect":
|
||
output = await self._execute_full_detect(task)
|
||
elif task.task_type == "citation_detect_single":
|
||
output = await self._execute_single_detect(task)
|
||
else:
|
||
raise ValueError(f"Unsupported task type: {task.task_type}")
|
||
|
||
elapsed = time.monotonic() - start_time
|
||
return TaskResult(
|
||
task_id=task.task_id,
|
||
agent_name=self.name,
|
||
status=TaskStatus.COMPLETED,
|
||
output_data=output,
|
||
error_message=None,
|
||
started_at=started_at,
|
||
completed_at=datetime.now(timezone.utc),
|
||
metrics={
|
||
"elapsed_seconds": round(elapsed, 2),
|
||
"task_type": task.task_type,
|
||
},
|
||
)
|
||
|
||
except Exception as e:
|
||
elapsed = time.monotonic() - start_time
|
||
logger.error(f"CitationDetector task {task.task_id} failed: {e}")
|
||
return TaskResult(
|
||
task_id=task.task_id,
|
||
agent_name=self.name,
|
||
status=TaskStatus.FAILED,
|
||
output_data=None,
|
||
error_message=str(e),
|
||
started_at=started_at,
|
||
completed_at=datetime.now(timezone.utc),
|
||
metrics={
|
||
"elapsed_seconds": round(elapsed, 2),
|
||
"task_type": task.task_type,
|
||
},
|
||
)
|
||
|
||
async def _execute_full_detect(self, task: TaskMessage) -> dict:
|
||
"""
|
||
执行完整的引用检测（遍历 query 的所有平台）。
|
||
input_data 需包含: query_id (str)
|
||
"""
|
||
query_id = task.input_data.get("query_id")
|
||
if not query_id:
|
||
raise ValueError("input_data must contain 'query_id'")
|
||
|
||
async with AsyncSessionLocal() as db:
|
||
from sqlalchemy import select
|
||
|
||
stmt = select(Query).where(Query.id == query_id)
|
||
result = await db.execute(stmt)
|
||
query = result.scalar_one_or_none()
|
||
|
||
if not query:
|
||
raise ValueError(f"Query {query_id} not found")
|
||
|
||
# 上报进度：开始检测
|
||
await self.report_progress(
|
||
task_id=task.task_id,
|
||
progress=0.1,
|
||
message=f"Starting citation detection for query '{query.keyword}'",
|
||
)
|
||
|
||
records = await self._engine.execute_query(query, db)
|
||
|
||
# 上报进度：检测完成
|
||
await self.report_progress(
|
||
task_id=task.task_id,
|
||
progress=1.0,
|
||
message=f"Detection completed: {len(records)} records found",
|
||
)
|
||
|
||
# 构建输出
|
||
record_summaries = []
|
||
for r in records:
|
||
record_summaries.append({
|
||
"id": str(r.id),
|
||
"platform": r.platform,
|
||
"cited": r.cited,
|
||
"confidence": r.confidence,
|
||
"match_type": r.match_type,
|
||
})
|
||
|
||
return {
|
||
"query_id": str(query_id),
|
||
"keyword": query.keyword,
|
||
"total_records": len(records),
|
||
"cited_count": sum(1 for r in records if r.cited),
|
||
"records": record_summaries,
|
||
}
|
||
|
||
async def _execute_single_detect(self, task: TaskMessage) -> dict:
|
||
"""
|
||
执行单个平台的引用检测。
|
||
input_data 需包含: keyword, platform, target_brand, brand_aliases(optional)
|
||
"""
|
||
keyword = task.input_data.get("keyword")
|
||
platform = task.input_data.get("platform")
|
||
target_brand = task.input_data.get("target_brand")
|
||
brand_aliases = task.input_data.get("brand_aliases", [])
|
||
|
||
if not all([keyword, platform, target_brand]):
|
||
raise ValueError(
|
||
"input_data must contain 'keyword', 'platform', 'target_brand'"
|
||
)
|
||
|
||
# 上报进度
|
||
await self.report_progress(
|
||
task_id=task.task_id,
|
||
progress=0.2,
|
||
message=f"Querying platform '{platform}' with keyword '{keyword}'",
|
||
)
|
||
|
||
result = await self._engine.execute_single_platform(
|
||
keyword=keyword,
|
||
platform=platform,
|
||
target_brand=target_brand,
|
||
brand_aliases=brand_aliases,
|
||
)
|
||
|
||
# 上报进度：完成
|
||
await self.report_progress(
|
||
task_id=task.task_id,
|
||
progress=1.0,
|
||
message=f"Single platform detection completed on '{platform}'",
|
||
)
|
||
|
||
# 清理 raw_response 以避免返回过大的数据
|
||
output = {k: v for k, v in result.items() if k != "raw_response"}
|
||
return output
|
||
|
||
# -----------------------------------------------------------------------
|
||
# 向后兼容：保留原有 CitationEngine 的同步调用接口
|
||
# -----------------------------------------------------------------------
|
||
|
||
async def execute_query_compat(self, query: Query, db) -> list[CitationRecord]:
|
||
"""
|
||
向后兼容方法：供现有 scheduler 继续调用。
|
||
签名与 CitationEngine.execute_query 完全一致。
|
||
"""
|
||
return await self._engine.execute_query(query, db)
|
||
|
||
async def execute_single_platform_compat(
|
||
self, keyword: str, platform: str, target_brand: str, brand_aliases: list
|
||
) -> dict:
|
||
"""
|
||
向后兼容方法：供现有 scheduler 继续调用。
|
||
签名与 CitationEngine.execute_single_platform 完全一致。
|
||
"""
|
||
return await self._engine.execute_single_platform(
|
||
keyword=keyword,
|
||
platform=platform,
|
||
target_brand=target_brand,
|
||
brand_aliases=brand_aliases,
|
||
)
|
||
|
||
async def close(self):
|
||
"""关闭底层引擎"""
|
||
await self._engine.close()
|