ibm-watsonx-gov 1.3.3__cp313-cp313-macosx_11_0_arm64.whl

ibm_watsonx_gov/tools/core/tool_loader.py

@@ -0,0 +1,237 @@
+# ----------------------------------------------------------------------------------------------------
+# IBM Confidential
+# Licensed Materials - Property of IBM
+# 5737-H76, 5900-A3Q
+# © Copyright IBM Corp. 2025  All Rights Reserved.
+# US Government Users Restricted Rights - Use, duplication or disclosure restricted by
+# GSA ADPSchedule Contract with IBM Corp.
+# ----------------------------------------------------------------------------------------------------
+
+import ast
+from typing import Callable, Union
+
+import pandas as pd
+from langchain_core.tools import StructuredTool
+from pydantic import BaseModel
+
+from ibm_watsonx_gov.utils.python_utils import get
+
+from ..clients import get_tool_info
+from ..utils import (TOOL_REGISTRY, get_base64_decoding, get_pydantic_model,
+                     validate_envs)
+from ..utils.constants import CustomToolType, ServiceProviderType
+from ..utils.package_utils import install_and_import_packages
+
+
+def load_tool(tool_name: str, **kwargs):
+    """
+        Method to load a tool based on tool name and other parameters
+
+        Args:
+            tool_name(str) : 
+                    Name of the registered tool
+            headers(Optional[headers]): 
+                    - Dynamic headers to be supplied to the tool
+                    - If the user supplies static headers (whose value do not change) no need of supplying this header while loading.
+                    - User is expected to provide the information related to dynamic headers while tool loading . Types of dynamic headers
+                        - If the tool is registered with header value $DYNAMIC HEADERS it is mandatory to supply the information while tool loading 
+                        - If the tool is registered with convention `headers_key: $header_key` this information should be supplied as part of environment
+            **kwargs: 
+                    Optional configuration based on the tool .(Eg: Configuration  properties for OOTB tool)
+
+        Returns:
+            Tool : A Langchain/Langgraph compatible tool object
+
+        Examples:
+        -------
+        1. Basic example to load a tool
+            .. code-block:: python
+
+                from ibm_watsonx_gov.tools import load_tool
+                tool = load_tool(tool_name=<TOOL_NAME>)
+
+        2. Example of tool loading that expects configuration
+            .. code-block:: python
+
+                from ibm_watsonx_gov.tools import load_tool
+                from ibm_watsonx_gov.tools.clients import get_tool_info
+
+                tool_details = get_tool_info(tool_name=<TOOL_NAME>)
+                tool_config = tool_details.get("entity").get("config)
+                config = {
+                    <Provide the properties as per the details in tool config>
+                }
+                tool = load_tool(tool_name=<TOOL_NAME>,**config)
+
+        3. Example of tool loading that expects headers 
+            .. code-block:: python
+
+                from ibm_watsonx_gov.tools import load_tool
+                from ibm_watsonx_gov.tools.clients import get_tool_info
+
+                tool_details = get_tool_info(tool_name=<TOOL_NAME>)
+                headers = tool_details.get("entity").get("endpoint").get("headers")
+
+                headers = {
+                    "Authorization":"Bearer __TOKEN__"  #Need to supply dynamic headers
+                }
+                tool = load_tool(tool_name=<TOOL_NAME>,headers=headers)
+    """
+    # Get the tool information
+    tool_details = get_tool_info(tool_name=tool_name)
+
+    # Get service provider type
+    service_provider_type = get(tool_details, "entity.service_provider_type")
+    tool_name = get(tool_details, "entity.tool_name")
+    env_list = get(tool_details, "entity.environment_variables", [])
+
+    # Validate the env for loading the tool
+    validate_envs(env_list)
+
+    # Load the tool based on the service provider type
+    try:
+        if service_provider_type == ServiceProviderType.CUSTOM.value:
+            tool_description = get(tool_details, "entity.description")
+            tool_schema = get(tool_details, "entity.schema")
+            tool_type = get(tool_details, "entity.tool_type")
+            dependencies = get(
+                tool_details, "entity.dependencies.run_time_packages")
+            if len(dependencies) > 0:
+                dependencies = [package for package in dependencies if not package.startswith(
+                    "ibm-watsonx-gov")]
+                install_and_import_packages(dependencies)
+
+            if tool_type == CustomToolType.CODE.value:
+                # This is a code based tool
+                encode_tool_code = get(
+                    tool_details, "entity.code.source_code_base64")
+                tool_code = get_base64_decoding(encode_tool_code)
+                return _load_custom_tool(tool_name,
+                                         tool_description,
+                                         tool_code,
+                                         tool_schema,
+                                         **kwargs)
+            elif tool_type == CustomToolType.ENDPOINT.value:
+                tool_endpoint = get(tool_details, "entity.endpoint.url")
+                headers = get(tool_details, "entity.endpoint.headers", {})
+                arg_headers = kwargs.get("headers", {})
+                tool_headers = headers | arg_headers
+                method = get(tool_details, "entity.endpoint.method", "POST")
+                return _load_restapi_tool(tool_name,
+                                          tool_description,
+                                          tool_endpoint,
+                                          tool_schema,
+                                          tool_headers,
+                                          method=method)
+        elif service_provider_type == ServiceProviderType.IBM.value:
+            return _load_ootb_tool(tool_name, **kwargs)
+        else:
+            raise Exception(
+                f"service_provide_type:{service_provider_type} is not supporteds. Acceptable values:{ServiceProviderType.values()}")
+    except Exception as ex:
+        raise Exception(f"Error loading tool {tool_name}. Details:{str(ex)}")
+
+
+def _load_ootb_tool(tool_name: str = None, **kwargs) -> Callable:
+    """
+    Load the OOTB tool using the tool name
+
+    Args:
+        tool_name (str): Name of the OOTB tool
+
+    Returns:
+        Callable: OOTB Tool loaded from the tool registry
+    """
+    if tool_name is None:
+        raise Exception("'tool_name' parameter is missing.")
+
+    tool_class_path = TOOL_REGISTRY.get(tool_name, None)
+    if tool_class_path is None:
+        raise Exception(f"OOTB tool  {tool_name} is not available")
+
+    module_path, _, class_name = tool_class_path.rpartition('.')
+    tool_module = __import__(module_path, fromlist=[f"{class_name}"])
+
+    return getattr(tool_module, class_name)(**kwargs)
+
+
+def _load_custom_tool(tool_name: str,
+                      tool_description: str,
+                      tool_code: str,
+                      tool_schema: dict,
+                      **kwargs) -> Callable:
+    """Method to create custom tool 
+
+    Args:
+        tool_name (str): Name of the agent tool
+        tool_description (str): Description of agent tool
+        tool_code (str): Tool code in the form of doc string 
+        tool_schema (dict): JSON args schema to be supplied for converting tool
+
+    Returns:
+        Callable: Custom tool created using the tool code
+    """
+
+    def call_tool(**kwargs):
+        tree = ast.parse(tool_code, mode="exec")
+        function_name = tree.body[0].name
+        compiled_code = compile(tree, 'custom_tool', 'exec')
+        namespace = {}
+        exec(compiled_code, namespace)
+        return namespace[function_name](**kwargs)
+
+    title = tool_name + \
+        "_schema" if tool_schema.get(
+            'title') is None else tool_schema.get('title')
+
+    tool = StructuredTool(
+        name=tool_name,
+        description=tool_description,
+        func=call_tool,
+        args_schema=get_pydantic_model(title, tool_schema)
+    )
+    return tool
+
+
+def _load_restapi_tool(tool_name: str,
+                       tool_description: str,
+                       tool_endpoint: str,
+                       tool_schema: Union[dict, BaseModel],
+                       headers: dict = {},
+                       method='POST') -> Callable:
+    """Method to create custom tool
+
+    Args:
+        tool_name (str): Name of the agent tool 
+        tool_description (str): Description of the agent tool 
+        tool_endpoint (str): RestEndpoint to be used as a tool 
+        tool_schema (Union[dict, BaseModel]): Tool schema 
+        headers (dict,optional): Headers to be used while calling the tool .Defaults to {}
+        method (str, optional): _description_. Defaults to 'POST'.
+
+    Returns:
+        Callable: RESTAPI tool created using the template tool
+    """
+    from ..rest_api import RestApiTool, load_headers
+    if isinstance(tool_schema, dict):
+
+        title = tool_name + \
+            "_schema" if tool_schema.get(
+                'title') is None else tool_schema.get('title')
+
+        args_schema = get_pydantic_model(title, tool_schema)
+    else:
+        args_schema = tool_schema
+
+    # Load headers with envs as needed
+    headers = load_headers(tool_name, headers)
+
+    tool = RestApiTool(
+        name=tool_name,
+        description=tool_description,
+        endpoint=tool_endpoint,
+        method=method,
+        headers=headers,
+        args_schema=args_schema
+    )
+    return tool

