ibm-watsonx-orchestrate 1.9.0b2__py3-none-any.whl → 1.10.0b1__py3-none-any.whl

ibm_watsonx_orchestrate/cli/commands/agents/agents_controller.py

@@ -29,13 +29,17 @@ from ibm_watsonx_orchestrate.client.agents.agent_client import AgentClient, Agen
 from ibm_watsonx_orchestrate.client.agents.external_agent_client import ExternalAgentClient
 from ibm_watsonx_orchestrate.client.agents.assistant_agent_client import AssistantAgentClient
 from ibm_watsonx_orchestrate.client.tools.tool_client import ToolClient
+from ibm_watsonx_orchestrate.client.voice_configurations.voice_configurations_client import VoiceConfigurationsClient
 from ibm_watsonx_orchestrate.utils.exceptions import BadRequest
 from ibm_watsonx_orchestrate.client.connections import get_connections_client
 from ibm_watsonx_orchestrate.client.knowledge_bases.knowledge_base_client import KnowledgeBaseClient
 
-from ibm_watsonx_orchestrate.client.utils import instantiate_client
+from ibm_watsonx_orchestrate.client.utils import instantiate_client, is_local_dev
 from ibm_watsonx_orchestrate.utils.utils import check_file_in_zip
 
+from rich.console import Console
+from rich.progress import Progress, SpinnerColumn, TextColumn
+
 logger = logging.getLogger(__name__)
 
 # Helper generic type for any agent
@@ -197,6 +201,7 @@ def get_app_id_from_conn_id(conn_id: str) -> str:
         exit(1)
     return app_id
 
+
 def get_agent_details(name: str, client: AgentClient | ExternalAgentClient | AssistantAgentClient) -> dict:
     agent_specs = client.get_draft_by_name(name)
     if len(agent_specs) > 1:
@@ -219,6 +224,7 @@ class AgentsController:
         self.assistant_client = None
         self.tool_client = None
         self.knowledge_base_client = None
+        self.voice_configuration_client = None
 
     def get_native_client(self):
         if not self.native_client:
@@ -245,6 +251,11 @@ class AgentsController:
             self.knowledge_base_client = instantiate_client(KnowledgeBaseClient)
         return self.knowledge_base_client
     
+    def get_voice_configuration_client(self):
+        if not self.voice_configuration_client:
+            self.voice_configuration_client = instantiate_client(VoiceConfigurationsClient)
+        return self.voice_configuration_client
+    
     @staticmethod
     def import_agent(file: str, app_id: str) -> List[Agent | ExternalAgent | AssistantAgent]:
         agents = parse_file(file)
@@ -520,6 +531,38 @@ class AgentsController:
                 guideline.tool = name
 
         return ref_agent
+    
+    def get_voice_config_name_from_id(self, voice_config_id: str) -> str | None:
+        client = self.get_voice_configuration_client()
+        config = client.get_by_id(voice_config_id)
+        return config.name if config else None
+
+    def get_voice_config_id_from_name(self, voice_config_name: str) -> str | None:
+        client = self.get_voice_configuration_client()
+        configs = client.get_by_name(voice_config_name)
+
+        if len(configs) == 0:
+            logger.error(f"No voice_configs with the name '{voice_config_name}' found. Failed to get config")
+            sys.exit(1)
+        
+        if len(configs) > 1:
+            logger.error(f"Multiple voice_configs with the name '{voice_config_name}' found. Failed to get config")
+            sys.exit(1)
+        
+        return configs[0].voice_configuration_id
+
+
+    def reference_voice_config(self,agent: Agent):
+        deref_agent = deepcopy(agent)
+        deref_agent.voice_configuration = self.get_voice_config_name_from_id(agent.voice_configuration_id)
+        del deref_agent.voice_configuration_id
+        return deref_agent
+
+    def dereference_voice_config(self,agent: Agent):
+        ref_agent = deepcopy(agent)
+        ref_agent.voice_configuration_id = self.get_voice_config_id_from_name(agent.voice_configuration)
+        del ref_agent.voice_configuration
+        return ref_agent
 
     @staticmethod
     def dereference_app_id(agent: ExternalAgent | AssistantAgent) -> ExternalAgent | AssistantAgent:
@@ -540,7 +583,18 @@ class AgentsController:
             agent.config.connection_id = None
 
         return agent
