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

ibm_watsonx_orchestrate/__init__.py

@@ -0,0 +1,28 @@
+#  -----------------------------------------------------------------------------------------
+#  (C) Copyright IBM Corp. 2023-2024.
+#  https://opensource.org/licenses/BSD-3-Clause
+#  -----------------------------------------------------------------------------------------
+
+pkg_name = "ibm-watsonx-orchestrate"
+
+__version__ = "1.0.0"
+
+
+
+try:
+    from importlib.metadata import version
+
+    ver = version(pkg_name)
+
+except (ModuleNotFoundError, AttributeError):
+    from importlib_metadata import version as imp_lib_ver
+
+    ver = imp_lib_ver(pkg_name)
+
+from ibm_watsonx_orchestrate.client.client import Client
+from ibm_watsonx_orchestrate.utils.logging.logger import setup_logging
+
+Client.version = ver
+__version__ = ver
+setup_logging()
+

ibm_watsonx_orchestrate/agent_builder/agents/__init__.py

@@ -0,0 +1,5 @@
+from .agent import Agent, AgentSpec
+from .external_agent import ExternalAgent, ExternalAgentSpec
+from .assistant_agent import AssistantAgent, AssistantAgentSpec
+# from .types import AgentKind, AgentStyle, ExternalAgentConfig, AssistantAgentConfig, SpecVersion
+from .types import AgentKind, AgentStyle, SpecVersion, ExternalAgentAuthScheme, AgentProvider, AssistantAgentConfig

ibm_watsonx_orchestrate/agent_builder/agents/agent.py

@@ -0,0 +1,27 @@
+import json
+from ibm_watsonx_orchestrate.utils.utils import yaml_safe_load
+from .types import AgentSpec
+
+
+class Agent(AgentSpec):
+
+    @staticmethod
+    def from_spec(file: str) -> 'Agent':
+        with open(file, 'r') as f:
+            if file.endswith('.yaml') or file.endswith('.yml'):
+                content = yaml_safe_load(f)
+            elif file.endswith('.json'):
+                content = json.load(f)
+            else:
+                raise ValueError('file must end in .json, .yaml, or .yml')
+            if not content.get("spec_version"):
+                raise ValueError(f"Field 'spec_version' not provided. Please ensure provided spec conforms to a valid spec format")
+            agent = Agent.model_validate(content)
+
+        return agent
+    
+    def __repr__(self):
+        return f"Agent(name='{self.name}', description='{self.description}')"
+
+    def __str__(self):
+        return self.__repr__()

ibm_watsonx_orchestrate/agent_builder/agents/assistant_agent.py

@@ -0,0 +1,28 @@
+import json
+from ibm_watsonx_orchestrate.utils.utils import yaml_safe_load
+from .types import AssistantAgentSpec
+
+
+class AssistantAgent(AssistantAgentSpec):
+
+    @staticmethod
+    def from_spec(file: str) -> 'AssistantAgent':
+        with open(file, 'r') as f:
+            if file.endswith('.yaml') or file.endswith('.yml'):
+                content = yaml_safe_load(f)
+            elif file.endswith('.json'):
+                content = json.load(f)
+            else:
+                raise ValueError('file must end in .json, .yaml, or .yml')
+            
+            if not content.get("spec_version"):
+                raise ValueError(f"Field 'spec_version' not provided. Please ensure provided spec conforms to a valid spec format")
+            agent = AssistantAgent.model_validate(content)
+
+        return agent
+    
+    def __repr__(self):
+        return f"AssistantAgent(name='{self.name}', description='{self.description}')"
+
+    def __str__(self):
+        return self.__repr__()

ibm_watsonx_orchestrate/agent_builder/agents/external_agent.py

