mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
test(ocr): replace per-provider OCR test classes with a declarative provider x auth x input matrix
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yujong Lee
parent
20e1e6f2a9
commit
9767878425
7 changed files with 328 additions and 502 deletions
|
|
@ -1,203 +0,0 @@
|
|||
"""
|
||||
Base test class for OCR functionality across different providers.
|
||||
|
||||
This follows the same pattern as BaseLLMChatTest in tests/llm_translation/base_llm_unit_tests.py
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import litellm
|
||||
import os
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
|
||||
# Test resources
|
||||
TEST_IMAGE_PATH = "test_image_edit.png"
|
||||
# Tiny in-repo PDF served via jsdelivr (sha-pinned, immutable). The arxiv
|
||||
# PDF previously used here was several MB — once base64-encoded into the
|
||||
# Vertex OCR request it ballooned cassettes past 100 MB per test. Keep
|
||||
# the URL stable across runs so cassettes don't churn.
|
||||
TEST_PDF_URL = (
|
||||
"https://cdn.jsdelivr.net/gh/BerriAI/litellm"
|
||||
"@d769e81c90d453240c61fc572cdb27fae06a89d0"
|
||||
"/tests/llm_translation/fixtures/dummy.pdf"
|
||||
)
|
||||
|
||||
|
||||
class BaseOCRTest(ABC):
|
||||
"""
|
||||
Abstract base test class that enforces common OCR tests across all providers.
|
||||
|
||||
Each provider-specific test class should inherit from this and implement
|
||||
get_base_ocr_call_args() to return provider-specific configuration.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def get_base_ocr_call_args(self) -> dict:
|
||||
"""Must return the base OCR call args for the specific provider"""
|
||||
pass
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_basic_ocr_with_url(self, sync_mode):
|
||||
"""
|
||||
Test basic OCR with a public URL.
|
||||
"""
|
||||
litellm._turn_on_debug()
|
||||
base_ocr_call_args = self.get_base_ocr_call_args()
|
||||
print("BASE OCR Call args=", base_ocr_call_args)
|
||||
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
|
||||
litellm.model_cost = litellm.get_model_cost_map(url="")
|
||||
|
||||
try:
|
||||
if sync_mode:
|
||||
response = litellm.ocr(
|
||||
document={"type": "document_url", "document_url": TEST_PDF_URL},
|
||||
**base_ocr_call_args,
|
||||
)
|
||||
else:
|
||||
response = await litellm.aocr(
|
||||
document={"type": "document_url", "document_url": TEST_PDF_URL},
|
||||
**base_ocr_call_args,
|
||||
)
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print(f"Sync Mode: {sync_mode}")
|
||||
print(f"Response type: {type(response)}")
|
||||
print(
|
||||
f"Response object: {response.object if hasattr(response, 'object') else 'N/A'}"
|
||||
)
|
||||
|
||||
# Check if response has expected OCR format
|
||||
assert hasattr(response, "pages"), "Response should have 'pages' attribute"
|
||||
assert hasattr(response, "model"), "Response should have 'model' attribute"
|
||||
assert hasattr(
|
||||
response, "object"
|
||||
), "Response should have 'object' attribute"
|
||||
assert (
|
||||
response.object == "ocr"
|
||||
), f"Expected object='ocr', got '{response.object}'"
|
||||
|
||||
# Validate pages structure
|
||||
assert isinstance(response.pages, list), "pages should be a list"
|
||||
assert len(response.pages) > 0, "Should have at least one page"
|
||||
|
||||
# Check first page structure
|
||||
first_page = response.pages[0]
|
||||
assert hasattr(first_page, "index"), "Page should have 'index' attribute"
|
||||
assert hasattr(
|
||||
first_page, "markdown"
|
||||
), "Page should have 'markdown' attribute"
|
||||
|
||||
# Extract text from all pages for validation
|
||||
total_text = "\n\n".join(
|
||||
page.markdown for page in response.pages if page.markdown
|
||||
)
|
||||
print(f"Total pages: {len(response.pages)}")
|
||||
print(f"Total extracted text length: {len(total_text)} characters")
|
||||
print(f"First 200 chars: {total_text[:200]}")
|
||||
print(f"Model: {response.model}")
|
||||
if response.usage_info:
|
||||
print(f"Pages processed: {response.usage_info.pages_processed}")
|
||||
print(f"{'='*80}\n")
|
||||
|
||||
assert len(total_text) > 0, "Should extract some text from the document"
|
||||
|
||||
#########################################################
|
||||
# validate we get a response cost in hidden parameters
|
||||
#########################################################
|
||||
hidden_params = response._hidden_params
|
||||
assert isinstance(
|
||||
hidden_params, dict
|
||||
), "Hidden parameters should be a dictionary"
|
||||
|
||||
print("response usage_info:", response.usage_info)
|
||||
|
||||
response_cost = hidden_params.get("response_cost")
|
||||
assert (
|
||||
response_cost is not None
|
||||
), "Response cost should be in hidden parameters"
|
||||
assert response_cost > 0, "Response cost should be greater than 0"
|
||||
print("response_cost=", response_cost)
|
||||
|
||||
except litellm.RateLimitError as e:
|
||||
error_msg = str(e)
|
||||
if "Quota exceeded" in error_msg or "RESOURCE_EXHAUSTED" in error_msg:
|
||||
pytest.skip(f"Quota exceeded - {error_msg}")
|
||||
else:
|
||||
pytest.skip(f"Rate limit exceeded - {error_msg}")
|
||||
except litellm.InternalServerError:
|
||||
pytest.skip("Model is overloaded")
|
||||
except litellm.BadRequestError as e:
|
||||
error_msg = str(e)
|
||||
if (
|
||||
"URL_REJECTED" in error_msg
|
||||
or "Cannot fetch content from the provided URL" in error_msg
|
||||
):
|
||||
pytest.skip(f"URL rejected by provider - {error_msg}")
|
||||
else:
|
||||
pytest.fail(f"OCR call failed: {str(e)}")
|
||||
except Exception as e:
|
||||
pytest.fail(f"OCR call failed: {str(e)}")
|
||||
|
||||
def test_ocr_response_structure(self):
|
||||
"""
|
||||
Test that the OCR response has the correct structure.
|
||||
"""
|
||||
litellm.set_verbose = True
|
||||
base_ocr_call_args = self.get_base_ocr_call_args()
|
||||
|
||||
try:
|
||||
response = litellm.ocr(
|
||||
document={"type": "document_url", "document_url": TEST_PDF_URL},
|
||||
**base_ocr_call_args,
|
||||
)
|
||||
|
||||
# Validate response structure
|
||||
assert hasattr(response, "pages"), "Response should have 'pages' attribute"
|
||||
assert hasattr(response, "model"), "Response should have 'model' attribute"
|
||||
assert hasattr(
|
||||
response, "object"
|
||||
), "Response should have 'object' attribute"
|
||||
assert hasattr(
|
||||
response, "usage_info"
|
||||
), "Response should have 'usage_info' attribute"
|
||||
|
||||
assert isinstance(response.pages, list), "pages should be a list"
|
||||
assert len(response.pages) > 0, "Should have at least one page"
|
||||
assert response.object == "ocr", "object should be 'ocr'"
|
||||
|
||||
# Validate first page structure
|
||||
first_page = response.pages[0]
|
||||
assert hasattr(first_page, "index"), "Page should have 'index' attribute"
|
||||
assert hasattr(
|
||||
first_page, "markdown"
|
||||
), "Page should have 'markdown' attribute"
|
||||
assert isinstance(first_page.markdown, str), "markdown should be a string"
|
||||
|
||||
print(f"\nResponse structure validated:")
|
||||
print(f" - object: {response.object}")
|
||||
print(f" - model: {response.model}")
|
||||
print(f" - pages: {len(response.pages)}")
|
||||
if response.usage_info:
|
||||
print(f" - pages_processed: {response.usage_info.pages_processed}")
|
||||
print(f" - doc_size_bytes: {response.usage_info.doc_size_bytes}")
|
||||
|
||||
except litellm.RateLimitError as e:
|
||||
error_msg = str(e)
|
||||
if "Quota exceeded" in error_msg or "RESOURCE_EXHAUSTED" in error_msg:
|
||||
pytest.skip(f"Quota exceeded - {error_msg}")
|
||||
else:
|
||||
pytest.skip(f"Rate limit exceeded - {error_msg}")
|
||||
except litellm.InternalServerError:
|
||||
pytest.skip("Model is overloaded")
|
||||
except litellm.BadRequestError as e:
|
||||
error_msg = str(e)
|
||||
if (
|
||||
"URL_REJECTED" in error_msg
|
||||
or "Cannot fetch content from the provided URL" in error_msg
|
||||
):
|
||||
pytest.skip(f"URL rejected by provider - {error_msg}")
|
||||
else:
|
||||
pytest.fail(f"OCR response structure test failed: {str(e)}")
|
||||
except Exception as e:
|
||||
pytest.fail(f"OCR response structure test failed: {str(e)}")
|
||||
|
|
@ -1,29 +0,0 @@
|
|||
"""
|
||||
Test OCR functionality with Azure AI API.
|
||||
|
||||
Note: Azure AI OCR automatically converts URLs to base64 data URIs since
|
||||
the Azure AI endpoint doesn't have internet access.
|
||||
"""
|
||||
|
||||
import os
|
||||
from base_ocr_unit_tests import BaseOCRTest
|
||||
|
||||
|
||||
class TestAzureAIOCR(BaseOCRTest):
|
||||
"""
|
||||
Test class for Azure AI OCR functionality.
|
||||
Inherits from BaseOCRTest and provides Azure AI-specific configuration.
|
||||
|
||||
Note: For Azure AI, LiteLLM will automatically convert URLs to base64 data URIs before
|
||||
sending to the API, since Azure AI OCR endpoint doesn't have internet access.
|
||||
"""
|
||||
|
||||
def get_base_ocr_call_args(self) -> dict:
|
||||
"""
|
||||
Return the base OCR call args for Azure AI.
|
||||
"""
|
||||
return {
|
||||
"model": "azure_ai/mistral-document-ai-2512",
|
||||
"api_key": os.getenv("AZURE_API_KEY"),
|
||||
"api_base": os.getenv("AZURE_API_BASE"),
|
||||
}
|
||||
|
|
@ -1,53 +1,13 @@
|
|||
"""
|
||||
Test OCR functionality with Azure Document Intelligence API.
|
||||
|
||||
Azure Document Intelligence provides advanced document analysis capabilities
|
||||
using the v4.0 (2024-11-30) API.
|
||||
"""
|
||||
|
||||
import os
|
||||
"""Azure Document Intelligence request transformation: Mistral-shaped `pages` to Azure's query string."""
|
||||
|
||||
import pytest
|
||||
|
||||
from base_ocr_unit_tests import BaseOCRTest
|
||||
from litellm.constants import AZURE_DOCUMENT_INTELLIGENCE_API_VERSION
|
||||
from litellm.llms.azure_ai.ocr.document_intelligence.transformation import (
|
||||
AzureDocumentIntelligenceOCRConfig,
|
||||
)
|
||||
|
||||
|
||||
class TestAzureDocumentIntelligenceOCR(BaseOCRTest):
|
||||
"""
|
||||
Test class for Azure Document Intelligence OCR functionality.
|
||||
|
||||
Inherits from BaseOCRTest and provides Azure Document Intelligence-specific configuration.
|
||||
|
||||
Tests the azure_ai/doc-intelligence/<model> provider route.
|
||||
"""
|
||||
|
||||
def get_base_ocr_call_args(self) -> dict:
|
||||
"""
|
||||
Return the base OCR call args for Azure Document Intelligence.
|
||||
|
||||
Uses prebuilt-layout model which is closest to Mistral OCR format.
|
||||
"""
|
||||
# Check for required environment variables
|
||||
api_key = os.environ.get("AZURE_DOCUMENT_INTELLIGENCE_API_KEY")
|
||||
endpoint = os.environ.get("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT")
|
||||
|
||||
if not api_key or not endpoint:
|
||||
pytest.skip(
|
||||
"AZURE_DOCUMENT_INTELLIGENCE_API_KEY and AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT "
|
||||
"environment variables are required for Azure Document Intelligence tests"
|
||||
)
|
||||
|
||||
return {
|
||||
"model": "azure_ai/doc-intelligence/prebuilt-layout",
|
||||
"api_key": api_key,
|
||||
"api_base": endpoint,
|
||||
}
|
||||
|
||||
|
||||
class TestAzureDocumentIntelligencePagesParam:
|
||||
"""
|
||||
Unit tests for the Mistral-compatible `pages` parameter translation to
|
||||
|
|
@ -101,7 +61,7 @@ class TestAzureDocumentIntelligencePagesParam:
|
|||
cfg.map_ocr_params({"pages": [True, False]}, {}, "prebuilt-layout")
|
||||
|
||||
def test_map_ocr_params_unsupported_type_raises(self, cfg):
|
||||
with pytest.raises(ValueError, match='based, Mistral-style\\) or a string like'):
|
||||
with pytest.raises(ValueError, match="based, Mistral-style\\) or a string like"):
|
||||
cfg.map_ocr_params({"pages": 5}, {}, "prebuilt-layout")
|
||||
|
||||
def test_get_complete_url_appends_pages_query(self, cfg):
|
||||
|
|
@ -110,9 +70,7 @@ class TestAzureDocumentIntelligencePagesParam:
|
|||
model="azure_ai/doc-intelligence/prebuilt-layout",
|
||||
optional_params={"pages": "1-3,5"},
|
||||
)
|
||||
assert (
|
||||
f"api-version={AZURE_DOCUMENT_INTELLIGENCE_API_VERSION}" in url
|
||||
), url
|
||||
assert f"api-version={AZURE_DOCUMENT_INTELLIGENCE_API_VERSION}" in url, url
|
||||
assert "pages=1-3,5" in url, url
|
||||
assert "/documentintelligence/documentModels/prebuilt-layout:analyze" in url
|
||||
|
||||
|
|
@ -168,4 +126,3 @@ class TestAzureDocumentIntelligencePagesParam:
|
|||
|
||||
assert "pages=3,4,5,6,7,8,9" in url
|
||||
assert req.data == {"urlSource": "https://example.com/x.pdf"}
|
||||
|
||||
|
|
|
|||
317
tests/ocr_tests/test_ocr_matrix.py
Normal file
317
tests/ocr_tests/test_ocr_matrix.py
Normal file
|
|
@ -0,0 +1,317 @@
|
|||
"""Live provider x auth x input coverage for ``litellm.ocr`` / ``litellm.aocr``.
|
||||
|
||||
Each ``Case`` is one hand-picked cell, not the full cross product: every provider
|
||||
exercises each of its credential kinds in both ``explicit`` (kwargs) and ``env``
|
||||
(monkeypatched environment) mode at least once, and every input kind a provider
|
||||
accepts is exercised at least once. Sync and async are spread across the cells.
|
||||
Every cell also checks the success callback saw the same response and cost.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import io
|
||||
import os
|
||||
import re
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRResponse
|
||||
|
||||
Document = Mapping[str, object]
|
||||
AuthMode = Literal["explicit", "env"]
|
||||
CallStyle = Literal["sync", "async"]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class LoggedCall:
|
||||
payload: Mapping[str, object]
|
||||
response: object
|
||||
|
||||
|
||||
class RecordingLogger(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
super().__init__() # pyright: ignore[reportUnknownMemberType] # CustomLogger.__init__ is untyped
|
||||
self.calls: Final[list[LoggedCall]] = [] # mutable-ok: append-only sink the callback hooks write into
|
||||
|
||||
def _record(self, kwargs: Mapping[str, object], response_obj: object) -> None:
|
||||
payload: Final = _string_keyed(kwargs.get("standard_logging_object"))
|
||||
self.calls.append(LoggedCall(payload, response_obj))
|
||||
|
||||
def log_success_event(
|
||||
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
self._record(kwargs, response_obj)
|
||||
|
||||
async def async_log_success_event(
|
||||
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
self._record(kwargs, response_obj)
|
||||
|
||||
async def wait_for_call(self, timeout: float = 10.0) -> LoggedCall:
|
||||
deadline: Final = asyncio.get_running_loop().time() + timeout
|
||||
while not self.calls:
|
||||
assert asyncio.get_running_loop().time() < deadline, "success callback never fired"
|
||||
await asyncio.sleep(0.05)
|
||||
assert len(self.calls) == 1, self.calls
|
||||
return self.calls[0]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def logger(monkeypatch: pytest.MonkeyPatch) -> RecordingLogger:
|
||||
recorder: Final = RecordingLogger()
|
||||
monkeypatch.setattr(litellm, "callbacks", [recorder])
|
||||
for registry in ("success_callback", "_async_success_callback", "failure_callback", "_async_failure_callback"):
|
||||
monkeypatch.setattr(litellm, registry, [])
|
||||
return recorder
|
||||
|
||||
|
||||
TESTS_DIR: Final = Path(__file__).resolve().parents[1]
|
||||
PDF_PATH: Final = TESTS_DIR / "llm_translation" / "fixtures" / "dummy.pdf"
|
||||
PNG_PATH: Final = TESTS_DIR / "image_gen_tests" / "test_image.png"
|
||||
PINNED_CDN: Final = "https://cdn.jsdelivr.net/gh/BerriAI/litellm@d769e81c90d453240c61fc572cdb27fae06a89d0"
|
||||
PDF_URL: Final = f"{PINNED_CDN}/tests/llm_translation/fixtures/dummy.pdf"
|
||||
PNG_URL: Final = f"{PINNED_CDN}/tests/image_gen_tests/test_image.png"
|
||||
PDF_TEXT: Final = "Test PDF File"
|
||||
PNG_TEXT: Final = "LiteLLM"
|
||||
|
||||
|
||||
class _NamedReader(io.BytesIO):
|
||||
def __init__(self, path: Path) -> None:
|
||||
super().__init__(path.read_bytes())
|
||||
self.name: Final = path.name
|
||||
|
||||
|
||||
def _data_uri(path: Path, mime: str) -> str:
|
||||
return f"data:{mime};base64,{base64.b64encode(path.read_bytes()).decode()}"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Input:
|
||||
id: str
|
||||
build: Callable[[], Document]
|
||||
expected_text: str
|
||||
|
||||
|
||||
PDF_BY_URL: Final = Input("pdf_url", lambda: {"type": "document_url", "document_url": PDF_URL}, PDF_TEXT)
|
||||
PNG_BY_URL: Final = Input("image_url", lambda: {"type": "image_url", "image_url": PNG_URL}, PNG_TEXT)
|
||||
PDF_DATA_URI: Final = Input(
|
||||
"pdf_data_uri",
|
||||
lambda: {"type": "document_url", "document_url": _data_uri(PDF_PATH, "application/pdf")},
|
||||
PDF_TEXT,
|
||||
)
|
||||
PNG_DATA_URI: Final = Input(
|
||||
"image_data_uri", lambda: {"type": "image_url", "image_url": _data_uri(PNG_PATH, "image/png")}, PNG_TEXT
|
||||
)
|
||||
PDF_AS_PATH: Final = Input("pdf_path", lambda: {"type": "file", "file": PDF_PATH}, PDF_TEXT)
|
||||
PDF_AS_BYTES: Final = Input(
|
||||
"pdf_bytes", lambda: {"type": "file", "file": PDF_PATH.read_bytes(), "mime_type": "application/pdf"}, PDF_TEXT
|
||||
)
|
||||
PNG_AS_BYTES: Final = Input(
|
||||
"image_bytes", lambda: {"type": "file", "file": PNG_PATH.read_bytes(), "mime_type": "image/png"}, PNG_TEXT
|
||||
)
|
||||
PNG_AS_FILE_OBJECT: Final = Input(
|
||||
"image_file_object", lambda: {"type": "file", "file": _NamedReader(PNG_PATH)}, PNG_TEXT
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Secret:
|
||||
"""One credential value: the ``litellm.ocr`` kwarg it travels in, the env var litellm reads
|
||||
when the kwarg is omitted, and the env var that holds the value in the test process."""
|
||||
|
||||
kwarg: str
|
||||
env: str
|
||||
source: str | None = None
|
||||
|
||||
@property
|
||||
def source_env(self) -> str:
|
||||
return self.source or self.env
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Credential:
|
||||
id: str
|
||||
secrets: tuple[Secret, ...]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Provider:
|
||||
id: str
|
||||
model: str
|
||||
credentials: tuple[Credential, ...]
|
||||
params: Mapping[str, str] = MappingProxyType({})
|
||||
|
||||
@property
|
||||
def env_vars(self) -> frozenset[str]:
|
||||
return frozenset(secret.env for credential in self.credentials for secret in credential.secrets)
|
||||
|
||||
|
||||
MISTRAL_KEY: Final = Credential("api_key", (Secret("api_key", "MISTRAL_API_KEY"),))
|
||||
COHERE_KEY: Final = Credential("api_key", (Secret("api_key", "COHERE_API_KEY"),))
|
||||
REDUCTO_KEY: Final = Credential("api_key", (Secret("api_key", "REDUCTO_API_KEY"),))
|
||||
|
||||
AZURE_ENTRA_SECRETS: Final = (
|
||||
Secret("tenant_id", "AZURE_TENANT_ID", "AZURE_FOUNDRY_TENANT_ID"),
|
||||
Secret("client_id", "AZURE_CLIENT_ID", "AZURE_FOUNDRY_ADMIN_CLIENT_ID"),
|
||||
Secret("client_secret", "AZURE_CLIENT_SECRET", "AZURE_FOUNDRY_ADMIN_CLIENT_SECRET"),
|
||||
)
|
||||
AZURE_AI_BASE: Final = Secret("api_base", "AZURE_AI_API_BASE")
|
||||
AZURE_AI_KEY: Final = Credential("api_key", (AZURE_AI_BASE, Secret("api_key", "AZURE_AI_API_KEY")))
|
||||
AZURE_AI_ENTRA: Final = Credential("entra", (AZURE_AI_BASE, *AZURE_ENTRA_SECRETS))
|
||||
|
||||
AZURE_DI_BASE: Final = Secret("api_base", "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT")
|
||||
AZURE_DI_KEY: Final = Credential("api_key", (AZURE_DI_BASE, Secret("api_key", "AZURE_DOCUMENT_INTELLIGENCE_API_KEY")))
|
||||
AZURE_DI_ENTRA: Final = Credential("entra", (AZURE_DI_BASE, *AZURE_ENTRA_SECRETS))
|
||||
|
||||
VERTEX_SERVICE_ACCOUNT: Final = Credential(
|
||||
"service_account",
|
||||
(Secret("vertex_credentials", "VERTEXAI_CREDENTIALS"), Secret("vertex_project", "VERTEXAI_PROJECT")),
|
||||
)
|
||||
|
||||
MISTRAL: Final = Provider("mistral", "mistral/mistral-ocr-latest", (MISTRAL_KEY,))
|
||||
AZURE_AI_MISTRAL: Final = Provider(
|
||||
"azure_ai_mistral", "azure_ai/mistral-document-ai-2512", (AZURE_AI_KEY, AZURE_AI_ENTRA)
|
||||
)
|
||||
AZURE_DOC_INTELLIGENCE: Final = Provider(
|
||||
"azure_doc_intelligence", "azure_ai/doc-intelligence/prebuilt-layout", (AZURE_DI_KEY, AZURE_DI_ENTRA)
|
||||
)
|
||||
COHERE: Final = Provider("cohere", "cohere/parse-v5.0", (COHERE_KEY,))
|
||||
REDUCTO_V3: Final = Provider("reducto_v3", "reducto/parse-v3", (REDUCTO_KEY,))
|
||||
REDUCTO_LEGACY: Final = Provider("reducto_legacy", "reducto/parse-legacy", (REDUCTO_KEY,))
|
||||
VERTEX_MISTRAL: Final = Provider(
|
||||
"vertex_mistral",
|
||||
"vertex_ai/mistral-ocr-2505",
|
||||
(VERTEX_SERVICE_ACCOUNT,),
|
||||
MappingProxyType({"vertex_location": "us-central1"}),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Case:
|
||||
provider: Provider
|
||||
credential: Credential
|
||||
auth: AuthMode
|
||||
document: Input
|
||||
call: CallStyle
|
||||
|
||||
@property
|
||||
def id(self) -> str:
|
||||
return f"{self.provider.id}-{self.credential.id}-{self.auth}-{self.document.id}-{self.call}"
|
||||
|
||||
def bind_credentials(self, monkeypatch: pytest.MonkeyPatch) -> Mapping[str, str]:
|
||||
"""Clear every env var the provider could fall back to, then supply this case's values via kwargs or env."""
|
||||
values: Final = {secret: os.environ.get(secret.source_env) for secret in self.credential.secrets}
|
||||
missing: Final = tuple(secret.source_env for secret, value in values.items() if not value)
|
||||
if missing:
|
||||
pytest.skip(f"{', '.join(missing)} not set")
|
||||
for env_var in self.provider.env_vars:
|
||||
monkeypatch.delenv(env_var, raising=False)
|
||||
if self.auth == "explicit":
|
||||
return {secret.kwarg: value for secret, value in values.items() if value}
|
||||
for secret, value in values.items():
|
||||
monkeypatch.setenv(secret.env, value or "")
|
||||
return {}
|
||||
|
||||
async def run(self, credentials: Mapping[str, str]) -> OCRResponse:
|
||||
kwargs: Final = {**self.provider.params, **credentials}
|
||||
document: Final = self.document.build()
|
||||
response: Final = (
|
||||
await litellm.aocr(model=self.provider.model, document=document, **kwargs) # pyright: ignore[reportUnknownMemberType] # @client erases the signature
|
||||
if self.call == "async"
|
||||
else litellm.ocr(model=self.provider.model, document=document, **kwargs)
|
||||
)
|
||||
assert isinstance(response, OCRResponse)
|
||||
return response
|
||||
|
||||
|
||||
CASES: Final = (
|
||||
Case(MISTRAL, MISTRAL_KEY, "explicit", PDF_BY_URL, "sync"),
|
||||
Case(MISTRAL, MISTRAL_KEY, "env", PNG_BY_URL, "async"),
|
||||
Case(MISTRAL, MISTRAL_KEY, "explicit", PDF_AS_PATH, "sync"),
|
||||
Case(MISTRAL, MISTRAL_KEY, "explicit", PNG_AS_BYTES, "async"),
|
||||
Case(MISTRAL, MISTRAL_KEY, "explicit", PNG_AS_FILE_OBJECT, "sync"),
|
||||
Case(AZURE_AI_MISTRAL, AZURE_AI_KEY, "explicit", PDF_BY_URL, "sync"),
|
||||
Case(AZURE_AI_MISTRAL, AZURE_AI_KEY, "env", PNG_BY_URL, "async"),
|
||||
Case(AZURE_AI_MISTRAL, AZURE_AI_ENTRA, "explicit", PDF_AS_PATH, "sync"),
|
||||
Case(AZURE_AI_MISTRAL, AZURE_AI_ENTRA, "env", PDF_DATA_URI, "async"),
|
||||
Case(AZURE_DOC_INTELLIGENCE, AZURE_DI_KEY, "explicit", PDF_BY_URL, "sync"),
|
||||
Case(AZURE_DOC_INTELLIGENCE, AZURE_DI_KEY, "env", PNG_AS_BYTES, "async"),
|
||||
Case(AZURE_DOC_INTELLIGENCE, AZURE_DI_ENTRA, "explicit", PNG_BY_URL, "async"),
|
||||
Case(AZURE_DOC_INTELLIGENCE, AZURE_DI_ENTRA, "env", PDF_AS_PATH, "sync"),
|
||||
Case(COHERE, COHERE_KEY, "explicit", PNG_BY_URL, "sync"),
|
||||
Case(COHERE, COHERE_KEY, "env", PNG_DATA_URI, "async"),
|
||||
Case(REDUCTO_V3, REDUCTO_KEY, "explicit", PDF_AS_PATH, "sync"),
|
||||
Case(REDUCTO_V3, REDUCTO_KEY, "env", PNG_AS_BYTES, "async"),
|
||||
Case(REDUCTO_V3, REDUCTO_KEY, "explicit", PDF_DATA_URI, "async"),
|
||||
Case(REDUCTO_LEGACY, REDUCTO_KEY, "explicit", PDF_AS_BYTES, "sync"),
|
||||
Case(VERTEX_MISTRAL, VERTEX_SERVICE_ACCOUNT, "explicit", PDF_BY_URL, "sync"),
|
||||
Case(VERTEX_MISTRAL, VERTEX_SERVICE_ACCOUNT, "env", PNG_BY_URL, "async"),
|
||||
)
|
||||
|
||||
|
||||
def _response_cost(response: OCRResponse) -> float:
|
||||
response_cost: Final[object] = response._hidden_params.get("response_cost") # pyright: ignore[reportPrivateUsage, reportUnknownMemberType, reportUnknownVariableType] # response_cost is only surfaced on _hidden_params
|
||||
assert isinstance(response_cost, float) and response_cost > 0
|
||||
return response_cost
|
||||
|
||||
|
||||
def _assert_ocr_response(response: OCRResponse, model: str, expected_text: str) -> None:
|
||||
assert response.object == "ocr"
|
||||
assert response.model == model.split("/", 1)[1]
|
||||
assert [page.index for page in response.pages] == list(range(len(response.pages)))
|
||||
text: Final = re.sub(r"\s+", " ", " ".join(page.markdown for page in response.pages))
|
||||
assert expected_text.lower() in text.lower(), text
|
||||
assert response.usage_info is not None
|
||||
assert response.usage_info.pages_processed == len(response.pages)
|
||||
_response_cost(response)
|
||||
|
||||
|
||||
def _string_keyed(value: object) -> Mapping[str, object]:
|
||||
assert isinstance(value, Mapping), type(value)
|
||||
items: Final = tuple(value.items()) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType, reportUnknownArgumentType] # narrowed from object
|
||||
return MappingProxyType({str(key): value for key, value in items}) # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # narrowed from object
|
||||
|
||||
|
||||
def _assert_logged(logged: LoggedCall, response: OCRResponse, model: str, logged_model: str, call: CallStyle) -> None:
|
||||
assert isinstance(logged.response, OCRResponse)
|
||||
assert logged.response.pages == response.pages
|
||||
assert logged.payload["status"] == "success"
|
||||
assert logged.payload["call_type"] == ("aocr" if call == "async" else "ocr")
|
||||
assert logged.payload["custom_llm_provider"] == model.split("/", 1)[0]
|
||||
assert logged.payload["model"] == logged_model
|
||||
assert logged.payload["response_cost"] == _response_cost(response)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("case", CASES, ids=[case.id for case in CASES])
|
||||
async def test_ocr(case: Case, monkeypatch: pytest.MonkeyPatch, logger: RecordingLogger) -> None:
|
||||
credentials: Final = case.bind_credentials(monkeypatch)
|
||||
response: Final = await case.run(credentials)
|
||||
_assert_ocr_response(response, case.provider.model, case.document.expected_text)
|
||||
_assert_logged(await logger.wait_for_call(), response, case.provider.model, response.model, case.call)
|
||||
|
||||
|
||||
async def test_router_aocr(monkeypatch: pytest.MonkeyPatch, logger: RecordingLogger) -> None:
|
||||
case: Final = Case(MISTRAL, MISTRAL_KEY, "explicit", PDF_BY_URL, "async")
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "ocr-alias",
|
||||
"litellm_params": {"model": MISTRAL.model, **case.bind_credentials(monkeypatch)},
|
||||
}
|
||||
]
|
||||
)
|
||||
response: Final = await router.aocr(model="ocr-alias", document=PDF_BY_URL.build()) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # Router.aocr is untyped
|
||||
assert isinstance(response, OCRResponse)
|
||||
_assert_ocr_response(response, MISTRAL.model, PDF_TEXT)
|
||||
_assert_logged(await logger.wait_for_call(), response, MISTRAL.model, MISTRAL.model, case.call)
|
||||
|
|
@ -1,92 +0,0 @@
|
|||
"""
|
||||
Test OCR functionality with Mistral API.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import pytest
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from base_ocr_unit_tests import BaseOCRTest, TEST_PDF_URL
|
||||
|
||||
|
||||
class TestMistralOCR(BaseOCRTest):
|
||||
"""
|
||||
Test class for Mistral OCR functionality.
|
||||
"""
|
||||
|
||||
def get_base_ocr_call_args(self) -> dict:
|
||||
"""Return the base OCR call args for Mistral"""
|
||||
return {
|
||||
"model": "mistral/mistral-ocr-latest",
|
||||
"api_key": os.getenv("MISTRAL_API_KEY"),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_router_aocr_with_mistral():
|
||||
"""
|
||||
Test OCR with Router using Mistral OCR deployment.
|
||||
"""
|
||||
litellm.set_verbose = True
|
||||
|
||||
# Create router with Mistral OCR deployment
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "mistral-ocr",
|
||||
"litellm_params": {
|
||||
"model": "mistral/mistral-ocr-latest",
|
||||
"api_key": os.getenv("MISTRAL_API_KEY"),
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
try:
|
||||
# Call OCR through router
|
||||
response = await router.aocr(
|
||||
model="mistral-ocr",
|
||||
document={"type": "document_url", "document_url": TEST_PDF_URL},
|
||||
)
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("Router OCR Test")
|
||||
print(f"Response type: {type(response)}")
|
||||
print(
|
||||
f"Response object: {response.object if hasattr(response, 'object') else 'N/A'}"
|
||||
)
|
||||
|
||||
# Check if response has expected Mistral OCR format
|
||||
assert hasattr(response, "pages"), "Response should have 'pages' attribute"
|
||||
assert hasattr(response, "model"), "Response should have 'model' attribute"
|
||||
assert hasattr(response, "object"), "Response should have 'object' attribute"
|
||||
assert (
|
||||
response.object == "ocr"
|
||||
), f"Expected object='ocr', got '{response.object}'"
|
||||
|
||||
# Validate pages structure
|
||||
assert isinstance(response.pages, list), "pages should be a list"
|
||||
assert len(response.pages) > 0, "Should have at least one page"
|
||||
|
||||
# Check first page structure
|
||||
first_page = response.pages[0]
|
||||
assert hasattr(first_page, "index"), "Page should have 'index' attribute"
|
||||
assert hasattr(first_page, "markdown"), "Page should have 'markdown' attribute"
|
||||
|
||||
# Extract text from all pages for validation
|
||||
total_text = "\n\n".join(
|
||||
page.markdown for page in response.pages if page.markdown
|
||||
)
|
||||
print(f"Total pages: {len(response.pages)}")
|
||||
print(f"Total extracted text length: {len(total_text)} characters")
|
||||
print(f"First 200 chars: {total_text[:200]}")
|
||||
print(f"Model: {response.model}")
|
||||
if response.usage_info:
|
||||
print(f"Pages processed: {response.usage_info.pages_processed}")
|
||||
print(f"{'='*80}\n")
|
||||
|
||||
assert len(total_text) > 0, "Should extract some text from the document"
|
||||
|
||||
except Exception as e:
|
||||
pytest.fail(f"Router OCR call failed: {str(e)}")
|
||||
|
|
@ -1,117 +1,8 @@
|
|||
"""
|
||||
Test OCR functionality with Vertex AI OCR APIs (Mistral and DeepSeek).
|
||||
"""Vertex AI OCR config routing and DeepSeek request shaping (no network)."""
|
||||
|
||||
Note: Vertex AI OCR automatically converts URLs to base64 data URIs since
|
||||
the Vertex AI endpoint doesn't have internet access.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from base_ocr_unit_tests import BaseOCRTest
|
||||
|
||||
|
||||
def load_vertex_ai_credentials():
|
||||
"""Load Vertex AI credentials for tests"""
|
||||
# Define the path to the vertex_key.json file
|
||||
print("loading vertex ai credentials")
|
||||
filepath = os.path.dirname(os.path.abspath(__file__))
|
||||
vertex_key_path = filepath + "/vertex_key.json"
|
||||
|
||||
# Read the existing content of the file or create an empty dictionary
|
||||
try:
|
||||
with open(vertex_key_path, "r") as file:
|
||||
# Read the file content
|
||||
print("Read vertexai file path")
|
||||
content = file.read()
|
||||
|
||||
# If the file is empty or not valid JSON, create an empty dictionary
|
||||
if not content or not content.strip():
|
||||
service_account_key_data = {}
|
||||
else:
|
||||
# Attempt to load the existing JSON content
|
||||
file.seek(0)
|
||||
service_account_key_data = json.load(file)
|
||||
except FileNotFoundError:
|
||||
# If the file doesn't exist, create an empty dictionary
|
||||
service_account_key_data = {}
|
||||
|
||||
# Update the service_account_key_data with environment variables
|
||||
private_key_id = os.environ.get("VERTEX_AI_PRIVATE_KEY_ID", "")
|
||||
private_key = os.environ.get("VERTEX_AI_PRIVATE_KEY", "")
|
||||
private_key = private_key.replace("\\n", "\n")
|
||||
service_account_key_data["private_key_id"] = private_key_id
|
||||
service_account_key_data["private_key"] = private_key
|
||||
|
||||
# Create a temporary file
|
||||
with tempfile.NamedTemporaryFile(mode="w+", delete=False) as temp_file:
|
||||
# Write the updated content to the temporary files
|
||||
json.dump(service_account_key_data, temp_file, indent=2)
|
||||
|
||||
# Export the temporary file as GOOGLE_APPLICATION_CREDENTIALS
|
||||
os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = os.path.abspath(temp_file.name)
|
||||
|
||||
|
||||
class TestVertexAIMistralOCR(BaseOCRTest):
|
||||
"""
|
||||
Test class for Vertex AI Mistral OCR functionality.
|
||||
Inherits from BaseOCRTest and provides Vertex AI-specific configuration.
|
||||
|
||||
Note: For Vertex AI, LiteLLM will automatically convert URLs to base64 data URIs before
|
||||
sending to the API, since Vertex AI OCR endpoint doesn't have internet access.
|
||||
"""
|
||||
|
||||
def setup_method(self):
|
||||
if os.environ.get("LITELLM_RUN_LIVE_VERTEX_MISTRAL_OCR_TESTS") != "1":
|
||||
pytest.skip("Live Vertex AI Mistral OCR E2E tests are opt-in")
|
||||
if os.environ.get("CASSETTE_REDIS_URL"):
|
||||
pytest.skip(
|
||||
"Live Vertex AI Mistral OCR E2E tests cannot run under VCR replay"
|
||||
)
|
||||
|
||||
def get_base_ocr_call_args(self) -> dict:
|
||||
"""
|
||||
Return the base OCR call args for Vertex AI Mistral OCR.
|
||||
"""
|
||||
load_vertex_ai_credentials()
|
||||
return {
|
||||
"model": "vertex_ai/mistral-ocr-2505",
|
||||
"vertex_location": "us-central1",
|
||||
}
|
||||
|
||||
|
||||
class TestVertexAIDeepSeekOCR(BaseOCRTest):
|
||||
"""
|
||||
Test class for Vertex AI DeepSeek OCR functionality.
|
||||
Inherits from BaseOCRTest and provides Vertex AI-specific configuration.
|
||||
|
||||
Note: DeepSeek OCR uses the chat completion API format through the openapi endpoint.
|
||||
Note: DeepSeek OCR does not support PDF URLs - only image URLs and base64 data.
|
||||
"""
|
||||
|
||||
def get_base_ocr_call_args(self) -> dict:
|
||||
"""
|
||||
Return the base OCR call args for Vertex AI DeepSeek OCR.
|
||||
"""
|
||||
load_vertex_ai_credentials()
|
||||
return {
|
||||
"model": "vertex_ai/deepseek-ocr-maas",
|
||||
"vertex_location": "us-central1",
|
||||
}
|
||||
|
||||
# Skip PDF URL tests for DeepSeek OCR as it doesn't support PDF URLs
|
||||
@pytest.mark.skip(reason="DeepSeek OCR does not support PDF URLs")
|
||||
async def test_basic_ocr_with_url(self, sync_mode):
|
||||
"""Skip this test for DeepSeek OCR - PDF URLs not supported"""
|
||||
pass
|
||||
|
||||
@pytest.mark.skip(reason="DeepSeek OCR does not support PDF URLs")
|
||||
def test_ocr_response_structure(self):
|
||||
"""Skip this test for DeepSeek OCR - PDF URLs not supported"""
|
||||
pass
|
||||
|
||||
|
||||
def test_vertex_ai_ocr_routing():
|
||||
|
|
@ -126,21 +17,19 @@ def test_vertex_ai_ocr_routing():
|
|||
|
||||
# Test DeepSeek OCR routing
|
||||
deepseek_config = get_vertex_ai_ocr_config("vertex_ai/deepseek-ocr-maas")
|
||||
assert isinstance(
|
||||
deepseek_config, VertexAIDeepSeekOCRConfig
|
||||
), "DeepSeek model should route to VertexAIDeepSeekOCRConfig"
|
||||
assert isinstance(deepseek_config, VertexAIDeepSeekOCRConfig), (
|
||||
"DeepSeek model should route to VertexAIDeepSeekOCRConfig"
|
||||
)
|
||||
|
||||
# Test Mistral OCR routing (should use default VertexAIOCRConfig)
|
||||
mistral_config = get_vertex_ai_ocr_config("vertex_ai/mistral-ocr-2505")
|
||||
assert isinstance(
|
||||
mistral_config, VertexAIOCRConfig
|
||||
), "Mistral model should route to VertexAIOCRConfig"
|
||||
assert isinstance(mistral_config, VertexAIOCRConfig), "Mistral model should route to VertexAIOCRConfig"
|
||||
|
||||
# Test other DeepSeek variants
|
||||
deepseek_variant = get_vertex_ai_ocr_config("vertex_ai/deepseek-ocr-maas")
|
||||
assert isinstance(
|
||||
deepseek_variant, VertexAIDeepSeekOCRConfig
|
||||
), "DeepSeek variant should route to VertexAIDeepSeekOCRConfig"
|
||||
assert isinstance(deepseek_variant, VertexAIDeepSeekOCRConfig), (
|
||||
"DeepSeek variant should route to VertexAIDeepSeekOCRConfig"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", ("deepseek-ocr-maas", "deepseek-ai/deepseek-ocr-maas"))
|
||||
|
|
|
|||
|
|
@ -1,13 +0,0 @@
|
|||
{
|
||||
"type": "service_account",
|
||||
"project_id": "litellm-ci-cd",
|
||||
"private_key_id": "",
|
||||
"private_key": "",
|
||||
"client_email": "test-litellm-ci-cd@litellm-ci-cd.iam.gserviceaccount.com",
|
||||
"client_id": "116563532503305622785",
|
||||
"auth_uri": "https://accounts.google.com/o/oauth2/auth",
|
||||
"token_uri": "https://oauth2.googleapis.com/token",
|
||||
"auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
|
||||
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test-litellm-ci-cd%40litellm-ci-cd.iam.gserviceaccount.com",
|
||||
"universe_domain": "googleapis.com"
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue