ibm-watsonx-orchestrate 1.2.0__py3-none-any.whl → 1.4.2__py3-none-any.whl

ibm_watsonx_orchestrate/cli/init_helper.py

@@ -0,0 +1,43 @@
+import importlib.metadata
+from importlib import resources
+from typing import Optional
+from rich import print as pprint
+from dotenv import dotenv_values
+import typer
+
+from ibm_watsonx_orchestrate.cli.config import Config, PYTHON_REGISTRY_HEADER, \
+    PYTHON_REGISTRY_TEST_PACKAGE_VERSION_OVERRIDE_OPT
+
+
+def version_callback(checkVersion: bool=True):
+    if checkVersion:
+        __version__ = importlib.metadata.version('ibm-watsonx-orchestrate')
+        default_env = dotenv_values(resources.files("ibm_watsonx_orchestrate.docker").joinpath("default.env"))
+        cfg = Config()
+        pypi_override = cfg.read(PYTHON_REGISTRY_HEADER, PYTHON_REGISTRY_TEST_PACKAGE_VERSION_OVERRIDE_OPT)
+
+        adk_version_str = f"[bold]ADK Version[/bold]: {__version__}"
+        if pypi_override is not None:
+            adk_version_str += f" [red bold](override: {pypi_override})[/red bold]"
+        pprint(adk_version_str)
+
+
+        pprint("[bold]Developer Edition Image Tags[/bold] [italic](if not overridden in env file)[/italic]")
+        for key, value in default_env.items():
+            if key.endswith('_TAG') or key == 'DBTAG':
+                pprint(f"  [bold]{key}[/bold]: {value}")
+
+        raise typer.Exit()
+
+    
+
+def init_callback(
+    ctx: typer.Context,
+    version: Optional[bool] = typer.Option(
+      None, 
+      "--version",
+      help="Show the installed version of the ADK and Developer Edition Tags",
+      callback=version_callback
+    )
+):
+    pass

ibm_watsonx_orchestrate/cli/main.py

@@ -12,10 +12,12 @@ from ibm_watsonx_orchestrate.cli.commands.environment.environment_command import
 from ibm_watsonx_orchestrate.cli.commands.channels.channels_command import channel_app
 from ibm_watsonx_orchestrate.cli.commands.knowledge_bases.knowledge_bases_command import knowledge_bases_app
 from ibm_watsonx_orchestrate.cli.commands.toolkit.toolkit_command import toolkits_app
+from ibm_watsonx_orchestrate.cli.init_helper import init_callback
 
 app = typer.Typer(
     no_args_is_help=True,
-    pretty_exceptions_enable=False
+    pretty_exceptions_enable=False,
+    callback=init_callback
 )
 app.add_typer(login_app)
 app.add_typer(environment_app, name="env", help='Add, remove, or select the activate env other commands will interact with (either your local server or a production instance)')

ibm_watsonx_orchestrate/client/agents/agent_client.py

@@ -42,5 +42,8 @@ class AgentClient(BaseAPIClient):
                 if e.response.status_code == 404 and "not found with the given name" in e.response.text:
                     return ""
                 raise(e)
-
+    
+    def get_drafts_by_ids(self, agent_ids: List[str]) -> List[dict]:
+        formatted_agent_ids = [f"ids={x}" for x  in agent_ids]
+        return self._get(f"{self.base_endpoint}?{'&'.join(formatted_agent_ids)}")
     

ibm_watsonx_orchestrate/client/agents/assistant_agent_client.py

@@ -25,7 +25,7 @@ class AssistantAgentClient(BaseAPIClient):
         formatted_agent_names = [f"names={x}" for x  in agent_names]
         return self._get(f"/assistants/watsonx?{'&'.join(formatted_agent_names)}")
     
-    def get_draft_by_id(self, agent_id: str) -> List[dict]:
+    def get_draft_by_id(self, agent_id: str) -> dict | str:
         if agent_id is None:
             return ""
         else:
@@ -36,3 +36,7 @@ class AssistantAgentClient(BaseAPIClient):
                 if e.response.status_code == 404 and "Assistant not found" in e.response.text:
                     return ""
                 raise(e)
+    
+    def get_drafts_by_ids(self, agent_ids: List[str]) -> List[dict]:
+        formatted_agent_ids = [f"ids={x}" for x  in agent_ids]
+        return self._get(f"/assistants/watsonx?{'&'.join(formatted_agent_ids)}")

