Repository: privacera/paig (language: python)

Issue:
[ENH]: Extend AWS Bedrock Guardrail support to take properties from UI
### Contact Details

amrakshay

### Feature Description

In #120 we added support for AWS Bedrock Guardrail by setting the properties.
We need to add / extend support of setting the Guardrail id and version for the application on UI.
Shield Service to check the guardrail id and version from database and make call to AWS Bedrock Guardrail to authorize request. 

### Code of Conduct

- [X] I agree to follow this project's Code of Conduct
Create base version
Create base version - initial commit of PAIG OSS


Discussion / hints:

Fix patch:
diff --git a/paig-server/backend/paig/alembic_db/versions/a95b604c47fb_added_guardrail_details_column_to_ai_app.py b/paig-server/backend/paig/alembic_db/versions/a95b604c47fb_added_guardrail_details_column_to_ai_app.py
new file mode 100644
index 00000000..250f0d50
--- /dev/null
+++ b/paig-server/backend/paig/alembic_db/versions/a95b604c47fb_added_guardrail_details_column_to_ai_app.py
@@ -0,0 +1,30 @@
+"""added guardrail_details column to ai_app table
+
+Revision ID: a95b604c47fb
+Revises: 5e805b526efa
+Create Date: 2024-10-25 13:30:46.711052
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = 'a95b604c47fb'
+down_revision: Union[str, None] = '5e805b526efa'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column('ai_application', sa.Column('guardrail_details', sa.String(length=255), nullable=True))
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column('ai_application', 'guardrail_details')
+    # ### end Alembic commands ###
diff --git a/paig-server/backend/paig/api/governance/api_schemas/ai_app.py b/paig-server/backend/paig/api/governance/api_schemas/ai_app.py
index 9b46411e..619a642f 100644
--- a/paig-server/backend/paig/api/governance/api_schemas/ai_app.py
+++ b/paig-server/backend/paig/api/governance/api_schemas/ai_app.py
@@ -26,6 +26,7 @@ class AIApplicationView(BaseView):
     description: Optional[str] = Field(default=None, description="The description of the AI application")
     application_key: Optional[str] = Field(None, description="The application key", alias="applicationKey")
     vector_dbs: Optional[List[str]] = Field([], description="The vector databases associated with the AI application", alias="vectorDBs")
+    guardrail_details: Optional[str] = Field(None, description="The guardrail details", alias="guardrailDetails")
 
     vector_db_id: Optional[int] = Field(None, description="The vector databases id with the AI application",
                                         alias="vectorDBId")
diff --git a/paig-server/backend/paig/api/governance/database/db_models/ai_app_model.py b/paig-server/backend/paig/api/governance/database/db_models/ai_app_model.py
index c15fd15c..7c8960db 100644
--- a/paig-server/backend/paig/api/governance/database/db_models/ai_app_model.py
+++ b/paig-server/backend/paig/api/governance/database/db_models/ai_app_model.py
@@ -28,6 +28,7 @@ class AIApplicationModel(BaseSQLModel):
     description = Column(String(4000), nullable=True)
     application_key = Column(String(255), nullable=False)
     vector_dbs = Column(CommaSeparatedList(255), nullable=True)
+    guardrail_details = Column(String(255), nullable=True)
 
     app_config = relationship("AIApplicationConfigModel", back_populates="ai_app", uselist=False, cascade="all, delete-orphan")
     app_policies = relationship("AIApplicationPolicyModel", back_populates="ai_app", cascade="all, delete-orphan")
diff --git a/paig-server/backend/paig/api/shield/client/http_governance_service_client.py b/paig-server/backend/paig/api/shield/client/http_governance_service_client.py
new file mode 100644
index 00000000..59bc2706
--- /dev/null
+++ b/paig-server/backend/paig/api/shield/client/http_governance_service_client.py
@@ -0,0 +1,88 @@
+import logging
+
+
+from paig_common.async_base_rest_http_client import AsyncBaseRESTHttpClient
+from api.shield.utils.custom_exceptions import ShieldException
+
+from api.shield.utils import config_utils
+from api.shield.interfaces.governance_service_interface import IGovernanceServiceClient
+
+logger = logging.getLogger(__name__)
+
+
+class HttpGovernanceServiceClient(AsyncBaseRESTHttpClient, IGovernanceServiceClient):
+    """
+       HTTP client for interacting with an governance service API, implementing the IGovernanceServiceClient interface.
+
+       This client supports fetching AWS Bedrock guardrail details from a remote service using HTTP requests.
+
+       Methods:
+           init_singleton():
+               Initializes the singleton instance of the client with the base URL from configuration.
+
+           get_headers(tenant_id: str):
+               Returns headers for API requests based on the provided tenant_id.
+
+           get_all_encryption_keys(tenant_id):
+               Retrieves all AWS Bedrock guardrail details for a specific tenant from the governance service API.
+       """
+
+    def __init__(self):
+        """
+        Initialize the client instance with the base URL from configuration.
+        """
+        super().__init__(config_utils.get_property_value("governance_service_base_url"))
+
+    @staticmethod
+    def get_headers(tenant_id: str):
+        """
+        Return headers for API requests based on the provided tenant_id.
+
+        Args:
+            tenant_id (str): Identifier of the tenant for which headers are generated.
+
+        Returns:
+            dict: Headers dictionary containing necessary headers for API requests.
+        """
+        paig_key = config_utils.get_property_value('paig_api_key')
+        if paig_key is None:
+            return {"x-tenant-id": tenant_id, "x-user-role": "INTERNAL_SERVICE"}
+
+        return {"x-paig-api-key": paig_key}
+
+    async def get_aws_bedrock_guardrail_info(self, tenant_id, application_key):
+        """
+        Retrieve all AWS Bedrock guardrail details for a specific tenant and application from the governance service API.
+
+        Args:
+            tenant_id (str): Identifier of the tenant whose AWS Bedrock guardrail details are to be fetched.
+            application_key (str): Identifier of the application for which AWS Bedrock guardrail details are to be fetched.
+
+        Returns:
+            dict: Dictionary containing AWS Bedrock guardrail details associated with the specified tenant_id and application_key.
+
+        Raises:
+            ShieldException: If the API request fails or returns a non-200 status code.
+        """
+        url = config_utils.get_property_value("governance_service_get_key_endpoint")
+        url = f'{url}/{application_key}'
+        logger.debug(f"Using base-url={self.baseUrl} uri={url} for tenant-id={tenant_id} and application-key={application_key}")
+
+        try:
+
+            response = await self.get(
+                url=url,
+                headers=self.get_headers(tenant_id),
+            )
+
+            logger.debug(f"response received: {response.__str__()}")
+
+            if response.status_code == 200:
+                return response.json()
+            else:
+                error_message = (f"Request get_aws_bedrock_guardrail_info({tenant_id}, {application_key}) failed "
+                                 f"with status code {response.status_code}: {response.text}")
+                raise ShieldException(error_message)
+        except Exception as ex:
+            logger.error(f"Request get_aws_bedrock_guardrail_info({tenant_id}, {application_key}) failed with error: {ex}")
+            raise ShieldException(ex.args[0])
diff --git a/paig-server/backend/paig/api/shield/client/local_governance_service_client.py b/paig-server/backend/paig/api/shield/client/local_governance_service_client.py
new file mode 100644
index 00000000..80443ec5
--- /dev/null
+++ b/paig-server/backend/paig/api/shield/client/local_governance_service_client.py
@@ -0,0 +1,37 @@
+import json
+from api.governance.services.ai_app_service import AIAppService
+from api.shield.interfaces.governance_service_interface import IGovernanceServiceClient
+from core.utils import SingletonDepends
+
+
+class LocalGovernanceServiceClient(IGovernanceServiceClient):
+    """
+    Local client implementation for fetching AWS Bedrock Guardrail info.
+
+    Methods:
+        get_aws_bedrock_guardrail_info(tenant_id, application_key):
+            Retrieves AWS Bedrock Guardrail details for the specified tenant and application key from a local data source.
+    """
+
+    def __init__(self, ai_application_service: AIAppService = SingletonDepends(AIAppService)):
+        """
+        Initialize the local governance service client.
+        """
+        self.ai_application_service = ai_application_service
+
+    async def get_aws_bedrock_guardrail_info(self, tenant_id, application_key):
+        """
+        Retrieve all encryption keys for a specific tenant from the underlying data source.
+
+        Args:
+            tenant_id (str): The identifier of the tenant whose AWS Bedrock Guardrail details are to be fetched.
+            application_key (str): The identifier of the application for which AWS Bedrock Guardrail details are to be fetched.
+
+        Returns:
+            dict: AWS Bedrock Guardrail details associated with the specified `tenant_id` and 'application_key'.
+        """
+        result = await self.ai_application_service.get_ai_application_by_application_key(application_key)
+        if result.guardrail_details:
+            return json.loads(result.guardrail_details)
+        else:
+            return {}
diff --git a/paig-server/backend/paig/api/shield/factory/account_service_factory.py b/paig-server/backend/paig/api/shield/factory/account_service_factory.py
index d793542c..e6b238b3 100644
--- a/paig-server/backend/paig/api/shield/factory/account_service_factory.py
+++ b/paig-server/backend/paig/api/shield/factory/account_service_factory.py
@@ -1,5 +1,4 @@
 
-from api.encryption.services.encryption_key_service import EncryptionKeyService
 from api.shield.client.local_account_service_client import LocalAccountServiceClient
 from api.shield.client.http_account_service_client import HttpAccountServiceClient
 from api.shield.utils import config_utils
diff --git a/paig-server/backend/paig/api/shield/factory/governance_service_factory.py b/paig-server/backend/paig/api/shield/factory/governance_service_factory.py
new file mode 100644
index 00000000..954f22e6
--- /dev/null
+++ b/paig-server/backend/paig/api/shield/factory/governance_service_factory.py
@@ -0,0 +1,37 @@
+
+from api.shield.client.local_governance_service_client import LocalGovernanceServiceClient
+from api.shield.client.http_governance_service_client import HttpGovernanceServiceClient
+from api.shield.utils import config_utils
+from core.utils import SingletonDepends
+
+
+class GovernanceServiceFactory:
+    """
+    Factory class for creating governance service client instances based on configuration.
+
+    Methods:
+        get_governance_service_client():
+            Returns an instance of an governance service client based on the configured client type.
+    """
+
+    def get_governance_service_client(self):
+        """
+        Returns an instance of an governance service client based on the configured client type.
+
+        Returns:
+            object: Instance of an governance service client.
+
+        Raises:
+            Exception: If an invalid service type is configured in the application.
+        """
+        client_type = config_utils.get_property_value('governance_service_client', 'local')
+        match client_type:
+            case "http":
+                return SingletonDepends(HttpGovernanceServiceClient)
+            case "local":
+                return SingletonDepends(LocalGovernanceServiceClient)
+            case _:
+                raise Exception(
+                    f"Invalid service type: '{client_type}'. Expected 'http' or 'local'. "
+                    "Please configure the 'governance_service_client' property with a valid service type."
+                )
diff --git a/paig-server/backend/paig/api/shield/interfaces/governance_service_interface.py b/paig-server/backend/paig/api/shield/interfaces/governance_service_interface.py
new file mode 100644
index 00000000..6f6ff91c
--- /dev/null
+++ b/paig-server/backend/paig/api/shield/interfaces/governance_service_interface.py
@@ -0,0 +1,15 @@
+from abc import ABC, abstractmethod
+
+
+class IGovernanceServiceClient(ABC):
+    """
+      Interface for fetching AWS Bedrock Guardrail details.
+
+      Methods:
+          get_aws_bedrock_guardrail_info(tenant_id, application_key):
+              Fetches all encryption keys for the specified `tenant_id` and 'application_key'.
+      """
+    @abstractmethod
+    def get_aws_bedrock_guardrail_info(self, tenant_id):
+        """Abstract method to be implemented by subclasses."""
+        pass
diff --git a/paig-server/backend/paig/api/shield/scanners/AWSBedrockGuardrailScanner.py b/paig-server/backend/paig/api/shield/scanners/AWSBedrockGuardrailScanner.py
index dab95911..97b85ad3 100644
--- a/paig-server/backend/paig/api/shield/scanners/AWSBedrockGuardrailScanner.py
+++ b/paig-server/backend/paig/api/shield/scanners/AWSBedrockGuardrailScanner.py
@@ -32,7 +32,7 @@ def __init__(self, **kwargs):
         """
         super().__init__(**kwargs)
 
