backends/ollama.py

from __future__ import annotations

import base64
from pathlib import Path

import httpx

from app.config import Settings
from backends.base import OCRResult
from backends.openai_compat import OpenAICompatEmbed, OpenAICompatLLM


def _is_lightning(model: str) -> bool:
    return "lightning" in model.lower()


class OllamaLLM(OpenAICompatLLM):
    def __init__(self, settings: Settings, *, client: httpx.Client | None = None) -> None:
        accepts = settings.llm_accepts_images and not _is_lightning(settings.llm_model)
        super().__init__(
            settings,
            name="ollama",
            accepts_images=accepts,
            extra_body={},
            client=client,
        )


class OllamaEmbed(OpenAICompatEmbed):
    def __init__(self, settings: Settings, *, client: httpx.Client | None = None) -> None:
        super().__init__(settings, name="ollama-embed", client=client)


class OllamaOCR:
    name = "ollama-ocr"

    def __init__(self, settings: Settings, *, client: httpx.Client | None = None) -> None:
        self.settings = settings
        self._client = client

    def ocr(self, path: Path) -> OCRResult:
        if not self.settings.ocr_model:
            raise NotImplementedError("set RECEIPT_OCR_MODEL for Ollama vision OCR")
        jpeg = path.read_bytes()
        b64 = base64.b64encode(jpeg).decode("ascii")
        own = self._client is None
        http = self._client or httpx.Client(
            base_url=self.settings.ocr_base_url.rstrip("/"),
            timeout=self.settings.llm_timeout_s,
            headers={"Authorization": f"Bearer {self.settings.ocr_api_key}"},
        )
        try:
            response = http.post(
                "/chat/completions",
                json={
                    "model": self.settings.ocr_model,
                    "messages": [
                        {
                            "role": "user",
                            "content": [
                                {
                                    "type": "image_url",
                                    "image_url": {
                                        "url": f"data:image/jpeg;base64,{b64}"
                                    },
                                },
                                {
                                    "type": "text",
                                    "text": "Transcribe this document verbatim.",
                                },
                            ],
                        }
                    ],
                    "temperature": 0,
                    "max_tokens": 4096,
                },
            )
            response.raise_for_status()
            text = response.json()["choices"][0]["message"]["content"]
        finally:
            if own:
                http.close()
        return OCRResult(text=text or "", engine=self.name)