ibm_watsonx_orchestrate/client/agents/external_agent_client.py

@@ -33,6 +33,10 @@ class ExternalAgentClient(BaseAPIClient):
                 agent = self._get(f"/agents/external-chat/{agent_id}")
                 return agent
             except ClientAPIException as e:
-                if e.response.status_code == 404 and "not found with the given name" in e.response.text:
+                if e.response.status_code == 404 and ("not found with the given name" in e.response.text or "Assistant not found" in e.response.text):
                     return ""
                 raise(e)
+
+    def get_drafts_by_ids(self, agent_ids: List[str]) -> List[dict]:
+        formatted_agent_ids = [f"ids={x}" for x  in agent_ids]
+        return self._get(f"/agents/external-chat?{'&'.join(formatted_agent_ids)}")

ibm_watsonx_orchestrate/client/analytics/llm/analytics_llm_client.py

@@ -32,12 +32,8 @@ class AnalyticsLLMClient(BaseAPIClient):
     def create(self):
         raise RuntimeError('unimplemented')
 
-    def get(self, project_id: Optional[str] = None) -> AnalyticsLLMConfig:
-        params = {}
-        if project_id:
-            params['project_id'] = project_id
-
-        response = self._get(f"/analytics/llm", params=params)
+    def get(self) -> AnalyticsLLMConfig:
+        response = self._get(f"/analytics/llm")
 
         return AnalyticsLLMConfig.model_validate(response)
 

ibm_watsonx_orchestrate/client/base_api_client.py

@@ -45,12 +45,15 @@ class BaseAPIClient:
             headers["Authorization"] = f"Bearer {self.authenticator.token_manager.get_token()}"
         return headers
 
-    def _get(self, path: str, params: dict = None, data=None) -> dict:
+    def _get(self, path: str, params: dict = None, data=None, return_raw=False) -> dict:
 
         url = f"{self.base_url}{path}"
         response = requests.get(url, headers=self._get_headers(), params=params, data=data)
         self._check_response(response)
-        return response.json()
+        if not return_raw:
+            return response.json()
+        else:
+            return response
 
     def _post(self, path: str, data: dict = None, files: dict = None) -> dict:
         url = f"{self.base_url}{path}"

ibm_watsonx_orchestrate/client/connections/connections_client.py

@@ -51,7 +51,7 @@ class ConnectionsClient(BaseAPIClient):
     # DELETE api/v1/connections/applications/{app_id}
     def delete(self, app_id: str) -> dict:
         return self._delete(f"/connections/applications/{app_id}")
-    
+
     # GET /api/v1/connections/applications/{app_id}
     def get(self, app_id: str) -> GetConnectionResponse:
         try:
@@ -60,7 +60,7 @@ class ConnectionsClient(BaseAPIClient):
             if e.response.status_code == 404:
                 return None
             raise e
-    
+
 
     # GET api/v1/connections/applications
     def list(self) -> List[ListConfigsResponse]:
@@ -75,15 +75,15 @@ class ConnectionsClient(BaseAPIClient):
                 return []
             raise e
 
-    
+
     # POST /api/v1/connections/applications/{app_id}/configurations
     def create_config(self, app_id: str, payload: dict) -> None:
         self._post(f"/connections/applications/{app_id}/configurations", data=payload)
-    
+
     # PATCH /api/v1/connections/applications/{app_id}/configurations/{env}
     def update_config(self, app_id: str, env: ConnectionEnvironment, payload: dict) -> None:
         self._patch(f"/connections/applications/{app_id}/configurations/{env}", data=payload)
-    
+
     # `GET /api/v1/connections/applications/{app_id}/configurations/{env}'
     def get_config(self, app_id: str, env: ConnectionEnvironment) -> GetConfigResponse:
         try:
@@ -93,7 +93,7 @@ class ConnectionsClient(BaseAPIClient):
             if e.response.status_code == 404:
                 return None
             raise e
-    
+
     # POST /api/v1/connections/applications/{app_id}/configs/{env}/credentials
     # POST /api/v1/connections/applications/{app_id}/configs/{env}/runtime_credentials
     def create_credentials(self, app_id: str, env: ConnectionEnvironment, payload: dict, use_sso: bool) -> None:
