agent/tools/eval_progress.py - master - Vibek/Agent_GG

agent/tools/eval_progress.py
185 строк5 KB
Vibek
прогресс бар оценки добавлен и работает корректно
24 апр 2026, 19:10
24 апр 2026, 19:101e17055
from __future__ import annotations

from contextlib import contextmanager
from contextvars import ContextVar
from dataclasses import dataclass, field
from threading import Lock
from time import time
from typing import Any

_TTL_SECONDS = 15 * 60
_active_progress_id: ContextVar[str | None] = ContextVar("gg_eval_progress_id", default=None)


def _now_ts() -> float:
    return time()


def _clamp_percent(value: int | float | None) -> int:
    if value is None:
        return 0
    return max(0, min(100, int(round(value))))


@dataclass(slots=True)
class EvaluationProgressState:
    progress_id: str
    active: bool = True
    status: str = "running"
    percent: int = 5
    stage: str = "Запуск оценки"
    detail: str = "Подготавливаю данные для аудита объекта..."
    created_at: float = field(default_factory=_now_ts)
    updated_at: float = field(default_factory=_now_ts)
    meta: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        return {
            "active": self.active,
            "status": self.status,
            "percent": self.percent,
            "stage": self.stage,
            "detail": self.detail,
            "created_at": self.created_at,
            "updated_at": self.updated_at,
        }


class EvaluationProgressTracker:
    def __init__(self) -> None:
        self._items: dict[str, EvaluationProgressState] = {}
        self._lock = Lock()

    def _cleanup_locked(self, now_ts: float) -> None:
        stale_ids = [
            progress_id
            for progress_id, state in self._items.items()
            if (now_ts - state.updated_at) > _TTL_SECONDS
        ]
        for progress_id in stale_ids:
            self._items.pop(progress_id, None)

    def get(self, progress_id: str) -> dict[str, Any]:
        now_ts = _now_ts()
        with self._lock:
            self._cleanup_locked(now_ts)
            state = self._items.get(progress_id)
            if state is None:
                return {
                    "active": False,
                    "status": "idle",
                    "percent": 0,
                    "stage": "",
                    "detail": "",
                    "created_at": now_ts,
                    "updated_at": now_ts,
                }
            return state.to_dict()

    def start(self, progress_id: str, *, listing_id: int) -> None:
        now_ts = _now_ts()
        with self._lock:
            self._cleanup_locked(now_ts)
            self._items[progress_id] = EvaluationProgressState(
                progress_id=progress_id,
                percent=5,
                stage="Запуск оценки",
                detail=f"Инициализирую аудит для объекта #{listing_id}.",
                created_at=now_ts,
                updated_at=now_ts,
                meta={"listing_id": listing_id},
            )

    def _update(self, progress_id: str, **kwargs: Any) -> None:
        now_ts = _now_ts()
        with self._lock:
            self._cleanup_locked(now_ts)
            state = self._items.get(progress_id)
            if state is None:
                return
            percent = kwargs.get("percent")
            if percent is not None:
                state.percent = _clamp_percent(percent)
            stage = kwargs.get("stage")
            if stage:
                state.stage = stage
            detail = kwargs.get("detail")
            if detail:
                state.detail = detail
            status = kwargs.get("status")
            if status:
                state.status = status
            meta = kwargs.get("meta")
            if isinstance(meta, dict):
                state.meta.update(meta)
            active = kwargs.get("active")
            if active is not None:
                state.active = bool(active)
            state.updated_at = now_ts

    def update(
        self,
        progress_id: str,
        *,
        percent: int | float | None = None,
        stage: str | None = None,
        detail: str | None = None,
        status: str | None = None,
        meta: dict[str, Any] | None = None,
    ) -> None:
        kwargs: dict[str, Any] = {}
        if percent is not None:
            kwargs["percent"] = percent
        if stage:
            kwargs["stage"] = stage
        if detail:
            kwargs["detail"] = detail
        if status:
            kwargs["status"] = status
        if meta is not None:
            kwargs["meta"] = meta
        if kwargs:
            self._update(progress_id, **kwargs)

    def complete(self, progress_id: str, *, detail: str) -> None:
        self._update(
            progress_id,
            percent=100,
            stage="Оценка завершена",
            detail=detail,
            status="completed",
            active=True,
        )

    def fail(self, progress_id: str, *, detail: str) -> None:
        self._update(
            progress_id,
            percent=100,
            stage="Ошибка оценки",
            detail=detail,
            status="failed",
            active=True,
        )


_TRACKER = EvaluationProgressTracker()


def get_eval_progress_tracker() -> EvaluationProgressTracker:
    return _TRACKER


def get_active_eval_progress_id() -> str | None:
    return _active_progress_id.get()


@contextmanager
def bind_eval_progress(progress_id: str | None):
    if not progress_id:
        yield
        return
    token = _active_progress_id.set(progress_id)
    try:
        yield
    finally:
        _active_progress_id.reset(token)