Skip to content

azad.ainetwork.callback_manager Module

azad.ainetwork.callback_manager

Callback management for LiteLLM.

This module provides a context manager for safely adding and removing callbacks to LiteLLM, ensuring proper cleanup even in the case of exceptions.

Attributes

logger module-attribute 

logger = getLogger(__name__)

Functions

clear_all_callbacks 

clear_all_callbacks()

Utility function to manually clear all registered callbacks.

Use this as a last resort if the application is experiencing issues with callbacks.

Source code in azad/ainetwork/callback_manager.py 
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
def clear_all_callbacks():
    """
    Utility function to manually clear all registered callbacks.

    Use this as a last resort if the application is experiencing issues with callbacks.
    """
    if hasattr(litellm, "logging_callback_manager"):
        litellm.logging_callback_manager._reset_all_callbacks()
    else:
        # Fallback to clearing individual lists
        callback_lists = ["success_callback", "failure_callback", "_async_success_callback", 
                         "_async_failure_callback", "callbacks"]
        for list_name in callback_lists:
            if hasattr(litellm, list_name):
                setattr(litellm, list_name, [])

    logger.warning("All LiteLLM callbacks have been cleared")

remove_callback_from_all_lists 

remove_callback_from_all_lists(callback: Callable)

Remove a callback function from all litellm callback lists where it might exist.

Source code in azad/ainetwork/callback_manager.py 
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
def remove_callback_from_all_lists(callback: Callable):
    """
    Remove a callback function from all litellm callback lists where it might exist.
    """
    # List of all callback lists to check
    callback_lists = ["success_callback", "failure_callback", "_async_success_callback",
                     "_async_failure_callback", "callbacks"]

    # Remove the callback from each list if present
    removed_count = 0
    for list_name in callback_lists:
        if hasattr(litellm, list_name):
            callback_list = getattr(litellm, list_name)
            # Remove all instances of the callback
            while callback in callback_list:
                callback_list.remove(callback)
                removed_count += 1

    if removed_count > 0:
        logger.debug(f"Removed callback from {removed_count} litellm callback lists")

    return removed_count

managed_callbacks async 

managed_callbacks(success_callback: Optional[Callable] = None, failure_callback: Optional[Callable] = None) -> AsyncGenerator[None, None]

Context manager to safely add and remove callbacks to LiteLLM.

This ensures that callbacks are properly registered and cleaned up, even in the case of exceptions.

Parameters:

  • success_callback (Optional[Callable], default: None ) –

    Optional callback for successful completions

  • failure_callback (Optional[Callable], default: None ) –

    Optional callback for failed completions

Yields:

  • AsyncGenerator[None, None] –

    None

Source code in azad/ainetwork/callback_manager.py 
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
@asynccontextmanager
async def managed_callbacks(
    success_callback: Optional[Callable] = None,
    failure_callback: Optional[Callable] = None
) -> AsyncGenerator[None, None]:
    """
    Context manager to safely add and remove callbacks to LiteLLM.

    This ensures that callbacks are properly registered and cleaned up,
    even in the case of exceptions.

    Args:
        success_callback: Optional callback for successful completions
        failure_callback: Optional callback for failed completions

    Yields:
        None
    """
    try:
        # Check if we need to clear callbacks to prevent hitting the limit
        if hasattr(litellm, "logging_callback_manager"):
            all_callbacks = litellm.logging_callback_manager._get_all_callbacks()
            if len(all_callbacks) >= 25:  # Safety threshold before hitting 30
                logger.warning(f"Approaching MAX_CALLBACKS limit ({len(all_callbacks)}/30). Clearing all callbacks.")
                clear_all_callbacks()

        # Add success callback if provided
        if success_callback:
            if not hasattr(litellm, "success_callback"):
                litellm.success_callback = []

            if success_callback not in litellm.success_callback:
                litellm._async_success_callback.append(success_callback)

        # Add failure callback if provided
        if failure_callback:
            if not hasattr(litellm, "failure_callback"):
                litellm.failure_callback = []

            if failure_callback not in litellm.failure_callback:
                litellm.failure_callback.append(failure_callback)

        # Yield control back to the caller
        yield

    except Exception as e:
        logger.error(f"Error in callback management: {str(e)}")
        raise

    finally:
        # Always attempt to remove our callbacks
        try:
            if success_callback:
                remove_callback_from_all_lists(success_callback)

            if failure_callback:
                remove_callback_from_all_lists(failure_callback)

        except Exception as e:
            logger.error(f"Error removing callbacks: {str(e)}")
            # Reset as last resort if error during removal
            clear_all_callbacks()