Source code for aidputils.agents.tools.utils

from typing import Optional
import os
import oci
import requests
import logging
from tenacity import (
    retry,
    wait_exponential,
    stop_after_attempt,
    retry_if_exception_type,
)

from aidputils.agents.auth.signer.custom_remote_signer import CustomRemoteSigner

logger = logging.getLogger(__name__)




def call_tool_by_class(tool_class_name: str, tool_conf: dict, runtime_params: dict, **context_vars):
    from aidputils.agents.tools.base_tool import BaseTool
    tool_cls = BaseTool.tool_class_registry.get(tool_class_name)

    if tool_cls:
        response = tool_cls.invoke(tool_conf, runtime_params, **context_vars)
        return response
    else:
        return f"Tool class '{tool_class_name}' not registered."





def call_tool_by_name(tool_name: str, tool_registry, runtime_params: dict, **context_vars):
    tool_info = tool_registry.get(tool_name)
    if not tool_info:
        logger.error(f"Tool '{tool_name}' not found in registry.")
        return

    tool_class_name = tool_info["class"]
    tool_conf = tool_info.get("conf", {})
    return call_tool_by_class(tool_class_name, tool_conf, runtime_params, **context_vars)



class Constants:
    AUTH_PROFILE = "auth_profile"
    AUTH_PROFILE_DEFAULT = "DEFAULT"

    # Signer types
    AUTH_TYPE = "auth_type"
    AUTH_TYPE_SECURITY_TOKEN = "security_token"
    AUTH_TYPE_INSTANCE_PRINCIPAL = "instance_principal"
    AUTH_TYPE_RESOURCE_PRINCIPAL = "resource_principal"
    AUTH_TYPE_REMOTE = "remote"

    # Connection properties
    CONNECTION_PROPS_KEY = "connectionProperties"
    USER_NAME = "user.name"
    USER_CREDENTIAL = "password"
    TNS = "tns"
    WALLET_CONTENT = "wallet.content"
    WALLET_CREDENTIAL = "wallet.password"
    WALLET_EXTRACT_DIR = "wallet_extract"
    WALLET_LOC_KEY = "wallet.zip"

    # Header keys and values
    ACCEPT_JSON = "application/json"
    CONTENT_TYPE_JSON = "application/json"
    DH_USER_PRINCIPAL_KEY = "dh-user-principal"
    DH_USER_PRINCIPAL_ENV_KEY = "DH_USER_PRINCIPAL"

    ADW_23_AI = "ADW_23_AI"
    GEN_AI = "GEN_AI"




class SystemUtils:


    @classmethod
    def get_env_var(cls, var_name: str, default: Optional[str] = None) -> Optional[str]:
        return os.getenv(var_name, default)




    @classmethod
    def make_security_token_signer(cls, oci_config):
        pk = oci.signer.load_private_key_from_file(oci_config.get("key_file"), None)
        with open(oci_config.get("security_token_file")) as f:
            st_string = f.read()
        return oci.auth.signers.SecurityTokenSigner(st_string, pk)




    @classmethod
    def get_signer(
            cls,
            signer_type=Constants.AUTH_TYPE_INSTANCE_PRINCIPAL,
            config_path=None,
            config_profile=Constants.AUTH_PROFILE_DEFAULT,
            **context_vars
    ):
        if signer_type == Constants.AUTH_TYPE_SECURITY_TOKEN:
            if config_path is None:
                config_path = os.getenv(
                    "OCI_CONFIG_FILE", os.path.expanduser("~/.oci/config")
                )
            config = oci.config.from_file(config_path, config_profile)
            return cls.make_security_token_signer(config)
        elif signer_type == Constants.AUTH_TYPE_INSTANCE_PRINCIPAL:
            return oci.auth.signers.InstancePrincipalsSecurityTokenSigner()
        elif signer_type == Constants.AUTH_TYPE_RESOURCE_PRINCIPAL:
            return oci.auth.signers.get_resource_principals_signer()
        elif signer_type == Constants.AUTH_TYPE_REMOTE:
            return CustomRemoteSigner()
        else:
            raise ValueError(
                "Invalid signer_type. Must be  one of security_token, instance, remote, instance_principal, resource_principal")






class HttpUtil:
    @staticmethod
    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=1, max=10),
        retry=retry_if_exception_type(
            (requests.exceptions.RequestException, requests.exceptions.HTTPError)
        ),
    )
    def get_request(endpoint, signer=None, headers=None):
        resp = requests.get(endpoint, headers=headers, auth=signer)
        if resp.status_code >= 400:
            logger.error(f"HTTP Error: {resp.status_code} {resp.headers} {resp.text}")
            raise requests.exceptions.HTTPError(f"Status code: {resp.status_code}")
        return resp.json()

    @staticmethod
    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=1, max=10),
        retry=retry_if_exception_type(
            (requests.exceptions.RequestException, requests.exceptions.HTTPError)
        ),
    )
    def post_request(endpoint, data=None, signer=None, headers=None):
        resp = requests.post(endpoint, json=data, headers=headers, auth=signer)
        if resp.status_code >= 400:
            logger.error(f"HTTP Error: {resp.status_code} {resp.text}")
            raise requests.exceptions.HTTPError(f"Status code: {resp.status_code}")
        return resp.json()

    @staticmethod
    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=1, max=10),
        retry=retry_if_exception_type(
            (requests.exceptions.RequestException, requests.exceptions.HTTPError)
        ),
    )
    def put_request(endpoint, data=None, signer=None, headers=None):
        resp = requests.put(endpoint, json=data, headers=headers, auth=signer)
        if resp.status_code >= 400:
            logger.error(f"HTTP Error: {resp.status_code} {resp.text}")
            raise requests.exceptions.HTTPError(f"Status code: {resp.status_code}")
        return resp.json()

    @staticmethod
    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=1, max=10),
        retry=retry_if_exception_type(
            (requests.exceptions.RequestException, requests.exceptions.HTTPError)
        ),
    )
    def delete_request(endpoint, signer=None, headers=None):
        resp = requests.delete(endpoint, headers=headers, auth=signer)
        if resp.status_code >= 400:
            logger.error(f"HTTP Error: {resp.status_code} {resp.text}")
            raise requests.exceptions.HTTPError(f"Status code: {resp.status_code}")
        return resp.json()