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

ibm_watsonx_gov/tools/ootb/vectordb/chromadb_retriever_tool.py

@@ -0,0 +1,111 @@
+# ----------------------------------------------------------------------------------------------------
+# 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, Optional, Type
+
+from langchain.callbacks.manager import CallbackManagerForToolRun
+from langchain.tools import BaseTool
+from pydantic import BaseModel, Field, PrivateAttr
+
+from ibm_watsonx_gov.tools.schemas.vectordb_retrieval_schema import \
+    ChromaDbRetrieverConfig
+
+
+class RetrievalToolInput(BaseModel):
+    query: Annotated[str, Field(description="Search query")]
+    n_results: Annotated[Optional[int], Field(
+        description="Number of results to return", default=5)]
+    where: Annotated[Optional[dict], Field(
+        description="A dict used to filter results by the metadata. E.g. {'source': 'student_info'} ", default=None)]
+    where_document: Annotated[Optional[dict], Field(
+        description="A dict used to filter by the documents. E.g. {$contains: 'My Text'}.", default=None)]
+
+
+class ChromaDBRetrievalTool(BaseTool):
+    """
+    Tool for searching the ChromaDB vector database.
+
+    Examples:
+        Basic usage
+            .. code-block:: python
+
+                case-1 : With path 
+                retriever_tool_config = {
+                    "path": "./medium_db",
+                    "collection_name":  "medium_articles",
+                    "embedding_function": embedding_functions.OpenAIEmbeddingFunction(),
+                    "n_results": 2
+                }
+}
+                chromadb_tool = ChromaDBRetrievalTool(**retriever_tool_config)
+                chromadb_tool.invoke({"query": "What is concept drift?"})
+    """
+
+    name: str = "chromadb_retrieval_tool"
+    description: str = "Search ChromaDB vector database and return the top-k results"
+    args_schema: Type[BaseModel] = RetrievalToolInput
+
+    _retriever: any = PrivateAttr()
+    _config = PrivateAttr()
+
+    def __init__(self, **kwargs):
+        """
+        Arguments for creating a Chromadb retrieval tool which connects to a local or remote ChromaDB server.
+
+        Args:
+            path: Path to the database
+            host: The hostname of the Chroma server. Defaults to "localhost". Required if path is not provided.
+            port: The port of the Chroma server. Defaults to "8000". Required if path is not provided
+            ssl: Whether to use SSL to connect to the Chroma server. Defaults to False.
+            headers: A dictionary of headers to send to the Chroma server. Defaults to {}.
+            settings: A dictionary of settings to communicate with the chroma server. Chromadb Settings object
+            tenant: The tenant to use for this client. Defaults to the default tenant.
+            database: The database to use for this client. Defaults to the default database.
+            collection_name: Name of the collection where documents are stored
+            embedding_function: Optional function to use to embed documents.
+                                Uses the default embedding function if not provided.
+            n_results: Optional number of results to return                    
+        """
+        super().__init__(**kwargs)
+
+        from ibm_watsonx_gov.tools.utils.package_utils import \
+            install_and_import_packages
+
+        install_and_import_packages(["chromadb"])
+        self._config = ChromaDbRetrieverConfig(**kwargs)
+        self._validate_and_create_retriever(self._config)
+
+        self._collection = self._retriever.get_collection(
+            name=self._config.collection_name, embedding_function=self._config.embedding_function)
+
+    def _run(self, query, run_manager: Optional[CallbackManagerForToolRun] = None, **kwargs):
+        """Performs a Chromadb vector search and returns the top n_results content"""
+
+        where = None
+        where_document = None
+
+        results = self._collection.query(
+            query_texts=[query],
+            where=where,
+            where_document=where_document,
+            n_results=self._config.n_results
+        )
+
+        return results['documents']
+
+    def _validate_and_create_retriever(self, config: ChromaDbRetrieverConfig):
+        import chromadb
+        if config.path != None:
+            self._retriever = chromadb.PersistentClient(config.path)
+        else:
+            from chromadb.config import DEFAULT_DATABASE, DEFAULT_TENANT
+            tenant = config.tenant or DEFAULT_TENANT
+            database = config.database or DEFAULT_DATABASE
+            self._retriever = chromadb.HttpClient(host=config.host, port=config.port, ssl=config.ssl, headers=config.headers,
+                                                  settings=config.settings, tenant=tenant, database=database)

