captcha_solver/solvers/base.py

"""Base solver class."""

from __future__ import annotations

import abc
import time
from dataclasses import dataclass, field
from typing import Optional

from captcha_solver.engines import (
    WhisperEngine,
    FlorenceEngine,
    MoondreamEngine,
    QwenEngine,
    OllamaEngine,
)


@dataclass
class SolveAttempt:
    """One solver strategy result."""

    answer: str
    confidence: float
    solver_name: str
    elapsed_ms: int = 0
    error: Optional[str] = None
    metadata: dict = field(default_factory=dict)


@dataclass
class SolveContext:
    """Shared resources passed to every solver."""

    whisper: WhisperEngine
    florence: FlorenceEngine
    moondream: MoondreamEngine
    qwen: QwenEngine
    ollama: OllamaEngine


class BaseSolver(abc.ABC):
    """Abstract captcha solver.

    Each solver is registered with the router. `name` is unique
    per solver. `attempts` returns an ordered list of strategies to try;
    the first one that yields a confident answer wins.
    """

    name: str = "base"
    captcha_type: str = "base"

    def __init__(self, ctx: SolveContext) -> None:
        self.ctx = ctx

    @abc.abstractmethod
    def attempts(self) -> list[callable]:
        """Return a list of zero-arg callables, each producing a SolveAttempt.

        Each callable should be self-contained: catch its own errors, set
        `error` on the attempt if it failed, and return a result. The
        router picks the first confident (>= min_confidence) success.
        """
        raise NotImplementedError

    def run_all(self, min_confidence: float = 0.4) -> SolveAttempt:
        """Run every strategy, return the first confident one.

        On no confident result, returns the highest-confidence attempt
        (even if it failed). Never raises.
        """
        best: Optional[SolveAttempt] = None
        for fn in self.attempts():
            t0 = time.time()
            try:
                attempt = fn()
            except Exception as exc:
                attempt = SolveAttempt(
                    answer="",
                    confidence=0.0,
                    solver_name=f"{self.name}.{fn.__name__}",
                    elapsed_ms=int((time.time() - t0) * 1000),
                    error=str(exc),
                )
            attempt.elapsed_ms = int((time.time() - t0) * 1000)
            attempt.solver_name = f"{self.name}.{fn.__name__}"
            if attempt.answer and attempt.confidence >= min_confidence:
                return attempt
            if best is None or attempt.confidence > best.confidence:
                best = attempt
        return best or SolveAttempt(
            answer="",
            confidence=0.0,
            solver_name=f"{self.name}.none",
            error="no attempts",
        )