@@ -101,7 +101,7 @@ class ConnectionsClient(BaseAPIClient):
             self._post(f"/connections/applications/{app_id}/configs/{env}/credentials", data=payload)
         else:
             self._post(f"/connections/applications/{app_id}/configs/{env}/runtime_credentials", data=payload)
-    
+
     # PATCH /api/v1/connections/applications/{app_id}/configs/{env}/credentials
     # PATCH /api/v1/connections/applications/{app_id}/configs/{env}/runtime_credentials
     def update_credentials(self, app_id: str, env: ConnectionEnvironment, payload: dict, use_sso: bool) -> None:
@@ -109,7 +109,7 @@ class ConnectionsClient(BaseAPIClient):
             self._patch(f"/connections/applications/{app_id}/configs/{env}/credentials", data=payload)
         else:
             self._patch(f"/connections/applications/{app_id}/configs/{env}/runtime_credentials", data=payload)
-    
+
     # GET /api/v1/connections/applications/{app_id}/configs/credentials?env={env}
     # GET /api/v1/connections/applications/{app_id}/configs/runtime_credentials?env={env}
     def get_credentials(self, app_id: str, env: ConnectionEnvironment, use_sso: bool) -> dict:
@@ -122,7 +122,7 @@ class ConnectionsClient(BaseAPIClient):
             if e.response.status_code == 404:
                 return None
             raise e
-    
+
     # DELETE /api/v1/connections/applications/{app_id}/configs/{env}/credentials
     # DELETE /api/v1/connections/applications/{app_id}/configs/{env}/runtime_credentials
     def delete_credentials(self, app_id: str, env: ConnectionEnvironment, use_sso: bool) -> None:
@@ -130,7 +130,7 @@ class ConnectionsClient(BaseAPIClient):
             self._delete(f"/connections/applications/{app_id}/configs/{env}/credentials")
         else:
             self._delete(f"/connections/applications/{app_id}/configs/{env}/runtime_credentials")
-    
+
     def get_draft_by_app_id(self, app_id: str) -> GetConnectionResponse:
         return self.get(app_id=app_id)
 
@@ -141,22 +141,16 @@ class ConnectionsClient(BaseAPIClient):
             if connection:
                 connections += connection
         return connections
-    
+
     def get_draft_by_id(self, conn_id) -> str:
         """Retrieve the app ID for a given connection ID."""
         if conn_id is None:
             return ""
         try:
-            connections = self.list()
+            app_details = self._get(f"/connections/applications/id/{conn_id}")
+            return app_details.get("app_id")
         except ClientAPIException as e:
             if e.response.status_code == 404:
                 logger.warning(f"Connections not found. Returning connection ID: {conn_id}")
-                return conn_id 
-            raise
-
-        app_id = next((conn.app_id for conn in connections if conn.connection_id == conn_id), None)
-
-        if app_id is None:
-            logger.warning(f"Connection with ID {conn_id} not found. Returning connection ID.")
-            return conn_id  
-        return app_id
+                return conn_id
+            raise e

ibm_watsonx_orchestrate/client/model_policies/model_policies_client.py

@@ -0,0 +1,47 @@
+from typing import List
+
+from pydantic import ValidationError
+from ibm_watsonx_orchestrate.client.base_api_client import BaseAPIClient, ClientAPIException
+
+import logging
+
+from ibm_watsonx_orchestrate.client.model_policies.types import ModelPolicy
+from ibm_watsonx_orchestrate.client.models.types import ListVirtualModel, CreateVirtualModel
+
+logger = logging.getLogger(__name__)
+
+
+
+
+
+class ModelPoliciesClient(BaseAPIClient):
+    """
+    Client to handle CRUD operations for ModelPolicies endpoint
+    """
+    # POST api/v1/model_policy
+    def create(self, model: ModelPolicy) -> None:
+        self._post("/model_policy", data=model.model_dump())
+
+    # DELETE api/v1/model_policy/{model_policy_id}
+    def delete(self, model_policy_id: str) -> dict:
+        return self._delete(f"/model_policy/{model_policy_id}")
+
+    # GET /api/v1/model_policy/{app_id}
+    def get(self, model_policy_id: str):
+        raise NotImplementedError
+
+
+    # GET api/v1/model_policy
+    def list(self) -> List[ModelPolicy]:
+        try:
+            res = self._get(f"/model_policy")
+            return [ModelPolicy.model_validate(policy) for policy in res]
+        except ValidationError as e:
+            logger.error("Received unexpected response from server")
+            raise e
+        except ClientAPIException as e:
+            if e.response.status_code == 404:
+                return []
+            raise e
+
+

