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

ibm_watsonx_orchestrate/cli/commands/knowledge_bases/knowledge_bases_controller.py

@@ -49,6 +49,15 @@ def get_file_name(path: str):
     # return to_column_name(path.split("/")[-1].split(".")[0]) 
     return path.split("/")[-1]
 
+def get_relative_file_path(path, dir):
+    if path.startswith("/"):
+        return path
+    elif path.startswith("./"):
+        return f"{dir}{path.removeprefix('.')}"
+    else:
+        return f"{dir}/{path}"
+
+
 class KnowledgeBaseController:
     def __init__(self):
         self.client = None
@@ -68,8 +77,9 @@ class KnowledgeBaseController:
                 kb.validate_documents_or_index_exists()
                 if kb.documents:
                     file_dir = "/".join(file.split("/")[:-1])
-                    files = [('files', (get_file_name(file_path), open(file_path if file_path.startswith("/") else (file_path if not file_dir else f"{file_dir}/{file_path}"), 'rb'))) for file_path in 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]
                     
+                    kb.prioritize_built_in_index = True
                     payload = kb.model_dump(exclude_none=True);
                     payload.pop('documents');
 
@@ -78,6 +88,7 @@ class KnowledgeBaseController:
                     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
@@ -90,6 +101,7 @@ class KnowledgeBaseController:
                             connection_id = connections.connection_id
                             kb.conversational_search_tool.index_config[0].connection_id = connection_id
 
+                    kb.prioritize_built_in_index = False
                     client.create(payload=kb.model_dump(exclude_none=True))
                 
                 logger.info(f"Successfully imported knowledge base '{kb.name}'")
@@ -126,13 +138,16 @@ class KnowledgeBaseController:
 
         if update_request.documents:
             file_dir = "/".join(file.split("/")[:-1])
-            files = [('files', (get_file_name(file_path), open(file_path if file_path.startswith("/") else f"{file_dir}/{file_path}", 'rb'))) for file_path in update_request.documents]
+            files = [('files', (get_file_name(file_path), open(get_relative_file_path(file_path, file_dir), 'rb'))) for file_path in update_request.documents]
             
+            update_request.prioritize_built_in_index = True
             payload = update_request.model_dump(exclude_none=True);
             payload.pop('documents');
 
             self.get_client().update_with_documents(knowledge_base_id, payload=payload, files=files)
         else:
+            if update_request.conversational_search_tool and update_request.conversational_search_tool.index_config:
+                update_request.prioritize_built_in_index = False
             self.get_client().update(knowledge_base_id, update_request.model_dump(exclude_none=True))
 
         logEnding = f"with ID '{id}'" if id else f"'{name}'"

ibm_watsonx_orchestrate/cli/commands/models/env_file_model_provider_mapper.py

