ibm-watsonx-orchestrate 1.0.0__py3-none-any.whl

ibm_watsonx_orchestrate/client/knowledge_bases/knowledge_base_client.py

@@ -0,0 +1,46 @@
+from ibm_watsonx_orchestrate.client.base_api_client import BaseAPIClient
+import json
+from typing_extensions import List
+from ibm_watsonx_orchestrate.client.base_api_client import ClientAPIException
+
+
+class KnowledgeBaseClient(BaseAPIClient):
+    """
+    Client to handle CRUD operations for Native Knowledge Base endpoint
+    """
+
+    def create(self, payload: dict) -> dict:
+        return self._post_form_data("/api/v1/orchestrate/knowledge-bases/documents", data={ "knowledge_base" : json.dumps(payload) })
+    
+    def create_built_in(self, payload: dict, files: list) -> dict:
+        return self._post_form_data("/api/v1/orchestrate/knowledge-bases/documents", data={ "knowledge_base" : json.dumps(payload) }, files=files)
+
+    def get(self) -> dict:
+        return self._get("/api/v1/orchestrate/knowledge-bases")
+    
+    def get_by_name(self, name: str) -> List[dict]:
+        kbs = self.get_by_names([name])
+        return None if len(kbs) == 0 else kbs[0]
+    
+    def get_by_id(self, knowledge_base_id: str) -> dict:
+        return self._get(f"/api/v1/orchestrate/knowledge-bases/{knowledge_base_id}")
+
+    def get_by_names(self, name: List[str]) -> List[dict]:
+        formatted_names = [f"names={x}" for x in name]
+        return self._get(f"/api/v1/orchestrate/knowledge-bases?{'&'.join(formatted_names)}")
+    
+    def status(self, knowledge_base_id: str) -> dict:
+        return self._get(f"/api/v1/orchestrate/knowledge-bases/{knowledge_base_id}/status")
+
+    def update(self, knowledge_base_id: str, payload: dict) -> dict:
+        return self._patch_form_data(f"/api/v1/orchestrate/knowledge-bases/{knowledge_base_id}/documents", data={ "knowledge_base" : json.dumps(payload) })
+    
+    def update_with_documents(self, knowledge_base_id: str, payload: dict, files: list) -> dict:
+        return self._patch_form_data(f"/api/v1/orchestrate/knowledge-bases/{knowledge_base_id}/documents", data={ "knowledge_base" : json.dumps(payload) }, files=files)
+
+    def delete(self, knowledge_base_id: str,) -> dict:
+        return self._delete(f"/api/v1/orchestrate/knowledge-bases/{knowledge_base_id}")
+        
+    
+
+    

ibm_watsonx_orchestrate/client/local_service_instance.py

@@ -0,0 +1,91 @@
+from ibm_watsonx_orchestrate.client.base_service_instance import BaseServiceInstance
+import logging
+import requests
+from ibm_watsonx_orchestrate.client.credentials import Credentials
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_TENANT = {
+    "name": "wxo-dev",
+    "title": "WatsonX Orchestrate Development",
+    "tags": ["test"]
+}
+
+DEFAULT_USER = {"username": "wxo.archer@ibm.com", "password": "watsonx"}
+DEFAULT_LOCAL_SERVICE_URL = "http://localhost:4321"
+DEFAULT_LOCAL_AUTH_ENDPOINT = f"{DEFAULT_LOCAL_SERVICE_URL}/api/v1/auth/token"
+DEFAULT_LOCAL_TENANT_URL = f"{DEFAULT_LOCAL_SERVICE_URL}/tenants"
+DEFAULT_LOCAL_TENANT_AUTH_ENDPOINT = "{}/api/v1/auth/token?tenant_id={}"
+
+
+class LocalServiceInstance(BaseServiceInstance):
+    """lite service instance for local development"""
+
+    def __init__(self, client) -> None:
+
+        self._logger = logging.getLogger(__name__)
+        self._client = client
+        self._credentials: Credentials = client.credentials
+        self._credentials.local_global_token = self._get_user_auth_token()
+
+        self.tenant_id = self._create_default_tenant_if_not_exist()
+        self.tenant_access_token = self._get_tenant_token(self.tenant_id)
+        # the local token does not have exp claim.
+        self._client.token = self.tenant_access_token
+        super().__init__()
+
+    @staticmethod
+    def get_default_tenant(apikey):
+        headers = {"Authorization": f"Bearer {apikey}",
+                   "Content-Type": "application/json"}
+        resp = requests.get(DEFAULT_LOCAL_TENANT_URL, headers=headers)
+        if resp.status_code == 200:
+            tenant_config = resp.json()
+            for tenant in tenant_config:
+                if tenant["name"] == DEFAULT_TENANT["name"]:
+                    return tenant
+            return {}
+        else:
+            resp.raise_for_status()
+
+    @staticmethod
+    def create_default_tenant(apikey):
+        headers = {"Authorization": f"Bearer {apikey}",
+                   "Content-Type": "application/json"}
+        resp = requests.post(DEFAULT_LOCAL_TENANT_URL, headers=headers, json=DEFAULT_TENANT)
+        if resp.status_code == 201:
+            return True
+        else:
+            resp.raise_for_status()
+
+    def _create_default_tenant_if_not_exist(self) -> str:
+        user_auth_token = self._credentials.local_global_token
+        default_tenant = self.get_default_tenant(user_auth_token)
+
+        if not default_tenant:
+            logger.info("no local tenant found. A default tenant is created")
+            self.create_default_tenant(user_auth_token)
+            default_tenant = self.get_default_tenant(user_auth_token)
+        else:
+            logger.info("local tenant found")
+        tenant_id = default_tenant["id"]
+        return tenant_id
+
+    def _get_user_auth_token(self):
+        resp = requests.post(DEFAULT_LOCAL_AUTH_ENDPOINT, data=DEFAULT_USER)
+        if resp.status_code == 200:
+            return resp.json()["access_token"]
+        else:
+            resp.raise_for_status()
+
+    def _get_tenant_token(self, tenant_id: str):
+        resp = requests.post(DEFAULT_LOCAL_TENANT_AUTH_ENDPOINT.format(DEFAULT_LOCAL_SERVICE_URL, tenant_id),
+                             data=DEFAULT_USER)
+        if resp.status_code == 200:
+            return resp.json()["access_token"]
+        else:
+            resp.raise_for_status()
+
+    def _create_token(self) -> str:
+
+        return self._get_tenant_token(self.tenant_id)

