Merge pull request #4831 from BerriAI/litellm_api_control_team_logging

[Feat] - API Endpoints to control logging callbacks per Team

docs/my-website/docs/proxy/logging.md

@@ -216,6 +216,9 @@ print(response)
 
 ### Team based Logging to Langfuse
 
+[👉 Tutorial - Allow each team to use their own Langfuse Project / custom callbacks](team_logging)
+<!-- 
+
 **Example:**
 
 This config would send langfuse logs to 2 different langfuse projects, based on the team id 
@@ -242,7 +245,7 @@ curl -X POST 'http://0.0.0.0:4000/key/generate' \
 -d '{"team_id": "ishaans-secret-project"}'
 ```
 
-All requests made with these keys will log data to their team-specific logging.
+All requests made with these keys will log data to their team-specific logging. -->
 
 ### Redacting Messages, Response Content from Langfuse Logging 
 

docs/my-website/docs/proxy/team_based_routing.md

@@ -71,7 +71,13 @@ curl --location 'http://0.0.0.0:4000/v1/chat/completions' \
 }'
 ```
 
+## Team Based Logging
 
+[👉 Tutorial - Allow each team to use their own Langfuse Project / custom callbacks](team_logging.md)
+
+
+
+<!-- 
 ## Logging / Caching
 
 Turn on/off logging and caching for a specific team id. 
