backends/base.py

from __future__ import annotations

from pathlib import Path
from typing import Literal, Protocol, runtime_checkable

InputType = Literal["query", "passage"]


class OCRResult:
    __slots__ = ("text", "engine")

    def __init__(self, text: str, engine: str) -> None:
        self.text = text
        self.engine = engine


@runtime_checkable
class LLMBackend(Protocol):
    accepts_images: bool
    name: str

    def complete_json(
        self,
        *,
        system: str,
        user: str,
        image_jpeg: bytes | None = None,
    ) -> str: ...

    def health(self) -> bool: ...


@runtime_checkable
class EmbedBackend(Protocol):
    name: str
    dim: int

    def embed(self, texts: list[str], *, input_type: InputType) -> list[list[float]]: ...

    def health(self) -> bool: ...


@runtime_checkable
class OCRBackend(Protocol):
    name: str

    def ocr(self, path: Path) -> OCRResult: ...