Skip to content
For Self-Managed AI Platform users running v11.1, see the on-premise platform documentation
Agentic AI > Develop agentic workflows > Add tools to agents

Add tools to agents

To add local tools to your agents, you can start by modifying the agent.py file or adding new files to the custom_model directory; the following examples will help you get started. Note that the structure and implementation of a tool is framework-specific. The examples provided here are for CrewAI, LangGraph, LlamaIndex, and NAT (NVIDIA NeMo Agent Toolkit). If you are using another framework, you will need to refer to the framework's documentation for details on how to implement tools.

Call tools from an agent

Once you have defined your tool, you can call it from your agent by adding it to the list of tools available to the agent. This is typically done in the agent.py file where the agent is defined. You will need to import the tool class and add it to the agent's tool list. The following simple examples illustrate one way of doing this for each framework.

To add a tool to a CrewAI agent, you can modify an agent in the agent.py file in the custom_model directory. An example of modifying the agent_planner agent to use the sample tool from the Local Datetime Tool Example is shown below:

@property
def agent_planner(self) -> Agent:
    datetime_tool = DateTimeTool()  # Import and instantiate your tool here

    return Agent(
        role="Content Planner",
        goal="Plan engaging and factually accurate content on {topic}",
        backstory="...",  # truncated for brevity in this example
        allow_delegation=False,
        verbose=self.verbose,
        llm=self.llm,
        tools=[datetime_tool] # Add your tool to the tools list here
    )

To add a tool to a LangGraph agent, you can modify an agent in the agent.py file in the custom_model directory. An example of modifying the agent_planner agent to use the sample tool from the Local Datetime Tool Example is shown below:

@property
def agent_planner(self) -> CompiledGraph:
    datetime_tool = DateTimeTool()  # Import and instantiate your tool here
    return create_react_agent(
        self.llm,
        tools=[datetime_tool], # Add your tool to the tools list here
        prompt=self.make_system_prompt(
            "...", # truncated for brevity in this example
        ),
    )

To add a tool to a LlamaIndex agent, you can modify an agent in the agent.py file in the custom_model directory. An example of modifying the research_agent agent to use the sample tool from the Local Datetime Tool Example is shown below:

@property
def research_agent(self) -> FunctionAgent:
    datetime_tool = DateTimeTool()  # Import and instantiate your tool here
    return FunctionAgent(
        name="ResearchAgent",
        description="...", # truncated for brevity in this example
        system_prompt=(
            "..." # truncated for brevity in this example
        ),
        llm=self.llm,
        tools=[self.record_notes, datetime_tool], # Add your tool to the tools list here
        can_handoff_to=["WriteAgent"],
    )

In NAT templates, tools are defined as functions in the workflow.yaml file. Functions are defined in the functions section with _type: chat_completion and referenced in the workflow's tool_list. Refer to the NVIDIA NeMo Agent Toolkit documentation for details on implementing tools in NAT.

Use predefined tools

Some frameworks provide predefined tools that you can use directly in your agents. For example, CrewAI provides a SearchTool that can be used to perform web searches. You can refer to the framework documentation for a list of predefined tools and how to use them. These tools can be added to your agent by simply importing them from the framework and adding them to the agent's tool list as shown in the examples above.

LangChain compatibility

Most agentic workflow frameworks (including CrewAI, LangGraph, and many others) are natively compatible with LangChain tools. This means you can import and use them without any modifications. You can refer to the LangChain Tools Documentation for more information on using LangChain tools in your agents.

Local datetime tool

The following examples show how to create a custom local datetime tool for your agents. This tool returns the current date and time, allowing the agent to be aware of and use the current date and time in its responses or actions. This tool does not require any network or file access and can be implemented and run without any additional permissions or credentials.

Starting point for local tools

This example can be used as a starting point for creating other local tools that do not require external access.

To add a local datetime tool to a CrewAI agent, you can modify the agent.py file in the custom_model directory. You can add the following code to define the datetime tool:

from datetime import datetime
from zoneinfo import ZoneInfo
from crewai.tools import BaseTool

