File size: 2,865 Bytes
2edb151 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | 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)
|