ibm_watsonx_orchestrate/client/service_instance.py

@@ -0,0 +1,73 @@
+#  -----------------------------------------------------------------------------------------
+#  (C) Copyright IBM Corp. 2024.
+#  https://opensource.org/licenses/BSD-3-Clause
+#  -----------------------------------------------------------------------------------------
+
+from __future__ import annotations
+
+from ibm_cloud_sdk_core.authenticators import MCSPAuthenticator 
+from ibm_cloud_sdk_core.authenticators import IAMAuthenticator
+
+from ibm_watsonx_orchestrate.client.utils import check_token_validity
+from ibm_watsonx_orchestrate.client.base_service_instance import BaseServiceInstance
+from ibm_watsonx_orchestrate.cli.commands.environment.types import EnvironmentAuthType
+
+from ibm_watsonx_orchestrate.client.client_errors import (
+    ClientError,
+)
+
+
+class ServiceInstance(BaseServiceInstance):
+    """Connect, get details, and check usage of a Watson Machine Learning service instance."""
+
+    def __init__(self, client) -> None:
+        super().__init__()
+        self._client = client
+        self._credentials = client.credentials
+        self._client.token = self._get_token()
+
+    def _get_token(self) -> str:
+        # If no token is set
+        if self._client.token is None:
+            return self._create_token()
+
+        # Refresh is possible and token is expired
+        if self._is_token_refresh_possible() and self._check_token_expiry():
+            return self._create_token()
+
+        return self._client.token
+    
+    def _create_token(self) -> str:
+        if not self._credentials.auth_type:
+            if ".cloud.ibm.com" in self._credentials.url:
+                return self._authenticate(EnvironmentAuthType.IBM_CLOUD_IAM)
+            else:
+                return self._authenticate(EnvironmentAuthType.MCSP)
+        else:
+            return self._authenticate(self._credentials.auth_type)
+
+    def _authenticate(self, auth_type: str) -> str:
+        """Handles authentication based on the auth_type."""
+        try:
+            match auth_type:
+                case EnvironmentAuthType.MCSP:
+                    authenticator = MCSPAuthenticator(apikey=self._credentials.api_key, url=self._credentials.iam_url)
+                case EnvironmentAuthType.IBM_CLOUD_IAM:
+                    authenticator = IAMAuthenticator(apikey=self._credentials.api_key, url=self._credentials.iam_url)
+                case _:
+                    raise ClientError(f"Unsupported authentication type: {auth_type}")
+
+            return authenticator.token_manager.get_token()
+        except Exception as e:
+            raise ClientError(f"Error getting {auth_type.upper()} Token", e)
+
+    
+    def _is_token_refresh_possible(self) -> bool:
+        if self._credentials.api_key:
+            return True
+        return False
+    
+    def _check_token_expiry(self):
+        token = self._client.token
+
+        return not check_token_validity(token)

ibm_watsonx_orchestrate/client/tools/tool_client.py