ibm_watsonx_orchestrate/client/model_policies/types.py

@@ -0,0 +1,36 @@
+from enum import Enum
+from typing import Optional, List, Dict, Any, Union
+
+from pydantic import Field, BaseModel, ConfigDict
+
+class ModelPolicyStrategyMode(str, Enum):
+    LOAD_BALANCED = "loadbalance"
+    FALL_BACK = "fallback"
+    SINGLE = "single"
+
+class ModelPolicyStrategy(BaseModel):
+    mode: ModelPolicyStrategyMode = None
+    on_status_codes: List[int] = None
+
+class ModelPolicyRetry(BaseModel):
+    attempts: int = None
+    on_status_codes: List[int] = None
+
+class ModelPolicyTarget(BaseModel):
+    model_id: str = None
+    weight: int = None
+
+
+class ModelPolicyInner(BaseModel):
+    strategy: ModelPolicyStrategy = None
+    retry: ModelPolicyRetry = None
+    targets: List[Union[ModelPolicyTarget, 'ModelPolicyInner']] = None
+
+
+class ModelPolicy(BaseModel):
+    model_config = ConfigDict(extra='allow')
+
+    name: str
+    display_name: str
+    policy: ModelPolicyInner
+

ibm_watsonx_orchestrate/client/models/models_client.py

@@ -0,0 +1,46 @@
+from typing import List
+
+from pydantic import ValidationError
+from ibm_watsonx_orchestrate.client.base_api_client import BaseAPIClient, ClientAPIException
+
+import logging
+
+from ibm_watsonx_orchestrate.client.models.types import ListVirtualModel, CreateVirtualModel
+
+logger = logging.getLogger(__name__)
+
+
+
+
+
+class ModelsClient(BaseAPIClient):
+    """
+    Client to handle CRUD operations for Models endpoint
+    """
+    # POST api/v1/models
+    def create(self, model: CreateVirtualModel) -> None:
+        self._post("/models", data=model.model_dump())
+
+    # DELETE api/v1/models/{model_id}
+    def delete(self, model_id: str) -> dict:
+        return self._delete(f"/models/{model_id}")
+
+    # GET /api/v1/models/{app_id}
+    def get(self, model_id: str):
+        raise NotImplementedError
+
+
+    # GET api/v1/models
+    def list(self) -> List[ListVirtualModel]:
+        try:
+            res = self._get(f"/models")
+            return [ListVirtualModel.model_validate(conn) for conn in res]
+        except ValidationError as e:
+            logger.error("Received unexpected response from server")
+            raise e
+        except ClientAPIException as e:
+            if e.response.status_code == 404:
+                return []
+            raise e
+
+

ibm_watsonx_orchestrate/client/models/types.py

