feat(orch): migrate OCR vision routing to task-based allocation with direct-allocate fallback

This commit is contained in:
pyr0ball 2026-05-13 10:46:07 -07:00
parent dd39418bc8
commit cdbc24240a
2 changed files with 112 additions and 1 deletions

View file

@ -32,6 +32,29 @@ def _try_docuvision(image_path: str | Path) -> str | None:
cf_orch_url = os.environ.get("CF_ORCH_URL") cf_orch_url = os.environ.get("CF_ORCH_URL")
if not cf_orch_url: if not cf_orch_url:
return None 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: try:
from circuitforge_orch.client import CFOrchClient from circuitforge_orch.client import CFOrchClient
from app.services.ocr.docuvision_client import DocuvisionClient 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) result = doc_client.extract_text(image_path)
return result.text if result.text else None return result.text if result.text else None
except Exception as exc: 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 return None

View file

@ -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