【NVIDIA AIQ】自定义函数实践

NVIDIA Agent Intelligence Toolkit自定义函数实践

定义工具函数

import loggingfrom pydantic import BaseModel
from pydantic import Field
from pydantic import HttpUrlfrom 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 FunctionBaseConfiglogger = 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

from . import custom_tool

配置引用

# NVIDIA NeMo Agent Toolkit 黑客松配置
# 支持用户自定义OpenAI兼容APIgeneral:use_uvloop: truefunctions:custom_tool:_type: custom_tooldescription: "自定义工具测试"workflow:_type: react_agenttool_names:- custom_toolllm_name: default_llmverbose: trueparse_agent_response_max_retries: 3max_iterations: 10

测试验证

输入关键字，触发AI调用

终端调用打印

参考资源

NVIDIA Agent Intelligence Toolkit

Custom Functions