comments and small fix

Signed-off-by: B-Step62 <yuki.watanabe@databricks.com>
This commit is contained in:
B-Step62 2025-02-28 18:39:15 +09:00
parent 0d9a3dd50c
commit 5d8b359384
5 changed files with 141 additions and 35 deletions

View file

@ -749,9 +749,7 @@ async def _client_async_logging_helper(
f"Async Wrapper: Completed Call, calling async_success_handler: {logging_obj.async_success_handler}"
)
# check if user does not want this to be logged
asyncio.create_task(
logging_obj.async_success_handler(result, start_time, end_time)
)
await logging_obj.async_success_handler(result, start_time, end_time)
logging_obj.handle_sync_success_callbacks_for_async_calls(
result=result,
start_time=start_time,
@ -1121,7 +1119,13 @@ def client(original_function): # noqa: PLR0915
# LOG SUCCESS - handle streaming success logging in the _next_ object, remove `handle_success` once it's deprecated
verbose_logger.info("Wrapper: Completed Call, calling success_handler")
logging_obj.success_handler(result, start_time, end_time)
if litellm.sync_logging:
logging_obj.success_handler(result, start_time, end_time)
else:
executor.submit(
# NB: We already run this in a TPE so the handler itself should run sync
logging_obj.success_handler, result, start_time, end_time, synchronous=True,
)
# RETURN RESULT
update_response_metadata(
@ -1288,15 +1292,18 @@ def client(original_function): # noqa: PLR0915
)
# LOG SUCCESS - handle streaming success logging in the _next_ object
asyncio.create_task(
_client_async_logging_helper(
logging_obj=logging_obj,
result=result,
start_time=start_time,
end_time=end_time,
is_completion_with_fallbacks=is_completion_with_fallbacks,
)
async_logging_params = dict(
logging_obj=logging_obj,
result=result,
start_time=start_time,
end_time=end_time,
is_completion_with_fallbacks=is_completion_with_fallbacks,
)
if litellm.sync_logging:
await _client_async_logging_helper(**async_logging_params)
else:
asyncio.create_task(_client_async_logging_helper(**async_logging_params))
logging_obj.handle_sync_success_callbacks_for_async_calls(
result=result,
start_time=start_time,