class DateTimeTool(BaseTool):
    name = "datetime_tool"
    description = (
        "Returns the current date and time. Optionally accepts a timezone parameter as an IANA timezone "
        "string (e.g., 'America/New_York', 'Europe/London', 'Asia/Tokyo'). If no timezone is provided, "
        "returns local time.")

    def run(self, query: str = "", timezone: str = None) -> str:
        try:
            # If the agent provides the timezone parameter, use it to get the current time in that timezone
            if timezone:
                # Use the specified timezone
                tz = ZoneInfo(timezone)
                current_time = datetime.now(tz)
                return f"{current_time.strftime('%Y-%m-%d %H:%M:%S')} ({timezone})"

            # Return the current local time if the agent does not have or provide the timezone parameter
            else:
                # Use local time
                return datetime.now().strftime("%Y-%m-%d %H:%M:%S")

        # Gracefully handle errors and exceptions so the agent can understand the error
        # and attempt to correct in a followup call of the tool if needed.
        except Exception as e:
            return (
                f"Error: Invalid timezone '{timezone}'. "
                f"Please use a valid IANA timezone string (e.g., 'America/New_York', 'Europe/London')."
            )

To add a local datetime tool to a LangGraph agent, you can modify the agent.py file in the custom_model directory. You can add the following code to define the datetime tool:

from datetime import datetime
from zoneinfo import ZoneInfo
from langchain_core.tools import Tool

class DateTimeTool(Tool):
    name = "datetime_tool"
    description = (
        "Returns the current date and time. Optionally accepts a timezone parameter as an IANA timezone "
        "string (e.g., 'America/New_York', 'Europe/London', 'Asia/Tokyo'). If no timezone is provided, "
        "returns local time.")

    def run(self, query: str = "", timezone: str = None) -> str:
        try:
            # If the agent provides the timezone parameter, use it to get the current time in that timezone
            if timezone:
                # Use the specified timezone
                tz = ZoneInfo(timezone)
                current_time = datetime.now(tz)
                return f"{current_time.strftime('%Y-%m-%d %H:%M:%S')} ({timezone})"

            # Return the current local time if the agent does not have or provide the timezone parameter
            else:
                # Use local time
                return datetime.now().strftime("%Y-%m-%d %H:%M:%S")

        # Gracefully handle errors and exceptions so the agent can understand the error
        # and attempt to correct in a followup call of the tool if needed.
        except Exception as e:
            return (
                f"Error: Invalid timezone '{timezone}'. "
                f"Please use a valid IANA timezone string (e.g., 'America/New_York', 'Europe/London')."
            )

To add a local datetime tool to a LlamaIndex agent, you can modify the agent.py file in the custom_model directory. You can add the following code to define the datetime tool:

from datetime import datetime
from zoneinfo import ZoneInfo
from llama_index.tools import BaseTool

class DateTimeTool(BaseTool):
    name = "datetime_tool"
    description = (
        "Returns the current date and time. Optionally accepts a timezone parameter as an IANA timezone "
        "string (e.g., 'America/New_York', 'Europe/London', 'Asia/Tokyo'). If no timezone is provided, "
        "returns local time.")

    def _run(self, query: str = "", timezone: str = None) -> str:
        try:
            # If the agent provides the timezone parameter, use it to get the current time in that timezone
            if timezone:
                # Use the specified timezone
                tz = ZoneInfo(timezone)
                current_time = datetime.now(tz)
                return f"{current_time.strftime('%Y-%m-%d %H:%M:%S')} ({timezone})"

            # Return the current local time if the agent does not have or provide the timezone parameter
            else:
                # Use local time
                return datetime.now().strftime("%Y-%m-%d %H:%M:%S")

        # Gracefully handle errors and exceptions so the agent can understand the error
        # and attempt to correct in a followup call of the tool if needed.
        except Exception as e:
            return (
                f"Error: Invalid timezone '{timezone}'. "
                f"Please use a valid IANA timezone string (e.g., 'America/New_York', 'Europe/London')."
            )

Refer to the NVIDIA NeMo Agent Toolkit documentation for details on implementing local tools in NAT.

Weather API tool

The following examples show how to create a custom tool that fetches weather information from a public API. This tool will require network access to fetch the weather data. You will need to ensure that your agent has the necessary permissions to make network requests and the appropriate API credentials required by the weather service.

Starting point for external API tools

This example can be used as a starting point for tools that need to communicate with external APIs.

To add a weather API tool to a CrewAI agent, you can modify the agent.py file in the custom_model directory. You can add the following code to define the weather tool:

