diff --git a/app/services/ocr/vl_model.py b/app/services/ocr/vl_model.py index b7c459c..df615c8 100644 --- a/app/services/ocr/vl_model.py +++ b/app/services/ocr/vl_model.py @@ -32,6 +32,29 @@ def _try_docuvision(image_path: str | Path) -> str | None: cf_orch_url = os.environ.get("CF_ORCH_URL") if not cf_orch_url: return None + + # Tier 1: task-based routing — coordinator owns model selection. + try: + from app.services.task_inference import task_allocate, TaskNotRegistered + from app.services.ocr.docuvision_client import DocuvisionClient + try: + with task_allocate( + "kiwi", "ocr", + service_hint="cf-docuvision", + ttl_s=60.0, + ) as alloc: + doc_client = DocuvisionClient(alloc.url) + result = doc_client.extract_text(image_path) + return result.text if result.text else None + except TaskNotRegistered: + logger.debug( + "kiwi.ocr not in coordinator assignments — " + "falling back to direct cf-docuvision allocation" + ) + except Exception as exc: + logger.debug("task allocation path failed, trying direct allocate: %s", exc) + + # Tier 2: direct allocation — hardcoded service type. try: from circuitforge_orch.client import CFOrchClient from app.services.ocr.docuvision_client import DocuvisionClient @@ -49,7 +72,7 @@ def _try_docuvision(image_path: str | Path) -> str | None: result = doc_client.extract_text(image_path) return result.text if result.text else None except Exception as exc: - logger.debug("cf-docuvision fast-path failed, falling back: %s", exc) + logger.debug("cf-docuvision fast-path failed, falling back to local VLM: %s", exc) return None diff --git a/tests/services/test_vl_model_task.py b/tests/services/test_vl_model_task.py new file mode 100644 index 0000000..d75e3ef --- /dev/null +++ b/tests/services/test_vl_model_task.py @@ -0,0 +1,88 @@ +"""Tests for task-based routing added to _try_docuvision().""" +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest + + +def _mock_doc_result(text: str = "RECEIPT TEXT") -> MagicMock: + r = MagicMock() + r.text = text + return r + + +def _make_task_ctx(url: str = "http://node:9010") -> MagicMock: + alloc = MagicMock() + alloc.url = url + alloc.allocation_id = "alloc-vis-1" + alloc.service = "cf-docuvision" + ctx = MagicMock() + ctx.__enter__ = MagicMock(return_value=alloc) + ctx.__exit__ = MagicMock(return_value=False) + return ctx + + +def _make_task_not_registered() -> MagicMock: + from app.services.task_inference import TaskNotRegistered + ctx = MagicMock() + ctx.__enter__ = MagicMock(side_effect=TaskNotRegistered("not registered")) + ctx.__exit__ = MagicMock(return_value=False) + return ctx + + +def _make_direct_alloc(url: str = "http://node:9011") -> MagicMock: + alloc = MagicMock() + alloc.url = url + ctx = MagicMock() + ctx.__enter__ = MagicMock(return_value=alloc) + ctx.__exit__ = MagicMock(return_value=False) + return ctx + + +def test_try_docuvision_task_path_returns_text(monkeypatch, tmp_path): + """_try_docuvision() uses task allocation and returns extracted text on success.""" + monkeypatch.setenv("CF_ORCH_URL", "http://coord:7700") + fake_image = tmp_path / "receipt.jpg" + fake_image.write_bytes(b"fake") + + with patch("app.services.task_inference.task_allocate", + return_value=_make_task_ctx(url="http://node:9010")), \ + patch("app.services.ocr.docuvision_client.DocuvisionClient") as MockDoc: + MockDoc.return_value.extract_text.return_value = _mock_doc_result("STORE $12.34") + from app.services.ocr.vl_model import _try_docuvision + result = _try_docuvision(str(fake_image)) + + assert result == "STORE $12.34" + MockDoc.assert_called_once_with("http://node:9010") + + +def test_try_docuvision_falls_back_to_direct_on_task_not_registered(monkeypatch, tmp_path): + """_try_docuvision() falls back to direct cf-docuvision allocation on TaskNotRegistered.""" + monkeypatch.setenv("CF_ORCH_URL", "http://coord:7700") + fake_image = tmp_path / "receipt.jpg" + fake_image.write_bytes(b"fake") + + with patch("app.services.task_inference.task_allocate", + return_value=_make_task_not_registered()), \ + patch("circuitforge_orch.client.CFOrchClient") as MockClient, \ + patch("app.services.ocr.docuvision_client.DocuvisionClient") as MockDoc: + MockClient.return_value.allocate.return_value = _make_direct_alloc("http://node:9011") + MockDoc.return_value.extract_text.return_value = _mock_doc_result("FALLBACK TEXT") + from app.services.ocr.vl_model import _try_docuvision + result = _try_docuvision(str(fake_image)) + + assert result == "FALLBACK TEXT" + MockDoc.assert_called_once_with("http://node:9011") + + +def test_try_docuvision_returns_none_without_cf_orch_url(monkeypatch, tmp_path): + """_try_docuvision() returns None immediately when CF_ORCH_URL is not set.""" + monkeypatch.delenv("CF_ORCH_URL", raising=False) + fake_image = tmp_path / "receipt.jpg" + fake_image.write_bytes(b"fake") + + from app.services.ocr.vl_model import _try_docuvision + result = _try_docuvision(str(fake_image)) + + assert result is None