Files
deerflow-code/offline-backend-20260512/backend/packages/harness/deerflow/persistence/agents/sql.py
T
吉鑫升andClaude Opus 4.7 78f6e37f76 新增智能体收藏、定时任务必须生成MD校验、问答时间显示及账号绑定的常用提示词
- 智能体支持从广场收藏到「我的」(新增 agent_favorites 表与收藏/取消收藏接口)
- 定时任务新增「必须生成Markdown文件」开关:注入提示词并校验生成的 .md 文件且正文超过100字,否则记为失败(存于 execution_context,无需迁移)
- 聊天问答展示时间(年月日 时分):问题在右下角、回答在左下角
- 常用提示词改为与账号绑定(新增 user_custom_prompts 表与 /api/user-prompts 接口),仅在首页输入框上方以居中标签 + 轮播展示,齿轮置于最右侧管理

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-28 18:29:32 +08:00

279 lines
12 KiB
Python

"""SQLAlchemy-backed custom agent repository."""
from __future__ import annotations
from datetime import UTC, datetime
from typing import Any
from uuid import uuid4
from sqlalchemy import and_, delete, func, or_, select
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
from deerflow.persistence.agents.base import AgentStore
from deerflow.persistence.agents.model import AgentFavoriteRow, AgentRow
def _row_to_dict(row: AgentRow) -> dict[str, Any]:
data = row.to_dict()
data["builtin"] = data.get("user_id") is None
for key in ("created_at", "updated_at"):
value = data.get(key)
if isinstance(value, datetime):
data[key] = value.isoformat()
return data
class AgentRepository(AgentStore):
def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
self._sf = session_factory
@staticmethod
def _visible_filter(user_id: str):
"""Visibility for single-record fetches (``get_visible``).
Lenient: a caller who already knows the agent id can fetch it even
when it's flagged ``is_private``. The private-square password is the
*discovery* gate (listings); individual reads are not re-gated, so a
bookmarked / linked private agent stays openable.
"""
return or_(AgentRow.user_id.is_(None), AgentRow.user_id == user_id, AgentRow.published.is_(True))
@staticmethod
def _listing_filter(user_id: str):
"""Visibility for *listings* — hides private-square agents from
non-owners so they only surface inside the password-gated endpoint.
Owners keep seeing their own private agents in '我的'.
"""
return or_(
AgentRow.user_id == user_id,
and_(
or_(AgentRow.user_id.is_(None), AgentRow.published.is_(True)),
AgentRow.is_private.is_(False),
),
)
async def list_visible(self, user_id: str) -> list[dict[str, Any]]:
stmt = select(AgentRow).where(self._listing_filter(user_id)).order_by(AgentRow.created_at.asc())
async with self._sf() as session:
result = await session.execute(stmt)
return [_row_to_dict(row) for row in result.scalars()]
async def get_visible(self, agent_id: str, user_id: str) -> dict[str, Any] | None:
stmt = select(AgentRow).where(AgentRow.id == agent_id, self._visible_filter(user_id))
async with self._sf() as session:
row = (await session.execute(stmt)).scalar_one_or_none()
return _row_to_dict(row) if row is not None else None
async def get_any(self, agent_id: str) -> dict[str, Any] | None:
async with self._sf() as session:
row = await session.get(AgentRow, agent_id)
return _row_to_dict(row) if row is not None else None
async def get_owned(self, agent_id: str, user_id: str) -> dict[str, Any] | None:
stmt = select(AgentRow).where(AgentRow.id == agent_id, AgentRow.user_id == user_id)
async with self._sf() as session:
row = (await session.execute(stmt)).scalar_one_or_none()
return _row_to_dict(row) if row is not None else None
async def create(self, data: dict[str, Any]) -> dict[str, Any]:
now = datetime.now(UTC)
row = AgentRow(
id=data["id"],
name=data["name"],
description=data.get("description") or "",
user_id=data.get("user_id"),
published=bool(data.get("published", False)),
is_private=bool(data.get("is_private", False)),
created_at=now,
updated_at=now,
)
async with self._sf() as session:
session.add(row)
await session.commit()
await session.refresh(row)
return _row_to_dict(row)
async def ensure_builtin(self, data: dict[str, Any]) -> dict[str, Any]:
async with self._sf() as session:
row = await session.get(AgentRow, data["id"])
if row is None:
now = datetime.now(UTC)
row = AgentRow(
id=data["id"],
name=data["name"],
description=data.get("description") or "",
user_id=None,
published=bool(data.get("published", True)),
created_at=now,
updated_at=now,
)
session.add(row)
elif row.user_id is None:
row.name = data["name"]
row.description = data.get("description") or ""
row.published = bool(data.get("published", True))
row.updated_at = datetime.now(UTC)
await session.commit()
await session.refresh(row)
return _row_to_dict(row)
async def update(self, agent_id: str, user_id: str, data: dict[str, Any]) -> dict[str, Any] | None:
async with self._sf() as session:
row = await session.get(AgentRow, agent_id)
if row is None or row.user_id != user_id:
return None
for key in ("name", "description", "published", "is_private"):
if key in data:
setattr(row, key, data[key])
row.updated_at = datetime.now(UTC)
await session.commit()
await session.refresh(row)
return _row_to_dict(row)
async def update_builtin(self, agent_id: str, data: dict[str, Any]) -> dict[str, Any] | None:
async with self._sf() as session:
row = await session.get(AgentRow, agent_id)
if row is None or row.user_id is not None:
return None
for key in ("name", "description", "published", "is_private"):
if key in data:
setattr(row, key, data[key])
row.updated_at = datetime.now(UTC)
await session.commit()
await session.refresh(row)
return _row_to_dict(row)
async def delete(self, agent_id: str, user_id: str) -> bool:
async with self._sf() as session:
row = await session.get(AgentRow, agent_id)
if row is None or row.user_id != user_id:
return False
await session.delete(row)
await session.commit()
return True
async def delete_builtin(self, agent_id: str) -> bool:
async with self._sf() as session:
row = await session.get(AgentRow, agent_id)
if row is None or row.user_id is not None:
return False
await session.delete(row)
await session.commit()
return True
async def count_owned(self, user_id: str) -> int:
stmt = select(func.count()).select_from(AgentRow).where(AgentRow.user_id == user_id)
async with self._sf() as session:
return await session.scalar(stmt) or 0
async def set_featured(self, agent_id: str, order: int | None) -> dict[str, Any] | None:
async with self._sf() as session:
row = await session.get(AgentRow, agent_id)
if row is None:
return None
row.featured_order = order
row.updated_at = datetime.now(UTC)
await session.commit()
await session.refresh(row)
return _row_to_dict(row)
async def list_featured(self, user_id: str) -> list[dict[str, Any]]:
stmt = (
select(AgentRow)
.where(AgentRow.featured_order.is_not(None), self._listing_filter(user_id))
.order_by(AgentRow.featured_order.asc(), AgentRow.created_at.asc())
)
async with self._sf() as session:
result = await session.execute(stmt)
return [_row_to_dict(row) for row in result.scalars()]
async def list_by_ids(self, agent_ids: list[str], user_id: str) -> list[dict[str, Any]]:
if not agent_ids:
return []
order_map = {aid: idx for idx, aid in enumerate(agent_ids)}
stmt = select(AgentRow).where(AgentRow.id.in_(agent_ids), self._listing_filter(user_id))
async with self._sf() as session:
result = await session.execute(stmt)
rows = list(result.scalars())
rows.sort(key=lambda r: order_map.get(r.id, len(order_map)))
return [_row_to_dict(r) for r in rows]
async def list_private(self) -> list[dict[str, Any]]:
stmt = select(AgentRow).where(AgentRow.is_private.is_(True)).order_by(AgentRow.created_at.asc())
async with self._sf() as session:
result = await session.execute(stmt)
return [_row_to_dict(row) for row in result.scalars()]
async def list_all(self) -> list[dict[str, Any]]:
stmt = select(AgentRow).order_by(AgentRow.created_at.asc())
async with self._sf() as session:
result = await session.execute(stmt)
return [_row_to_dict(row) for row in result.scalars()]
async def update_any(self, agent_id: str, data: dict[str, Any]) -> dict[str, Any] | None:
async with self._sf() as session:
row = await session.get(AgentRow, agent_id)
if row is None:
return None
for key in ("name", "description", "published", "is_private"):
if key in data:
setattr(row, key, data[key])
row.updated_at = datetime.now(UTC)
await session.commit()
await session.refresh(row)
return _row_to_dict(row)
async def delete_any(self, agent_id: str) -> bool:
async with self._sf() as session:
row = await session.get(AgentRow, agent_id)
if row is None:
return False
await session.delete(row)
await session.commit()
return True
# --- Favorites ("收藏") ---
async def add_favorite(self, user_id: str, agent_id: str) -> None:
async with self._sf() as session:
existing = (
await session.execute(
select(AgentFavoriteRow.id)
.where(AgentFavoriteRow.user_id == user_id, AgentFavoriteRow.agent_id == agent_id)
.limit(1)
)
).scalar_one_or_none()
if existing is not None:
return
session.add(AgentFavoriteRow(id=uuid4().hex, user_id=user_id, agent_id=agent_id))
try:
await session.commit()
except IntegrityError:
# Lost a race against a concurrent favorite of the same agent —
# the unique constraint already guarantees the desired state.
await session.rollback()
async def remove_favorite(self, user_id: str, agent_id: str) -> bool:
async with self._sf() as session:
result = await session.execute(
delete(AgentFavoriteRow).where(
AgentFavoriteRow.user_id == user_id,
AgentFavoriteRow.agent_id == agent_id,
)
)
await session.commit()
return (result.rowcount or 0) > 0
async def list_favorite_ids(self, user_id: str) -> set[str]:
async with self._sf() as session:
result = await session.execute(
select(AgentFavoriteRow.agent_id).where(AgentFavoriteRow.user_id == user_id)
)
return {row[0] for row in result.all()}
async def delete_favorites_for_agent(self, agent_id: str) -> None:
async with self._sf() as session:
await session.execute(delete(AgentFavoriteRow).where(AgentFavoriteRow.agent_id == agent_id))
await session.commit()