b77e151387
Phase 20 - 提升 Tool Calling 精準度 50% → 83.3% 新增: - src/models/nvidia.py: Pydantic Schema - src/services/nvidia_provider.py: NvidiaProvider 類別 - tests/test_nvidia_provider.py: 15 項單元測試 (全部通過) 修改: - ai_router.py: AIProvider.NVIDIA + route_tool_calling() - ai_rate_limiter.py: NVIDIA 限制 (5 RPM, 100/day) - models.json: NVIDIA 配置 - cd.yaml: Secrets 注入 NVIDIA_API_KEY 路由策略: - Tool Calling: Nemotron → Gemini → Claude - 一般對話: Ollama → Gemini → Claude (不變) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
317 lines
10 KiB
Python
317 lines
10 KiB
Python
"""
|
|
NVIDIA Provider Tests - ADR-036
|
|
===============================
|
|
測試 Nemotron Tool Calling 整合
|
|
|
|
注意: 這些是單元測試，不需要真實的 NVIDIA API Key
|
|
"""
|
|
|
|
import json
|
|
|
|
import pytest
|
|
|
|
from src.models.nvidia import (
|
|
NvidiaChoice,
|
|
NvidiaMessage,
|
|
NvidiaProviderResult,
|
|
NvidiaResponse,
|
|
NvidiaUsage,
|
|
ToolCall,
|
|
ToolCallValidationResult,
|
|
ToolDefinition,
|
|
ToolFunction,
|
|
)
|
|
from src.services.nvidia_provider import (
|
|
HIGH_RISK_TOOLS,
|
|
NvidiaProvider,
|
|
create_tool_definition,
|
|
get_nvidia_provider,
|
|
reset_nvidia_provider,
|
|
)
|
|
|
|
|
|
class TestNvidiaModels:
|
|
"""測試 NVIDIA Pydantic Models"""
|
|
|
|
def test_tool_function_model(self):
|
|
"""測試 ToolFunction 模型"""
|
|
func = ToolFunction(
|
|
name="restart_pod",
|
|
arguments='{"pod_name": "api-server", "namespace": "default"}',
|
|
)
|
|
assert func.name == "restart_pod"
|
|
assert '"pod_name"' in func.arguments
|
|
|
|
def test_tool_call_model(self):
|
|
"""測試 ToolCall 模型"""
|
|
tc = ToolCall(
|
|
id="call_123",
|
|
type="function",
|
|
function=ToolFunction(
|
|
name="scale_deployment",
|
|
arguments='{"replicas": 3}',
|
|
),
|
|
)
|
|
assert tc.id == "call_123"
|
|
assert tc.function.name == "scale_deployment"
|
|
|
|
def test_nvidia_response_model(self):
|
|
"""測試 NvidiaResponse 模型"""
|
|
response = NvidiaResponse(
|
|
id="resp_123",
|
|
created=1234567890,
|
|
model="nvidia/llama-3.1-nemotron-70b-instruct",
|
|
choices=[
|
|
NvidiaChoice(
|
|
index=0,
|
|
message=NvidiaMessage(
|
|
role="assistant",
|
|
content=None,
|
|
tool_calls=[
|
|
ToolCall(
|
|
id="call_1",
|
|
function=ToolFunction(
|
|
name="restart_pod",
|
|
arguments='{"pod": "api"}',
|
|
),
|
|
)
|
|
],
|
|
),
|
|
finish_reason="tool_calls",
|
|
)
|
|
],
|
|
usage=NvidiaUsage(
|
|
prompt_tokens=100,
|
|
completion_tokens=50,
|
|
total_tokens=150,
|
|
),
|
|
)
|
|
|
|
assert response.id == "resp_123"
|
|
assert len(response.choices) == 1
|
|
assert response.choices[0].message.tool_calls is not None
|
|
assert len(response.choices[0].message.tool_calls) == 1
|
|
assert response.usage.total_tokens == 150
|
|
|
|
def test_tool_call_validation_result(self):
|
|
"""測試 ToolCallValidationResult 模型"""
|
|
result = ToolCallValidationResult(
|
|
valid=True,
|
|
tool_name="restart_pod",
|
|
arguments={"pod_name": "api", "namespace": "default"},
|
|
)
|
|
assert result.valid
|
|
assert result.tool_name == "restart_pod"
|
|
assert result.arguments["namespace"] == "default"
|
|
|
|
def test_tool_definition(self):
|
|
"""測試 ToolDefinition 模型"""
|
|
definition = ToolDefinition(
|
|
type="function",
|
|
function={
|
|
"name": "restart_pod",
|
|
"description": "Restart a Kubernetes pod",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"pod_name": {"type": "string"},
|
|
"namespace": {"type": "string"},
|
|
},
|
|
"required": ["pod_name"],
|
|
},
|
|
},
|
|
)
|
|
assert definition.type == "function"
|
|
assert definition.function["name"] == "restart_pod"
|
|
|
|
|
|
class TestNvidiaProvider:
|
|
"""測試 NvidiaProvider 類別"""
|
|
|
|
def test_singleton(self):
|
|
"""測試單例模式"""
|
|
reset_nvidia_provider()
|
|
p1 = get_nvidia_provider()
|
|
p2 = get_nvidia_provider()
|
|
assert p1 is p2
|
|
reset_nvidia_provider()
|
|
|
|
def test_high_risk_tool_detection(self):
|
|
"""測試高風險 Tool 檢測"""
|
|
provider = NvidiaProvider()
|
|
|
|
# 高風險操作
|
|
assert provider.is_high_risk_tool("delete_pod")
|
|
assert provider.is_high_risk_tool("DELETE_POD") # 大寫也應該匹配
|
|
assert provider.is_high_risk_tool("delete_deployment")
|
|
assert provider.is_high_risk_tool("scale_to_zero")
|
|
assert provider.is_high_risk_tool("drain_node")
|
|
|
|
# 非高風險操作
|
|
assert not provider.is_high_risk_tool("restart_pod")
|
|
assert not provider.is_high_risk_tool("scale_deployment")
|
|
assert not provider.is_high_risk_tool("get_logs")
|
|
|
|
def test_filter_high_risk_tools(self):
|
|
"""測試過濾高風險 Tool Calls"""
|
|
provider = NvidiaProvider()
|
|
|
|
tool_calls = [
|
|
ToolCallValidationResult(
|
|
valid=True,
|
|
tool_name="restart_pod",
|
|
arguments={"pod": "api"},
|
|
),
|
|
ToolCallValidationResult(
|
|
valid=True,
|
|
tool_name="delete_pod",
|
|
arguments={"pod": "test"},
|
|
),
|
|
ToolCallValidationResult(
|
|
valid=False,
|
|
tool_name="invalid_tool",
|
|
error="Parse error",
|
|
),
|
|
]
|
|
|
|
high_risk = provider.get_high_risk_tools(tool_calls)
|
|
|
|
assert len(high_risk) == 1
|
|
assert high_risk[0].tool_name == "delete_pod"
|
|
|
|
def test_validate_tool_calls(self):
|
|
"""測試 Tool Call 驗證"""
|
|
provider = NvidiaProvider()
|
|
|
|
# 建立模擬回應
|
|
response = NvidiaResponse(
|
|
id="resp_123",
|
|
created=1234567890,
|
|
model="nvidia/llama-3.1-nemotron-70b-instruct",
|
|
choices=[
|
|
NvidiaChoice(
|
|
index=0,
|
|
message=NvidiaMessage(
|
|
role="assistant",
|
|
tool_calls=[
|
|
ToolCall(
|
|
id="call_1",
|
|
function=ToolFunction(
|
|
name="restart_pod",
|
|
arguments='{"pod_name": "api", "namespace": "default"}',
|
|
),
|
|
),
|
|
ToolCall(
|
|
id="call_2",
|
|
function=ToolFunction(
|
|
name="invalid_tool",
|
|
arguments="not valid json{", # 無效 JSON
|
|
),
|
|
),
|
|
],
|
|
),
|
|
)
|
|
],
|
|
)
|
|
|
|
results = provider._validate_tool_calls(response)
|
|
|
|
assert len(results) == 2
|
|
assert results[0].valid
|
|
assert results[0].tool_name == "restart_pod"
|
|
assert results[0].arguments["pod_name"] == "api"
|
|
assert not results[1].valid
|
|
assert "JSON" in results[1].error
|
|
|
|
|
|
class TestCreateToolDefinition:
|
|
"""測試 Tool 定義建立函數"""
|
|
|
|
def test_create_tool_definition(self):
|
|
"""測試建立 Tool 定義"""
|
|
definition = create_tool_definition(
|
|
name="scale_deployment",
|
|
description="Scale a Kubernetes deployment",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {
|
|
"deployment": {"type": "string"},
|
|
"replicas": {"type": "integer"},
|
|
},
|
|
"required": ["deployment", "replicas"],
|
|
},
|
|
)
|
|
|
|
assert definition.type == "function"
|
|
assert definition.function["name"] == "scale_deployment"
|
|
assert definition.function["description"] == "Scale a Kubernetes deployment"
|
|
assert "replicas" in definition.function["parameters"]["properties"]
|
|
|
|
|
|
class TestHighRiskTools:
|
|
"""測試高風險 Tool 清單"""
|
|
|
|
def test_high_risk_tools_list(self):
|
|
"""確認高風險 Tool 清單包含所有必要操作"""
|
|
assert "delete_pod" in HIGH_RISK_TOOLS
|
|
assert "delete_deployment" in HIGH_RISK_TOOLS
|
|
assert "delete_namespace" in HIGH_RISK_TOOLS
|
|
assert "scale_to_zero" in HIGH_RISK_TOOLS
|
|
assert "drain_node" in HIGH_RISK_TOOLS
|
|
assert "cordon_node" in HIGH_RISK_TOOLS
|
|
|
|
def test_restart_not_high_risk(self):
|
|
"""確認 restart 不在高風險清單"""
|
|
assert "restart_pod" not in HIGH_RISK_TOOLS
|
|
assert "restart_deployment" not in HIGH_RISK_TOOLS
|
|
|
|
|
|
class TestAIRouterNvidiaIntegration:
|
|
"""測試 AIRouter NVIDIA 整合"""
|
|
|
|
def test_nvidia_provider_in_router(self):
|
|
"""測試 AIProvider 包含 NVIDIA"""
|
|
from src.services.ai_router import AIProvider
|
|
|
|
assert hasattr(AIProvider, "NVIDIA")
|
|
assert AIProvider.NVIDIA.value == "nvidia"
|
|
|
|
def test_tool_calling_route(self):
|
|
"""測試 Tool Calling 路由"""
|
|
from src.services.ai_router import get_ai_router, AIProvider, reset_ai_router
|
|
|
|
reset_ai_router()
|
|
router = get_ai_router()
|
|
|
|
provider, model, fallback_chain = router.route_tool_calling()
|
|
|
|
assert provider == AIProvider.NVIDIA
|
|
assert "nvidia" in model.lower() or "nemotron" in model.lower()
|
|
# Fallback 應該包含 Gemini 和 Claude
|
|
fallback_providers = [p for p, _ in fallback_chain]
|
|
assert AIProvider.GEMINI in fallback_providers
|
|
assert AIProvider.CLAUDE in fallback_providers
|
|
|
|
reset_ai_router()
|
|
|
|
def test_existing_routing_not_affected(self):
|
|
"""測試現有路由規則不受影響"""
|
|
from src.services.ai_router import get_ai_router, AIProvider, reset_ai_router
|
|
|
|
reset_ai_router()
|
|
router = get_ai_router()
|
|
|
|
# 測試同步路由 (不涉及 NVIDIA)
|
|
decision = router.route_sync("重啟 api pod")
|
|
|
|
# 應該還是使用 Ollama (低複雜度)
|
|
assert decision.selected_provider in [
|
|
AIProvider.OLLAMA,
|
|
AIProvider.GEMINI,
|
|
AIProvider.CLAUDE,
|
|
]
|
|
# NVIDIA 不應該出現在一般路由中
|
|
assert decision.selected_provider != AIProvider.NVIDIA
|
|
|
|
reset_ai_router()
|