wehub-resource-sync
566 lines
18 KiB
Python
566 lines
18 KiB
Python
from unittest import mock
|
|
|
|
import pydantic
|
|
import pytest
|
|
|
|
from mlflow.genai.agent_tester import (
|
|
_DEFAULT_NUM_TEST_CASES,
|
|
_DEFAULT_TESTING_GUIDANCE,
|
|
AgentTestResult,
|
|
_AgentDescription,
|
|
_describe_agent_from_response,
|
|
_describe_agent_from_traces,
|
|
_generate_test_cases,
|
|
_get_agent_response_text,
|
|
_load_traces,
|
|
_resolve_agent_description,
|
|
_TestCase,
|
|
_TestCaseList,
|
|
)
|
|
from mlflow.genai.agent_tester import test_agent as run_test_agent
|
|
|
|
_MODEL = "openai:/gpt-4o"
|
|
|
|
|
|
def _make_agent_desc(description="A helper", capabilities=("assist",), limitations=()):
|
|
return _AgentDescription(
|
|
description=description,
|
|
capabilities=list(capabilities),
|
|
limitations=list(limitations),
|
|
)
|
|
|
|
|
|
def _make_test_case_list(n=2):
|
|
return _TestCaseList(
|
|
test_cases=[
|
|
_TestCase(
|
|
goal=f"Goal {i}",
|
|
persona=f"Persona {i}",
|
|
simulation_guidelines=[f"guideline {i}"],
|
|
)
|
|
for i in range(n)
|
|
]
|
|
)
|
|
|
|
|
|
def _llm_dispatcher(agent_desc, test_case_list):
|
|
"""Return the right object based on output_schema."""
|
|
|
|
def dispatch(**kwargs):
|
|
schema = kwargs.get("output_schema")
|
|
if schema is _AgentDescription:
|
|
return agent_desc
|
|
if schema is _TestCaseList:
|
|
return test_case_list
|
|
raise ValueError(f"Unexpected schema: {schema}")
|
|
|
|
return dispatch
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_llm():
|
|
with mock.patch("mlflow.genai.agent_tester.get_chat_completions_with_structured_output") as m:
|
|
yield m
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_no_active_trace():
|
|
with mock.patch("mlflow.get_last_active_trace_id", return_value=None):
|
|
yield
|
|
|
|
|
|
def test_agent_description_str():
|
|
desc = _AgentDescription(
|
|
description="A weather assistant",
|
|
capabilities=["forecast", "alerts"],
|
|
limitations=["no historical data"],
|
|
)
|
|
result = str(desc)
|
|
assert "A weather assistant" in result
|
|
assert "forecast" in result
|
|
assert "no historical data" in result
|
|
|
|
|
|
def test_agent_description_str_structure():
|
|
desc = _AgentDescription(
|
|
description="A coding assistant",
|
|
capabilities=["write code", "debug"],
|
|
limitations=["no internet access"],
|
|
)
|
|
result = str(desc)
|
|
assert result.startswith("Agent description:")
|
|
assert "Capabilities:" in result
|
|
assert "Limitations:" in result
|
|
assert "- write code" in result
|
|
assert "- no internet access" in result
|
|
|
|
|
|
def test_agent_description_validation_requires_fields():
|
|
with pytest.raises(pydantic.ValidationError, match="capabilities"):
|
|
_AgentDescription.model_validate({})
|
|
|
|
|
|
def test_agent_test_result_str():
|
|
issue = mock.MagicMock()
|
|
issue.severity = "high"
|
|
issue.name = "Issue: Confusing responses"
|
|
issue.description = "Agent gives contradictory answers."
|
|
issues_result = mock.MagicMock()
|
|
issues_result.issues = [issue]
|
|
result = AgentTestResult(
|
|
test_cases=[],
|
|
agent_description="A helpful assistant",
|
|
simulation_traces=[],
|
|
issues_result=issues_result,
|
|
)
|
|
text = str(result)
|
|
assert "A helpful assistant" in text
|
|
assert "Issues found: 1" in text
|
|
assert "[high] Issue: Confusing responses" in text
|
|
assert "Agent gives contradictory answers." in text
|
|
|
|
|
|
def test_test_case_requires_all_fields():
|
|
with pytest.raises(pydantic.ValidationError, match="persona"):
|
|
_TestCase.model_validate({"goal": "Do something"})
|
|
|
|
|
|
def test_test_case_list_empty_is_valid():
|
|
result = _TestCaseList.model_validate({"test_cases": []})
|
|
assert result.test_cases == []
|
|
|
|
|
|
def test_get_agent_response_text_returns_string_directly():
|
|
def predict(messages):
|
|
return "I am a helpful assistant."
|
|
|
|
result = _get_agent_response_text(predict)
|
|
assert result == "I am a helpful assistant."
|
|
|
|
|
|
def test_get_agent_response_text_dispatches_messages_kwarg():
|
|
received = {}
|
|
|
|
def agent(messages):
|
|
received["messages"] = messages
|
|
return "response"
|
|
|
|
_get_agent_response_text(agent)
|
|
assert "messages" in received
|
|
assert received["messages"][0]["role"] == "user"
|
|
|
|
|
|
def test_get_agent_response_text_dispatches_input_kwarg():
|
|
received = {}
|
|
|
|
def agent(input):
|
|
received["input"] = input
|
|
return "response"
|
|
|
|
_get_agent_response_text(agent)
|
|
assert "input" in received
|
|
|
|
|
|
def test_get_agent_response_text_returns_none_on_exception(mock_no_active_trace):
|
|
def predict(messages):
|
|
raise RuntimeError("boom")
|
|
|
|
result = _get_agent_response_text(predict)
|
|
assert result is None
|
|
|
|
|
|
def test_get_agent_response_text_falls_back_to_trace():
|
|
def predict(messages):
|
|
return None
|
|
|
|
mock_trace = mock.MagicMock()
|
|
with (
|
|
mock.patch("mlflow.get_last_active_trace_id", return_value="trace-123"),
|
|
mock.patch("mlflow.get_trace", return_value=mock_trace) as mock_get_trace,
|
|
mock.patch(
|
|
"mlflow.genai.utils.trace_utils.extract_outputs_from_trace",
|
|
return_value={"content": "trace text"},
|
|
),
|
|
mock.patch(
|
|
"mlflow.genai.utils.trace_utils.parse_outputs_to_str",
|
|
side_effect=[None, "trace text"],
|
|
),
|
|
):
|
|
result = _get_agent_response_text(predict)
|
|
|
|
mock_get_trace.assert_called_once_with("trace-123")
|
|
assert result == "trace text"
|
|
|
|
|
|
def test_get_agent_response_text_returns_none_when_no_output(mock_no_active_trace):
|
|
def predict(messages):
|
|
return None
|
|
|
|
result = _get_agent_response_text(predict)
|
|
assert result is None
|
|
|
|
|
|
def test_describe_agent_from_response_calls_llm(mock_llm):
|
|
expected = _make_agent_desc()
|
|
mock_llm.return_value = expected
|
|
|
|
result = _describe_agent_from_response("I help with questions.", model=_MODEL)
|
|
|
|
assert result == expected
|
|
mock_llm.assert_called_once()
|
|
call_kwargs = mock_llm.call_args.kwargs
|
|
assert call_kwargs["model_uri"] == _MODEL
|
|
assert call_kwargs["output_schema"] is _AgentDescription
|
|
assert any("I help with questions." in m.content for m in call_kwargs["messages"])
|
|
|
|
|
|
def test_describe_agent_from_traces_with_no_sessions(mock_llm):
|
|
expected = _make_agent_desc()
|
|
mock_llm.return_value = expected
|
|
|
|
with (
|
|
mock.patch("mlflow.genai.discovery.utils.group_traces_by_session", return_value={}),
|
|
mock.patch(
|
|
"mlflow.genai.utils.trace_utils.extract_available_tools_from_trace",
|
|
return_value=None,
|
|
),
|
|
):
|
|
result = _describe_agent_from_traces([], model=_MODEL)
|
|
|
|
assert result == expected
|
|
assert "(no traces)" in mock_llm.call_args.kwargs["messages"][-1].content
|
|
|
|
|
|
def test_describe_agent_from_traces_includes_conversation(mock_llm):
|
|
trace = mock.MagicMock()
|
|
expected = _make_agent_desc(description="A coding assistant", capabilities=["write code"])
|
|
mock_llm.return_value = expected
|
|
|
|
with (
|
|
mock.patch(
|
|
"mlflow.genai.discovery.utils.group_traces_by_session",
|
|
return_value={"sess-1": [trace]},
|
|
),
|
|
mock.patch(
|
|
"mlflow.genai.utils.trace_utils.resolve_conversation_from_session",
|
|
return_value=[{"role": "user", "content": "Write a function"}],
|
|
),
|
|
mock.patch(
|
|
"mlflow.genai.discovery.extraction.extract_execution_paths_for_session",
|
|
return_value="(no routing)",
|
|
),
|
|
mock.patch(
|
|
"mlflow.genai.utils.trace_utils.extract_available_tools_from_trace",
|
|
return_value=None,
|
|
),
|
|
):
|
|
result = _describe_agent_from_traces([trace], model=_MODEL)
|
|
|
|
assert result == expected
|
|
assert "Write a function" in mock_llm.call_args.kwargs["messages"][-1].content
|
|
|
|
|
|
def test_describe_agent_from_traces_includes_tools(mock_llm):
|
|
trace = mock.MagicMock()
|
|
tool = mock.MagicMock()
|
|
tool.function.name = "search_web"
|
|
expected = _make_agent_desc(description="A research agent", capabilities=["search_web"])
|
|
mock_llm.return_value = expected
|
|
|
|
with (
|
|
mock.patch("mlflow.genai.discovery.utils.group_traces_by_session", return_value={}),
|
|
mock.patch(
|
|
"mlflow.genai.utils.trace_utils.extract_available_tools_from_trace",
|
|
return_value=[tool],
|
|
),
|
|
):
|
|
result = _describe_agent_from_traces([trace], model=_MODEL)
|
|
|
|
assert result == expected
|
|
assert "search_web" in mock_llm.call_args.kwargs["messages"][-1].content
|
|
|
|
|
|
def test_generate_test_cases_uses_default_count_and_guidance(mock_llm):
|
|
agent_desc = _make_agent_desc(capabilities=["answer questions"])
|
|
mock_llm.return_value = _make_test_case_list(7)
|
|
|
|
result = _generate_test_cases(agent_desc, model=_MODEL)
|
|
|
|
assert len(result) == 7
|
|
user_content = mock_llm.call_args.kwargs["messages"][-1].content
|
|
system_content = mock_llm.call_args.kwargs["messages"][0].content
|
|
assert f"Generate {_DEFAULT_NUM_TEST_CASES}" in user_content
|
|
assert _DEFAULT_TESTING_GUIDANCE in system_content
|
|
|
|
|
|
def test_generate_test_cases_uses_provided_count_and_guidance(mock_llm):
|
|
agent_desc = _make_agent_desc(capabilities=["answer questions"])
|
|
mock_llm.return_value = _make_test_case_list(3)
|
|
|
|
result = _generate_test_cases(
|
|
agent_desc, model=_MODEL, num_test_cases=3, guidance="Focus on edge cases"
|
|
)
|
|
|
|
assert len(result) == 3
|
|
assert "Generate 3" in mock_llm.call_args.kwargs["messages"][-1].content
|
|
assert "Focus on edge cases" in mock_llm.call_args.kwargs["messages"][0].content
|
|
|
|
|
|
def test_generate_test_cases_raises_on_invalid_count(mock_llm):
|
|
agent_desc = _make_agent_desc()
|
|
with pytest.raises(ValueError, match="num_test_cases must be >= 1"):
|
|
_generate_test_cases(agent_desc, model=_MODEL, num_test_cases=0)
|
|
|
|
|
|
def test_generate_test_cases_returns_dicts(mock_llm):
|
|
agent_desc = _make_agent_desc()
|
|
mock_llm.return_value = _make_test_case_list(1)
|
|
|
|
result = _generate_test_cases(agent_desc, model=_MODEL)
|
|
|
|
assert isinstance(result[0], dict)
|
|
assert result[0]["goal"] == "Goal 0"
|
|
|
|
|
|
def test_load_traces_returns_none_when_no_experiment_id():
|
|
result = _load_traces(experiment_id=None)
|
|
assert result is None
|
|
|
|
|
|
def test_load_traces_searches_experiment():
|
|
mock_traces = [mock.MagicMock(), mock.MagicMock()]
|
|
|
|
with mock.patch("mlflow.search_traces", return_value=mock_traces) as mock_search:
|
|
result = _load_traces(experiment_id="exp-123")
|
|
|
|
assert result == mock_traces
|
|
mock_search.assert_called_once_with(
|
|
locations=["exp-123"],
|
|
max_results=50,
|
|
return_type="list",
|
|
)
|
|
|
|
|
|
def test_resolve_agent_description_uses_self_description(mock_llm):
|
|
agent_desc = _make_agent_desc(capabilities=["assist"])
|
|
mock_llm.return_value = agent_desc
|
|
|
|
def predict(messages):
|
|
return "I am a helpful assistant."
|
|
|
|
result = _resolve_agent_description(predict, None, None, _MODEL)
|
|
assert result == agent_desc
|
|
|
|
|
|
def test_resolve_agent_description_falls_back_to_traces(mock_llm):
|
|
traces = [mock.MagicMock()]
|
|
agent_desc = _make_agent_desc(capabilities=["assist"])
|
|
mock_llm.return_value = agent_desc
|
|
|
|
def predict(messages):
|
|
raise RuntimeError("cannot self-describe")
|
|
|
|
with (
|
|
mock.patch("mlflow.genai.discovery.utils.group_traces_by_session", return_value={}),
|
|
mock.patch(
|
|
"mlflow.genai.utils.trace_utils.extract_available_tools_from_trace",
|
|
return_value=None,
|
|
),
|
|
):
|
|
result = _resolve_agent_description(predict, None, traces, _MODEL)
|
|
|
|
assert result == agent_desc
|
|
|
|
|
|
def test_resolve_agent_description_loads_traces_from_experiment(mock_llm):
|
|
agent_desc = _make_agent_desc(capabilities=["assist"])
|
|
loaded_traces = [mock.MagicMock()]
|
|
mock_llm.return_value = agent_desc
|
|
|
|
def predict(messages):
|
|
raise RuntimeError("cannot self-describe")
|
|
|
|
with (
|
|
mock.patch("mlflow.search_traces", return_value=loaded_traces) as mock_search,
|
|
mock.patch("mlflow.genai.discovery.utils.group_traces_by_session", return_value={}),
|
|
mock.patch(
|
|
"mlflow.genai.utils.trace_utils.extract_available_tools_from_trace",
|
|
return_value=None,
|
|
),
|
|
):
|
|
result = _resolve_agent_description(predict, "exp-456", None, _MODEL)
|
|
|
|
assert result == agent_desc
|
|
mock_search.assert_called_once_with(locations=["exp-456"], max_results=50, return_type="list")
|
|
|
|
|
|
def test_resolve_agent_description_ignores_experiment_id_when_traces_provided(mock_llm):
|
|
traces = [mock.MagicMock()]
|
|
agent_desc = _make_agent_desc(capabilities=["assist"])
|
|
mock_llm.return_value = agent_desc
|
|
|
|
def predict(messages):
|
|
raise RuntimeError("cannot self-describe")
|
|
|
|
with (
|
|
mock.patch("mlflow.search_traces") as mock_search,
|
|
mock.patch("mlflow.genai.discovery.utils.group_traces_by_session", return_value={}),
|
|
mock.patch(
|
|
"mlflow.genai.utils.trace_utils.extract_available_tools_from_trace",
|
|
return_value=None,
|
|
),
|
|
):
|
|
result = _resolve_agent_description(predict, "exp-456", traces, _MODEL)
|
|
|
|
assert result == agent_desc
|
|
mock_search.assert_not_called()
|
|
|
|
|
|
def test_resolve_agent_description_falls_back_when_llm_raises(mock_llm):
|
|
traces = [mock.MagicMock()]
|
|
agent_desc = _make_agent_desc(capabilities=["assist"])
|
|
# First call (self-description) raises, second call (traces) succeeds
|
|
mock_llm.side_effect = [RuntimeError("LLM error"), agent_desc]
|
|
|
|
def predict(messages):
|
|
return "I am a helpful assistant."
|
|
|
|
with (
|
|
mock.patch("mlflow.genai.discovery.utils.group_traces_by_session", return_value={}),
|
|
mock.patch(
|
|
"mlflow.genai.utils.trace_utils.extract_available_tools_from_trace",
|
|
return_value=None,
|
|
),
|
|
):
|
|
result = _resolve_agent_description(predict, None, traces, _MODEL)
|
|
|
|
assert result == agent_desc
|
|
|
|
|
|
def test_resolve_agent_description_returns_default_when_all_fail():
|
|
def predict(messages):
|
|
raise RuntimeError("cannot self-describe")
|
|
|
|
with mock.patch("mlflow.search_traces", return_value=None):
|
|
result = _resolve_agent_description(predict, None, None, _MODEL)
|
|
|
|
assert result.description == "A conversational AI agent"
|
|
assert result.capabilities == ["general conversation"]
|
|
|
|
|
|
def test_test_agent_uses_default_model(mock_llm):
|
|
agent_desc = _make_agent_desc()
|
|
mock_llm.side_effect = _llm_dispatcher(agent_desc, _make_test_case_list())
|
|
mock_issues = mock.MagicMock()
|
|
|
|
def predict(messages):
|
|
return "I am a helpful assistant."
|
|
|
|
with (
|
|
mock.patch(
|
|
# Patching the definition site works here because `test_agent` imports
|
|
# `get_default_simulation_model` via a local `from ... import` inside the
|
|
# function body, so the name is re-bound from the (already-patched) module
|
|
# object each time `test_agent` is called.
|
|
"mlflow.genai.simulators.utils.get_default_simulation_model",
|
|
return_value=_MODEL,
|
|
) as mock_get_model,
|
|
mock.patch("mlflow.genai.simulators.ConversationSimulator") as mock_sim_cls,
|
|
mock.patch("mlflow.genai.discovery.pipeline.discover_issues", return_value=mock_issues),
|
|
):
|
|
mock_sim_cls.return_value.simulate.return_value = [[mock.MagicMock()]]
|
|
result = run_test_agent(predict)
|
|
|
|
mock_get_model.assert_called_once()
|
|
assert result.agent_description == str(agent_desc)
|
|
|
|
|
|
def test_test_agent_passes_model_to_simulator_and_discovery(mock_llm):
|
|
agent_desc = _make_agent_desc()
|
|
mock_llm.side_effect = _llm_dispatcher(agent_desc, _make_test_case_list())
|
|
mock_issues = mock.MagicMock()
|
|
|
|
def predict(messages):
|
|
return "I am a helpful assistant."
|
|
|
|
with (
|
|
mock.patch("mlflow.genai.simulators.ConversationSimulator") as mock_sim_cls,
|
|
mock.patch(
|
|
"mlflow.genai.discovery.pipeline.discover_issues", return_value=mock_issues
|
|
) as mock_discover,
|
|
):
|
|
mock_sim_cls.return_value.simulate.return_value = [[]]
|
|
run_test_agent(predict, model=_MODEL)
|
|
|
|
assert all(call.kwargs["model_uri"] == _MODEL for call in mock_llm.call_args_list)
|
|
assert mock_sim_cls.call_args.kwargs["user_model"] == _MODEL
|
|
assert mock_discover.call_args.kwargs["model"] == _MODEL
|
|
|
|
|
|
def test_test_agent_returns_correct_result(mock_llm):
|
|
agent_desc = _make_agent_desc()
|
|
test_case_list = _make_test_case_list(2)
|
|
mock_llm.side_effect = _llm_dispatcher(agent_desc, test_case_list)
|
|
mock_sim_traces = [[mock.MagicMock()], [mock.MagicMock()]]
|
|
mock_issues = mock.MagicMock()
|
|
|
|
def predict(messages):
|
|
return "I am a helpful assistant."
|
|
|
|
with (
|
|
mock.patch("mlflow.genai.simulators.ConversationSimulator") as mock_sim_cls,
|
|
mock.patch("mlflow.genai.discovery.pipeline.discover_issues", return_value=mock_issues),
|
|
):
|
|
mock_sim_cls.return_value.simulate.return_value = mock_sim_traces
|
|
result = run_test_agent(predict, model=_MODEL)
|
|
|
|
assert len(result.test_cases) == 2
|
|
assert result.agent_description == str(agent_desc)
|
|
assert result.simulation_traces == mock_sim_traces
|
|
assert result.issues_result is mock_issues
|
|
|
|
|
|
def test_test_agent_flattens_traces_for_issue_detection(mock_llm):
|
|
agent_desc = _make_agent_desc()
|
|
mock_llm.side_effect = _llm_dispatcher(agent_desc, _make_test_case_list(2))
|
|
mock_issues = mock.MagicMock()
|
|
t1 = mock.MagicMock()
|
|
t2 = mock.MagicMock()
|
|
t3 = mock.MagicMock()
|
|
|
|
def predict(messages):
|
|
return "I am a helpful assistant."
|
|
|
|
with (
|
|
mock.patch("mlflow.genai.simulators.ConversationSimulator") as mock_sim_cls,
|
|
mock.patch(
|
|
"mlflow.genai.discovery.pipeline.discover_issues", return_value=mock_issues
|
|
) as mock_discover,
|
|
):
|
|
mock_sim_cls.return_value.simulate.return_value = [[t1, t2], [t3]]
|
|
run_test_agent(predict, model=_MODEL)
|
|
|
|
flat_traces = mock_discover.call_args.kwargs["traces"]
|
|
assert flat_traces == [t1, t2, t3]
|
|
|
|
|
|
def test_test_agent_passes_max_turns_and_max_issues(mock_llm):
|
|
agent_desc = _make_agent_desc()
|
|
mock_llm.side_effect = _llm_dispatcher(agent_desc, _make_test_case_list(1))
|
|
mock_issues = mock.MagicMock()
|
|
|
|
def predict(messages):
|
|
return "I am a helpful assistant."
|
|
|
|
with (
|
|
mock.patch("mlflow.genai.simulators.ConversationSimulator") as mock_sim_cls,
|
|
mock.patch(
|
|
"mlflow.genai.discovery.pipeline.discover_issues", return_value=mock_issues
|
|
) as mock_discover,
|
|
):
|
|
mock_sim_cls.return_value.simulate.return_value = [[]]
|
|
run_test_agent(predict, model=_MODEL, max_turns=5, max_issues=10)
|
|
|
|
assert mock_sim_cls.call_args.kwargs["max_turns"] == 5
|
|
assert mock_discover.call_args.kwargs["max_issues"] == 10
|