Merge pull request #2501 from BerriAI/litellm_fix_using_enterprise_docker

(fix) using enterprise folder on litellm Docker
This commit is contained in:
Ishaan Jaff 2024-03-13 14:26:21 -07:00 committed by GitHub
commit 8a886c6e93
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 58 additions and 14 deletions

View file

@ -141,6 +141,14 @@ import json
import logging
from typing import Union
# import enterprise folder
try:
# when using litellm cli
import litellm.proxy.enterprise as enterprise
except:
# when using litellm docker image
import enterprise # type: ignore
ui_link = f"/ui/"
ui_message = (
f"👉 [```LiteLLM Admin Panel on /ui```]({ui_link}). Create, Edit Keys with SSO"
@ -1626,7 +1634,7 @@ class ProxyConfig:
isinstance(callback, str)
and callback == "llamaguard_moderations"
):
from litellm.proxy.enterprise.enterprise_hooks.llama_guard import (
from enterprise.enterprise_hooks.llama_guard import (
_ENTERPRISE_LlamaGuard,
)
@ -1636,7 +1644,7 @@ class ProxyConfig:
isinstance(callback, str)
and callback == "google_text_moderation"
):
from litellm.proxy.enterprise.enterprise_hooks.google_text_moderation import (
from enterprise.enterprise_hooks.google_text_moderation import (
_ENTERPRISE_GoogleTextModeration,
)
@ -1648,7 +1656,7 @@ class ProxyConfig:
isinstance(callback, str)
and callback == "llmguard_moderations"
):
from litellm.proxy.enterprise.enterprise_hooks.llm_guard import (
from enterprise.enterprise_hooks.llm_guard import (
_ENTERPRISE_LLMGuard,
)
@ -1658,7 +1666,7 @@ class ProxyConfig:
isinstance(callback, str)
and callback == "blocked_user_check"
):
from litellm.proxy.enterprise.enterprise_hooks.blocked_user_list import (
from enterprise.enterprise_hooks.blocked_user_list import (
_ENTERPRISE_BlockedUserList,
)
@ -1668,7 +1676,7 @@ class ProxyConfig:
isinstance(callback, str)
and callback == "banned_keywords"
):
from litellm.proxy.enterprise.enterprise_hooks.banned_keywords import (
from enterprise.enterprise_hooks.banned_keywords import (
_ENTERPRISE_BannedKeywords,
)
@ -4119,7 +4127,7 @@ async def view_spend_tags(
```
"""
from litellm.proxy.enterprise.utils import get_spend_by_tags
from enterprise.utils import get_spend_by_tags
global prisma_client
try:
@ -4528,12 +4536,8 @@ async def global_spend_models(
dependencies=[Depends(user_api_key_auth)],
)
async def global_predict_spend_logs(request: Request):
try:
# when using litellm package
from litellm.proxy.enterprise.utils import _forecast_daily_cost
except:
# when using litellm docker image
from enterprise.utils import _forecast_daily_cost
from enterprise.utils import _forecast_daily_cost
data = await request.json()
data = data.get("data")
return _forecast_daily_cost(data)
@ -4991,7 +4995,7 @@ async def block_user(data: BlockUsers):
}'
```
"""
from litellm.proxy.enterprise.enterprise_hooks.blocked_user_list import (
from enterprise.enterprise_hooks.blocked_user_list import (
_ENTERPRISE_BlockedUserList,
)
@ -5032,7 +5036,7 @@ async def unblock_user(data: BlockUsers):
}'
```
"""
from litellm.proxy.enterprise.enterprise_hooks.blocked_user_list import (
from enterprise.enterprise_hooks.blocked_user_list import (
_ENTERPRISE_BlockedUserList,
)