mlflow--mlflow
wehub-resource-sync
137 行
4.5 KiB
Python
137 行
4.5 KiB
Python
import json
|
|
from dataclasses import dataclass
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
|
|
from mlflow.entities.trace import Trace
|
|
from mlflow.entities.trace_info import TraceInfo
|
|
from mlflow.entities.trace_location import TraceLocation
|
|
from mlflow.entities.trace_state import TraceState
|
|
from mlflow.genai.judges.utils.tool_calling_utils import _process_tool_calls
|
|
from mlflow.types.llm import ChatMessage, ToolCall
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_trace():
|
|
trace_info = TraceInfo(
|
|
trace_id="test-trace",
|
|
trace_location=TraceLocation.from_experiment_id("0"),
|
|
request_time=1234567890,
|
|
state=TraceState.OK,
|
|
)
|
|
return Trace(info=trace_info, data=None)
|
|
|
|
|
|
def _make_tool_call(call_id, name, arguments="{}"):
|
|
return ToolCall(
|
|
id=call_id,
|
|
function={"name": name, "arguments": arguments},
|
|
)
|
|
|
|
|
|
def test_process_tool_calls_success(mock_trace):
|
|
tool_call = _make_tool_call("call_123", "test_tool", '{"arg": "value"}')
|
|
|
|
with mock.patch(
|
|
"mlflow.genai.judges.tools.registry._judge_tool_registry.invoke"
|
|
) as mock_invoke:
|
|
mock_invoke.return_value = {"result": "success"}
|
|
|
|
result = _process_tool_calls(tool_calls=[tool_call], trace=mock_trace)
|
|
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], ChatMessage)
|
|
assert result[0].role == "tool"
|
|
assert result[0].tool_call_id == "call_123"
|
|
assert result[0].name == "test_tool"
|
|
assert json.loads(result[0].content) == {"result": "success"}
|
|
|
|
|
|
def test_process_tool_calls_with_error(mock_trace):
|
|
tool_call = _make_tool_call("call_456", "failing_tool")
|
|
|
|
with mock.patch(
|
|
"mlflow.genai.judges.tools.registry._judge_tool_registry.invoke"
|
|
) as mock_invoke:
|
|
mock_invoke.side_effect = RuntimeError("Tool execution failed")
|
|
|
|
result = _process_tool_calls(tool_calls=[tool_call], trace=mock_trace)
|
|
|
|
assert len(result) == 1
|
|
assert result[0].role == "tool"
|
|
assert result[0].tool_call_id == "call_456"
|
|
assert "Error: Tool execution failed" in result[0].content
|
|
|
|
|
|
def test_process_tool_calls_multiple(mock_trace):
|
|
tool_call_1 = _make_tool_call("call_1", "tool_1")
|
|
tool_call_2 = _make_tool_call("call_2", "tool_2")
|
|
|
|
with mock.patch(
|
|
"mlflow.genai.judges.tools.registry._judge_tool_registry.invoke"
|
|
) as mock_invoke:
|
|
mock_invoke.side_effect = [{"result": "first"}, {"result": "second"}]
|
|
|
|
result = _process_tool_calls(tool_calls=[tool_call_1, tool_call_2], trace=mock_trace)
|
|
|
|
assert len(result) == 2
|
|
assert result[0].tool_call_id == "call_1"
|
|
assert result[1].tool_call_id == "call_2"
|
|
assert json.loads(result[0].content) == {"result": "first"}
|
|
assert json.loads(result[1].content) == {"result": "second"}
|
|
|
|
|
|
def test_process_tool_calls_with_dataclass(mock_trace):
|
|
@dataclass
|
|
class ToolResult:
|
|
status: str
|
|
count: int
|
|
|
|
tool_call = _make_tool_call("call_789", "dataclass_tool")
|
|
|
|
with mock.patch(
|
|
"mlflow.genai.judges.tools.registry._judge_tool_registry.invoke"
|
|
) as mock_invoke:
|
|
mock_invoke.return_value = ToolResult(status="ok", count=42)
|
|
|
|
result = _process_tool_calls(tool_calls=[tool_call], trace=mock_trace)
|
|
|
|
assert len(result) == 1
|
|
assert result[0].role == "tool"
|
|
content = json.loads(result[0].content)
|
|
assert content == {"status": "ok", "count": 42}
|
|
|
|
|
|
def test_process_tool_calls_with_string_result(mock_trace):
|
|
tool_call = _make_tool_call("call_str", "string_tool")
|
|
|
|
with mock.patch(
|
|
"mlflow.genai.judges.tools.registry._judge_tool_registry.invoke"
|
|
) as mock_invoke:
|
|
mock_invoke.return_value = "Plain string result"
|
|
|
|
result = _process_tool_calls(tool_calls=[tool_call], trace=mock_trace)
|
|
|
|
assert len(result) == 1
|
|
assert result[0].role == "tool"
|
|
assert result[0].content == "Plain string result"
|
|
|
|
|
|
def test_process_tool_calls_mixed_success_and_error(mock_trace):
|
|
tool_call_1 = _make_tool_call("call_success", "success_tool")
|
|
tool_call_2 = _make_tool_call("call_error", "error_tool")
|
|
|
|
with mock.patch(
|
|
"mlflow.genai.judges.tools.registry._judge_tool_registry.invoke"
|
|
) as mock_invoke:
|
|
mock_invoke.side_effect = [{"result": "success"}, RuntimeError("Failed")]
|
|
|
|
result = _process_tool_calls(tool_calls=[tool_call_1, tool_call_2], trace=mock_trace)
|
|
|
|
assert len(result) == 2
|
|
assert result[0].tool_call_id == "call_success"
|
|
assert json.loads(result[0].content) == {"result": "success"}
|
|
assert result[1].tool_call_id == "call_error"
|
|
assert "Error: Failed" in result[1].content
|