666 lines
27 KiB
Python
666 lines
27 KiB
Python
"""SQLAlchemy-backed custom agent repository."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
from datetime import UTC, datetime
|
|
from typing import Any
|
|
from uuid import uuid4
|
|
|
|
from sqlalchemy import and_, case, delete, func, or_, select
|
|
from sqlalchemy.exc import IntegrityError
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
|
|
|
from deerflow.config.agent_name_desensitize import source_terms_for_needle
|
|
from deerflow.persistence.agents.base import AgentStore
|
|
from deerflow.persistence.agents.model import (
|
|
AgentExtrasRow,
|
|
AgentFavoriteRow,
|
|
AgentPinRow,
|
|
AgentRow,
|
|
AgentSkillRow,
|
|
)
|
|
from deerflow.persistence.roundtable_chains.model import RoundtableChainRow
|
|
from deerflow.persistence.tags.model import TagAssignmentRow
|
|
|
|
# Seats are stored as JSON text. A correlated LIKE against that blob, once per
|
|
# agent row, is what made 「我的」 spike: the plan sometimes nested a full chain
|
|
# scan inside the favorite EXISTS. Extract ids once instead.
|
|
_SEAT_AGENT_ID = re.compile(r'"agent_id"\s*:\s*"([^"]+)"')
|
|
|
|
|
|
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
|
|
|
|
async def _load_chain_seat_agent_ids(self) -> set[str]:
|
|
stmt = select(RoundtableChainRow.seats)
|
|
async with self._sf() as session:
|
|
blobs = (await session.execute(stmt)).scalars().all()
|
|
found: set[str] = set()
|
|
for blob in blobs:
|
|
if not blob:
|
|
continue
|
|
found.update(_SEAT_AGENT_ID.findall(str(blob)))
|
|
return found
|
|
|
|
async def _favorite_chain_exclusions(self, user_id: str) -> set[str] | None:
|
|
"""Chain-seat ids to drop from position favorites.
|
|
|
|
``None`` means this user has no position-granted favorites, so the
|
|
list does not need to read chain seats at all.
|
|
"""
|
|
async with self._sf() as session:
|
|
position_favorite = await session.scalar(
|
|
select(AgentFavoriteRow.id)
|
|
.where(
|
|
AgentFavoriteRow.user_id == user_id,
|
|
AgentFavoriteRow.origin == "position",
|
|
)
|
|
.limit(1)
|
|
)
|
|
if position_favorite is None:
|
|
return None
|
|
return await self._load_chain_seat_agent_ids()
|
|
|
|
@staticmethod
|
|
def _valid_favorite_filter(chain_agent_ids: set[str] | None):
|
|
if chain_agent_ids is None:
|
|
return AgentFavoriteRow.origin == "user"
|
|
position = [
|
|
AgentFavoriteRow.origin == "position",
|
|
AgentRow.user_id.is_not(None),
|
|
]
|
|
if chain_agent_ids:
|
|
position.append(AgentFavoriteRow.agent_id.notin_(sorted(chain_agent_ids)))
|
|
return or_(AgentFavoriteRow.origin == "user", and_(*position))
|
|
|
|
@staticmethod
|
|
def _visible_filter(user_id: str, chain_agent_ids: set[str] | None):
|
|
"""Visibility for single-record fetches (``get_visible``).
|
|
|
|
Lenient — must cover the union of every listing surface, so any agent
|
|
a user can see in a list always opens (detail + chat) instead of
|
|
404ing:
|
|
|
|
- built-in / owned / published — the baseline gallery visibility;
|
|
- ``is_private`` — the private-square password is the *discovery*
|
|
gate (listings); individual reads are not re-gated, so a private
|
|
agent opened from the square (or a bookmark/link) stays openable.
|
|
Note moving an agent into the private square forces
|
|
``published=false``, so this clause cannot lean on ``published``;
|
|
- favorited — the "我的" listing keeps favorited agents even after
|
|
the owner unpublishes them, so reads keep honoring the favorite.
|
|
"""
|
|
valid_favorite_exists = (
|
|
select(AgentFavoriteRow.id)
|
|
.where(
|
|
AgentFavoriteRow.user_id == user_id,
|
|
AgentFavoriteRow.agent_id == AgentRow.id,
|
|
AgentRepository._valid_favorite_filter(chain_agent_ids),
|
|
)
|
|
.exists()
|
|
)
|
|
return or_(
|
|
AgentRow.user_id.is_(None),
|
|
AgentRow.user_id == user_id,
|
|
AgentRow.published.is_(True),
|
|
AgentRow.is_private.is_(True),
|
|
valid_favorite_exists,
|
|
)
|
|
|
|
@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 '我的'.
|
|
Admin-hidden agents (``hidden=True``) are excluded from every
|
|
non-admin listing; the admin console reads them via ``list_all`` /
|
|
``list_paginated(admin=True)``.
|
|
"""
|
|
return and_(
|
|
or_(
|
|
AgentRow.user_id == user_id,
|
|
and_(
|
|
or_(
|
|
AgentRow.user_id.is_(None),
|
|
AgentRow.published.is_(True),
|
|
),
|
|
AgentRow.is_private.is_(False),
|
|
),
|
|
),
|
|
AgentRow.hidden.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:
|
|
chain_agent_ids = await self._favorite_chain_exclusions(user_id)
|
|
stmt = select(AgentRow).where(AgentRow.id == agent_id, self._visible_filter(user_id, chain_agent_ids))
|
|
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)),
|
|
square_id=str(data.get("square_id") or ""),
|
|
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", "square_id", "hidden"):
|
|
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", "square_id", "hidden"):
|
|
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_sidebar(self, user_id: str, *, limit: int = 3) -> list[dict[str, Any]]:
|
|
"""Fetch two tiny id/name projections instead of assembling gallery records."""
|
|
chain_agent_ids = await self._favorite_chain_exclusions(user_id)
|
|
mine_filter = self._scope_filter(
|
|
user_id=user_id,
|
|
scope="mine",
|
|
admin=False,
|
|
chain_agent_ids=chain_agent_ids,
|
|
)
|
|
square_filter = self._scope_filter(user_id=user_id, scope="square", admin=False)
|
|
pin_exists = select(AgentPinRow.id).where(
|
|
AgentPinRow.user_id == user_id,
|
|
AgentPinRow.agent_id == AgentRow.id,
|
|
).exists()
|
|
pin_rank = case((pin_exists, 0), else_=1)
|
|
|
|
is_mine = mine_filter
|
|
stmt = (
|
|
select(AgentRow.id, AgentRow.name)
|
|
.where(or_(mine_filter, square_filter), AgentRow.hidden.is_(False))
|
|
.order_by(
|
|
pin_rank.asc(),
|
|
case((is_mine, 0), else_=1).asc(),
|
|
AgentRow.created_at.desc(),
|
|
AgentRow.id.asc(),
|
|
)
|
|
.limit(max(1, min(int(limit), 10)))
|
|
)
|
|
async with self._sf() as session:
|
|
rows = (await session.execute(stmt)).all()
|
|
return [{"id": str(agent_id), "name": str(name or "")} for agent_id, name in rows]
|
|
|
|
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), AgentRow.hidden.is_(False))
|
|
.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", "square_id", "hidden"):
|
|
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
|
|
|
|
# --- Paginated listing ---
|
|
|
|
def _scope_filter(self, *, user_id: str, scope: str | None, admin: bool, chain_agent_ids: set[str] | None = None):
|
|
"""Build the visibility predicate for a paginated listing."""
|
|
if admin:
|
|
return None # no visibility filter — every agent
|
|
if scope == "square":
|
|
return and_(
|
|
or_(AgentRow.user_id.is_(None), AgentRow.published.is_(True)),
|
|
AgentRow.is_private.is_(False),
|
|
)
|
|
if scope == "mine":
|
|
# owned OR valid favorites. Built-in agents and business-chain
|
|
# seats only enter "mine" when explicitly favorited by the user;
|
|
# the original position-granted behavior stays for ordinary agents.
|
|
valid_favorite_exists = (
|
|
select(AgentFavoriteRow.id)
|
|
.where(
|
|
AgentFavoriteRow.user_id == user_id,
|
|
AgentFavoriteRow.agent_id == AgentRow.id,
|
|
self._valid_favorite_filter(chain_agent_ids),
|
|
)
|
|
.exists()
|
|
)
|
|
return or_(AgentRow.user_id == user_id, valid_favorite_exists)
|
|
if scope == "private":
|
|
return AgentRow.is_private.is_(True)
|
|
if scope == "builtin":
|
|
# 内置智能体 = user_id IS NULL。仅管理员可用此 scope（路由层已 admin 闸门）。
|
|
return AgentRow.user_id.is_(None)
|
|
return self._listing_filter(user_id)
|
|
|
|
@staticmethod
|
|
def _search_clause(search: str):
|
|
"""Match the raw name OR (via the desensitize mirror) its display form."""
|
|
needle = search.strip().lower()
|
|
clauses = [func.lower(AgentRow.name).like(f"%{needle}%")]
|
|
for term in source_terms_for_needle(needle):
|
|
clauses.append(AgentRow.name.like(f"%{term}%"))
|
|
return or_(*clauses)
|
|
|
|
@staticmethod
|
|
def _tag_clause(tag_ids: list[str]):
|
|
subq = select(TagAssignmentRow.target_id).where(
|
|
TagAssignmentRow.target_type == "agent",
|
|
TagAssignmentRow.tag_id.in_(tag_ids),
|
|
)
|
|
return AgentRow.id.in_(subq)
|
|
|
|
async def list_paginated(
|
|
self,
|
|
*,
|
|
user_id: str,
|
|
scope: str | None,
|
|
search: str | None,
|
|
tag_ids: list[str] | None,
|
|
page: int,
|
|
page_size: int,
|
|
admin: bool = False,
|
|
sort: str = "time_desc",
|
|
owner: str | None = None,
|
|
square_id: str | None = None,
|
|
square_is_default: bool = False,
|
|
hidden: bool | None = None,
|
|
) -> tuple[list[dict[str, Any]], int]:
|
|
filters = []
|
|
chain_agent_ids = await self._favorite_chain_exclusions(user_id) if scope == "mine" and not admin else None
|
|
scope_filter = self._scope_filter(user_id=user_id, scope=scope, admin=admin, chain_agent_ids=chain_agent_ids)
|
|
if scope_filter is not None:
|
|
filters.append(scope_filter)
|
|
if not admin:
|
|
# Admin-hidden agents vanish from every non-admin listing scope.
|
|
filters.append(AgentRow.hidden.is_(False))
|
|
elif hidden is not None:
|
|
# Admin console filter: keep only hidden / only visible agents.
|
|
filters.append(AgentRow.hidden.is_(hidden))
|
|
if search and search.strip():
|
|
filters.append(self._search_clause(search))
|
|
if tag_ids:
|
|
filters.append(self._tag_clause(tag_ids))
|
|
if owner and owner.strip():
|
|
# Filter to one publisher. Built-in agents (user_id IS NULL) never match.
|
|
filters.append(AgentRow.user_id == owner.strip())
|
|
if square_id and square_id.strip():
|
|
sid = square_id.strip()
|
|
if square_is_default:
|
|
# Legacy / unassigned rows (square_id == "") belong to the
|
|
# default square, so include them alongside explicit matches.
|
|
filters.append(or_(AgentRow.square_id == sid, AgentRow.square_id == ""))
|
|
else:
|
|
filters.append(AgentRow.square_id == sid)
|
|
|
|
count_stmt = select(func.count()).select_from(AgentRow)
|
|
pin_subq = select(AgentPinRow.agent_id).where(AgentPinRow.user_id == user_id)
|
|
pin_rank = case((AgentRow.id.in_(pin_subq), 0), else_=1)
|
|
|
|
if sort == "time_asc":
|
|
order_by = (AgentRow.created_at.asc(), AgentRow.id.asc())
|
|
elif sort == "updated_desc":
|
|
order_by = (AgentRow.updated_at.desc(), AgentRow.id.asc())
|
|
elif sort == "updated_asc":
|
|
order_by = (AgentRow.updated_at.asc(), AgentRow.id.asc())
|
|
elif sort == "name_asc":
|
|
order_by = (func.lower(AgentRow.name).asc(), AgentRow.id.asc())
|
|
elif sort == "name_desc":
|
|
order_by = (func.lower(AgentRow.name).desc(), AgentRow.id.asc())
|
|
else:
|
|
order_by = (AgentRow.created_at.desc(), AgentRow.id.asc())
|
|
|
|
page_stmt = (
|
|
select(AgentRow)
|
|
.order_by(pin_rank.asc(), *order_by)
|
|
.limit(page_size)
|
|
.offset(max(page - 1, 0) * page_size)
|
|
)
|
|
if filters:
|
|
where = and_(*filters)
|
|
count_stmt = count_stmt.where(where)
|
|
page_stmt = page_stmt.where(where)
|
|
|
|
async with self._sf() as session:
|
|
total = await session.scalar(count_stmt) or 0
|
|
result = await session.execute(page_stmt)
|
|
rows = [_row_to_dict(row) for row in result.scalars()]
|
|
return rows, total
|
|
|
|
# --- Listing cache (skills / tool_groups / model) ---
|
|
|
|
async def replace_skills(self, agent_id: str, skills: list[str] | None) -> None:
|
|
async with self._sf() as session:
|
|
await session.execute(delete(AgentSkillRow).where(AgentSkillRow.agent_id == agent_id))
|
|
for position, name in enumerate(skills or []):
|
|
if not isinstance(name, str):
|
|
continue
|
|
session.add(
|
|
AgentSkillRow(
|
|
id=uuid4().hex,
|
|
agent_id=agent_id,
|
|
skill_name=name,
|
|
position=position,
|
|
)
|
|
)
|
|
await session.commit()
|
|
|
|
async def set_extras(
|
|
self,
|
|
agent_id: str,
|
|
*,
|
|
tool_groups: list[str] | None,
|
|
model: str | None,
|
|
chat_link: Any = None,
|
|
update_chat_link: bool = False,
|
|
) -> None:
|
|
async with self._sf() as session:
|
|
row = await session.get(AgentExtrasRow, agent_id)
|
|
now = datetime.now(UTC)
|
|
if row is None:
|
|
session.add(
|
|
AgentExtrasRow(
|
|
agent_id=agent_id,
|
|
tool_groups=tool_groups,
|
|
model=model,
|
|
chat_link=chat_link if update_chat_link else None,
|
|
skills_synced_at=now,
|
|
)
|
|
)
|
|
else:
|
|
row.tool_groups = tool_groups
|
|
row.model = model
|
|
if update_chat_link:
|
|
row.chat_link = chat_link
|
|
row.skills_synced_at = now
|
|
await session.commit()
|
|
|
|
async def get_extras_for(self, agent_ids: list[str]) -> dict[str, dict[str, Any]]:
|
|
if not agent_ids:
|
|
return {}
|
|
out: dict[str, dict[str, Any]] = {}
|
|
async with self._sf() as session:
|
|
extras = await session.execute(
|
|
select(AgentExtrasRow).where(AgentExtrasRow.agent_id.in_(agent_ids))
|
|
)
|
|
for row in extras.scalars():
|
|
out[row.agent_id] = {
|
|
"skills": [], # default; filled below if skill rows exist
|
|
"tool_groups": row.tool_groups,
|
|
"model": row.model,
|
|
"chat_link": row.chat_link,
|
|
}
|
|
skill_rows = await session.execute(
|
|
select(AgentSkillRow.agent_id, AgentSkillRow.skill_name)
|
|
.where(AgentSkillRow.agent_id.in_(agent_ids))
|
|
.order_by(AgentSkillRow.agent_id.asc(), AgentSkillRow.position.asc())
|
|
)
|
|
for agent_id, skill_name in skill_rows.all():
|
|
entry = out.setdefault(
|
|
agent_id,
|
|
{"skills": [], "tool_groups": None, "model": None, "chat_link": None},
|
|
)
|
|
entry["skills"].append(skill_name)
|
|
return out
|
|
|
|
async def delete_extras_for_agent(self, agent_id: str) -> None:
|
|
async with self._sf() as session:
|
|
await session.execute(delete(AgentSkillRow).where(AgentSkillRow.agent_id == agent_id))
|
|
await session.execute(delete(AgentExtrasRow).where(AgentExtrasRow.agent_id == agent_id))
|
|
await session.commit()
|
|
|
|
# --- 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)
|
|
.where(AgentFavoriteRow.user_id == user_id, AgentFavoriteRow.agent_id == agent_id)
|
|
.limit(1)
|
|
)
|
|
).scalar_one_or_none()
|
|
if existing is not None:
|
|
if existing.origin != "user" or existing.position_id is not None:
|
|
existing.origin = "user"
|
|
existing.position_id = None
|
|
await session.commit()
|
|
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]:
|
|
chain_agent_ids = await self._favorite_chain_exclusions(user_id)
|
|
async with self._sf() as session:
|
|
result = await session.execute(
|
|
select(AgentFavoriteRow.agent_id)
|
|
.join(AgentRow, AgentRow.id == AgentFavoriteRow.agent_id)
|
|
.where(
|
|
AgentFavoriteRow.user_id == user_id,
|
|
self._valid_favorite_filter(chain_agent_ids),
|
|
)
|
|
)
|
|
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()
|
|
|
|
# --- Pins ("置顶") ---
|
|
|
|
async def add_pin(self, user_id: str, agent_id: str) -> None:
|
|
async with self._sf() as session:
|
|
existing = (
|
|
await session.execute(
|
|
select(AgentPinRow.id)
|
|
.where(AgentPinRow.user_id == user_id, AgentPinRow.agent_id == agent_id)
|
|
.limit(1)
|
|
)
|
|
).scalar_one_or_none()
|
|
if existing is not None:
|
|
return
|
|
session.add(AgentPinRow(id=uuid4().hex, user_id=user_id, agent_id=agent_id))
|
|
try:
|
|
await session.commit()
|
|
except IntegrityError:
|
|
await session.rollback()
|
|
|
|
async def remove_pin(self, user_id: str, agent_id: str) -> bool:
|
|
async with self._sf() as session:
|
|
result = await session.execute(
|
|
delete(AgentPinRow).where(
|
|
AgentPinRow.user_id == user_id,
|
|
AgentPinRow.agent_id == agent_id,
|
|
)
|
|
)
|
|
await session.commit()
|
|
return (result.rowcount or 0) > 0
|
|
|
|
async def list_pin_ids(self, user_id: str) -> set[str]:
|
|
async with self._sf() as session:
|
|
result = await session.execute(
|
|
select(AgentPinRow.agent_id).where(AgentPinRow.user_id == user_id)
|
|
)
|
|
return {row[0] for row in result.all()}
|
|
|
|
async def delete_pins_for_agent(self, agent_id: str) -> None:
|
|
async with self._sf() as session:
|
|
await session.execute(delete(AgentPinRow).where(AgentPinRow.agent_id == agent_id))
|
|
await session.commit()
|