@@ -0,0 +1,41 @@
+from ibm_watsonx_orchestrate.client.base_api_client import BaseAPIClient, ClientAPIException
+from typing_extensions import List
+
+class ToolClient(BaseAPIClient):
+    """
+    Client to handle CRUD operations for Tool endpoint
+    """
+
+    def create(self, payload: dict) -> dict:
+        return self._post("/tools", data=payload)
+
+    def get(self) -> dict:
+        return self._get("/tools")
+
+    def update(self, agent_id: str, data: dict) -> dict:
+        return self._put(f"/tools/{agent_id}", data=data)
+
+    def delete(self, tool_id: str) -> dict:
+        return self._delete(f"/tools/{tool_id}")
+
+    def upload_tools_artifact(self, tool_id: str, file_path: str) -> dict:
+        return self._post(f"/tools/{tool_id}/upload", files={"file": (f"{tool_id}.zip", open(file_path, "rb"), "application/zip", {"Expires": "0"})})
+    
+    def get_draft_by_name(self, tool_name: str) -> List[dict]:
+        return self.get_drafts_by_names([tool_name])
+
+    def get_drafts_by_names(self, tool_names: List[str]) -> List[dict]:
+        formatted_tool_names = [f"names={x}" for x in tool_names]
+        return self._get(f"/tools?{'&'.join(formatted_tool_names)}")
+    
+    def get_draft_by_id(self, tool_id: str) -> List[dict]:
+        if tool_id is None:
+            return ""
+        else:
+            try:
+                tool = self._get(f"/tools/{tool_id}")
+                return tool
+            except ClientAPIException as e:
+                if e.response.status_code == 404 and "not found with the given name" in e.response.text:
+                    return ""
+                raise(e)

ibm_watsonx_orchestrate/client/utils.py

@@ -0,0 +1,95 @@
+from ibm_watsonx_orchestrate.cli.config import (
+    Config,
+    DEFAULT_CONFIG_FILE_FOLDER,
+    DEFAULT_CONFIG_FILE,
+    AUTH_CONFIG_FILE_FOLDER,
+    AUTH_CONFIG_FILE,
+    AUTH_SECTION_HEADER,
+    AUTH_MCSP_TOKEN_OPT,
+    CONTEXT_SECTION_HEADER,
+    CONTEXT_ACTIVE_ENV_OPT,
+    ENVIRONMENTS_SECTION_HEADER,
+    ENV_WXO_URL_OPT
+)
+from threading import Lock
+from ibm_watsonx_orchestrate.client.base_api_client import BaseAPIClient
+from ibm_watsonx_orchestrate.utils.utils import yaml_safe_load
+import logging
+from typing import TypeVar
+import os
+import jwt
+import time
+
+logger = logging.getLogger(__name__)
+LOCK = Lock()
+T = TypeVar("T", bound=BaseAPIClient)
+
+
+def is_local_dev(url: str | None = None) -> bool:
+    if url is None:
+        cfg = Config()
+        active_env = cfg.read(CONTEXT_SECTION_HEADER, CONTEXT_ACTIVE_ENV_OPT)
+        url = cfg.get(ENVIRONMENTS_SECTION_HEADER, active_env, ENV_WXO_URL_OPT)
+
+    if url.startswith("http://localhost"):
+        return True
+
+    if url.startswith("http://127.0.0.1"):
+        return True
+
+    if url.startswith("http://[::1]"):
+        return True
+
+    if url.startswith("http://0.0.0.0"):
+        return True
+
+    return False
+
+def check_token_validity(token: str) -> bool:
+    try:
+        token_claimset = jwt.decode(token, options={"verify_signature": False})
+        expiry = token_claimset.get('exp')
+
+        current_timestamp = int(time.time())
+        # Check if the token is not expired (or will not be expired in 10 minutes)
+        if not expiry or current_timestamp < expiry - 600:
+            return True
+        return False
+    except:
+        return False
+
+
+def instantiate_client(client: type[T] , url: str | None=None) -> T:
+    try:
+        with LOCK:
+            with open(os.path.join(DEFAULT_CONFIG_FILE_FOLDER, DEFAULT_CONFIG_FILE), "r") as f:
+                config = yaml_safe_load(f)
+            active_env = config.get(CONTEXT_SECTION_HEADER, {}).get(CONTEXT_ACTIVE_ENV_OPT)
+
+            if not url:
+                url = config.get(ENVIRONMENTS_SECTION_HEADER, {}).get(active_env, {}).get(ENV_WXO_URL_OPT)
+
+            with open(os.path.join(AUTH_CONFIG_FILE_FOLDER, AUTH_CONFIG_FILE), "r") as f:
+                auth_config = yaml_safe_load(f)
+            auth_settings = auth_config.get(AUTH_SECTION_HEADER, {}).get(active_env, {})
+
+            if not active_env:
+                logger.error("No active environment set. Use `orchestrate env activate` to activate an environment")
+                exit(1)
+            if not url:
+                logger.error(f"No URL found for environment '{active_env}'. Use `orchestrate env list` to view existing environments and `orchesrtate env add` to reset the URL")
+                exit(1)
+            if not auth_settings:
+                logger.error(f"No credentials found for active env '{active_env}'. Use `orchestrate env activate {active_env}` to refresh your credentials")
+                exit(1)
+            token = auth_settings.get(AUTH_MCSP_TOKEN_OPT)
+            if not check_token_validity(token):
+                logger.error(f"The token found for environment '{active_env}' is missing or expired. Use `orchestrate env activate {active_env}` to fetch a new one")
+                exit(1)
+            client_instance = client(base_url=url, api_key=token, is_local=is_local_dev(url))
+
+        return client_instance
+    except FileNotFoundError as e:
+        message = "No active environment found. Please run `orchestrate env activate` to activate an environment"
+        logger.error(message)
+        raise FileNotFoundError(message)