feat(langsmith.py): add langsmith sampling rate

Closes LIT-879

Author: krrishdholakia

litellm/proxy/litellm_pre_call_utils.py

@@ -21,7 +21,9 @@
 )
 
 # Cache special headers as a frozenset for O(1) lookup performance
-_SPECIAL_HEADERS_CACHE = frozenset(v.value.lower() for v in SpecialHeaders._member_map_.values())
+_SPECIAL_HEADERS_CACHE = frozenset(
+    v.value.lower() for v in SpecialHeaders._member_map_.values()
+)
 from litellm.proxy.auth.route_checks import RouteChecks
 from litellm.router import Router
 from litellm.types.llms.anthropic import ANTHROPIC_API_HEADERS
@@ -64,6 +66,7 @@ def parse_cache_control(cache_control):
     "files",
 )
 
+
 def _get_metadata_variable_name(request: Request) -> str:
     """
     Helper to return what the "metadata" field should be called in the request data
@@ -157,6 +160,7 @@ def get_key_dynamic_logging_settings(user_api_key_dict: UserAPIKeyAuth):
 
     @staticmethod
     def get_team_dynamic_logging_settings(user_api_key_dict: UserAPIKeyAuth):
+
         if (
             user_api_key_dict.team_metadata is not None
             and "logging" in user_api_key_dict.team_metadata
@@ -169,12 +173,12 @@ def _get_dynamic_logging_metadata(
     user_api_key_dict: UserAPIKeyAuth, proxy_config: ProxyConfig
 ) -> Optional[TeamCallbackMetadata]:
     callback_settings_obj: Optional[TeamCallbackMetadata] = None
-    key_dynamic_logging_settings: Optional[
-        dict
-    ] = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict)
-    team_dynamic_logging_settings: Optional[
-        dict
-    ] = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(user_api_key_dict)
+    key_dynamic_logging_settings: Optional[dict] = (
+        KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict)
+    )
+    team_dynamic_logging_settings: Optional[dict] = (
+        KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(user_api_key_dict)
+    )
     #########################################################################################
     # Key-based callbacks
     #########################################################################################
@@ -234,12 +238,16 @@ def clean_headers(
     Removes litellm api key from headers
     """
     clean_headers = {}
-    litellm_key_lower = litellm_key_header_name.lower() if litellm_key_header_name is not None else None
-    
+    litellm_key_lower = (
+        litellm_key_header_name.lower() if litellm_key_header_name is not None else None
+    )
+
     for header, value in headers.items():
         header_lower = header.lower()
         # Check if header should be excluded: either in special headers cache or matches custom litellm key
-        if (header_lower not in _SPECIAL_HEADERS_CACHE and (litellm_key_lower is None or header_lower != litellm_key_lower)):
+        if header_lower not in _SPECIAL_HEADERS_CACHE and (
+            litellm_key_lower is None or header_lower != litellm_key_lower
+        ):
             clean_headers[header] = value
     return clean_headers
 
@@ -614,11 +622,11 @@ def add_key_level_controls(
 
         ## KEY-LEVEL SPEND LOGS / TAGS
         if "tags" in key_metadata and key_metadata["tags"] is not None:
-            data[_metadata_variable_name][
-                "tags"
-            ] = LiteLLMProxyRequestSetup._merge_tags(
-                request_tags=data[_metadata_variable_name].get("tags"),
-                tags_to_add=key_metadata["tags"],
+            data[_metadata_variable_name]["tags"] = (
+                LiteLLMProxyRequestSetup._merge_tags(
+                    request_tags=data[_metadata_variable_name].get("tags"),
+                    tags_to_add=key_metadata["tags"],
+                )
             )
         if "spend_logs_metadata" in key_metadata and isinstance(
             key_metadata["spend_logs_metadata"], dict
@@ -847,9 +855,9 @@ async def add_litellm_data_to_request(  # noqa: PLR0915
     data[_metadata_variable_name]["litellm_api_version"] = version
 
     if general_settings is not None:
-        data[_metadata_variable_name][
-            "global_max_parallel_requests"
-        ] = general_settings.get("global_max_parallel_requests", None)
+        data[_metadata_variable_name]["global_max_parallel_requests"] = (
+            general_settings.get("global_max_parallel_requests", None)
+        )
 
     ### KEY-LEVEL Controls
     key_metadata = user_api_key_dict.metadata

litellm/proxy/management_endpoints/team_callback_endpoints.py

@@ -79,10 +79,14 @@ async def add_team_callbacks(
 
     """
     try:
+        from litellm.proxy._types import CommonProxyErrors
         from litellm.proxy.proxy_server import prisma_client
 
         if prisma_client is None:
-            raise HTTPException(status_code=500, detail={"error": "No db connected"})
+            raise HTTPException(
+                status_code=500,
+                detail={"error": CommonProxyErrors.db_not_connected_error.value},
+            )
 
         # Check if team_id exists already
         _existing_team = await prisma_client.get_data(
@@ -101,13 +105,14 @@ async def add_team_callbacks(
         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}",
+                    message=f"callback_name = {data.callback_name} already exists in success_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",
@@ -168,22 +173,16 @@ async def add_team_callbacks(
             "data": new_team_row,
         }
 
+    except HTTPException as e:
+        raise e
+    except ProxyException as e:
+        raise e
     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,

litellm/types/utils.py

@@ -2140,6 +2140,7 @@ class StandardCallbackDynamicParams(TypedDict, total=False):
     langsmith_api_key: Optional[str]
     langsmith_project: Optional[str]
     langsmith_base_url: Optional[str]
+    langsmith_sampling_rate: Optional[float]
 
     # Humanloop dynamic params
     humanloop_api_key: Optional[str]