fix checking if _known_custom_logger_compatible_callbacks

This commit is contained in:
Ishaan Jaff 2024-07-22 15:43:43 -07:00
parent d8d08a1ba3
commit 15c109f023
2 changed files with 7 additions and 7 deletions

View file

@ -23,11 +23,11 @@ def initialize_callbacks_on_proxy(
) )
if isinstance(value, list): if isinstance(value, list):
imported_list: List[Any] = [] imported_list: List[Any] = []
known_compatible_callbacks = list(
get_args(litellm._custom_logger_compatible_callbacks_literal)
)
for callback in value: # ["presidio", <my-custom-callback>] for callback in value: # ["presidio", <my-custom-callback>]
if isinstance(callback, str) and callback in known_compatible_callbacks: if (
isinstance(callback, str)
and callback in litellm._known_custom_logger_compatible_callbacks
):
imported_list.append(callback) imported_list.append(callback)
elif isinstance(callback, str) and callback == "otel": elif isinstance(callback, str) and callback == "otel":
from litellm.integrations.opentelemetry import OpenTelemetry from litellm.integrations.opentelemetry import OpenTelemetry

View file

@ -158,6 +158,7 @@ from typing import (
Tuple, Tuple,
Union, Union,
cast, cast,
get_args,
) )
from .caching import Cache from .caching import Cache
@ -405,7 +406,6 @@ def function_setup(
# Pop the async items from input_callback in reverse order to avoid index issues # Pop the async items from input_callback in reverse order to avoid index issues
for index in reversed(removed_async_items): for index in reversed(removed_async_items):
litellm.input_callback.pop(index) litellm.input_callback.pop(index)
if len(litellm.success_callback) > 0: if len(litellm.success_callback) > 0:
removed_async_items = [] removed_async_items = []
for index, callback in enumerate(litellm.success_callback): # type: ignore for index, callback in enumerate(litellm.success_callback): # type: ignore
@ -417,9 +417,9 @@ def function_setup(
# we only support async dynamo db logging for acompletion/aembedding since that's used on proxy # we only support async dynamo db logging for acompletion/aembedding since that's used on proxy
litellm._async_success_callback.append(callback) litellm._async_success_callback.append(callback)
removed_async_items.append(index) removed_async_items.append(index)
elif callback == "langsmith": elif callback in litellm._known_custom_logger_compatible_callbacks:
callback_class = litellm.litellm_core_utils.litellm_logging._init_custom_logger_compatible_class( # type: ignore callback_class = litellm.litellm_core_utils.litellm_logging._init_custom_logger_compatible_class( # type: ignore
callback, internal_usage_cache=None, llm_router=None callback, internal_usage_cache=None, llm_router=None # type: ignore
) )
# don't double add a callback # don't double add a callback