diff --git a/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py b/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py new file mode 100644 index 0000000000..43369630a4 --- /dev/null +++ b/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py @@ -0,0 +1,169 @@ +""" +AUDIT LOGGING + +All /audit logging endpoints. Attempting to write these as CRUD endpoints. + +GET - /audit/{id} - Get audit log by id +GET - /audit - Get all audit logs +""" + +from datetime import datetime +from typing import Any, Dict, List, Optional + +#### AUDIT LOGGING #### +from fastapi import APIRouter, Depends, HTTPException, Query +from litellm_enterprise.types.proxy.audit_logging_endpoints import ( + AuditLogResponse, + PaginatedAuditLogResponse, +) + +from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + +router = APIRouter() + + +@router.get( + "/audit", + tags=["Audit Logging"], + dependencies=[Depends(user_api_key_auth)], + response_model=PaginatedAuditLogResponse, +) +async def get_audit_logs( + page: int = Query(1, ge=1), + page_size: int = Query(10, ge=1, le=100), + # Filter parameters + changed_by: Optional[str] = Query( + None, description="Filter by user or system that performed the action" + ), + changed_by_api_key: Optional[str] = Query( + None, description="Filter by API key hash that performed the action" + ), + action: Optional[str] = Query( + None, description="Filter by action type (create, update, delete)" + ), + table_name: Optional[str] = Query( + None, description="Filter by table name that was modified" + ), + object_id: Optional[str] = Query( + None, description="Filter by ID of the object that was modified" + ), + start_date: Optional[str] = Query(None, description="Filter logs after this date"), + end_date: Optional[str] = Query(None, description="Filter logs before this date"), + # Sorting parameters + sort_by: Optional[str] = Query( + None, + description="Column to sort by (e.g. 'updated_at', 'action', 'table_name')", + ), + sort_order: str = Query("desc", description="Sort order ('asc' or 'desc')"), +): + """ + Get all audit logs with filtering and pagination. + + Returns a paginated response of audit logs matching the specified filters. + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"message": CommonProxyErrors.db_not_connected_error.value}, + ) + + # Build filter conditions + where_conditions: Dict[str, Any] = {} + if changed_by: + where_conditions["changed_by"] = changed_by + if changed_by_api_key: + where_conditions["changed_by_api_key"] = changed_by_api_key + if action: + where_conditions["action"] = action + if table_name: + where_conditions["table_name"] = table_name + if object_id: + where_conditions["object_id"] = object_id + if start_date or end_date: + date_filter = {} + if start_date: + date_filter["gte"] = start_date + if end_date: + date_filter["lte"] = end_date + where_conditions["updated_at"] = date_filter + + # Build sort conditions + order_by = {} + if sort_by and isinstance(sort_by, str): + order_by[sort_by] = sort_order + elif sort_order and isinstance(sort_order, str): + order_by["updated_at"] = sort_order # Default sort by updated_at + + # Get paginated results + audit_logs = await prisma_client.db.litellm_auditlog.find_many( + where=where_conditions, + order=order_by, + skip=(page - 1) * page_size, + take=page_size, + ) + + # Get total count for pagination + total_count = await prisma_client.db.litellm_auditlog.count(where=where_conditions) + total_pages = -(-total_count // page_size) # Ceiling division + + # Return paginated response + return PaginatedAuditLogResponse( + audit_logs=[ + AuditLogResponse(**audit_log.model_dump()) for audit_log in audit_logs + ] + if audit_logs + else [], + total=total_count, + page=page, + page_size=page_size, + total_pages=total_pages, + ) + + +@router.get( + "/audit/{id}", + tags=["Audit Logging"], + dependencies=[Depends(user_api_key_auth)], + response_model=AuditLogResponse, + responses={ + 404: {"description": "Audit log not found"}, + 500: {"description": "Database connection error"}, + }, +) +async def get_audit_log_by_id( + id: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth) +): + """ + Get detailed information about a specific audit log entry by its ID. + + Args: + id (str): The unique identifier of the audit log entry + + Returns: + AuditLogResponse: Detailed information about the audit log entry + + Raises: + HTTPException: If the audit log is not found or if there's a database connection error + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"message": CommonProxyErrors.db_not_connected_error.value}, + ) + + # Get the audit log by ID + audit_log = await prisma_client.db.litellm_auditlog.find_unique(where={"id": id}) + + if audit_log is None: + raise HTTPException( + status_code=404, detail={"message": f"Audit log with ID {id} not found"} + ) + + # Convert to response model + print(audit_log) + return AuditLogResponse(**audit_log.model_dump()) diff --git a/enterprise/litellm_enterprise/proxy/enterprise_routes.py b/enterprise/litellm_enterprise/proxy/enterprise_routes.py index 2420ad2e05..1e4ed58061 100644 --- a/enterprise/litellm_enterprise/proxy/enterprise_routes.py +++ b/enterprise/litellm_enterprise/proxy/enterprise_routes.py @@ -4,6 +4,7 @@ from litellm_enterprise.enterprise_callbacks.send_emails.endpoints import ( router as email_events_router, ) +from .audit_logging_endpoints import router as audit_logging_router from .guardrails.endpoints import router as guardrails_router from .utils import _should_block_robots from .vector_stores.endpoints import router as vector_stores_router @@ -12,6 +13,7 @@ router = APIRouter() router.include_router(vector_stores_router) router.include_router(guardrails_router) router.include_router(email_events_router) +router.include_router(audit_logging_router) @router.get("/robots.txt") diff --git a/enterprise/litellm_enterprise/proxy/readme.md b/enterprise/litellm_enterprise/proxy/readme.md index 9ec611fa83..60b07cf49a 100644 --- a/enterprise/litellm_enterprise/proxy/readme.md +++ b/enterprise/litellm_enterprise/proxy/readme.md @@ -4,3 +4,8 @@ This directory contains enterprise features used on the LiteLLM proxy. +## Format + +Create a file for every group of endpoints (e.g. `key_management_endpoints.py`, `user_management_endpoints.py`, etc.) + +If there is a broader semantic group of endpoints, create a folder for that group (e.g. `management_endpoints`, `auth_endpoints`, etc.) diff --git a/enterprise/litellm_enterprise/types/proxy/audit_logging_endpoints.py b/enterprise/litellm_enterprise/types/proxy/audit_logging_endpoints.py new file mode 100644 index 0000000000..4615bde2b1 --- /dev/null +++ b/enterprise/litellm_enterprise/types/proxy/audit_logging_endpoints.py @@ -0,0 +1,30 @@ +from datetime import datetime +from typing import Any, Dict, List, Optional + +from pydantic import BaseModel, Field + + +class AuditLogResponse(BaseModel): + """Response model for a single audit log entry""" + + id: str + updated_at: datetime + changed_by: str + changed_by_api_key: str + action: str + table_name: str + object_id: str + before_value: Optional[Dict[str, Any]] = None + updated_values: Optional[Dict[str, Any]] = None + + +class PaginatedAuditLogResponse(BaseModel): + """Response model for paginated audit logs""" + + audit_logs: List[AuditLogResponse] + total: int = Field( + ..., description="Total number of audit logs matching the filters" + ) + page: int = Field(..., description="Current page number") + page_size: int = Field(..., description="Number of items per page") + total_pages: int = Field(..., description="Total number of pages") diff --git a/litellm/integrations/deepeval/__init__.py b/litellm/integrations/deepeval/__init__.py new file mode 100644 index 0000000000..e074075b88 --- /dev/null +++ b/litellm/integrations/deepeval/__init__.py @@ -0,0 +1,3 @@ +from .deepeval import DeepEvalLogger + +__all__ = ["DeepEvalLogger"] diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index a255f32f5e..90a44b1f27 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -78,6 +78,9 @@ model_list: prompt_label: "latest" api_key: os.environ/OPENAI_API_KEY +litellm_settings: + store_audit_logs: true + general_settings: store_model_in_db: true store_prompts_in_spend_logs: true diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c1feb3dc33..06ac41d6a1 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1790,7 +1790,7 @@ class LiteLLM_ErrorLogs(LiteLLMPydanticObjectBase): endTime: Union[str, datetime, None] -AUDIT_ACTIONS = Literal["created", "updated", "deleted", "blocked"] +AUDIT_ACTIONS = Literal["created", "updated", "deleted", "blocked", "rotated"] class LiteLLM_AuditLogs(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py index 082b506ddd..dad1cf3910 100644 --- a/litellm/proxy/hooks/key_management_event_hooks.py +++ b/litellm/proxy/hooks/key_management_event_hooks.py @@ -4,8 +4,6 @@ import uuid from datetime import datetime, timezone from typing import Any, List, Optional -from fastapi import status - import litellm from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( @@ -16,8 +14,6 @@ from litellm.proxy._types import ( Litellm_EntityType, LiteLLM_VerificationToken, LitellmTableNames, - ProxyErrorTypes, - ProxyException, RegenerateKeyRequest, UpdateKeyRequest, UserAPIKeyAuth, @@ -126,11 +122,16 @@ class KeyManagementEventHooks: @staticmethod async def async_key_rotated_hook( data: Optional[RegenerateKeyRequest], - existing_key_row: Any, + existing_key_row: LiteLLM_VerificationToken, response: GenerateKeyResponse, user_api_key_dict: UserAPIKeyAuth, litellm_changed_by: Optional[str] = None, ): + from litellm.proxy.management_helpers.audit_logs import ( + create_audit_log_for_update, + ) + from litellm.proxy.proxy_server import litellm_proxy_admin_name + # store the generated key in the secret manager if data is not None and response.token_id is not None: initial_secret_name = ( @@ -142,6 +143,28 @@ class KeyManagementEventHooks: new_secret_value=response.key, ) + # store the audit log + if litellm.store_audit_logs is True and existing_key_row.token is not None: + asyncio.create_task( + create_audit_log_for_update( + request_data=LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=litellm_changed_by + or user_api_key_dict.user_id + or litellm_proxy_admin_name, + changed_by_api_key=user_api_key_dict.token, + table_name=LitellmTableNames.KEY_TABLE_NAME, + object_id=existing_key_row.token, + action="rotated", + updated_values=response.model_dump_json(exclude_none=True), + before_value=existing_key_row.model_dump_json( + exclude_none=True + ), + ) + ) + ) + @staticmethod async def async_key_deleted_hook( data: KeyRequest, @@ -159,27 +182,16 @@ class KeyManagementEventHooks: from litellm.proxy.management_helpers.audit_logs import ( create_audit_log_for_update, ) - from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client + from litellm.proxy.proxy_server import litellm_proxy_admin_name # Enterprise Feature - Audit Logging. Enable with litellm.store_audit_logs = True # we do this after the first for loop, since first for loop is for validation. we only want this inserted after validation passes if litellm.store_audit_logs is True and data.keys is not None: - # make an audit log for each team deleted - for key in data.keys: - key_row = await prisma_client.get_data( # type: ignore - token=key, table_name="key", query_type="find_unique" - ) - - if key_row is None: - raise ProxyException( - message=f"Key {key} not found", - type=ProxyErrorTypes.bad_request_error, - param="key", - code=status.HTTP_404_NOT_FOUND, - ) - - key_row = key_row.json(exclude_none=True) - _key_row = json.dumps(key_row, default=str) + # make an audit log for each key deleted + for key in keys_being_deleted: + if key.token is None: + continue + _key_row = key.model_dump_json(exclude_none=True) asyncio.create_task( create_audit_log_for_update( @@ -189,9 +201,9 @@ class KeyManagementEventHooks: changed_by=litellm_changed_by or user_api_key_dict.user_id or litellm_proxy_admin_name, - changed_by_api_key=user_api_key_dict.api_key, + changed_by_api_key=user_api_key_dict.token, table_name=LitellmTableNames.KEY_TABLE_NAME, - object_id=key, + object_id=key.token, action="deleted", updated_values="{}", before_value=_key_row, diff --git a/tests/enterprise/litellm_enterprise/proxy/test_audit_logging_endpoints.py b/tests/enterprise/litellm_enterprise/proxy/test_audit_logging_endpoints.py new file mode 100644 index 0000000000..ffc163dea1 --- /dev/null +++ b/tests/enterprise/litellm_enterprise/proxy/test_audit_logging_endpoints.py @@ -0,0 +1,128 @@ +from datetime import datetime, timedelta +from unittest.mock import AsyncMock, patch + +import pytest +from fastapi.testclient import TestClient +from litellm_enterprise.types.proxy.audit_logging_endpoints import AuditLogResponse + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.proxy_server import app + +client = TestClient(app) + +# Mock data for testing +MOCK_AUDIT_LOG = { + "id": "test-audit-log-1", + "updated_at": datetime.utcnow(), + "changed_by": "test-user", + "changed_by_api_key": "test-api-key-hash", + "action": "create", + "table_name": "test_table", + "object_id": "test-object-1", + "before_value": None, + "updated_values": {"name": "test", "value": 123}, +} + + +@pytest.fixture +def mock_prisma_client(): + with patch("litellm.proxy.proxy_server.prisma_client") as mock: + mock.db.litellm_auditlog.find_many = AsyncMock() + mock.db.litellm_auditlog.find_unique = AsyncMock() + mock.db.litellm_auditlog.count = AsyncMock() + yield mock + + +@pytest.mark.asyncio +async def test_get_audit_logs(mock_prisma_client): + """Test successful retrieval of audit logs with pagination""" + # Mock the database responses + mock_prisma_client.db.litellm_auditlog.find_many.return_value = [ + AuditLogResponse(**MOCK_AUDIT_LOG) + ] + mock_prisma_client.db.litellm_auditlog.count.return_value = 1 + + # Mock the auth dependency + with patch("litellm.proxy.auth.user_api_key_auth.user_api_key_auth") as mock_auth: + mock_auth.return_value = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id=None, + organization_id=None, + user_role="proxy_admin", + ) + + # Make the request + response = client.get("/audit?page=1&page_size=10") + + # Assert response + assert response.status_code == 200 + data = response.json() + assert "audit_logs" in data + assert len(data["audit_logs"]) == 1 + assert data["total"] == 1 + assert data["page"] == 1 + assert data["page_size"] == 10 + assert data["total_pages"] == 1 + + # Verify the audit log data + audit_log = data["audit_logs"][0] + assert audit_log["id"] == MOCK_AUDIT_LOG["id"] + assert audit_log["action"] == MOCK_AUDIT_LOG["action"] + assert audit_log["table_name"] == MOCK_AUDIT_LOG["table_name"] + + +@pytest.mark.asyncio +async def test_get_audit_log_by_id(mock_prisma_client): + """Test successful retrieval of a specific audit log by ID""" + # Mock the database response + mock_prisma_client.db.litellm_auditlog.find_unique.return_value = AuditLogResponse( + **MOCK_AUDIT_LOG + ) + + # Mock the auth dependency + with patch("litellm.proxy.auth.user_api_key_auth.user_api_key_auth") as mock_auth: + mock_auth.return_value = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id=None, + organization_id=None, + user_role="proxy_admin", + ) + + # Make the request + response = client.get(f"/audit/{MOCK_AUDIT_LOG['id']}") + + # Assert response + assert response.status_code == 200 + data = response.json() + assert data["id"] == MOCK_AUDIT_LOG["id"] + assert data["action"] == MOCK_AUDIT_LOG["action"] + assert data["table_name"] == MOCK_AUDIT_LOG["table_name"] + assert data["object_id"] == MOCK_AUDIT_LOG["object_id"] + + +@pytest.mark.asyncio +async def test_get_audit_log_by_id_not_found(mock_prisma_client): + """Test error handling when audit log is not found""" + # Mock the database response to return None + mock_prisma_client.db.litellm_auditlog.find_unique.return_value = None + + # Mock the auth dependency + with patch("litellm.proxy.auth.user_api_key_auth.user_api_key_auth") as mock_auth: + mock_auth.return_value = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id=None, + organization_id=None, + user_role="proxy_admin", + ) + + # Make the request + response = client.get("/audit/non-existent-id") + + # Assert response + assert response.status_code == 404 + data = response.json() + assert "message" in data["detail"] + assert "not found" in data["detail"]["message"].lower() diff --git a/tests/logging_callback_tests/test_unit_tests_init_callbacks.py b/tests/logging_callback_tests/test_unit_tests_init_callbacks.py index f6535c8ea4..e4915861b7 100644 --- a/tests/logging_callback_tests/test_unit_tests_init_callbacks.py +++ b/tests/logging_callback_tests/test_unit_tests_init_callbacks.py @@ -18,6 +18,7 @@ from litellm._logging import verbose_logger from prometheus_client import REGISTRY, CollectorRegistry from litellm.integrations.lago import LagoLogger +from litellm.integrations.deepeval import DeepEvalLogger from litellm.integrations.openmeter import OpenMeterLogger from litellm.integrations.braintrust_logging import BraintrustLogger from litellm.integrations.galileo import GalileoObserve