ibm-watsonx-orchestrate 1.5.1__py3-none-any.whl → 1.7.0a0__py3-none-any.whl

ibm_watsonx_orchestrate/cli/commands/connections/connections_controller.py

@@ -15,19 +15,19 @@ from ibm_watsonx_orchestrate.agent_builder.connections.types import (
     ConnectionType,
     IdpConfigData,
     IdpConfigDataBody,
-    AppConfigData, 
+    AppConfigData,
     BasicAuthCredentials,
     BearerTokenAuthCredentials,
     APIKeyAuthCredentials,
     # OAuth2AuthCodeCredentials,
-    # OAuth2ClientCredentials,
+    OAuth2ClientCredentials,
     # OAuth2ImplicitCredentials,
     # OAuth2PasswordCredentials,
     OAuthOnBehalfOfCredentials,
     KeyValueConnectionCredentials,
     CREDENTIALS,
     IdentityProviderCredentials,
-    OAUTH_CONNECTION_TYPES
+    OAUTH_CONNECTION_TYPES, OAuth2AuthCodeCredentials
 
 )
 
@@ -80,7 +80,6 @@ def _create_connection_from_spec(content: dict) -> None:
         config = environments.get(environment)
         config["environment"] = environment
         config["app_id"] = app_id
-        config["environment"] = environment
         config = ConnectionConfiguration.model_validate(config)
         add_configuration(config)
     
@@ -137,40 +136,31 @@ def _validate_connection_params(type: ConnectionType, **args) -> None:
             f"Missing flags --api-key is required for type {type}"
         )
 
-    # if type in OAUTH_CONNECTION_TYPES and args.get('client_id') is None:
-    #     raise typer.BadParameter(
-    #         f"Missing flags --client-id is required for type {type}"
-    #     )
-
-    # if type in (OAUTH_CONNECTION_TYPES.difference({ConnectionType.OAUTH2_IMPLICIT, ConnectionType.OAUTH_ON_BEHALF_OF_FLOW})) and args.get('client_secret') is None:
-    #     raise typer.BadParameter(
-    #         f"Missing flags --client-secret is required for type {type}"
-    #     )
-    
-    # if type in (OAUTH_CONNECTION_TYPES.difference({ConnectionType.OAUTH2_IMPLICIT})) and args.get('token_url') is None:
-    #     raise typer.BadParameter(
-    #         f"Missing flags --token-url is required for type {type}"
-    #     )
+    if type in {ConnectionType.OAUTH2_CLIENT_CREDS, ConnectionType.OAUTH2_AUTH_CODE} and args.get('client_secret') is None:
+        raise typer.BadParameter(
+            f"Missing flags --client-secret is required for type {type}"
+        )
     
-    # if type in (OAUTH_CONNECTION_TYPES.difference({ConnectionType.OAUTH2_CLIENT_CREDS, ConnectionType.OAUTH_ON_BEHALF_OF_FLOW})) and args.get('auth_url') is None:
-    #     raise typer.BadParameter(
-    #         f"Missing flags --auth-url is required for type {type}"
-    #     )
+    if type in {ConnectionType.OAUTH2_AUTH_CODE} and args.get('auth_url') is None:
+        raise typer.BadParameter(
+            f"Missing flags --auth-url is required for type {type}"
+        )
 
