captcha_solver/learning/collector.py

"""Hook that collects solve results and stores them in the learning DB.

Call `capture()` from any solver to record every attempt."""
from __future__ import annotations

import hashlib
import time
from typing import Any

from .db import LearningDB

_db = LearningDB()


def image_hash(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()[:16]


class Collector:
    """Collector that wraps a solver's solve method to record results."""

    def __init__(self, solver_name: str) -> None:
        self.solver_name = solver_name
        self._last_run: dict[str, Any] = {}

    def capture(
        self,
        captcha_type: str,
        answer: str | None,
        start_time: float,
        expected: str | None = None,
        correct: bool | None = None,
        confidence: float | None = None,
        hint: str | None = None,
        image_bytes: bytes | None = None,
        preprocess_steps: str | None = "",
        run_source: str = "api",
    ) -> int:
        latency = int((time.time() - start_time) * 1000)
        img_hash = image_hash(image_bytes) if image_bytes else None
        self._last_run = {
            "type": captcha_type,
            "answer": answer,
            "expected": expected,
            "correct": correct,
            "latency_ms": latency,
        }
        return _db.record_attempt(
            captcha_type=captcha_type,
            solver_used=self.solver_name,
            answer=answer,
            expected=expected,
            correct=correct,
            confidence=confidence,
            latency_ms=latency,
            hint=hint,
            image_hash=img_hash,
            preprocess_steps=preprocess_steps,
            run_source=run_source,
        )

    @property
    def last_run(self) -> dict[str, Any]:
        return self._last_run


def get_stats() -> dict:
    return _db.summary()


def get_solver_ranking(captcha_type: str | None = None) -> list[dict]:
    return _db.get_solver_ranking(captcha_type)


def get_best_solver(captcha_type: str) -> str | None:
    best = _db.get_best_solver(captcha_type, min_samples=3)
    return best["solver_name"] if best else None