ibm_watsonx_gov/tools/rest_api/__init__.py

@@ -0,0 +1,10 @@
+# ----------------------------------------------------------------------------------------------------
+# 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 .restapi_tool import RestApiTool, load_headers

ibm_watsonx_gov/tools/rest_api/restapi_tool.py

@@ -0,0 +1,72 @@
+# ----------------------------------------------------------------------------------------------------
+# 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 json
+import os
+from typing import Any, Dict, Type
+
+import requests
+from langchain.tools import BaseTool
+from pydantic import BaseModel
+
+
+class RestApiTool(BaseTool):
+    """
+      Class to construct a tool using  the rest api , schema dynamically
+    """
+    name: str
+    description: str
+    method: str = "POST"
+    endpoint: str = ""
+    headers: Dict[str, str] = {}
+    args_schema: Type[BaseModel]
+
+    def _run(self, **kwargs) -> Any:
+        if self.method.upper() == "GET":
+            response = requests.get(
+                self.endpoint, params=kwargs, headers=self.headers)
+        elif self.method.upper() == "POST":
+            parsed = self.args_schema(**kwargs)
+            payload = json.loads(parsed.model_dump_json())
+            response = requests.post(
+                self.endpoint, json=payload, headers=self.headers)
+        else:
+            raise ValueError(f"Unsupported method {self.method}")
+        return response.json()
+
+
+def load_headers(tool_name, headers: dict = {}) -> dict:
+    """Method to load the header using evv information
+
+    Args:
+        tool_name (_type_): Name of the tool 
+        headers (dict, optional): Headers associated. Defaults to {}.
+
+    Raises:
+        Exception: _description_
+
+    Returns:
+        dict: _description_
+    """
+    missing_keys = []
+    for key, value in headers.items():
+        if value == f"${key}":
+            value = os.getenv(key)
+            if value is None:
+                missing_keys.append(key)
+            else:
+                headers[key] = value
+        elif value == "$DYNAMIC_HEADER":
+            missing_keys.append(key)
+
+    if len(missing_keys) > 0:
+        raise Exception(
+            f"Missing header information while loading tool :{tool_name}. Details :{missing_keys}")
+    return headers

ibm_watsonx_gov/tools/schemas/__init__.py

@@ -0,0 +1,10 @@
+# ----------------------------------------------------------------------------------------------------
+# 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 .search_tool_schema import SearchToolConfig, SearchToolInput, GoogleSearchToolConfig
+from .vectordb_retrieval_schema import ChromaDbRetrieverConfig

ibm_watsonx_gov/tools/schemas/search_tool_schema.py

@@ -0,0 +1,46 @@
+# ----------------------------------------------------------------------------------------------------
+# 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, Optional
+
+from pydantic import BaseModel, Field
+
+
+# Define the args schema using Pydantic
+class SearchToolInput(BaseModel):
+    """
+        Model that can be used for setting input args for search tools
+    """
+    query: Annotated[str,
+                     Field(..., description="Search the input from the given search tool")]
+    top_k_results: Annotated[Optional[int],  Field(
+        3, description="Number of search results to retrieve")]
+
+
+class SearchToolConfig(BaseModel):
+    """
+        Model to be used for getting the initialization parameters for search tools
+    """
+    top_k_results: Annotated[Optional[int],  Field(
+        3, description="Number of search results to retrieve")]
+
+
+class GoogleSearchToolConfig(BaseModel):
+    """
+            Model to be used for getting the initialization parameters for Google search tool
+    """
+    top_k_results: Annotated[Optional[int], Field(
+        3, description="Number of search results to retrieve")]
+    serpapi_key: Annotated[
+        str,
+        Field(
+            ...,
+            description="API key used to authenticate requests to SerpAPI. You can generate your API key from https://serpapi.com/manage-api-key."
+        )
+    ]

