litellm-enterprise 0.1.14__tar.gz → 0.1.16__tar.gz

{litellm_enterprise-0.1.14 → litellm_enterprise-0.1.16}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: litellm-enterprise
-Version: 0.1.14
+Version: 0.1.16
 Summary: Package for LiteLLM Enterprise features
 Author: BerriAI
 Requires-Python: >=3.8, !=2.7.*, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*, !=3.5.*, !=3.6.*, !=3.7.*

{litellm_enterprise-0.1.14 → litellm_enterprise-0.1.16}/litellm_enterprise/enterprise_callbacks/callback_controls.py

@@ -1,3 +1,5 @@
+from typing import List, Optional
+
 import litellm
 from litellm._logging import verbose_logger
 from litellm.constants import X_LITELLM_DISABLE_CALLBACKS
@@ -6,15 +8,18 @@ from litellm.litellm_core_utils.llm_request_utils import (
     get_proxy_server_request_headers,
 )
 from litellm.proxy._types import CommonProxyErrors
+from litellm.types.utils import StandardCallbackDynamicParams
 
 
 class EnterpriseCallbackControls:
     @staticmethod
-    def is_callback_disabled_via_headers(
-            callback: litellm.CALLBACK_TYPES, litellm_params: dict
+    def is_callback_disabled_dynamically(
+            callback: litellm.CALLBACK_TYPES, 
+            litellm_params: dict,
+            standard_callback_dynamic_params: StandardCallbackDynamicParams
         ) -> bool:
             """
-            Check if a callback is disabled via the x-litellm-disable-callbacks header.
+            Check if a callback is disabled via the x-litellm-disable-callbacks header or via `litellm_disabled_callbacks` in standard_callback_dynamic_params.
             
             Args:
                 callback: The callback to check (can be string, CustomLogger instance, or callable)
@@ -28,8 +33,7 @@ class EnterpriseCallbackControls:
             )
 
             try:
-                request_headers = get_proxy_server_request_headers(litellm_params)
-                disabled_callbacks = request_headers.get(X_LITELLM_DISABLE_CALLBACKS, None)
+                disabled_callbacks = EnterpriseCallbackControls.get_disabled_callbacks(litellm_params, standard_callback_dynamic_params)
                 verbose_logger.debug(f"Dynamically disabled callbacks from {X_LITELLM_DISABLE_CALLBACKS}: {disabled_callbacks}")
                 verbose_logger.debug(f"Checking if {callback} is disabled via headers. Disable callbacks from headers: {disabled_callbacks}")
                 if disabled_callbacks is not None:
@@ -39,7 +43,6 @@ class EnterpriseCallbackControls:
                     if not EnterpriseCallbackControls._premium_user_check():
                         return False
                     #########################################################
-                    disabled_callbacks = set([cb.strip().lower() for cb in disabled_callbacks.split(",")])
                     if isinstance(callback, str):
                         if callback.lower() in disabled_callbacks:
                             verbose_logger.debug(f"Not logging to {callback} because it is disabled via {X_LITELLM_DISABLE_CALLBACKS}")
@@ -56,6 +59,29 @@ class EnterpriseCallbackControls:
                     f"Error checking disabled callbacks header: {str(e)}"
                 )
                 return False
+    @staticmethod
+    def get_disabled_callbacks(litellm_params: dict, standard_callback_dynamic_params: StandardCallbackDynamicParams) -> Optional[List[str]]:
+        """
+        Get the disabled callbacks from the standard callback dynamic params.
+        """
+
+        #########################################################
+        # check if disabled via headers
+        #########################################################
+        request_headers = get_proxy_server_request_headers(litellm_params)
+        disabled_callbacks = request_headers.get(X_LITELLM_DISABLE_CALLBACKS, None)
+        if disabled_callbacks is not None:
+            disabled_callbacks = set([cb.strip().lower() for cb in disabled_callbacks.split(",")])
+            return list(disabled_callbacks)
+        
+
+        #########################################################
+        # check if disabled via request body
+        #########################################################
+        if standard_callback_dynamic_params.get("litellm_disabled_callbacks", None) is not None:
+            return standard_callback_dynamic_params.get("litellm_disabled_callbacks", None)
+        
+        return None
     
     @staticmethod
     def _premium_user_check():

{litellm_enterprise-0.1.14 → litellm_enterprise-0.1.16}/litellm_enterprise/integrations/prometheus.py

@@ -1,4 +1,4 @@
- # used for /metrics endpoint on LiteLLM Proxy
+# used for /metrics endpoint on LiteLLM Proxy
 #### What this does ####
 #    On success, log events to Prometheus
 import sys
@@ -1231,8 +1231,15 @@ class PrometheusLogger(CustomLogger):
                     "team_alias",
                 ] + EXCEPTION_LABELS,
         """
+        from litellm.litellm_core_utils.litellm_logging import (
+            StandardLoggingPayloadSetup,
+        )
+
         try:
-            _tags = cast(List[str], request_data.get("tags") or [])
+            _tags = StandardLoggingPayloadSetup._get_request_tags(
+                request_data.get("metadata", {}),
+                request_data.get("proxy_server_request", {}),
+            )
             enum_values = UserAPIKeyLabelValues(
                 end_user=user_api_key_dict.end_user_id,
                 user=user_api_key_dict.user_id,
@@ -1277,6 +1284,10 @@ class PrometheusLogger(CustomLogger):
         Proxy level tracking - triggered when the proxy responds with a success response to the client
         """
         try:
+            from litellm.litellm_core_utils.litellm_logging import (
+                StandardLoggingPayloadSetup,
+            )
+
             enum_values = UserAPIKeyLabelValues(
                 end_user=user_api_key_dict.end_user_id,
                 hashed_api_key=user_api_key_dict.api_key,
@@ -1288,6 +1299,9 @@ class PrometheusLogger(CustomLogger):
                 user_email=user_api_key_dict.user_email,
                 status_code="200",
                 route=user_api_key_dict.request_route,
+                tags=StandardLoggingPayloadSetup._get_request_tags(
+                    data.get("metadata", {}), data.get("proxy_server_request", {})
+                ),
             )
             _labels = prometheus_label_factory(
                 supported_enum_labels=self.get_labels_for_metric(
@@ -1352,6 +1366,7 @@ class PrometheusLogger(CustomLogger):
                 team_alias=standard_logging_payload["metadata"][
                     "user_api_key_team_alias"
                 ],
+                tags=standard_logging_payload.get("request_tags", []),
             )
 
             """
@@ -2148,9 +2163,11 @@ class PrometheusLogger(CustomLogger):
 
         It emits the current remaining budget metrics for all Keys and Teams.
         """
+        from enterprise.litellm_enterprise.integrations.prometheus import (
+            PrometheusLogger,
+        )
         from litellm.constants import PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES
         from litellm.integrations.custom_logger import CustomLogger
-        from enterprise.litellm_enterprise.integrations.prometheus import PrometheusLogger
 
         prometheus_loggers: List[CustomLogger] = (
             litellm.logging_callback_manager.get_custom_loggers_for_type(
@@ -2280,6 +2297,8 @@ def get_custom_labels_from_tags(tags: List[str]) -> Dict[str, str]:
     """
     Get custom labels from tags based on admin configuration
     """
+    from litellm.types.integrations.prometheus import _sanitize_prometheus_label_name
+
     configured_tags = litellm.custom_prometheus_tags
     if configured_tags is None or len(configured_tags) == 0:
         return {}
@@ -2289,7 +2308,7 @@ def get_custom_labels_from_tags(tags: List[str]) -> Dict[str, str]:
     # Map each configured tag to its presence in the request tags
     for configured_tag in configured_tags:
         # Create a safe prometheus label name
-        label_name = f"tag_{configured_tag}".replace("-", "_").replace(".", "_")
+        label_name = _sanitize_prometheus_label_name(f"tag_{configured_tag}")
 
         # Check if this tag is present in the request tags
         if configured_tag in tags:

{litellm_enterprise-0.1.14 → litellm_enterprise-0.1.16}/litellm_enterprise/proxy/common_utils/check_batch_cost.py

@@ -83,15 +83,25 @@ class CheckBatchCost:
                 )
                 continue
 
-            response = await self.llm_router.aretrieve_batch(
-                model=model_id,
-                batch_id=batch_id,
-                litellm_metadata={
-                    "user_api_key_user_id": job.created_by or "default-user-id",
-                    "batch_ignore_default_logging": True,
-                },
+            verbose_proxy_logger.info(
+                f"Querying model ID: {model_id} for cost and usage of batch ID: {batch_id}"
             )
 
+            try:
+                response = await self.llm_router.aretrieve_batch(
+                    model=model_id,
+                    batch_id=batch_id,
+                    litellm_metadata={
+                        "user_api_key_user_id": job.created_by or "default-user-id",
+                        "batch_ignore_default_logging": True,
+                    },
+                )
+            except Exception as e:
+                verbose_proxy_logger.info(
+                    f"Skipping job {unified_object_id} because of error querying model ID: {model_id} for cost and usage of batch ID: {batch_id}: {e}"
+                )
+                continue
+
             ## RETRIEVE THE BATCH JOB OUTPUT FILE
             managed_files_obj = cast(
                 Optional[_PROXY_LiteLLMManagedFiles],
@@ -102,6 +112,9 @@ class CheckBatchCost:
                 and response.output_file_id is not None
                 and managed_files_obj is not None
             ):
+                verbose_proxy_logger.info(
+                    f"Batch ID: {batch_id} is complete, tracking cost and usage"
+                )
                 # track cost
                 model_file_id_mapping = {
                     response.output_file_id: {model_id: response.output_file_id}

{litellm_enterprise-0.1.14 → litellm_enterprise-0.1.16}/litellm_enterprise/proxy/vector_stores/endpoints.py

@@ -9,9 +9,9 @@ All /vector_store management endpoints
 """
 
 import copy
-from typing import List
+from typing import List, Optional
 
-from fastapi import APIRouter, Depends, HTTPException
+from fastapi import APIRouter, Depends, HTTPException, Request, Response
 
 import litellm
 from litellm._logging import verbose_proxy_logger
@@ -22,12 +22,16 @@ from litellm.types.vector_stores import (
     LiteLLM_ManagedVectorStore,
     LiteLLM_ManagedVectorStoreListResponse,
     VectorStoreDeleteRequest,
+    VectorStoreInfoRequest,
+    VectorStoreUpdateRequest,
 )
 from litellm.vector_stores.vector_store_registry import VectorStoreRegistry
 
 router = APIRouter()
 
-
+########################################################
+# Management Endpoints
+########################################################
 @router.post(
     "/vector_store/new",
     tags=["vector store management"],
@@ -48,6 +52,7 @@ async def new_vector_store(
     - vector_store_metadata: Optional[Dict] - Additional metadata for the vector store
     """
     from litellm.proxy.proxy_server import prisma_client
+    from litellm.types.router import GenericLiteLLMParams
 
     if prisma_client is None:
         raise HTTPException(status_code=500, detail="Database not connected")
@@ -70,9 +75,20 @@ async def new_vector_store(
                 vector_store.get("vector_store_metadata")
             )
 
+        # Safely handle JSON serialization of litellm_params
+        litellm_params_json: Optional[str] = None
+        _input_litellm_params: dict = vector_store.get("litellm_params", {}) or {}
+        if _input_litellm_params is not None:
+            litellm_params_dict = GenericLiteLLMParams(**_input_litellm_params).model_dump(exclude_none=True)
+            litellm_params_json = safe_dumps(litellm_params_dict)
+            del vector_store["litellm_params"]
+
         _new_vector_store = (
             await prisma_client.db.litellm_managedvectorstorestable.create(
-                data=vector_store
+                data={
+                    **vector_store,
+                    "litellm_params": litellm_params_json,
+                }
             )
         )
 
@@ -205,3 +221,75 @@ async def delete_vector_store(
         return {"message": f"Vector store {data.vector_store_id} deleted successfully"}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
+
+
+@router.post(
+    "/vector_store/info",
+    tags=["vector store management"],
+    dependencies=[Depends(user_api_key_auth)],
+)
+async def get_vector_store_info(
+    data: VectorStoreInfoRequest,
+    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
+):
+    """Return a single vector store's details"""
+    from litellm.proxy.proxy_server import prisma_client
+
+    if prisma_client is None:
+        raise HTTPException(status_code=500, detail="Database not connected")
+
+    try:
+        vector_store = await prisma_client.db.litellm_managedvectorstorestable.find_unique(
+            where={"vector_store_id": data.vector_store_id}
+        )
+        if vector_store is None:
+            raise HTTPException(
+                status_code=404,
+                detail=f"Vector store with ID {data.vector_store_id} not found",
+            )
+
+        vector_store_dict = vector_store.model_dump()
+        return {"vector_store": vector_store_dict}
+    except Exception as e:
+        verbose_proxy_logger.exception(f"Error getting vector store info: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@router.post(
+    "/vector_store/update",
+    tags=["vector store management"],
+    dependencies=[Depends(user_api_key_auth)],
+)
+async def update_vector_store(
+    data: VectorStoreUpdateRequest,
+    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
+):
+    """Update vector store details"""
+    from litellm.proxy.proxy_server import prisma_client
+
+    if prisma_client is None:
+        raise HTTPException(status_code=500, detail="Database not connected")
+
+    try:
+        update_data = data.model_dump(exclude_unset=True)
+        vector_store_id = update_data.pop("vector_store_id")
+        if update_data.get("vector_store_metadata") is not None:
+            update_data["vector_store_metadata"] = safe_dumps(update_data["vector_store_metadata"])
+
+        updated = await prisma_client.db.litellm_managedvectorstorestable.update(
+            where={"vector_store_id": vector_store_id},
+            data=update_data,
+        )
+
+        updated_vs = LiteLLM_ManagedVectorStore(**updated.model_dump())
+
+        if litellm.vector_store_registry is not None:
+            litellm.vector_store_registry.update_vector_store_in_registry(
+                vector_store_id=vector_store_id,
+                updated_data=updated_vs,
+            )
+
+        return {"vector_store": updated_vs}
+    except Exception as e:
+        verbose_proxy_logger.exception(f"Error updating vector store: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))

{litellm_enterprise-0.1.14 → litellm_enterprise-0.1.16}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "litellm-enterprise"
-version = "0.1.14"
+version = "0.1.16"
 description = "Package for LiteLLM Enterprise features"
 authors = ["BerriAI"]
 readme = "README.md"
@@ -22,7 +22,7 @@ requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.commitizen]
-version = "0.1.14"
+version = "0.1.16"
 version_files = [
     "pyproject.toml:version",
     "../requirements.txt:litellm-enterprise==",