litellm/enterprise/enterprise_hooks/__init__.py

from typing import Dict, Literal, Type, Union

from litellm_enterprise.proxy.hooks.managed_files import PROXY_LiteLLMManagedFiles
from litellm_enterprise.proxy.hooks.managed_vector_stores import (
    PROXY_LiteLLMManagedVectorStores,
)

from litellm.integrations.custom_logger import CustomLogger

ENTERPRISE_PROXY_HOOKS: Dict[str, Type[CustomLogger]] = {
    "managed_files": PROXY_LiteLLMManagedFiles,
    "managed_vector_stores": PROXY_LiteLLMManagedVectorStores,
}


def get_enterprise_proxy_hook(
    hook_name: Union[
        Literal[
            "managed_files",
            "managed_vector_stores",
            "max_parallel_requests",
        ],
        str,
    ],
):
    """
    Factory method to get a enterprise hook instance by name
    """
    if hook_name not in ENTERPRISE_PROXY_HOOKS:
        raise ValueError(
            f"Unknown hook: {hook_name}. Available hooks: {list(ENTERPRISE_PROXY_HOOKS.keys())}"
        )
    return ENTERPRISE_PROXY_HOOKS[hook_name]