Fix OCR and Fireworks capability handling

This commit is contained in:
Cursor Agent 2026-05-01 22:59:03 +00:00 • committed by Sameer Kankute
parent 005f67bff3
commit 186e32290e
No known key found for this signature in database
4 changed files with 40 additions and 2 deletions

View file

@ -1895,7 +1895,9 @@ def ocr_cost(
pages_processed = response.usage_info.pages_processed
if pages_processed is None:
return 0.0, 0.0
if custom_llm_provider == "reducto" or model.startswith("reducto/"):
return 0.0, 0.0
raise ValueError("OCR response pages_processed is None")
ocr_cost_per_page: float = 0.0
if model_info is not None:

View file

@ -285,7 +285,7 @@ class FireworksAIConfig(OpenAIGPTConfig):
)
# Only include supports_reasoning if True
if supports_reasoning_value is not None:
if supports_reasoning_value:
provider_specific_model_info["supports_reasoning"] = (
supports_reasoning_value
)

View file

@ -6,6 +6,8 @@ from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
import litellm
sys.path.insert(
0, os.path.abspath("../../../../..")
) # Adds the parent directory to the system path
@ -144,6 +146,17 @@ def test_get_model_info_respects_explicit_fireworks_capabilities():
assert model_info["supports_tool_choice"] is False
def test_get_provider_info_omits_false_supports_reasoning(monkeypatch):
"""Test that Fireworks only overrides supports_reasoning for supported models."""
config = FireworksAIConfig()
model = "fireworks_ai/test-reasoning-false"
monkeypatch.setitem(litellm.model_cost, model, {"supports_reasoning": False})
info = config.get_provider_info(model)
assert "supports_reasoning" not in info
def test_add_transform_inline_image_block_skips_data_urls():
"""
data: URLs must not have #transform=inline appended — doing so corrupts the

View file

@ -1,4 +1,5 @@
import litellm
import pytest
from litellm.cost_calculator import completion_cost
from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo
@ -97,3 +98,25 @@ def test_ocr_cost_returns_zero_when_no_pricing_and_no_pages(monkeypatch):
)
assert cost == 0.0
def test_ocr_cost_raises_when_pages_processed_missing_for_page_pricing(monkeypatch):
monkeypatch.setattr(
litellm,
"get_model_info",
lambda model, custom_llm_provider=None: {"ocr_cost_per_page": 0.5},
)
response = OCRResponse(
pages=[OCRPage(index=0, markdown="missing pages")],
model="mistral-ocr-latest",
usage_info=OCRUsageInfo(pages_processed=None),
)
with pytest.raises(ValueError, match="OCR response pages_processed is None"):
completion_cost(
completion_response=response,
model="mistral/mistral-ocr-latest",
custom_llm_provider="mistral",
call_type="ocr",
)