@@ -102,4 +108,4 @@ curl -X POST 'http://0.0.0.0:4000/key/generate' \
 -d '{"team_id": "ishaans-secret-project"}'
 ```
 
-All requests made with these keys will log data to their team-specific logging.
+All requests made with these keys will log data to their team-specific logging. -->

docs/my-website/docs/proxy/team_logging.md

@@ -0,0 +1,81 @@
+import Image from '@theme/IdealImage';
+import Tabs from '@theme/Tabs';
+import TabItem from '@theme/TabItem';
+
+# 👥📊 Team Based Logging
+
+Allow each team to use their own Langfuse Project / custom callbacks
+
+**This allows you to do the following**
+```
+Team 1 -> Logs to Langfuse Project 1 
+Team 2 -> Logs to Langfuse Project 2
+Team 3 -> Logs to Langsmith
+```
+
+## Quick Start
+
+## 1. Set callback for team 
+
+```shell
+curl -X POST 'http:/localhost:4000/team/dbe2f686-a686-4896-864a-4c3924458709/callback' \
+-H 'Content-Type: application/json' \
+-H 'Authorization: Bearer sk-1234' \
+-d '{
+  "callback_name": "langfuse",
+  "callback_type": "success",
+  "callback_vars": {
+    "langfuse_public_key": "pk", 
+    "langfuse_secret_key": "sk_", 
+    "langfuse_host": "https://cloud.langfuse.com"
+    }
+  
+}'
+```
+
+#### Supported Values
+
+| Field | Supported Values | Notes |
+|-------|------------------|-------|
+| `callback_name` | `"langfuse"` | Currently only supports "langfuse" |
+| `callback_type` | `"success"`, `"failure"`, `"success_and_failure"` | |
+| `callback_vars` | | dict of callback settings |
+|     `langfuse_public_key` | string | Required |
+|     `langfuse_secret_key` | string | Required |
+|     `langfuse_host` | string | Optional (defaults to https://cloud.langfuse.com) |
+
+## 2. Create key for team
+
+All keys created for team `dbe2f686-a686-4896-864a-4c3924458709` will log to langfuse project specified on [Step 1. Set callback for team](#1-set-callback-for-team)
+
+
+```shell
+curl --location 'http://0.0.0.0:4000/key/generate' \
+    --header 'Authorization: Bearer sk-1234' \
+    --header 'Content-Type: application/json' \
+    --data '{
+        "team_id": "dbe2f686-a686-4896-864a-4c3924458709"
+}'
+```
+
+
+## 3. Make `/chat/completion` request for team
+
+```shell
+curl -i http://localhost:4000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -H "Authorization: Bearer sk-KbUuE0WNptC0jXapyMmLBA" \
+  -d '{
+    "model": "gpt-4",
+    "messages": [
+      {"role": "user", "content": "Hello, Claude gm!"}
+    ]
+}'
+```
+
+Expect this to be logged on the langfuse project specified on [Step 1. Set callback for team](#1-set-callback-for-team)
+
+## 
+
+
+

docs/my-website/sidebars.js

@@ -44,19 +44,20 @@ const sidebars = {
         "proxy/cost_tracking",
         "proxy/self_serve",
         "proxy/virtual_keys",
-        "proxy/tag_routing",
-        "proxy/users",
-        "proxy/team_budgets",
-        "proxy/customers",
-        "proxy/billing",
-        "proxy/guardrails",
-        "proxy/token_auth",
-        "proxy/alerting",
         {
           type: "category",
           label: "🪢 Logging",
           items: ["proxy/logging", "proxy/streaming_logging"],
         },
+        "proxy/team_logging",
+        "proxy/guardrails",
+        "proxy/tag_routing",
+        "proxy/users",
+        "proxy/team_budgets",
+        "proxy/customers",
+        "proxy/billing",
+        "proxy/token_auth",
+        "proxy/alerting",
         "proxy/ui",
         "proxy/prometheus",
         "proxy/pass_through",

litellm/proxy/_types.py

@@ -884,6 +884,26 @@ class BlockTeamRequest(LiteLLMBase):
     team_id: str  # required
 
 
+class AddTeamCallback(LiteLLMBase):
+    callback_name: str
+    callback_type: Literal["success", "failure", "success_and_failure"]
+    # for now - only supported for langfuse
+    callback_vars: Dict[
+        Literal["langfuse_public_key", "langfuse_secret_key", "langfuse_host"], str
+    ]
+
+
+class TeamCallbackMetadata(LiteLLMBase):
+    success_callback: Optional[List[str]] = []
+    failure_callback: Optional[List[str]] = []
+    # for now - only supported for langfuse
+    callback_vars: Optional[
+        Dict[
+            Literal["langfuse_public_key", "langfuse_secret_key", "langfuse_host"], str
+        ]
+    ] = {}
+
+
 class LiteLLM_TeamTable(TeamBase):
     spend: Optional[float] = None
     max_parallel_requests: Optional[int] = None
@@ -1236,6 +1256,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken):
     soft_budget: Optional[float] = None
     team_model_aliases: Optional[Dict] = None
     team_member_spend: Optional[float] = None
+    team_metadata: Optional[Dict] = None
 
     # End User Params
     end_user_id: Optional[str] = None
@@ -1681,3 +1702,5 @@ class ProxyErrorTypes(str, enum.Enum):
     budget_exceeded = "budget_exceeded"
     expired_key = "expired_key"
     auth_error = "auth_error"
+    internal_server_error = "internal_server_error"
+    bad_request_error = "bad_request_error"

litellm/proxy/litellm_pre_call_utils.py

@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Any, Dict, Optional
 from fastapi import Request
 
 from litellm._logging import verbose_logger, verbose_proxy_logger
-from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
+from litellm.proxy._types import CommonProxyErrors, TeamCallbackMetadata, UserAPIKeyAuth
 from litellm.types.utils import SupportedCacheControls
 
 if TYPE_CHECKING:
@@ -207,6 +207,29 @@ async def add_litellm_data_to_request(
                 **data,
             }  # add the team-specific configs to the completion call
 
+    # Team Callbacks controls
+    if user_api_key_dict.team_metadata is not None:
+        team_metadata = user_api_key_dict.team_metadata
+        if "callback_settings" in team_metadata:
+            callback_settings = team_metadata.get("callback_settings", None) or {}
+            callback_settings_obj = TeamCallbackMetadata(**callback_settings)
+            """
+            callback_settings = {
+              {
+                'callback_vars': {'langfuse_public_key': 'pk', 'langfuse_secret_key': 'sk_'}, 
+                'failure_callback': [], 
+                'success_callback': ['langfuse', 'langfuse']
+            }
+            }
+            """
+            data["success_callback"] = callback_settings_obj.success_callback
+            data["failure_callback"] = callback_settings_obj.failure_callback
+
+            if callback_settings_obj.callback_vars is not None:
+                # unpack callback_vars in data
+                for k, v in callback_settings_obj.callback_vars.items():
+                    data[k] = v
+
     return data
 
 

litellm/proxy/management_endpoints/team_callback_endpoints.py

@@ -0,0 +1,279 @@
+"""
+Endpoints to control callbacks per team
+
+Use this when each team should control its own callbacks
+"""
+
+import asyncio
+import copy
+import json
+import traceback
+import uuid
+from datetime import datetime, timedelta, timezone
+from typing import List, Optional
+
+import fastapi
+from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
+
+import litellm
+from litellm._logging import verbose_proxy_logger
+from litellm.proxy._types import (
+    AddTeamCallback,
+    LiteLLM_TeamTable,
+    ProxyErrorTypes,
+    ProxyException,
+    TeamCallbackMetadata,
+    UserAPIKeyAuth,
+)
+from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+from litellm.proxy.management_helpers.utils import (
+    add_new_member,
+    management_endpoint_wrapper,
+)
+
+router = APIRouter()
+
+
+@router.post(
+    "/team/{team_id:path}/callback",
+    tags=["team management"],
+    dependencies=[Depends(user_api_key_auth)],
+)
+@management_endpoint_wrapper
+async def add_team_callbacks(
+    data: AddTeamCallback,
+    http_request: Request,
+    team_id: str,
+    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
+    litellm_changed_by: Optional[str] = Header(
+        None,
+        description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability",
+    ),
+):
+    """
+    Add a success/failure callback to a team
+
+    Use this if if you want different teams to have different success/failure callbacks
+
+    Example curl:
+    ```
+    curl -X POST 'http:/localhost:4000/team/dbe2f686-a686-4896-864a-4c3924458709/callback' \
+        -H 'Content-Type: application/json' \
+        -H 'Authorization: Bearer sk-1234' \
+        -d '{
+        "callback_name": "langfuse",
+        "callback_type": "success",
+        "callback_vars": {"langfuse_public_key": "pk-lf-xxxx1", "langfuse_secret_key": "sk-xxxxx"}
+        
+    }'
+    ```
+
+    This means for the team where team_id = dbe2f686-a686-4896-864a-4c3924458709, all LLM calls will be logged to langfuse using the public key pk-lf-xxxx1 and the secret key sk-xxxxx
+
+    """
+    try:
+        from litellm.proxy.proxy_server import (
+            _duration_in_seconds,
+            create_audit_log_for_update,
+            litellm_proxy_admin_name,
+            prisma_client,
+        )
+
+        if prisma_client is None:
+            raise HTTPException(status_code=500, detail={"error": "No db connected"})
+
+        # Check if team_id exists already
+        _existing_team = await prisma_client.get_data(
+            team_id=team_id, table_name="team", query_type="find_unique"
+        )
+        if _existing_team is None:
+            raise HTTPException(
+                status_code=400,
+                detail={
+                    "error": f"Team id = {team_id} does not exist. Please use a different team id."
+                },
+            )
+
+        # store team callback settings in metadata
+        team_metadata = _existing_team.metadata
+        team_callback_settings = team_metadata.get("callback_settings", {})
+        # expect callback settings to be
+        team_callback_settings_obj = TeamCallbackMetadata(**team_callback_settings)
+        if data.callback_type == "success":
+            if team_callback_settings_obj.success_callback is None:
+                team_callback_settings_obj.success_callback = []
+
+            if data.callback_name in team_callback_settings_obj.success_callback:
+                raise ProxyException(
+                    message=f"callback_name = {data.callback_name} already exists in failure_callback, for team_id = {team_id}. \n Existing failure_callback = {team_callback_settings_obj.success_callback}",
+                    code=status.HTTP_400_BAD_REQUEST,
+                    type=ProxyErrorTypes.bad_request_error,
+                    param="callback_name",
+                )
+
+            team_callback_settings_obj.success_callback.append(data.callback_name)
+        elif data.callback_type == "failure":
+            if team_callback_settings_obj.failure_callback is None:
+                team_callback_settings_obj.failure_callback = []
+
+            if data.callback_name in team_callback_settings_obj.failure_callback:
+                raise ProxyException(
+                    message=f"callback_name = {data.callback_name} already exists in failure_callback, for team_id = {team_id}. \n Existing failure_callback = {team_callback_settings_obj.failure_callback}",
+                    code=status.HTTP_400_BAD_REQUEST,
+                    type=ProxyErrorTypes.bad_request_error,
+                    param="callback_name",
+                )
+            team_callback_settings_obj.failure_callback.append(data.callback_name)
+        elif data.callback_type == "success_and_failure":
+            if team_callback_settings_obj.success_callback is None:
+                team_callback_settings_obj.success_callback = []
+            if team_callback_settings_obj.failure_callback is None:
+                team_callback_settings_obj.failure_callback = []
+            if data.callback_name in team_callback_settings_obj.success_callback:
+                raise ProxyException(
+                    message=f"callback_name = {data.callback_name} already exists in success_callback, for team_id = {team_id}. \n Existing success_callback = {team_callback_settings_obj.success_callback}",
+                    code=status.HTTP_400_BAD_REQUEST,
+                    type=ProxyErrorTypes.bad_request_error,
+                    param="callback_name",
+                )
+
+            if data.callback_name in team_callback_settings_obj.failure_callback:
+                raise ProxyException(
+                    message=f"callback_name = {data.callback_name} already exists in failure_callback, for team_id = {team_id}. \n Existing failure_callback = {team_callback_settings_obj.failure_callback}",
+                    code=status.HTTP_400_BAD_REQUEST,
+                    type=ProxyErrorTypes.bad_request_error,
+                    param="callback_name",
+                )
+
+            team_callback_settings_obj.success_callback.append(data.callback_name)
+            team_callback_settings_obj.failure_callback.append(data.callback_name)
+        for var, value in data.callback_vars.items():
+            if team_callback_settings_obj.callback_vars is None:
+                team_callback_settings_obj.callback_vars = {}
+            team_callback_settings_obj.callback_vars[var] = value
+
+        team_callback_settings_obj_dict = team_callback_settings_obj.model_dump()
+
+        team_metadata["callback_settings"] = team_callback_settings_obj_dict
+        team_metadata_json = json.dumps(team_metadata)  # update team_metadata
+
+        new_team_row = await prisma_client.db.litellm_teamtable.update(
+            where={"team_id": team_id}, data={"metadata": team_metadata_json}  # type: ignore
+        )
+
+        return {
+            "status": "success",
+            "data": new_team_row,
+        }
+
+    except Exception as e:
+        verbose_proxy_logger.error(
+            "litellm.proxy.proxy_server.add_team_callbacks(): Exception occured - {}".format(
+                str(e)
+            )
+        )
+        verbose_proxy_logger.debug(traceback.format_exc())
+        if isinstance(e, HTTPException):
+            raise ProxyException(
+                message=getattr(e, "detail", f"Internal Server Error({str(e)})"),
+                type=ProxyErrorTypes.internal_server_error.value,
+                param=getattr(e, "param", "None"),
+                code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR),
+            )
+        elif isinstance(e, ProxyException):
+            raise e
+        raise ProxyException(
+            message="Internal Server Error, " + str(e),
+            type=ProxyErrorTypes.internal_server_error.value,
+            param=getattr(e, "param", "None"),
+            code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+        )
+
+
+@router.get(
+    "/team/{team_id:path}/callback",
+    tags=["team management"],
+    dependencies=[Depends(user_api_key_auth)],
+)
+@management_endpoint_wrapper
+async def get_team_callbacks(
+    http_request: Request,
+    team_id: str,
+    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
+):
+    """
+    Get the success/failure callbacks and variables for a team
+
+    Example curl:
+    ```
+    curl -X GET 'http://localhost:4000/team/dbe2f686-a686-4896-864a-4c3924458709/callback' \
+        -H 'Authorization: Bearer sk-1234'
+    ```
+
+    This will return the callback settings for the team with id dbe2f686-a686-4896-864a-4c3924458709
+
+    Returns {
+            "status": "success",
+            "data": {
+                "team_id": team_id,
+                "success_callbacks": team_callback_settings_obj.success_callback,
+                "failure_callbacks": team_callback_settings_obj.failure_callback,
+                "callback_vars": team_callback_settings_obj.callback_vars,
+            },
+        }
+    """
+    try:
+        from litellm.proxy.proxy_server import prisma_client
+
+        if prisma_client is None:
+            raise HTTPException(status_code=500, detail={"error": "No db connected"})
+
+        # Check if team_id exists
+        _existing_team = await prisma_client.get_data(
+            team_id=team_id, table_name="team", query_type="find_unique"
+        )
+        if _existing_team is None:
+            raise HTTPException(
+                status_code=404,
+                detail={"error": f"Team id = {team_id} does not exist."},
+            )
+
+        # Retrieve team callback settings from metadata
+        team_metadata = _existing_team.metadata
+        team_callback_settings = team_metadata.get("callback_settings", {})
+
+        # Convert to TeamCallbackMetadata object for consistent structure
+        team_callback_settings_obj = TeamCallbackMetadata(**team_callback_settings)
+
+        return {
+            "status": "success",
+            "data": {
+                "team_id": team_id,
+                "success_callbacks": team_callback_settings_obj.success_callback,
+                "failure_callbacks": team_callback_settings_obj.failure_callback,
+                "callback_vars": team_callback_settings_obj.callback_vars,
+            },
+        }
+
+    except Exception as e:
+        verbose_proxy_logger.error(
+            "litellm.proxy.proxy_server.get_team_callbacks(): Exception occurred - {}".format(
+                str(e)
+            )
+        )
+        verbose_proxy_logger.debug(traceback.format_exc())
+        if isinstance(e, HTTPException):
+            raise ProxyException(
+                message=getattr(e, "detail", f"Internal Server Error({str(e)})"),
+                type=ProxyErrorTypes.internal_server_error.value,
+                param=getattr(e, "param", "None"),
+                code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR),
+            )
+        elif isinstance(e, ProxyException):
+            raise e
+        raise ProxyException(
+            message="Internal Server Error, " + str(e),
+            type=ProxyErrorTypes.internal_server_error.value,
+            param=getattr(e, "param", "None"),
+            code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+        )

litellm/proxy/proxy_server.py

@@ -170,6 +170,9 @@ from litellm.proxy.management_endpoints.key_management_endpoints import (
 from litellm.proxy.management_endpoints.key_management_endpoints import (
     router as key_management_router,
 )
+from litellm.proxy.management_endpoints.team_callback_endpoints import (
+    router as team_callback_router,
+)
 from litellm.proxy.management_endpoints.team_endpoints import router as team_router
 from litellm.proxy.openai_files_endpoints.files_endpoints import (
     router as openai_files_router,
@@ -9457,3 +9460,4 @@ app.include_router(analytics_router)
 app.include_router(debugging_endpoints_router)
 app.include_router(ui_crud_endpoints_router)
 app.include_router(openai_files_router)
+app.include_router(team_callback_router)

litellm/proxy/utils.py

@@ -1313,6 +1313,7 @@ class PrismaClient:
                     t.tpm_limit AS team_tpm_limit,
                     t.rpm_limit AS team_rpm_limit,
                     t.models AS team_models,
+                    t.metadata AS team_metadata,
                     t.blocked AS team_blocked,
                     t.team_alias AS team_alias,
                     tm.spend AS team_member_spend,