import requests
from crewai.tools import BaseTool

class WeatherTool(BaseTool):
    name = "weather_tool"
    description = (
        "Fetches the current weather for a specified city. Usage: weather_tool(city='City Name'). "
        "Requires an API key from OpenWeatherMap. Sign up at https://openweathermap.org/api.")

    def run(self, city: str) -> str:
        api_key = "YOUR_API_KEY"  # Replace with your OpenWeatherMap API key
        base_url = "http://api.openweathermap.org/data/2.5/weather"
        params = {"q": city, "appid": api_key, "units": "metric"}

        try:
            # Submit a query to an API using requests
            response = requests.get(base_url, params=params, timeout=10)
            response.raise_for_status()

            # Collect and format the response
            data = response.json()
            weather = data['weather'][0]
            main = data['main']

            # Format and return the response to the agent
            return (
                f"Current weather in {data['name']}, {data['sys']['country']}:\n"
                f"Temperature: {main['temp']}°C (feels like {main['feels_like']}°C)\n"
                f"Condition: {weather['main']} - {weather['description']}\n"
                f"Humidity: {main['humidity']}%\n"
                f"Pressure: {main['pressure']} hPa"
            )

        # Gracefully handle errors and exceptions so the agent can understand the error
        # and attempt to correct in a followup call of the tool if needed.
        except requests.exceptions.RequestException as e:
            return f"Error fetching weather data: {str(e)}"
        except KeyError as e:
            return f"Error parsing weather data: Missing key {str(e)}"
        except Exception as e:
            return f"Unexpected error: {str(e)}"

To add a weather API tool to a LangGraph agent, you can modify the agent.py file in the custom_model directory. You can add the following code to define the weather tool:

import requests
from langchain_core.tools import Tool

class WeatherTool(Tool):
    name = "weather_tool"
    description = (
        "Fetches the current weather for a specified city. Usage: weather_tool(city='City Name'). "
        "Requires an API key from OpenWeatherMap. Sign up at https://openweathermap.org/api.")

    def run(self, city: str) -> str:
        api_key = "YOUR_API_KEY"  # Replace with your OpenWeatherMap API key
        base_url = "http://api.openweathermap.org/data/2.5/weather"
        params = {"q": city, "appid": api_key, "units": "metric"}

        try:
            # Submit a query to an API using requests
            response = requests.get(base_url, params=params, timeout=10)
            response.raise_for_status()

            # Collect and format the response
            data = response.json()
            weather = data['weather'][0]
            main = data['main']

            # Format and return the response to the agent
            return (
                f"Current weather in {data['name']}, {data['sys']['country']}:\n"
                f"Temperature: {main['temp']}°C (feels like {main['feels_like']}°C)\n"
                f"Condition: {weather['main']} - {weather['description']}\n"
                f"Humidity: {main['humidity']}%\n"
                f"Pressure: {main['pressure']} hPa"
            )

        # Gracefully handle errors and exceptions so the agent can understand the error
        # and attempt to correct in a followup call of the tool if needed.
        except requests.exceptions.RequestException as e:
            return f"Error fetching weather data: {str(e)}"
        except KeyError as e:
            return f"Error parsing weather data: Missing key {str(e)}"
        except Exception as e:
            return f"Unexpected error: {str(e)}"

To add a weather API tool to a LlamaIndex agent, you can modify the agent.py file in the custom_model directory. You can add the following code to define the weather tool:

import requests
from llama_index.tools import BaseTool

class WeatherTool(BaseTool):
    name = "weather_tool"
    description = (
        "Fetches the current weather for a specified city. Usage: weather_tool(city='City Name'). "
        "Requires an API key from OpenWeatherMap. Sign up at https://openweathermap.org/api.")

    def _run(self, city: str) -> str:
        api_key = "YOUR_API_KEY"  # Replace with your OpenWeatherMap API key
        base_url = "http://api.openweathermap.org/data/2.5/weather"
        params = {"q": city, "appid": api_key, "units": "metric"}

        try:
            # Submit a query to an API using requests
            response = requests.get(base_url, params=params, timeout=10)
            response.raise_for_status()

            # Collect and format the response
            data = response.json()
            weather = data['weather'][0]
            main = data['main']

            # Format and return the response to the agent
            return (
                f"Current weather in {data['name']}, {data['sys']['country']}:\n"
                f"Temperature: {main['temp']}°C (feels like {main['feels_like']}°C)\n"
                f"Condition: {weather['main']} - {weather['description']}\n"
                f"Humidity: {main['humidity']}%\n"
                f"Pressure: {main['pressure']} hPa"
            )

        # Gracefully handle errors and exceptions so the agent can understand the error
        # and attempt to correct in a followup call of the tool if needed.
        except requests.exceptions.RequestException as e:
            return f"Error fetching weather data: {str(e)}"
        except KeyError as e:
            return f"Error parsing weather data: Missing key {str(e)}"
        except Exception as e:
            return f"Unexpected error: {str(e)}"