@@ -0,0 +1,180 @@
+# https://portkey.ai/
+from dotenv import dotenv_values
+
+from ibm_watsonx_orchestrate.client.models.types import ProviderConfig, ModelProvider
+
+def _bpp(name: str, provider_config=None) -> dict:
+    if provider_config is None:
+        provider_config = {}
+
+    name = name.upper()
+
+    provider_config.update({
+        f"{name}_API_KEY": 'api_key',
+        f"{name}_CUSTOM_HOST": 'custom_host',
+        f"{name}_URL_TO_FETCH": 'url_to_fetch',
+        f"{name}_FORWARD_HEADERS": lambda value, config: setattr(config, 'forward_headers', list(map(lambda v: v.strip(), value.split(';')))),
+        f"{name}_REQUEST_TIMEOUT": lambda value, config: setattr(config, 'request_timeout', int(value)),
+        f"{name}_TRANSFORM_TO_FORM_DATA": 'transform_to_form_data'
+    })
+
+    return provider_config
+
+# model provider prefix => ENV_VAR key => provider config key
+PROVIDER_PROPERTIES_LUT = {
+    ModelProvider.OPENAI: _bpp('OPENAI', {}),
+    # ModelProvider.A21: _bpp('A21', {}),
+    ModelProvider.ANTHROPIC: _bpp('ANTHROPIC', {
+        'ANTHROPIC_BETA': 'anthropic_beta',
+        'ANTHROPIC_VERSION': 'anthropic_version'
+    }),
+    # ModelProvider.ANYSCALE: _bpp('ANYSCALE', {}),
+    # ModelProvider.AZURE_OPENAI: _bpp('AZURE_OPENAI', {}),
+    # ModelProvider.AZURE_AI: _bpp('AZURE', {
+    #     'AZURE_AI_RESOURCE_NAME': 'azure_resource_name',
+    #     'AZURE_AI_DEPLOYMENT_ID': 'azure_deployment_id',
+    #     'AZURE_AI_API_VERSION': 'azure_api_version',
+    #     'AZURE_AI_AD_AUTH': 'ad_auth',
+    #     'AZURE_AI_AUTH_MODE': 'azure_auth_mode',
+    #     'AZURE_AI_MANAGED_CLIENT_ID': 'azure_managed_client_id',
+    #     'AZURE_AI_ENTRA_CLIENT_ID': 'azure_entra_client_id',
+    #     'AZURE_AI_ENTRA_CLIENT_SECRET': 'azure_entra_client_secret',
+    #     'AZURE_AI_ENTRA_TENANT_ID': 'azure_entra_tenant_id',
+    #     'AZURE_AI_AD_TOKEN': 'azure_ad_token',
+    #     'AZURE_AI_MODEL_NAME': 'azure_model_name',
+    # }),
+    # ModelProvider.BEDROCK: _bpp('BEDROCK', {
+    #     'AWS_SECRET_ACCESS_KEY': 'aws_secret_access_key',
+    #     'AWS_ACCESS_KEY_ID': 'aws_access_key_id',
+    #     'AWS_SESSION_TOKEN': 'aws_session_token',
+    #     'AWS_REGION': 'aws_region',
+    #     'AWS_AUTH_TYPE': 'aws_auth_type',
+    #     'AWS_ROLE_ARN': 'aws_role_arn',
+    #     'AWS_EXTERNAL_ID': 'aws_external_id',
+    #     'AWS_S3_BUCKET': 'aws_s3_bucket',
+    #     'AWS_S3_OBJECT_KEY': 'aws_s3_object_key',
+    #     'AWS_BEDROCK_MODEL': 'aws_bedrock_model',
+    #     'AWS_SERVER_SIDE_ENCRYPTION': 'aws_server_side_encryption',
+    #     'AWS_SERVER_SIDE_ENCRYPTION_KMS_KEY_ID': 'aws_server_side_encryption_kms_key_id',
+    # }),
+    # ModelProvider.CEREBRAS: _bpp('COHERE', {}),
+    # ModelProvider.COHERE: _bpp('COHERE', {}),
+    ModelProvider.GOOGLE: _bpp('GOOGLE', {}),
+    # ModelProvider.VERTEX_AI: _bpp('GOOGLE_VERTEX_AI', {
+    #     'GOOGLE_VERTEX_AI_REGION': 'vertex_region',
+    #     'GOOGLE_VERTEX_AI_PROJECT_ID': 'vertex_project_id',
+    #     'GOOGLE_VERTEX_AI_SERVICE_ACCOUNT_JSON': 'vertex_service_account_json',
+    #     'GOOGLE_VERTEX_AI_STORAGE_BUCKET_NAME': 'vertex_storage_bucket_name',
+    #     'GOOGLE_VERTEX_AI_MODEL_NAME': 'vertex_model_name',
+    #     'GOOGLE_VERTEX_AI_FILENAME': 'filename'
+    # }),
+    # ModelProvider.GROQ: _bpp('GROQ', {}),
+    # ModelProvider.HUGGINGFACE: _bpp('HUGGINGFACE', {
+    #     'HUGGINGFACE_BASE_URL': 'huggingfaceBaseUrl'
+    # }),
+    ModelProvider.MISTRAL_AI: _bpp('MISTRAL', {
+        'MISTRAL_FIM_COMPLETION': 'mistral_fim_completion'
+    }),
+    # ModelProvider.JINA: _bpp('JINA', {}),
+    ModelProvider.OLLAMA: _bpp('OLLAMA', {}),
+    ModelProvider.OPENROUTER: _bpp('OPENROUTER', {}),
+    # ModelProvider.STABILITY_AI: _bpp('STABILITY', {
+    #     'STABILITY_CLIENT_ID': 'stability_client_id',
+    #     'STABILITY_CLIENT_USER_ID': 'stability_client_user_id',
+    #     'STABILITY_CLIENT_VERSION': 'stability_client_version'
+    # }),
+    # ModelProvider.TOGETHER_AI: _bpp('TOGETHER_AI', {}),
+    ModelProvider.WATSONX: _bpp('WATSONX', {
+        'WATSONX_VERSION': 'watsonx_version',
+        'WATSONX_SPACE_ID': 'watsonx_space_id',
+        'WATSONX_PROJECT_ID': 'watsonx_project_id',
+        'WATSONX_APIKEY': 'api_key'
+    })
+
+    # 'palm': _bpp('PALM', {}),
+    # 'nomic': _bpp('NOMIC', {}),
+    # 'perplexity-ai': _bpp('PERPLEXITY_AI', {}),
+    # 'segmind': _bpp('SEGMIND', {}),
+    # 'deepinfra': _bpp('DEEPINFRA', {}),
+    # 'novita-ai': _bpp('NOVITA_AI', {}),
+    # 'fireworks-ai': _bpp('FIREWORKS',{
+    #     'FIREWORKS_ACCOUNT_ID': 'fireworks_account_id'
+    # }),
+    # 'deepseek': _bpp('DEEPSEEK', {}),
+    # 'voyage': _bpp('VOYAGE', {}),
+    # 'moonshot': _bpp('MOONSHOT', {}),
+    # 'lingyi': _bpp('LINGYI', {}),
+    # 'zhipu': _bpp('ZHIPU', {}),
+    # 'monsterapi': _bpp('MONSTERAPI', {}),
+    # 'predibase': _bpp('PREDIBASE', {}),
+
+    # 'github': _bpp('GITHUB', {}),
+    # 'deepbricks': _bpp('DEEPBRICKS', {}),
+    # 'sagemaker': _bpp('AMZN_SAGEMAKER', {
+    #     'AMZN_SAGEMAKER_CUSTOM_ATTRIBUTES': 'amzn_sagemaker_custom_attributes',
+    #     'AMZN_SAGEMAKER_TARGET_MODEL': 'amzn_sagemaker_target_model',
+    #     'AMZN_SAGEMAKER_TARGET_VARIANT': 'amzn_sagemaker_target_variant',
+    #     'AMZN_SAGEMAKER_TARGET_CONTAINER_HOSTNAME': 'amzn_sagemaker_target_container_hostname',
+    #     'AMZN_SAGEMAKER_INFERENCE_ID': 'amzn_sagemaker_inference_id',
+    #     'AMZN_SAGEMAKER_ENABLE_EXPLANATIONS': 'amzn_sagemaker_enable_explanations',
+    #     'AMZN_SAGEMAKER_INFERENCE_COMPONENT': 'amzn_sagemaker_inference_component',
+    #     'AMZN_SAGEMAKER_SESSION_ID': 'amzn_sagemaker_session_id',
+    #     'AMZN_SAGEMAKER_MODEL_NAME': 'amzn_sagemaker_model_name'
+    # }),
+    # '@cf': _bpp('WORKERS_AI', { # workers ai
+    #     'WORKERS_AI_ACCOUNT_ID': 'workers_ai_account_id'
+    # }),
+    # 'snowflake': _bpp('SNOWFLAKE', { # no provider prefix found
+    #     'SNOWFLAKE_ACCOUNT': 'snowflake_account'
+    # })
+}
+PROVIDER_PROPERTIES_RLUT= {}
+for provider in PROVIDER_PROPERTIES_LUT.keys():
+    PROVIDER_PROPERTIES_RLUT[provider] = {v:k for k,v in PROVIDER_PROPERTIES_LUT[provider].items()}
+
+PROVIDER_LUT = {k:k for k in PROVIDER_PROPERTIES_LUT.keys() }
+PROVIDER_LUT.update({
+    # any overrides for the provider prefix to provider name can be provided here on PROVIDER_LUT
+})
+
+
+
+PROVIDER_REQUIRED_FIELDS = {k:['api_key'] if k not in ['ollama']  else [] for k in PROVIDER_PROPERTIES_LUT.keys()}
+PROVIDER_REQUIRED_FIELDS.update({
+    # Mark the required fields for a provider
+})
+
+def env_file_to_model_ProviderConfig(model_name: str, env_file_path: str) -> ProviderConfig | None:
+    provider = next(filter(lambda x: x not in ('virtual-policy', 'virtual-model'), model_name.split('/')))
+    if provider not in PROVIDER_LUT:
+        raise ValueError(f"Unsupported model provider {provider}")
+
+    values = dotenv_values(str(env_file_path))
+
+    if values is None:
+        raise ValueError(f"No provider configuration in env file {env_file_path}")
+
+    cfg = ProviderConfig()
+    cfg.provider = PROVIDER_LUT[provider]
+
+    cred_lut = PROVIDER_PROPERTIES_LUT[provider]
+
+
+    consumed_credentials = []
+    for key, value in values.items():
+        if key in cred_lut:
+            k = cred_lut[key]
+            consumed_credentials.append(k)
+            setattr(cfg, k, value)
+
+
+    required_creds = PROVIDER_REQUIRED_FIELDS[provider]
+    missing_credentials = []
+    for cred in required_creds:
+        if cred not in consumed_credentials:
+            missing_credentials.append(cred)
+
+    if len(missing_credentials) > 0:
+        raise ValueError(f"Missing environment variable(s) {', '.join(map(lambda c: PROVIDER_PROPERTIES_RLUT[provider][c], missing_credentials))} required for the provider {provider}")
+
+    return cfg

