176 lines
7.7 KiB
Python
176 lines
7.7 KiB
Python
"""SQLAlchemy-backed stream-truncation event repository."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from datetime import UTC, datetime, timedelta
|
|
from typing import Any
|
|
|
|
from sqlalchemy import delete, desc, func, select
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker
|
|
|
|
from deerflow.persistence.stream_truncation.base import (
|
|
RETENTION_DAYS,
|
|
StreamTruncationEvent,
|
|
StreamTruncationQuery,
|
|
StreamTruncationStore,
|
|
)
|
|
from deerflow.persistence.stream_truncation.model import StreamTruncationEventRow
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class SqlStreamTruncationStore(StreamTruncationStore):
|
|
"""Persist truncation events with a hard one-month retention cap."""
|
|
|
|
def __init__(self, session_factory: async_sessionmaker) -> None:
|
|
self._session_factory = session_factory
|
|
|
|
async def _ensure_table(self, session) -> None:
|
|
# Table creation follows the repo's create_all-on-first-use convention
|
|
# (see e.g. task_buttons); safe to call repeatedly.
|
|
try:
|
|
await session.run_sync(
|
|
lambda sync_session: StreamTruncationEventRow.__table__.create(
|
|
sync_session.bind, checkfirst=True
|
|
)
|
|
)
|
|
except Exception:
|
|
logger.debug("stream_truncation table create skipped", exc_info=True)
|
|
|
|
async def record(self, event: StreamTruncationEvent) -> None:
|
|
try:
|
|
async with self._session_factory() as session:
|
|
await self._ensure_table(session)
|
|
session.add(
|
|
StreamTruncationEventRow(
|
|
id=event.id,
|
|
created_at=event.created_at,
|
|
model_name=event.model_name,
|
|
provider_completion_id=event.provider_completion_id,
|
|
run_id=event.run_id,
|
|
event_type=event.event_type,
|
|
strict_mode=event.strict_mode,
|
|
finish_reason=event.finish_reason,
|
|
status=event.status,
|
|
error_type=event.error_type,
|
|
chunk_count=event.chunk_count,
|
|
duration_ms=event.duration_ms,
|
|
first_chunk_ms=event.first_chunk_ms,
|
|
detail=event.detail,
|
|
)
|
|
)
|
|
await session.commit()
|
|
# Same session/transaction scope: enforce the retention cap on
|
|
# every insert. Truncations are rare, so this stays cheap and
|
|
# keeps the one-month promise without any scheduler.
|
|
await self._enforce_retention(session)
|
|
except Exception:
|
|
logger.warning(
|
|
"Failed to persist stream truncation event model=%s type=%s",
|
|
event.model_name,
|
|
event.event_type,
|
|
exc_info=True,
|
|
)
|
|
|
|
async def _enforce_retention(self, session) -> None:
|
|
"""Delete rows beyond the retention cap. Best-effort, never raises."""
|
|
try:
|
|
cutoff = datetime.now(UTC) - timedelta(days=RETENTION_DAYS)
|
|
await session.execute(
|
|
delete(StreamTruncationEventRow).where(
|
|
StreamTruncationEventRow.created_at < cutoff
|
|
)
|
|
)
|
|
await session.commit()
|
|
except Exception:
|
|
logger.debug("stream truncation retention sweep failed", exc_info=True)
|
|
|
|
def _apply_filters(self, query: StreamTruncationQuery):
|
|
stmt = select(StreamTruncationEventRow)
|
|
if query.model_name:
|
|
stmt = stmt.where(StreamTruncationEventRow.model_name == query.model_name)
|
|
if query.event_type:
|
|
stmt = stmt.where(StreamTruncationEventRow.event_type == query.event_type)
|
|
if query.since is not None:
|
|
stmt = stmt.where(StreamTruncationEventRow.created_at >= query.since)
|
|
if query.until is not None:
|
|
stmt = stmt.where(StreamTruncationEventRow.created_at <= query.until)
|
|
return stmt
|
|
|
|
async def list(self, query: StreamTruncationQuery) -> list[dict[str, Any]]:
|
|
try:
|
|
async with self._session_factory() as session:
|
|
await self._ensure_table(session)
|
|
stmt = self._apply_filters(query).order_by(
|
|
desc(StreamTruncationEventRow.created_at)
|
|
).limit(query.limit).offset(query.offset)
|
|
rows = (await session.execute(stmt)).scalars().all()
|
|
return [row.to_dict() for row in rows]
|
|
except Exception:
|
|
logger.warning("Failed to list stream truncation events", exc_info=True)
|
|
return []
|
|
|
|
async def count(self, query: StreamTruncationQuery) -> int:
|
|
try:
|
|
async with self._session_factory() as session:
|
|
await self._ensure_table(session)
|
|
stmt = self._apply_filters(query).with_only_columns(
|
|
func.count(StreamTruncationEventRow.id)
|
|
)
|
|
return int((await session.execute(stmt)).scalar() or 0)
|
|
except Exception:
|
|
logger.warning("Failed to count stream truncation events", exc_info=True)
|
|
return 0
|
|
|
|
async def model_breakdown(self, query: StreamTruncationQuery) -> list[dict[str, Any]]:
|
|
try:
|
|
async with self._session_factory() as session:
|
|
await self._ensure_table(session)
|
|
stmt = self._apply_filters(query).with_only_columns(
|
|
StreamTruncationEventRow.model_name,
|
|
StreamTruncationEventRow.event_type,
|
|
func.count(StreamTruncationEventRow.id).label("total"),
|
|
func.max(StreamTruncationEventRow.created_at).label("last_at"),
|
|
).group_by(StreamTruncationEventRow.model_name, StreamTruncationEventRow.event_type)
|
|
rows = (await session.execute(stmt)).all()
|
|
merged: dict[str, dict[str, Any]] = {}
|
|
for row in rows:
|
|
name = row.model_name or "(unknown)"
|
|
bucket = merged.setdefault(
|
|
name, {"model_name": name, "total": 0, "last_at": None}
|
|
)
|
|
bucket[row.event_type] = int(row.total or 0)
|
|
bucket["total"] += int(row.total or 0)
|
|
row_last = row.last_at
|
|
if row_last is not None and (
|
|
bucket["last_at"] is None or row_last > bucket["last_at"]
|
|
):
|
|
bucket["last_at"] = row_last
|
|
result = sorted(
|
|
merged.values(), key=lambda item: item["last_at"] or datetime.min.replace(tzinfo=UTC), reverse=True
|
|
)
|
|
for item in result:
|
|
if isinstance(item["last_at"], datetime):
|
|
item["last_at"] = item["last_at"].isoformat()
|
|
return result
|
|
except Exception:
|
|
logger.warning("Failed to aggregate stream truncation events", exc_info=True)
|
|
return []
|
|
|
|
async def purge_before(self, cutoff: datetime) -> int:
|
|
try:
|
|
async with self._session_factory() as session:
|
|
await self._ensure_table(session)
|
|
result = await session.execute(
|
|
delete(StreamTruncationEventRow).where(
|
|
StreamTruncationEventRow.created_at < cutoff
|
|
)
|
|
)
|
|
await session.commit()
|
|
rowcount = result.rowcount
|
|
return 0 if rowcount is None else int(rowcount)
|
|
except Exception:
|
|
logger.warning("Failed to purge stream truncation events", exc_info=True)
|
|
return 0
|