@@ -0,0 +1,28 @@
+import json
+from ibm_watsonx_orchestrate.utils.utils import yaml_safe_load
+from .types import ExternalAgentSpec
+
+
+class ExternalAgent(ExternalAgentSpec):
+
+    @staticmethod
+    def from_spec(file: str) -> 'ExternalAgent':
+        with open(file, 'r') as f:
+            if file.endswith('.yaml') or file.endswith('.yml'):
+                content = yaml_safe_load(f)
+            elif file.endswith('.json'):
+                content = json.load(f)
+            else:
+                raise ValueError('file must end in .json, .yaml, or .yml')
+            
+            if not content.get("spec_version"):
+                raise ValueError(f"Field 'spec_version' not provided. Please ensure provided spec conforms to a valid spec format")
+            agent = ExternalAgent.model_validate(content)
+
+        return agent
+    
+    def __repr__(self):
+        return f"ExternalAgent(name='{self.name}', description='{self.description}')"
+
+    def __str__(self):
+        return self.__repr__()

ibm_watsonx_orchestrate/agent_builder/agents/types.py

@@ -0,0 +1,204 @@
+import json
+import yaml
+from enum import Enum
+from typing import List, Optional, Dict
+from pydantic import BaseModel, model_validator, ConfigDict
+from ibm_watsonx_orchestrate.agent_builder.tools import BaseTool
+from ibm_watsonx_orchestrate.agent_builder.knowledge_bases.types import KnowledgeBaseSpec
+from pydantic import Field, AliasChoices
+from typing import Annotated
+
+# TO-DO: this is just a placeholder. Will update this later to align with backend
+DEFAULT_LLM = "watsonx/meta-llama/llama-3-1-70b-instruct"
+
+class SpecVersion(str, Enum):
+    V1 = "v1"
+
+
+class AgentKind(str, Enum):
+    NATIVE = "native"
+    EXTERNAL = "external"
+    ASSISTANT = "assistant"
+
+class ExternalAgentAuthScheme(str, Enum):
+    BEARER_TOKEN = 'BEARER_TOKEN'
+    API_KEY = "API_KEY"
+    NONE = 'NONE'
+
+class AgentProvider(str, Enum):
+    WXAI = "wx.ai"
+    EXT_CHAT = "external_chat"
+    SALESFORCE = "salesforce"
+    WATSONX = "watsonx" #provider type returned from an assistant agent
+
+
+class AssistantAgentAuthType(str, Enum):
+    ICP_IAM = "ICP_IAM"
+    IBM_CLOUD_IAM = "IBM_CLOUD_IAM"
+    MCSP = "MCSP"
+    BEARER_TOKEN = "BEARER_TOKEN"
+    HIDDEN = "<hidden>"
+
+
+class BaseAgentSpec(BaseModel):
+    spec_version: SpecVersion = None
+    kind: AgentKind
+    id: Optional[Annotated[str, Field(json_schema_extra={"min_length_str": 1})]] = None
+    name: Annotated[str, Field(json_schema_extra={"min_length_str":1})]
+    description: Annotated[str, Field(json_schema_extra={"min_length_str":1})]
+
+    def dump_spec(self, file: str) -> None:
+        dumped = self.model_dump(mode='json', exclude_unset=True, exclude_none=True)
+        with open(file, 'w') as f:
+            if file.endswith('.yaml') or file.endswith('.yml'):
+                yaml.dump(dumped, f, sort_keys=False)
+            elif file.endswith('.json'):
+                json.dump(dumped, f, indent=2)
+            else:
+                raise ValueError('file must end in .json, .yaml, or .yml')
+
+    def dumps_spec(self) -> str:
+        dumped = self.model_dump(mode='json', exclude_none=True)
+        return json.dumps(dumped, indent=2)
+
+# ===============================
+#      NATIVE AGENT TYPES
+# ===============================
+
+class AgentStyle(str, Enum):
+    DEFAULT = "default"
+    REACT = "react"
+
+class AgentSpec(BaseAgentSpec):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    kind: AgentKind = AgentKind.NATIVE
+    llm: str = DEFAULT_LLM
+    style: AgentStyle = AgentStyle.DEFAULT
+    instructions: Annotated[Optional[str], Field(json_schema_extra={"min_length_str":1})] = None
+    collaborators: Optional[List[str]] | Optional[List['BaseAgentSpec']] = []
+    tools: Optional[List[str]] | Optional[List['BaseTool']] = []
+    hidden: bool = False
+    knowledge_base: Optional[List[str]] | Optional[List['KnowledgeBaseSpec']] = []
+
+
+    def __init__(self, *args, **kwargs):
+        if "tools" in kwargs and kwargs["tools"]:
+            kwargs["tools"] = [x.__tool_spec__.name if isinstance(x, BaseTool) else x for x in kwargs["tools"]]
+        if "collaborators" in kwargs and kwargs["collaborators"]:
+            kwargs["collaborators"] = [x.name if isinstance(x, BaseAgentSpec) else x for x in kwargs["collaborators"]]
+        super().__init__(*args, **kwargs)
+
+    @model_validator(mode="before")
+    def validate_fields(cls, values):
+        return validate_agent_fields(values)
+    
+    @model_validator(mode="after")
+    def validate_kind(self):
+        if self.kind != AgentKind.NATIVE:
+            raise ValueError(f"The specified kind '{self.kind}' cannot be used to create a native agent.")
+        return self
+
+def validate_agent_fields(values: dict) -> dict:
+    # Check for empty strings or whitespace
+    for field in ["id", "name", "kind", "description", "collaborators", "tools"]:
+        value = values.get(field)
+        if value and not str(value).strip():
+            raise ValueError(f"{field} cannot be empty or just whitespace")
+    
+    name = values.get("name")
+    collaborators = values.get("collaborators", [])  if values.get("collaborators", []) else []
+    for collaborator in collaborators:
+        if collaborator == name:
+            raise ValueError(f"Circular reference detected. The agent '{name}' cannot contain itself as a collaborator")
+
+
+    return values
+
+# ===============================
+#      EXTERNAL AGENT TYPES
+# ===============================
+
+class ExternalAgentConfig(BaseModel):
+    hidden: bool = False
+    enable_cot: bool = False
+
+class ExternalAgentSpec(BaseAgentSpec):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    kind: AgentKind = AgentKind.EXTERNAL
+    title: Annotated[str, Field(json_schema_extra={"min_length_str":1})]
+    tags: Optional[List[str]] = None
+    api_url: Annotated[str, Field(json_schema_extra={"min_length_str":1})]
+    auth_scheme: ExternalAgentAuthScheme = ExternalAgentAuthScheme.NONE
+    auth_config: dict = {}
+    provider: AgentProvider = AgentProvider.EXT_CHAT
+    chat_params: dict = None
+    config: ExternalAgentConfig = ExternalAgentConfig()
+    nickname: Annotated[str | None, Field(json_schema_extra={"min_length_str":1})] = None
+    app_id: Annotated[str | None, Field(json_schema_extra={"min_length_str":1})] = None
+    connection_id: Annotated[str | None, Field(json_schema_extra={"min_length_str":1})] = None
+
+    @model_validator(mode="before")
+    def validate_fields_for_external(cls, values):
+        return validate_external_agent_fields(values)
+
+    @model_validator(mode="after")
+    def validate_kind_for_external(self):
+        if self.kind != AgentKind.EXTERNAL:
+            raise ValueError(f"The specified kind '{self.kind}' cannot be used to create an external agent.")
+        return self
+
+def validate_external_agent_fields(values: dict) -> dict:
+    # Check for empty strings or whitespace
+    for field in ["name", "kind", "description", "title", "tags", "api_url", "chat_params", "nickname", "app_id"]:
+        value = values.get(field)
+        if value and not str(value).strip():
+            raise ValueError(f"{field} cannot be empty or just whitespace")
+
+    return values
+
+# # ===============================
+# #      ASSISTANT AGENT TYPES
+# # ===============================
+
+class AssistantAgentConfig(BaseModel):
+    api_version: Annotated[str | None, Field(json_schema_extra={"min_length_str":1})] = None
+    assistant_id: Annotated[str | None, Field(json_schema_extra={"min_length_str":1})] = None
+    crn: Annotated[str | None, Field(json_schema_extra={"min_length_str":1})] = None
+    service_instance_url: Annotated[str | None, Field(validation_alias=AliasChoices('instance_url', 'service_instance_url'), serialization_alias='service_instance_url')] = None
+    environment_id: Annotated[str | None, Field(json_schema_extra={"min_length_str":1})] = None
+    auth_type: Annotated[str | None, Field(json_schema_extra={"min_length_str":1})] = None
+    connection_id: Annotated[str | None, Field(json_schema_extra={"min_length_str":1})] = None
+    api_key: Annotated[str | None, Field(json_schema_extra={"min_length_str":1})] = None
+    authorization_url: Annotated[str | None, Field(json_schema_extra={"min_length_str":1})] = None
+    auth_type: AssistantAgentAuthType = AssistantAgentAuthType.MCSP
+
+class AssistantAgentSpec(BaseAgentSpec):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    kind: AgentKind = AgentKind.ASSISTANT
+    title: Annotated[str, Field(json_schema_extra={"min_length_str":1})]
+    tags: Optional[List[str]] = None
+    config: AssistantAgentConfig = AssistantAgentConfig()
+    nickname: Annotated[str | None, Field(json_schema_extra={"min_length_str":1})] = None
+    connection_id: Annotated[str | None, Field(json_schema_extra={"min_length_str":1})] = None
+
+    @model_validator(mode="before")
+    def validate_fields_for_external(cls, values):
+        return validate_assistant_agent_fields(values)
+
+    @model_validator(mode="after")
+    def validate_kind_for_external(self):
+        if self.kind != AgentKind.ASSISTANT:
+            raise ValueError(f"The specified kind '{self.kind}' cannot be used to create an assistant agent.")
+        return self
+
+def validate_assistant_agent_fields(values: dict) -> dict:
+    # Check for empty strings or whitespace
+    for field in ["name", "kind", "description", "title", "tags", "nickname", "app_id"]:
+        value = values.get(field)
+        if value and not str(value).strip():
+            raise ValueError(f"{field} cannot be empty or just whitespace")
+
+    return values

