File size: 4,652 Bytes
087e4a0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | from __future__ import annotations
import json
import importlib
import tempfile
import unittest
from pathlib import Path
class _FakeModel:
def __init__(self, outputs: list[str]) -> None:
self.outputs = list(outputs)
self.states: list[dict[str, object]] = []
def generate(self, prompt: str, *, state: dict[str, object]) -> str:
self.states.append(dict(state))
if not self.outputs:
raise AssertionError("unexpected model call")
return self.outputs.pop(0)
def _chunk() -> dict[str, object]:
question = (
"Mercury completes one revolution in 88 Earth days. "
"How many Earth years is that? Return a decimal number."
)
return {
"task_id": "smoke/mercury",
"hcl_interface": {
"evidence": {
"summary": "Mercury takes 88 Earth days per revolution.",
"items": [{"modality": "text", "content": question, "source": "question"}],
},
"recognizable_goal": {"content": "Convert 88 days to Earth years."},
"recognizable_constraints": [],
},
"task_context": {"files": [], "capability_inputs": {}},
}
class RouterManagedArithmeticArgumentsTest(unittest.TestCase):
@classmethod
def setUpClass(cls) -> None:
global Router, RouterConfig, CapabilityExecutionContext
router_module = importlib.import_module("hcl.router")
capability_module = importlib.import_module("hcl.capability")
Router = router_module.Router
RouterConfig = router_module.RouterConfig
CapabilityExecutionContext = capability_module.CapabilityExecutionContext
def test_builds_executes_and_caches_safe_arithmetic_arguments(self) -> None:
model = _FakeModel([
json.dumps({
"tool_name": "arithmetic_calculator",
"arguments": {"expression": "88 / 365.25", "precision": 6},
})
])
with tempfile.TemporaryDirectory() as temp_dir:
router = Router(
model,
RouterConfig(
capability_enabled=True,
capability_execution_mode="router_managed",
tool_argument_cache_path=str(Path(temp_dir) / "arguments.jsonl"),
),
)
workflow = {
"workflow_name": "conversion",
"steps": ["Convert 88 days to years."],
"use_memory": False,
"use_tools": True,
"memory_query": "",
"tool_needs": ["arithmetic_calculator"],
"rationale": "Compute 88 / 365.25.",
}
selected = [{"tool_name": "arithmetic_calculator"}]
context = CapabilityExecutionContext(task_id="smoke/mercury")
first = router.execute_selected_capabilities(
_chunk(), workflow, selected, execution_context=context
)
second = router.execute_selected_capabilities(
_chunk(), workflow, selected, execution_context=context
)
self.assertEqual(first[0]["status"], "success")
self.assertEqual(first[0]["arguments"]["expression"], "88 / 365.25")
self.assertEqual(first[0]["output"]["result"], "0.240931")
self.assertEqual(second, first)
self.assertEqual(len(model.states), 1)
self.assertEqual(model.states[0]["router_phase"], "build_tool_arguments")
def test_rejects_unsafe_expression(self) -> None:
model = _FakeModel([
json.dumps({
"tool_name": "arithmetic_calculator",
"arguments": {"expression": "__import__('os').system('id')", "precision": 2},
})
])
router = Router(
model,
RouterConfig(
capability_enabled=True,
capability_execution_mode="router_managed",
tool_argument_invalid_output_retries=0,
),
)
result = router.execute_selected_capabilities(
_chunk(),
{
"workflow_name": "unsafe",
"steps": [],
"use_memory": False,
"use_tools": True,
"memory_query": "",
"tool_needs": ["arithmetic_calculator"],
"rationale": "",
},
[{"tool_name": "arithmetic_calculator"}],
execution_context=CapabilityExecutionContext(task_id="smoke/unsafe"),
)
self.assertEqual(result[0]["status"], "skipped")
if __name__ == "__main__":
unittest.main()
|