ibm-watsonx-orchestrate 1.10.0b0__py3-none-any.whl → 1.11.0b0__py3-none-any.whl

ibm_watsonx_orchestrate/__init__.py

@@ -5,9 +5,7 @@
 
 pkg_name = "ibm-watsonx-orchestrate"
 
-__version__ = "1.10.0b0"
-
-
+__version__ = "1.11.0b0"
 
 
 from ibm_watsonx_orchestrate.utils.logging.logger import setup_logging

ibm_watsonx_orchestrate/agent_builder/connections/__init__.py

@@ -18,7 +18,7 @@ from .types import (
     OAuth2AuthCodeCredentials,
     OAuth2ClientCredentials,
     # OAuth2ImplicitCredentials,
-    # OAuth2PasswordCredentials,
+    OAuth2PasswordCredentials,
     OAuthOnBehalfOfCredentials,
     KeyValueConnectionCredentials,
     CONNECTION_KIND_SCHEME_MAPPING,

ibm_watsonx_orchestrate/agent_builder/connections/connections.py

@@ -37,7 +37,7 @@ connection_type_security_schema_map = {
     ConnectionType.OAUTH2_CLIENT_CREDS: ConnectionSecurityScheme.OAUTH2,
     ConnectionType.OAUTH_ON_BEHALF_OF_FLOW: ConnectionSecurityScheme.OAUTH2,
     # ConnectionType.OAUTH2_IMPLICIT: ConnectionSecurityScheme.OAUTH2,
-    # ConnectionType.OAUTH2_PASSWORD: ConnectionSecurityScheme.OAUTH2
+    ConnectionType.OAUTH2_PASSWORD: ConnectionSecurityScheme.OAUTH2
 }
 
 def _clean_env_vars(vars: dict[str:str], requirements: List[str], app_id: str) -> dict[str,str]:

ibm_watsonx_orchestrate/agent_builder/connections/types.py

@@ -8,7 +8,7 @@ class ConnectionKind(str, Enum):
     api_key = 'api_key'
     oauth_auth_code_flow = 'oauth_auth_code_flow'
     # oauth_auth_implicit_flow = 'oauth_auth_implicit_flow'
-    # oauth_auth_password_flow = 'oauth_auth_password_flow'
+    oauth_auth_password_flow = 'oauth_auth_password_flow'
     oauth_auth_client_credentials_flow = 'oauth_auth_client_credentials_flow'
     oauth_auth_on_behalf_of_flow = 'oauth_auth_on_behalf_of_flow'
     key_value = 'key_value'
@@ -34,7 +34,7 @@ class ConnectionPreference(str, Enum):
 class ConnectionAuthType(str, Enum):
     OAUTH2_AUTH_CODE = 'oauth2_auth_code'
     # OAUTH2_IMPLICIT = 'oauth2_implicit'
-    # OAUTH2_PASSWORD = 'oauth2_password'
+    OAUTH2_PASSWORD = 'oauth2_password'
     OAUTH2_CLIENT_CREDS = 'oauth2_client_creds'
     OAUTH_ON_BEHALF_OF_FLOW = 'oauth_on_behalf_of_flow'
 
@@ -65,7 +65,7 @@ class ConnectionType(str, Enum):
     API_KEY_AUTH = ConnectionSecurityScheme.API_KEY_AUTH.value
     OAUTH2_AUTH_CODE = ConnectionAuthType.OAUTH2_AUTH_CODE.value
     # OAUTH2_IMPLICIT = ConnectionAuthType.OAUTH2_IMPLICIT.value
-    # OAUTH2_PASSWORD = ConnectionAuthType.OAUTH2_PASSWORD.value
+    OAUTH2_PASSWORD = ConnectionAuthType.OAUTH2_PASSWORD.value
     OAUTH2_CLIENT_CREDS = ConnectionAuthType.OAUTH2_CLIENT_CREDS.value
     OAUTH_ON_BEHALF_OF_FLOW = ConnectionAuthType.OAUTH_ON_BEHALF_OF_FLOW.value
     KEY_VALUE = ConnectionSecurityScheme.KEY_VALUE.value
@@ -90,7 +90,7 @@ OAUTH_CONNECTION_TYPES = {
     ConnectionType.OAUTH2_AUTH_CODE,
     ConnectionType.OAUTH2_CLIENT_CREDS,
     # ConnectionType.OAUTH2_IMPLICIT,
-    # ConnectionType.OAUTH2_PASSWORD,
+    ConnectionType.OAUTH2_PASSWORD,
     ConnectionType.OAUTH_ON_BEHALF_OF_FLOW,
 }
 
@@ -165,6 +165,53 @@ class ConnectionConfiguration(BaseModel):
             raise ValueError("Connection of type 'key_value' cannot be configured at the 'member' level. Key value connections must be of type 'team'")
         return self
 
+class ConnectionCredentialsEntryLocation(str, Enum):
+    BODY = 'body'
+    HEADER = 'header',
+    QUERY = 'query'
+
+    def __str__(self):
+        return self.value
+
+class ConnectionCredentialsEntry(BaseModel):
+    key: str
+    value: str
+    location: ConnectionCredentialsEntryLocation
+
+    def __str__(self):
+        return f"<ConnectionCredentialsEntry: {self.location}:{self.key}={self.value}>"
+
+class BaseOAuthCredentials(BaseModel):
+    custom_token_query: Optional[dict] = None
+    custom_token_header: Optional[dict] = None
+    custom_token_body: Optional[dict] = None
+    custom_auth_query: Optional[dict] = None
+
+class ConnectionCredentialsCustomFields(BaseOAuthCredentials):
+    def add_field(self, entry: ConnectionCredentialsEntry, is_token:bool=True) -> None:
+        match entry.location:
+            case ConnectionCredentialsEntryLocation.HEADER:
+                if not is_token:
+                    return
+                attribute = "custom_token_header"
+            case ConnectionCredentialsEntryLocation.BODY:
+                if not is_token:
+                    return
+                attribute = "custom_token_body"
+            case ConnectionCredentialsEntryLocation.QUERY:
+                if is_token:
+                     attribute = "custom_token_query"
+                else:
+                    attribute = "custom_auth_query"
+            case _:
+                return
+        
+        fields = getattr(self, attribute)
+        if not fields:
+            setattr(self, attribute, {})
+            fields = getattr(self, attribute)
+        fields[entry.key] = entry.value
+
 class BasicAuthCredentials(BaseModel):
     username: str
     password: str
@@ -182,7 +229,7 @@ class OAuth2TokenCredentials(BaseModel):
     access_token: str
     url: Optional[str] = None
 
-class OAuth2AuthCodeCredentials(BaseModel):
+class OAuth2AuthCodeCredentials(BaseOAuthCredentials):
     client_id: str
     client_secret: str
     token_url: str
@@ -193,13 +240,17 @@ class OAuth2AuthCodeCredentials(BaseModel):
 #     client_id: str
 #     authorization_url: str
 
-# class OAuth2PasswordCredentials(BaseModel):
-#     client_id: str
-#     client_secret: str
-#     token_url: str
-#     authorization_url: str
+class OAuth2PasswordCredentials(BaseOAuthCredentials):
+    username: str
+    password: str
+    client_id: str
+    client_secret: str
+    token_url: str
+    scope: Optional[str] = None
+    grant_type: str = "password"
+    
 
-class OAuth2ClientCredentials(BaseModel):
+class OAuth2ClientCredentials(BaseOAuthCredentials):
     client_id: str
     client_secret: str
     token_url: str
@@ -207,7 +258,7 @@ class OAuth2ClientCredentials(BaseModel):
     send_via: ConnectionSendVia = ConnectionSendVia.HEADER
     grant_type: str = "client_credentials"
 
-class OAuthOnBehalfOfCredentials(BaseModel):
+class OAuthOnBehalfOfCredentials(BaseOAuthCredentials):
     client_id: str
     access_token_url: str
     grant_type: str
@@ -226,7 +277,7 @@ CREDENTIALS_SET = Union[
     APIKeyAuthCredentials,
     OAuth2AuthCodeCredentials,
     # OAuth2ImplicitCredentials,
-    # OAuth2PasswordCredentials,
+    OAuth2PasswordCredentials,
     OAuth2ClientCredentials,
     OAuthOnBehalfOfCredentials,
     KeyValueConnectionCredentials
@@ -240,7 +291,7 @@ CONNECTION_KIND_SCHEME_MAPPING = {
     ConnectionKind.api_key: ConnectionSecurityScheme.API_KEY_AUTH,
     ConnectionKind.oauth_auth_code_flow: ConnectionSecurityScheme.OAUTH2,
     # ConnectionKind.oauth_auth_implicit_flow: ConnectionSecurityScheme.OAUTH2,
-    # ConnectionKind.oauth_auth_password_flow: ConnectionSecurityScheme.OAUTH2,
+    ConnectionKind.oauth_auth_password_flow: ConnectionSecurityScheme.OAUTH2,
     ConnectionKind.oauth_auth_client_credentials_flow: ConnectionSecurityScheme.OAUTH2,
     ConnectionKind.oauth_auth_on_behalf_of_flow: ConnectionSecurityScheme.OAUTH2,
     ConnectionKind.key_value: ConnectionSecurityScheme.KEY_VALUE,
@@ -250,7 +301,7 @@ CONNECTION_KIND_SCHEME_MAPPING = {
 CONNECTION_KIND_OAUTH_TYPE_MAPPING = {
     ConnectionKind.oauth_auth_code_flow: ConnectionAuthType.OAUTH2_AUTH_CODE,
     # ConnectionKind.oauth_auth_implicit_flow: ConnectionAuthType.OAUTH2_IMPLICIT,
-    # ConnectionKind.oauth_auth_password_flow: ConnectionAuthType.OAUTH2_PASSWORD,
+    ConnectionKind.oauth_auth_password_flow: ConnectionAuthType.OAUTH2_PASSWORD,
     ConnectionKind.oauth_auth_client_credentials_flow: ConnectionAuthType.OAUTH2_CLIENT_CREDS,
     ConnectionKind.oauth_auth_on_behalf_of_flow: ConnectionAuthType.OAUTH_ON_BEHALF_OF_FLOW,
 }
@@ -263,7 +314,7 @@ CONNECTION_TYPE_CREDENTIAL_MAPPING = {
     ConnectionSecurityScheme.KEY_VALUE: KeyValueConnectionCredentials,
 }
 
-class IdentityProviderCredentials(BaseModel):
+class IdentityProviderCredentials(BaseOAuthCredentials):
     idp_url: str = Field(validation_alias=AliasChoices('idp_url', 'url'), serialization_alias='idp_url')
     client_id: str
     client_secret: str
@@ -272,4 +323,4 @@ class IdentityProviderCredentials(BaseModel):
 
 class ExpectedCredentials(BaseModel):
     app_id: str
-    type: ConnectionType | List[ConnectionType]
+    type: ConnectionType | List[ConnectionType]

ibm_watsonx_orchestrate/agent_builder/knowledge_bases/types.py

@@ -11,6 +11,7 @@ class SpecVersion(str, Enum):
 class KnowledgeBaseKind(str, Enum):
     KNOWLEDGE_BASE = "knowledge_base"
 class RetrievalConfidenceThreshold(str, Enum):
+    Off = "Off"
     Lowest = "Lowest"
     Low = "Low"
     High = "High"
@@ -24,6 +25,7 @@ class GeneratedResponseLength(str, Enum):
 
 
 class ResponseConfidenceThreshold(str, Enum):
+    Off = "Off"
     Lowest = "Lowest"
     Low = "Low"
     High = "High"
@@ -86,6 +88,8 @@ class GenerationConfiguration(BaseModel):
     {
         "model_id": "meta-llama/llama-3-1-70b-instruct",
         "prompt_instruction": "When the documents are in different languages, you should respond in english.",
+        "max_docs_passed_to_llm": 10,
+        "retrieval_confidence_threshold": "Lowest",
         "generated_response_length": "Moderate",
         "display_text_no_results_found": "no docs found",
         "display_text_connectivity_issue": "conn failed",
@@ -95,6 +99,7 @@ class GenerationConfiguration(BaseModel):
 
     model_id: Optional[str] = None
     prompt_instruction: Optional[str] = None
+    max_docs_passed_to_llm: Optional[int] = None
     generated_response_length: Optional[GeneratedResponseLength] = None
     display_text_no_results_found: Optional[str] = None
     display_text_connectivity_issue: Optional[str] = None
@@ -134,12 +139,13 @@ class MilvusConnection(BaseModel):
     }
     """
     grpc_host: Optional[str] = None
+    grpc_port: Optional[str] = None
+    server_cert: Optional[str] = None
     database: Optional[str] = None
     collection: Optional[str] = None
     index: Optional[str] = None
     embedding_model_id: Optional[str] = None
     limit : Optional[int] = None
-    grpc_port: Optional[str] = None
     filter: Optional[str] = None
     field_mapping: Optional[FieldMapping] = None
 
@@ -197,12 +203,46 @@ class CustomSearchConnection(BaseModel):
     filter: Optional[str] = None
     metadata: Optional[dict] = None
 
+class AstraDBConnection(BaseModel):
+    """
+    example:
+    {
+        "api_endpoint": "https://xxx-us-east-2.apps.astra.datastax.com",
+        "key_space": "default_keyspace",
+        "collection": "search_wa_docs",
+        "embedding_model_id": "sentence-transformers/all-minilm-l12-v2",
+        "port": "443",
+        "filter": "productType: \"boots\"",
+        "limit": 5,
+        "field_mapping": {
+                        "title": "title",
+                        "body": "text",
+                        "url": "some-url"
+                    }
+    }
+    """
+    api_endpoint: str
+    port: Optional[str] = None
+    server_cert: Optional[str] = None
+    keyspace: Optional[str]
+    data_type: str
+    collection: Optional[str]
+    table: Optional[str]
+    index_column: Optional[str]
+    embedding_mode: str
+    embedding_model_id: Optional[str]
+    credentials: dict
+    search_mode: str
+    limit: Optional[int] = 5
+    filter: Optional[str] = None
+    field_mapping: Optional[FieldMapping] = None
+
 class IndexConnection(BaseModel):
     connection_id: Optional[str] = None
     milvus: Optional[MilvusConnection] = None
     elastic_search: Optional[ElasticSearchConnection] = None
     custom_search: Optional[CustomSearchConnection] = None
-
+    astradb: Optional[AstraDBConnection] = None
     
 class ConversationalSearchConfig(BaseModel):
     language: Optional[str] = None
@@ -218,6 +258,10 @@ class KnowledgeBaseBuiltInVectorIndexConfig(BaseModel):
     chunk_size: Optional[int] = None
     chunk_overlap: Optional[int] = None
     limit: Optional[int] = None
+
+class FileUpload(BaseModel):
+    path: str
+    url: Optional[str] = None
     
 class KnowledgeBaseSpec(BaseModel):
     """Schema for a complete knowledge-base."""
@@ -236,4 +280,4 @@ class KnowledgeBaseSpec(BaseModel):
     created_on: Optional[datetime] = None 
     updated_at: Optional[datetime] = None
     # For import/update
-    documents: list[str] = None
+    documents: list[str] | list[FileUpload] = None

ibm_watsonx_orchestrate/agent_builder/toolkits/types.py

@@ -1,27 +1,36 @@
-from typing import List, Dict, Optional
+from typing import List, Dict, Optional, Union
 from enum import Enum
-from pydantic import BaseModel, model_validator
+from pydantic import BaseModel 
 
 class ToolkitKind(str, Enum):
     MCP = "mcp"
 
-class Language(str, Enum):
-    NODE = "node"
-    PYTHON ="python"
-
 class ToolkitSource(str, Enum):
     FILES = "files"
     PUBLIC_REGISTRY = "public-registry"
 
+class ToolkitTransportKind(str, Enum):
+    STREAMABLE_HTTP = "streamable_http"
+    SSE = "sse"
 
+class Language(str, Enum):
+    NODE = "node"
+    PYTHON ="python"
 
-class McpModel(BaseModel):
-    source: str
+class LocalMcpModel(BaseModel):
+    source: ToolkitSource
     command: str
     args: List[str]
     tools: List[str]
     connections: Dict[str, str]
 
+class RemoteMcpModel(BaseModel):
+    server_url: str
+    transport: ToolkitTransportKind
+    tools: List[str]
+    connections: Dict[str, str]
+
+McpModel = Union[LocalMcpModel, RemoteMcpModel]
 
 class ToolkitSpec(BaseModel):
     id: str
@@ -33,10 +42,4 @@ class ToolkitSpec(BaseModel):
     created_by: str
     created_by_username: str
     tools: List[str] | None
-    mcp: McpModel
-
-    @model_validator(mode='after')
-    def validate_tools_and_mcp(self) -> 'ToolkitSpec':
-        if self.mcp.source not in {"files", "public-registry"}:
-            raise ValueError("MCP source must be either 'files' or 'public-registry'.")
-        return self
+    mcp: McpModel

ibm_watsonx_orchestrate/agent_builder/tools/types.py

@@ -152,7 +152,7 @@ class ClientSideToolBinding(BaseModel):
 
 class McpToolBinding(BaseModel):
     server_url: Optional[str] = None
-    source: str
+    source: str | None
     connections: Dict[str, str] | None
 
 class FlowToolBinding(BaseModel):

ibm_watsonx_orchestrate/cli/commands/connections/connections_command.py

@@ -1,6 +1,6 @@
 import typer
 from typing_extensions import Annotated, List
-from ibm_watsonx_orchestrate.agent_builder.connections.types import ConnectionEnvironment, ConnectionPreference, ConnectionKind
+from ibm_watsonx_orchestrate.agent_builder.connections.types import ConnectionEnvironment, ConnectionPreference, ConnectionKind, ConnectionCredentialsEntry
 from ibm_watsonx_orchestrate.cli.commands.connections.connections_controller import (
     add_connection,
     remove_connection,
@@ -8,7 +8,9 @@ from ibm_watsonx_orchestrate.cli.commands.connections.connections_controller imp
     import_connection,
     configure_connection,
     set_credentials_connection,
-    set_identity_provider_connection
+    set_identity_provider_connection,
+    token_entry_connection_credentials_parse,
+    auth_entry_connection_credentials_parse
 )
 
 connections_app = typer.Typer(no_args_is_help=True)
@@ -160,7 +162,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 +170,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 +193,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 +215,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 +229,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[
@@ -244,6 +246,22 @@ def set_credentials_connection_command(
             help="For key_value, a key value pair in the form '<key>=<value>'. Multiple values can be passed using `-e key1=value1 -e key2=value2`"
         )
     ] = None,
+    token_entries: Annotated[
+        List[ConnectionCredentialsEntry],
+        typer.Option(
+            '--token-entries', "-t",
+            parser=token_entry_connection_credentials_parse,
+            help="Custom field options for oauth types token request, a key value location option in the form 'location:<key>=<value>' or '<key>=<value>' with location defaulting to 'header'. Multiple values can be passed using `-t key1=value1 -t location:key2=value2`"
+        )
+    ] = None,
+    auth_entries: Annotated[
+        List[ConnectionCredentialsEntry],
+        typer.Option(
+            '--auth-entries',
+            parser=auth_entry_connection_credentials_parse,
+            help="Custom field options for oauth_auth_code_flow auth server request, a key value location option in the form 'location:<key>=<value>' or '<key>=<value>' with location defaulting to 'query'. Note only 'query' is a valid location. Multiple values can be passed using `--auth-entries key1=value1 --auth-entries location:key2=value2`"
+        )
+    ] = None,
 ):
     set_credentials_connection(
         app_id=app_id,
@@ -259,7 +277,9 @@ def set_credentials_connection_command(
         auth_url=auth_url,
         grant_type=grant_type,
         scope=scope,
-        entries=entries
+        entries=entries,
+        token_entries=token_entries,
+        auth_entries=auth_entries
     )
 
 @connections_app.command(name="set-identity-provider")
@@ -311,6 +331,14 @@ def set_identity_provider_connection_command(
             help='The grant-type of the the identity provider'
         )
     ],
+    token_entries: Annotated[
+        List[ConnectionCredentialsEntry],
+        typer.Option(
+            '--token-entries', "-t",
+            parser=token_entry_connection_credentials_parse,
+            help="Custom field options for oauth types token request, a key value location option in the form 'location:<key>=<value>' or '<key>=<value>' with location defaulting to 'header'. Multiple values can be passed using `-t key1=value1 -t location:key2=value2`"
+        )
+    ] = None,
 ):
     set_identity_provider_connection(
         app_id=app_id,
@@ -319,5 +347,6 @@ def set_identity_provider_connection_command(
         client_id=client_id,
         client_secret=client_secret,
         scope=scope,
-        grant_type=grant_type
+        grant_type=grant_type,
+        token_entries=token_entries
     )