From fc6ece0193f4627debb2c77adfbf07953a944f75 Mon Sep 17 00:00:00 2001 From: justin212407 Date: Mon, 17 Aug 2026 14:52:03 +0530 Subject: [PATCH 1/3] feat:nao test: assert the agent performed a specific action Signed-off-by: justin212407 --- apps/backend/src/agents/tools/index.ts | 17 +- apps/backend/src/routes/test.ts | 3 +- cli/README.md | 14 +- cli/nao_core/commands/test/assertions.py | 149 ++++++++++++++++++ cli/nao_core/commands/test/case.py | 8 +- cli/nao_core/commands/test/client.py | 3 +- cli/nao_core/commands/test/runner.py | 26 ++- .../nao_core/commands/test_assertions.py | 130 +++++++++++++++ cli/tests/nao_core/commands/test_case.py | 38 ++++- cli/tests/nao_core/commands/test_runner.py | 124 +++++++++++++++ .../create-context-tests/templates/test.yaml | 9 ++ 11 files changed, 505 insertions(+), 16 deletions(-) create mode 100644 cli/nao_core/commands/test/assertions.py create mode 100644 cli/tests/nao_core/commands/test_assertions.py diff --git a/apps/backend/src/agents/tools/index.ts b/apps/backend/src/agents/tools/index.ts index 9953d26fb..e4f60cf2c 100644 --- a/apps/backend/src/agents/tools/index.ts +++ b/apps/backend/src/agents/tools/index.ts @@ -55,6 +55,11 @@ export const getTools = ( agentSettings: AgentSettings | null, extraTools?: Record, options: { + /** + * @deprecated No longer strips tools. Kept so existing callers that pass + * `testMode` continue to typecheck. Clarification stays available so + * `nao test` can assert intermediate actions (e.g. follow-up questions). + */ testMode?: boolean; mcpEnabled?: boolean; mcpServers?: string[] | null; @@ -89,14 +94,9 @@ export const getTools = ( } : {}; - const { - execute_python, - execute_sandboxed_code, - clarification: clarificationTool, - suggest_follow_ups, - write: writeTool, - ...rest - } = tools; + const { execute_python, execute_sandboxed_code, suggest_follow_ups, write: writeTool, ...rest } = tools; + // Keep clarification available in testMode so `nao test` can assert follow-up + // questions and other intermediate actions on the recorded tool-call trace. const baseTools = { ...rest, ...(isStorageEnabled() && { write: writeTool }), @@ -105,7 +105,6 @@ export const getTools = ( const allTools = { ...baseTools, - ...(!options.testMode && { clarification: clarificationTool }), ...mcpTools, ...(agentSettings?.experimental?.pythonSandboxing && execute_python && { execute_python }), ...(agentSettings?.experimental?.sandboxes && execute_sandboxed_code && { execute_sandboxed_code }), diff --git a/apps/backend/src/routes/test.ts b/apps/backend/src/routes/test.ts index 9f9685e59..e2456cd43 100644 --- a/apps/backend/src/routes/test.ts +++ b/apps/backend/src/routes/test.ts @@ -39,7 +39,8 @@ export const testRoutes = async (app: App) => { body: z.object({ prompt: z.string(), model: llmSelectedModelSchema, - sql: z.string(), + // Optional: assertion-only tests omit reference SQL + sql: z.string().optional().default(''), meta: z .object({ costs: customModelCostSchema, diff --git a/cli/README.md b/cli/README.md index 87d7af767..056adc80e 100644 --- a/cli/README.md +++ b/cli/README.md @@ -212,7 +212,19 @@ databases: nao test ``` -Runs test cases defined as YAML files in `tests/`. Each test has a `name`, `prompt`, and expected `sql`. Results are saved to `tests/outputs/`. +Runs test cases defined as YAML files in `tests/`. Each test has a `name`, `prompt`, and optional expected `sql` and/or `assertions`. Results are saved to `tests/outputs/`. + +Final-output checks use reference `sql` (dataframe equality). Intermediate agent actions use `assertions` against the run's tool-call trace — for example, requiring a clarifying follow-up: + +```yaml +name: ambiguous_revenue_period +prompt: What was the revenue? +assertions: + - type: tool_call + tool: clarification +``` + +`tool_call` assertions can also require a specific tool (e.g. `execute_sql`), optional arg subset match via `args`, and `min_count`. SQL verification and assertions can be combined; the run passes only if every check passes. Options: diff --git a/cli/nao_core/commands/test/assertions.py b/cli/nao_core/commands/test/assertions.py new file mode 100644 index 000000000..a67b067f4 --- /dev/null +++ b/cli/nao_core/commands/test/assertions.py @@ -0,0 +1,149 @@ +"""Extensible action/step assertions for `nao test`. + +Assertions check intermediate agent behavior (tool calls, steps) independently +of final-output dataframe verification. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, cast + + +class AssertionConfigError(ValueError): + """Raised when an assertion definition in a test YAML is invalid.""" + + +@dataclass(frozen=True) +class ToolCallAssertion: + """Require that a named tool was invoked during the agentic loop. + + Optional ``args`` values must appear as a subset of the tool call's args + (nested dicts are matched recursively; lists require equality). + """ + + tool: str + args: dict[str, Any] | None = None + min_count: int = 1 + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> ToolCallAssertion: + tool = data.get("tool") + if not isinstance(tool, str) or not tool.strip(): + raise AssertionConfigError("tool_call assertion requires a non-empty 'tool' string") + + args = data.get("args") + if args is not None and not isinstance(args, dict): + raise AssertionConfigError("tool_call assertion 'args' must be a mapping when provided") + + min_count = data.get("min_count", 1) + if not isinstance(min_count, int) or isinstance(min_count, bool) or min_count < 1: + raise AssertionConfigError("tool_call assertion 'min_count' must be an integer >= 1") + + unknown = set(data) - {"type", "tool", "args", "min_count"} + if unknown: + raise AssertionConfigError(f"unknown tool_call assertion fields: {sorted(unknown)}") + + return cls(tool=tool.strip(), args=args, min_count=min_count) + + +Assertion = ToolCallAssertion + + +def parse_assertions(raw: Any) -> list[Assertion]: + """Parse the optional ``assertions`` list from a test YAML document.""" + if raw is None: + return [] + if not isinstance(raw, list): + raise AssertionConfigError("'assertions' must be a list") + + assertions: list[Assertion] = [] + for index, item in enumerate(raw): + if not isinstance(item, dict): + raise AssertionConfigError(f"assertions[{index}] must be a mapping") + entry = cast(dict[str, Any], item) + assertion_type = entry.get("type") + if assertion_type == "tool_call": + assertions.append(ToolCallAssertion.from_dict(entry)) + elif assertion_type is None: + raise AssertionConfigError(f"assertions[{index}] is missing 'type'") + else: + raise AssertionConfigError(f"assertions[{index}] has unknown type {assertion_type!r}; supported: tool_call") + return assertions + + +def _args_match(expected: Any, actual: Any) -> bool: + """Return True when ``expected`` is a subset of ``actual`` (dicts recursive).""" + if isinstance(expected, dict): + if not isinstance(actual, dict): + return False + return all(key in actual and _args_match(value, actual[key]) for key, value in expected.items()) + return expected == actual + + +def evaluate_tool_call_assertion( + assertion: ToolCallAssertion, + tool_calls: list[dict[str, Any]] | None, +) -> tuple[bool, str]: + """Evaluate a single tool_call assertion against recorded tool calls.""" + calls = tool_calls or [] + matches = [ + call + for call in calls + if call.get("toolName") == assertion.tool + and (assertion.args is None or _args_match(assertion.args, call.get("args") or {})) + ] + count = len(matches) + if count >= assertion.min_count: + if assertion.min_count == 1 and assertion.args is None: + return True, f"tool_call: {assertion.tool}" + detail = f"tool_call: {assertion.tool} (x{count}" + if assertion.min_count > 1: + detail += f", min {assertion.min_count}" + if assertion.args is not None: + detail += ", args matched" + detail += ")" + return True, detail + + if assertion.args is not None: + same_tool = sum(1 for call in calls if call.get("toolName") == assertion.tool) + if same_tool: + return ( + False, + f"missing tool_call: {assertion.tool} with args {assertion.args} " + f"(found {same_tool} call(s) without matching args)", + ) + return False, f"missing tool_call: {assertion.tool} with args {assertion.args}" + + if assertion.min_count > 1: + return ( + False, + f"missing tool_call: {assertion.tool} (found {count}, need >= {assertion.min_count})", + ) + return False, f"missing tool_call: {assertion.tool}" + + +def evaluate_assertions( + assertions: list[Assertion], + tool_calls: list[dict[str, Any]] | None, +) -> tuple[bool, str]: + """Evaluate all assertions. Returns (passed, combined message).""" + if not assertions: + return True, "" + + messages: list[str] = [] + all_passed = True + for assertion in assertions: + if isinstance(assertion, ToolCallAssertion): + passed, message = evaluate_tool_call_assertion(assertion, tool_calls) + else: # pragma: no cover - exhaustive for current Assertion union + passed, message = False, f"unsupported assertion: {assertion!r}" + messages.append(message) + all_passed = all_passed and passed + + return all_passed, "; ".join(messages) + + +def combine_check_messages(*parts: str) -> str: + """Join non-empty check messages with '; '.""" + return "; ".join(part for part in parts if part) diff --git a/cli/nao_core/commands/test/case.py b/cli/nao_core/commands/test/case.py index 5d30c407d..8e9295820 100644 --- a/cli/nao_core/commands/test/case.py +++ b/cli/nao_core/commands/test/case.py @@ -1,10 +1,12 @@ -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path import yaml from nao_core.ui import UI +from .assertions import Assertion, parse_assertions + TESTS_FOLDER = "tests/" @@ -15,7 +17,8 @@ class TestCase: name: str prompt: str file_path: Path - sql: str + sql: str | None = None + assertions: list[Assertion] = field(default_factory=list) @classmethod def from_yaml(cls, file_path: Path) -> "TestCase": @@ -28,6 +31,7 @@ def from_yaml(cls, file_path: Path) -> "TestCase": prompt=data["prompt"], sql=data.get("sql"), file_path=file_path, + assertions=parse_assertions(data.get("assertions")), ) diff --git a/cli/nao_core/commands/test/client.py b/cli/nao_core/commands/test/client.py index c0b82b83a..76bf8f395 100644 --- a/cli/nao_core/commands/test/client.py +++ b/cli/nao_core/commands/test/client.py @@ -128,7 +128,8 @@ def run_test( "modelId": model_id, }, "prompt": test_case.prompt, - "sql": test_case.sql, + # sql is optional: assertion-only tests may omit reference SQL + "sql": test_case.sql or "", } cost_payload = serialize_model_costs(costs) diff --git a/cli/nao_core/commands/test/runner.py b/cli/nao_core/commands/test/runner.py index d0a35da45..f730bd280 100644 --- a/cli/nao_core/commands/test/runner.py +++ b/cli/nao_core/commands/test/runner.py @@ -18,6 +18,7 @@ from nao_core.config.test import ComparisonConfig, TestConfig from nao_core.ui import UI +from .assertions import combine_check_messages, evaluate_assertions from .case import TESTS_FOLDER, TestCase, discover_tests from .client import BACKEND_URL, AgentClientError, VerificationResult, get_client from .compare import normalize_dataframe_numbers @@ -214,16 +215,20 @@ def run_test( UI.print(f"[dim] Cost: ${result.cost.totalCost}[/dim]") UI.print(f"[dim] Time: {result.duration_ms}ms[/dim]") + assertions_passed, assertions_msg = evaluate_assertions(test_case.assertions, result.tool_calls) + if result.verification: tolerances = comparison or ComparisonConfig() if result.verification.sql: UI.print(f"[dim] Verification SQL: {result.verification.sql}[/dim]") - passed, msg, diff = check_dataframe( + data_passed, data_msg, diff = check_dataframe( result.verification, rtol=tolerances.rtol, atol=tolerances.atol, decimals=tolerances.decimals, ) + passed = assertions_passed and data_passed + msg = combine_check_messages(assertions_msg, data_msg) status = "[green]✓[/green]" if passed else "[red]✗[/red]" UI.print(f" {status} {msg}") return TestRunResult( @@ -246,6 +251,25 @@ def run_test( ), ) + if test_case.assertions: + status = "[green]✓[/green]" if assertions_passed else "[red]✗[/red]" + UI.print(f" {status} {assertions_msg}") + return TestRunResult( + name=test_case.name, + model=str(model), + passed=assertions_passed, + message=assertions_msg, + tokens=result.usage.totalTokens, + cost=result.cost.totalCost, + duration_ms=result.duration_ms, + tool_call_count=tool_call_count, + details=TestRunDetails( + response_text=result.text, + tool_calls=result.tool_calls, + reference_sql=test_case.sql, + ), + ) + UI.print("[yellow] ⚠ no verification data[/yellow]") return TestRunResult( name=test_case.name, diff --git a/cli/tests/nao_core/commands/test_assertions.py b/cli/tests/nao_core/commands/test_assertions.py new file mode 100644 index 000000000..f9c383b89 --- /dev/null +++ b/cli/tests/nao_core/commands/test_assertions.py @@ -0,0 +1,130 @@ +import pytest + +from nao_core.commands.test.assertions import ( + AssertionConfigError, + ToolCallAssertion, + combine_check_messages, + evaluate_assertions, + evaluate_tool_call_assertion, + parse_assertions, +) + + +def test_parse_assertions_empty(): + assert parse_assertions(None) == [] + assert parse_assertions([]) == [] + + +def test_parse_tool_call_assertion(): + assertions = parse_assertions([{"type": "tool_call", "tool": "clarification"}]) + + assert assertions == [ToolCallAssertion(tool="clarification")] + + +def test_parse_tool_call_with_args_and_min_count(): + assertions = parse_assertions( + [ + { + "type": "tool_call", + "tool": "execute_sql", + "args": {"sql_query": "SELECT 1"}, + "min_count": 2, + } + ] + ) + + assert assertions == [ToolCallAssertion(tool="execute_sql", args={"sql_query": "SELECT 1"}, min_count=2)] + + +def test_parse_rejects_unknown_type(): + with pytest.raises(AssertionConfigError, match="unknown type 'text'"): + parse_assertions([{"type": "text", "pattern": "hello"}]) + + +def test_parse_rejects_missing_tool(): + with pytest.raises(AssertionConfigError, match="non-empty 'tool'"): + parse_assertions([{"type": "tool_call"}]) + + +def test_evaluate_tool_call_passes_when_present(): + passed, msg = evaluate_tool_call_assertion( + ToolCallAssertion(tool="clarification"), + [{"toolName": "clarification", "args": {"question": "Which period?"}}], + ) + + assert passed is True + assert msg == "tool_call: clarification" + + +def test_evaluate_tool_call_fails_when_missing(): + # Models the issue case: agent answered numerically via SQL instead of asking. + passed, msg = evaluate_tool_call_assertion( + ToolCallAssertion(tool="clarification"), + [{"toolName": "execute_sql", "args": {"sql_query": "SELECT SUM(amount) FROM orders"}}], + ) + + assert passed is False + assert msg == "missing tool_call: clarification" + + +def test_evaluate_tool_call_fails_on_empty_trace(): + passed, msg = evaluate_tool_call_assertion(ToolCallAssertion(tool="clarification"), []) + + assert passed is False + assert msg == "missing tool_call: clarification" + + +def test_evaluate_tool_call_args_subset_match(): + assertion = ToolCallAssertion(tool="execute_sql", args={"sql_query": "SELECT 1"}) + passed, msg = evaluate_tool_call_assertion( + assertion, + [ + { + "toolName": "execute_sql", + "args": {"sql_query": "SELECT 1", "limit": 100}, + } + ], + ) + + assert passed is True + assert "args matched" in msg + + +def test_evaluate_tool_call_args_mismatch(): + assertion = ToolCallAssertion(tool="execute_sql", args={"sql_query": "SELECT 1"}) + passed, msg = evaluate_tool_call_assertion( + assertion, + [{"toolName": "execute_sql", "args": {"sql_query": "SELECT 2"}}], + ) + + assert passed is False + assert "without matching args" in msg + + +def test_evaluate_tool_call_min_count(): + assertion = ToolCallAssertion(tool="read", min_count=2) + one_call = [{"toolName": "read", "args": {"path": "a"}}] + two_calls = one_call + [{"toolName": "read", "args": {"path": "b"}}] + + assert evaluate_tool_call_assertion(assertion, one_call)[0] is False + assert evaluate_tool_call_assertion(assertion, two_calls)[0] is True + + +def test_evaluate_assertions_combines_messages(): + assertions = [ + ToolCallAssertion(tool="clarification"), + ToolCallAssertion(tool="execute_sql"), + ] + passed, msg = evaluate_assertions( + assertions, + [{"toolName": "clarification", "args": {"question": "?"}}], + ) + + assert passed is False + assert msg == "tool_call: clarification; missing tool_call: execute_sql" + + +def test_combine_check_messages_skips_empty(): + assert combine_check_messages("tool_call: clarification", "match") == "tool_call: clarification; match" + assert combine_check_messages("", "match") == "match" + assert combine_check_messages("tool_call: clarification", "") == "tool_call: clarification" diff --git a/cli/tests/nao_core/commands/test_case.py b/cli/tests/nao_core/commands/test_case.py index 80de67a64..a23937bdf 100644 --- a/cli/tests/nao_core/commands/test_case.py +++ b/cli/tests/nao_core/commands/test_case.py @@ -1,4 +1,9 @@ -from nao_core.commands.test.case import discover_tests +from pathlib import Path + +import pytest + +from nao_core.commands.test.assertions import ToolCallAssertion +from nao_core.commands.test.case import TestCase, discover_tests def test_discover_tests_is_recursive(tmp_path): @@ -22,3 +27,34 @@ def test_discover_tests_ignores_outputs_dir(tmp_path): cases = discover_tests(tmp_path) assert {c.name for c in cases} == {"real"} + + +def test_from_yaml_loads_tool_call_assertions(tmp_path): + path = tmp_path / "ambiguous_revenue.yml" + path.write_text( + "\n".join( + [ + "name: ambiguous_revenue_period", + "prompt: What was the revenue?", + "assertions:", + " - type: tool_call", + " tool: clarification", + "", + ] + ) + ) + + case = TestCase.from_yaml(path) + + assert case.name == "ambiguous_revenue_period" + assert case.prompt == "What was the revenue?" + assert case.sql is None + assert case.assertions == [ToolCallAssertion(tool="clarification")] + + +def test_from_yaml_rejects_invalid_assertions(tmp_path): + path = tmp_path / "bad.yml" + path.write_text("prompt: hi\nassertions:\n - type: nope\n") + + with pytest.raises(Exception, match="unknown type"): + TestCase.from_yaml(Path(path)) diff --git a/cli/tests/nao_core/commands/test_runner.py b/cli/tests/nao_core/commands/test_runner.py index e62fb5954..63ee31eca 100644 --- a/cli/tests/nao_core/commands/test_runner.py +++ b/cli/tests/nao_core/commands/test_runner.py @@ -6,6 +6,7 @@ import pytest +from nao_core.commands.test.assertions import ToolCallAssertion from nao_core.commands.test.case import TestCase as NaoTestCase from nao_core.commands.test.client import ( AgentClientError, @@ -253,6 +254,129 @@ def test_run_test_records_reference_sql_on_client_error(monkeypatch): assert result.details.reference_sql == "select 1" +def test_run_test_assertion_only_passes_when_tool_called(monkeypatch): + """Issue #1261: assert the agent asked a follow-up (clarification tool).""" + test_case = NaoTestCase( + name="ambiguous_revenue_period", + prompt="What was the revenue?", + file_path=Path("tests/ambiguous_revenue_period.yml"), + sql=None, + assertions=[ToolCallAssertion(tool="clarification")], + ) + model = ModelConfig(provider="openai", model_id="custom-model") + client = Mock() + client.run_test.return_value = AgentTestResult( + text="Which time period should I use?", + tool_calls=[{"toolName": "clarification", "args": {"question": "Which time period?"}}], + usage=TokenUsage(totalTokens=10), + cost=TokenCost(totalCost=0.01), + finish_reason="stop", + duration_ms=5, + ) + monkeypatch.setattr(test_runner_module, "get_client", lambda **_: client) + + result = run_test(test_case, model) + + assert result.passed is True + assert result.message == "tool_call: clarification" + + +def test_run_test_assertion_only_fails_when_agent_returns_numeric_answer(monkeypatch): + """Issue #1261 fail direction: numeric SQL answer instead of a follow-up.""" + test_case = NaoTestCase( + name="ambiguous_revenue_period", + prompt="What was the revenue?", + file_path=Path("tests/ambiguous_revenue_period.yml"), + sql=None, + assertions=[ToolCallAssertion(tool="clarification")], + ) + model = ModelConfig(provider="openai", model_id="custom-model") + client = Mock() + client.run_test.return_value = AgentTestResult( + text="Total revenue is 52123123", + tool_calls=[ + { + "toolName": "execute_sql", + "args": {"sql_query": "SELECT SUM(amount) AS total_revenue FROM orders"}, + } + ], + usage=TokenUsage(totalTokens=20), + cost=TokenCost(totalCost=0.02), + finish_reason="stop", + duration_ms=8, + ) + monkeypatch.setattr(test_runner_module, "get_client", lambda **_: client) + + result = run_test(test_case, model) + + assert result.passed is False + assert result.message == "missing tool_call: clarification" + + +def test_run_test_combines_sql_verification_and_tool_assertion(monkeypatch): + test_case = NaoTestCase( + name="orders", + prompt="total revenue", + file_path=Path("tests/orders.yml"), + sql="select 1", + assertions=[ToolCallAssertion(tool="execute_sql")], + ) + model = ModelConfig(provider="openai", model_id="custom-model") + client = Mock() + client.run_test.return_value = AgentTestResult( + text="ok", + tool_calls=[{"toolName": "execute_sql", "args": {"sql_query": "SELECT 1"}}], + usage=TokenUsage(totalTokens=5), + cost=TokenCost(totalCost=0.0), + finish_reason="stop", + duration_ms=2, + verification=VerificationResult( + data=[{"total": 1}], + expectedData=[{"total": 1}], + expectedColumns=["total"], + sql="SELECT total FROM query_abc", + ), + ) + monkeypatch.setattr(test_runner_module, "get_client", lambda **_: client) + + result = run_test(test_case, model) + + assert result.passed is True + assert result.message == "tool_call: execute_sql; match" + + +def test_run_test_fails_when_sql_matches_but_tool_assertion_missing(monkeypatch): + test_case = NaoTestCase( + name="orders", + prompt="total revenue", + file_path=Path("tests/orders.yml"), + sql="select 1", + assertions=[ToolCallAssertion(tool="clarification")], + ) + model = ModelConfig(provider="openai", model_id="custom-model") + client = Mock() + client.run_test.return_value = AgentTestResult( + text="ok", + tool_calls=[{"toolName": "execute_sql", "args": {}}], + usage=TokenUsage(totalTokens=5), + cost=TokenCost(totalCost=0.0), + finish_reason="stop", + duration_ms=2, + verification=VerificationResult( + data=[{"total": 1}], + expectedData=[{"total": 1}], + expectedColumns=["total"], + ), + ) + monkeypatch.setattr(test_runner_module, "get_client", lambda **_: client) + + result = run_test(test_case, model) + + assert result.passed is False + assert "missing tool_call: clarification" in result.message + assert "match" in result.message + + def test_filter_test_cases_by_folder(tmp_path): tests_dir = tmp_path / "tests" tc_orders = NaoTestCase(name="orders", prompt="p1", file_path=tests_dir / "revenue" / "orders.yml", sql="select 1") diff --git a/skills/create-context-tests/templates/test.yaml b/skills/create-context-tests/templates/test.yaml index 5e5a41edd..4f8b2d8df 100644 --- a/skills/create-context-tests/templates/test.yaml +++ b/skills/create-context-tests/templates/test.yaml @@ -25,3 +25,12 @@ sql: | # category: revenue | activity | conversion | churn | retention | ... # difficulty: easy | medium | hard # notes: why this test matters / what failure mode it catches +# +# Assert intermediate agent actions (independent of final SQL output): +# assertions: +# - type: tool_call +# tool: clarification # e.g. agent must ask a follow-up +# # - type: tool_call +# # tool: execute_sql +# # args: { sql_query: "..." } # optional subset match +# # min_count: 1 From 4880eea9921dcb2d09051034edaa3f3e4c6b2bab Mon Sep 17 00:00:00 2001 From: justin212407 Date: Thu, 27 Aug 2026 16:07:11 +0530 Subject: [PATCH 2/3] fix: enable clarification assertions in test mode Signed-off-by: justin212407 --- apps/backend/src/agents/tools/index.ts | 8 -- .../src/components/ai/system-prompt.tsx | 18 +-- .../src/handlers/automation.handler.ts | 3 +- apps/backend/src/services/agent.ts | 17 +-- .../src/services/test-agent.service.ts | 1 - .../nao_core/commands/test_assertions.py | 130 ------------------ 6 files changed, 14 insertions(+), 163 deletions(-) delete mode 100644 cli/tests/nao_core/commands/test_assertions.py diff --git a/apps/backend/src/agents/tools/index.ts b/apps/backend/src/agents/tools/index.ts index e4f60cf2c..d1cb518e7 100644 --- a/apps/backend/src/agents/tools/index.ts +++ b/apps/backend/src/agents/tools/index.ts @@ -55,12 +55,6 @@ export const getTools = ( agentSettings: AgentSettings | null, extraTools?: Record, options: { - /** - * @deprecated No longer strips tools. Kept so existing callers that pass - * `testMode` continue to typecheck. Clarification stays available so - * `nao test` can assert intermediate actions (e.g. follow-up questions). - */ - testMode?: boolean; mcpEnabled?: boolean; mcpServers?: string[] | null; excludeFollowUps?: boolean; @@ -95,8 +89,6 @@ export const getTools = ( : {}; const { execute_python, execute_sandboxed_code, suggest_follow_ups, write: writeTool, ...rest } = tools; - // Keep clarification available in testMode so `nao test` can assert follow-up - // questions and other intermediate actions on the recorded tool-call trace. const baseTools = { ...rest, ...(isStorageEnabled() && { write: writeTool }), diff --git a/apps/backend/src/components/ai/system-prompt.tsx b/apps/backend/src/components/ai/system-prompt.tsx index 455f1caae..387d76c29 100644 --- a/apps/backend/src/components/ai/system-prompt.tsx +++ b/apps/backend/src/components/ai/system-prompt.tsx @@ -29,7 +29,6 @@ type SystemPromptProps = { /** Names of MCP servers the agent is allowed to call (tools discovered as on-disk specs). */ mcpServers?: string[]; timezone?: string; - testMode?: boolean; /** Names of the tools in the run's tool set — rules for surface-dependent tools (e.g. display_map) are only emitted when the tool is present. Omit to include every rule. */ toolNames?: string[]; options?: SystemPromptOptions; @@ -52,7 +51,6 @@ export function SystemPrompt({ customCharts = [], mcpServers = [], timezone, - testMode, toolNames, options = {}, }: SystemPromptProps) { @@ -106,15 +104,13 @@ export function SystemPrompt({ researching. , If you can execute a SQL query, use the execute_sql tool for it., - !testMode && ( - - Use the clarification tool when the user's request is genuinely ambiguous and - proceeding would likely produce the wrong result (e.g. multiple plausible tables, unclear - time range, undefined metric). If you need to ask another clarifying question after the user - answers, call the clarification tool again instead of asking in plain text, - bullet lists, or examples. - - ), + + Use the clarification tool when the user's request is genuinely ambiguous and + proceeding would likely produce the wrong result (e.g. multiple plausible tables, unclear time + range, undefined metric). If you need to ask another clarifying question after the user answers, + call the clarification tool again instead of asking in plain text, bullet lists, or + examples. + , ...dialectToolCallRules, ]} diff --git a/apps/backend/src/handlers/automation.handler.ts b/apps/backend/src/handlers/automation.handler.ts index e067ca79d..bcb5faf42 100644 --- a/apps/backend/src/handlers/automation.handler.ts +++ b/apps/backend/src/handlers/automation.handler.ts @@ -124,7 +124,7 @@ async function finishAutomationRun(automation: AutomationWithSchedule, run: DBAu { excludeFollowUps: true, supportsCustomCharts: false, - tools: ({ chat: agentChat, agentSettings, webTools }) => + tools: ({ agentSettings, webTools }) => getTools( agentSettings, { @@ -139,7 +139,6 @@ async function finishAutomationRun(automation: AutomationWithSchedule, run: DBAu }), }, { - testMode: agentChat.testMode, mcpEnabled: automation.mcpEnabled, mcpServers: automation.mcpServers, excludeFollowUps: true, diff --git a/apps/backend/src/services/agent.ts b/apps/backend/src/services/agent.ts index 173c58bde..f6843a04e 100644 --- a/apps/backend/src/services/agent.ts +++ b/apps/backend/src/services/agent.ts @@ -101,7 +101,6 @@ export interface AgentRunResult { export type AgentChat = Pick & { forkMetadata?: ForkMetadata | null; - testMode?: boolean; }; /** Dependencies a tool resolver receives once a run's context has been resolved. */ @@ -119,26 +118,25 @@ export interface AgentToolsContext { export type AgentToolsResolver = (context: AgentToolsContext) => AgentTools | Promise; /** Default tool set for interactive runs: all built-ins, MCP tools and web search. */ -export const defaultAgentTools: AgentToolsResolver = ({ chat, agentSettings, webTools, customBoundaries }) => - getTools(agentSettings, webTools ?? {}, { testMode: chat.testMode, customBoundaries }); +export const defaultAgentTools: AgentToolsResolver = ({ agentSettings, webTools, customBoundaries }) => + getTools(agentSettings, webTools ?? {}, { customBoundaries }); /** Default tool set minus the given built-ins — for runs whose surface cannot render them. */ export const defaultAgentToolsExcluding = (excludeBuiltinTools: string[]): AgentToolsResolver => - ({ chat, agentSettings, webTools, customBoundaries }) => - getTools(agentSettings, webTools ?? {}, { testMode: chat.testMode, excludeBuiltinTools, customBoundaries }); + ({ agentSettings, webTools, customBoundaries }) => + getTools(agentSettings, webTools ?? {}, { excludeBuiltinTools, customBoundaries }); /** * Admin-mode tool set: the same `execute_sql` tool the chat already uses (it * runs against nao's own app database when `ToolContext.adminMode` is set), * plus charting and follow-ups. Excludes the filesystem context tools. */ -export const adminAgentTools: AgentToolsResolver = ({ chat, agentSettings }) => +export const adminAgentTools: AgentToolsResolver = ({ agentSettings }) => getTools( agentSettings, {}, { - testMode: chat.testMode, builtinToolAllowlist: [ 'execute_sql', 'read_query_result', @@ -273,9 +271,7 @@ export class AgentService { const agentTools = await resolveTools({ chat, agentSettings, toolContext, webTools, customBoundaries }); const stopWhen: StopCondition[] = options.excludeFollowUps ? [stepCountIs(options.maxSteps ?? 20)] - : chat.testMode - ? [hasToolCall('suggest_follow_ups')] - : [hasToolCall('suggest_follow_ups'), hasToolCall('clarification')]; + : [hasToolCall('suggest_follow_ups'), hasToolCall('clarification')]; const agent = new AgentManager( chat, modelConfig, @@ -613,7 +609,6 @@ class AgentManager { customCharts, mcpServers, timezone, - testMode: this.chat.testMode, toolNames: Object.keys(this._agentTools), options: { canGrepSavedFiles: canGrepUserFiles() }, }), diff --git a/apps/backend/src/services/test-agent.service.ts b/apps/backend/src/services/test-agent.service.ts index aa0488106..73b33eef9 100644 --- a/apps/backend/src/services/test-agent.service.ts +++ b/apps/backend/src/services/test-agent.service.ts @@ -54,7 +54,6 @@ export class TestAgentService extends AgentService { messages: [userMessage], userId: 'test', projectId, - testMode: true, }; const agent = await this.create(tempChat, modelSelection); diff --git a/cli/tests/nao_core/commands/test_assertions.py b/cli/tests/nao_core/commands/test_assertions.py deleted file mode 100644 index f9c383b89..000000000 --- a/cli/tests/nao_core/commands/test_assertions.py +++ /dev/null @@ -1,130 +0,0 @@ -import pytest - -from nao_core.commands.test.assertions import ( - AssertionConfigError, - ToolCallAssertion, - combine_check_messages, - evaluate_assertions, - evaluate_tool_call_assertion, - parse_assertions, -) - - -def test_parse_assertions_empty(): - assert parse_assertions(None) == [] - assert parse_assertions([]) == [] - - -def test_parse_tool_call_assertion(): - assertions = parse_assertions([{"type": "tool_call", "tool": "clarification"}]) - - assert assertions == [ToolCallAssertion(tool="clarification")] - - -def test_parse_tool_call_with_args_and_min_count(): - assertions = parse_assertions( - [ - { - "type": "tool_call", - "tool": "execute_sql", - "args": {"sql_query": "SELECT 1"}, - "min_count": 2, - } - ] - ) - - assert assertions == [ToolCallAssertion(tool="execute_sql", args={"sql_query": "SELECT 1"}, min_count=2)] - - -def test_parse_rejects_unknown_type(): - with pytest.raises(AssertionConfigError, match="unknown type 'text'"): - parse_assertions([{"type": "text", "pattern": "hello"}]) - - -def test_parse_rejects_missing_tool(): - with pytest.raises(AssertionConfigError, match="non-empty 'tool'"): - parse_assertions([{"type": "tool_call"}]) - - -def test_evaluate_tool_call_passes_when_present(): - passed, msg = evaluate_tool_call_assertion( - ToolCallAssertion(tool="clarification"), - [{"toolName": "clarification", "args": {"question": "Which period?"}}], - ) - - assert passed is True - assert msg == "tool_call: clarification" - - -def test_evaluate_tool_call_fails_when_missing(): - # Models the issue case: agent answered numerically via SQL instead of asking. - passed, msg = evaluate_tool_call_assertion( - ToolCallAssertion(tool="clarification"), - [{"toolName": "execute_sql", "args": {"sql_query": "SELECT SUM(amount) FROM orders"}}], - ) - - assert passed is False - assert msg == "missing tool_call: clarification" - - -def test_evaluate_tool_call_fails_on_empty_trace(): - passed, msg = evaluate_tool_call_assertion(ToolCallAssertion(tool="clarification"), []) - - assert passed is False - assert msg == "missing tool_call: clarification" - - -def test_evaluate_tool_call_args_subset_match(): - assertion = ToolCallAssertion(tool="execute_sql", args={"sql_query": "SELECT 1"}) - passed, msg = evaluate_tool_call_assertion( - assertion, - [ - { - "toolName": "execute_sql", - "args": {"sql_query": "SELECT 1", "limit": 100}, - } - ], - ) - - assert passed is True - assert "args matched" in msg - - -def test_evaluate_tool_call_args_mismatch(): - assertion = ToolCallAssertion(tool="execute_sql", args={"sql_query": "SELECT 1"}) - passed, msg = evaluate_tool_call_assertion( - assertion, - [{"toolName": "execute_sql", "args": {"sql_query": "SELECT 2"}}], - ) - - assert passed is False - assert "without matching args" in msg - - -def test_evaluate_tool_call_min_count(): - assertion = ToolCallAssertion(tool="read", min_count=2) - one_call = [{"toolName": "read", "args": {"path": "a"}}] - two_calls = one_call + [{"toolName": "read", "args": {"path": "b"}}] - - assert evaluate_tool_call_assertion(assertion, one_call)[0] is False - assert evaluate_tool_call_assertion(assertion, two_calls)[0] is True - - -def test_evaluate_assertions_combines_messages(): - assertions = [ - ToolCallAssertion(tool="clarification"), - ToolCallAssertion(tool="execute_sql"), - ] - passed, msg = evaluate_assertions( - assertions, - [{"toolName": "clarification", "args": {"question": "?"}}], - ) - - assert passed is False - assert msg == "tool_call: clarification; missing tool_call: execute_sql" - - -def test_combine_check_messages_skips_empty(): - assert combine_check_messages("tool_call: clarification", "match") == "tool_call: clarification; match" - assert combine_check_messages("", "match") == "match" - assert combine_check_messages("tool_call: clarification", "") == "tool_call: clarification" From f0d6e615b0d6537ba9df331f61b18fb67e63a373 Mon Sep 17 00:00:00 2001 From: justin212407 Date: Fri, 28 Aug 2026 16:10:41 +0530 Subject: [PATCH 3/3] fix: address review feedback on test assertions and prompt gating Signed-off-by: justin212407 --- .../src/components/ai/system-prompt.tsx | 18 +++++---- cli/nao_core/commands/test/assertions.py | 2 +- cli/nao_core/commands/test/runner.py | 21 +++++++++- cli/nao_core/config/databases/duckdb.py | 38 ++++++++++++++++++- cli/tests/nao_core/commands/test_case.py | 22 ++++++++--- cli/tests/nao_core/commands/test_runner.py | 28 ++++++++++++++ 6 files changed, 114 insertions(+), 15 deletions(-) diff --git a/apps/backend/src/components/ai/system-prompt.tsx b/apps/backend/src/components/ai/system-prompt.tsx index 387d76c29..9b712b06f 100644 --- a/apps/backend/src/components/ai/system-prompt.tsx +++ b/apps/backend/src/components/ai/system-prompt.tsx @@ -104,13 +104,17 @@ export function SystemPrompt({ researching. , If you can execute a SQL query, use the execute_sql tool for it., - - Use the clarification tool when the user's request is genuinely ambiguous and - proceeding would likely produce the wrong result (e.g. multiple plausible tables, unclear time - range, undefined metric). If you need to ask another clarifying question after the user answers, - call the clarification tool again instead of asking in plain text, bullet lists, or - examples. - , + ...(hasTool('clarification') + ? [ + + Use the clarification tool when the user's request is genuinely + ambiguous and proceeding would likely produce the wrong result (e.g. multiple + plausible tables, unclear time range, undefined metric). If you need to ask another + clarifying question after the user answers, call the clarification tool + again instead of asking in plain text, bullet lists, or examples. + , + ] + : []), ...dialectToolCallRules, ]} diff --git a/cli/nao_core/commands/test/assertions.py b/cli/nao_core/commands/test/assertions.py index a67b067f4..6e1fe45bb 100644 --- a/cli/nao_core/commands/test/assertions.py +++ b/cli/nao_core/commands/test/assertions.py @@ -111,7 +111,7 @@ def evaluate_tool_call_assertion( return ( False, f"missing tool_call: {assertion.tool} with args {assertion.args} " - f"(found {same_tool} call(s) without matching args)", + f"(found {count} matching, need >= {assertion.min_count}; {same_tool} total)", ) return False, f"missing tool_call: {assertion.tool} with args {assertion.args}" diff --git a/cli/nao_core/commands/test/runner.py b/cli/nao_core/commands/test/runner.py index f730bd280..62fa2bf54 100644 --- a/cli/nao_core/commands/test/runner.py +++ b/cli/nao_core/commands/test/runner.py @@ -251,7 +251,7 @@ def run_test( ), ) - if test_case.assertions: + if test_case.assertions and not test_case.sql: status = "[green]✓[/green]" if assertions_passed else "[red]✗[/red]" UI.print(f" {status} {assertions_msg}") return TestRunResult( @@ -270,6 +270,25 @@ def run_test( ), ) + if test_case.assertions and test_case.sql: + msg = combine_check_messages(assertions_msg, "no verification data") + UI.print(f" [red]✗[/red] {msg}") + return TestRunResult( + name=test_case.name, + model=str(model), + passed=False, + message=msg, + tokens=result.usage.totalTokens, + cost=result.cost.totalCost, + duration_ms=result.duration_ms, + tool_call_count=tool_call_count, + details=TestRunDetails( + response_text=result.text, + tool_calls=result.tool_calls, + reference_sql=test_case.sql, + ), + ) + UI.print("[yellow] ⚠ no verification data[/yellow]") return TestRunResult( name=test_case.name, diff --git a/cli/nao_core/config/databases/duckdb.py b/cli/nao_core/config/databases/duckdb.py index 673314438..a4f037c11 100644 --- a/cli/nao_core/config/databases/duckdb.py +++ b/cli/nao_core/config/databases/duckdb.py @@ -1,7 +1,7 @@ from __future__ import annotations from pathlib import Path -from typing import TYPE_CHECKING, Literal +from typing import TYPE_CHECKING, Any, Literal from pydantic import Field @@ -15,6 +15,42 @@ class DuckDBDatabaseContext(DatabaseContext): + """DuckDB context with table/column comment discovery via duckdb metadata functions.""" + + def description(self) -> str | None: + try: + query = ( + "SELECT comment FROM duckdb_tables() " + f"WHERE schema_name = '{self._schema}' AND table_name = '{self._table_name}'" + ) + row = self._fetchone(self._conn.raw_sql(query)) # type: ignore[union-attr] + if row and row[0] is not None: + text = str(row[0]).strip() + if text: + return text + except Exception: + pass + return None + + def columns(self) -> list[dict[str, Any]]: + cols = super().columns() + try: + col_descs = self._fetch_column_descriptions() + for col in cols: + if desc := col_descs.get(col["name"]): + col["description"] = desc + except Exception: + pass + return cols + + def _fetch_column_descriptions(self) -> dict[str, str]: + query = ( + "SELECT column_name, comment FROM duckdb_columns() " + f"WHERE schema_name = '{self._schema}' AND table_name = '{self._table_name}'" + ) + rows = self._fetchall(self._conn.raw_sql(query)) # type: ignore[union-attr] + return {row[0]: str(row[1]) for row in rows if row[1] is not None} + def _cast_complex_to_string(self, col_sql: str) -> str: return f"CAST({col_sql} AS VARCHAR)" diff --git a/cli/tests/nao_core/commands/test_case.py b/cli/tests/nao_core/commands/test_case.py index a23937bdf..773626728 100644 --- a/cli/tests/nao_core/commands/test_case.py +++ b/cli/tests/nao_core/commands/test_case.py @@ -1,8 +1,6 @@ -from pathlib import Path - import pytest -from nao_core.commands.test.assertions import ToolCallAssertion +from nao_core.commands.test.assertions import AssertionConfigError, ToolCallAssertion from nao_core.commands.test.case import TestCase, discover_tests @@ -56,5 +54,19 @@ def test_from_yaml_rejects_invalid_assertions(tmp_path): path = tmp_path / "bad.yml" path.write_text("prompt: hi\nassertions:\n - type: nope\n") - with pytest.raises(Exception, match="unknown type"): - TestCase.from_yaml(Path(path)) + with pytest.raises(AssertionConfigError, match="unknown type"): + TestCase.from_yaml(path) + + +def test_evaluate_tool_call_assertion_reports_matched_count_when_min_count_exceeds(): + from nao_core.commands.test.assertions import evaluate_tool_call_assertion + + assertion = ToolCallAssertion(tool="fetch", args={"type": "user"}, min_count=2) + tool_calls = [ + {"toolName": "fetch", "args": {"type": "user"}}, + {"toolName": "fetch", "args": {"type": "other"}}, + ] + passed, msg = evaluate_tool_call_assertion(assertion, tool_calls) + + assert passed is False + assert "(found 1 matching, need >= 2; 2 total)" in msg diff --git a/cli/tests/nao_core/commands/test_runner.py b/cli/tests/nao_core/commands/test_runner.py index 63ee31eca..e913760a8 100644 --- a/cli/tests/nao_core/commands/test_runner.py +++ b/cli/tests/nao_core/commands/test_runner.py @@ -377,6 +377,34 @@ def test_run_test_fails_when_sql_matches_but_tool_assertion_missing(monkeypatch) assert "match" in result.message +def test_run_test_fails_when_both_sql_and_assertions_configured_but_no_verification_data(monkeypatch): + test_case = NaoTestCase( + name="orders", + prompt="total revenue", + file_path=Path("tests/orders.yml"), + sql="select 1", + assertions=[ToolCallAssertion(tool="execute_sql")], + ) + model = ModelConfig(provider="openai", model_id="custom-model") + client = Mock() + client.run_test.return_value = AgentTestResult( + text="ok", + tool_calls=[{"toolName": "execute_sql", "args": {"sql_query": "SELECT 1"}}], + usage=TokenUsage(totalTokens=5), + cost=TokenCost(totalCost=0.0), + finish_reason="stop", + duration_ms=2, + verification=None, + ) + monkeypatch.setattr(test_runner_module, "get_client", lambda **_: client) + + result = run_test(test_case, model) + + assert result.passed is False + assert "tool_call: execute_sql" in result.message + assert "no verification data" in result.message + + def test_filter_test_cases_by_folder(tmp_path): tests_dir = tmp_path / "tests" tc_orders = NaoTestCase(name="orders", prompt="p1", file_path=tests_dir / "revenue" / "orders.yml", sql="select 1")