ibm_watsonx_gov/tools/entities/__init__.py

@@ -0,0 +1,8 @@
+# ----------------------------------------------------------------------------------------------------
+# IBM Confidential
+# Licensed Materials - Property of IBM
+# 5737-H76, 5900-A3Q
+# © Copyright IBM Corp. 2025  All Rights Reserved.
+# US Government Users Restricted Rights - Use, duplication or disclosure restricted by
+# GSA ADPSchedule Contract with IBM Corp.
+# ----------------------------------------------------------------------------------------------------

ibm_watsonx_gov/tools/entities/ai_tools.py

@@ -0,0 +1,435 @@
+# ----------------------------------------------------------------------------------------------------
+# IBM Confidential
+# Licensed Materials - Property of IBM
+# 5737-H76, 5900-A3Q
+# © Copyright IBM Corp. 2025  All Rights Reserved.
+# US Government Users Restricted Rights - Use, duplication or disclosure restricted by
+# GSA ADPSchedule Contract with IBM Corp.
+# ----------------------------------------------------------------------------------------------------
+
+from typing import Annotated, Any, Dict, List, Optional
+
+from pydantic import BaseModel, Field, model_validator
+from pydantic_core import PydanticCustomError
+
+from ..utils.constants import (Categories, CustomToolType, Framework,
+                               PatchOperationTypes, ServiceProviderType)
+from ..utils.python_utils import get_base64_encoding
+
+
+class AIToolRuntimeDetails(BaseModel):
+    engine: Annotated[
+        str,
+        Field(
+            ...,
+            description="Python version required by the tool."
+        )
+    ]
+    cpu_capacity: Annotated[
+        Optional[str],
+        Field(
+            default=None,
+            description="CPU capacity required by the tool."
+        )
+    ]
+    memory: Annotated[
+        Optional[str],
+        Field(
+            default=None,
+            description="Memory required by the tool."
+        )
+    ]
+
+
+class AIToolCodePayload(BaseModel):
+    source: Annotated[
+        Optional[str],
+        Field(
+            default=None,
+            description="Source location of the tool's code."
+        )
+    ]
+    commit: Annotated[
+        Optional[str],
+        Field(
+            default=None
+        )
+    ]
+    language: Annotated[
+        Optional[str],
+        Field(
+            default="python",
+            description="Language used in the code"
+        )
+    ]
+    source_code_url: Annotated[
+        Optional[str],
+        Field(
+            default=None,
+            description="Remote URL for accessing the source code. Currently, only GitHub URLs are supported."
+        )
+    ]
+    source_code_base64: Annotated[
+        Optional[str],
+        Field(
+            ...,
+            description="Source code for invoking the tool.",
+        )
+    ]
+    run_time_details: Annotated[
+        Optional[AIToolRuntimeDetails],
+        Field(
+            ...,
+            description="Runtime details required for executing the source code."
+        )
+    ]
+
+    @model_validator(mode='before')
+    @classmethod
+    def accept_multiple_field_names(cls, data):
+        if "source_code" in data and "source_code_base64" not in data:
+            data["source_code_base64"] = data["source_code"]
+        return data
+
+
+class AIToolEndpointPayload(BaseModel):
+    url: Annotated[
+        str,
+        Field(
+            ...,
+            description="URL to access the endpoint."
+        )
+    ]
+    method: Annotated[
+        str,
+        Field(
+            ...,
+            description="Method to access the endpoint."
+        )
+    ]
+    headers: Annotated[
+        Dict[str, Any],
+        Field(
+            ...,
+            description="Headers required to call the endpoint."
+        )
+    ]
+
+
+class AIToolUsedInApplication(BaseModel):
+    inventory_id: Annotated[
+        str,
+        Field(
+            ...,
+            description="Inventory ID of the application where the tool is used."
+        )
+    ]
+    agent_asset_id: Annotated[
+        str,
+        Field(
+            ...,
+            description="Agent ID of the application where the tool is used."
+        )
+    ]
+
+
+class AIToolDependencies(BaseModel):
+    remote_services: Annotated[
+        list[str],
+        Field(
+            default=None,
+        )
+    ]
+    run_time_packages: Annotated[
+        list[str],
+        Field(
+            default=None,
+        )
+    ]
+
+
+class AIToolBenchmark(BaseModel):
+    avg_latency: Annotated[
+        str,
+        Field(
+            default=None,
+        )]
+    dataset: Annotated[
+        str,
+        Field(
+            ...,
+            description="GitHub URL for the dataset."
+        )]
+    records: Annotated[
+        int,
+        Field(
+            ...,
+            description="Number of records contained in the dataset."
+        )]
+
+
+class AIToolSchema(BaseModel):
+    title: Annotated[
+        Optional[str],
+        Field(
+            default=None,
+            description="Title for schema."
+        )]
+    type: Annotated[
+        Optional[str],
+        Field(
+            default=None,
+            description="Type of schema."
+        )]
+    properties: Annotated[
+        dict,
+        Field(
+            ...,
+            description="Properties for schema."
+        )]
+    required: Annotated[
+        Optional[list],
+        Field(
+            default=None,
+            description="Required for schema."
+        )]
+
+
+class ToolRegistrationPayload(BaseModel):
+    tool_name: Annotated[
+        str,
+        Field(
+            ...,
+            description="Unique name for the tool"
+        )
+    ]
+    display_name: Annotated[
+        Optional[str],
+        Field(
+            default=None,
+            description="Unique name for displaying the tool."
+        )
+    ]
+    service_provider_type: Annotated[
+        Optional[ServiceProviderType],
+        Field(
+            default=ServiceProviderType.CUSTOM,
+            description=f"Service provider type. Allowed values are: {ServiceProviderType.values()}."
+        )
+    ]
+    tool_type: Annotated[
+        Optional[CustomToolType],
+        Field(
+            default=CustomToolType.CODE,
+            description=f"Type of the tool. Allowed tool types are: {CustomToolType.values()}."
+        )
+    ]
+    description: Annotated[
+        str,
+        Field(
+            ...,
+            description="Short description about the tool."
+        )
+    ]
+    summary: Annotated[
+        Optional[str],
+        Field(
+            default=None,
+            description="Summary of the tool."
+        )
+    ]
+    development_uri_link: Annotated[
+        Optional[str],
+        Field(
+            default=None
+        )
+    ]
+    inventory_id: Annotated[
+        Optional[str],
+        Field(
+            default=None,
+            description="Inventory ID in which the asset needs to be created."
+        )
+    ]
+    reusable: Annotated[
+        bool,
+        Field(
+            default=True,
+            description="Specify the tool will be reusable."
+        )
+    ]
+    framework: Annotated[
+        list[Framework],
+        Field(
+            default_factory=lambda: [Framework.LANGCHAIN, Framework.LANGGRAPH],
+            description="Specify the framework used by the tool."
+        )
+    ]
+    category: Annotated[
+        List[Categories],
+        Field(
+            default_factory=lambda: [Categories.OTHER],
+            description=f"Specify the category under which the tool will be classified. Allowed values are {Categories.values()}."
+        )
+    ]
+    used_in_applications: Annotated[
+        Optional[List[AIToolUsedInApplication]],
+        Field(
+            default=None,
+            description="List of applications where this tool is used."
+        )
+    ]
+    code: Annotated[
+        Optional[AIToolCodePayload],
+        Field(
+            default=None,
+            description="Code-related information for the tool."
+        )
+    ]
+    endpoint: Annotated[
+        Optional[AIToolEndpointPayload],
+        Field(
+            default=None,
+            description="Endpoint-related information for the tool."
+        )
+    ]
+    dependencies: Annotated[
+        Optional[AIToolDependencies],
+        Field(
+            default=None,
+            description="Dependencies required by this tool."
+        )
+    ]
+    config: Annotated[
+        Optional[Dict],
+        Field(
+            default=None,
+            description="Configuration information for this tool."
+        )
+    ]
+    schema_: Annotated[
+        AIToolSchema,
+        Field(
+            ...,
+            description="Schema information for this tool.",
+            alias="schema"
+        )
+    ]
+    environment_variables: Annotated[
+        Optional[list[str]],
+        Field(
+            default=None,
+            description="Environment variables required by this tool."
+        )
+    ]
+    metrics: Annotated[
+        Optional[Dict],
+        Field(
+            default=None,
+            description="Metrics and their respective values."
+        )
+    ]
+    pricing: Annotated[
+        Optional[Dict],
+        Field(
+            default=None,
+            description="Type of currency and its corresponding price."
+        )
+    ]
+    benchmark_test: Annotated[
+        Optional[AIToolBenchmark],
+        Field(
+            default=None,
+            description="Benchmark test details."
+        )
+    ]
+
+    @model_validator(mode='before')
+    @classmethod
+    def check_tool_type(cls, data):
+        if "tool_type" not in data:
+            data["tool_type"] = CustomToolType.CODE.value
+        if data["tool_type"] == CustomToolType.CODE.value and "code" in data and "endpoint" in data:
+            data["endpoint"] = None
+        if data["tool_type"] == CustomToolType.ENDPOINT.value and "code" in data and "endpoint" in data:
+            data["code"] = None
+        return data
+
+    @model_validator(mode="after")
+    def validate_post_payload(self):
+        if self.tool_type == CustomToolType.CODE and not self.code:
+            raise PydanticCustomError(
+                "missing code_payload", 'tool_type is "code", but code field missing in payload.'
+            )
+
+        if self.tool_type == CustomToolType.CODE and not self.code.source_code_url and not self.code.source_code_base64:
+            raise PydanticCustomError(
+                "missing_source_code ", 'Missing both source_code_url and source_code_base64.'
+            )
+
+        if self.code and self.code.source_code_base64:
+            self.code.source_code_base64 = get_base64_encoding(
+                tool_code=self.code.source_code_base64)
+
+        if self.tool_type == CustomToolType.ENDPOINT and not self.endpoint:
+            raise PydanticCustomError(
+                "missing_endpoint_payload", 'tool_type is "endpoint", but endpoint field missing in payload.'
+            )
+
+        if self.service_provider_type.value == ServiceProviderType.IBM.value:
+            raise Exception(f"The service provider type 'IBM' is not supported. Supported service provider type are: 'custom'")
+
+        if not self.display_name:
+            self.display_name = self.tool_name
+
+        return self
+
+
+class PatchPayload(BaseModel):
+    op: Annotated[
+        PatchOperationTypes,
+        Field(
+            ...,
+            description="Type of operation to be performed."
+        )
+    ]
+    path: Annotated[
+        str,
+        Field(
+            ...,
+            description="Path of the field on which the operation needs to be performed."
+        )
+    ]
+    value: Annotated[
+        Any,
+        Field(
+            ...,
+            description="The value to be applied during the specified operation on the given field path."
+        )
+    ]
+
+    @model_validator(mode="after")
+    def validate_patch_payload(self):
+        if "source_code_base64" in self.path:
+            self.value = get_base64_encoding(tool_code=self.value)
+        return self
+
+
+class ToolUpdatePayload(BaseModel):
+    payload: Annotated[
+        list[PatchPayload],
+        Field(
+            ...,
+            description="List of patch payloads used to perform the patch operation."
+        )
+    ]
+
+    @model_validator(mode='before')
+    @classmethod
+    def validate_ai_tool_patch_payload(cls, values):
+        if "payload" not in values or not isinstance(values.get("payload"), list):
+            raise PydanticCustomError(
+                "Invalid_payload",
+                "The payload is either missing or not in a valid list format."
+            )
+        return values