ibm_watsonx_orchestrate/agent_builder/connections/__init__.py

@@ -0,0 +1,27 @@
+from .connections import (
+    get_application_connection_credentials,
+    get_connection_type
+)
+from .types import (
+    ConnectionType,
+    ConnectionKind,
+    ConnectionAuthType,
+    ConnectionConfiguration,
+    ConnectionEnvironment,
+    ConnectionPreference,
+    ConnectionSecurityScheme,
+    CREDENTIALS,
+    BasicAuthCredentials,
+    BearerTokenAuthCredentials,
+    APIKeyAuthCredentials,
+    OAuth2TokenCredentials,
+    # OAuth2AuthCodeCredentials,
+    # OAuth2ClientCredentials,
+    # OAuth2ImplicitCredentials,
+    # OAuth2PasswordCredentials,
+    OAuthOnBehalfOfCredentials,
+    KeyValueConnectionCredentials,
+    CONNECTION_KIND_SCHEME_MAPPING,
+    CONNECTION_TYPE_CREDENTIAL_MAPPING,
+    ExpectedCredentials
+)

ibm_watsonx_orchestrate/agent_builder/connections/connections.py

@@ -0,0 +1,123 @@
+import os
+import logging
+from typing import List
+from ibm_watsonx_orchestrate.agent_builder.connections.types import (
+    BasicAuthCredentials,
+    BearerTokenAuthCredentials,
+    APIKeyAuthCredentials,
+    OAuth2TokenCredentials,
+    KeyValueConnectionCredentials,
+    ConnectionType,
+    CREDENTIALS,
+    CONNECTION_TYPE_CREDENTIAL_MAPPING
+)
+
+from ibm_watsonx_orchestrate.utils.utils import sanatize_app_id
+
+logger = logging.getLogger(__name__)
+
+_PREFIX_TEMPLATE = "WXO_CONNECTION_{app_id}_"
+
+connection_type_requirements_mapping = {
+    BasicAuthCredentials: ["username", "password"],
+    BearerTokenAuthCredentials: ["token"],
+    APIKeyAuthCredentials: ["api_key"],
+    OAuth2TokenCredentials: ["access_token"],
+    KeyValueConnectionCredentials: None
+}
+
+def _clean_env_vars(vars: dict[str:str], requirements: List[str], app_id: str) -> dict[str,str]:
+    base_prefix = _PREFIX_TEMPLATE.format(app_id=app_id)
+
+    required_env_vars = {}
+    missing_requirements = []
+    for requirement in requirements:
+        key = base_prefix + requirement
+        value = vars.get(key)
+        if value:
+            required_env_vars[key] = value
+        else:
+            missing_requirements.append(key)
+    
+    # Get value from optional url env var
+    key = base_prefix + "url"
+    value = vars.get(key)
+    required_env_vars[key] = value
+    
+    if len(missing_requirements) > 0:
+        missing_requirements_str = ", ".join(missing_requirements)
+        message = f"Missing requirement environment variables '{missing_requirements_str}' for connection '{app_id}'"
+        logger.error(message)
+        raise ValueError(message)
+    
+    return required_env_vars
+
+def _build_credentials_model(credentials_type: type[CREDENTIALS], vars: dict[str,str], base_prefix: str) -> type[CREDENTIALS]:
+    requirements = connection_type_requirements_mapping[credentials_type]
+
+    if requirements:
+        requirements.append("url")
+        model_dict={}
+
+        for requirement in requirements:
+            model_dict[requirement] = vars[base_prefix+requirement]
+        return credentials_type(
+            **model_dict
+        )
+    else:
+        # Strip the default prefix
+        model_dict = {}
+        for key in vars:
+            new_key = key.removeprefix(base_prefix)
+            model_dict[new_key] = vars[key]
+        return credentials_type(
+            model_dict
+        )
+
+
+def _validate_schema_type(requested_type: ConnectionType, expected_type: ConnectionType) -> bool:
+        return expected_type == requested_type
+
+def _get_credentials_model(connection_type: ConnectionType, app_id: str) -> type[CREDENTIALS]:
+    base_prefix = _PREFIX_TEMPLATE.format(app_id=app_id)
+    variables = {}
+    for key, value in os.environ.items():
+        if key.startswith(base_prefix):
+            variables[key] = value
+    
+    credentials_type = CONNECTION_TYPE_CREDENTIAL_MAPPING[connection_type]
+
+    requirements = connection_type_requirements_mapping.get(credentials_type)
+    if requirements:
+        variables = _clean_env_vars(vars=variables, requirements=requirements, app_id=app_id)
+
+    return _build_credentials_model(credentials_type=credentials_type, vars=variables, base_prefix=base_prefix)
+
+def get_connection_type(app_id: str) -> ConnectionType:
+    sanitized_app_id = sanatize_app_id(app_id=app_id)
+    expected_schema_key = f"WXO_SECURITY_SCHEMA_{sanitized_app_id}"
+    expected_schema = os.environ.get(expected_schema_key)
+
+    if not expected_schema:
+        message = f"No credentials found for connections '{app_id}'"
+        logger.error(message)
+        raise ValueError(message)
+
+    auth_types = {e.value for e in ConnectionType}
+    if expected_schema not in auth_types:
+        message = f"The expected type '{expected_schema}' cannot be resolved into a valid connection auth type ({', '.join(list(auth_types))})"
+        logger.error(message)
+        raise ValueError(message)
+
+    return expected_schema
+
+def get_application_connection_credentials(type: ConnectionType, app_id: str) -> CREDENTIALS:
+    sanitized_app_id = sanatize_app_id(app_id=app_id)
+    expected_schema = get_connection_type(app_id=app_id)
+
+    if not _validate_schema_type(requested_type=type, expected_type=expected_schema):
+        message = f"The requested type '{type.__name__}' does not match the type '{expected_schema}' for the connection '{app_id}'"
+        logger.error(message)
+        raise ValueError(message)
+
+    return _get_credentials_model(connection_type=type, app_id=sanitized_app_id)