Refer to the NVIDIA NeMo Agent Toolkit documentation for details on implementing external API tools in NAT.

Integrate DataRobot deployed tools

DataRobot provides a set of global tools that can be used across different agent frameworks. These tools are designed to interact with DataRobot's platform and services. The following examples demonstrate a global tool that searches the DataRobot Data Registry for datasets. This tool will require network access to fetch the data from DataRobot. Before integrating these tools into an agent, deploy them in DataRobot.

DataRobot service integration

This example can be modified to create any tool that interacts with DataRobot's services.

Source code for many global tools can be found in the agent-tool-templates repository.

When building agents with external tools—either global tools or custom tools created as custom models in Workshop—the following components are required to call deployed tools from an agent:

  • The ToolClient class instance from the datarobot-genai package.
  • The authorization context initialization function from the datarobot-genai package.
  • A deployed tool's deployment_id, defined in model-metadata.yaml.
  • The tool's metadata, defined in a tool module; for example, tool_ai_catalog_search.py.

The examples in this documentation show how to add tool integration capabilities to agent templates. The base agent templates provided in this repository do not include tool implementations by default—these examples demonstrate the patterns you can follow to integrate deployed tools into your own agent implementations.

To assemble an agentic workflow using agentic tools deployed from Registry, an example workflow could include the following files:

File Contents
__init__.py Python package initialization file, making the directory a Python package.
custom.py The custom model code implementing the Bolt-on Governance API (the chat hook) to call the LLM and also passing those parameters to the agent (defined in agent.py).
agent.py The agent code, implementing the agentic workflow in the MyAgent class with the required invoke method. Tool integration properties can be added to this class to interface with deployed tools.
config.py The code for loading the configuration from environment variables, runtime parameters, and DataRobot credentials.
mcp_client.py The code providing MCP server connection management for tool integration (optional, only needed when using MCP tools).
tool_deployment.py The BaseTool class code, containing all necessary metadata for implementing tools.
tool.py The code for interfacing with the deployed tool, defining the input arguments and schema. Often, this file won't be named tool.py, as you may implement more than one tool. In this example, this functionality is defined in tool_ai_catalog_search.py.
model-metadata.yaml The custom model metadata and runtime parameters required by the agentic workflow.
pyproject.toml The libraries (and versions) required by the agentic workflow, using modern Python packaging standards.

Implement the ToolClient class instance

Every agent template and framework requires the ToolClient class from the datarobot-genai package to offload tool call processing to deployed global tools. The tool client calls the deployed tool and returns the results to the agent. To import the ToolClient module into an agent.py file, use the following import statement:

agent.py
from datarobot_genai.core.chat.client import ToolClient

The ToolClient is available in the datarobot-genai package. It defines the API endpoint and deployment ID for the deployed tool, gets the authorization context (if required), and provides interfaces for the score, score_unstructured, and chat hooks.

After you import the ToolClient into agent.py, you can add a tool_client property to your MyAgent class. The ToolClient automatically uses environment variables DATAROBOT_API_TOKEN and DATAROBOT_ENDPOINT for authentication, but you can also pass these explicitly if needed.

agent.py
from datarobot_genai.core.chat.client import ToolClient

class MyAgent:

    # More agentic workflow code.

    @property
    def tool_client(self):
        """ToolClient instance for calling deployed tools."""
        return ToolClient(
            api_key=self.api_key,
            base_url=self.api_base,
        )

    # More agentic workflow code.

(Optional) Initialize authorization context for external tools

Authorization context is required to allow downstream agents and tools to retrieve access tokens when connecting to external services. The authorization context functionality is available in the datarobot-genai package alongside the ToolClient class.