ibm_watsonx_gov/tools/schemas/vectordb_retrieval_schema.py

@@ -0,0 +1,55 @@
+# ----------------------------------------------------------------------------------------------------
+# 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, Callable, Optional
+
+from pydantic import BaseModel, Field, model_validator
+
+
+class ChromaDbRetrieverConfig(BaseModel):
+    """
+        Configuration for retrieval tool
+    """
+
+    path: Annotated[str, Field(
+        description="Path to the database")]
+    host: Annotated[Optional[str], Field(
+        description="The hostname of the Chroma server. Defaults to 'localhost'. Required if path is not provided", default='localhost')]
+    port: Annotated[Optional[int], Field(
+        description="The port of the Chroma server. Defaults to 8000. Required if path is not provided", default=8000)]
+    ssl: Annotated[Optional[bool], Field(
+        description="Whether to use SSL to connect to the Chroma server. Defaults to False", default=False)]
+    headers: Annotated[Optional[dict], Field(
+        description="A dictionary of headers to send to the Chroma server. Defaults to {}", default={})]
+    settings: Annotated[Optional[dict], Field(
+        description="A dictionary of settings to communicate with the chroma server. Chromadb Settings object", default={})]
+    tenant: Annotated[Optional[str], Field(
+        description="The tenant to use for this Chroma server. Defaults to the default tenant.", default=None)]
+    database: Annotated[Optional[str], Field(
+        description="The database to use for this client. Defaults to the default database.", default=None)]
+    collection_name: Annotated[str, Field(
+        description="Name of the collection where documents are stored")]
+    embedding_function: Annotated[Callable, Field(
+        description="Function to use for embedded documents")]
+    n_results: Annotated[Optional[int], Field(
+        description="Optional number of results to return", default=3)]
+
+    @model_validator(mode="after")
+    def validate_config(self) -> "ChromaDbRetrieverConfig":
+        if self.path is None or (self.host is None and self.port is None):
+            raise Exception(
+                "Connection information for Chromadb is missing. Provide either the path to the database or provide host and port where it is running.")
+
+        if self.collection_name is None:
+            raise Exception("Provide name of the collection to work with.")
+
+        if self.embedding_function is None:
+            raise Exception("Missing embedding function")
+
+        return self

ibm_watsonx_gov/tools/utils/__init__.py

@@ -0,0 +1,14 @@
+# ----------------------------------------------------------------------------------------------------
+# 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 .display_utils import display_tools
+from .package_utils import install_and_import_packages
+from .python_utils import (get_base64_decoding, get_base64_encoding,
+                           validate_envs)
+from .tool_utils import TOOL_REGISTRY, get_pydantic_model, list_ootb_tools

ibm_watsonx_gov/tools/utils/constants.py

@@ -0,0 +1,69 @@
+# ----------------------------------------------------------------------------------------------------
+# 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 enum import Enum
+
+
+class CustomToolType(Enum):
+    CODE = "code"
+    ENDPOINT = "endpoint"
+
+    @classmethod
+    def values(cls):
+        return [member.value for member in cls]
+
+
+class ServiceProviderType(Enum):
+    IBM = "IBM"
+    CUSTOM = "custom"
+
+    @classmethod
+    def values(cls):
+        return [member.value for member in cls]
+
+
+class Categories(Enum):
+    SEARCH = "Search"
+    GUARDRAILS = "Guardrails"
+    QUERY = "Query"
+    RAG = "RAG"
+    EVALUATION = "Evaluation"
+    OTHER = "Other"
+
+    @classmethod
+    def values(cls):
+        return [member.value for member in cls]
+
+
+class PatchOperationTypes(Enum):
+    ADD = "add"
+    REMOVE = "remove"
+    REPLACE = "replace"
+
+    @classmethod
+    def values(cls):
+        return [member.value for member in cls]
+
+
+class ComponentTypes(Enum):
+    AI_TOOL = "ai_tool"
+    AI_AGENT = "ai_agent"
+
+    @classmethod
+    def values(cls):
+        return [member.value for member in cls]
+
+
+class Framework(Enum):
+    LANGCHAIN = "langchain"
+    LANGGRAPH = "langgraph"
+
+    @classmethod
+    def values(cls):
+        return [member.value for member in cls]

