Skip to content

Add tools to agents

エージェントにローカルツールを追加するには、まずagent.pyファイルを修正するか、custom_modelディレクトリに新しいファイルを追加します。以下の例は、作業を開始するのに役立ちます。 Note that the structure and implementation of a tool is framework-specific. ここで紹介するのは、CrewAI、LangGraph、LlamaIndex、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"],
    ) 

NATのテンプレートでは、ツールはworkflow.yamlファイル内の関数として定義されます。 関数はfunctionsセクション内で_type: chat_completionとともに定義され、ワークフローのtool_listで参照されます。 NATでのツールの実装の詳細については、NVIDIA NeMo Agent Toolkitのドキュメントを参照してください。

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')."
            ) 

NATでのローカルツールの実装の詳細については、NVIDIA NeMo Agent Toolkitのドキュメントを参照してください。

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)}" 

NATでの外部APIツールの実装の詳細については、NVIDIA NeMo Agent Toolkitのドキュメントを参照してください。

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.
  • model-metadata.yamlで定義されているデプロイ済みツールのdeployment_id
  • ツールモジュールで定義されているツールのメタデータ。たとえば、tool_ai_catalog_search.py

このドキュメントの例では、エージェントテンプレートにツール統合機能を追加する方法を紹介します。 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.

レジストリからデプロイされたエージェントツールを使用してエージェントワークフローを構築する場合、ワークフローの例として以下のファイルがあります。

ファイル 内容
__init__.py Pythonパッケージ初期化ファイル。ディレクトリをPythonパッケージにします。
custom.py LLMを呼び出すためのボルトオンのガバナンスAPIchatフック)を実装し、それらのパラメーターを(agent.pyで定義されている)エージェントに渡すカスタムモデルコード
agent.py エージェントコード。必須のinvokeメソッドを使用して、MyAgentクラスにエージェントワークフローを実装します。 このクラスにツール統合プロパティを追加して、デプロイされたツールとやり取りできます。
config.py 環境変数、ランタイムパラメーター、およびDataRobotの資格情報から設定を読み込むためのコード。
mcp_client.py ツール連携のためのMCPサーバー接続管理を提供するコード(オプション。MCPツールを使用する場合にのみ必要です)。
tool_deployment.py BaseToolクラスコード。ツールの実装に必要なすべてのメタデータが含まれています。
tool.py デプロイされたツールとのインターフェイスのためのコードで、入力引数やスキーマを定義します。 多くの場合、このファイルはtool.pyという名前にはなりません。これは、複数のツールを実装する場合があるためです。 この例では、この機能はtool_ai_catalog_search.pyで定義されています。
model-metadata.yaml エージェントのワークフローに必要なカスタムモデルのメタデータとランタイムパラメーター
pyproject.toml 最新のPythonパッケージング標準を使用した、エージェントのワークフローに必要なライブラリ(とバージョン)

ToolClientクラスのインスタンスを実装する

Every agent template and framework requires the ToolClient class from the datarobot-genai package to offload tool call processing to deployed global tools. ツールクライアントは、デプロイされたツールを呼び出し、その結果をエージェントに返します。 ToolClientモジュールをagent.pyファイルにインポートするには、以下のimport文を使用します。

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

The ToolClient is available in the datarobot-genai package. デプロイされたツールのAPIエンドポイントとデプロイIDを定義し、(必要な場合)認可コンテキストを取得して、scorescore_unstructuredchatフックのインターフェイスを提供します。

ToolClientagent.pyにインポートしたら、MyAgentクラスにtool_clientプロパティを追加できます。 ToolClientは、認証に環境変数DATAROBOT_API_TOKENおよびDATAROBOT_ENDPOINTを自動的に使用しますが、必要に応じてこれらを明示的に渡すこともできます。

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. 

(オプション)外部ツールの認可コンテキストを初期化する

認可コンテキストは、下流のエージェントやツールが外部サービスに接続する際にアクセストークンを取得できるようにするために必要です。 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. この関数は、DataRobot Python APIクライアントのユーティリティメソッドを使用します。 * set_authorization_context:現在のプロセスに認可コンテキストを設定するメソッド。 * get_authorization_context:現在のプロセスに認可コンテキストを取得するメソッド。