ibm_watsonx_orchestrate/cli/commands/models/models_command.py

@@ -1,30 +1,42 @@
 import logging
 import os
-import requests
 import sys
-import rich.highlighter
-import typer
+from typing import List
+
+import requests
 import rich
+import rich.highlighter
 import typer
 from typing_extensions import Annotated
+
 from ibm_watsonx_orchestrate.cli.commands.server.server_command import get_default_env_file, merge_env
+from ibm_watsonx_orchestrate.client.model_policies.model_policies_client import ModelPoliciesClient
+from ibm_watsonx_orchestrate.client.model_policies.types import ModelPolicy, ModelPolicyInner, \
+    ModelPolicyRetry, ModelPolicyStrategy, ModelPolicyStrategyMode, ModelPolicyTarget
+from ibm_watsonx_orchestrate.client.models.models_client import ModelsClient
+from ibm_watsonx_orchestrate.client.models.types import CreateVirtualModel, ANTHROPIC_DEFAULT_MAX_TOKENS
+from ibm_watsonx_orchestrate.client.utils import instantiate_client
 
 logger = logging.getLogger(__name__)
 models_app = typer.Typer(no_args_is_help=True)
+models_policy_app = typer.Typer(no_args_is_help=True)
+# models_app.add_typer(models_policy_app, name='policy', help='Add or remove pseudo models which route traffic between multiple downstream models')
 
 WATSONX_URL = os.getenv("WATSONX_URL")
 
 class ModelHighlighter(rich.highlighter.RegexHighlighter):
     base_style = "model."