@@ -0,0 +1,177 @@
+from typing import Optional, List, Dict, Any, Union
+from enum import Enum
+from pydantic import Field, BaseModel, ConfigDict
+
+class  ModelProvider(str, Enum):
+    OPENAI = 'openai'
+    A21 = 'a21'
+    ANTHROPIC = 'anthropic'
+    ANYSCALE = 'anyscale'
+    AZURE_OPENAI = 'azure-openai'
+    AZURE_AI = 'azure-ai'
+    BEDROCK = 'bedrock'
+    CEREBRAS = 'cerebras'
+    COHERE = 'cohere'
+    GOOGLE = 'google'
+    VERTEX_AI = 'vertex-ai'
+    GROQ = 'groq'
+    HUGGINGFACE = 'huggingface'
+    MISTRAL_AI = 'mistral-ai'
+    JINA = 'jina'
+    OLLAMA = 'ollama'
+    OPENROUTER = 'openrouter'
+    STABILITY_AI = 'stability-ai'
+    TOGETHER_AI = 'together-ai'
+    WATSONX = 'watsonx'
+
+    def __str__(self):
+        return self.value
+
+    def __repr__(self):
+        return self.value
+    
+
+class ProviderConfig(BaseModel):
+    # Required fields
+    provider: Optional[str]=''
+
+
+    api_key: Optional[str] = '' # this is not optional
+    url_to_fetch: Optional[str] = Field(None, alias="urlToFetch")
+
+    # Misc
+    custom_host: Optional[str] = Field(None, alias="customHost")
+    forward_headers: Optional[List[str]] = Field(None, alias="forwardHeaders")
+    index: Optional[int] = None
+    cache: Optional[Union[str, Dict[str, Any]]] = None  # Define CacheSettings if needed
+    metadata: Optional[Dict[str, str]] = None
+    request_timeout: Optional[int] = Field(None, alias="requestTimeout")
+    transform_to_form_data: Optional[bool] = Field(None, alias="transformToFormData")
+
+    # Azure specific
+    azure_resource_name: Optional[str] = Field(None, alias="resourceName")
+    azure_deployment_id: Optional[str] = Field(None, alias="deploymentId")
+    azure_api_version: Optional[str] = Field(None, alias="apiVersion")
+    ad_auth: Optional[str] = Field(None, alias="adAuth")
+    azure_auth_mode: Optional[str] = Field(None, alias="azureAuthMode")
+    azure_managed_client_id: Optional[str] = Field(None, alias="azureManagedClientId")
+    azure_entra_client_id: Optional[str] = Field(None, alias="azureEntraClientId")
+    azure_entra_client_secret: Optional[str] = Field(None, alias="azureEntraClientSecret")
+    azure_entra_tenant_id: Optional[str] = Field(None, alias="azureEntraTenantId")
+    azure_ad_token: Optional[str] = Field(None, alias="azureAdToken")
+    azure_model_name: Optional[str] = Field(None, alias="azureModelName")
+
+    # Workers AI specific
+    workers_ai_account_id: Optional[str] = Field(None, alias="workersAiAccountId")
+
+
+    # AWS
+    aws_secret_access_key: Optional[str] = Field(None, alias="awsSecretAccessKey")
+    aws_access_key_id: Optional[str] = Field(None, alias="awsAccessKeyId")
+    aws_session_token: Optional[str] = Field(None, alias="awsSessionToken")
+    aws_region: Optional[str] = Field(None, alias="awsRegion")
+    aws_auth_type: Optional[str] = Field(None, alias="awsAuthType")
+    aws_role_arn: Optional[str] = Field(None, alias="awsRoleArn")
+    aws_external_id: Optional[str] = Field(None, alias="awsExternalId")
+    aws_s3_bucket: Optional[str] = Field(None, alias="awsS3Bucket")
+    aws_s3_object_key: Optional[str] = Field(None, alias="awsS3ObjectKey")
+    aws_bedrock_model: Optional[str] = Field(None, alias="awsBedrockModel")
+    aws_server_side_encryption: Optional[str] = Field(None, alias="awsServerSideEncryption")
+    aws_server_side_encryption_kms_key_id: Optional[str] = Field(None, alias="awsServerSideEncryptionKMSKeyId")
+
+    # Sagemaker
+    amzn_sagemaker_custom_attributes: Optional[str] = Field(None, alias="amznSagemakerCustomAttributes")
+    amzn_sagemaker_target_model: Optional[str] = Field(None, alias="amznSagemakerTargetModel")
+    amzn_sagemaker_target_variant: Optional[str] = Field(None, alias="amznSagemakerTargetVariant")
+    amzn_sagemaker_target_container_hostname: Optional[str] = Field(None, alias="amznSagemakerTargetContainerHostname")
+    amzn_sagemaker_inference_id: Optional[str] = Field(None, alias="amznSagemakerInferenceId")
+    amzn_sagemaker_enable_explanations: Optional[str] = Field(None, alias="amznSagemakerEnableExplanations")
+    amzn_sagemaker_inference_component: Optional[str] = Field(None, alias="amznSagemakerInferenceComponent")
+    amzn_sagemaker_session_id: Optional[str] = Field(None, alias="amznSagemakerSessionId")
+    amzn_sagemaker_model_name: Optional[str] = Field(None, alias="amznSagemakerModelName")
+
+    # Stability AI
+    stability_client_id: Optional[str] = Field(None, alias="stabilityClientId")
+    stability_client_user_id: Optional[str] = Field(None, alias="stabilityClientUserId")
+    stability_client_version: Optional[str] = Field(None, alias="stabilityClientVersion")
+
+    # Hugging Face
+    huggingface_base_url: Optional[str] = Field(None, alias="huggingfaceBaseUrl")
+
+    # Google Vertex AI
+    vertex_region: Optional[str] = Field(None, alias="vertexRegion")
+    vertex_project_id: Optional[str] = Field(None, alias="vertexProjectId")
+    vertex_service_account_json: Optional[Dict[str, Any]] = Field(None, alias="vertexServiceAccountJson")
+    vertex_storage_bucket_name: Optional[str] = Field(None, alias="vertexStorageBucketName")
+    vertex_model_name: Optional[str] = Field(None, alias="vertexModelName")
+
+    filename: Optional[str] = None
+
+    after_request_hooks: Optional[List[Dict[str, Any]]] = Field(None, alias="afterRequestHooks")
+    before_request_hooks: Optional[List[Dict[str, Any]]] = Field(None, alias="beforeRequestHooks")
+    default_input_guardrails: Optional[List[Dict[str, Any]]] = Field(None, alias="defaultInputGuardrails")
+    default_output_guardrails: Optional[List[Dict[str, Any]]] = Field(None, alias="defaultOutputGuardrails")
+
+    # OpenAI
+    openai_project: Optional[str] = Field(None, alias="openaiProject")
+    openai_organization: Optional[str] = Field(None, alias="openaiOrganization")
+    openai_beta: Optional[str] = Field(None, alias="openaiBeta")
+
+    # Azure Inference
+    azure_deployment_name: Optional[str] = Field(None, alias="azureDeploymentName")
+    azure_api_version: Optional[str] = Field(None, alias="azureApiVersion")
+    azure_extra_params: Optional[str] = Field(None, alias="azureExtraParams")
+    azure_foundry_url: Optional[str] = Field(None, alias="azureFoundryUrl")
+
+    strict_open_ai_compliance: Optional[bool] = Field(False, alias="strictOpenAiCompliance")
+    mistral_fim_completion: Optional[str] = Field(None, alias="mistralFimCompletion")
+
+    # Anthropic
+    anthropic_beta: Optional[str] = Field(None, alias="anthropicBeta")
+    anthropic_version: Optional[str] = Field(None, alias="anthropicVersion")
+
+    # Fireworks
+    fireworks_account_id: Optional[str] = Field(None, alias="fireworksAccountId")
+
+    # Cortex
+    snowflake_account: Optional[str] = Field(None, alias="snowflakeAccount")
+
+    # WatsonX
+    watsonx_version: Optional[str] = Field(None, alias="watsonxVersion")
+    watsonx_space_id: Optional[str] = Field(None, alias="watsonxSpaceId")
+    watsonx_project_id: Optional[str] = Field(None, alias="watsonxProjectId")
+
+    model_config = {
+        "populate_by_name": True,  # Replaces allow_population_by_field_name
+        "extra": "forbid",         # Same as before
+        "json_schema_extra": lambda schema: schema.get("properties", {}).pop("provider", None)
+    }
+
+
+
+class CreateVirtualModel(BaseModel):
+    model_config = ConfigDict(extra='allow')
+
+    name: str
+    display_name: Optional[str]
+    description: Optional[str]
+    config: Optional[dict] = None
+    provider_config: ProviderConfig
+    tags: List[str]
+    model_type: str = 'chat'
+
+
+class ListVirtualModel(BaseModel):
+    model_config = ConfigDict(extra='allow')
+
+    id: Optional[str] = None
+    name: Optional[str] = None
+    display_name: Optional[str] = None
+    description: Optional[str] = None
+    provider: Optional[str] = None
+    # model_configuration: Optional[dict] = None
+    provider_config: Optional[ProviderConfig] = None
+    tags: Optional[List[str]] = None
+    model_type: Optional[str] = None
+
+ANTHROPIC_DEFAULT_MAX_TOKENS = 4096