ibm_watsonx_gov/tools/utils/display_utils.py

@@ -0,0 +1,38 @@
+# ----------------------------------------------------------------------------------------------------
+# 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 pandas as pd
+
+try:
+    import itables
+    from itables import init_notebook_mode
+    init_notebook_mode(all_interactive=True)
+except Exception as e:
+    pass
+
+
+def display_tools(column_list: list = [], limit: int = None):
+    from ..clients.ai_tool_client import list_tools
+    if len(column_list) == 0:
+        column_list = ["asset_name", "service_provider_type",
+                       "category"]
+
+    details = list_tools(limit=limit)
+    tools = details['tools'] or []
+    if len(tools) > 0:
+        df = pd.DataFrame(tools)['entity'].values
+        df1 = pd.DataFrame(list(df))[column_list]
+        df1 = df1.dropna()
+
+        # Rename the column for ease of use
+        df1 = df1.rename(columns={'asset_name': 'tool_name'})
+    else:
+        df1 = pd.DataFrame()
+
+    itables.show(df1)

ibm_watsonx_gov/tools/utils/environment.py

@@ -0,0 +1,108 @@
+# ----------------------------------------------------------------------------------------------------
+# 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 os
+import warnings
+
+from ...entities.enums import Region
+from ...utils.python_utils import get_environment_variable_value
+from .platform_url_mapping import ALLOWED_PLATFORM_URL, PROD_PLATFORM_URL
+from ...utils.url_mapping import WATSONX_REGION_URLS
+
+
+def get_property_value(property_name, default=None):
+    if os.environ.get(property_name):
+        return os.environ.get(property_name)
+    else:
+        return default
+
+
+def get_base_url():
+    is_cpd = get_is_cpd()
+
+    if is_cpd:
+        watsonx_url = get_environment_variable_value(
+            possible_env_variables=["WATSONX_URL", "PLATFORM_URL", "WXG_URL"])
+        if not watsonx_url:
+            raise Exception(
+                "The WATSONX_URL cannot be empty for CPD environment")
+    else:
+        watsonx_region = get_property_value(property_name="WATSONX_REGION")
+        platform_url = get_property_value(property_name="PLATFORM_URL")
+
+        if watsonx_region:
+            watsonx_region_map = WATSONX_REGION_URLS.get(watsonx_region)
+            if not watsonx_region_map:
+                raise Exception(
+                    f"\nThe WATSONX_REGION '{watsonx_region}' is invalid or not supported.\n"
+                    f"Supported WATSONX_REGION are: {', '.join(Region.values())}."
+                )
+            watsonx_url = watsonx_region_map.dai_url
+        elif platform_url:
+            if platform_url not in ALLOWED_PLATFORM_URL:
+                raise Exception(
+                    f"\nThe platform URL '{platform_url}' is invalid or not supported.\n"
+                    f"Supported platform URLs are:\n{', '.join(PROD_PLATFORM_URL.keys())}."
+                )
+            watsonx_url = platform_url
+        else:
+            warnings.warn(
+                "Since WATSONX_REGION is not provided in the environment variable, the Dallas region will be used as the default.",
+                UserWarning)
+            # Setting Dallas region as default if not provided.
+            os.environ["PLATFORM_URL"] = "https://api.dataplatform.cloud.ibm.com"
+            watsonx_url = get_property_value(property_name="PLATFORM_URL")
+
+    return watsonx_url
+
+
+def get_authenticator_url():
+    base_url = get_base_url()
+    if get_is_cpd():
+        return base_url
+    else:
+        return ALLOWED_PLATFORM_URL.get(base_url)
+
+
+def get_api_key():
+    return get_property_value(property_name="WATSONX_APIKEY")
+
+
+def get_cpd_password():
+    return get_property_value(property_name="WATSONX_PASSWORD")
+
+
+def get_cpd_username():
+    return get_property_value(property_name="WATSONX_USERNAME")
+
+
+def get_is_cpd():
+    watsonx_version = get_property_value(property_name="WATSONX_VERSION")
+    if watsonx_version:
+        return True
+    else:
+        return False
+
+
+def get_ssl_verification():
+    verify = get_property_value(
+        property_name="WATSONX_DISABLE_SSL", default=False)
+    return str(verify).lower() in ("false", "0", "no")
+
+
+def get_service_instance_id():
+    return get_property_value(property_name="WXG_SERVICE_INSTANCE_ID")
+
+
+def get_wxai_url():
+    return get_property_value(property_name="WXAI_URL")
+
+
+def get_watsonx_region():
+    return get_property_value(property_name="WATSONX_REGION")

