Initial commit for RAG in py-ml

CHANGELOG.md

@@ -46,6 +46,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - updating listing file with three v2 sparse model - by @dhrubo-os ([#412](https://github.com/opensearch-project/opensearch-py-ml/pull/412))
 - Update model upload history -  opensearch-project/opensearch-neural-sparse-encoding-doc-v2-mini (v.1.0.0)(TORCH_SCRIPT) by @dhrubo-os ([#417](https://github.com/opensearch-project/opensearch-py-ml/pull/417))
 - Update model upload history -  opensearch-project/opensearch-neural-sparse-encoding-v2-distill (v.1.0.0)(TORCH_SCRIPT) by @dhrubo-os ([#419](https://github.com/opensearch-project/opensearch-py-ml/pull/419))
+- Added RAG functionality into `opensearch-py-ml` by @hmumtazz in ([#427](https://github.com/opensearch-project/opensearch-py-ml/pull/427))
 
 ### Fixed
 - Fix the wrong final zip file name in model_uploader workflow, now will name it by the upload_prefix alse.([#413](https://github.com/opensearch-project/opensearch-py-ml/pull/413/files))

opensearch_py_ml/ml_commons/IAMRoleHelper.py

@@ -0,0 +1,272 @@
+# SPDX-License-Identifier: Apache-2.0
+# The OpenSearch Contributors require contributions made to
+# this file be licensed under the Apache-2.0 license or a
+# compatible open source license.
+# Any modifications Copyright OpenSearch Contributors. See
+# GitHub history for details.
+
+import json
+import logging
+import uuid
+from datetime import datetime
+
+import boto3
+from botocore.exceptions import ClientError
+
+
+class IAMRoleHelper:
+    """
+    Helper class for managing IAM roles and their interactions with OpenSearch.
+    """
+
+    def __init__(
+        self,
+        region,
+        opensearch_domain_url=None,
+        opensearch_domain_username=None,
+        opensearch_domain_password=None,
+    ):
+        """
+        Initialize the IAMRoleHelper with AWS and OpenSearch configurations.
+
+        :param region: AWS region.
+        :param opensearch_domain_url: URL of the OpenSearch domain.
+        :param opensearch_domain_username: Username for OpenSearch domain authentication.
+        :param opensearch_domain_password: Password for OpenSearch domain authentication.
+        """
+        self.region = region
+        self.opensearch_domain_url = opensearch_domain_url
+        self.opensearch_domain_username = opensearch_domain_username
+        self.opensearch_domain_password = opensearch_domain_password
+
+        self.iam_client = boto3.client("iam", region_name=self.region)
+        self.sts_client = boto3.client("sts", region_name=self.region)
+
+    def role_exists(self, role_name):
+        """
+        Check if an IAM role exists.
+
+        :param role_name: Name of the IAM role.
+        :return: True if the role exists, False otherwise.
+        """
+        try:
+            self.iam_client.get_role(RoleName=role_name)
+            return True
+        except ClientError as e:
+            if e.response["Error"]["Code"] == "NoSuchEntity":
+                print(f"The requested role '{role_name}' does not exist.")
+            else:
+                print(f"An error occurred: {e}")
+            return False
+
+    def delete_role(self, role_name):
+        """
+        Delete an IAM role along with its attached policies.
+
+        :param role_name: Name of the IAM role to delete.
+        """
+        try:
+            # Detach any managed policies from the role
+            policies = self.iam_client.list_attached_role_policies(RoleName=role_name)[
+                "AttachedPolicies"
+            ]
+            for policy in policies:
+                self.iam_client.detach_role_policy(
+                    RoleName=role_name, PolicyArn=policy["PolicyArn"]
+                )
+            print(f"All managed policies detached from role '{role_name}'.")
+
+            # Delete inline policies associated with the role
+            inline_policies = self.iam_client.list_role_policies(RoleName=role_name)[
+                "PolicyNames"
+            ]
+            for policy_name in inline_policies:
+                self.iam_client.delete_role_policy(
+                    RoleName=role_name, PolicyName=policy_name
+                )
+            print(f"All inline policies deleted from role '{role_name}'.")
+
+            # Finally, delete the IAM role
+            self.iam_client.delete_role(RoleName=role_name)
+            print(f"Role '{role_name}' deleted.")
+
+        except ClientError as e:
+            if e.response["Error"]["Code"] == "NoSuchEntity":
+                print(f"Role '{role_name}' does not exist.")
+            else:
+                print(f"An error occurred: {e}")
+
+    def create_iam_role(
+        self,
+        role_name,
+        trust_policy_json,
+        inline_policy_json,
+        policy_name=None,
+    ):
+        """
+        Create a new IAM role with specified trust and inline policies.
+
+        :param role_name: Name of the IAM role to create.
+        :param trust_policy_json: Trust policy document in JSON format.
+        :param inline_policy_json: Inline policy document in JSON format.
+        :param policy_name: Optional. If not provided, a unique one will be generated.
+        :return: ARN of the created role or None if creation failed.
+        """
+        try:
+            # Create the role with the provided trust policy
+            create_role_response = self.iam_client.create_role(
+                RoleName=role_name,
+                AssumeRolePolicyDocument=json.dumps(trust_policy_json),
+                Description="Role with custom trust and inline policies",
+            )
+
+            # Retrieve the ARN of the newly created role
+            role_arn = create_role_response["Role"]["Arn"]
+
+            # If policy_name is not provided, generate a unique one
+            if not policy_name:
+                timestamp = datetime.utcnow().strftime("%Y%m%d%H%M%S")
+                policy_name = f"InlinePolicy-{role_name}-{timestamp}"
+
+            # Attach the inline policy to the role
+            self.iam_client.put_role_policy(
+                RoleName=role_name,
+                PolicyName=policy_name,
+                PolicyDocument=json.dumps(inline_policy_json),
+            )
+
+            print(f"Created role: {role_name} with inline policy: {policy_name}")
+            return role_arn
+
+        except ClientError as e:
+            print(f"Error creating the role: {e}")
+            return None
+
+    def get_role_info(self, role_name, include_details=False):
+        """
+        Retrieve information about an IAM role.
+
+        :param role_name: Name of the IAM role.
+        :param include_details: If False, returns only the role's ARN.
+                               If True, returns a dictionary with full role details.
+        :return: ARN or dict of role details. Returns None if not found.
+        """
+        if not role_name:
+            return None
+
+        try:
+            response = self.iam_client.get_role(RoleName=role_name)
+            role = response["Role"]
+            role_arn = role["Arn"]
+
+            if not include_details:
+                return role_arn
+
+            # Build a detailed dictionary
+            role_details = {
+                "RoleName": role["RoleName"],
+                "RoleId": role["RoleId"],
+                "Arn": role_arn,
+                "CreationDate": role["CreateDate"],
+                "AssumeRolePolicyDocument": role["AssumeRolePolicyDocument"],
+                "InlinePolicies": {},
+            }
+
+            # List and retrieve any inline policies
+            list_role_policies_response = self.iam_client.list_role_policies(
+                RoleName=role_name
+            )
+            for policy_name in list_role_policies_response["PolicyNames"]:
+                get_role_policy_response = self.iam_client.get_role_policy(
+                    RoleName=role_name, PolicyName=policy_name
+                )
+                role_details["InlinePolicies"][policy_name] = get_role_policy_response[
+                    "PolicyDocument"
+                ]
+
+            return role_details
+
+        except ClientError as e:
+            if e.response["Error"]["Code"] == "NoSuchEntity":
+                print(f"Role '{role_name}' does not exist.")
+            else:
+                print(f"An error occurred: {e}")
+            return None
+
+    def get_user_arn(self, username):
+        """
+        Retrieve the ARN of an IAM user.
+
+        :param username: Name of the IAM user.
+        :return: ARN of the user or None if not found.
+        """
+        if not username:
+            return None
+        try:
+            response = self.iam_client.get_user(UserName=username)
+            return response["User"]["Arn"]
+        except ClientError as e:
+            if e.response["Error"]["Code"] == "NoSuchEntity":
+                print(f"IAM user '{username}' not found.")
+            else:
+                print(f"An error occurred: {e}")
+            return None
+
+    def assume_role(self, role_arn, role_session_name=None, session=None):
+        """
+        Assume an IAM role and obtain temporary security credentials.
+
+        :param role_arn: ARN of the IAM role to assume.
+        :param role_session_name: Identifier for the assumed role session.
+        :param session: Optional boto3 session object. Defaults to the class-level sts_client.
+        :return: Dictionary with temporary security credentials and metadata, or None on failure.
+        """
+        if not role_arn:
+            logging.error("Role ARN is required.")
+            return None
+
+        # Use the provided session's STS client or fall back to the class-level sts_client
+        sts_client = session.client("sts") if session else self.sts_client
+
+        # Generate a default session name if none is provided
+        role_session_name = role_session_name or f"session-{uuid.uuid4()}"
+
+        try:
+            assumed_role_object = sts_client.assume_role(
+                RoleArn=role_arn,
+                RoleSessionName=role_session_name,
+            )
+
+            temp_credentials = assumed_role_object["Credentials"]
+            expiration = temp_credentials["Expiration"]
+
+            logging.info(
+                f"Assumed role: {role_arn}. Temporary credentials valid until: {expiration}"
+            )
+
+            return {
+                "credentials": {
+                    "AccessKeyId": temp_credentials["AccessKeyId"],
+                    "SecretAccessKey": temp_credentials["SecretAccessKey"],
+                    "SessionToken": temp_credentials["SessionToken"],
+                },
+                "expiration": expiration,
+                "session_name": role_session_name,
+            }
+
+        except ClientError as e:
+            error_code = e.response["Error"]["Code"]
+            logging.error(f"Error assuming role {role_arn}: {error_code} - {e}")
+            return None
+
+    def get_iam_user_name_from_arn(self, iam_principal_arn):
+        """
+        Extract the IAM user name from an IAM principal ARN.
+
+        :param iam_principal_arn: ARN of the IAM principal.
+        :return: IAM user name or None if extraction fails.
+        """
+        # IAM user ARN format: arn:aws:iam::123456789012:user/user-name
+        if iam_principal_arn and ":user/" in iam_principal_arn:
+            return iam_principal_arn.split(":user/")[-1]
+        return None

opensearch_py_ml/ml_commons/SecretsHelper.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# The OpenSearch Contributors require contributions made to
+# this file be licensed under the Apache-2.0 license or a
+# compatible open source license.
+# Any modifications Copyright OpenSearch Contributors. See
+# GitHub history for details.
+
+import json
+import logging
+
+import boto3
+from botocore.exceptions import ClientError
+
+# Configure the logger for this module
+logger = logging.getLogger(__name__)
+
+
+class SecretHelper:
+    """
+    Helper class for managing secrets in AWS Secrets Manager.
+    Provides methods to check existence, retrieve details, and create secrets.
+    """
+
+    def __init__(self, region: str):
+        """
+        Initialize the SecretHelper with the specified AWS region.
+        :param region: AWS region where the Secrets Manager is located.
+        """
+        self.region = region
+        # Create the Secrets Manager client once at the class level
+        self.secretsmanager = boto3.client("secretsmanager", region_name=self.region)
+
+    def secret_exists(self, secret_name: str) -> bool:
+        """
+        Check if a secret with the given name exists in AWS Secrets Manager.
+        :param secret_name: Name of the secret to check.
+        :return: True if the secret exists, False otherwise.
+        """
+        try:
+            # Attempt to retrieve the secret value
+            self.secretsmanager.get_secret_value(SecretId=secret_name)
+            return True
+        except ClientError as e:
+            # If the secret does not exist, return False
+            if e.response["Error"]["Code"] == "ResourceNotFoundException":
+                return False
+            else:
+                # Log other client errors and return False
+                logger.error(f"An error occurred: {e}")
+                return False
+
+    def get_secret_details(self, secret_name: str, fetch_value: bool = False) -> dict:
+        """
+        Retrieve details of a secret from AWS Secrets Manager.
+        Optionally fetch the secret value as well.
+
+        :param secret_name: Name of the secret.
+        :param fetch_value: Whether to also fetch the secret value (default is False).
+        :return: A dictionary with secret details (ARN and optionally the secret value)
+                 or an error dictionary if something went wrong.
+        """
+        try:
+            # Describe the secret to get its ARN and metadata
+            describe_response = self.secretsmanager.describe_secret(
+                SecretId=secret_name
+            )
+
+            secret_details = {
+                "ARN": describe_response["ARN"],
+                # You can add more fields from `describe_response` if needed
+            }
+
+            # Fetch the secret value if requested
+            if fetch_value:
+                value_response = self.secretsmanager.get_secret_value(
+                    SecretId=secret_name
+                )
+                secret_details["SecretValue"] = value_response.get("SecretString")
+
+            return secret_details
+
+        except ClientError as e:
+            error_code = e.response["Error"]["Code"]
+            if error_code == "ResourceNotFoundException":
+                logger.warning(f"The requested secret '{secret_name}' was not found")
+            else:
+                logger.error(
+                    f"An error occurred while fetching secret '{secret_name}': {e}"
+                )
+            # Return a dictionary with error details
+            return {"error": str(e), "error_code": error_code}
+
+    def create_secret(self, secret_name: str, secret_value: dict) -> str:
+        """
+        Create a new secret in AWS Secrets Manager.
+        :param secret_name: Name of the secret to create.
+        :param secret_value: Dictionary containing the secret data.
+        :return: ARN of the created secret if successful, None otherwise.
+        """
+        try:
+            # Create the secret with the provided name and value
+            response = self.secretsmanager.create_secret(
+                Name=secret_name,
+                SecretString=json.dumps(secret_value),
+            )
+            # Log success and return the secret's ARN
+            logger.info(f"Secret '{secret_name}' created successfully.")
+            return response["ARN"]
+        except ClientError as e:
+            # Log errors during secret creation and return None
+            logger.error(f"Error creating secret '{secret_name}': {e}")
+            return None