From ab15028850b6b123d36ecbc75dbfbbb389cc73ec Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 25 Oct 2024 13:31:02 +0400 Subject: [PATCH 1/7] use separate file for create_audit_log_for_update --- .../proxy/management_helpers/audit_logs.py | 43 +++++++++++++++++++ litellm/proxy/proxy_server.py | 1 + 2 files changed, 44 insertions(+) create mode 100644 litellm/proxy/management_helpers/audit_logs.py diff --git a/litellm/proxy/management_helpers/audit_logs.py b/litellm/proxy/management_helpers/audit_logs.py new file mode 100644 index 000000000..27da9911a --- /dev/null +++ b/litellm/proxy/management_helpers/audit_logs.py @@ -0,0 +1,43 @@ +""" +Functions to create audit logs for LiteLLM Proxy +""" + +import json + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLM_AuditLogs + + +async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs): + from litellm.proxy.proxy_server import premium_user, prisma_client + + if premium_user is not True: + return + + if litellm.store_audit_logs is not True: + return + if prisma_client is None: + raise Exception("prisma_client is None, no DB connected") + + verbose_proxy_logger.debug("creating audit log for %s", request_data) + + if isinstance(request_data.updated_values, dict): + request_data.updated_values = json.dumps(request_data.updated_values) + + if isinstance(request_data.before_value, dict): + request_data.before_value = json.dumps(request_data.before_value) + + _request_data = request_data.dict(exclude_none=True) + + try: + await prisma_client.db.litellm_auditlog.create( + data={ + **_request_data, # type: ignore + } + ) + except Exception as e: + # [Non-Blocking Exception. Do not allow blocking LLM API call] + verbose_proxy_logger.error(f"Failed Creating audit log {e}") + + return diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index bcdd4e86c..26ff04d87 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -194,6 +194,7 @@ from litellm.proxy.management_endpoints.team_callback_endpoints import ( ) from litellm.proxy.management_endpoints.team_endpoints import router as team_router from litellm.proxy.management_endpoints.ui_sso import router as ui_sso_router +from litellm.proxy.management_helpers.audit_logs import create_audit_log_for_update from litellm.proxy.openai_files_endpoints.files_endpoints import is_known_model from litellm.proxy.openai_files_endpoints.files_endpoints import ( router as openai_files_router, From b0bf182db9d8e6034f045b6a4a03abef8650bc41 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 25 Oct 2024 16:47:17 +0400 Subject: [PATCH 2/7] fix LitellmTableNames type --- litellm/proxy/_types.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 629e002b5..743bd66aa 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -104,7 +104,7 @@ class LitellmUserRoles(str, enum.Enum): return ui_labels.get(self.value, "") -class LitellmTableNames(enum.Enum): +class LitellmTableNames(str, enum.Enum): """ Enum for Table Names used by LiteLLM """ From 61e28ebb3eff34b84617fea96a43a8fe88178c93 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 25 Oct 2024 16:47:35 +0400 Subject: [PATCH 3/7] fix StandardLoggingMetadata with user_api_key_org_id --- litellm/litellm_core_utils/litellm_logging.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 5201bfe1e..206cb235e 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2515,6 +2515,7 @@ class StandardLoggingPayloadSetup: user_api_key_hash=None, user_api_key_alias=None, user_api_key_team_id=None, + user_api_key_org_id=None, user_api_key_user_id=None, user_api_key_team_alias=None, spend_logs_metadata=None, From eb24ce25ea6bdd93eb2766afbdf157824222b222 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 25 Oct 2024 16:48:25 +0400 Subject: [PATCH 4/7] fix create_audit_log_for_update --- .../proxy/management_helpers/audit_logs.py | 2 +- litellm/proxy/proxy_config.yaml | 2 +- litellm/proxy/proxy_server.py | 32 ------------------- 3 files changed, 2 insertions(+), 34 deletions(-) diff --git a/litellm/proxy/management_helpers/audit_logs.py b/litellm/proxy/management_helpers/audit_logs.py index 27da9911a..b023e9096 100644 --- a/litellm/proxy/management_helpers/audit_logs.py +++ b/litellm/proxy/management_helpers/audit_logs.py @@ -28,7 +28,7 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs): if isinstance(request_data.before_value, dict): request_data.before_value = json.dumps(request_data.before_value) - _request_data = request_data.dict(exclude_none=True) + _request_data = request_data.model_dump(exclude_none=True) try: await prisma_client.db.litellm_auditlog.create( diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 8f22d0d78..18cc262b4 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -13,7 +13,7 @@ general_settings: proxy_batch_write_at: 60 # Batch write spend updates every 60s litellm_settings: - success_callback: ["langfuse"] + store_audit_logs: true # https://docs.litellm.ai/docs/proxy/reliability#default-fallbacks default_fallbacks: ["gpt-4o-2024-08-06", "claude-3-5-sonnet-20240620"] diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 26ff04d87..7918543b7 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6434,38 +6434,6 @@ async def list_end_user( return returned_response -async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs): - if premium_user is not True: - return - - if litellm.store_audit_logs is not True: - return - if prisma_client is None: - raise Exception("prisma_client is None, no DB connected") - - verbose_proxy_logger.debug("creating audit log for %s", request_data) - - if isinstance(request_data.updated_values, dict): - request_data.updated_values = json.dumps(request_data.updated_values) - - if isinstance(request_data.before_value, dict): - request_data.before_value = json.dumps(request_data.before_value) - - _request_data = request_data.dict(exclude_none=True) - - try: - await prisma_client.db.litellm_auditlog.create( - data={ - **_request_data, # type: ignore - } - ) - except Exception as e: - # [Non-Blocking Exception. Do not allow blocking LLM API call] - verbose_proxy_logger.error(f"Failed Creating audit log {e}") - - return - - #### BUDGET TABLE MANAGEMENT #### From c27555677e39b0a2d22ed8cbce2986fa5e40dc09 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 25 Oct 2024 16:50:24 +0400 Subject: [PATCH 5/7] fix code quality --- litellm/proxy/proxy_server.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7918543b7..cadb6063c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6399,11 +6399,7 @@ async def list_end_user( --header 'Authorization: Bearer sk-1234' ``` """ - from litellm.proxy.proxy_server import ( - create_audit_log_for_update, - litellm_proxy_admin_name, - prisma_client, - ) + from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client if ( user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN From 8cf0191cf48fdbdf3987b7cc7d2eb0f700f8df65 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 25 Oct 2024 17:28:37 +0400 Subject: [PATCH 6/7] unit testing test_create_audit_log_in_db --- tests/local_testing/test_audit_logs_proxy.py | 151 +++++++++++++++++++ 1 file changed, 151 insertions(+) create mode 100644 tests/local_testing/test_audit_logs_proxy.py diff --git a/tests/local_testing/test_audit_logs_proxy.py b/tests/local_testing/test_audit_logs_proxy.py new file mode 100644 index 000000000..48187e9b2 --- /dev/null +++ b/tests/local_testing/test_audit_logs_proxy.py @@ -0,0 +1,151 @@ +import os +import sys +import traceback +import uuid +from datetime import datetime + +from dotenv import load_dotenv +from fastapi import Request +from fastapi.routing import APIRoute + + +import io +import os +import time + +# this file is to test litellm/proxy + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +import asyncio +import logging + +load_dotenv() + +import pytest +import uuid +import litellm +from litellm._logging import verbose_proxy_logger + +from litellm.proxy.proxy_server import ( + LitellmUserRoles, + audio_transcriptions, + chat_completion, + completion, + embeddings, + image_generation, + model_list, + moderations, + new_end_user, + user_api_key_auth, +) + +from litellm.proxy.utils import PrismaClient, ProxyLogging, hash_token, update_spend + +verbose_proxy_logger.setLevel(level=logging.DEBUG) + +from starlette.datastructures import URL + +from litellm.proxy.management_helpers.audit_logs import create_audit_log_for_update +from litellm.proxy._types import LiteLLM_AuditLogs, LitellmTableNames +from litellm.caching.caching import DualCache + +proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) +import json + + +@pytest.mark.asyncio +async def test_create_audit_log_for_update_premium_user(): + """ + Basic unit test for create_audit_log_for_update + + Test that the audit log is created when a premium user updates a team + """ + with patch("litellm.proxy.proxy_server.premium_user", True), patch( + "litellm.store_audit_logs", True + ), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + + mock_prisma.db.litellm_auditlog.create = AsyncMock() + + request_data = LiteLLM_AuditLogs( + id="test_id", + updated_at=datetime.now(), + changed_by="test_changed_by", + action="updated", + table_name=LitellmTableNames.TEAM_TABLE_NAME, + object_id="test_object_id", + updated_values=json.dumps({"key": "value"}), + before_value=json.dumps({"old_key": "old_value"}), + ) + + await create_audit_log_for_update(request_data) + + mock_prisma.db.litellm_auditlog.create.assert_called_once_with( + data={ + "id": "test_id", + "updated_at": request_data.updated_at, + "changed_by": request_data.changed_by, + "action": request_data.action, + "table_name": request_data.table_name, + "object_id": request_data.object_id, + "updated_values": request_data.updated_values, + "before_value": request_data.before_value, + } + ) + + +@pytest.fixture +def prisma_client(): + from litellm.proxy.proxy_cli import append_query_params + + ### add connection pool + pool timeout args + params = {"connection_limit": 100, "pool_timeout": 60} + database_url = os.getenv("DATABASE_URL") + modified_url = append_query_params(database_url, params) + os.environ["DATABASE_URL"] = modified_url + + # Assuming PrismaClient is a class that needs to be instantiated + prisma_client = PrismaClient( + database_url=os.environ["DATABASE_URL"], proxy_logging_obj=proxy_logging_obj + ) + + return prisma_client + + +@pytest.mark.asyncio() +async def test_create_audit_log_in_db(prisma_client): + print("prisma client=", prisma_client) + + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + setattr(litellm.proxy.proxy_server, "premium_user", True) + setattr(litellm, "store_audit_logs", True) + + await litellm.proxy.proxy_server.prisma_client.connect() + audit_log_id = f"audit_log_id_{uuid.uuid4()}" + + # create a audit log for /key/generate + request_data = LiteLLM_AuditLogs( + id=audit_log_id, + updated_at=datetime.now(), + changed_by="test_changed_by", + action="updated", + table_name=LitellmTableNames.TEAM_TABLE_NAME, + object_id="test_object_id", + updated_values=json.dumps({"key": "value"}), + before_value=json.dumps({"old_key": "old_value"}), + ) + + await create_audit_log_for_update(request_data) + + await asyncio.sleep(1) + + # now read the last log from the db + last_log = await prisma_client.db.litellm_auditlog.find_first( + where={"id": audit_log_id} + ) + + assert last_log.id == audit_log_id + + setattr(litellm, "store_audit_logs", False) From 6f1c06f7ae36bc386be4dbe3f1babd57fe4dae50 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 25 Oct 2024 18:21:17 +0400 Subject: [PATCH 7/7] fix test audit logs --- tests/local_testing/test_audit_logs_proxy.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/local_testing/test_audit_logs_proxy.py b/tests/local_testing/test_audit_logs_proxy.py index 48187e9b2..275d48670 100644 --- a/tests/local_testing/test_audit_logs_proxy.py +++ b/tests/local_testing/test_audit_logs_proxy.py @@ -50,6 +50,7 @@ from starlette.datastructures import URL from litellm.proxy.management_helpers.audit_logs import create_audit_log_for_update from litellm.proxy._types import LiteLLM_AuditLogs, LitellmTableNames from litellm.caching.caching import DualCache +from unittest.mock import patch, AsyncMock proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) import json