ibm_watsonx_orchestrate/client/toolkit/toolkit_client.py

@@ -1,9 +1,18 @@
 from ibm_watsonx_orchestrate.client.base_api_client import BaseAPIClient, ClientAPIException
+from ibm_watsonx_orchestrate.client.utils import is_local_dev
 from typing_extensions import List
 import os
 import json
 
 class ToolKitClient(BaseAPIClient):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.base_endpoint = "/orchestrate/toolkits" if is_local_dev(self.base_url) else "/toolkits"
+
+    def get(self) -> dict:
+        return self._get(self.base_endpoint)
+
     # POST /toolkits/prepare/list-tools
     def list_tools(self, zip_file_path: str, command: str, args: List[str]) -> List[str]:
         """
@@ -24,7 +33,7 @@ class ToolKitClient(BaseAPIClient):
                 "file": (filename, f, "application/zip"),
             }
 
-            response = self._post("/orchestrate/toolkits/prepare/list-tools", files=files)
+            response = self._post(f"{self.base_endpoint}/prepare/list-tools", files=files)
 
         return response.get("tools", [])
 
@@ -35,7 +44,7 @@ class ToolKitClient(BaseAPIClient):
         Creates new toolkit metadata
         """
         try:
-            return self._post("/orchestrate/toolkits", data=payload)
+            return self._post(self.base_endpoint, data=payload)
         except ClientAPIException as e:
             if e.response.status_code == 400 and "already exists" in e.response.text:
                 raise ClientAPIException(
@@ -54,25 +63,25 @@ class ToolKitClient(BaseAPIClient):
             files = {
                 "file": (filename, f, "application/zip", {"Expires": "0"})
             }
-            return self._post(f"/orchestrate/toolkits/{toolkit_id}/upload", files=files)
+            return self._post(f"{self.base_endpoint}/{toolkit_id}/upload", files=files)
         
     # DELETE /toolkits/{toolkit-id} 
     def delete(self, toolkit_id: str) -> dict:
-        return self._delete(f"/orchestrate/toolkits/{toolkit_id}")
+        return self._delete(f"{self.base_endpoint}/{toolkit_id}")
 
     def get_draft_by_name(self, toolkit_name: str) -> List[dict]:
         return self.get_drafts_by_names([toolkit_name])
 
     def get_drafts_by_names(self, toolkit_names: List[str]) -> List[dict]:
         formatted_toolkit_names = [f"names={x}" for x in toolkit_names]
-        return self._get(f"/orchestrate/toolkits?{'&'.join(formatted_toolkit_names)}")
+        return self._get(f"{self.base_endpoint}?{'&'.join(formatted_toolkit_names)}")
     
     def get_draft_by_id(self, toolkit_id: str) -> dict:
         if toolkit_id is None:
             return ""
         else:
             try:
-                toolkit = self._get(f"/orchestrate/toolkits/{toolkit_id}")
+                toolkit = self._get(f"{self.base_endpoint}/{toolkit_id}")
                 return toolkit
             except ClientAPIException as e:
                 if e.response.status_code == 404 and "not found with the given name" in e.response.text:

ibm_watsonx_orchestrate/client/tools/tempus_client.py

@@ -0,0 +1,40 @@
+from ibm_watsonx_orchestrate.client.base_api_client import BaseAPIClient, ClientAPIException
+from typing_extensions import List
+from urllib.parse import urlparse, urlunparse
+from ibm_cloud_sdk_core.authenticators import MCSPAuthenticator
+
+DEFAULT_TEMPUS_PORT= 9044
+
+class TempusClient(BaseAPIClient):
+    """
+    Client to handle CRUD operations for Tempus endpoint
+
+    This may be temporary and may want to create a proxy API in wxo-server 
+    to redirect to the internal tempus runtime, and add a new operation in the ToolClient instead
+    """
+    def __init__(self, base_url: str, api_key: str = None, is_local: bool = False, authenticator: MCSPAuthenticator = None):
+        parsed_url = urlparse(base_url)
+       
+        # Reconstruct netloc with new port - use default above - eventually we need to open up a way through the wxo-server API
+        new_netloc = f"{parsed_url.hostname}:{DEFAULT_TEMPUS_PORT}"
+
+        # Replace netloc and rebuild the URL
+        new_url = urlunparse(parsed_url._replace(netloc=new_netloc))
+        # remove trailing slash
+
+        super().__init__(
+            base_url=new_url,
+            api_key=api_key,
+            is_local=is_local,
+            authenticator=authenticator
+        )
+        
+    def create_update_flow_model(self, flow_id: str, model: dict) -> dict:
+        return self._post(f"/v1/flow-models/{flow_id}", data=model)
+    
+    def run_flow(self, flow_id: str, input: dict) -> dict:
+        return self._post(f"/v1/flows/{flow_id}/versions/TIP/run", data=input)
+    
+    def arun_flow(self, flow_id: str, input: dict) -> dict:
+        return self._post(f"/v1/flows/{flow_id}/versions/TIP/run/async", data=input)
+