From 73acdf373608bd5fb766c492b4533d5a9c729703 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 17 Feb 2024 17:52:32 -0800 Subject: [PATCH] refactor: rename enterprise specific hooks and callbacks to be more precise --- .../example_logging_api.py | 0 .../generic_api_callback.py | 0 .../google_text_moderation.py | 0 .../{hooks => enterprise_hooks}/llama_guard.py | 0 litellm/proxy/proxy_server.py | 2 +- litellm/utils.py | 14 +++++--------- 6 files changed, 6 insertions(+), 10 deletions(-) rename enterprise/{callbacks => enterprise_callbacks}/example_logging_api.py (100%) rename enterprise/{callbacks => enterprise_callbacks}/generic_api_callback.py (100%) rename enterprise/{hooks => enterprise_hooks}/google_text_moderation.py (100%) rename enterprise/{hooks => enterprise_hooks}/llama_guard.py (100%) diff --git a/enterprise/callbacks/example_logging_api.py b/enterprise/enterprise_callbacks/example_logging_api.py similarity index 100% rename from enterprise/callbacks/example_logging_api.py rename to enterprise/enterprise_callbacks/example_logging_api.py diff --git a/enterprise/callbacks/generic_api_callback.py b/enterprise/enterprise_callbacks/generic_api_callback.py similarity index 100% rename from enterprise/callbacks/generic_api_callback.py rename to enterprise/enterprise_callbacks/generic_api_callback.py diff --git a/enterprise/hooks/google_text_moderation.py b/enterprise/enterprise_hooks/google_text_moderation.py similarity index 100% rename from enterprise/hooks/google_text_moderation.py rename to enterprise/enterprise_hooks/google_text_moderation.py diff --git a/enterprise/hooks/llama_guard.py b/enterprise/enterprise_hooks/llama_guard.py similarity index 100% rename from enterprise/hooks/llama_guard.py rename to enterprise/enterprise_hooks/llama_guard.py diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index bd30a621f..0f4f3d5ef 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1421,7 +1421,7 @@ class ProxyConfig: isinstance(callback, str) and callback == "llamaguard_moderations" ): - from litellm.proxy.enterprise.hooks.llama_guard import ( + from litellm.proxy.enterprise.enterprise_hooks.llama_guard import ( _ENTERPRISE_LlamaGuard, ) diff --git a/litellm/utils.py b/litellm/utils.py index d173b0092..faa464448 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -34,6 +34,7 @@ from dataclasses import ( try: # this works in python 3.8 import pkg_resources + filename = pkg_resources.resource_filename(__name__, "llms/tokenizers") # try: # filename = str( @@ -42,6 +43,7 @@ try: except: # this works in python 3.9+ from importlib import resources + filename = str( resources.files(litellm).joinpath("llms/tokenizers") # for python 3.10 ) # for python 3.10+ @@ -87,16 +89,10 @@ from .exceptions import ( UnprocessableEntityError, ) -# Import Enterprise features -project_path = abspath(join(dirname(__file__), "..", "..")) -# Add the "enterprise" directory to sys.path -verbose_logger.debug(f"current project_path: {project_path}") -enterprise_path = abspath(join(project_path, "enterprise")) -sys.path.append(enterprise_path) - -verbose_logger.debug(f"sys.path: {sys.path}") try: - from enterprise.callbacks.generic_api_callback import GenericAPILogger + from .proxy.enterprise.enterprise_callbacks.generic_api_callback import ( + GenericAPILogger, + ) except Exception as e: verbose_logger.debug(f"Exception import enterprise features {str(e)}")