NVIDIA Agent Intelligence Toolkit自定义函数实践
定义工具函数
python
import logging
from pydantic import BaseModel
from pydantic import Field
from pydantic import HttpUrl
from aiq.builder.builder import Builder
from aiq.builder.function_info import FunctionInfo
from aiq.cli.register_workflow import register_function
from aiq.data_models.function import FunctionBaseConfig
logger = logging.getLogger(__name__)
# 工具类
class CustomToolConfig(FunctionBaseConfig, name="custom_tool"):
""" 测试自定义工具调用 """
# 注册工具函数
@register_function(config_type=CustomToolConfig)
async def custom_tool(config: CustomToolConfig, builder: Builder):
# 工具函数逻辑处理
async def _custom_tool_test(text: str) -> str:
return '这是一个自定义工具调用测试${text}'
# 逻辑导出调用
yield FunctionInfo.from_fn(
_custom_tool_test,
description="这是一个自定义工具调用测试",
)
注册: 工具导出
- aiq/tool/register.py
python
from . import custom_tool
配置引用
yaml
# NVIDIA NeMo Agent Toolkit 黑客松配置
# 支持用户自定义OpenAI兼容API
general:
use_uvloop: true
functions:
custom_tool:
_type: custom_tool
description: "自定义工具测试"
workflow:
_type: react_agent
tool_names:
- custom_tool
llm_name: default_llm
verbose: true
parse_agent_response_max_retries: 3
max_iterations: 10
测试验证
输入关键字,触发AI调用
终端调用打印