mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
* fix(logging): fire one CustomLogger hook per async call outside the async flag list aspeech, atranscription, aimage_generation, aimage_edit, aocr and asearch are missing from is_sync_litellm_request's flag list, so the sync success and failure handlers called a CustomLogger's sync hook on top of the async hook the async handler had already called for the same request. The sync handlers now skip a CustomLogger's log_success_event and log_failure_event when that run is paired with an async handler run (for_async_request) and the async handler gives the logger its own async hook. The async client wrapper, dispatch_failure_handlers, the proxy-only failure dispatch and the native bridge's sync failure step for an async call mark their sync failure run as paired. Resolves LIT-9486 * fix(logging): keep the sync CustomLogger hook when no async success runs The async client wrapper skips async success logging for calls carrying a fallbacks kwarg, and the native bridge skips it for internal and fallbacks calls, yet the paired sync success run still dropped a both-hooks logger's sync hook, so that logger got no success event at all. The sync success run now suppresses the sync hook only when the async success handler was actually scheduled for the call: the client wrapper passes async_handler_runs=not is_completion_with_fallbacks, and the native host records per call whether it enqueued or deferred the async success. Tests that pinned the doubled failure hook on native aocr, messages and responses calls and on proxy-rejected requests now expect the async hook alone. The integration test adds /v1/messages and /v1/responses cells and runs the key-auth rejection on every route. * fix(logging): type dispatch_success_handlers kwargs The new async_handler_runs keyword on success_handler made the untyped **kwargs that dispatch_success_handlers forwards an unknown argument for it, which grew the reportUnknownArgumentType count past its merge-base ceiling. Typing the forwarded kwargs as Any, as _success_handler_body already does, removes those diagnostics without changing what is forwarded. * test(integration): cover CustomLogger hook counts across modes, chaos and router strategies Moves the hook count rig into tests/integration/_support/hook_counts.py and extends the cells it drives: - streamed chat, Messages and Responses requests, success and failure - loggers the async handler never reaches (sync-only overrides, sync event lists only) keep one sync hook per request on every surface outside the async flag list, key-auth rejections included - a concurrent burst while the logger's destination stalls and then fails, and a burst while a worker is killed, never give a request two hooks - least-busy keeps counting a request in flight while async speech requests finish around it, and usage-based routing v2 counts an async image request's tokens once * test(integration): Datadog and LangSmith deliver one record per async request Points the built-in Datadog, LangSmith and OpenTelemetry loggers at a local sink and checks every async surface outside the flag list: one Datadog log entry and one LangSmith run per request (two on the merge base), and never two OpenTelemetry spans, which its own emit-once guard already ensured. * fix(logging): give a both-hooks logger its sync hook when the async run never reaches it The sync hook was skipped whenever the logger sat in the async callback list, so a logger the async run never reached got no hook at all: a sibling stalled past the logging worker timeout, a batch poll that skips default logging, or an async failure run cut off. The sync run now holds the hook and the async run hands it back unless it reached the logger * fix(logging): run the held sync success hooks when the proxy drops a blocked response's success A post-call guardrail block drops the parked async success, which left the sync success hook held for it unrun. Base gave a both-hooks logger that hook on these call types, so the drop now ends the pairing and the held hook runs as before * fix(logging): run held sync hooks only after every sync run of the request has closed A held hook released by the async run could run on another thread while the sync run was still walking its other callbacks. The pairing now waits until the async run has ended and every sync run has closed, and the last sync run to close runs the held hooks itself * fix(logging): end the async run when its coroutine is dropped before it starts A logging queue that drops the async success coroutine, or a deferred success that fails to schedule, left the sync hooks held for it unrun. The enqueued async run now ends when its coroutine is discarded, and the proxy ends it when firing the deferred success raises * fix(logging): leave a still-queued async run to the exit flush The discard finalizer ran at interpreter exit before the logging worker's exit flush, so a queued request could get its held sync hook and then its async hook * test: keep the router and file-stream tests in step with main Main gives each owned proxy a private Redis by default, so the router cells share the worker's Redis to read its counters, and the file-stream stub takes the hook pairing the stream now passes * fix(logging): run held sync hooks when a during_call guardrail raises or the exit flush gives up A during_call guardrail that raised after the LLM call parked its deferred success skipped the proxy's flush, and async runs the exit flush never started were never ended, so a both-hooks logger lost its sync hook in both cases * fix(logging): skip the sync failure hook when the async failure run logs to the logger An async call outside is_sync_litellm_request's flag list ran both failure handlers, so a CustomLogger with both failure hooks got both. The sync run of an async request now skips log_failure_event for a logger that overrides async_log_failure_event and is registered for async failures. The four sync runs of async failures (dispatch_failure_handlers, the @client async wrapper, the proxy-only failure thread and the Rust host) pass for_async_request=True. Drops the HookPairing fallback. Successes are left to #47192, which already gives such a logger only async_log_success_event. --------- Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com>
596 lines
21 KiB
Python
596 lines
21 KiB
Python
import asyncio
|
|
import copy
|
|
import gc
|
|
import queue
|
|
import threading
|
|
from collections.abc import Mapping
|
|
from types import MappingProxyType
|
|
from typing import Final
|
|
|
|
import pytest
|
|
from hypothesis import HealthCheck, given, settings
|
|
from hypothesis import strategies as st
|
|
|
|
import litellm
|
|
from litellm.integrations.custom_guardrail import CustomGuardrail
|
|
from litellm.integrations.custom_logger import CustomLogger
|
|
from litellm.llms.base_llm.ocr.transformation import OCRResponse
|
|
from tests.test_litellm_rust.ocr.data import (
|
|
OCR_DOCUMENT,
|
|
OCR_RESPONSE,
|
|
call,
|
|
call_aocr,
|
|
call_ocr,
|
|
)
|
|
from tests.test_litellm_rust.support.callback_recorder import (
|
|
RecordingLogger,
|
|
drain_logging,
|
|
request_body,
|
|
request_headers,
|
|
)
|
|
from tests.test_litellm_rust.support.isolation import isolated_callback_registries, rebound
|
|
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
|
|
|
|
pytestmark = pytest.mark.requires_rust_extension
|
|
|
|
|
|
def test_native_ocr_pre_call_callback_receives_transformed_provider_request(server: RecordingServer) -> None:
|
|
observations: Final = []
|
|
|
|
class Observe(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
observations.append((model, copy.deepcopy(kwargs["additional_args"])))
|
|
|
|
call_ocr(server, callbacks=[Observe()], pages=[0])
|
|
|
|
assert len(observations) == 1
|
|
model, additional_args = observations[0]
|
|
assert model == "mistral-ocr-latest"
|
|
assert additional_args["api_base"] == f"{server.base_url}/v1/ocr"
|
|
assert additional_args["complete_input_dict"] == {
|
|
"model": "mistral-ocr-latest",
|
|
"document": OCR_DOCUMENT,
|
|
"pages": [0],
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_native_ocr_does_not_restore_keywords_deleted_by_a_pre_call_hook(server: RecordingServer) -> None:
|
|
class DropKeywords(CustomLogger):
|
|
async def async_pre_call_deployment_hook(
|
|
self, kwargs: dict[str, object], call_type: object
|
|
) -> dict[str, object]:
|
|
return {name: value for name, value in kwargs.items() if name not in ("pages", "extra_headers")}
|
|
|
|
with rebound(litellm, "callbacks", [DropKeywords()]):
|
|
await call_aocr(server, pages=[0], extra_headers={"x-deleted": "remove this header"})
|
|
|
|
assert len(server.requests) == 1
|
|
sent: Final = server.requests[0]
|
|
assert sent.body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT}
|
|
assert sent.headers.get("x-deleted") is None
|
|
|
|
|
|
@pytest.mark.parametrize("raise_after_edit", [False, True], ids=["callback-returns", "callback-raises"])
|
|
def test_native_ocr_pre_call_body_edit_reaches_next_callback_and_provider(
|
|
server: RecordingServer, raise_after_edit: bool
|
|
) -> None:
|
|
observed: Final = []
|
|
|
|
class Edit(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
request_body(kwargs)["include_image_base64"] = True
|
|
if raise_after_edit:
|
|
raise RuntimeError("pre-call callback failed")
|
|
|
|
class Observe(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
observed.append(copy.deepcopy(request_body(kwargs)))
|
|
|
|
call_ocr(server, callbacks=[Edit(), Observe()], include_image_base64=False)
|
|
|
|
assert observed[0]["include_image_base64"] is True
|
|
assert server.requests[0].body["include_image_base64"] is True
|
|
|
|
|
|
def test_native_ocr_pre_call_header_edit_reaches_next_callback_and_provider(server: RecordingServer) -> None:
|
|
observed: Final = []
|
|
|
|
class Edit(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
request_headers(kwargs)["x-audit-tag"] = "reviewed"
|
|
|
|
class Observe(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
observed.append(dict(request_headers(kwargs)))
|
|
|
|
call_ocr(server, callbacks=[Edit(), Observe()])
|
|
|
|
assert observed[0]["x-audit-tag"] == "reviewed"
|
|
assert server.requests[0].headers["x-audit-tag"] == "reviewed"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
|
async def test_native_ocr_pre_call_nested_document_edit_updates_caller_callback_and_provider_references(
|
|
server: RecordingServer, asynchronous: bool
|
|
) -> None:
|
|
original: Final = dict(OCR_DOCUMENT)
|
|
replacement_url: Final = "data:application/pdf;base64,ZGVm"
|
|
retained: Final = []
|
|
aliases: Final = []
|
|
|
|
class Retain(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
aliases.append(request_body(kwargs)["document"] is original)
|
|
retained.append(request_body(kwargs)["document"])
|
|
|
|
class Edit(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
original["document_url"] = replacement_url
|
|
|
|
arguments: Final = {
|
|
"model": "mistral/mistral-ocr-latest",
|
|
"document": original,
|
|
"api_key": "test-key",
|
|
"api_base": server.base_url,
|
|
"callbacks": [Retain(), Edit()],
|
|
}
|
|
response: Final = await call_aocr(server, **arguments) if asynchronous else call_ocr(server, **arguments)
|
|
|
|
assert aliases == [True]
|
|
assert retained[0]["document_url"] == replacement_url
|
|
assert original["document_url"] == replacement_url
|
|
assert server.requests[0].body["document"]["document_url"] == replacement_url
|
|
assert response.pages[0].markdown == "native OCR response"
|
|
|
|
|
|
def test_native_ocr_pre_call_body_rebinding_is_visible_to_callbacks_but_not_provider(
|
|
server: RecordingServer,
|
|
) -> None:
|
|
observed: Final = []
|
|
|
|
class Rebind(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
kwargs["additional_args"]["complete_input_dict"] = {"replacement": True}
|
|
|
|
class Observe(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
observed.append(request_body(kwargs))
|
|
|
|
call_ocr(server, callbacks=[Rebind(), Observe()])
|
|
|
|
assert observed == [{"replacement": True}]
|
|
assert server.requests[0].body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT}
|
|
|
|
|
|
def test_native_ocr_callback_retained_body_observes_later_callback_mutation(server: RecordingServer) -> None:
|
|
queued: Final = []
|
|
|
|
class QueuePayload(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
queued.append(request_body(kwargs))
|
|
|
|
class Edit(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
request_body(kwargs)["queued-edit"] = True
|
|
|
|
call_ocr(server, callbacks=[QueuePayload(), Edit()])
|
|
|
|
assert queued[0]["queued-edit"] is True
|
|
|
|
|
|
def test_native_ocr_success_callback_receives_state_added_by_pre_call_callback(server: RecordingServer) -> None:
|
|
token: Final = object()
|
|
terminal_tokens: queue.SimpleQueue[object] = queue.SimpleQueue()
|
|
finished: Final = threading.Event()
|
|
|
|
class Stash(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
kwargs["test-token"] = token
|
|
|
|
def log_success_event(self, kwargs, response_obj, start_time, end_time):
|
|
terminal_tokens.put(kwargs["test-token"])
|
|
finished.set()
|
|
|
|
call_ocr(server, callbacks=[Stash()])
|
|
|
|
assert finished.wait(10)
|
|
assert terminal_tokens.get_nowait() is token
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_native_aocr_success_callback_receives_call_id_metadata_and_response(
|
|
server: RecordingServer,
|
|
) -> None:
|
|
recorder: Final = RecordingLogger()
|
|
|
|
await call_aocr(
|
|
server,
|
|
callbacks=[recorder],
|
|
litellm_call_id="ocr-success",
|
|
metadata={"source": "callback-test"},
|
|
)
|
|
events: Final = await recorder.wait_for_async("async_log_success_event")
|
|
|
|
assert len(events) == 1
|
|
assert events[0].call_type == "aocr"
|
|
assert events[0].kwargs["litellm_call_id"] == "ocr-success"
|
|
assert events[0].kwargs["litellm_params"]["metadata"]["source"] == "callback-test"
|
|
assert events[0].response.pages[0].markdown == "native OCR response"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_native_aocr_failure_callbacks_receive_call_type_error_and_no_response(
|
|
server: RecordingServer,
|
|
) -> None:
|
|
server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500))
|
|
observations: Final = []
|
|
|
|
class Observe(CustomLogger):
|
|
def log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
|
observations.append(("sync", kwargs["call_type"], kwargs["exception"], response_obj))
|
|
|
|
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
|
observations.append(("async", kwargs["call_type"], kwargs["exception"], response_obj))
|
|
|
|
with pytest.raises(litellm.InternalServerError):
|
|
await call_aocr(server, callbacks=[Observe()])
|
|
|
|
assert [observation[0] for observation in observations] == ["async"]
|
|
assert all(observation[1] == "aocr" for observation in observations)
|
|
assert all(isinstance(observation[2], litellm.InternalServerError) for observation in observations)
|
|
assert all(observation[3] is None for observation in observations)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_native_aocr_pre_call_callback_runs_on_caller_loop_and_thread(server: RecordingServer) -> None:
|
|
caller_loop: Final = asyncio.get_running_loop()
|
|
caller_thread: Final = threading.current_thread()
|
|
recorder: Final = RecordingLogger()
|
|
|
|
await call_aocr(server, callbacks=[recorder])
|
|
|
|
events: Final = await recorder.wait_for_async("log_pre_api_call")
|
|
assert len(events) == 1
|
|
assert events[0].loop is caller_loop
|
|
assert events[0].thread is caller_thread
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_native_aocr_failure_callbacks_receive_state_added_by_pre_call_callback(
|
|
server: RecordingServer,
|
|
) -> None:
|
|
server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500))
|
|
token: Final = object()
|
|
observed: Final = []
|
|
|
|
class TrackInFlightRequest(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
kwargs["request-token"] = token
|
|
|
|
def log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
|
observed.append(("sync", kwargs["request-token"]))
|
|
|
|
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
|
observed.append(("async", kwargs["request-token"]))
|
|
|
|
with pytest.raises(litellm.InternalServerError):
|
|
await call_aocr(server, callbacks=[TrackInFlightRequest()])
|
|
|
|
assert [event for event, _ in observed] == ["async"]
|
|
assert all(observed_token is token for _, observed_token in observed)
|
|
|
|
|
|
def test_native_ocr_dispatches_each_callback_phase_once_when_logger_is_registered_multiple_times(
|
|
server: RecordingServer,
|
|
) -> None:
|
|
recorder: Final = RecordingLogger()
|
|
|
|
call_ocr(
|
|
server,
|
|
callbacks=[recorder, recorder],
|
|
success_callback=[recorder],
|
|
failure_callback=[recorder],
|
|
)
|
|
recorder.wait_for("log_success_event")
|
|
|
|
assert recorder.names.count("log_pre_api_call") == 1
|
|
assert recorder.names.count("logging_hook") == 1
|
|
assert recorder.names.count("log_success_event") == 1
|
|
assert "log_failure_event" not in recorder.names
|
|
|
|
|
|
JSON_SCALARS: Final = (
|
|
st.none()
|
|
| st.booleans()
|
|
| st.integers(min_value=-(2**63), max_value=2**63 - 1)
|
|
| st.floats(allow_nan=False, allow_infinity=False)
|
|
| st.text(max_size=8)
|
|
)
|
|
JSON_VALUES: Final = st.recursive(
|
|
JSON_SCALARS,
|
|
lambda children: st.lists(children, max_size=3) | st.dictionaries(st.text(max_size=6), children, max_size=3),
|
|
max_leaves=8,
|
|
)
|
|
|
|
|
|
class ApplyEdits(CustomLogger):
|
|
def __init__(self, edits: Mapping[str, object]) -> None:
|
|
super().__init__()
|
|
self.edits: Final = edits
|
|
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
request_body(kwargs).update(copy.deepcopy(dict(self.edits)))
|
|
|
|
|
|
@settings(max_examples=25, deadline=None, suppress_health_check=[HealthCheck.function_scoped_fixture])
|
|
@given(edits=st.dictionaries(st.from_regex(r"x_[a-z]{1,6}", fullmatch=True), JSON_VALUES, max_size=3))
|
|
def test_native_ocr_provider_receives_the_body_exactly_as_pre_call_callbacks_left_it(
|
|
server: RecordingServer, edits: dict[str, object]
|
|
) -> None:
|
|
server.expected_requests = None
|
|
|
|
with isolated_callback_registries():
|
|
call_ocr(server, callbacks=[ApplyEdits(MappingProxyType(edits))])
|
|
|
|
assert server.requests[-1].body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT, **edits}
|
|
|
|
|
|
@pytest.mark.parametrize("hook", ["log_pre_api_call", "logging_hook", "log_success_event"])
|
|
def test_native_ocr_sync_hooks_see_no_running_event_loop(server: RecordingServer, hook: str) -> None:
|
|
recorder: Final = RecordingLogger()
|
|
|
|
call_ocr(server, callbacks=[recorder])
|
|
|
|
[event] = recorder.wait_for(hook)
|
|
assert event.loop is None
|
|
assert (event.thread is threading.current_thread()) == (hook == "log_pre_api_call")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
|
async def test_native_ocr_payload_a_callback_retains_outlives_the_call_intact(
|
|
server: RecordingServer, asynchronous: bool
|
|
) -> None:
|
|
retained: Final = []
|
|
|
|
class Retain(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
retained.append((kwargs, request_body(kwargs), request_headers(kwargs)))
|
|
|
|
await call(server, asynchronous, callbacks=[Retain()])
|
|
await drain_logging()
|
|
gc.collect()
|
|
|
|
[(details, body, headers)] = retained
|
|
assert body == server.requests[0].body
|
|
assert headers
|
|
assert all(server.requests[0].headers[name.lower()] == value for name, value in headers.items())
|
|
assert details["additional_args"]["complete_input_dict"] is body
|
|
assert details["additional_args"]["headers"] is headers
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("family", ["sync", "async"])
|
|
async def test_native_ocr_success_callbacks_share_one_logging_payload(server: RecordingServer, family: str) -> None:
|
|
queued: Final = []
|
|
finished: Final = threading.Event()
|
|
|
|
def queue_payload(kwargs: dict[str, object]) -> None:
|
|
queued.append(kwargs["standard_logging_object"])
|
|
|
|
def strip_payload(kwargs: dict[str, object]) -> None:
|
|
payload: Final = kwargs["standard_logging_object"]
|
|
assert isinstance(payload, dict)
|
|
payload["stripped-by-a-later-callback"] = True
|
|
finished.set()
|
|
|
|
class QueuePayload(CustomLogger):
|
|
if family == "sync":
|
|
|
|
def log_success_event(self, kwargs, response_obj, start_time, end_time):
|
|
queue_payload(kwargs)
|
|
|
|
else:
|
|
|
|
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
|
queue_payload(kwargs)
|
|
|
|
class StripPayload(CustomLogger):
|
|
if family == "sync":
|
|
|
|
def log_success_event(self, kwargs, response_obj, start_time, end_time):
|
|
strip_payload(kwargs)
|
|
|
|
else:
|
|
|
|
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
|
strip_payload(kwargs)
|
|
|
|
await call(server, family == "async", callbacks=[QueuePayload(), StripPayload()])
|
|
await drain_logging()
|
|
|
|
assert await asyncio.to_thread(finished.wait, 10)
|
|
assert [payload["stripped-by-a-later-callback"] for payload in queued] == [True]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_native_aocr_state_stashed_before_a_blocking_hook_raises_reaches_failure_callbacks(
|
|
server: RecordingServer,
|
|
) -> None:
|
|
token: Final = object()
|
|
observed: Final = []
|
|
|
|
class Blocked(Exception):
|
|
pass
|
|
|
|
class Block(CustomGuardrail):
|
|
async def async_post_call_success_deployment_hook(self, request_data, response, call_type):
|
|
request_data["litellm_logging_obj"].model_call_details["blocked-by"] = token
|
|
raise Blocked("blocked after the provider answered")
|
|
|
|
def log_success_event(self, kwargs, response_obj, start_time, end_time):
|
|
observed.append(("success", None, None))
|
|
|
|
def log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
|
observed.append(("sync", kwargs.get("blocked-by"), kwargs["exception"]))
|
|
|
|
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
|
observed.append(("async", kwargs.get("blocked-by"), kwargs["exception"]))
|
|
|
|
litellm.callbacks.append(Block())
|
|
|
|
with pytest.raises(Blocked) as raised:
|
|
await call_aocr(server)
|
|
await drain_logging()
|
|
|
|
assert observed == [("async", token, raised.value)]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
|
async def test_native_azure_ocr_resolves_token_before_pre_call_on_caller_context(
|
|
server: RecordingServer,
|
|
isolated_azure_auth: None,
|
|
asynchronous: bool,
|
|
) -> None:
|
|
from contextvars import ContextVar
|
|
|
|
context: Final = ContextVar("azure-token-context", default="missing")
|
|
context.set("caller")
|
|
caller_thread: Final = threading.current_thread()
|
|
caller_loop: Final = asyncio.get_running_loop()
|
|
observations: Final = []
|
|
|
|
class Provider:
|
|
def __call__(self) -> str:
|
|
assert context.get() == "caller"
|
|
assert threading.current_thread() is caller_thread
|
|
assert asyncio.get_running_loop() is caller_loop
|
|
observations.append("token")
|
|
return "caller-token"
|
|
|
|
class Edit(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
assert request_headers(kwargs)["Authorization"] == "Bearer caller-token"
|
|
observations.append("pre_call")
|
|
request_headers(kwargs)["Authorization"] = "Bearer edited"
|
|
|
|
provider: Final = Provider()
|
|
arguments: Final = {
|
|
"model": "azure_ai/mistral-ocr-latest",
|
|
"api_key": None,
|
|
"azure_ad_token_provider": provider,
|
|
"callbacks": [Edit()],
|
|
}
|
|
response: Final = await call_aocr(server, **arguments) if asynchronous else call_ocr(server, **arguments)
|
|
assert response.pages[0].markdown == "native OCR response"
|
|
assert observations == ["token", "pre_call"]
|
|
assert server.requests[0].headers["authorization"] == "Bearer edited"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
|
async def test_native_azure_ocr_token_provider_can_make_nested_native_ocr_call(
|
|
server: RecordingServer,
|
|
isolated_azure_auth: None,
|
|
asynchronous: bool,
|
|
) -> None:
|
|
server.expected_requests = 2
|
|
calls: Final = []
|
|
|
|
def provider() -> str:
|
|
calls.append("token")
|
|
nested: Final = call_ocr(server)
|
|
assert nested.pages[0].markdown == "native OCR response"
|
|
return "outer-token"
|
|
|
|
arguments: Final = {
|
|
"model": "azure_ai/mistral-ocr-latest",
|
|
"api_key": None,
|
|
"azure_ad_token_provider": provider,
|
|
}
|
|
response: Final = await call_aocr(server, **arguments) if asynchronous else call_ocr(server, **arguments)
|
|
assert response.pages[0].markdown == "native OCR response"
|
|
assert calls == ["token"]
|
|
assert [request.headers["authorization"] for request in server.requests] == [
|
|
"Bearer test-key",
|
|
"Bearer outer-token",
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_concurrent_native_azure_ocr_calls_isolate_token_results_and_error(
|
|
server: RecordingServer,
|
|
isolated_azure_auth: None,
|
|
) -> None:
|
|
server.expected_requests = 2
|
|
|
|
async def request(token: str, fail: bool) -> object:
|
|
def provider() -> str:
|
|
if fail:
|
|
raise ValueError(token)
|
|
return token
|
|
|
|
return await call_aocr(
|
|
server,
|
|
model="azure_ai/mistral-ocr-latest",
|
|
api_key=None,
|
|
azure_ad_token_provider=provider,
|
|
)
|
|
|
|
responses: Final = await asyncio.gather(
|
|
request("first", False),
|
|
request("failed", True),
|
|
request("second", False),
|
|
return_exceptions=True,
|
|
)
|
|
assert isinstance(responses[0], OCRResponse)
|
|
assert isinstance(responses[1], litellm.APIConnectionError)
|
|
assert "Failed to get Azure AD token: failed" in str(responses[1])
|
|
assert isinstance(responses[2], OCRResponse)
|
|
assert sorted(request.headers["authorization"] for request in server.requests) == [
|
|
"Bearer first",
|
|
"Bearer second",
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_native_azure_ocr_releases_token_provider_after_cancellation(
|
|
server: RecordingServer,
|
|
isolated_azure_auth: None,
|
|
) -> None:
|
|
import gc
|
|
import weakref
|
|
|
|
from tests.test_litellm_rust.support.callback_recorder import drain_logging
|
|
|
|
class Provider:
|
|
def __call__(self) -> str:
|
|
return "caller-token"
|
|
|
|
async def invoke() -> weakref.ReferenceType[Provider]:
|
|
provider: Final = Provider()
|
|
reference: Final = weakref.ref(provider)
|
|
server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.1))
|
|
task: Final = asyncio.create_task(
|
|
call_aocr(
|
|
server,
|
|
model="azure_ai/mistral-ocr-latest",
|
|
api_key=None,
|
|
azure_ad_token_provider=provider,
|
|
)
|
|
)
|
|
await server.wait_for_requests(1)
|
|
assert reference() is provider
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
return reference
|
|
|
|
reference: Final = await invoke()
|
|
await drain_logging()
|
|
await asyncio.sleep(0)
|
|
gc.collect()
|
|
assert reference() is None
|