+    
+    def dereference_common_agent_dependencies(self, agent: AnyAgentT) -> AnyAgentT:
+        if agent.voice_configuration:
+            agent = self.dereference_voice_config(agent)
+
+        return agent  
+
+    def reference_common_agent_dependencies(self, agent: AnyAgentT) -> AnyAgentT:
+        if agent.voice_configuration_id:
+            agent = self.reference_voice_config(agent)
 
+        return agent
 
     def dereference_native_agent_dependencies(self, agent: Agent) -> Agent:
         if agent.collaborators and len(agent.collaborators):
@@ -584,6 +638,8 @@ class AgentsController:
     
     # Convert all names used in an agent to the corresponding ids
     def dereference_agent_dependencies(self, agent: AnyAgentT) -> AnyAgentT:
+
+        agent = self.dereference_common_agent_dependencies(agent)
         if isinstance(agent, Agent):
             return self.dereference_native_agent_dependencies(agent)
         if isinstance(agent, ExternalAgent) or isinstance(agent, AssistantAgent):
@@ -591,6 +647,8 @@ class AgentsController:
 
     # Convert all ids used in an agent to the corresponding names
     def reference_agent_dependencies(self, agent: AnyAgentT) -> AnyAgentT:
+
+        agent = self.reference_common_agent_dependencies(agent)
         if isinstance(agent, Agent):
             return self.reference_native_agent_dependencies(agent)
         if isinstance(agent, ExternalAgent) or isinstance(agent, AssistantAgent):
@@ -1111,3 +1169,114 @@ class AgentsController:
             logger.info(f"Successfully wrote agents and tools to '{output_path}'")
             zip_file_out.close()
 
+
+    def deploy_agent(self, name: str):
+        if is_local_dev():
+            logger.error("Agents cannot be deployed in Developer Edition")
+            sys.exit(1)
+        native_client = self.get_native_client()
+        external_client = self.get_external_client()
+        assistant_client = self.get_assistant_client()
+
+        existing_native_agents = native_client.get_draft_by_name(name)
+        existing_external_agents = external_client.get_draft_by_name(name)
+        existing_assistant_agents = assistant_client.get_draft_by_name(name)
+
+        if len(existing_native_agents) == 0 and (len(existing_external_agents) >= 1 or len(existing_assistant_agents) >= 1):
+            logger.error(f"No native agent found with name '{name}'. Only Native Agents can be deployed to a Live Environment")
+            sys.exit(1)
+        if len(existing_native_agents) > 1:
+            logger.error(f"Multiple native agents with the name '{name}' found. Failed to get agent")
+            sys.exit(1)
+        if len(existing_native_agents) == 0:
+            logger.error(f"No native agents with the name '{name}' found. Failed to get agent")
+            sys.exit(1)
+            
+
+        agent_details = existing_native_agents[0]
+        agent_id = agent_details.get("id")
+
+        environments = native_client.get_environments_for_agent(agent_id)
+
+        live_environment = [env for env in environments if env.get("name") == "live"]
+        if live_environment is None:
+            logger.error("No live environment found for this tenant")
+            sys.exit(1)
+
+        live_env_id = live_environment[0].get("id")
+
+        console = Console()
+        with Progress(
+            SpinnerColumn(spinner_name="dots"),
+            TextColumn("[progress.description]{task.description}"),
+            transient=True,
+            console=console,
+                ) as progress:
+                    progress.add_task(description="Deploying agent to Live envrionment", total=None)
+
+                    status = native_client.deploy(agent_id, live_env_id)
+
+        if status:
+            logger.info(f"Successfully deployed agent {name}")
+        else:
+            logger.error(f"Error deploying agent {name}")
+
+    def undeploy_agent(self, name: str):
+        if is_local_dev():
+            logger.error("Agents cannot be undeployed in Developer Edition")
+            sys.exit(1)
+        
+        native_client = self.get_native_client()
+        external_client = self.get_external_client()
+        assistant_client = self.get_assistant_client()
+
+        existing_native_agents = native_client.get_draft_by_name(name)
+        existing_external_agents = external_client.get_draft_by_name(name)
+        existing_assistant_agents = assistant_client.get_draft_by_name(name)
+
+        if len(existing_native_agents) == 0 and (len(existing_external_agents) >= 1 or len(existing_assistant_agents) >= 1):
+            logger.error(f"No native agent found with name '{name}'. Only Native Agents can be undeployed from a Live Environment")
+            sys.exit(1)
+        if len(existing_native_agents) > 1:
+            logger.error(f"Multiple native agents with the name '{name}' found. Failed to get agent")
+            sys.exit(1)
+        if len(existing_native_agents) == 0:
+            logger.error(f"No native agents with the name '{name}' found. Failed to get agent")
+            sys.exit(1)
+
+        agent_details = existing_native_agents[0]
+        agent_id = agent_details.get("id")
+
+        environments = native_client.get_environments_for_agent(agent_id)
+        live_environment = [env for env in environments if env.get("name") == "live"]
+        if live_environment is None:
+            logger.error("No live environment found for this tenant")
+            sys.exit(1)
+        version_id = live_environment[0].get("current_version")
+
+        if version_id is None:
+            agent_name = agent_details.get("name")
+            logger.error(f"Agent {agent_name} does not exist in a Live environment")
+            sys.exit(1)
+
+        draft_environment = [env for env in environments if env.get("name") == "draft"]
+        if draft_environment is None:
+            logger.error("No draft environment found for this tenant")
+            sys.exit(1)
+        draft_env_id = draft_environment[0].get("id")
+
+        console = Console()
+        with Progress(
+            SpinnerColumn(spinner_name="dots"),
+            TextColumn("[progress.description]{task.description}"),
+            transient=True,
+            console=console,
+                ) as progress:
+                    progress.add_task(description="Undeploying agent to Draft envrionment", total=None)
+
+                    status = native_client.undeploy(agent_id, version_id, draft_env_id)
+        if status:
+            logger.info(f"Successfully undeployed agent {name}")
+        else:
+            logger.error(f"Error undeploying agent {name}")
+

