diff --git a/litellm/utils.py b/litellm/utils.py index 5970ea8876a..48db9c08e0d 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -65,6 +65,7 @@ from litellm.constants import ( DEFAULT_EMBEDDING_PARAM_VALUES, DEFAULT_MAX_LRU_CACHE_SIZE, DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT, + DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, DEFAULT_TRIM_RATIO, FUNCTION_DEFINITION_TOKEN_COUNT, @@ -278,7 +279,7 @@ except (ImportError, AttributeError, TypeError): # Convert to str (if necessary) claude_json_str = json.dumps(json_data) import importlib.metadata -from collections.abc import Callable, Iterable, Mapping, Sequence +from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast, get_args from litellm import utils as litellm_utils @@ -1646,9 +1647,6 @@ def client(original_function): kwargs=kwargs, ) - verbose_logger.info("Wrapper: Completed Call, calling success_handler") - completion.success(result, start_time, end_time) - # RETURN RESULT update_response_metadata = getattr(sys.modules[__name__], "update_response_metadata") update_response_metadata( result=result, @@ -1658,6 +1656,8 @@ def client(original_function): start_time=start_time, end_time=end_time, ) + verbose_logger.info("Wrapper: Completed Call, calling success_handler") + completion.success(result, start_time, end_time) return result except Exception as e: call_type = original_function.__name__ @@ -1934,13 +1934,13 @@ def client(original_function): args=args, ) - completion.success(result, start_time, end_time) # REBUILD EMBEDDING CACHING if ( isinstance(result, EmbeddingResponse) and _caching_handler_response is not None and _caching_handler_response.final_embedding_cached_response is not None ): + completion.success(result, start_time, end_time) return _llm_caching_handler._combine_cached_embedding_response_with_api_result( _caching_handler_response=_caching_handler_response, embedding_response=result, @@ -1956,6 +1956,7 @@ def client(original_function): start_time=start_time, end_time=end_time, ) + completion.success(result, start_time, end_time) return result except Exception as e: @@ -6953,7 +6954,26 @@ class TextCompletionStreamWrapper: raise StopAsyncIteration -def mock_completion_streaming_obj(model_response, mock_response, model, n: int | None = None): +def mock_stream_usage_chunk(model_response: ModelResponseStream, model: str, prompt_tokens: int) -> ModelResponseStream: + return ModelResponseStream( + id=model_response.id, + choices=[], # mutable-ok: ModelResponseStream only treats a list as explicit choices, a tuple gets a default choice + model=model, + usage=Usage( + prompt_tokens=prompt_tokens, + completion_tokens=DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, + total_tokens=prompt_tokens + DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, + ), + ) + + +def mock_completion_streaming_obj( + model_response: ModelResponseStream, + mock_response: str | MockException | ModelResponseStream, + model: str, + n: int | None = None, + prompt_tokens: int | None = None, +) -> Iterator[ModelResponseStream]: if isinstance(mock_response, litellm.MockException): raise mock_response if isinstance(mock_response, ModelResponseStream): @@ -6973,14 +6993,17 @@ def mock_completion_streaming_obj(model_response, mock_response, model, n: int | _all_choices.append(_streaming_choice) model_response.choices = _all_choices yield model_response + if prompt_tokens is not None: + yield mock_stream_usage_chunk(model_response, model=model, prompt_tokens=prompt_tokens) async def async_mock_completion_streaming_obj( - model_response, + model_response: ModelResponseStream, mock_response: str | MockException | ModelResponseStream, - model, + model: str, n: int | None = None, -): + prompt_tokens: int | None = None, +) -> AsyncIterator[ModelResponseStream]: if isinstance(mock_response, litellm.MockException): raise mock_response if isinstance(mock_response, ModelResponseStream): @@ -7000,6 +7023,8 @@ async def async_mock_completion_streaming_obj( _all_choices.append(_streaming_choice) model_response.choices = _all_choices yield model_response + if prompt_tokens is not None: + yield mock_stream_usage_chunk(model_response, model=model, prompt_tokens=prompt_tokens) ########## Reading Config File ############################ diff --git a/tests/test_litellm/litellm_core_utils/test_call_completion.py b/tests/test_litellm/litellm_core_utils/test_call_completion.py index 78006070c3f..dbf6bf0112f 100644 --- a/tests/test_litellm/litellm_core_utils/test_call_completion.py +++ b/tests/test_litellm/litellm_core_utils/test_call_completion.py @@ -282,7 +282,7 @@ async def test_async_ocr_wrapper_sends_final_failure_to_attached_completion( @pytest.mark.asyncio -async def test_async_ocr_wrapper_retains_completion_until_metadata_finishes( +async def test_async_ocr_wrapper_reports_metadata_failure_without_success( monkeypatch: pytest.MonkeyPatch, ) -> None: native_completion: Final = RecordingCompletion() @@ -307,7 +307,7 @@ async def test_async_ocr_wrapper_retains_completion_until_metadata_finishes( await wrapped() assert caught.value is metadata_error - assert native_completion.successes == [response] + assert native_completion.successes == [] assert native_completion.failures == [metadata_error, metadata_error] @@ -343,9 +343,23 @@ async def test_ocr_completion_stays_separate_from_marshaled_provider_options( ) -> OCRResponse: return run(request, resolve_secret, convert_file_document) - monkeypatch.setattr(import_module("litellm.ocr.main"), "rust_enabled", lambda: True) - monkeypatch.setattr(rust_ocr_bridge, "run", run) - monkeypatch.setattr(rust_ocr_bridge, "arun", arun) + def run_bridge( + request: rust_ocr_bridge.LiteLLMOcrRequest, + resolve_api_key: Callable[[str], str | None], + ) -> OCRResponse: + return run(request, resolve_api_key, lambda document: {}) + + async def arun_bridge( + request: rust_ocr_bridge.LiteLLMOcrRequest, + resolve_api_key: Callable[[str], str | None], + ) -> OCRResponse: + return await arun(request, resolve_api_key, lambda document: {}) + + ocr_main: Final = import_module("litellm.ocr.main") + monkeypatch.setattr(ocr_main, "rust_enabled", lambda: True) + monkeypatch.setattr(ocr_main, "_rust_ocr_supported", lambda request: True) + monkeypatch.setattr(ocr_main, "_run_rust_ocr", run_bridge) + monkeypatch.setattr(ocr_main, "_run_rust_aocr", arun_bridge) arguments: Final = { "model": "mistral/mistral-ocr-latest", "document": {"type": "document_url", "document_url": "https://example.com/doc.pdf"},