mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-27 21:50:36 +00:00
4e9f6c68ea
* fix(persistence): retry feedback and conversation upserts that lose a race FeedbackRepository.upsert and ChannelConnectionRepository.set_thread_id looked up the row and inserted it when missing, with no conflict handling. Two first writers for one key both missed the lookup, and the second INSERT violated uq_feedback_thread_run_user or uq_channel_conversation_connection_external: a 500 for the second of two concurrent first ratings, and a failed message when a conversation's first message raced /new. Both now catch IntegrityError, roll back and retry, so the losing writer re-reads the committed row and updates it, as upsert_connection already does. A retry rather than ON CONFLICT keeps the NULL-user feedback path updating its single row, since a NULL key never conflicts. * docs(changelog): reference #6852 for the upsert race fix * test(persistence): pin upsert retry exhaustion and share the miss-select double Both retry loops now have a test that makes every SELECT miss against a pre-seeded row, so all attempts hit the unique key: the original IntegrityError must propagate after exactly _UPSERT_MAX_ATTEMPTS attempts and the stored row must stay unchanged. The "first SELECT misses" session double lived in four copies (one per retry test, including upsert_connection's). It now lives once in tests/_repository_race_helpers.py, with a misses count so the same double drives both the single lost race and exhaustion.
296 lines
12 KiB
Python
296 lines
12 KiB
Python
"""SQLAlchemy-backed feedback storage.
|
|
|
|
Each method acquires its own short-lived session.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import uuid
|
|
from datetime import UTC, datetime
|
|
|
|
from sqlalchemy import case, delete, func, select
|
|
from sqlalchemy.exc import IntegrityError
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
|
|
|
from deerflow.persistence.feedback.model import FeedbackRow
|
|
from deerflow.runtime.user_context import AUTO, _AutoSentinel, resolve_user_id
|
|
from deerflow.utils.time import coerce_iso
|
|
|
|
_UPSERT_MAX_ATTEMPTS = 3
|
|
|
|
|
|
class FeedbackRepository:
|
|
def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
|
|
self._sf = session_factory
|
|
|
|
@staticmethod
|
|
def _row_to_dict(row: FeedbackRow) -> dict:
|
|
d = row.to_dict()
|
|
val = d.get("created_at")
|
|
if isinstance(val, datetime):
|
|
# SQLite drops tzinfo on read; normalize via ``coerce_iso`` so output is always tz-aware.
|
|
d["created_at"] = coerce_iso(val)
|
|
return d
|
|
|
|
async def create(
|
|
self,
|
|
*,
|
|
run_id: str,
|
|
thread_id: str,
|
|
rating: int,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
message_id: str | None = None,
|
|
comment: str | None = None,
|
|
) -> dict:
|
|
"""Create a feedback record. rating must be +1 or -1."""
|
|
if rating not in (1, -1):
|
|
raise ValueError(f"rating must be +1 or -1, got {rating}")
|
|
resolved_user_id = resolve_user_id(user_id, method_name="FeedbackRepository.create")
|
|
row = FeedbackRow(
|
|
feedback_id=str(uuid.uuid4()),
|
|
run_id=run_id,
|
|
thread_id=thread_id,
|
|
user_id=resolved_user_id,
|
|
message_id=message_id,
|
|
rating=rating,
|
|
comment=comment,
|
|
created_at=datetime.now(UTC),
|
|
)
|
|
async with self._sf() as session:
|
|
session.add(row)
|
|
await session.commit()
|
|
await session.refresh(row)
|
|
return self._row_to_dict(row)
|
|
|
|
async def get(
|
|
self,
|
|
feedback_id: str,
|
|
*,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
) -> dict | None:
|
|
resolved_user_id = resolve_user_id(user_id, method_name="FeedbackRepository.get")
|
|
async with self._sf() as session:
|
|
row = await session.get(FeedbackRow, feedback_id)
|
|
if row is None:
|
|
return None
|
|
if resolved_user_id is not None and row.user_id != resolved_user_id:
|
|
return None
|
|
return self._row_to_dict(row)
|
|
|
|
async def list_by_run(
|
|
self,
|
|
thread_id: str,
|
|
run_id: str,
|
|
*,
|
|
limit: int = 100,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
) -> list[dict]:
|
|
resolved_user_id = resolve_user_id(user_id, method_name="FeedbackRepository.list_by_run")
|
|
stmt = select(FeedbackRow).where(FeedbackRow.thread_id == thread_id, FeedbackRow.run_id == run_id)
|
|
if resolved_user_id is not None:
|
|
stmt = stmt.where(FeedbackRow.user_id == resolved_user_id)
|
|
stmt = stmt.order_by(FeedbackRow.created_at.asc()).limit(limit)
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
return [self._row_to_dict(r) for r in result.scalars()]
|
|
|
|
async def list_by_thread(
|
|
self,
|
|
thread_id: str,
|
|
*,
|
|
limit: int = 100,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
) -> list[dict]:
|
|
resolved_user_id = resolve_user_id(user_id, method_name="FeedbackRepository.list_by_thread")
|
|
stmt = select(FeedbackRow).where(FeedbackRow.thread_id == thread_id)
|
|
if resolved_user_id is not None:
|
|
stmt = stmt.where(FeedbackRow.user_id == resolved_user_id)
|
|
stmt = stmt.order_by(FeedbackRow.created_at.asc()).limit(limit)
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
return [self._row_to_dict(r) for r in result.scalars()]
|
|
|
|
async def delete(
|
|
self,
|
|
feedback_id: str,
|
|
*,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
) -> bool:
|
|
resolved_user_id = resolve_user_id(user_id, method_name="FeedbackRepository.delete")
|
|
async with self._sf() as session:
|
|
row = await session.get(FeedbackRow, feedback_id)
|
|
if row is None:
|
|
return False
|
|
if resolved_user_id is not None and row.user_id != resolved_user_id:
|
|
return False
|
|
await session.delete(row)
|
|
await session.commit()
|
|
return True
|
|
|
|
async def upsert(
|
|
self,
|
|
*,
|
|
run_id: str,
|
|
thread_id: str,
|
|
rating: int,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
comment: str | None = None,
|
|
) -> dict:
|
|
"""Create or update feedback for (thread_id, run_id, user_id). rating must be +1 or -1.
|
|
|
|
Two first submissions for one key can both miss the SELECT, and the
|
|
later INSERT then violates ``uq_feedback_thread_run_user``. That writer
|
|
rolls back and retries, so it re-reads the committed row and updates it
|
|
instead of failing the request.
|
|
"""
|
|
if rating not in (1, -1):
|
|
raise ValueError(f"rating must be +1 or -1, got {rating}")
|
|
resolved_user_id = resolve_user_id(user_id, method_name="FeedbackRepository.upsert")
|
|
stmt = select(FeedbackRow).where(
|
|
FeedbackRow.thread_id == thread_id,
|
|
FeedbackRow.run_id == run_id,
|
|
FeedbackRow.user_id == resolved_user_id,
|
|
)
|
|
async with self._sf() as session:
|
|
last_error: IntegrityError | None = None
|
|
for _ in range(_UPSERT_MAX_ATTEMPTS):
|
|
try:
|
|
row = (await session.execute(stmt)).scalar_one_or_none()
|
|
if row is not None:
|
|
row.rating = rating
|
|
row.comment = comment
|
|
row.created_at = datetime.now(UTC)
|
|
else:
|
|
row = FeedbackRow(
|
|
feedback_id=str(uuid.uuid4()),
|
|
run_id=run_id,
|
|
thread_id=thread_id,
|
|
user_id=resolved_user_id,
|
|
rating=rating,
|
|
comment=comment,
|
|
created_at=datetime.now(UTC),
|
|
)
|
|
session.add(row)
|
|
await session.commit()
|
|
await session.refresh(row)
|
|
return self._row_to_dict(row)
|
|
except IntegrityError as exc:
|
|
last_error = exc
|
|
await session.rollback()
|
|
raise last_error # type: ignore[misc] # loop runs at least once
|
|
|
|
async def delete_by_run(
|
|
self,
|
|
*,
|
|
thread_id: str,
|
|
run_id: str,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
) -> bool:
|
|
"""Delete the current user's feedback for a run. Returns True if a record was deleted."""
|
|
resolved_user_id = resolve_user_id(user_id, method_name="FeedbackRepository.delete_by_run")
|
|
async with self._sf() as session:
|
|
stmt = select(FeedbackRow).where(
|
|
FeedbackRow.thread_id == thread_id,
|
|
FeedbackRow.run_id == run_id,
|
|
FeedbackRow.user_id == resolved_user_id,
|
|
)
|
|
result = await session.execute(stmt)
|
|
row = result.scalar_one_or_none()
|
|
if row is None:
|
|
return False
|
|
await session.delete(row)
|
|
await session.commit()
|
|
return True
|
|
|
|
async def delete_by_thread(
|
|
self,
|
|
thread_id: str,
|
|
*,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
) -> int:
|
|
"""Delete the owner's feedback for every run of a thread.
|
|
|
|
``user_id`` keeps the repository's three-state convention: ``AUTO``
|
|
resolves the request context, an explicit id scopes the delete to that
|
|
owner, and ``None`` removes every owner's rows (migration/CLI callers).
|
|
"""
|
|
resolved_user_id = resolve_user_id(user_id, method_name="FeedbackRepository.delete_by_thread")
|
|
|
|
conditions = [FeedbackRow.thread_id == thread_id]
|
|
if resolved_user_id is not None:
|
|
conditions.append(FeedbackRow.user_id == resolved_user_id)
|
|
|
|
async with self._sf() as session:
|
|
count = await session.scalar(select(func.count()).select_from(FeedbackRow).where(*conditions)) or 0
|
|
if count:
|
|
await session.execute(delete(FeedbackRow).where(*conditions))
|
|
await session.commit()
|
|
return count
|
|
|
|
async def list_by_thread_grouped(
|
|
self,
|
|
thread_id: str,
|
|
*,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
) -> dict[str, dict]:
|
|
"""Return feedback grouped by run_id for a thread: {run_id: feedback_dict}.
|
|
|
|
With an explicit ``None`` user id (unfiltered reads) several users may
|
|
hold feedback on the same run, so order deterministically — the
|
|
per-run collapse below keeps the last row per ``run_id``, i.e. the
|
|
most recently written feedback (``created_at`` is refreshed on
|
|
update), with ``feedback_id`` breaking ties.
|
|
"""
|
|
resolved_user_id = resolve_user_id(user_id, method_name="FeedbackRepository.list_by_thread_grouped")
|
|
stmt = select(FeedbackRow).where(FeedbackRow.thread_id == thread_id)
|
|
if resolved_user_id is not None:
|
|
stmt = stmt.where(FeedbackRow.user_id == resolved_user_id)
|
|
stmt = stmt.order_by(FeedbackRow.created_at.asc(), FeedbackRow.feedback_id.asc())
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
return {row.run_id: self._row_to_dict(row) for row in result.scalars()}
|
|
|
|
async def list_by_run_ids(
|
|
self,
|
|
thread_id: str,
|
|
run_ids: set[str],
|
|
*,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
) -> dict[str, dict]:
|
|
"""Return feedback for only the selected runs in one thread.
|
|
|
|
Same deterministic ordering as :meth:`list_by_thread_grouped`: with an
|
|
explicit ``None`` user id the per-run collapse keeps the most recently
|
|
written feedback (``created_at`` is refreshed on update), ties broken
|
|
by ``feedback_id``.
|
|
"""
|
|
if not run_ids:
|
|
return {}
|
|
resolved_user_id = resolve_user_id(user_id, method_name="FeedbackRepository.list_by_run_ids")
|
|
stmt = select(FeedbackRow).where(
|
|
FeedbackRow.thread_id == thread_id,
|
|
FeedbackRow.run_id.in_(run_ids),
|
|
)
|
|
if resolved_user_id is not None:
|
|
stmt = stmt.where(FeedbackRow.user_id == resolved_user_id)
|
|
stmt = stmt.order_by(FeedbackRow.created_at.asc(), FeedbackRow.feedback_id.asc())
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
return {row.run_id: self._row_to_dict(row) for row in result.scalars()}
|
|
|
|
async def aggregate_by_run(self, thread_id: str, run_id: str) -> dict:
|
|
"""Aggregate feedback stats for a run using database-side counting."""
|
|
stmt = select(
|
|
func.count().label("total"),
|
|
func.coalesce(func.sum(case((FeedbackRow.rating == 1, 1), else_=0)), 0).label("positive"),
|
|
func.coalesce(func.sum(case((FeedbackRow.rating == -1, 1), else_=0)), 0).label("negative"),
|
|
).where(FeedbackRow.thread_id == thread_id, FeedbackRow.run_id == run_id)
|
|
async with self._sf() as session:
|
|
row = (await session.execute(stmt)).one()
|
|
return {
|
|
"run_id": run_id,
|
|
"total": row.total,
|
|
"positive": row.positive,
|
|
"negative": row.negative,
|
|
}
|