-    highlights = [r"(?P<name>watsonx\/.+\/.+):"]
+    highlights = [r"(?P<name>(watsonx|virtual[-]model|virtual[-]policy)\/.+\/.+):"]
 
-@models_app.command(name="list")
+@models_app.command(name="list", help="List available models")
 def model_list(
     print_raw: Annotated[
         bool,
         typer.Option("--raw", "-r", help="Display the list of models in a non-tabular format"),
     ] = False,
 ):
+    models_client: ModelsClient = instantiate_client(ModelsClient)
+    model_policies_client: ModelPoliciesClient = instantiate_client(ModelPoliciesClient)
     global WATSONX_URL
     default_env_path = get_default_env_file()
     merged_env_dict = merge_env(
@@ -40,9 +52,18 @@ def model_list(
         logger.error("Error: WATSONX_URL is required in the environment.")
         sys.exit(1)
 
+    logger.info("Retrieving virtual-model models list...")
+    virtual_models = models_client.list()
+
+
+
+    logger.info("Retrieving virtual-policies models list...")
+    virtual_model_policies = model_policies_client.list()
+
     logger.info("Retrieving watsonx.ai models list...")
     found_models = _get_wxai_foundational_models()
 
+
     preferred_str = merged_env_dict.get('PREFERRED_MODELS', '')
     incompatible_str = merged_env_dict.get('INCOMPATIBLE_MODELS', '') 
 
@@ -64,7 +85,7 @@ def model_list(
                 continue
             filtered_models.append(model)
         
-        # Sort to put preferred first
+        # Sort to put the preferred first
         def sort_key(model):
             model_id = model.get("model_id", "").lower()
             is_preferred = any(pref in model_id for pref in preferred_list)
@@ -76,6 +97,13 @@ def model_list(
             theme = rich.theme.Theme({"model.name": "bold cyan"})
             console = rich.console.Console(highlighter=ModelHighlighter(), theme=theme)
             console.print("[bold]Available Models:[/bold]")
+
+            for model in virtual_models:
+                console.print(f"- ✨️ {model.name}:", model.description or 'No description provided.')
+
+            for model in virtual_model_policies:
+                console.print(f"- ✨️ {model.name}:", 'No description provided.')
+
             for model in sorted_models:
                 model_id = model.get("model_id", "N/A")
                 short_desc = model.get("short_description", "No description provided.")
@@ -83,17 +111,23 @@ def model_list(
                 marker = "★ " if any(pref in model_id.lower() for pref in preferred_list) else ""
                 console.print(f"- [yellow]{marker}[/yellow]{full_model_name}")
 
-            console.print("[yellow]★[/yellow] [italic dim]indicates a supported and preferred model[/italic dim]" )
+            console.print("[yellow]★[/yellow] [italic dim]indicates a supported and preferred model[/italic dim]\n[blue dim]✨️[/blue dim] [italic dim]indicates a model from a custom provider[/italic dim]" )
         else:
             table = rich.table.Table(
                 show_header=True,
                 title="[bold]Available Models[/bold]",
-                caption="[yellow]★[/yellow] indicates a supported and preferred model",
+                caption="[yellow]★ [/yellow] indicates a supported and preferred model from watsonx\n[blue]✨️[/blue] indicates a model from a custom provider",
                 show_lines=True)
             columns = ["Model", "Description"]
             for col in columns:
                 table.add_column(col)
 
+            for model in virtual_models:
+                table.add_row(f"✨️ {model.name}", model.description or 'No description provided.')
+
+            for model in virtual_model_policies:
+                table.add_row(f"✨️ {model.name}", 'No description provided.')
+
             for model in sorted_models:
                 model_id = model.get("model_id", "N/A")
                 short_desc = model.get("short_description", "No description provided.")
@@ -102,6 +136,158 @@ def model_list(
         
             rich.print(table)
 
+
+
+@models_app.command(name="add", help="Add an llm from a custom provider")
+def models_add(
+    name: Annotated[
+        str,
+        typer.Option("--name", "-n", help="The name of the model to add"),
+    ],
+    env_file: Annotated[
+        str,
+        typer.Option('--env-file', '-e', help='The path to an .env file containing the credentials for your llm provider'),
+    ],
+    description: Annotated[
+        str,
+        typer.Option('--description', '-d', help='The description of the model to add'),
+    ] = None,
+    display_name: Annotated[
+        str,
+        typer.Option('--display-name', help='What name should this llm appear as within the ui'),
+    ] = None,
+
+):
+    from ibm_watsonx_orchestrate.cli.commands.models.env_file_model_provider_mapper import env_file_to_model_ProviderConfig # lazily import this because the lut building is expensive
+
+    models_client: ModelsClient = instantiate_client(ModelsClient)
+    provider_config = env_file_to_model_ProviderConfig(model_name=name, env_file_path=env_file)
+    if not name.startswith('virtual-model/'):
+        name = f"virtual-model/{name}"
+    
+    config=None
+    # Anthropic has no default for max_tokens
+    if "anthropic" in name:
+        config = {
+            "max_tokens": ANTHROPIC_DEFAULT_MAX_TOKENS
+        }
+
+    model = CreateVirtualModel(
+        name=name,
+        display_name=display_name or name,
+        description=description,
+        tags=[],
+        provider_config=provider_config,
+        config=config
+    )
+
+    models_client.create(model)
+    logger.info(f"Successfully added the model '{name}'")
+
+
+
+@models_app.command(name="remove", help="Remove an llm from a custom provider")
+def models_remove(
+        name: Annotated[
+            str,
+            typer.Option("--name", "-n", help="The name of the model to remove"),
+        ]
+):
+    models_client: ModelsClient = instantiate_client(ModelsClient)
+    models = models_client.list()
+    model = next(filter(lambda x: x.name == name or x.name == f"virtual-model/{name}", models), None)
+    if not model:
+        logger.error(f"No model found with the name '{name}'")
+        sys.exit(1)
+    
+    models_client.delete(model_id=model.id)
+    logger.info(f"Successfully removed the model '{name}'")
+        
+
+# @models_policy_app.command(name='add', help='Add a model policy')
+# def models_policy_add(
+#         name: Annotated[
+#             str,
+#             typer.Option("--name", "-n", help="The name of the model to remove"),
+#         ],
+#         models: Annotated[
+#             List[str],
+#             typer.Option('--model', '-m', help='The name of the model to add'),
+#         ],
+#         strategy: Annotated[
+#             ModelPolicyStrategyMode,
+#             typer.Option('--strategy', '-s', help='How to spread traffic across models'),
+#         ],
+#         strategy_on_code: Annotated[
+#             List[int],
+#             typer.Option('--strategy-on-code', help='The http status to consider invoking the strategy'),
+#         ],
+#         retry_on_code: Annotated[
+#             List[int],
+#             typer.Option('--retry-on-code', help='The http status to consider retrying the llm call'),
+#         ],
+#         retry_attempts: Annotated[
+#             int,
+#             typer.Option('--retry-attempts', help='The number of attempts to retry'),
+#         ],
+#         display_name: Annotated[
+#             str,
+#             typer.Option('--display-name', help='What name should this llm appear as within the ui'),
+#         ] = None
+# ):
+#     model_policies_client: ModelPoliciesClient = instantiate_client(ModelPoliciesClient)
+#     model_client: ModelsClient = instantiate_client(ModelsClient)
+#     model_lut = {m.name: m.id for m in model_client.list()}
+#     for m in models:
+#         if m not in model_lut:
+#             logger.error(f"No model found with the name '{m}'")
+#             exit(1)
+
+#     inner = ModelPolicyInner()
+#     inner.strategy = ModelPolicyStrategy(
+#         mode=strategy,
+#         on_status_codes=strategy_on_code
+#     )
+#     inner.targets = [ModelPolicyTarget(model_id=model_lut[m], weight=1) for m in models]
+#     if retry_on_code:
+#         inner.retry = ModelPolicyRetry(
+#             on_status_codes=retry_on_code,
+#             attempts=retry_attempts
+#         )
+
+#     if not display_name:
+#         display_name = name
+
+
+#     policy = ModelPolicy(
+#         name=name,
+#         display_name=display_name,
+#         policy=inner
+#     )
+#     model_policies_client.create(policy)
+#     logger.info(f"Successfully added the model policy '{name}'")
+
+
+
+# @models_policy_app.command(name='remove', help='Remove a model policy')
+# def models_policy_remove(
+#         name: Annotated[
+#             str,
+#             typer.Option("--name", "-n", help="The name of the model policy to remove"),
+#         ]
+# ):
+#     model_policies_client: ModelPoliciesClient = instantiate_client(ModelPoliciesClient)
+#     model_policies = model_policies_client.list()
+
+#     policy = next(filter(lambda x: x.name == name or x.name == f"virtual-policy/{name}", model_policies), None)
+#     if not policy:
+#         logger.error(f"No model found with the name '{name}'")
+#         exit(1)
+
+#     model_policies_client.delete(model_policy_id=policy.id)
+#     logger.info(f"Successfully removed the model '{name}'")
+
+
 def _get_wxai_foundational_models():
     foundation_models_url = WATSONX_URL + "/ml/v1/foundation_model_specs?version=2024-05-01"