captcha_solver/solvers/router.py

"""Solver router: picks the right solver for the captcha type."""

from __future__ import annotations

import time
from typing import Optional

from captcha_solver.engines import (
    WhisperEngine,
    FlorenceEngine,
    MoondreamEngine,
    QwenEngine,
    OllamaEngine,
)
from captcha_solver.solvers.base import BaseSolver, SolveContext, SolveAttempt
from captcha_solver.solvers.math_solver import MathSolver
from captcha_solver.solvers.audio_solver import AudioSolver
from captcha_solver.solvers.text_ocr_solver import TextOcrSolver
from captcha_solver.solvers.image_grid_solver import ImageGridSolver
from captcha_solver.utils.detector import detect_type


class SolverRouter:
    """Routes a SolveRequest to the right solver and returns the result."""

    def __init__(self) -> None:
        self._engines: dict | None = None
        self._solvers: dict[str, BaseSolver] = {}

    def init(self) -> None:
        """Load all engines (lazy, in a controlled order)."""
        whisper = WhisperEngine()
        florence = FlorenceEngine()
        moondream = MoondreamEngine()
        qwen = QwenEngine()
        ollama = OllamaEngine()
        for eng in (whisper, florence, moondream, qwen, ollama):
            try:
                eng.load()
            except Exception:
                pass
        ctx = SolveContext(
            whisper=whisper,
            florence=florence,
            moondream=moondream,
            qwen=qwen,
            ollama=ollama,
        )
        self._engines = {
            "whisper": whisper,
            "florence2": florence,
            "moondream2": moondream,
            "qwen": qwen,
            "ollama": ollama,
        }
        for s in (MathSolver, AudioSolver, TextOcrSolver, ImageGridSolver):
            self._solvers[s.captcha_type] = s(ctx)

    def engine_status(self) -> dict[str, str]:
        if self._engines is None:
            return {}
        return {name: eng.status for name, eng in self._engines.items()}

    def solve(
        self,
        captcha_type: str,
        image_b64: Optional[str],
        audio_b64: Optional[str],
        hint: Optional[str] = None,
        min_confidence: float = 0.4,
    ) -> SolveAttempt:
        if self._engines is None:
            self.init()
        if captcha_type == "auto":
            captcha_type = detect_type(image_b64, audio_b64)
        solver = self._solvers.get(captcha_type)
        if solver is None:
            return SolveAttempt(
                answer="",
                confidence=0.0,
                solver_name="router.unknown",
                error=f"unknown captcha type: {captcha_type}",
            )
        solver.prepare(image_b64, audio_b64, hint)
        return solver.run_all(min_confidence=min_confidence)