ibm_watsonx_gov/tools/utils/package_utils.py

@@ -0,0 +1,40 @@
+# ----------------------------------------------------------------------------------------------------
+# 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 subprocess
+import sys
+from typing import List
+
+import pkg_resources
+
+
+def is_package_installed(package_name):
+    try:
+        pkg_resources.get_distribution(package_name)
+        return True
+    except pkg_resources.DistributionNotFound:
+        return False
+
+
+def install_and_import_packages(packages: List[str]):
+    try:
+        # Find missing packages
+        missing_packages = []
+        for p in packages:
+            if not is_package_installed(p):
+                missing_packages.append(p)
+
+        # Install missing packages if any
+        if missing_packages:
+            print(
+                f"Installing missing packages: {', '.join(missing_packages)}")
+            subprocess.check_call([sys.executable, "-m", "pip",
+                                   "install"] + missing_packages)
+    except Exception as ex:
+        raise ImportError(f"Error while installing {packages}.Details:{ex}")

ibm_watsonx_gov/tools/utils/python_utils.py

@@ -0,0 +1,68 @@
+# ----------------------------------------------------------------------------------------------------
+# 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.
+# ----------------------------------------------------------------------------------------------------
+
+# Convert string to bytes and encode
+import base64
+import os
+import json
+
+from ibm_watsonx_gov.utils.python_utils import get
+
+
+# Source code encoder
+def get_base64_encoding(tool_code: str):
+    return base64.b64encode(tool_code.encode()).decode()
+
+
+# Decode the base64 string
+def get_base64_decoding(encoded_code: str):
+    return base64.b64decode(encoded_code).decode()
+
+
+# Validate envs:
+def validate_envs(tool_name: str, env_list: list = []):
+    if len(env_list) == 0:
+        return env_list
+
+    import os
+    missing_envs = []
+    for env in env_list:
+        value = os.getenv(env)
+        if value is None:
+            missing_envs.append(env)
+
+    if len(missing_envs) > 0:
+        raise Exception(
+            f"Please set environment values :{missing_envs} for using the tool:{tool_name}")
+
+
+# Extracting bss_account id from user_token
+def get_bss_account_id(token: str):
+    payload_b64 = token.split('.')[1]
+    padding = '=' * (-len(payload_b64) % 4)
+    payload_b64 = base64.urlsafe_b64decode(payload_b64 + padding)
+    payload = json.loads(payload_b64)
+    return get(payload, "bss", default="")
+
+
+# Processing endpoint response
+def process_result(result):
+    try:
+        if result.text:
+            try:
+                result = json.loads(result.text)
+            except json.JSONDecodeError:
+                raise Exception(result.text)
+
+            if result.get("errors"):
+                raise Exception(result.get("errors"))
+            return result
+        raise Exception(f"Empty response received. Status code: {result.status_code}")
+    except Exception as e:
+        raise Exception(str(e))