利用可能なOAuthユーティリティメソッド

これらのユーティリティメソッドは、バージョン3.8.0からDataRobot Python APIクライアントで利用できるようになりました。

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. ToolClient.call()メソッドは、デプロイされたツールを呼び出すときに、認可コンテキストを自動的に含めます。

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.

外部サービスに依存するカスタムツールを実装している場合、@datarobot_tool_authデコレーターを使用すると、認可コンテキストの取得、関連データの抽出、およびDataRobotで設定されたOAuthプロバイダーからOAuthアクセストークンを取得するためのDataRobot APIへの接続のプロセスを効率化できます。 OAuthプロバイダーが1つしか設定されていない場合、デコレーターはproviderパラメーターを必要としません。これは、利用可能なプロバイダーが1つだけであるためです。ただし、複数のプロバイダーが利用可能(または利用可能になる)場合は、このパラメーターを定義する必要があります。

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. 

ツールのデプロイとのやり取り

グローバルツールとレジストリワークショップでカスタム構築されたツールは、エージェントがそれらを呼び出すためにデプロイする必要があります。 これらのツールがデプロイされている場合、それらと通信するにはデプロイIDが必要です。これは、DataRobot APIを介したツールとのやり取りに使用されます。 デプロイされたツールのデプロイIDをエージェントに提供する主な方法は、エージェントのメタデータでランタイムパラメーターとして定義されている環境変数を使用することです。 このメタデータを提供するには、model-metadata.yamlファイルを作成または修正して、エージェントが通信する必要がある各デプロイツールのランタイムパラメーターを追加します。 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. 以下の例のフィールド名は、これらのツールを実装するDataRobotエージェントテンプレートで使用されます。ただし、fieldNameは変更可能であり、agent.pyファイルにあるエージェントのコードの実装と一致する必要があります。

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 

必要な値をdefaultValueとして指定していない場合は、エージェントのコードにデプロイIDを提供するために、ランタイムパラメーターを設定する必要があります。 これには2つの方法があります。

  • 手動:エージェントをデプロイする前に、UIのレジストリワークショップで値を設定します。
  • 自動:Pulumiスクリプトで値を設定し、カスタムモデルに渡します。

パラメーターが設定されると、RuntimeParametersクラスがdatarobot_drumからインポートされている限り、エージェントコードからアクセスできます。 以下の例では、ツールのプロパティをMyAgentクラスに追加して、デプロイされたツールとやり取りする方法を示します。

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. 

ツールのメタデータを定義する

エージェントのツールを構築する場合、メタデータはエージェントLLMがツールを呼び出す方法を定義します。 メタデータが提供する詳細が多いほど、LLMはツールをより効果的に使用します。 メタデータには、ツールの説明、各引数のスキーマ、および関連する説明が含まれます。 各フレームワークには、このメタデータを定義する独自の方法があります。ただし、ほとんどの場合、pydanticを利用してBaseModelをインポートし、ツールの引数を定義できます。

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."
    ) 

以下の例では、CrewAI向けに簡単なBaseToolクラスを実装しています。これは、tool_deployment.pyに実装され、必要なすべてのメタデータを含み、複数のCrewAIツールで再利用できます。

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.""" 

tool_ai_catalog_search.pyで定義されているSearchAICatalogToolは、BaseToolWithDeploymentモジュールを構築するためにtool_deploymentを呼び出します。

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", []) 

以下の例では、CrewAIフレームワークを使用し、BaseToolAgentTaskクラスを通じてツールを実装しています。 以下のメソッドでは、MyAgentクラスにツールのプロパティを追加して、データレジストリの検索ツールを初期化する、データレジストリを検索するLLMエージェントを定義する、エージェントのタスクを定義する、という順序で実行する方法を示します。

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. 

フレームワーク固有のドキュメント

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

エージェントツールに関する注意事項

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.