Source code for langchain_ibm.agent_toolkits.utility.toolkit
"""IBM watsonx.ai Toolkit wrapper."""
from typing import Any, cast
from ibm_watsonx_ai import APIClient # type: ignore[import-untyped]
from ibm_watsonx_ai.foundation_models.utils import ( # type: ignore[import-untyped]
Tool,
Toolkit,
)
from langchain_core.callbacks import CallbackManagerForToolRun
from langchain_core.tools.base import BaseTool, BaseToolkit
from langchain_core.utils.utils import secret_from_env
from pydantic import (
BaseModel,
ConfigDict,
Field,
PrivateAttr,
SecretStr,
create_model,
model_validator,
)
from typing_extensions import Self
from langchain_ibm.agent_toolkits.utility.utils import convert_to_watsonx_tool
from langchain_ibm.utils import resolve_watsonx_credentials
[docs]
class WatsonxTool(BaseTool):
"""IBM watsonx.ai Tool."""
name: str
"""Name of the tool."""
description: str
"""Description of what the tool is used for."""
agent_description: str | None = None
"""The precise instruction to agent LLMs
and should be treated as part of the system prompt."""
tool_input_schema: dict | None = None
"""Schema of the input that is provided when running the tool if applicable."""
tool_config_schema: dict | None = None
"""Schema of the config that can be provided when running the tool if applicable."""
tool_config: dict | None = None
"""Config properties to be used when running a tool if applicable."""
args_schema: type[BaseModel] = BaseModel
_watsonx_tool: Tool | None = PrivateAttr(default=None) #: :meta private:
watsonx_client: APIClient = Field(exclude=True)
@model_validator(mode="after")
def validate_tool(self) -> Self:
"""Validate tool."""
self._watsonx_tool = Tool(
api_client=self.watsonx_client,
name=self.name,
description=self.description,
agent_description=self.agent_description,
input_schema=self.tool_input_schema,
config_schema=self.tool_config_schema,
)
converted_tool = convert_to_watsonx_tool(self)
json_schema = converted_tool["function"]["parameters"]
self.args_schema = _json_schema_to_pydantic_model(
name="ToolArgsSchema",
schema=json_schema,
)
return self
def _run(
self,
*args: Any,
run_manager: CallbackManagerForToolRun | None = None,
**kwargs: Any,
) -> dict:
"""Run the tool."""
if self.tool_input_schema is None:
input = kwargs.get("input") or args[0]
else:
input = {
k: v
for k, v in kwargs.items()
if k in self.tool_input_schema["properties"]
}
return cast("dict", self._watsonx_tool.run(input, self.tool_config)) # type: ignore[union-attr]
[docs]
def set_tool_config(self, tool_config: dict) -> None:
"""Set tool config properties.
???+ example "Example"
```python
google_search = watsonx_toolkit.get_tool("GoogleSearch")
print(google_search.tool_config_schema)
tool_config = {"maxResults": 3}
google_search.set_tool_config(tool_config)
```
"""
self.tool_config = tool_config
[docs]
class WatsonxToolkit(BaseToolkit):
"""IBM watsonx.ai Toolkit.
???+ info "Setup"
To use, you should have `langchain_ibm` python package installed,
and the environment variable `WATSONX_APIKEY` set with your API key, or pass
it as a named parameter `apikey` to the constructor.
```bash
pip install -U langchain-ibm
# or using uv
uv add langchain-ibm
```
```bash
export WATSONX_APIKEY="your-api-key"
```
??? info "Instantiate"
IBM watsonx.ai for IBM Cloud:
```python
from langchain_ibm.agent_toolkits.utility import WatsonxToolkit
watsonx_toolkit = WatsonxToolkit(
url="https://us-south.ml.cloud.ibm.com",
project_id="*****", # or `space_id`
apikey="*****", # not needed if `WATSONX_APIKEY` is set
)
```
IBM watsonx.ai software:
```python
from langchain_ibm.agent_toolkits.utility import WatsonxToolkit
watsonx_toolkit = WatsonxToolkit(
url="<CPD_URL>",
project_id="*****", # or `space_id`
username="*****",
password="*****",
instance_id="*****",
version="*****", # optional
)
```
??? info "Invoke"
```python
tools = watsonx_toolkit.get_tools()
google_search = watsonx_toolkit.get_tool(tool_name="GoogleSearch")
tool_config = {
"maxResults": 3,
}
google_search.set_tool_config(tool_config)
input = {
"input": "Search IBM",
}
search_result = google_search.invoke(input)
```
??? info "Run"
```python
rag_query = watsonx_toolkit.get_tool(tool_name="RAGQuery")
rag_query.set_tool_config(
{
"vectorIndexId": "<vector-index-id>",
"projectId": "<project-id>",
}
)
res = rag_query.run("How to initialize APIClient?")
```
"""
project_id: str | None = None
"""ID of the watsonx.ai Studio project."""
space_id: str | None = None
"""ID of the watsonx.ai Studio space."""
url: SecretStr = Field(
alias="url",
default_factory=secret_from_env("WATSONX_URL", default=None), # type: ignore[assignment]
)
"""URL to the watsonx.ai Runtime."""
apikey: SecretStr | None = Field(
alias="apikey",
default_factory=secret_from_env("WATSONX_APIKEY", default=None),
)
"""API key to the watsonx.ai Runtime."""
token: SecretStr | None = Field(
alias="token",
default_factory=secret_from_env("WATSONX_TOKEN", default=None),
)
"""Token to the watsonx.ai Runtime."""
password: SecretStr | None = Field(
alias="password",
default_factory=secret_from_env("WATSONX_PASSWORD", default=None),
)
"""Password to the CPD instance."""
username: SecretStr | None = Field(
alias="username",
default_factory=secret_from_env("WATSONX_USERNAME", default=None),
)
"""Username to the CPD instance."""
instance_id: SecretStr | None = Field(
alias="instance_id",
default_factory=secret_from_env("WATSONX_INSTANCE_ID", default=None),
)
"""Instance_id of the CPD instance."""
version: SecretStr | None = None
"""Version of the CPD instance."""
verify: str | bool | None = None
"""You can pass one of following as verify:
* the path to a CA_BUNDLE file
* the path of directory with certificates of trusted CAs
* True - default path to truststore will be taken
* False - no verification will be made"""
_tools: list[WatsonxTool] | None = None
"""Tools in the toolkit."""
_watsonx_toolkit: Toolkit | None = PrivateAttr(default=None) #: :meta private:
watsonx_client: APIClient | None = Field(default=None, exclude=True)
model_config = ConfigDict(arbitrary_types_allowed=True)
@model_validator(mode="after")
def validate_environment(self) -> Self:
"""Validate that credentials and python package exists in environment."""
if isinstance(self.watsonx_client, APIClient):
self._watsonx_toolkit = Toolkit(self.watsonx_client)
else:
credentials = resolve_watsonx_credentials(
url=self.url,
apikey=self.apikey,
token=self.token,
password=self.password,
username=self.username,
instance_id=self.instance_id,
version=self.version,
verify=self.verify,
)
self.watsonx_client = APIClient(
credentials=credentials,
project_id=self.project_id,
space_id=self.space_id,
)
self._watsonx_toolkit = Toolkit(self.watsonx_client)
self._tools = [
WatsonxTool(
watsonx_client=self.watsonx_client,
name=tool["name"],
description=tool["description"],
agent_description=tool.get("agent_description"),
tool_input_schema=tool.get("input_schema"),
tool_config_schema=tool.get("config_schema"),
)
for tool in self._watsonx_toolkit.get_tools()
]
return self
[docs]
def get_tools(self) -> list[WatsonxTool]: # type: ignore[override]
"""Get the tools in the toolkit."""
return self._tools # type: ignore[return-value]
[docs]
def get_tool(self, tool_name: str) -> WatsonxTool:
"""Get the tool with a given name."""
for tool in self.get_tools():
if tool.name == tool_name:
return tool
error_msg = f"A tool with the given name ({tool_name}) was not found."
raise ValueError(error_msg)
def _json_schema_to_pydantic_model(
name: str,
schema: dict[str, Any],
) -> type[BaseModel]:
properties = schema.get("properties", {})
fields = {}
type_mapping = {
"string": str,
"integer": int,
"number": float,
"boolean": bool,
"array": list,
"object": dict,
}
for field_name, field_schema in properties.items():
field_type = field_schema.get("type", "string")
is_required = field_name in schema.get("required", [])
py_type = type_mapping.get(field_type, Any)
fields[field_name] = (py_type, ... if is_required else None)
return cast("type[BaseModel]", create_model(name, **fields)) # type: ignore[call-overload]