From 081084816b41e70fc19171e5fbd697efc2a3d9de Mon Sep 17 00:00:00 2001 From: Florian Egger Date: Fri, 21 Aug 2026 16:44:14 +0200 Subject: [PATCH] Improve VLM client lifecycle, retries, metrics, and cache keys - Adds abstract aclose() and shared load_prompt() helper to BaseVLMClient. - OllamaClient now logs latency, tokens, and prompt hash per call. - Retries cover HTTP 429 and 5xx in addition to network/timeout errors. - Cache key now includes response_format and max_tokens. - Warns when JSON schema generation falls back to plain 'json' format. - Adds FakeVLMClient.aclose() and cache unit tests. --- src/odoo_ocr/clients/base.py | 24 ++++++++++++++++++- src/odoo_ocr/clients/ollama.py | 39 +++++++++++++++++++++++++++--- src/odoo_ocr/utils/cache.py | 33 ++++++++++++++++++------- tests/conftest.py | 4 ++++ tests/test_cache.py | 44 ++++++++++++++++++++++++++++++++++ 5 files changed, 132 insertions(+), 12 deletions(-) create mode 100644 tests/test_cache.py diff --git a/src/odoo_ocr/clients/base.py b/src/odoo_ocr/clients/base.py index a09c3d6..f9c264f 100644 --- a/src/odoo_ocr/clients/base.py +++ b/src/odoo_ocr/clients/base.py @@ -35,6 +35,11 @@ class BaseVLMClient(ABC): """Send a chat request with optional images and return parsed text.""" ... + @abstractmethod + async def aclose(self) -> None: + """Release client resources (connections, file handles, ...).""" + ... + @staticmethod def parse_json(content: str, model: type[T]) -> T: """Strip markdown fences and parse/validate JSON against a Pydantic model.""" @@ -54,7 +59,24 @@ class BaseVLMClient(ABC): @staticmethod def _load_prompt_file(path: Path | str) -> str: - """Load a plain-text prompt from the prompts directory.""" + """Load a plain-text prompt file.""" path = Path(path) with path.open("r", encoding="utf-8") as f: return f.read() + + +def load_prompt(name: str) -> str: + """Load a plain-text system prompt from the project's ``prompts/`` directory. + + Tries the current working directory first (running from the repo root), + then falls back to the repository root derived from this package's + location, so the CLI keeps working from any working directory. + """ + candidates = [ + Path("prompts") / name, + Path(__file__).resolve().parents[3] / "prompts" / name, + ] + for candidate in candidates: + if candidate.is_file(): + return candidate.read_text(encoding="utf-8") + raise FileNotFoundError(f"Prompt file not found: {name!r} (tried: {candidates})") diff --git a/src/odoo_ocr/clients/ollama.py b/src/odoo_ocr/clients/ollama.py index ad05dd1..23993f4 100644 --- a/src/odoo_ocr/clients/ollama.py +++ b/src/odoo_ocr/clients/ollama.py @@ -3,6 +3,7 @@ from __future__ import annotations import base64 +import hashlib import io import logging import time @@ -13,7 +14,7 @@ from PIL import Image from pydantic import BaseModel from tenacity import ( retry, - retry_if_exception_type, + retry_if_exception, stop_after_attempt, wait_exponential, ) @@ -26,6 +27,16 @@ from odoo_ocr.utils import cache logger = logging.getLogger(__name__) +def _is_transient_retryable(exc: BaseException) -> bool: + """Retry network errors, timeouts, HTTP 429, and server errors (5xx).""" + if isinstance(exc, (httpx.NetworkError, httpx.TimeoutException)): + return True + if isinstance(exc, httpx.HTTPStatusError): + status = exc.response.status_code + return status == 429 or status >= 500 + return False + + class OllamaClient(BaseVLMClient): """Client for Ollama's /api/chat endpoint with vision support.""" @@ -35,6 +46,10 @@ class OllamaClient(BaseVLMClient): self.timeout = httpx.Timeout(120.0, connect=10.0) self.http = httpx.AsyncClient(timeout=self.timeout) + async def aclose(self) -> None: + """Release the underlying HTTP client.""" + await self.http.aclose() + @staticmethod def _encode_image(image: Image.Image, quality: int = 85) -> str: """Convert a PIL image to a base64-encoded PNG or JPEG string.""" @@ -62,6 +77,7 @@ class OllamaClient(BaseVLMClient): model_name = model or self.settings.models.ocr encoded_images = [self._encode_image(img) for img in images] image_blobs = [base64.b64decode(enc) for enc in encoded_images] + format_key = response_format.__name__ if response_format is not None else "text" cached = cache.get_cached( system_prompt, @@ -69,7 +85,9 @@ class OllamaClient(BaseVLMClient): image_blobs, model_name, temperature, + max_tokens, self.settings, + format_key=format_key, ) if cached: return VLMResponse.model_validate(cached) @@ -93,7 +111,12 @@ class OllamaClient(BaseVLMClient): if response_format is not None: try: body["format"] = response_format.model_json_schema() - except Exception: + except Exception as exc: + logger.warning( + "Could not build JSON schema for %s, falling back to 'json' format: %s", + response_format.__name__, + exc, + ) body["format"] = "json" start = time.perf_counter() @@ -118,15 +141,25 @@ class OllamaClient(BaseVLMClient): image_blobs, model_name, temperature, + max_tokens, vlm_response.model_dump(mode="json"), self.settings, + format_key=format_key, + ) + logger.info( + "Ollama %s: %d ms, prompt_tokens=%s, completion_tokens=%s, prompt_hash=%s", + model_name, + duration_ms, + vlm_response.prompt_tokens, + vlm_response.completion_tokens, + hashlib.sha256((system_prompt + user_prompt).encode("utf-8")).hexdigest(), ) return vlm_response @retry( stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=1, max=10), - retry=retry_if_exception_type((httpx.NetworkError, httpx.TimeoutException)), + retry=retry_if_exception(_is_transient_retryable), ) async def _post_chat(self, body: dict[str, Any]) -> dict[str, Any]: """Send the chat request and validate the response.""" diff --git a/src/odoo_ocr/utils/cache.py b/src/odoo_ocr/utils/cache.py index 43accd4..7c556bd 100644 --- a/src/odoo_ocr/utils/cache.py +++ b/src/odoo_ocr/utils/cache.py @@ -26,21 +26,37 @@ def _cache_path(settings: Settings, key: str) -> Path: return cache_dir / f"{key}.json" +def _cache_key_parts( + system_prompt: str, + user_prompt: str, + image_blobs: list[bytes], + model: str, + temperature: float, + max_tokens: int, + format_key: str, +) -> list[str]: + parts = [system_prompt, user_prompt, model, str(temperature), str(max_tokens), format_key] + for blob in image_blobs: + parts.append(hashlib.sha256(blob).hexdigest()) + return parts + + def get_cached( system_prompt: str, user_prompt: str, image_blobs: list[bytes], model: str, temperature: float, + max_tokens: int, settings: Settings, + format_key: str = "text", ) -> dict[str, Any] | None: """Return cached response dict if it exists, otherwise None.""" if not settings.cache.enabled: return None - parts = [system_prompt, user_prompt, model, str(temperature)] - for blob in image_blobs: - parts.append(hashlib.sha256(blob).hexdigest()) - key = _make_key(parts) + key = _make_key( + _cache_key_parts(system_prompt, user_prompt, image_blobs, model, temperature, max_tokens, format_key) + ) path = _cache_path(settings, key) if path.exists(): logger.debug("Cache hit: %s", key) @@ -55,16 +71,17 @@ def set_cached( image_blobs: list[bytes], model: str, temperature: float, + max_tokens: int, response: dict[str, Any], settings: Settings, + format_key: str = "text", ) -> None: """Store a response in the file cache.""" if not settings.cache.enabled: return - parts = [system_prompt, user_prompt, model, str(temperature)] - for blob in image_blobs: - parts.append(hashlib.sha256(blob).hexdigest()) - key = _make_key(parts) + key = _make_key( + _cache_key_parts(system_prompt, user_prompt, image_blobs, model, temperature, max_tokens, format_key) + ) path = _cache_path(settings, key) with path.open("w", encoding="utf-8") as f: json.dump(response, f, ensure_ascii=False, indent=2) diff --git a/tests/conftest.py b/tests/conftest.py index fa2531f..2f0361d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -52,3 +52,7 @@ class FakeVLMClient(BaseVLMClient): if key in system_prompt or key in user_prompt: return VLMResponse(content=json.dumps(value), model="fake") return VLMResponse(content=json.dumps({}), model="fake") + + async def aclose(self) -> None: + """The fake client holds no resources.""" + return None diff --git a/tests/test_cache.py b/tests/test_cache.py new file mode 100644 index 0000000..fac010e --- /dev/null +++ b/tests/test_cache.py @@ -0,0 +1,44 @@ +"""Tests for the LLM response cache.""" + +from __future__ import annotations + +import pytest + +from odoo_ocr.config import Settings +from odoo_ocr.utils import cache + + +@pytest.fixture +def cache_settings(tmp_path) -> Settings: + return Settings(cache={"enabled": True, "dir": tmp_path / "llm_cache"}) + + +def test_cache_roundtrip(cache_settings: Settings) -> None: + response = {"content": "hello", "model": "m", "prompt_tokens": 1} + cache.set_cached("sys", "user", [b"img"], "model", 0.1, 4096, response, cache_settings) + assert cache.get_cached("sys", "user", [b"img"], "model", 0.1, 4096, cache_settings) == response + + +def test_cache_miss(cache_settings: Settings) -> None: + assert cache.get_cached("a", "b", [], "m", 0.1, 4096, cache_settings) is None + + +def test_cache_disabled(tmp_path) -> None: + settings = Settings(cache={"enabled": False, "dir": tmp_path / "llm_cache"}) + cache.set_cached("sys", "user", [], "m", 0.1, 4096, {"content": "x"}, settings) + assert cache.get_cached("sys", "user", [], "m", 0.1, 4096, settings) is None + assert not (tmp_path / "llm_cache").exists() + + +def test_cache_key_includes_temperature_format_and_max_tokens(cache_settings: Settings) -> None: + response = {"content": "x"} + cache.set_cached("sys", "user", [], "m", 0.1, 4096, response, cache_settings, format_key="text") + # Different temperature must not hit the same entry. + assert cache.get_cached("sys", "user", [], "m", 0.2, 4096, cache_settings) is None + # Different response_format must not hit the same entry. + assert ( + cache.get_cached("sys", "user", [], "m", 0.1, 4096, cache_settings, format_key="ExtractedInvoice") + is None + ) + # Different max_tokens must not hit the same entry. + assert cache.get_cached("sys", "user", [], "m", 0.1, 2048, cache_settings) is None