mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
fix(native): restack shared eligibility contracts
This commit is contained in:
Yujong Lee
yujonglee
parent
3dd53c0dd0
commit
7e8613c123
6 changed files with 626 additions and 259 deletions
|
|
@ -1,38 +1,14 @@
|
|||
import base64
|
||||
from io import IOBase
|
||||
from typing import Final, NoReturn
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.litellm_core_utils.audio_utils.utils import process_audio_file
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.rust_bridge import transcription as rust_transcription_bridge
|
||||
from litellm.rust_bridge.dispatch import anative_first, native_first, provider_errors
|
||||
from litellm.rust_bridge.request import request_context
|
||||
from litellm.rust_bridge.runtime import DispatchResult, adapt_result
|
||||
from litellm.types.utils import FileTypes, TranscriptionResponse
|
||||
|
||||
|
||||
def _unavailable() -> NoReturn:
|
||||
raise RuntimeError("Rust audio transcription bridge is unavailable")
|
||||
|
||||
|
||||
async def _aunavailable() -> NoReturn:
|
||||
_unavailable()
|
||||
|
||||
|
||||
class BedrockAudioTranscriptionRustDispatch:
|
||||
@staticmethod
|
||||
def _input_source_kind(audio_file: FileTypes) -> str:
|
||||
content: Final = audio_file[1] if isinstance(audio_file, tuple) else audio_file
|
||||
if isinstance(content, (bytes, bytearray, memoryview)):
|
||||
return "bytes"
|
||||
if isinstance(content, IOBase):
|
||||
return "file"
|
||||
if isinstance(content, str):
|
||||
return "path"
|
||||
return "opaque"
|
||||
|
||||
@staticmethod
|
||||
def _audio_payload(audio_file: FileTypes) -> dict[str, object]:
|
||||
processed_audio: Final = process_audio_file(audio_file)
|
||||
|
|
@ -55,44 +31,6 @@ class BedrockAudioTranscriptionRustDispatch:
|
|||
"filename": processed_audio.filename,
|
||||
}
|
||||
|
||||
def _attempt_audio_transcriptions(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
audio_file: FileTypes,
|
||||
api_key: str | None,
|
||||
api_base: str | None,
|
||||
custom_llm_provider: str,
|
||||
extra_headers: dict[str, object] | None,
|
||||
optional_params: dict[str, object],
|
||||
timeout: float | httpx.Timeout | None,
|
||||
logging_obj: Logging | None = None,
|
||||
) -> DispatchResult[TranscriptionResponse]:
|
||||
result: Final = rust_transcription_bridge.transcription(
|
||||
model=model,
|
||||
audio=self._audio_payload(audio_file),
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
extra_headers=extra_headers,
|
||||
optional_params=optional_params,
|
||||
timeout=timeout,
|
||||
input_source_kind=self._input_source_kind(audio_file),
|
||||
context=request_context(
|
||||
logging_obj=logging_obj,
|
||||
request_model=logging_obj.model if logging_obj is not None else model,
|
||||
litellm_params=logging_obj.litellm_params if logging_obj is not None else None,
|
||||
),
|
||||
)
|
||||
return adapt_result(result, lambda response: TranscriptionResponse(**response))
|
||||
|
||||
@native_first(
|
||||
native=_attempt_audio_transcriptions,
|
||||
route="audio transcription",
|
||||
errors=lambda self, model, audio_file, api_key, api_base, custom_llm_provider, extra_headers, optional_params, timeout, logging_obj=None: (
|
||||
provider_errors(custom_llm_provider, model)
|
||||
),
|
||||
)
|
||||
def audio_transcriptions(
|
||||
self,
|
||||
*,
|
||||
|
|
@ -104,24 +42,8 @@ class BedrockAudioTranscriptionRustDispatch:
|
|||
extra_headers: dict[str, object] | None,
|
||||
optional_params: dict[str, object],
|
||||
timeout: float | httpx.Timeout | None,
|
||||
logging_obj: Logging | None = None,
|
||||
) -> TranscriptionResponse:
|
||||
_unavailable()
|
||||
|
||||
async def _attempt_async_audio_transcriptions(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
audio_file: FileTypes,
|
||||
api_key: str | None,
|
||||
api_base: str | None,
|
||||
custom_llm_provider: str,
|
||||
extra_headers: dict[str, object] | None,
|
||||
optional_params: dict[str, object],
|
||||
timeout: float | httpx.Timeout | None,
|
||||
logging_obj: Logging | None = None,
|
||||
) -> DispatchResult[TranscriptionResponse]:
|
||||
result: Final = await rust_transcription_bridge.atranscription(
|
||||
rust_response: Final = rust_transcription_bridge.transcription(
|
||||
model=model,
|
||||
audio=self._audio_payload(audio_file),
|
||||
api_key=api_key,
|
||||
|
|
@ -130,22 +52,11 @@ class BedrockAudioTranscriptionRustDispatch:
|
|||
extra_headers=extra_headers,
|
||||
optional_params=optional_params,
|
||||
timeout=timeout,
|
||||
input_source_kind=self._input_source_kind(audio_file),
|
||||
context=request_context(
|
||||
logging_obj=logging_obj,
|
||||
request_model=logging_obj.model if logging_obj is not None else model,
|
||||
litellm_params=logging_obj.litellm_params if logging_obj is not None else None,
|
||||
),
|
||||
)
|
||||
return adapt_result(result, lambda response: TranscriptionResponse(**response))
|
||||
if rust_response is None:
|
||||
raise RuntimeError("Rust audio transcription bridge is unavailable")
|
||||
return TranscriptionResponse(**rust_response)
|
||||
|
||||
@anative_first(
|
||||
native=_attempt_async_audio_transcriptions,
|
||||
route="audio transcription",
|
||||
errors=lambda self, model, audio_file, api_key, api_base, custom_llm_provider, extra_headers, optional_params, timeout, logging_obj=None: (
|
||||
provider_errors(custom_llm_provider, model)
|
||||
),
|
||||
)
|
||||
async def async_audio_transcriptions(
|
||||
self,
|
||||
*,
|
||||
|
|
@ -157,6 +68,17 @@ class BedrockAudioTranscriptionRustDispatch:
|
|||
extra_headers: dict[str, object] | None,
|
||||
optional_params: dict[str, object],
|
||||
timeout: float | httpx.Timeout | None,
|
||||
logging_obj: Logging | None = None,
|
||||
) -> TranscriptionResponse:
|
||||
await _aunavailable()
|
||||
rust_response: Final = await rust_transcription_bridge.atranscription(
|
||||
model=model,
|
||||
audio=self._audio_payload(audio_file),
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
extra_headers=extra_headers,
|
||||
optional_params=optional_params,
|
||||
timeout=timeout,
|
||||
)
|
||||
if rust_response is None:
|
||||
raise RuntimeError("Rust audio transcription bridge is unavailable")
|
||||
return TranscriptionResponse(**rust_response)
|
||||
|
|
|
|||
|
|
@ -10,7 +10,8 @@ from litellm.rust_bridge.bindings import (
|
|||
UNCHANGED,
|
||||
NativeBinding,
|
||||
Unchanged,
|
||||
native_exception_types,
|
||||
native_declined_types,
|
||||
native_upstream_types,
|
||||
)
|
||||
from litellm.rust_bridge.protocols import NativeModule, RustRouteDecline
|
||||
from litellm.rust_bridge.request import NativeRequestCapabilities, NativeRequestContext
|
||||
|
|
@ -32,6 +33,14 @@ class PythonFallbackReason(Enum):
|
|||
NATIVE_DECLINED = "native_declined"
|
||||
|
||||
|
||||
class NativeSkipReason(Enum):
|
||||
DISABLED = "disabled"
|
||||
INELIGIBLE = "ineligible"
|
||||
UNAVAILABLE = "unavailable"
|
||||
DECLINED = "declined"
|
||||
FAILED = "failed"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Handled(Generic[ResultT]):
|
||||
value: ResultT
|
||||
|
|
@ -43,7 +52,18 @@ class PythonFallback:
|
|||
detail: str | None = None
|
||||
|
||||
|
||||
DispatchResult: TypeAlias = Handled[ResultT] | PythonFallback
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class NativeSkipped:
|
||||
reason: NativeSkipReason
|
||||
detail: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class NativeFailed:
|
||||
error: Exception
|
||||
|
||||
|
||||
DispatchResult: TypeAlias = Handled[ResultT] | PythonFallback | NativeSkipped | NativeFailed
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -296,10 +316,10 @@ class EndpointBinding(Generic[BindingT]):
|
|||
adapt: Callable[[NativeT], ResultT],
|
||||
error_context: BridgeErrorContext,
|
||||
) -> DispatchResult[ResultT]:
|
||||
exceptions: Final = native_exception_types()
|
||||
if exceptions is None:
|
||||
declined: Final = native_declined_types()
|
||||
upstream: Final = native_upstream_types()
|
||||
if not declined or not upstream:
|
||||
return Handled(adapt(call()))
|
||||
declined, upstream = exceptions
|
||||
try:
|
||||
value: Final = call()
|
||||
except declined as error:
|
||||
|
|
@ -315,10 +335,10 @@ class EndpointBinding(Generic[BindingT]):
|
|||
adapt: Callable[[NativeT], ResultT],
|
||||
error_context: BridgeErrorContext,
|
||||
) -> DispatchResult[ResultT]:
|
||||
exceptions: Final = native_exception_types()
|
||||
if exceptions is None:
|
||||
declined: Final = native_declined_types()
|
||||
upstream: Final = native_upstream_types()
|
||||
if not declined or not upstream:
|
||||
return Handled(adapt(await call()))
|
||||
declined, upstream = exceptions
|
||||
try:
|
||||
value: Final = await call()
|
||||
except declined as error:
|
||||
|
|
@ -505,6 +525,12 @@ def identity(value: ResultT) -> ResultT:
|
|||
return value
|
||||
|
||||
|
||||
def adapt_result(result: DispatchResult[NativeT], adapt: Callable[[NativeT], ResultT]) -> DispatchResult[ResultT]:
|
||||
if isinstance(result, Handled):
|
||||
return Handled(adapt(result.value))
|
||||
return result
|
||||
|
||||
|
||||
async def async_none() -> None:
|
||||
return None
|
||||
|
||||
|
|
|
|||
|
|
@ -57,7 +57,6 @@ _PREFLIGHT: Final[EndpointBinding[RustRouteDecline]] = EndpointBinding.native(
|
|||
|
||||
|
||||
def configure_rust_transcription(
|
||||
enabled: bool = True,
|
||||
*,
|
||||
transcription: RustTranscription | None | Unchanged = UNCHANGED,
|
||||
atranscription: RustAtranscription | None | Unchanged = UNCHANGED,
|
||||
|
|
|
|||
|
|
@ -12,8 +12,13 @@ import pytest
|
|||
import litellm
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRResponse
|
||||
from litellm.rust_bridge import configuration
|
||||
from litellm.rust_bridge.request import NativeOCRRequest, NativeRequestContext, NativeRequestOptions
|
||||
from litellm.rust_bridge.runtime import Handled
|
||||
from litellm.rust_bridge.request import (
|
||||
NativeOCRRequest,
|
||||
NativeRequestContext,
|
||||
NativeRequestOptions,
|
||||
NativeVertexOptions,
|
||||
PreparedNativeCall,
|
||||
)
|
||||
from litellm.rust_bridge.timeouts import timeout_to_seconds
|
||||
|
||||
# `litellm/__init__.py` does `from .ocr.main import *`, which binds the `ocr`
|
||||
|
|
@ -47,7 +52,6 @@ class RecordingBridge:
|
|||
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[dict[str, object]] = []
|
||||
self.contexts: list[NativeRequestContext] = []
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
|
|
@ -65,10 +69,10 @@ class RecordingBridge:
|
|||
"custom_llm_provider": options.custom_llm_provider,
|
||||
"extra_headers": options.extra_headers,
|
||||
"optional_params": request.optional_params,
|
||||
"vertex": options.vertex,
|
||||
"timeout_seconds": options.timeout_seconds,
|
||||
}
|
||||
)
|
||||
self.contexts.append(context)
|
||||
return dict(FAKE_OCR_RESPONSE)
|
||||
|
||||
|
||||
|
|
@ -77,7 +81,6 @@ class RecordingAsyncBridge:
|
|||
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[dict[str, object]] = []
|
||||
self.contexts: list[NativeRequestContext] = []
|
||||
|
||||
async def __call__(
|
||||
self,
|
||||
|
|
@ -95,10 +98,10 @@ class RecordingAsyncBridge:
|
|||
"custom_llm_provider": options.custom_llm_provider,
|
||||
"extra_headers": options.extra_headers,
|
||||
"optional_params": request.optional_params,
|
||||
"vertex": options.vertex,
|
||||
"timeout_seconds": options.timeout_seconds,
|
||||
}
|
||||
)
|
||||
self.contexts.append(context)
|
||||
return dict(FAKE_OCR_RESPONSE)
|
||||
|
||||
|
||||
|
|
@ -192,7 +195,7 @@ def build_prepared_request(
|
|||
litellm_params: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = 12.5,
|
||||
) -> Any:
|
||||
return rust_bridge.PreparedOCRRequest(
|
||||
return ocr_main._PreparedOCRRequest(
|
||||
model=model,
|
||||
document=document,
|
||||
api_key=api_key,
|
||||
|
|
@ -397,13 +400,105 @@ def test_timeout_to_seconds_handles_float_timeout_and_none():
|
|||
assert timeout_to_seconds(httpx.Timeout(30.0, read=42.0)) == 42.0
|
||||
|
||||
|
||||
def test_bridge_wrapper_forwards_prepared_args_and_wraps_response():
|
||||
bridge = RecordingBridge()
|
||||
|
||||
litellm.rust(True)
|
||||
|
||||
rust_bridge.set_rust_ocr(ocr=bridge)
|
||||
response = rust_bridge.dispatch_ocr(
|
||||
prepare=lambda: PreparedNativeCall(
|
||||
request=NativeOCRRequest(
|
||||
model="mistral-ocr-latest",
|
||||
document=DOCUMENT,
|
||||
optional_params={"include_image_base64": True, "pages": [0]},
|
||||
),
|
||||
options=NativeRequestOptions(
|
||||
api_key="sk-test",
|
||||
api_base="https://proxy.internal",
|
||||
custom_llm_provider="mistral",
|
||||
extra_headers={"Authorization": "Bearer sk-test", "x-trace-id": "trace-1"},
|
||||
timeout_seconds=12.5,
|
||||
),
|
||||
),
|
||||
fallback=lambda: pytest.fail("unexpected Python fallback"),
|
||||
adapt=dict,
|
||||
model="mistral-ocr-latest",
|
||||
provider="mistral",
|
||||
eligible=True,
|
||||
)
|
||||
|
||||
assert response == FAKE_OCR_RESPONSE
|
||||
call = bridge.calls[0]
|
||||
assert call == {
|
||||
"model": "mistral-ocr-latest",
|
||||
"document": DOCUMENT,
|
||||
"api_key": "sk-test",
|
||||
"api_base": "https://proxy.internal",
|
||||
"custom_llm_provider": "mistral",
|
||||
"extra_headers": {
|
||||
"Authorization": "Bearer sk-test",
|
||||
"x-trace-id": "trace-1",
|
||||
},
|
||||
"optional_params": {"include_image_base64": True, "pages": [0]},
|
||||
"vertex": None,
|
||||
"timeout_seconds": 12.5,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response():
|
||||
bridge = RecordingAsyncBridge()
|
||||
|
||||
litellm.rust(True)
|
||||
|
||||
rust_bridge.set_rust_ocr(aocr=bridge)
|
||||
|
||||
async def unexpected_fallback():
|
||||
pytest.fail("unexpected Python fallback")
|
||||
|
||||
response = await rust_bridge.adispatch_ocr(
|
||||
prepare=lambda: PreparedNativeCall(
|
||||
request=NativeOCRRequest(
|
||||
model="mistral-ocr-maas",
|
||||
document=DOCUMENT,
|
||||
optional_params={},
|
||||
),
|
||||
options=NativeRequestOptions(
|
||||
custom_llm_provider="vertex_ai",
|
||||
vertex=NativeVertexOptions(project="project-1"),
|
||||
timeout_seconds=42.0,
|
||||
),
|
||||
),
|
||||
fallback=unexpected_fallback,
|
||||
adapt=dict,
|
||||
model="mistral-ocr-maas",
|
||||
provider="vertex_ai",
|
||||
eligible=True,
|
||||
)
|
||||
|
||||
assert response == FAKE_OCR_RESPONSE
|
||||
assert bridge.calls[0] == {
|
||||
"model": "mistral-ocr-maas",
|
||||
"document": DOCUMENT,
|
||||
"api_key": None,
|
||||
"api_base": None,
|
||||
"custom_llm_provider": "vertex_ai",
|
||||
"extra_headers": None,
|
||||
"optional_params": {},
|
||||
"vertex": NativeVertexOptions(project="project-1"),
|
||||
"timeout_seconds": 42.0,
|
||||
}
|
||||
|
||||
|
||||
def test_run_rust_ocr_prepares_request_and_wraps_response():
|
||||
bridge = RecordingBridge()
|
||||
logging_obj = RecordingLogging()
|
||||
litellm.rust(True)
|
||||
rust_bridge.set_rust_ocr(ocr=bridge)
|
||||
|
||||
response = rust_bridge.attempt_ocr(
|
||||
response = ocr_main._run_rust_ocr(
|
||||
fallback=lambda: pytest.fail("unexpected Python fallback"),
|
||||
prepared_request=build_prepared_request(
|
||||
logging_obj=logging_obj,
|
||||
api_base="https://proxy.internal",
|
||||
|
|
@ -414,13 +509,8 @@ def test_run_rust_ocr_prepares_request_and_wraps_response():
|
|||
resolve_api_key=lambda _name: None,
|
||||
)
|
||||
|
||||
assert isinstance(response, Handled)
|
||||
response = response.value
|
||||
assert isinstance(response, OCRResponse)
|
||||
assert response.pages[0].markdown == "hello world"
|
||||
assert bridge.contexts[0].capabilities.execution_mode == "sync"
|
||||
assert bridge.contexts[0].capabilities.input_source_kind == "document_url"
|
||||
assert bridge.contexts[0].capabilities.native_response_format is False
|
||||
assert bridge.calls[0] == {
|
||||
"model": "mistral-ocr-latest",
|
||||
"document": DOCUMENT,
|
||||
|
|
@ -432,6 +522,7 @@ def test_run_rust_ocr_prepares_request_and_wraps_response():
|
|||
"x-trace-id": "trace-1",
|
||||
},
|
||||
"optional_params": {"include_image_base64": True},
|
||||
"vertex": NativeVertexOptions(),
|
||||
"timeout_seconds": 12.5,
|
||||
}
|
||||
|
||||
|
|
@ -441,7 +532,8 @@ def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing():
|
|||
litellm.rust(True)
|
||||
rust_bridge.set_rust_ocr(ocr=bridge)
|
||||
|
||||
rust_bridge.attempt_ocr(
|
||||
ocr_main._run_rust_ocr(
|
||||
fallback=lambda: pytest.fail("unexpected Python fallback"),
|
||||
prepared_request=build_prepared_request(api_key=None, timeout=None),
|
||||
resolve_api_key=lambda name: "sk-from-vault" if name == "MISTRAL_API_KEY" else None,
|
||||
)
|
||||
|
|
@ -457,7 +549,8 @@ def test_run_rust_ocr_prefers_explicit_key_over_resolver():
|
|||
def _resolver(name: str) -> str | None:
|
||||
raise AssertionError(f"resolver should not be called for {name}")
|
||||
|
||||
rust_bridge.attempt_ocr(
|
||||
ocr_main._run_rust_ocr(
|
||||
fallback=lambda: pytest.fail("unexpected Python fallback"),
|
||||
prepared_request=build_prepared_request(
|
||||
api_key="sk-explicit",
|
||||
timeout=None,
|
||||
|
|
@ -478,7 +571,8 @@ def test_run_rust_ocr_uses_provider_api_key_env_var():
|
|||
resolver_calls.append(name)
|
||||
return "sk-provider-env"
|
||||
|
||||
rust_bridge.attempt_ocr(
|
||||
ocr_main._run_rust_ocr(
|
||||
fallback=lambda: pytest.fail("unexpected Python fallback"),
|
||||
prepared_request=build_prepared_request(
|
||||
provider_config=FakeOCRConfig(api_key_env_var="PROVIDER_OCR_API_KEY"),
|
||||
model="provider-ocr-model",
|
||||
|
|
@ -497,7 +591,8 @@ def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata():
|
|||
litellm.rust(True)
|
||||
rust_bridge.set_rust_ocr(ocr=bridge)
|
||||
|
||||
rust_bridge.attempt_ocr(
|
||||
ocr_main._run_rust_ocr(
|
||||
fallback=lambda: pytest.fail("unexpected Python fallback"),
|
||||
prepared_request=build_prepared_request(
|
||||
custom_llm_provider="vertex_ai",
|
||||
model="mistral-ocr-maas",
|
||||
|
|
@ -512,11 +607,8 @@ def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata():
|
|||
resolve_api_key=lambda _name: None,
|
||||
)
|
||||
|
||||
assert bridge.calls[0]["optional_params"] == {
|
||||
"include_image_base64": True,
|
||||
"vertex_project": "project-1",
|
||||
"vertex_location": "us-central1",
|
||||
}
|
||||
assert bridge.calls[0]["optional_params"] == {"include_image_base64": True}
|
||||
assert bridge.calls[0]["vertex"] == NativeVertexOptions(project="project-1", location="us-central1")
|
||||
|
||||
|
||||
def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager():
|
||||
|
|
@ -530,7 +622,8 @@ def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_mana
|
|||
"VERTEXAI_LOCATION": "us-east5",
|
||||
}.get(name)
|
||||
|
||||
rust_bridge.attempt_ocr(
|
||||
ocr_main._run_rust_ocr(
|
||||
fallback=lambda: pytest.fail("unexpected Python fallback"),
|
||||
prepared_request=build_prepared_request(
|
||||
custom_llm_provider="vertex_ai",
|
||||
model="mistral-ocr-maas",
|
||||
|
|
@ -539,8 +632,7 @@ def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_mana
|
|||
resolve_api_key=_resolver,
|
||||
)
|
||||
|
||||
assert bridge.calls[0]["optional_params"]["vertex_project"] == "project-from-secret"
|
||||
assert bridge.calls[0]["optional_params"]["vertex_location"] == "us-east5"
|
||||
assert bridge.calls[0]["vertex"] == NativeVertexOptions(project="project-from-secret", location="us-east5")
|
||||
|
||||
|
||||
def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager():
|
||||
|
|
@ -548,7 +640,8 @@ def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager():
|
|||
litellm.rust(True)
|
||||
rust_bridge.set_rust_ocr(ocr=bridge)
|
||||
|
||||
rust_bridge.attempt_ocr(
|
||||
ocr_main._run_rust_ocr(
|
||||
fallback=lambda: pytest.fail("unexpected Python fallback"),
|
||||
prepared_request=build_prepared_request(
|
||||
custom_llm_provider="azure_ai",
|
||||
model="pixtral-12b-2409",
|
||||
|
|
@ -566,7 +659,8 @@ def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint():
|
|||
litellm.rust(True)
|
||||
rust_bridge.set_rust_ocr(ocr=bridge)
|
||||
|
||||
rust_bridge.attempt_ocr(
|
||||
ocr_main._run_rust_ocr(
|
||||
fallback=lambda: pytest.fail("unexpected Python fallback"),
|
||||
prepared_request=build_prepared_request(
|
||||
custom_llm_provider="azure_ai",
|
||||
model="doc-intelligence/prebuilt-layout",
|
||||
|
|
@ -587,7 +681,8 @@ def test_run_rust_ocr_runs_pre_call_logging():
|
|||
litellm.rust(True)
|
||||
rust_bridge.set_rust_ocr(ocr=bridge)
|
||||
|
||||
rust_bridge.attempt_ocr(
|
||||
ocr_main._run_rust_ocr(
|
||||
fallback=lambda: pytest.fail("unexpected Python fallback"),
|
||||
prepared_request=build_prepared_request(
|
||||
logging_obj=logging_obj,
|
||||
api_base="https://api.mistral.ai/v1",
|
||||
|
|
@ -734,7 +829,7 @@ async def test_ocr_fallback_skips_native_preparation(
|
|||
def unexpected_preparation(*_args: object, **_kwargs: object) -> None:
|
||||
pytest.fail("Python fallback must not resolve native credentials or emit native pre_call")
|
||||
|
||||
monkeypatch.setattr(rust_bridge, "_prepare_rust_ocr_call", unexpected_preparation)
|
||||
monkeypatch.setattr(ocr_main, "_prepare_rust_ocr_call", unexpected_preparation)
|
||||
monkeypatch.setattr(ocr_main.base_llm_http_handler, "ocr", fallback)
|
||||
|
||||
response: Final = (
|
||||
|
|
@ -747,25 +842,6 @@ async def test_ocr_fallback_skips_native_preparation(
|
|||
fallback.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aocr_rejects_empty_python_fallback_response(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def fake_exception_type(**kwargs: object) -> CapturedException:
|
||||
captured.update(kwargs)
|
||||
return CapturedException("wrapped")
|
||||
|
||||
monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type)
|
||||
monkeypatch.setattr(ocr_main.base_llm_http_handler, "ocr", AsyncMock(return_value=None))
|
||||
|
||||
with pytest.raises(CapturedException, match="wrapped"):
|
||||
await litellm.aocr(model=MODEL, document=DOCUMENT, api_key="sk-test")
|
||||
|
||||
original: Final = captured["original_exception"]
|
||||
assert isinstance(original, ValueError)
|
||||
assert str(original) == "Got an unexpected None response from the OCR API: None"
|
||||
|
||||
|
||||
def test_ocr_provider_configs_expose_api_key_env_vars():
|
||||
from litellm.llms.azure_ai.ocr.document_intelligence.transformation import (
|
||||
AzureDocumentIntelligenceOCRConfig,
|
||||
|
|
|
|||
|
|
@ -1,117 +1,460 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.rust_bridge import runtime
|
||||
from litellm.exceptions import APIError, AuthenticationError, InternalServerError, RateLimitError
|
||||
from litellm.rust_bridge import bindings, runtime
|
||||
|
||||
|
||||
class RustBridgeDeclined(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class RustUpstreamError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def native_exceptions(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(
|
||||
bindings,
|
||||
"get_native_bridge",
|
||||
lambda: SimpleNamespace(
|
||||
RustBridgeDeclined=RustBridgeDeclined,
|
||||
RustUpstreamError=RustUpstreamError,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def context() -> runtime.BridgeErrorContext:
|
||||
return runtime.BridgeErrorContext(provider="anthropic", model="model")
|
||||
|
||||
|
||||
def enabled() -> bool:
|
||||
return True
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class FallbackCase:
|
||||
process_enabled: bool | None = None
|
||||
eligible: bool = True
|
||||
binding_available: bool = True
|
||||
declined: bool = False
|
||||
expected_events: tuple[str, ...] = ()
|
||||
|
||||
|
||||
FALLBACK_CASES: Final = (
|
||||
pytest.param(
|
||||
FallbackCase(process_enabled=False, expected_events=("python",)),
|
||||
id="process-disabled",
|
||||
),
|
||||
pytest.param(
|
||||
FallbackCase(eligible=False, expected_events=("python",)),
|
||||
id="request-ineligible",
|
||||
),
|
||||
pytest.param(
|
||||
FallbackCase(binding_available=False, expected_events=("load", "python")),
|
||||
id="bridge-unavailable",
|
||||
),
|
||||
pytest.param(
|
||||
FallbackCase(declined=True, expected_events=("load", "prepare", "rust", "python")),
|
||||
id="bridge-declined",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("case", FALLBACK_CASES)
|
||||
def test_invoke_falls_back_only_before_provider_success(case: FallbackCase) -> None:
|
||||
events: list[str] = []
|
||||
|
||||
def load() -> object | None:
|
||||
events.append("load")
|
||||
return object() if case.binding_available else None
|
||||
|
||||
def call(_binding: object, _request: object) -> int:
|
||||
events.append("rust")
|
||||
if case.declined:
|
||||
raise RustBridgeDeclined("unsupported")
|
||||
return 3
|
||||
|
||||
bridge: Final = runtime.EndpointBinding(
|
||||
route="messages", load=load, enabled=lambda: case.process_enabled is not False
|
||||
)
|
||||
result: Final = bridge.invoke(
|
||||
prepare=lambda: events.append("prepare"),
|
||||
call=call,
|
||||
fallback=lambda: events.append("python") or "fallback",
|
||||
adapt=str,
|
||||
error_context=context(),
|
||||
eligible=case.eligible,
|
||||
)
|
||||
|
||||
assert result == "fallback"
|
||||
assert tuple(events) == case.expected_events
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("case", FALLBACK_CASES)
|
||||
async def test_ainvoke_matches_sync_fallback_contract(case: FallbackCase) -> None:
|
||||
events: list[str] = []
|
||||
|
||||
def load() -> object | None:
|
||||
events.append("load")
|
||||
return object() if case.binding_available else None
|
||||
|
||||
async def call(_binding: object, _request: object) -> int:
|
||||
events.append("rust")
|
||||
if case.declined:
|
||||
raise RustBridgeDeclined("unsupported")
|
||||
return 3
|
||||
|
||||
async def fallback() -> str:
|
||||
events.append("python")
|
||||
return "fallback"
|
||||
|
||||
bridge: Final = runtime.EndpointBinding(
|
||||
route="messages", load=load, enabled=lambda: case.process_enabled is not False
|
||||
)
|
||||
result: Final = await bridge.ainvoke(
|
||||
prepare=lambda: events.append("prepare"),
|
||||
call=call,
|
||||
fallback=fallback,
|
||||
adapt=str,
|
||||
error_context=context(),
|
||||
eligible=case.eligible,
|
||||
)
|
||||
|
||||
assert result == "fallback"
|
||||
assert tuple(events) == case.expected_events
|
||||
|
||||
|
||||
def test_invoke_adapts_native_success_without_fallback() -> None:
|
||||
bridge: Final = runtime.EndpointBinding(route="messages", load=object, enabled=enabled)
|
||||
|
||||
result: Final = bridge.invoke(
|
||||
prepare=lambda: 3,
|
||||
call=lambda _binding, request: request * 2,
|
||||
fallback=lambda: pytest.fail("fallback must not run"),
|
||||
adapt=lambda value: f"adapted-{value}",
|
||||
error_context=context(),
|
||||
)
|
||||
|
||||
assert result == "adapted-6"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ainvoke_adapts_native_success_without_fallback() -> None:
|
||||
async def call(_binding: object, request: int) -> int:
|
||||
return request * 2
|
||||
|
||||
async def fallback() -> str:
|
||||
pytest.fail("fallback must not run")
|
||||
|
||||
bridge: Final = runtime.EndpointBinding(route="messages", load=object, enabled=enabled)
|
||||
result: Final = await bridge.ainvoke(
|
||||
prepare=lambda: 3,
|
||||
call=call,
|
||||
fallback=fallback,
|
||||
adapt=lambda value: f"adapted-{value}",
|
||||
error_context=context(),
|
||||
)
|
||||
|
||||
assert result == "adapted-6"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("error", "expected_type", "expected_status", "expected_message"),
|
||||
(
|
||||
pytest.param(RustUpstreamError(401, "unauthorized"), AuthenticationError, 401, "unauthorized", id="auth"),
|
||||
pytest.param(RustUpstreamError(429, "rate limited"), RateLimitError, 429, "rate limited", id="rate-limit"),
|
||||
pytest.param(RustUpstreamError(500, "failed"), InternalServerError, 500, "failed", id="server-error"),
|
||||
pytest.param(RustUpstreamError(0, "connection reset"), APIError, 500, "connection reset", id="transport"),
|
||||
pytest.param(RustUpstreamError(403, "forbidden"), APIError, 403, "forbidden", id="other-status"),
|
||||
),
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", (False, True))
|
||||
@pytest.mark.parametrize("state", ("disabled", "ineligible", "unavailable", "handled"))
|
||||
async def test_attempt_only_prepares_selected_requests(asynchronous: bool, state: str) -> None:
|
||||
events: Final[list[str]] = []
|
||||
async def test_upstream_failure_maps_to_api_error_without_fallback(
|
||||
asynchronous: bool,
|
||||
error: RustUpstreamError,
|
||||
expected_type: type[BaseException],
|
||||
expected_status: int,
|
||||
expected_message: str,
|
||||
) -> None:
|
||||
def fail(_binding: object, _request: object) -> object:
|
||||
raise error
|
||||
|
||||
async def afail(binding: object, request: object) -> object:
|
||||
return fail(binding, request)
|
||||
|
||||
async def fallback() -> str:
|
||||
pytest.fail("fallback must not run")
|
||||
|
||||
bridge: Final = runtime.EndpointBinding(route="messages", load=object, enabled=enabled)
|
||||
|
||||
async def invoke() -> None:
|
||||
if asynchronous:
|
||||
await bridge.ainvoke(
|
||||
prepare=lambda: None, call=afail, fallback=fallback, adapt=str, error_context=context()
|
||||
)
|
||||
else:
|
||||
bridge.invoke(
|
||||
prepare=lambda: None,
|
||||
call=fail,
|
||||
fallback=lambda: pytest.fail("fallback must not run"),
|
||||
adapt=str,
|
||||
error_context=context(),
|
||||
)
|
||||
|
||||
with pytest.raises(expected_type, match=expected_message) as caught:
|
||||
await invoke()
|
||||
|
||||
assert type(caught.value) is expected_type
|
||||
assert caught.value.status_code == expected_status
|
||||
assert caught.value.llm_provider == "anthropic"
|
||||
assert caught.value.model == "model"
|
||||
assert caught.value.__cause__ is error
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_upstream_failure_maps_to_api_error_without_fallback() -> None:
|
||||
async def fail(_binding: object, _request: object) -> object:
|
||||
raise RustUpstreamError(503, "overloaded")
|
||||
|
||||
async def fallback() -> object:
|
||||
pytest.fail("fallback must not run")
|
||||
|
||||
bridge: Final = runtime.EndpointBinding(route="messages", load=object, enabled=enabled)
|
||||
|
||||
with pytest.raises(APIError, match="overloaded") as caught:
|
||||
await bridge.ainvoke(prepare=lambda: None, call=fail, fallback=fallback, adapt=str, error_context=context())
|
||||
|
||||
assert caught.value.status_code == 503
|
||||
|
||||
|
||||
def test_unknown_failure_is_preserved_without_fallback() -> None:
|
||||
error: Final = RuntimeError("unknown")
|
||||
bridge: Final = runtime.EndpointBinding(route="messages", load=object, enabled=enabled)
|
||||
|
||||
with pytest.raises(RuntimeError, match="unknown") as caught:
|
||||
bridge.invoke(
|
||||
prepare=lambda: None,
|
||||
call=lambda _binding, _request: (_ for _ in ()).throw(error),
|
||||
fallback=lambda: pytest.fail("fallback must not run"),
|
||||
adapt=str,
|
||||
error_context=context(),
|
||||
)
|
||||
|
||||
assert caught.value is error
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("process_enabled", "binding_available", "declined", "expected_message"),
|
||||
(
|
||||
pytest.param(False, True, False, "native messages endpoint is disabled", id="disabled"),
|
||||
pytest.param(None, False, False, "native messages endpoint is unavailable", id="unavailable"),
|
||||
pytest.param(
|
||||
None,
|
||||
True,
|
||||
True,
|
||||
"native messages endpoint declined the request: unsupported",
|
||||
id="declined",
|
||||
),
|
||||
),
|
||||
)
|
||||
def test_require_explains_why_rust_did_not_handle_request(
|
||||
process_enabled: bool | None,
|
||||
binding_available: bool,
|
||||
declined: bool,
|
||||
expected_message: str,
|
||||
) -> None:
|
||||
def call(_binding: object, _request: object) -> object:
|
||||
if declined:
|
||||
raise RustBridgeDeclined("unsupported")
|
||||
return object()
|
||||
|
||||
bridge: Final = runtime.EndpointBinding(
|
||||
route="messages",
|
||||
load=object if binding_available else lambda: None,
|
||||
enabled=lambda: process_enabled is not False,
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match=f"^{expected_message}$"):
|
||||
bridge.require(
|
||||
prepare=lambda: None,
|
||||
call=call,
|
||||
adapt=str,
|
||||
error_context=context(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("state", "expected", "expected_events"),
|
||||
(
|
||||
pytest.param("disabled", False, (), id="disabled"),
|
||||
pytest.param("ineligible", False, (), id="ineligible"),
|
||||
pytest.param("unavailable", False, ("load",), id="unavailable"),
|
||||
pytest.param("available", True, ("load",), id="available"),
|
||||
),
|
||||
)
|
||||
def test_can_attempt_only_enabled_available_requests(
|
||||
state: str,
|
||||
expected: bool,
|
||||
expected_events: tuple[str, ...],
|
||||
) -> None:
|
||||
events: list[str] = []
|
||||
|
||||
def load() -> object | None:
|
||||
events.append("load")
|
||||
return None if state == "unavailable" else object()
|
||||
|
||||
def prepare() -> int:
|
||||
events.append("prepare")
|
||||
return 3
|
||||
bridge: Final = runtime.EndpointBinding(route="messages", load=load, enabled=lambda: state != "disabled")
|
||||
|
||||
def call(_binding: object, request: int) -> int:
|
||||
events.append("call")
|
||||
return request * 2
|
||||
|
||||
async def acall(binding: object, request: int) -> int:
|
||||
return call(binding, request)
|
||||
|
||||
def adapt(value: int) -> str:
|
||||
events.append("adapt")
|
||||
return str(value)
|
||||
|
||||
result: Final = (
|
||||
await runtime.aattempt(
|
||||
load=load,
|
||||
enabled=state != "disabled",
|
||||
assert (
|
||||
bridge.can_attempt(
|
||||
eligible=state != "ineligible",
|
||||
prepare=prepare,
|
||||
call=acall,
|
||||
adapt=adapt,
|
||||
)
|
||||
if asynchronous
|
||||
else runtime.attempt(
|
||||
load=load,
|
||||
enabled=state != "disabled",
|
||||
eligible=state != "ineligible",
|
||||
prepare=prepare,
|
||||
call=call,
|
||||
adapt=adapt,
|
||||
)
|
||||
is expected
|
||||
)
|
||||
if state == "handled":
|
||||
assert result == runtime.Handled("6")
|
||||
assert events == ["load", "prepare", "call", "adapt"]
|
||||
else:
|
||||
assert result == runtime.NativeSkipped(runtime.NativeSkipReason(state))
|
||||
assert events == (["load"] if state == "unavailable" else [])
|
||||
assert tuple(events) == expected_events
|
||||
|
||||
|
||||
def test_native_endpoint_applies_partial_overrides_and_reset(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def native_sync() -> str:
|
||||
return "native"
|
||||
|
||||
async def native_async() -> str:
|
||||
return "native async"
|
||||
|
||||
def replacement_sync() -> str:
|
||||
return "replacement"
|
||||
|
||||
monkeypatch.setattr(
|
||||
bindings,
|
||||
"get_native_bridge",
|
||||
lambda: SimpleNamespace(chat_completions=native_sync, achat_completions=native_async),
|
||||
)
|
||||
endpoint: Final[runtime.EndpointDispatch[object, object]] = runtime.EndpointDispatch.native(
|
||||
route="test",
|
||||
sync=lambda native: native.chat_completions,
|
||||
asynchronous=lambda native: native.achat_completions,
|
||||
enabled=enabled,
|
||||
)
|
||||
|
||||
assert endpoint.sync.load() is native_sync
|
||||
assert endpoint.asynchronous.load() is native_async
|
||||
endpoint.override(sync=replacement_sync)
|
||||
assert endpoint.sync.load() is replacement_sync
|
||||
assert endpoint.asynchronous.load() is native_async
|
||||
endpoint.override(asynchronous=None)
|
||||
assert endpoint.sync.load() is replacement_sync
|
||||
assert endpoint.asynchronous.load() is None
|
||||
endpoint.reset()
|
||||
assert endpoint.sync.load() is native_sync
|
||||
assert endpoint.asynchronous.load() is native_async
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", (False, True))
|
||||
@pytest.mark.parametrize("phase", ("prepare", "call"))
|
||||
async def test_attempt_reports_failure_without_deciding_retry(asynchronous: bool, phase: str) -> None:
|
||||
error: Final = RuntimeError("native failure")
|
||||
async def test_response_adaptation_failure_never_authorizes_fallback(asynchronous: bool) -> None:
|
||||
def adapt(value: str) -> str:
|
||||
assert value == "provider response"
|
||||
raise RustBridgeDeclined("adapter failed after provider response")
|
||||
|
||||
def prepare() -> int:
|
||||
if phase == "prepare":
|
||||
raise error
|
||||
return 3
|
||||
async def native(binding: object, request: object) -> str:
|
||||
return "provider response"
|
||||
|
||||
def call(_binding: object, request: int) -> int:
|
||||
raise error
|
||||
async def fallback() -> str:
|
||||
pytest.fail("a received response must not be retried")
|
||||
|
||||
async def acall(binding: object, request: int) -> int:
|
||||
return call(binding, request)
|
||||
bridge = runtime.EndpointBinding(route="messages", load=object, enabled=enabled)
|
||||
|
||||
def adapt(value: int) -> str:
|
||||
pytest.fail("failed attempts cannot be adapted")
|
||||
|
||||
result: Final = (
|
||||
await runtime.aattempt(load=object, enabled=True, eligible=True, prepare=prepare, call=acall, adapt=adapt)
|
||||
if asynchronous
|
||||
else runtime.attempt(load=object, enabled=True, eligible=True, prepare=prepare, call=call, adapt=adapt)
|
||||
)
|
||||
assert isinstance(result, runtime.NativeFailed)
|
||||
assert result.error is error
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", (False, True))
|
||||
async def test_adaptation_failure_remains_distinct_from_native_failure(asynchronous: bool) -> None:
|
||||
error: Final = ValueError("invalid response")
|
||||
|
||||
async def acall(_binding: object, request: int) -> int:
|
||||
return request
|
||||
|
||||
def adapt(value: int) -> str:
|
||||
raise error
|
||||
|
||||
async def run() -> None:
|
||||
async def invoke() -> None:
|
||||
if asynchronous:
|
||||
await runtime.aattempt(load=object, enabled=True, eligible=True, prepare=lambda: 3, call=acall, adapt=adapt)
|
||||
await bridge.ainvoke(
|
||||
prepare=lambda: None, call=native, fallback=fallback, adapt=adapt, error_context=context()
|
||||
)
|
||||
else:
|
||||
runtime.attempt(
|
||||
load=object,
|
||||
enabled=True,
|
||||
eligible=True,
|
||||
prepare=lambda: 3,
|
||||
call=lambda binding, request: request,
|
||||
bridge.invoke(
|
||||
prepare=lambda: None,
|
||||
call=lambda binding, request: "provider response",
|
||||
fallback=lambda: pytest.fail("a received response must not be retried"),
|
||||
adapt=adapt,
|
||||
error_context=context(),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="invalid response") as caught:
|
||||
await run()
|
||||
assert caught.value is error
|
||||
with pytest.raises(RustBridgeDeclined, match="adapter failed"):
|
||||
await invoke()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", (False, True))
|
||||
@pytest.mark.parametrize("available, accepted", ((False, False), (True, False), (True, True)))
|
||||
async def test_preflight_runs_after_binding_selection_before_preparation(
|
||||
asynchronous: bool, available: bool, accepted: bool
|
||||
) -> None:
|
||||
events: list[str] = []
|
||||
|
||||
def load() -> object | None:
|
||||
events.append("load")
|
||||
return object() if available else None
|
||||
|
||||
def preflight() -> runtime.PythonFallback | None:
|
||||
events.append("preflight")
|
||||
return None if accepted else runtime.PythonFallback(runtime.PythonFallbackReason.NATIVE_DECLINED)
|
||||
|
||||
def prepare() -> int:
|
||||
events.append("prepare")
|
||||
return 7
|
||||
|
||||
def call(binding: object, request: int) -> int:
|
||||
events.append("native")
|
||||
return request
|
||||
|
||||
async def acall(binding: object, request: int) -> int:
|
||||
return call(binding, request)
|
||||
|
||||
def fallback() -> str:
|
||||
events.append("python")
|
||||
return "3"
|
||||
|
||||
async def afallback() -> str:
|
||||
return fallback()
|
||||
|
||||
endpoint: Final = runtime.EndpointBinding(route="ocr", load=load, enabled=enabled)
|
||||
result: Final = (
|
||||
await endpoint.ainvoke(
|
||||
prepare=prepare, call=acall, fallback=afallback, adapt=str, error_context=context(), preflight=preflight
|
||||
)
|
||||
if asynchronous
|
||||
else endpoint.invoke(
|
||||
prepare=prepare, call=call, fallback=fallback, adapt=str, error_context=context(), preflight=preflight
|
||||
)
|
||||
)
|
||||
assert result == ("7" if available and accepted else "3")
|
||||
assert events == (
|
||||
["load", "preflight", "prepare", "native"]
|
||||
if available and accepted
|
||||
else ["load", "preflight", "python"] if available else ["load", "python"]
|
||||
)
|
||||
|
||||
|
||||
def test_preflight_failure_is_not_a_native_decline() -> None:
|
||||
endpoint: Final = runtime.EndpointBinding(route="ocr", load=object, enabled=enabled)
|
||||
|
||||
def preflight() -> runtime.PythonFallback | None:
|
||||
raise ValueError("invalid acceptance contract")
|
||||
|
||||
with pytest.raises(ValueError, match="invalid acceptance contract"):
|
||||
endpoint.invoke(
|
||||
prepare=lambda: pytest.fail("must not prepare"),
|
||||
call=lambda binding, request: pytest.fail("must not invoke"),
|
||||
fallback=lambda: pytest.fail("must not fall back"),
|
||||
adapt=str,
|
||||
error_context=context(),
|
||||
preflight=preflight,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -4,8 +4,11 @@ import pytest
|
|||
|
||||
import litellm
|
||||
from litellm.llms.bedrock.audio_transcription import BedrockAudioTranscriptionRustDispatch
|
||||
from litellm.rust_bridge.request import NativeRequestContext, NativeRequestOptions, NativeTranscriptionRequest
|
||||
from litellm.rust_bridge.runtime import Handled
|
||||
from litellm.rust_bridge.request import (
|
||||
NativeRequestContext,
|
||||
NativeRequestOptions,
|
||||
NativeTranscriptionRequest,
|
||||
)
|
||||
|
||||
rust_bridge = importlib.import_module("litellm.rust_bridge.transcription")
|
||||
|
||||
|
|
@ -24,7 +27,6 @@ def reset_rust_transcription() -> None:
|
|||
class SyncBridge:
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[dict[str, object]] = []
|
||||
self.contexts: list[NativeRequestContext] = []
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
|
|
@ -34,9 +36,13 @@ class SyncBridge:
|
|||
context: NativeRequestContext,
|
||||
) -> dict[str, object]:
|
||||
self.calls.append(
|
||||
{"model": request.model, "audio": request.audio, "optional_params": request.optional_params}
|
||||
{
|
||||
"model": request.model,
|
||||
"audio": request.audio,
|
||||
"optional_params": request.optional_params,
|
||||
"bedrock": options.bedrock,
|
||||
}
|
||||
)
|
||||
self.contexts.append(context)
|
||||
return {"text": "hello"}
|
||||
|
||||
|
||||
|
|
@ -63,17 +69,9 @@ def test_enabled_sync_bridge_receives_audio() -> None:
|
|||
extra_headers=None,
|
||||
optional_params={"temperature": 0},
|
||||
timeout=5.0,
|
||||
stream=True,
|
||||
has_custom_client=True,
|
||||
input_source_kind="file",
|
||||
)
|
||||
assert isinstance(result, Handled)
|
||||
assert result.value == {"text": "hello"}
|
||||
assert result == {"text": "hello"}
|
||||
assert bridge.calls[0]["audio"] == {"data": "AQI=", "format": "wav", "filename": "audio.wav"}
|
||||
assert bridge.contexts[0].capabilities.execution_mode == "sync"
|
||||
assert bridge.contexts[0].capabilities.stream is True
|
||||
assert bridge.contexts[0].capabilities.has_custom_client is True
|
||||
assert bridge.contexts[0].capabilities.input_source_kind == "file"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -89,7 +87,7 @@ async def test_enabled_async_bridge() -> None:
|
|||
optional_params={},
|
||||
timeout=None,
|
||||
)
|
||||
assert result == Handled({"text": "async"})
|
||||
assert result == {"text": "async"}
|
||||
|
||||
|
||||
def test_loader_returns_none_without_native_extension(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
|
|
@ -100,8 +98,7 @@ def test_loader_returns_none_without_native_extension(monkeypatch: pytest.Monkey
|
|||
|
||||
|
||||
def test_dispatch_sync_path_requires_bridge(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
rust_bridge.configure_rust_transcription(transcription=None)
|
||||
monkeypatch.setattr("litellm.rust_bridge.bindings.get_native_bridge", lambda: None)
|
||||
monkeypatch.setattr(rust_bridge, "transcription", lambda **_: None)
|
||||
|
||||
with pytest.raises(RuntimeError, match="bridge is unavailable"):
|
||||
BedrockAudioTranscriptionRustDispatch().audio_transcriptions(
|
||||
|
|
@ -118,8 +115,10 @@ def test_dispatch_sync_path_requires_bridge(monkeypatch: pytest.MonkeyPatch) ->
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dispatch_async_path_requires_bridge(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
rust_bridge.configure_rust_transcription(atranscription=None)
|
||||
monkeypatch.setattr("litellm.rust_bridge.bindings.get_native_bridge", lambda: None)
|
||||
async def unavailable(**_: object) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(rust_bridge, "atranscription", unavailable)
|
||||
|
||||
with pytest.raises(RuntimeError, match="bridge is unavailable"):
|
||||
await BedrockAudioTranscriptionRustDispatch().async_audio_transcriptions(
|
||||
|
|
@ -136,7 +135,7 @@ async def test_dispatch_async_path_requires_bridge(monkeypatch: pytest.MonkeyPat
|
|||
|
||||
def test_bedrock_transcription_uses_rust_only_path() -> None:
|
||||
rust_bridge.configure_rust_transcription(
|
||||
transcription=lambda *_args, **_: {"text": "rust"},
|
||||
transcription=lambda request, *, options, context: {"text": "rust"},
|
||||
atranscription=None,
|
||||
)
|
||||
try:
|
||||
|
|
@ -152,7 +151,9 @@ def test_bedrock_transcription_uses_rust_only_path() -> None:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bedrock_atranscription_uses_rust_only_path() -> None:
|
||||
async def rust_response(*_args: object, **_: object) -> dict[str, object]:
|
||||
async def rust_response(
|
||||
request: NativeTranscriptionRequest, *, options: object, context: NativeRequestContext
|
||||
) -> dict[str, object]:
|
||||
return {"text": "rust"}
|
||||
|
||||
rust_bridge.configure_rust_transcription(transcription=None, atranscription=rust_response)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue