deerflow-code/offline-backend-20260512/backend/packages/harness/deerflow/persistence/tool_metrics/sql.py

"""SQLAlchemy-backed tool metrics repository."""

from __future__ import annotations

import logging
from collections.abc import AsyncIterator
from datetime import UTC, datetime
from typing import Any

from sqlalchemy import desc, func, select
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

from deerflow.persistence.tool_metrics.base import ToolMetricRecord, ToolMetricsQuery, ToolMetricsStore
from deerflow.persistence.tool_metrics.model import ToolCallMetricRow

logger = logging.getLogger(__name__)


def _normalize_created_at(value: datetime) -> datetime:
    return value if value.tzinfo else value.replace(tzinfo=UTC)


class SqlToolMetricsStore(ToolMetricsStore):
    """SQLite/Postgres/MySQL-backed implementation; writes are fire-and-forget."""

    def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
        self._sf = session_factory

    async def record(self, metric: ToolMetricRecord) -> None:
        row = ToolCallMetricRow(
            id=metric.id,
            created_at=_normalize_created_at(metric.created_at),
            user_id=metric.user_id,
            thread_id=metric.thread_id,
            run_id=metric.run_id,
            agent_name=metric.agent_name,
            tool_name=metric.tool_name,
            skill_name=metric.skill_name,
            duration_ms=metric.duration_ms,
            status=metric.status,
            error_category=metric.error_category,
            error_message=metric.error_message,
        )
        try:
            async with self._sf() as session:
                session.add(row)
                await session.commit()
        except Exception:
            # Metrics are best-effort; never surface a write failure to the runtime.
            logger.warning("Failed to persist tool_call_metrics row %s", metric.id, exc_info=True)

    def _apply_filters(self, stmt, query: ToolMetricsQuery):
        if query.user_id:
            stmt = stmt.where(ToolCallMetricRow.user_id == query.user_id)
        if query.tool_name:
            stmt = stmt.where(ToolCallMetricRow.tool_name == query.tool_name)
        if query.skill_only:
            stmt = stmt.where(ToolCallMetricRow.skill_name.is_not(None))
        if query.skill_name:
            stmt = stmt.where(ToolCallMetricRow.skill_name == query.skill_name)
        if query.status:
            stmt = stmt.where(ToolCallMetricRow.status == query.status)
        if query.error_category:
            stmt = stmt.where(ToolCallMetricRow.error_category == query.error_category)
        if query.since is not None:
            stmt = stmt.where(ToolCallMetricRow.created_at >= _normalize_created_at(query.since))
        if query.until is not None:
            stmt = stmt.where(ToolCallMetricRow.created_at <= _normalize_created_at(query.until))
        return stmt

    async def list(self, query: ToolMetricsQuery) -> list[dict[str, Any]]:
        stmt = select(ToolCallMetricRow)
        stmt = self._apply_filters(stmt, query)
        stmt = stmt.order_by(desc(ToolCallMetricRow.created_at)).limit(max(1, min(query.limit, 1000))).offset(max(0, query.offset))
        async with self._sf() as session:
            result = await session.execute(stmt)
            return [row.to_dict() for row in result.scalars()]

    async def count(self, query: ToolMetricsQuery) -> int:
        stmt = select(func.count()).select_from(ToolCallMetricRow)
        stmt = self._apply_filters(stmt, query)
        async with self._sf() as session:
            return (await session.scalar(stmt)) or 0

    async def iter_all(self, query: ToolMetricsQuery) -> AsyncIterator[dict[str, Any]]:
        page_size = 500
        offset = 0
        while True:
            page = await self.list(
                ToolMetricsQuery(
                    user_id=query.user_id,
                    tool_name=query.tool_name,
                    skill_only=query.skill_only,
                    skill_name=query.skill_name,
                    status=query.status,
                    error_category=query.error_category,
                    since=query.since,
                    until=query.until,
                    limit=page_size,
                    offset=offset,
                )
            )
            if not page:
                return
            for item in page:
                yield item
            if len(page) < page_size:
                return
            offset += page_size

    async def skill_breakdown(self, query: ToolMetricsQuery) -> list[dict[str, Any]]:
        # Count success/error per skill. Force the skill-only filter so rows
        # without a skill_name (plain tool calls) never leak into the result.
        scoped = ToolMetricsQuery(
            user_id=query.user_id,
            skill_only=True,
            skill_name=query.skill_name,
            since=query.since,
            until=query.until,
        )
        stmt = select(
            ToolCallMetricRow.skill_name,
            ToolCallMetricRow.status,
            func.count().label("n"),
        )
        stmt = self._apply_filters(stmt, scoped).group_by(
            ToolCallMetricRow.skill_name, ToolCallMetricRow.status
        )
        agg: dict[str, dict[str, int]] = {}
        async with self._sf() as session:
            for skill_name, status, n in (await session.execute(stmt)).all():
                bucket = agg.setdefault(skill_name, {"total": 0, "success": 0, "error": 0})
                bucket["total"] += n
                bucket["success" if status == "success" else "error"] += n
        return [
            {"skill_name": name, **counts}
            for name, counts in sorted(agg.items(), key=lambda kv: kv[1]["total"], reverse=True)
        ]