ibm_watsonx_orchestrate/cli/commands/connections/connections_command.py

@@ -160,7 +160,7 @@ def set_credentials_connection_command(
         typer.Option(
             '--username',
             '-u',
-            help='For basic auth, the username to login with'
+            help='For basic auth and oauth_auth_password_flow, the username to login with'
         )
     ] = None,
     password: Annotated[
@@ -168,7 +168,7 @@ def set_credentials_connection_command(
         typer.Option(
             '--password',
             '-p',
-            help='For basic auth, the password to login with'
+            help='For basic auth and oauth_auth_password_flow, the password to login with'
         )
     ] = None,
     token: Annotated[
@@ -191,14 +191,14 @@ def set_credentials_connection_command(
         typer.Option(
             '--client-id',
             # help='For oauth_auth_on_behalf_of_flow, oauth_auth_code_flow, oauth_auth_implicit_flow, oauth_auth_password_flow and oauth_auth_client_credentials_flow, the client_id to authenticate against the application token server'
-            help='For oauth_auth_on_behalf_of_flow and oauth_auth_client_credentials_flow, the client_id to authenticate against the application token server'
+            help='For oauth_auth_on_behalf_of_flow, oauth_auth_password_flow and oauth_auth_client_credentials_flow, the client_id to authenticate against the application token server'
         )
     ] = None,
     client_secret: Annotated[
         str,
         typer.Option(
             '--client-secret',
-            help='For oauth_auth_client_credentials_flow, the client_secret to authenticate with'
+            help='For oauth_auth_client_credentials_flow and oauth_auth_password_flow, the client_secret to authenticate with'
         )
     ] = None,
     send_via: Annotated[
@@ -213,7 +213,7 @@ def set_credentials_connection_command(
         typer.Option(
             '--token-url',
             # help='For oauth_auth_on_behalf_of_flow, oauth_auth_code_flow, oauth_auth_password_flow and oauth_auth_client_credentials_flow, the url of the application token server'
-            help='For oauth_auth_on_behalf_of_flow and oauth_auth_client_credentials_flow, the url of the application token server'
+            help='For oauth_auth_on_behalf_of_flow, oauth_auth_password_flow and oauth_auth_client_credentials_flow, the url of the application token server'
         )
     ] = None,
     auth_url: Annotated[
@@ -227,14 +227,14 @@ def set_credentials_connection_command(
         str,
         typer.Option(
             '--grant-type',
-            help='For oauth_auth_on_behalf_of_flow and oauth_auth_client_credentials_flow, the grant type used by the application token server'
+            help='For oauth_auth_on_behalf_of_flow, oauth_auth_password_flow and oauth_auth_client_credentials_flow, the grant type used by the application token server'
         )
     ] = None,
     scope: Annotated[
         str,
         typer.Option(
             '--scope',
-            help='For oauth_auth_code_flow and oauth_auth_client_credentials_flow, the optional scopes used by the application token server. Should be in the form of a space seperated string.'
+            help='For oauth_auth_code_flow, oauth_auth_password_flow and oauth_auth_client_credentials_flow, the optional scopes used by the application token server. Should be in the form of a space seperated string.'
         )
     ] = None,
     entries: Annotated[

ibm_watsonx_orchestrate/cli/commands/connections/connections_controller.py

@@ -19,15 +19,15 @@ from ibm_watsonx_orchestrate.agent_builder.connections.types import (
     BasicAuthCredentials,
     BearerTokenAuthCredentials,
     APIKeyAuthCredentials,
-    # OAuth2AuthCodeCredentials,
+    OAuth2AuthCodeCredentials,
     OAuth2ClientCredentials,
     # OAuth2ImplicitCredentials,
-    # OAuth2PasswordCredentials,
+    OAuth2PasswordCredentials,
     OAuthOnBehalfOfCredentials,
     KeyValueConnectionCredentials,
     CREDENTIALS,
     IdentityProviderCredentials,
-    OAUTH_CONNECTION_TYPES, OAuth2AuthCodeCredentials
+    OAUTH_CONNECTION_TYPES
 
 )
 
@@ -115,7 +115,7 @@ def _format_token_headers(header_list: List) -> dict:
 
 def _validate_connection_params(type: ConnectionType, **args) -> None:
 
-    if type == ConnectionType.BASIC_AUTH and (
+    if type in {ConnectionType.BASIC_AUTH, ConnectionType.OAUTH2_PASSWORD} and (
             args.get('username') is None or args.get('password') is None
     ):
         raise typer.BadParameter(
@@ -136,7 +136,7 @@ def _validate_connection_params(type: ConnectionType, **args) -> None:
             f"Missing flags --api-key is required for type {type}"
         )
 
-    if type in {ConnectionType.OAUTH2_CLIENT_CREDS, ConnectionType.OAUTH2_AUTH_CODE} and args.get('client_secret') is None:
+    if type in {ConnectionType.OAUTH2_CLIENT_CREDS, ConnectionType.OAUTH2_AUTH_CODE, ConnectionType.OAUTH2_PASSWORD} and args.get('client_secret') is None:
         raise typer.BadParameter(
             f"Missing flags --client-secret is required for type {type}"
         )
@@ -146,14 +146,14 @@ def _validate_connection_params(type: ConnectionType, **args) -> None:
             f"Missing flags --auth-url is required for type {type}"
         )
 
-    if type in {ConnectionType.OAUTH_ON_BEHALF_OF_FLOW, ConnectionType.OAUTH2_CLIENT_CREDS, ConnectionType.OAUTH2_AUTH_CODE} and (
+    if type in {ConnectionType.OAUTH_ON_BEHALF_OF_FLOW, ConnectionType.OAUTH2_CLIENT_CREDS, ConnectionType.OAUTH2_AUTH_CODE, ConnectionType.OAUTH2_PASSWORD} and (
             args.get('client_id') is None
     ):
         raise typer.BadParameter(
             f"Missing flags --client-id is required for type {type}"
         )
     
-    if type in {ConnectionType.OAUTH_ON_BEHALF_OF_FLOW, ConnectionType.OAUTH2_CLIENT_CREDS, ConnectionType.OAUTH2_AUTH_CODE} and (
+    if type in {ConnectionType.OAUTH_ON_BEHALF_OF_FLOW, ConnectionType.OAUTH2_CLIENT_CREDS, ConnectionType.OAUTH2_AUTH_CODE, ConnectionType.OAUTH2_PASSWORD} and (
             args.get('token_url') is None
     ):
         raise typer.BadParameter(
@@ -209,13 +209,11 @@ def _get_credentials(type: ConnectionType, **kwargs):
         #         authorization_url=kwargs.get("auth_url"),
         #         client_id=kwargs.get("client_id"),
         #     )
-        # case ConnectionType.OAUTH2_PASSWORD:
-        #     return OAuth2PasswordCredentials(
-        #         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_PASSWORD:
+            keys = ["username", "password", "client_id","client_secret","token_url","grant_type", "scope"]
+            filtered_args = { key_name: kwargs[key_name] for key_name in keys if kwargs.get(key_name) }
+            return OAuth2PasswordCredentials(**filtered_args)
+        
         case ConnectionType.OAUTH_ON_BEHALF_OF_FLOW:
             return OAuthOnBehalfOfCredentials(
                 client_id=kwargs.get("client_id"),
@@ -283,25 +281,24 @@ def add_configuration(config: ConnectionConfiguration) -> None:
         logger.error(response_text)
         exit(1)
 
-def add_credentials(app_id: str, environment: ConnectionEnvironment, use_app_credentials: bool, credentials: CREDENTIALS) -> None:
+def add_credentials(app_id: str, environment: ConnectionEnvironment, use_app_credentials: bool, credentials: CREDENTIALS, payload: dict = None) -> None:
     client = get_connections_client()
     try:
         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)
-            }
-        else:
-            payload = {
-                "runtime_credentials": credentials.model_dump(exclude_none=True)
-            }
+        if not payload:
+            if use_app_credentials:
+                payload = {
+                    "app_credentials": credentials.model_dump(exclude_none=True)
+                }
+            else:
+                payload = {
+                    "runtime_credentials": credentials.model_dump(exclude_none=True)
+                }
         
-        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_app_credentials=use_app_credentials, payload=payload)
         else:
             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
         response_text = response.text
@@ -489,7 +486,20 @@ def set_credentials_connection(
     _validate_connection_params(type=conn_type, **kwargs)
     credentials = _get_credentials(type=conn_type, **kwargs)
 
-    add_credentials(app_id=app_id, environment=environment, use_app_credentials=use_app_credentials, credentials=credentials)
+    # Special handling for oauth2 password flow as it sends both app_creds and runtime_creds
+    logger.info(f"Setting credentials for environment '{environment}' on connection '{app_id}'")
+    if conn_type == ConnectionType.OAUTH2_PASSWORD:
+        credentials_model = credentials.model_dump(exclude_none=True)
+        runtime_cred_keys = {"username", "password"}
+        app_creds = {"app_credentials": {k: credentials_model[k] for k in credentials_model if k not in runtime_cred_keys}}
+        runtime_creds = {"runtime_credentials": {k: credentials_model[k] for k in credentials_model if k in runtime_cred_keys}}
+
+        add_credentials(app_id=app_id, environment=environment, use_app_credentials=True, credentials=credentials, payload=app_creds)
+        add_credentials(app_id=app_id, environment=environment, use_app_credentials=False, credentials=credentials, payload=runtime_creds)
+    else:
+        add_credentials(app_id=app_id, environment=environment, use_app_credentials=use_app_credentials, credentials=credentials)
+    
+    logger.info(f"Credentials successfully set for '{environment}' environment of connection '{app_id}'")
 
 def set_identity_provider_connection(
     app_id: str,

ibm_watsonx_orchestrate/cli/commands/knowledge_bases/knowledge_bases_controller.py

@@ -13,6 +13,7 @@ from ibm_watsonx_orchestrate.client.knowledge_bases.knowledge_base_client import
 from ibm_watsonx_orchestrate.client.base_api_client import ClientAPIException
 from ibm_watsonx_orchestrate.client.connections import get_connections_client
 from ibm_watsonx_orchestrate.client.utils import instantiate_client
+from ibm_watsonx_orchestrate.agent_builder.knowledge_bases.types import FileUpload
 
 logger = logging.getLogger(__name__)
 
@@ -43,7 +44,8 @@ def parse_file(file: str) -> List[KnowledgeBase]:
 def to_column_name(col: str):
     return " ".join([word.capitalize() if not word[0].isupper() else word for word in col.split("_")])
 
-def get_file_name(path: str):
+def get_file_name(file: str | FileUpload):
+    path = file.path if isinstance(file, FileUpload) else file
     # This name prettifying currently screws up file type detection on ingestion
     # return to_column_name(path.split("/")[-1].split(".")[0]) 
     return path.split("/")[-1]
@@ -55,7 +57,11 @@ def get_relative_file_path(path, dir):
         return f"{dir}{path.removeprefix('.')}"
     else:
         return f"{dir}/{path}"
-
+    
+def build_file_object(file_dir: str, file: str | FileUpload):
+    if isinstance(file, FileUpload):
+        return ('files', (get_file_name(file.path), open(get_relative_file_path(file.path, file_dir), 'rb')))
+    return ('files', (get_file_name(file), open(get_relative_file_path(file, file_dir), 'rb')))
 
 class KnowledgeBaseController:
     def __init__(self):
@@ -72,6 +78,21 @@ class KnowledgeBaseController:
 
         knowledge_bases = parse_file(file=file)
         
+        if app_id:
+            connections_client = get_connections_client()
+            connection_id = None
+                
+            connections = connections_client.get_draft_by_app_id(app_id=app_id)
+            if not connections:
+                logger.error(f"No connection exists with the app-id '{app_id}'")
+                exit(1)
+
+            connection_id = connections.connection_id
+
+            for kb in knowledge_bases:
+                if kb.conversational_search_tool and kb.conversational_search_tool.index_config and len(kb.conversational_search_tool.index_config) > 0:
+                    kb.conversational_search_tool.index_config[0].connection_id = connection_id
+        
         existing_knowledge_bases = client.get_by_names([kb.name for kb in knowledge_bases])
         
         for kb in knowledge_bases:
@@ -86,32 +107,32 @@ class KnowledgeBaseController:
 
                 kb.validate_documents_or_index_exists()
                 if kb.documents:
-                    files = [('files', (get_file_name(file_path), open(get_relative_file_path(file_path, file_dir), 'rb'))) for file_path in kb.documents]
+                    files = [build_file_object(file_dir, file) for file in kb.documents]
+                    file_urls = { get_file_name(file): file.url for file in kb.documents if isinstance(file, FileUpload) and file.url }
                     
                     kb.prioritize_built_in_index = True
                     payload = kb.model_dump(exclude_none=True);
                     payload.pop('documents');
 
-                    client.create_built_in(payload=payload, files=files)
+                    data = {
+                        'knowledge_base': json.dumps(payload),
+                        'file_urls': json.dumps(file_urls)
+                    }
+
+                    client.create_built_in(payload=data, files=files)
                 else:
                     if len(kb.conversational_search_tool.index_config) != 1:
                         raise ValueError(f"Must provide exactly one conversational_search_tool.index_config. Provided {len(kb.conversational_search_tool.index_config)}.")
                     
-                    
-                    if app_id:
-                        connections_client = get_connections_client()
-                        connection_id = None
-                        if app_id is not None:
-                            connections = connections_client.get_draft_by_app_id(app_id=app_id)
-                            if not connections:
-                                logger.error(f"No connection exists with the app-id '{app_id}'")
-                                exit(1)
-
-                            connection_id = connections.connection_id
-                            kb.conversational_search_tool.index_config[0].connection_id = connection_id
+                    if (kb.conversational_search_tool.index_config[0].milvus or \
+                        kb.conversational_search_tool.index_config[0].elastic_search) and \
+                            not kb.conversational_search_tool.index_config[0].connection_id:
+                        raise ValueError(f"Must provide credentials (via --app-id) when using milvus or elastic_search.")
 
                     kb.prioritize_built_in_index = False
-                    client.create(payload=kb.model_dump(exclude_none=True))
+                    data = { 'knowledge_base': json.dumps(kb.model_dump(exclude_none=True)) }
+
+                    client.create(payload=data)
                 
                 logger.info(f"Successfully imported knowledge base '{kb.name}'")
             except ClientAPIException as e:
@@ -144,8 +165,8 @@ class KnowledgeBaseController:
             existing_docs = [doc.get("metadata", {}).get("original_file_name", "") for doc in status.get("documents", [])]
             
             removed_docs = existing_docs[:]
-            for filepath in kb.documents:
-                filename = get_file_name(filepath)
+            for file in kb.documents:
+                filename = get_file_name(file)
 
                 if filename in existing_docs:
                     logger.warning(f'Document \"{filename}\" already exists in knowledge base. Updating...')
@@ -155,17 +176,25 @@ class KnowledgeBaseController:
                 logger.warning(f'Document \"{filename}\" removed from knowledge base.')
 
 
-            files = [('files', (get_file_name(file_path), open(get_relative_file_path(file_path, file_dir), 'rb'))) for file_path in kb.documents]
+            files = [build_file_object(file_dir, file) for file in kb.documents]
+            file_urls = { get_file_name(file): file.url for file in kb.documents if isinstance(file, FileUpload) and file.url }
             
             kb.prioritize_built_in_index = True
             payload = kb.model_dump(exclude_none=True);
             payload.pop('documents');
 
-            self.get_client().update_with_documents(knowledge_base_id, payload=payload, files=files)
+            data = {
+                'knowledge_base': json.dumps(payload),
+                'file_urls': json.dumps(file_urls)
+            }
+
+            self.get_client().update_with_documents(knowledge_base_id, payload=data, files=files)
         else:
             if kb.conversational_search_tool and kb.conversational_search_tool.index_config:
                 kb.prioritize_built_in_index = False
-            self.get_client().update(knowledge_base_id, kb.model_dump(exclude_none=True))
+
+            data = { 'knowledge_base': json.dumps(kb.model_dump(exclude_none=True)) }
+            self.get_client().update(knowledge_base_id, payload=data)
 
         logger.info(f"Knowledge base '{kb.name}' updated successfully")