The initialize_authorization_context function is available in the datarobot-genai package and handles setting the authorization context for the current process. This function uses utility methods from the DataRobot Python API client: * set_authorization_context: A method to set the authorization context for the current process. * get_authorization_context: A method to retrieve the authorization context for the current process.

OAuth utility method availability

These utility methods are available in the DataRobot Python API client starting with version 3.8.0.

You can review the initialize_authorization_context function in the datarobot-genai package.

The initialize_authorization_context function is available in the datarobot-genai package. It sets the authorization context for the agent, which is required for propagating information needed by downstream agents and tools to retrieve access tokens to connect to external services. When set, authorization context will be automatically propagated when using the ToolClient class.

In the custom.py example below, the chat() hook calls initialize_authorization_context (imported from datarobot-genai) each time a chat request is made to the agentic workflow, providing any credentials required for external tools.

custom.py
from typing import Iterator, Union

from agent import MyAgent
from datarobot_drum import RuntimeParameters
from datarobot_genai.core.chat import (
    CustomModelChatResponse,
    CustomModelStreamingResponse,
    initialize_authorization_context,
    to_custom_model_chat_response,
    to_custom_model_streaming_response,
)
from openai.types.chat import CompletionCreateParams
from openai.types.chat.completion_create_params import (
    CompletionCreateParamsNonStreaming,
    CompletionCreateParamsStreaming,
)


def load_model(code_dir: str) -> str:
    """The agent is instantiated in this function and returned."""
    _ = code_dir
    return "success"


def chat(
    completion_create_params: CompletionCreateParams
    | CompletionCreateParamsNonStreaming
    | CompletionCreateParamsStreaming,
    model: str,
) -> Union[CustomModelChatResponse, Iterator[CustomModelStreamingResponse]]:
    """When using the chat endpoint, this function is called.

    Agent inputs are in OpenAI message format and defined as the 'user' portion
    of the input prompt.
    """
    _ = model

    # Initialize the authorization context for downstream agents and tools to retrieve
    # access tokens for external services.
    initialize_authorization_context(completion_create_params)

    # Instantiate the agent, all fields from the completion_create_params are passed to the agent
    # allowing environment variables to be passed during execution
    agent = MyAgent(**completion_create_params)

    if completion_create_params.get("stream"):
        streaming_response_generator = agent.invoke(
            completion_create_params=completion_create_params
        )
        return to_custom_model_streaming_response(
            streaming_response_generator, model=completion_create_params.get("model")
        )
    else:
        # Synchronous non-streaming response
        response_text, pipeline_interactions, usage_metrics = agent.invoke(
            completion_create_params=completion_create_params
        )
        return to_custom_model_chat_response(
            response_text,
            pipeline_interactions,
            usage_metrics,
            model=completion_create_params.get("model"),
        )

When authorization context is set, it is automatically propagated by the ToolClient class from the datarobot-genai package. The ToolClient.call() method automatically includes the authorization context when calling deployed tools.

The ToolClient class provides methods to call the custom model tool using various hooks: score, score_unstructured, and chat. When the authorization_context is set, the client automatically propagates it to the agent tool. The authorization_context is required for retrieving access tokens to connect to external services.

If you're implementing a custom tool reliant on an external service, you can use the @datarobot_tool_auth decorator to streamline the process of retrieving the authorization context, extracting the relevant data, and connecting to the DataRobot API to obtain the OAuth access token from an OAuth provider configured in DataRobot. When only one OAuth provider is configured the decorator doesn't require the provider parameter, as it will use the only available provider; however, if multiple providers are (or will be) available, you should define this parameter.

tool.py
from datarobot.models.genai.agent.auth import datarobot_tool_auth, AuthType

# More tool code.

@datarobot_tool_auth(
  type=AuthType.OBO,  # on-behalf-of
  provider="google",  # required with multiple OAuth providers
)
def list_files_in_google_drive(folder_name: str, token: str = "") -> list[dict]:
  """The value for token parameter will be injected by the decorator."""

    # More tool code.

Interface with tool deployments

