【NVIDIA AIQ】自定义函数实践

程序猿炎义2025-09-06 20:14

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调用

终端调用打印

参考资源

NVIDIA Agent Intelligence Toolkit

Custom Functions