[Fix] QA Fixes - Vector Store Object Permissions (#11291)

* fix: QA for key,team,org permissions

* fix: add_vector_store_to_registry

* fix: refactor bedrock guard

* fix: refactor using us east 1 with vector stores

* fix: code QA checks

* fix: testing for mgmt endpoints

Author: ishaan-jaff

litellm/integrations/vector_stores/bedrock_vector_store.py

@@ -34,7 +34,6 @@
     VectorStoreSearchResponse,
     VectorStoreSearchResult,
 )
-from litellm.utils import load_credentials_from_list
 
 if TYPE_CHECKING:
     from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@@ -258,22 +257,49 @@ async def make_bedrock_kb_retrieve_request(
         from fastapi import HTTPException
 
         non_default_params = non_default_params or {}
-        load_credentials_from_list(kwargs=non_default_params)
+        credentials_dict: Dict[str, Any] = {}
+        if litellm.vector_store_registry is not None:
+            credentials_dict = (
+                litellm.vector_store_registry.get_credentials_for_vector_store(
+                    knowledge_base_id
+                )
+            )
+
         credentials = self.get_credentials(
-            aws_access_key_id=non_default_params.get("aws_access_key_id", None),
-            aws_secret_access_key=non_default_params.get("aws_secret_access_key", None),
-            aws_session_token=non_default_params.get("aws_session_token", None),
-            aws_region_name=non_default_params.get("aws_region_name", None),
-            aws_session_name=non_default_params.get("aws_session_name", None),
-            aws_profile_name=non_default_params.get("aws_profile_name", None),
-            aws_role_name=non_default_params.get("aws_role_name", None),
-            aws_web_identity_token=non_default_params.get(
-                "aws_web_identity_token", None
+            aws_access_key_id=credentials_dict.get(
+                "aws_access_key_id", non_default_params.get("aws_access_key_id", None)
+            ),
+            aws_secret_access_key=credentials_dict.get(
+                "aws_secret_access_key",
+                non_default_params.get("aws_secret_access_key", None),
+            ),
+            aws_session_token=credentials_dict.get(
+                "aws_session_token", non_default_params.get("aws_session_token", None)
+            ),
+            aws_region_name=credentials_dict.get(
+                "aws_region_name", non_default_params.get("aws_region_name", None)
+            ),
+            aws_session_name=credentials_dict.get(
+                "aws_session_name", non_default_params.get("aws_session_name", None)
+            ),
+            aws_profile_name=credentials_dict.get(
+                "aws_profile_name", non_default_params.get("aws_profile_name", None)
+            ),
+            aws_role_name=credentials_dict.get(
+                "aws_role_name", non_default_params.get("aws_role_name", None)
+            ),
+            aws_web_identity_token=credentials_dict.get(
+                "aws_web_identity_token",
+                non_default_params.get("aws_web_identity_token", None),
+            ),
+            aws_sts_endpoint=credentials_dict.get(
+                "aws_sts_endpoint", non_default_params.get("aws_sts_endpoint", None)
             ),
-            aws_sts_endpoint=non_default_params.get("aws_sts_endpoint", None),
         )
-        aws_region_name = self._get_aws_region_name(
-            optional_params=self.optional_params
+        aws_region_name = self.get_aws_region_name_for_non_llm_api_calls(
+            aws_region_name=credentials_dict.get(
+                "aws_region_name", non_default_params.get("aws_region_name", None)
+            ),
         )
 
         # Prepare request data

litellm/llms/bedrock/base_aws_llm.py

@@ -336,6 +336,36 @@ def _get_aws_region_name(
 
         return aws_region_name
 
+    def get_aws_region_name_for_non_llm_api_calls(
+        self,
+        aws_region_name: Optional[str] = None,
+    ):
+        """
+        Get the AWS region name for non-llm api calls.
+
+        LLM API calls check the model arn and end up using that as the region name.
+
+        For non-llm api calls eg. Guardrails, Vector Stores we just need to check the dynamic param or env vars.
+        """
+        if aws_region_name is None:
+            # check env #
+            litellm_aws_region_name = get_secret("AWS_REGION_NAME", None)
+
+            if litellm_aws_region_name is not None and isinstance(
+                litellm_aws_region_name, str
+            ):
+                aws_region_name = litellm_aws_region_name
+
+            standard_aws_region_name = get_secret("AWS_REGION", None)
+            if standard_aws_region_name is not None and isinstance(
+                standard_aws_region_name, str
+            ):
+                aws_region_name = standard_aws_region_name
+
+            if aws_region_name is None:
+                aws_region_name = "us-west-2"
+        return aws_region_name
+
     @tracer.wrap()
     def _auth_with_web_identity_token(
         self,

litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py

@@ -30,7 +30,6 @@
     httpxSpecialProvider,
 )
 from litellm.proxy._types import UserAPIKeyAuth
-from litellm.secret_managers.main import get_secret
 from litellm.types.guardrails import GuardrailEventHooks
 from litellm.types.llms.openai import AllMessageValues
 from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
@@ -129,23 +128,9 @@ def _load_credentials(
         aws_sts_endpoint = self.optional_params.get("aws_sts_endpoint", None)
 
         ### SET REGION NAME ###
-        if aws_region_name is None:
-            # check env #
-            litellm_aws_region_name = get_secret("AWS_REGION_NAME", None)
-
-            if litellm_aws_region_name is not None and isinstance(
-                litellm_aws_region_name, str
-            ):
-                aws_region_name = litellm_aws_region_name
-
-            standard_aws_region_name = get_secret("AWS_REGION", None)
-            if standard_aws_region_name is not None and isinstance(
-                standard_aws_region_name, str
-            ):
-                aws_region_name = standard_aws_region_name
-
-            if aws_region_name is None:
-                aws_region_name = "us-west-2"
+        aws_region_name = self.get_aws_region_name_for_non_llm_api_calls(
+            aws_region_name=aws_region_name,
+        )
 
         credentials: Credentials = self.get_credentials(
             aws_access_key_id=aws_access_key_id,

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -45,6 +45,9 @@
 from litellm.proxy.management_endpoints.model_management_endpoints import (
     _add_model_to_db,
 )
+from litellm.proxy.management_helpers.object_permission_utils import (
+    handle_update_object_permission_common,
+)
 from litellm.proxy.management_helpers.team_member_permission_checks import (
     TeamMemberPermissionChecks,
 )
@@ -733,48 +736,20 @@ async def _handle_update_object_permission(
     """
     from litellm.proxy.proxy_server import prisma_client
 
-    if prisma_client is None:
-        raise ValueError("Prisma client not found")
-
-    #########################################################
-    # Ensure `object_permission` is not added to the data_json
-    # We need to update the entity at the object_permission_id level in the LiteLLM_ObjectPermissionTable
-    #########################################################
-    new_object_permission = data_json.pop("object_permission")
-    if new_object_permission is None:
-        return data_json
-
-    # lookup existing object permission ID and update that entry
-    existing_object_permission_id = existing_key_row.object_permission_id
-    existing_object_permission = (
-        await prisma_client.db.litellm_objectpermissiontable.find_unique(
-            where={"object_permission_id": existing_object_permission_id},
-        )
+    # Use the common helper to handle the object permission update
+    object_permission_id = await handle_update_object_permission_common(
+        data_json=data_json,
+        existing_object_permission_id=existing_key_row.object_permission_id,
+        prisma_client=prisma_client,
     )
-    existing_object_permissions_dict: dict = {}
-    if existing_object_permission is not None:
-        # update the object permission
-        existing_object_permissions_dict = existing_object_permission.model_dump(
-            exclude_unset=True, exclude_none=True
-        )
-        existing_object_permissions_dict.update(dict(new_object_permission))
-
-    #########################################################
-    # Commit the update to the LiteLLM_ObjectPermissionTable
-    #########################################################
-    new_object_permission_row = (
-        await prisma_client.db.litellm_objectpermissiontable.upsert(
-            where={"object_permission_id": existing_object_permission_id},
-            data={
-                "create": existing_object_permissions_dict,
-                "update": existing_object_permissions_dict,
-            },
+
+    # Add the object_permission_id to data_json if one was created/updated
+    if object_permission_id is not None:
+        data_json["object_permission_id"] = object_permission_id
+        verbose_proxy_logger.debug(
+            f"updated object_permission_id: {object_permission_id}"
         )
-    )
-    verbose_proxy_logger.debug(
-        f"new_object_permission_row: {new_object_permission_row}"
-    )
-    data_json["object_permission_id"] = new_object_permission_row.object_permission_id
+
     return data_json
 
 

litellm/proxy/management_endpoints/organization_endpoints.py

@@ -24,6 +24,9 @@
     new_budget,
     update_budget,
 )
+from litellm.proxy.management_helpers.object_permission_utils import (
+    handle_update_object_permission_common,
+)
 from litellm.proxy.management_helpers.utils import (
     get_new_internal_user_defaults,
     management_endpoint_wrapper,
@@ -309,61 +312,20 @@ async def handle_update_object_permission(
     - Upserts the new object permission into the LiteLLM_ObjectPermissionTable
     - Adds object_permission_id to data_json (this gets added in the DB)
     - Pops the object_permission from data_json
-    -
     """
     from litellm.proxy.proxy_server import prisma_client
 
-    if prisma_client is None:
-        raise ValueError("Prisma client not found")
-
-    #########################################################
-    # Ensure `object_permission` is not added to the data_json
-    # We need to update the entity at the object_permission_id level in the LiteLLM_ObjectPermissionTable
-    #########################################################
-    new_object_permission: Union[dict, str] = data_json.pop("object_permission") or {}
-    if new_object_permission is None:
-        return data_json
-
-    # lookup existing object permission ID and update that entry
-    existing_object_permission_id = existing_organization_row.object_permission_id
-    existing_object_permissions_dict = {}
-
-    existing_object_permission = (
-        await prisma_client.db.litellm_objectpermissiontable.find_unique(
-            where={"object_permission_id": existing_object_permission_id},
-        )
+    # Use the common helper to handle the object permission update
+    object_permission_id = await handle_update_object_permission_common(
+        data_json=data_json,
+        existing_object_permission_id=existing_organization_row.object_permission_id,
+        prisma_client=prisma_client,
     )
 
-    # update the object permission
-    if existing_object_permission is not None:
-        existing_object_permissions_dict = existing_object_permission.model_dump(
-            exclude_unset=True, exclude_none=True
-        )
+    # Add the object_permission_id to data_json if one was created/updated
+    if object_permission_id is not None:
+        data_json["object_permission_id"] = object_permission_id
 
-    if isinstance(new_object_permission, str):
-        new_object_permission = json.loads(new_object_permission)
-
-    if isinstance(new_object_permission, dict):
-        existing_object_permissions_dict.update(new_object_permission)
-
-    #########################################################
-    # Commit the update to the LiteLLM_ObjectPermissionTable
-    #########################################################
-    created_object_permission_row = (
-        await prisma_client.db.litellm_objectpermissiontable.upsert(
-            where={"object_permission_id": existing_object_permission_id},
-            data={
-                "create": existing_object_permissions_dict,
-                "update": existing_object_permissions_dict,
-            },
-        )
-    )
-    data_json[
-        "object_permission_id"
-    ] = created_object_permission_row.object_permission_id
-    verbose_proxy_logger.debug(
-        f"created_object_permission_row: {created_object_permission_row}"
-    )
     return data_json
 
 

litellm/proxy/management_endpoints/team_endpoints.py

@@ -70,6 +70,9 @@
 from litellm.proxy.management_endpoints.tag_management_endpoints import (
     get_daily_activity,
 )
+from litellm.proxy.management_helpers.object_permission_utils import (
+    handle_update_object_permission_common,
+)
 from litellm.proxy.management_helpers.team_member_permission_checks import (
     TeamMemberPermissionChecks,
 )
@@ -752,48 +755,20 @@ async def handle_update_object_permission(
     """
     from litellm.proxy.proxy_server import prisma_client
 
-    if prisma_client is None:
-        raise ValueError("Prisma client not found")
-
-    #########################################################
-    # Ensure `object_permission` is not added to the data_json
-    # We need to update the entity at the object_permission_id level in the LiteLLM_ObjectPermissionTable
-    #########################################################
-    new_object_permission = data_json.pop("object_permission")
-    if new_object_permission is None:
-        return data_json
-
-    # lookup existing object permission ID and update that entry
-    existing_object_permission_id = existing_team_row.object_permission_id
-    existing_object_permission = (
-        await prisma_client.db.litellm_objectpermissiontable.find_unique(
-            where={"object_permission_id": existing_object_permission_id},
-        )
+    # Use the common helper to handle the object permission update
+    object_permission_id = await handle_update_object_permission_common(
+        data_json=data_json,
+        existing_object_permission_id=existing_team_row.object_permission_id,
+        prisma_client=prisma_client,
     )
 
-    existing_object_permissions_dict: Dict = {}
-
-    # update the object permission
-    if existing_object_permission is not None:
-        existing_object_permissions_dict = existing_object_permission.model_dump(
-            exclude_unset=True, exclude_none=True
+    # Add the object_permission_id to data_json if one was created/updated
+    if object_permission_id is not None:
+        data_json["object_permission_id"] = object_permission_id
+        verbose_proxy_logger.debug(
+            f"updated object_permission_id: {object_permission_id}"
         )
-    existing_object_permissions_dict.update(dict(new_object_permission))
-    created_object_permission_row = (
-        await prisma_client.db.litellm_objectpermissiontable.upsert(
-            where={"object_permission_id": existing_object_permission_id},
-            data={
-                "create": existing_object_permissions_dict,
-                "update": existing_object_permissions_dict,
-            },
-        )
-    )
-    data_json[
-        "object_permission_id"
-    ] = created_object_permission_row.object_permission_id
-    verbose_proxy_logger.debug(
-        f"created_object_permission_row: {created_object_permission_row}"
-    )
+
     return data_json