Global tools and tools custom-built in the Registry workshop must be deployed for the agent to call them. When these tools are deployed, communicating with them requires a deployment ID, used to interface with the tool through the DataRobot API. The primary method for providing a deployed tool's deployment ID to the agent is through environment variables, defined as runtime parameters in the agent's metadata. To provide this metadata, create or modify a model-metadata.yaml file to add the runtime parameter for each deployed tool the agent needs to communicate with. Define runtime parameters in runtimeParameterDefinitions.

model-metadata.yaml
runtimeParameterDefinitions:
- fieldName: AI_CATALOG_SEARCH_TOOL_DEPLOYMENT_ID
    defaultValue: SET_VIA_PULUMI_OR_MANUALLY
    type: string

The example below illustrates a model-metadata.yaml file configured for an agenticworkflow implementing two global tools, Search Data Registry and Get Data Registry Dataset. The field names in the example below are used by DataRobot agent templates implementing these tools; however, the fieldName is configurable, and must match the implementation in the agent's code, located in the agent.py file.

model-metadata.yaml
---
name: agent_with_tools
type: inference
targetType: agenticworkflow
runtimeParameterDefinitions:
  - fieldName: OTEL_SDK_ENABLED
    defaultValue: true
    type: boolean
  - fieldName: LLM_DEPLOYMENT_ID
    defaultValue: SET_VIA_PULUMI_OR_MANUALLY
    type: string
  - fieldName: DATA_REGISTRY_SEARCH_TOOL_DEPLOYMENT_ID
    defaultValue: SET_VIA_PULUMI_OR_MANUALLY
    type: string
  - fieldName: DATA_REGISTRY_READ_TOOL_DEPLOYMENT_ID
    defaultValue: SET_VIA_PULUMI_OR_MANUALLY
    type: string

Unless you provided the required values as the defaultValue, you must set the runtime parameter to provide the deployment IDs to the agent's code. You can do this in two ways:

  • Manually: Configuring the values in the UI, in the Registry workshop, before deploying the agent.
  • Automatically: Configuring the values in a Pulumi script and passing them to the custom model.

When the parameters are set, they're accessible in the agent code, as long as the RuntimeParameters class is imported from datarobot_drum. The following example shows how you can add tool properties to your MyAgent class to interface with deployed tools:

agent.py
import os
from datarobot_drum import RuntimeParameters
from datarobot_genai.core.chat.client import ToolClient

class MyAgent:

    # More agentic workflow code.

    @property
    def tool_client(self) -> ToolClient:
        """ToolClient instance for calling deployed tools."""
        return ToolClient(
            api_key=self.api_key,
            base_url=self.api_base,
        )

    @property
    def tool_ai_catalog_search(self) -> BaseTool:
        """Search AI Catalog tool using deployed tool."""
        deployment_id = os.environ.get("AI_CATALOG_SEARCH_TOOL_DEPLOYMENT_ID")
        if not deployment_id:
            deployment_id = RuntimeParameters.get("AI_CATALOG_SEARCH_TOOL_DEPLOYMENT_ID")

        return SearchAICatalogTool(
            tool_client=self.tool_client,
            deployment_id=deployment_id
        )

    @property
    def tool_ai_catalog_read(self) -> BaseTool:
        """Read AI Catalog tool using deployed tool."""
        deployment_id = os.environ.get("AI_CATALOG_READ_TOOL_DEPLOYMENT_ID")
        if not deployment_id:
            deployment_id = RuntimeParameters.get("AI_CATALOG_READ_TOOL_DEPLOYMENT_ID")

        return ReadAICatalogTool(
            tool_client=self.tool_client,
            deployment_id=deployment_id,
        )

    # More agentic workflow code.

Define tool metadata

When building tools for an agent, the metadata defines how the agent LLM should call the tool. The more details the metadata provides, the more effectively the LLM uses the tool. The metadata includes the tool description and each arguments' schema and related description. Each framework has a unique way to define this metadata; however, in most cases you can leverage pydantic to import BaseModel to define the tool's arguments.

tool_ai_catalog_search.py
from pydantic import BaseModel as PydanticBaseModel, Field

class SearchAICatalogArgs(PydanticBaseModel):
    search_terms: str = Field(
        default="",
        description="Terms for the search. Leave blank to return all datasets."
    )
    limit: int = Field(
        default=20,
        description="The maximum number of datasets to return. "
        "Set to -1 to return all."
    )

The example below implements a simple BaseTool class for CrewAI, implemented in tool_deployment.py, containing all necessary metadata and available for reuse across multiple CrewAI tools.

