Normalize Ollama demo tool arguments

This commit is contained in:
Thales Maciel 2026-03-08 17:16:28 -03:00
parent 75082467f9
commit f3e7d4aa3e
2 changed files with 106 additions and 12 deletions

View file

@ -95,6 +95,64 @@ def test_run_ollama_tool_demo_happy_path(monkeypatch: pytest.MonkeyPatch) -> Non
assert any(line == "[tool] result vm_run" for line in logs)
def test_run_ollama_tool_demo_accepts_legacy_profile_and_string_network(
monkeypatch: pytest.MonkeyPatch,
) -> None:
requests: list[dict[str, Any]] = []
def fake_post_chat_completion(base_url: str, payload: dict[str, Any]) -> dict[str, Any]:
assert base_url == "http://localhost:11434/v1"
requests.append(payload)
if len(requests) == 1:
return {
"choices": [
{
"message": {
"role": "assistant",
"content": "",
"tool_calls": [
{
"id": "1",
"function": {
"name": "vm_run",
"arguments": json.dumps(
{
"profile": "debian:12",
"command": "printf 'true\\n'",
"vcpu_count": 1,
"mem_mib": 512,
"network": "true",
}
),
},
}
],
}
}
]
}
return {
"choices": [
{
"message": {
"role": "assistant",
"content": "Executed git command in ephemeral VM.",
}
}
]
}
monkeypatch.setattr(ollama_demo, "_post_chat_completion", fake_post_chat_completion)
result = ollama_demo.run_ollama_tool_demo()
assert result["fallback_used"] is False
assert int(result["exec_result"]["exit_code"]) == 0
assert result["tool_events"][0]["success"] is True
assert result["tool_events"][0]["arguments"]["environment"] == "debian:12"
assert "profile" not in result["tool_events"][0]["arguments"]
def test_run_ollama_tool_demo_recovers_from_bad_vm_id(
monkeypatch: pytest.MonkeyPatch,
) -> None:
@ -467,14 +525,25 @@ def test_require_int_validation(value: Any) -> None:
ollama_demo._require_int({"k": value}, "k")
@pytest.mark.parametrize(("arguments", "expected"), [({}, False), ({"k": True}, True)])
@pytest.mark.parametrize(
("arguments", "expected"),
[
({}, False),
({"k": True}, True),
({"k": "true"}, True),
({"k": " false "}, False),
({"k": 1}, True),
({"k": 0}, False),
],
)
def test_require_bool(arguments: dict[str, Any], expected: bool) -> None:
assert ollama_demo._require_bool(arguments, "k", default=False) is expected
def test_require_bool_validation() -> None:
@pytest.mark.parametrize("value", ["", "maybe", 2])
def test_require_bool_validation(value: Any) -> None:
with pytest.raises(ValueError, match="must be a boolean"):
ollama_demo._require_bool({"k": "true"}, "k")
ollama_demo._require_bool({"k": value}, "k")
def test_post_chat_completion_success(monkeypatch: pytest.MonkeyPatch) -> None: