Codex
102 lines
3.7 KiB
Python
102 lines
3.7 KiB
Python
from __future__ import annotations
|
|
|
|
from hashlib import sha256
|
|
from typing import Optional
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.core.timezone import beijing_now
|
|
from app.db.models import AgentRun, Report, Signal, Topic
|
|
from app.schemas.agent import SignalItem
|
|
|
|
|
|
class AgentRepository:
|
|
"""Database boundary for the agent runner."""
|
|
|
|
def __init__(self, db: Session) -> None:
|
|
self.db = db
|
|
|
|
def list_enabled_topics(self) -> list[Topic]:
|
|
# 启用主题是Agent每次运行时使用的长期目标。
|
|
return list(self.db.scalars(select(Topic).where(Topic.enabled.is_(True))).all())
|
|
|
|
def create_run(self, request_payload: dict) -> AgentRun:
|
|
# 运行先进入running状态，情报和日报保存完成后再标记完成。
|
|
run = AgentRun(status="running", started_at=beijing_now(), request_payload=request_payload)
|
|
self.db.add(run)
|
|
self.db.flush()
|
|
return run
|
|
|
|
def mark_run_complete(self, run: AgentRun, raw_response: dict) -> None:
|
|
run.status = "completed"
|
|
run.finished_at = beijing_now()
|
|
run.raw_response = raw_response
|
|
self.db.flush()
|
|
|
|
def mark_run_failed(self, run: AgentRun, error: str) -> None:
|
|
run.status = "failed"
|
|
run.finished_at = beijing_now()
|
|
run.error = error
|
|
self.db.flush()
|
|
|
|
def save_signals(self, run: AgentRun, items: list[SignalItem]) -> list[Signal]:
|
|
# 相同来源链接只保存一次，避免日报和仪表盘重复出现同一条情报。
|
|
saved: list[Signal] = []
|
|
for item in items:
|
|
signal = self.save_or_update_signal(run=run, item=item)
|
|
if signal is None:
|
|
continue
|
|
saved.append(signal)
|
|
return saved
|
|
|
|
def save_or_update_signal(self, run: AgentRun, item: SignalItem) -> Optional[Signal]:
|
|
# 同一来源链接只保留一条情报，流式接口和普通接口共用这条去重规则。
|
|
source_url_hash = self._source_url_hash(item.source_url)
|
|
existing = self.db.scalar(select(Signal).where(Signal.source_url_hash == source_url_hash))
|
|
if existing and existing.run_id != run.id:
|
|
return None
|
|
if existing:
|
|
existing.topic = item.topic
|
|
existing.title = item.title
|
|
existing.summary = item.summary
|
|
existing.source_name = item.source_name
|
|
existing.published_at = item.published_at
|
|
existing.importance = item.importance
|
|
existing.entities = item.entities
|
|
self.db.flush()
|
|
return existing
|
|
signal = Signal(
|
|
run_id=run.id,
|
|
topic=item.topic,
|
|
title=item.title,
|
|
summary=item.summary,
|
|
source_url=item.source_url,
|
|
source_url_hash=source_url_hash,
|
|
source_name=item.source_name,
|
|
published_at=item.published_at,
|
|
importance=item.importance,
|
|
entities=item.entities,
|
|
created_at=beijing_now(),
|
|
)
|
|
self.db.add(signal)
|
|
self.db.flush()
|
|
return signal
|
|
|
|
def _source_url_hash(self, source_url: str) -> str:
|
|
# MySQL不适合直接唯一索引长URL，使用SHA-256作为去重键。
|
|
return sha256(source_url.encode("utf-8")).hexdigest()
|
|
|
|
def save_report(self, run: AgentRun, title: str, content_md: str, file_path: str) -> Report:
|
|
# 文件路径便于人工打开Markdown产物，MySQL保存可查询正文。
|
|
report = Report(
|
|
run_id=run.id,
|
|
title=title,
|
|
content_md=content_md,
|
|
file_path=file_path,
|
|
created_at=beijing_now(),
|
|
)
|
|
self.db.add(report)
|
|
self.db.flush()
|
|
return report
|