tool_deployment.py
from abc import ABC
from crewai.tools import BaseTool
from datarobot_genai.core.chat.client import ToolClient

class BaseToolWithDeployment(BaseTool, ABC):
    model_config = {
        "arbitrary_types_allowed": True
    }
    """Adds support for arbitrary types in Pydantic models, needed for the ToolClient."""

    tool_client: ToolClient
    """The tool client initialized by the agent with access to the ToolClient authorization context."""

    deployment_id: str
    """The DataRobot deployment ID of the custom model executing tool logic."""

The SearchAICatalogTool, defined in tool_ai_catalog_search.py, invokes tool_deployment to build off the BaseToolWithDeployment module.

tool_ai_catalog_search.py
import json
from typing import Dict, List, Type
from pydantic import BaseModel as PydanticBaseModel
from tool_deployment import BaseToolWithDeployment

class SearchAICatalogTool(BaseToolWithDeployment):
    name: str = "Search Data Registry"
    description: str = (
        "This tool provides a list of all available dataset names and their associated IDs from the Data Registry. "
        "You should always check to see if the dataset you are looking for can be found here. "
        "For future queries, you should use the associated dataset ID instead of the name to avoid ambiguity."
    )
    args_schema: Type[PydanticBaseModel] = SearchAICatalogArgs
    def _run(self, **kwargs) -> List[Dict[str, str]]:
        # Validate and parse the input arguments using the defined schema.
        validated_args = self.args_schema(**kwargs)
        # Call the tool deployment with the generated payload.
        result = self.tool_client.call(
            deployment_id=self.deployment_id,
            payload=validated_args.model_dump()
        )
        # Format and return the results.
        return json.loads(result.data).get("datasets", [])

The example below uses the CrewAI framework to implement a tool through the BaseTool, Agent, and Task classes. The following methods show how you can add tool properties to your MyAgent class to initialize the Data Registry searching tool, define an LLM agent to search the Data Registry, and then define a task for the agent:

agent.py
from crewai.tools import BaseTool
from crewai import Agent, Task
from tool_ai_catalog_search import SearchAICatalogTool
from datarobot_genai.core.chat.client import ToolClient

class MyAgent:

    # More agentic workflow code.

    @property
    def tool_client(self) -> ToolClient:
        """ToolClient instance for calling deployed tools."""
        return ToolClient(
            api_key=self.api_key,
            base_url=self.api_base,
        )

    @property
    def search_ai_catalog_tool(self) -> BaseTool:
        """Search AI Catalog tool using deployed tool."""
        deployment_id = self.search_ai_catalog_deployment_id
        if not deployment_id:
            raise ValueError("Configure a deployment ID for the Search Data Registry tool.")
        return SearchAICatalogTool(
            tool_client=self.tool_client,
            deployment_id=deployment_id
        )

    @property
    def agent_ai_catalog_searcher(self) -> Agent:
        """Agent configured to search the AI Catalog."""
        return Agent(
            role="Expert Data Registry Searcher",
            goal="Search for and retrieve relevant files from Data Registry.",
            backstory="You are a meticulous analyst that is skilled at examining lists of files and "
            "determining the most appropriate file based on the context.",
            verbose=self.verbose,
            allow_delegation=False,
            llm=self.llm_with_datarobot_llm_gateway,
    )

    @property
    def task_ai_catalog_search(self) -> Task:
        """Task for searching the AI Catalog."""
        return Task(
            description=(
                "You should search for a relevant dataset id in the Data Registry "
                "based on the provided dataset topic: {dataset_topic}."
            ),
            expected_output=(
                "Search for a list of relevant files in the Data Registry and "
                "determine the most relevant dataset id that matches the given topic. "
                "You should return the entire dataset id."
            ),
            agent=self.agent_ai_catalog_searcher,
            tools=[self.search_ai_catalog_tool],
        )

    # More agentic workflow code.

Framework-specific documentation

You can also refer to the framework repositories and documentation for more information on constructing more advanced tools and agents:

Agentic tool considerations

When deploying the application and agent separately from the agentic tool (for example, deploying the application and agent via Pulumi after local development and the tool manually in DataRobot), all components must be deployed by the same user. Custom models use the creator's API key and require a shared scope to store and retrieve authentication data from the OAuth Providers Service.