-        self.guardrail_id, self.guardrail_version, self.region = self.get_guardrail_details()
+        self.guardrail_id, self.guardrail_version, self.region = self._get_guardrail_details()
         self.bedrock_client = boto3.client(
             'bedrock-runtime',
             region_name=self.region
@@ -79,7 +79,7 @@ def scan(self, message: str) -> ScannerResult:
         logger.info("AWSBedrockGuardrailScanner: No action required for the message.")
         return ScannerResult(traits=[])
 
-    def get_guardrail_details(self) -> (str, str, str):
+    def _get_guardrail_details(self) -> (str, str, str):
         """
         Fetch guardrail details
         """
diff --git a/paig-server/backend/paig/api/shield/services/application_manager_service.py b/paig-server/backend/paig/api/shield/services/application_manager_service.py
index aae0eec9..07d2eae6 100644
--- a/paig-server/backend/paig/api/shield/services/application_manager_service.py
+++ b/paig-server/backend/paig/api/shield/services/application_manager_service.py
@@ -1,6 +1,7 @@
 import logging
 
 from api.shield.model.scanner_result import ScannerResult
+from api.shield.model.authorize_request import AuthorizeRequest
 from api.shield.scanners.BaseScanner import Scanner
 from api.shield.scanners.scanner_util import parse_properties
 from api.shield.cache.lru_cache import LRUCache
@@ -44,7 +45,7 @@ def load_scanners(self, application_key: str):
         self.application_key_scanners.put(application_key, scanner_list)
         logger.info(f"Found {scanner_list} scanners for application key: {application_key}")
 
-    def get_scanners(self, application_key: str, request_type: str, is_authz_scan: bool) -> list:
+    def get_scanners(self, application_key: str, request_type: str, is_authz_scan: bool, auth_req: AuthorizeRequest) -> list:
         """
         Get the scanners for the given application key.
         If the scanners are not in the cache, load them.
@@ -63,14 +64,26 @@ def get_scanners(self, application_key: str, request_type: str, is_authz_scan: b
         all_scanners = self.application_key_scanners.get(application_key)
 
         scanners_list = [
-            (setattr(scanner, 'scan_for_req_type', request_type) or scanner)
-            for scanner in all_scanners
+            scanner for scanner in all_scanners
             if getattr(scanner, 'enforce_access_control', False) == is_authz_scan and request_type in getattr(scanner, 'request_types', [])
         ]
 
+        for scanner in scanners_list:
+            setattr(scanner, 'scan_for_req_type', request_type)
+            setattr(scanner, 'application_key', application_key)
+
+            if getattr(scanner, 'name') == 'AWSBedrockGuardrailScanner':
+                guardrail_id = auth_req.context.get('guardrail_id')
+                guardrail_version = auth_req.context.get('guardrail_version')
+                region = auth_req.context.get('region')
+
+                setattr(scanner, 'guardrail_id', guardrail_id)
+                setattr(scanner, 'guardrail_version', guardrail_version)
+                setattr(scanner, 'region', region)
+
         return scanners_list
 
-    def scan_messages(self, application_key: str, message: str, request_type: str, is_authz_scan: bool) -> (dict[str, ScannerResult], dict[str, str]):
+    def scan_messages(self, message: str, auth_req: AuthorizeRequest, is_authz_scan: bool) -> (dict[str, ScannerResult], dict[str, str]):
         """
         Scan the given messages for all the scanners where the enforce access control flag is true.
 
@@ -84,7 +97,9 @@ def scan_messages(self, application_key: str, message: str, request_type: str, i
             tuple: A tuple containing the scan results and the access control results.
         """
 
-        scanners = self.get_scanners(application_key, request_type, is_authz_scan)
+        application_key = auth_req.application_key
+        request_type = auth_req.request_type
+        scanners = self.get_scanners(application_key, request_type, is_authz_scan, auth_req)
         logger.debug(f"Found {len(scanners)} scanners for application key: {application_key}")
 
         scan_results, scan_timings = {}, {}
diff --git a/paig-server/backend/paig/api/shield/services/auth_service.py b/paig-server/backend/paig/api/shield/services/auth_service.py
index 30b03910..f524d532 100644
--- a/paig-server/backend/paig/api/shield/services/auth_service.py
+++ b/paig-server/backend/paig/api/shield/services/auth_service.py
@@ -28,6 +28,7 @@
 from api.shield.logfile.log_message_in_s3 import LogMessageInS3File
 from api.shield.logfile.log_message_in_local import LogMessageInLocal
 from paig_common.paig_exception import DiskFullException, AuditEventQueueFullException
+from api.shield.factory.governance_service_factory import GovernanceServiceFactory
 
 from core.utils import SingletonDepends
 
@@ -46,7 +47,8 @@ class AuthService:
 
     def __init__(self, authz_service_client_factory=SingletonDepends(AuthzServiceClientFactory),
                  data_store_controller: DataStoreController = SingletonDepends(DataStoreController),
-                 account_service_factory: AccountServiceFactory = SingletonDepends(AccountServiceFactory)):
+                 account_service_factory: AccountServiceFactory = SingletonDepends(AccountServiceFactory),
+                 governance_service_factory: GovernanceServiceFactory = SingletonDepends(GovernanceServiceFactory)):
         """
         Initializes the AuthService class with various dependencies and configuration settings.
 
@@ -62,6 +64,7 @@ def __init__(self, authz_service_client_factory=SingletonDepends(AuthzServiceCli
         self.application_manager = ApplicationManager()
         self.authz_service_client = authz_service_client_factory.get_authz_service_client()
         self.account_service_client = account_service_factory.get_account_service_client()
+        self.governance_service_client = governance_service_factory.get_governance_service_client()
         self.fluentd_logger_client = FluentdRestHttpClient()
         self.tenant_data_encryptor_service = TenantDataEncryptorService(self.account_service_client)
         self.presidio_anonymizer_engine = PresidioAnonymizerEngine()
@@ -102,6 +105,9 @@ async def authorize(self, auth_req: AuthorizeRequest):
         logger.debug("Auth Request After Decrypt : " + json_utils.mask_json_fields(json.dumps(auth_req.__dict__),
                                                                                    ['messages']))
         decrypt_time = f"{((time.perf_counter() - decrypt_start_time) * 1000):.3f}"
+        
+        gov_result = await self.governance_service_client.get_aws_bedrock_guardrail_info(auth_req.tenant_id, auth_req.application_key)
+        auth_req.context.update(gov_result)
 
         # loop through the messages in request to scan for traits
         message_analyze_start_time = time.perf_counter()
@@ -209,7 +215,7 @@ def analyze_scan_messages(self, access_control_traits, all_result_traits, analyz
         for request_text in auth_req.messages:
             # Analyze traits
             scanners_result, message_scan_timings = self.application_manager.scan_messages(
-                auth_req.application_key, request_text, auth_req.request_type, is_authz_scan)
+                request_text, auth_req, is_authz_scan)
             scan_timings = {scanner_name: f"{message_scan_time}ms" for scanner_name, message_scan_time in
                             message_scan_timings.items()}
 


Test patch:
diff --git a/paig-server/backend/paig/tests/api/router/governance/test_ai_app_routers.py b/paig-server/backend/paig/tests/api/router/governance/test_ai_app_routers.py
index 8f8469cf..a6999cf0 100644
--- a/paig-server/backend/paig/tests/api/router/governance/test_ai_app_routers.py
+++ b/paig-server/backend/paig/tests/api/router/governance/test_ai_app_routers.py
@@ -156,3 +156,45 @@ async def test_ai_application_crud_negative_operations(self, client: AsyncClient
         assert response.status_code == 400
         assert response.json()['message'] == "Vector DB not found with names: ['invalid_vector_db']"
 
+    @pytest.mark.asyncio
+    async def test_ai_application_crud_operations_with_guardrail_details(self, client: AsyncClient, app: FastAPI):
+        app.dependency_overrides[get_auth_user] = self.auth_user
+
+        new_ai_app_dict = self.ai_application_dict.copy()
+        new_ai_app_dict['guardrailDetails'] = "{\"guardrail_id\": \"guardrail1\",  \"guardrail_version\": \"DRAFT\",\"region\": \"us-east-1\"}"
+
+        await client.post(
+            f"{governance_services_base_route}/application", content=json.dumps(new_ai_app_dict)
+        )
+
+        response = await client.get(
+            f"{governance_services_base_route}/application"
+        )
+        assert response.status_code == 200
+        assert response.json()['content'][0]['guardrailDetails'] == new_ai_app_dict['guardrailDetails']
+        app_id = str(response.json()['content'][0]['id'])
+
+        response = await client.get(
+            f"{governance_services_base_route}/application/" + app_id
+        )
+        assert response.status_code == 200
+        assert response.json()['guardrailDetails'] == new_ai_app_dict['guardrailDetails']
+
+        update_req = {
+            "status": 1,
+            "name": "test_guardrail_app1",
+            "description": "test application1 updated",
+            "applicationKey": response.json()['applicationKey'],
+            "vector_dbs": [],
+            "guardrailDetails": "{\"guardrail_id\": \"guardrail2\",  \"guardrail_version\": \"DRAFT\",\"region\": \"us-east-1\"}"
+        }
+        response = await client.put(
+            f"{governance_services_base_route}/application/" + app_id, content=json.dumps(update_req)
+        )
+        assert response.status_code == 200
+        assert response.json()['guardrailDetails'] == "{\"guardrail_id\": \"guardrail2\",  \"guardrail_version\": \"DRAFT\",\"region\": \"us-east-1\"}"
+
+        response = await client.delete(
+            f"{governance_services_base_route}/application/" + app_id
+        )
+        assert response.status_code == 204
diff --git a/paig-server/backend/paig/tests/api/shield/test_application_manager_service.py b/paig-server/backend/paig/tests/api/shield/test_application_manager_service.py
index d1f35050..0d5d467b 100644
--- a/paig-server/backend/paig/tests/api/shield/test_application_manager_service.py
+++ b/paig-server/backend/paig/tests/api/shield/test_application_manager_service.py
@@ -1,11 +1,22 @@
 from unittest.mock import patch, MagicMock
 import pytest
+import json
+from pathlib import Path
 
 from api.shield.model.scanner_result import ScannerResult
+from api.shield.model.authorize_request import AuthorizeRequest
 from api.shield.scanners.BaseScanner import Scanner
 from api.shield.services.application_manager_service import ApplicationManager, scan_with_scanner
 from api.shield.utils.custom_exceptions import ShieldException
 
+def authorize_req_data():
+    json_file_path = f"{Path(__file__).parent}/json_data/authorize_request.json"
+    with open(json_file_path, 'r') as json_file:
+        req_json = json.load(json_file)
+
+    return AuthorizeRequest(tenant_id='test_tenant', req_data=req_json, user_role='OWNER')
+
+auth_req = authorize_req_data()
 
 class TestScanner1(Scanner):
     def __init__(self, **kwargs):
@@ -48,7 +59,7 @@ class TestApplicationManager:
     def test_get_scanners_with_cache_hit(self, mock_load_scanners):
         manager = ApplicationManager()
         manager.application_key_scanners.put('app_key', ['scanner1', 'scanner2'])
-        manager.get_scanners('app_key', 'prompt', True)
+        manager.get_scanners('app_key', 'prompt', True, auth_req)
         mock_load_scanners.assert_not_called()
 
     @patch('api.shield.services.application_manager_service.parse_properties')
@@ -56,8 +67,8 @@ def test_load_scanners(self, mock_parse_properties, mock_scanners):
         mock_parse_properties.return_value = mock_scanners
         manager = ApplicationManager()
         manager.load_scanners('app_key')
-        assert manager.get_scanners('app_key', 'prompt', True)[0] == mock_scanners[0]
-        assert manager.get_scanners('app_key', 'prompt', False)[0] == mock_scanners[1]
+        assert manager.get_scanners('app_key', 'prompt', True, auth_req)[0] == mock_scanners[0]
+        assert manager.get_scanners('app_key', 'prompt', False, auth_req)[0] == mock_scanners[1]
 
     @patch('api.shield.services.application_manager_service.parse_properties')
     def test_scan_messages(self, mock_parse_properties):
@@ -68,7 +79,7 @@ def test_scan_messages(self, mock_parse_properties):
         mock_parse_properties.return_value = [scanner1, scanner2]
         manager = ApplicationManager()
         manager.load_scanners('app_key')
-        scan_results, scan_timing = manager.scan_messages('app_key', 'message', 'prompt', True)
+        scan_results, scan_timing = manager.scan_messages('message', auth_req, True)
         assert len(scan_results) == 1
         for key, value in scan_results.items():
             assert value == {'traits': ['trait1', 'trait2']}
@@ -100,7 +111,7 @@ def test_scan_messages_success(app_manager, mock_scanners):
         traits=[],
         analyzer_result=["result2"]
     )):
-        scan_results, scan_timings = app_manager.scan_messages(application_key, message, 'prompt', True)
+        scan_results, scan_timings = app_manager.scan_messages(message, auth_req, True)
 
     # Verify the results
     assert len(scan_results) == 2
@@ -118,7 +129,7 @@ def test_scan_messages_with_exception(app_manager, mock_scanners):
         with patch.object(mock_scanners[1], 'scan',
                           return_value=ScannerResult(traits=["trait1"], analyzer_result=["result1"])):
             with pytest.raises(ShieldException) as e:
-                scan_results, access_control_traits = app_manager.scan_messages(application_key, message, 'prompt', True)
+                scan_results, access_control_traits = app_manager.scan_messages(message, auth_req, True)
 
                 # Verify the results
                 assert len(scan_results) == 1  # scanner1 failed, so only one result
diff --git a/paig-server/backend/paig/tests/api/shield/test_auth_service.py b/paig-server/backend/paig/tests/api/shield/test_auth_service.py
index fd8699f5..cdb2ac70 100644
--- a/paig-server/backend/paig/tests/api/shield/test_auth_service.py
+++ b/paig-server/backend/paig/tests/api/shield/test_auth_service.py
@@ -134,6 +134,8 @@ async def test_decrypt_authorize_request(self, mocker):
         mocker.patch.object(auth_service.application_manager, 'scan_messages', return_value=({}, {}))
         mocker.patch('api.shield.services.auth_service.AuthService.audit', return_value=(0, 0))
 
+        mocker.patch.object(auth_service.governance_service_client, 'get_aws_bedrock_guardrail_info', new_callable=AsyncMock, return_value={})
+
         # Create a mock AuthorizeRequest object
         mock_auth_req = authorize_req_data()
 
@@ -164,6 +166,8 @@ async def test_analyze_messages(self, mocker):
         mocker.patch.object(auth_service.application_manager, 'scan_messages', return_value=({}, {}))
         mocker.patch('api.shield.services.auth_service.AuthService.audit', return_value=(0, 0))
 
+        mocker.patch.object(auth_service.governance_service_client, 'get_aws_bedrock_guardrail_info', new_callable=AsyncMock, return_value={})
+
         # Create a mock AuthorizeRequest object
         mock_auth_req = authorize_req_data()
 
@@ -312,6 +316,8 @@ async def test_not_allowed_authz_response(self, mocker):
         mocker.patch.object(auth_service.application_manager, 'scan_messages', return_value=({}, {}))
         mocker.patch('api.shield.services.auth_service.AuthService.audit', return_value=(0, 0))
 
+        mocker.patch.object(auth_service.governance_service_client, 'get_aws_bedrock_guardrail_info', new_callable=AsyncMock, return_value={})
+
         # Call the authorize method
         auth_res = await auth_service.authorize(auth_req)
         # Assertions
@@ -593,6 +599,8 @@ async def test_authorize_authreq_with_stream_id(self, mocker):
         mocker.patch.object(auth_service.application_manager, 'scan_messages', return_value=({}, {}))
         mocker.patch('api.shield.services.auth_service.AuthService.audit', return_value=(0, 0))
 
+        mocker.patch.object(auth_service.governance_service_client, 'get_aws_bedrock_guardrail_info', new_callable=AsyncMock, return_value={})
+
         # Create a mock AuthorizeRequest object
         mock_auth_req = authorize_req_data_with_streamid()