-    if type == ConnectionType.OAUTH_ON_BEHALF_OF_FLOW and (
+    if type in {ConnectionType.OAUTH_ON_BEHALF_OF_FLOW, ConnectionType.OAUTH2_CLIENT_CREDS, ConnectionType.OAUTH2_AUTH_CODE} and (
             args.get('client_id') is None
     ):
         raise typer.BadParameter(
             f"Missing flags --client-id is required for type {type}"
         )
     
-    if type == ConnectionType.OAUTH_ON_BEHALF_OF_FLOW and (
+    if type in {ConnectionType.OAUTH_ON_BEHALF_OF_FLOW, ConnectionType.OAUTH2_CLIENT_CREDS, ConnectionType.OAUTH2_AUTH_CODE} and (
             args.get('token_url') is None
     ):
         raise typer.BadParameter(
             f"Missing flags --token-url is required for type {type}"
         )
 
+
     if type == ConnectionType.OAUTH_ON_BEHALF_OF_FLOW and (
             args.get('grant_type') is None
     ):
@@ -201,19 +191,21 @@ def _get_credentials(type: ConnectionType, **kwargs):
             return APIKeyAuthCredentials(
                 api_key=kwargs.get("api_key")
             )
-        # case ConnectionType.OAUTH2_AUTH_CODE:
-        #     return OAuth2AuthCodeCredentials(
-        #         authorization_url=kwargs.get("auth_url"),
-        #         client_id=kwargs.get("client_id"),
-        #         client_secret=kwargs.get("client_secret"),
-        #         token_url=kwargs.get("token_url")
-        #     )
-        # case ConnectionType.OAUTH2_CLIENT_CREDS:
-        #     return OAuth2ClientCredentials(
-        #         client_id=kwargs.get("client_id"),
-        #         client_secret=kwargs.get("client_secret"),
-        #         token_url=kwargs.get("token_url")
-        #     )
+        case ConnectionType.OAUTH2_AUTH_CODE:
+            return OAuth2AuthCodeCredentials(
+                authorization_url=kwargs.get("auth_url"),
+                client_id=kwargs.get("client_id"),
+                client_secret=kwargs.get("client_secret"),
+                token_url=kwargs.get("token_url"),
+                scopes=kwargs.get("scopes")
+            )
+        case ConnectionType.OAUTH2_CLIENT_CREDS:
+            return OAuth2ClientCredentials(
+                client_id=kwargs.get("client_id"),
+                client_secret=kwargs.get("client_secret"),
+                token_url=kwargs.get("token_url"),
+                scopes=kwargs.get("scopes")
+            )
         # case ConnectionType.OAUTH2_IMPLICIT:
         #     return OAuth2ImplicitCredentials(
         #         authorization_url=kwargs.get("auth_url"),
@@ -268,11 +260,14 @@ def add_configuration(config: ConnectionConfiguration) -> None:
                 logger.warning(f"Detected a change in sso from '{existing_configuration.sso}' to '{config.sso}'. The associated credentials will be removed.")
                 should_delete_credentials = True
 
+            existing_conn_type = get_connection_type(security_scheme=existing_configuration.security_scheme, auth_type=existing_configuration.auth_type)
+            use_app_credentials = existing_conn_type in OAUTH_CONNECTION_TYPES
+
             if should_delete_credentials:
                 try:
-                    existing_credentials = client.get_credentials(app_id=app_id, env=environment, use_sso=existing_configuration.sso)
+                    existing_credentials = client.get_credentials(app_id=app_id, env=environment, use_app_credentials=use_app_credentials)
                     if existing_credentials:
-                        client.delete_credentials(app_id=app_id, env=environment, use_sso=existing_configuration.sso)
+                        client.delete_credentials(app_id=app_id, env=environment, use_app_credentials=use_app_credentials)
                 except:
                     logger.error(f"Error removing credentials for connection '{app_id}' in environment '{environment}'. No changes have been made to the configuration.")
                     sys.exit(1)
@@ -290,11 +285,11 @@ def add_configuration(config: ConnectionConfiguration) -> None:
         logger.error(response_text)
         exit(1)
 
-def add_credentials(app_id: str, environment: ConnectionEnvironment, use_sso: bool, credentials: CREDENTIALS) -> None:
+def add_credentials(app_id: str, environment: ConnectionEnvironment, use_app_credentials: bool, credentials: CREDENTIALS) -> None:
     client = get_connections_client()
     try:
-        existing_credentials = client.get_credentials(app_id=app_id, env=environment, use_sso=use_sso)
-        if use_sso:
+        existing_credentials = client.get_credentials(app_id=app_id, env=environment, use_app_credentials=use_app_credentials)
+        if use_app_credentials:
             payload = {
                 "app_credentials": credentials.model_dump(exclude_none=True)
             }
@@ -305,9 +300,9 @@ def add_credentials(app_id: str, environment: ConnectionEnvironment, use_sso: bo
         
         logger.info(f"Setting credentials for environment '{environment}' on connection '{app_id}'")
         if existing_credentials:
-            client.update_credentials(app_id=app_id, env=environment, use_sso=use_sso, payload=payload)
+            client.update_credentials(app_id=app_id, env=environment, use_app_credentials=use_app_credentials, payload=payload)
         else:
-            client.create_credentials(app_id=app_id,env=environment, use_sso=use_sso, payload=payload)
+            client.create_credentials(app_id=app_id,env=environment, use_app_credentials=use_app_credentials, payload=payload)
         logger.info(f"Credentials successfully set for '{environment}' environment of connection '{app_id}'")
     except requests.HTTPError as e:
         response = e.response
@@ -319,7 +314,7 @@ def add_identity_provider(app_id: str, environment: ConnectionEnvironment, idp:
     client = get_connections_client()
 
     try:
-        existing_credentials = client.get_credentials(app_id=app_id, env=environment, use_sso=True)
+        existing_credentials = client.get_credentials(app_id=app_id, env=environment, use_app_credentials=True)
         
         payload = {
             "idp_credentials": idp.model_dump()
@@ -327,9 +322,9 @@ def add_identity_provider(app_id: str, environment: ConnectionEnvironment, idp:
 
         logger.info(f"Setting identity provider for environment '{environment}' on connection '{app_id}'")
         if existing_credentials:
-            client.update_credentials(app_id=app_id, env=environment, use_sso=True, payload=payload)
+            client.update_credentials(app_id=app_id, env=environment, use_app_credentials=True, payload=payload)
         else:
-            client.create_credentials(app_id=app_id,env=environment, use_sso=True, payload=payload)
+            client.create_credentials(app_id=app_id,env=environment, use_app_credentials=True, payload=payload)
         logger.info(f"Identity provider successfully set for '{environment}' environment of connection '{app_id}'")
     except requests.HTTPError as e:
         response = e.response
@@ -388,11 +383,17 @@ def list_connections(environment: ConnectionEnvironment | None, verbose: bool =
         non_configured_table = rich.table.Table(show_header=True, header_style="bold white", show_lines=True, title="*Non-Configured")
         draft_table = rich.table.Table(show_header=True, header_style="bold white", show_lines=True, title="Draft")
         live_table = rich.table.Table(show_header=True, header_style="bold white", show_lines=True, title="Live")
-        columns = ["App ID", "Auth Type", "Type", "Credentials Set"]
-        for column in columns:
-            draft_table.add_column(column, justify='center', no_wrap=True)
-            live_table.add_column(column, justify='center', no_wrap=True)
-            non_configured_table.add_column(column, justify='center', no_wrap=True)
+        default_args = {"justify": "center", "no_wrap": True}
+        column_args = {
+            "App ID": {"overflow": "fold"}, 
+            "Auth Type": {}, 
+            "Type": {}, 
+            "Credentials Set/ Connected": {}
+        }
+        for column in column_args:
+            draft_table.add_column(column,**default_args, **column_args[column])
+            live_table.add_column(column,**default_args, **column_args[column])
+            non_configured_table.add_column(column,**default_args, **column_args[column])
         
         for conn in connections:
             if conn.environment is None:
@@ -437,6 +438,7 @@ def configure_connection(**kwargs) -> None:
         logger.error(f"Cannot create configuration for environment '{kwargs.get('environment')}'. Local development does not support any environments other than 'draft'.")
         sys.exit(1)
 
+    
     idp_config_body = None
     if kwargs.get("idp_token_type") or kwargs.get("idp_token_use"):
         idp_config_body = IdpConfigDataBody(
@@ -463,11 +465,6 @@ def configure_connection(**kwargs) -> None:
 
     config = ConnectionConfiguration.model_validate(kwargs)
 
-    # TODO: Remove once Oauth is supported on local
-    if config.security_scheme == ConnectionSecurityScheme.OAUTH2 and is_local_dev():
-        logger.error("Use of OAuth connections unsupported for local development at this time.")
-        sys.exit(1)
-
     add_configuration(config)
 
 def set_credentials_connection(
@@ -484,11 +481,12 @@ def set_credentials_connection(
 
     sso_enabled = config.sso
     conn_type = get_connection_type(security_scheme=config.security_scheme, auth_type=config.auth_type)
+    use_app_credentials = conn_type in OAUTH_CONNECTION_TYPES
 
     _validate_connection_params(type=conn_type, **kwargs)
     credentials = _get_credentials(type=conn_type, **kwargs)
 
-    add_credentials(app_id=app_id, environment=environment, use_sso=sso_enabled, credentials=credentials)
+    add_credentials(app_id=app_id, environment=environment, use_app_credentials=use_app_credentials, credentials=credentials)
 
 def set_identity_provider_connection(
     app_id: str,

ibm_watsonx_orchestrate/cli/commands/environment/environment_command.py

@@ -5,6 +5,7 @@ from ibm_watsonx_orchestrate.cli.commands.environment import environment_control
 from ibm_watsonx_orchestrate.cli.commands.environment.types import EnvironmentAuthType
 from ibm_watsonx_orchestrate.cli.commands.tools.types import RegistryType
 from ibm_watsonx_orchestrate.client.utils import is_local_dev
+import sys
 
 logger = logging.getLogger(__name__)
 
@@ -20,7 +21,19 @@ def activate_env(
         apikey: Annotated[
             str,
             typer.Option(
-                "--apikey", "-a", help="WXO API Key. Leave Blank if developing locally"
+                "--api-key", "-a", help="WXO or CPD API Key. Leave Blank if developing locally. For CPD, either a Passoword or Apikey is accepted for CPD, but not both."
+            ),
+        ] = None,
+        username: Annotated[
+            str,
+            typer.Option(
+                "--username", "-u", help="Username specifically for CPD Environments."
+            ),
+        ] = None,
+        password: Annotated[
+            str,
+            typer.Option(
+                "--password", "-p", help="Password specifically for CPD Environments. Either a Passoword or Apikey is accepted for CPD, but not both."
             ),
         ] = None,
         registry: Annotated[
@@ -32,7 +45,7 @@ def activate_env(
             typer.Option("--test-package-version-override", help="Which prereleased package version to reference when using --registry testpypi", hidden=True),
         ] = None
 ):
-    environment_controller.activate(name=name, apikey=apikey, registry=registry, test_package_version_override=test_package_version_override)
+    environment_controller.activate(name=name, apikey=apikey,username=username, password=password, registry=registry, test_package_version_override=test_package_version_override)
 
 
 @environment_app.command(name="add")
@@ -58,9 +71,21 @@ def add_env(
         type: Annotated[
             EnvironmentAuthType,
             typer.Option("--type", "-t", help="The type of auth you wish to use"),
-        ] = None
+        ] = None,
+        insecure: Annotated[
+            bool,
+            typer.Option("--insecure", help="Ignore SSL validation errors. Used for CPD Environments only"),
+        ] = False,
+        verify: Annotated[
+            str,
+            typer.Option("--verify", help="Path to SSL Cert. Used for CPD Environments only"),
+        ] = None,
 ):
-    environment_controller.add(name=name, url=url, should_activate=activate, iam_url=iam_url, type=type)
+    if insecure and verify:
+        logger.error("Please choose either '--insecure' or '--verify' but not both.")
+        sys.exit(1)
+
+    environment_controller.add(name=name, url=url, should_activate=activate, iam_url=iam_url, type=type, insecure=insecure, verify=verify)
 
 
 @environment_app.command(name="remove")

ibm_watsonx_orchestrate/cli/commands/environment/environment_controller.py

@@ -2,6 +2,7 @@ import logging
 import rich
 import jwt
 import getpass
+import sys
 
 from ibm_watsonx_orchestrate.cli.commands.tools.types import RegistryType
 from ibm_watsonx_orchestrate.cli.config import (
@@ -17,14 +18,15 @@ from ibm_watsonx_orchestrate.cli.config import (
     ENV_IAM_URL_OPT,
     ENVIRONMENTS_SECTION_HEADER,
     PROTECTED_ENV_NAME,
-    ENV_AUTH_TYPE, PYTHON_REGISTRY_HEADER, PYTHON_REGISTRY_TYPE_OPT, PYTHON_REGISTRY_TEST_PACKAGE_VERSION_OVERRIDE_OPT,
+    ENV_AUTH_TYPE, PYTHON_REGISTRY_HEADER, PYTHON_REGISTRY_TYPE_OPT, PYTHON_REGISTRY_TEST_PACKAGE_VERSION_OVERRIDE_OPT, BYPASS_SSL, VERIFY,
     DEFAULT_CONFIG_FILE_CONTENT
 )
 from ibm_watsonx_orchestrate.client.client import Client
 from ibm_watsonx_orchestrate.client.client_errors import ClientError
+from ibm_watsonx_orchestrate.client.agents.agent_client import AgentClient, ClientAPIException
 from ibm_watsonx_orchestrate.client.credentials import Credentials
 from threading import Lock
-from ibm_watsonx_orchestrate.client.utils import is_local_dev, check_token_validity
+from ibm_watsonx_orchestrate.client.utils import is_local_dev, check_token_validity, is_cpd_env
 from ibm_watsonx_orchestrate.cli.commands.environment.types import EnvironmentAuthType
 
 logger = logging.getLogger(__name__)
@@ -42,7 +44,33 @@ def _decode_token(token: str, is_local: bool = False) -> dict:
         logger.error("Invalid token format")
         raise e
 
-def _login(name: str, apikey: str = None) -> None:
+
+def _validate_token_functionality(token: str, url: str) -> None:
+    '''
+    Validates a token by making a request to GET /agents
+
+    Args: 
+        token: A JWT token
+        url: WXO instance URL
+    '''
+    is_cpd = is_cpd_env(url)
+    if is_cpd is True:
+        agent_client = AgentClient(base_url=url, api_key=token, is_local=is_local_dev(url), verify=False)
+    else:
+        agent_client = AgentClient(base_url=url, api_key=token, is_local=is_local_dev(url))
+    agent_client.api_key = token
+
+    try:
+        agent_client.get()
+    except ClientAPIException as e:
+        if e.response.status_code >= 400:
+            reason = e.response.reason
+            logger.error(f"Failed to authenticate to provided instance '{url}'. Reason: '{reason}'. Please ensure provider URL and API key are valid.")
+            sys.exit(1)
+        raise e
+
+
+def _login(name: str, apikey: str = None, username: str = None, password: str = None) -> None:
     cfg = Config()
     auth_cfg = Config(AUTH_CONFIG_FILE_FOLDER, AUTH_CONFIG_FILE)
 
@@ -55,14 +83,43 @@ def _login(name: str, apikey: str = None) -> None:
     except (KeyError, AttributeError):
         auth_type = None
 
+    username = username
+    apikey = apikey
+    password = password
+
+    if is_cpd_env(url):
+        if username is None:
+            username = getpass.getpass("Please enter CPD Username: ")
+
+        if not apikey and not password:
+            apikey = getpass.getpass("Enter CPD API key (or leave blank to use password): ")
+        if not apikey and not password:
+            password = getpass.getpass("Enter CPD password (or leave blank if you used API key): ")
+
+        if apikey and password:
+            logger.error("For CPD, please use either an Apikey or a Password but not both.")
+            sys.exit(1)
+
+        if not apikey and not password:
+            logger.error("For CPD, you must provide either an API key or a password.")
+            sys.exit(1)
+    
 
-    if apikey is None and not is_local:
+    if not apikey and not password and not is_local and auth_type != "cpd":
         apikey = getpass.getpass("Please enter WXO API key: ")
 
     try:
-        creds = Credentials(url=url, api_key=apikey, iam_url=iam_url, auth_type=auth_type)
+        creds = Credentials(
+            url=url, 
+            api_key=apikey, 
+            username=username, 
+            password=password, 
+            iam_url=iam_url, 
+            auth_type=auth_type
+        )
         client = Client(creds)
         token = _decode_token(client.token, is_local)
+        _validate_token_functionality(token=token.get(AUTH_MCSP_TOKEN_OPT), url=url)
         with lock:
             auth_cfg.save(
                 {
@@ -74,7 +131,7 @@ def _login(name: str, apikey: str = None) -> None:
     except ClientError as e:
         raise ClientError(e)
 
-def activate(name: str, apikey: str=None, registry: RegistryType=None, test_package_version_override=None) -> None:
+def activate(name: str, apikey: str=None, username: str=None, password: str=None, registry: RegistryType=None, test_package_version_override=None) -> None:
     cfg = Config()
     auth_cfg = Config(AUTH_CONFIG_FILE_FOLDER, AUTH_CONFIG_FILE)
     env_cfg = cfg.read(ENVIRONMENTS_SECTION_HEADER, name)
@@ -92,7 +149,7 @@ def activate(name: str, apikey: str=None, registry: RegistryType=None, test_pack
     existing_token = existing_auth_config.get(AUTH_MCSP_TOKEN_OPT) if existing_auth_config else None
 
     if not check_token_validity(existing_token) or is_local:
-        _login(name=name, apikey=apikey)
+        _login(name=name, apikey=apikey, username=username, password=password)
 
     with lock:
         cfg.write(CONTEXT_SECTION_HEADER, CONTEXT_ACTIVE_ENV_OPT, name)
@@ -104,8 +161,11 @@ def activate(name: str, apikey: str=None, registry: RegistryType=None, test_pack
             cfg.write(PYTHON_REGISTRY_HEADER, PYTHON_REGISTRY_TEST_PACKAGE_VERSION_OVERRIDE_OPT, test_package_version_override)
 
     logger.info(f"Environment '{name}' is now active")
+    is_cpd = is_cpd_env(url)
+    if is_cpd:
+        logger.warning("Support for CPD clusters is currently an early access preview")
 
-def add(name: str, url: str, should_activate: bool=False, iam_url: str=None, type: EnvironmentAuthType=None) -> None:
+def add(name: str, url: str, should_activate: bool=False, iam_url: str=None, type: EnvironmentAuthType=None, insecure: bool=None, verify: str=None) -> None:
     if name == PROTECTED_ENV_NAME:
         logger.error(f"The name '{PROTECTED_ENV_NAME}' is a reserved environment name. Please select a diffrent name or use `orchestrate env activate {PROTECTED_ENV_NAME}` to swap to '{PROTECTED_ENV_NAME}'")
         return
@@ -124,6 +184,12 @@ def add(name: str, url: str, should_activate: bool=False, iam_url: str=None, typ
             cfg.write(ENVIRONMENTS_SECTION_HEADER, name, {ENV_IAM_URL_OPT: iam_url})
         if type:
             cfg.write(ENVIRONMENTS_SECTION_HEADER, name, {ENV_AUTH_TYPE: str(type)})
+        if insecure:
+            cfg.write(ENVIRONMENTS_SECTION_HEADER, name, {BYPASS_SSL: insecure})
+            cfg.write(ENVIRONMENTS_SECTION_HEADER, name, {VERIFY: 'None'})
+        if verify:
+            cfg.write(ENVIRONMENTS_SECTION_HEADER, name, {VERIFY: verify})
+            cfg.write(ENVIRONMENTS_SECTION_HEADER, name, {BYPASS_SSL: False})
         
 
     logger.info(f"Environment '{name}' has been created")

ibm_watsonx_orchestrate/cli/commands/environment/types.py

@@ -4,6 +4,7 @@ from enum import Enum
 class EnvironmentAuthType(str, Enum):
     IBM_CLOUD_IAM = 'ibm_iam'
     MCSP = 'mcsp'
+    CPD = 'cpd'
 
     def __str__(self):
         return self.value