From 82246d8e30c3f0d52049d117fdddedbc59d6331a Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 13 Mar 2024 12:16:58 -0700 Subject: [PATCH 1/3] (fix) using enterprise folder on litellm --- litellm/proxy/proxy_server.py | 32 ++++++++++++++++++-------------- 1 file changed, 18 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ebc9a3bfd..cbe0e3913 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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" @@ -1617,7 +1625,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, ) @@ -1627,7 +1635,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, ) @@ -1639,7 +1647,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, ) @@ -1649,7 +1657,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, ) @@ -1659,7 +1667,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, ) @@ -4103,7 +4111,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: @@ -4512,12 +4520,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) @@ -4975,7 +4979,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, ) @@ -5016,7 +5020,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, ) From ba5cc19edf62e6461b1b73d3e7e454d8bbb55bd9 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 13 Mar 2024 14:00:04 -0700 Subject: [PATCH 2/3] (test) using get_predict_spend_logs --- tests/test_spend_logs.py | 40 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/tests/test_spend_logs.py b/tests/test_spend_logs.py index 4d7ad175f..d95ad00bb 100644 --- a/tests/test_spend_logs.py +++ b/tests/test_spend_logs.py @@ -113,6 +113,46 @@ async def test_spend_logs(): await get_spend_logs(session=session, request_id=response["id"]) +async def get_predict_spend_logs(session): + url = f"http://0.0.0.0:4035/global/predict/spend/logs" + headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} + data = { + "data": [ + { + "date": "2024-03-09", + "spend": 200000, + "api_key": "f19bdeb945164278fc11c1020d8dfd70465bffd931ed3cb2e1efa6326225b8b7", + } + ] + } + + async with session.post(url, headers=headers, json=data) as response: + status = response.status + response_text = await response.text() + + print(response_text) + print() + + if status != 200: + raise Exception(f"Request did not return a 200 status code: {status}") + return await response.json() + + +@pytest.mark.asyncio +async def test_get_predicted_spend_logs(): + """ + - Create key + - Make call (makes sure it's in spend logs) + - Get request id from logs + """ + async with aiohttp.ClientSession() as session: + result = await get_predict_spend_logs(session=session) + print(result) + + assert "response" in result + assert len(result["response"]) > 0 + + @pytest.mark.skip(reason="High traffic load test, meant to be run locally") @pytest.mark.asyncio async def test_spend_logs_high_traffic(): From 992aba2a773b015c95668eda1ce979b93e8f31f3 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 13 Mar 2024 14:09:51 -0700 Subject: [PATCH 3/3] (fix) predict/spend/logs test --- tests/test_spend_logs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_spend_logs.py b/tests/test_spend_logs.py index d95ad00bb..c6866317d 100644 --- a/tests/test_spend_logs.py +++ b/tests/test_spend_logs.py @@ -114,7 +114,7 @@ async def test_spend_logs(): async def get_predict_spend_logs(session): - url = f"http://0.0.0.0:4035/global/predict/spend/logs" + url = f"http://0.0.0.0:4000/global/predict/spend/logs" headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} data = { "data": [