wehub-resource-sync
108 lines
3.4 KiB
Python
108 lines
3.4 KiB
Python
from unittest import mock
|
|
|
|
import pytest
|
|
|
|
from mlflow.genai.simulators.utils import (
|
|
format_history,
|
|
get_default_simulation_model,
|
|
invoke_model_without_tracing,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("history", "expected"),
|
|
[
|
|
([], None),
|
|
([{"role": "user", "content": "Hello"}], "user: Hello"),
|
|
(
|
|
[
|
|
{"role": "user", "content": "Hello"},
|
|
{"role": "assistant", "content": "Hi there!"},
|
|
{"role": "user", "content": "How are you?"},
|
|
],
|
|
"user: Hello\nassistant: Hi there!\nuser: How are you?",
|
|
),
|
|
([{"content": "Hello"}], "unknown: Hello"),
|
|
([{"role": "user"}], "user: "),
|
|
([{"role": None, "content": None}], "unknown: "),
|
|
],
|
|
)
|
|
def test_format_history(history, expected):
|
|
assert format_history(history) == expected
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"model_uri",
|
|
[
|
|
"openai:/gpt-4o-mini",
|
|
"anthropic:/claude-3-haiku",
|
|
],
|
|
)
|
|
def test_invoke_model_without_tracing_with_provider(model_uri):
|
|
from mlflow.types.llm import ChatMessage
|
|
|
|
messages = [ChatMessage(role="user", content="Hello")]
|
|
|
|
with mock.patch(
|
|
"mlflow.genai.scorers.llm_backend.ScorerLLMClient.complete", return_value="Hi there!"
|
|
) as mock_invoke:
|
|
result = invoke_model_without_tracing(model_uri=model_uri, messages=messages)
|
|
|
|
assert result == "Hi there!"
|
|
mock_invoke.assert_called_once()
|
|
|
|
|
|
def test_invoke_model_without_tracing_with_inference_params():
|
|
from mlflow.types.llm import ChatMessage
|
|
|
|
messages = [ChatMessage(role="user", content="Hello")]
|
|
|
|
with mock.patch(
|
|
"mlflow.genai.scorers.llm_backend.ScorerLLMClient.complete", return_value="Response"
|
|
) as mock_invoke:
|
|
invoke_model_without_tracing(
|
|
model_uri="openai:/gpt-4o-mini",
|
|
messages=messages,
|
|
inference_params={"temperature": 0.5},
|
|
)
|
|
|
|
mock_invoke.assert_called_once_with(
|
|
[{"role": "user", "content": "Hello"}],
|
|
response_format=None,
|
|
num_retries=3,
|
|
temperature=0.5,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("model_uri", ["databricks", "gpt-oss-120b"])
|
|
def test_invoke_model_without_tracing_with_databricks(model_uri):
|
|
from mlflow.types.llm import ChatMessage
|
|
|
|
messages = [ChatMessage(role="user", content="Hello")]
|
|
|
|
with (
|
|
mock.patch("mlflow.genai.simulators.utils.call_chat_completions") as mock_call,
|
|
mock.patch(
|
|
"mlflow.genai.simulators.utils._create_message_from_databricks_response"
|
|
) as mock_create,
|
|
):
|
|
mock_call.return_value = mock.MagicMock(error_code=None, output_json='{"content": "Hi"}')
|
|
mock_create.return_value = mock.MagicMock(content="Hi from Databricks")
|
|
|
|
result = invoke_model_without_tracing(model_uri=model_uri, messages=messages)
|
|
|
|
assert result == "Hi from Databricks"
|
|
mock_call.assert_called_once()
|
|
|
|
|
|
def test_get_default_simulation_model_non_databricks():
|
|
with mock.patch("mlflow.genai.simulators.utils.is_databricks_uri", return_value=False):
|
|
model = get_default_simulation_model()
|
|
assert model == "openai:/gpt-5"
|
|
|
|
|
|
def test_get_default_simulation_model_databricks():
|
|
with mock.patch("mlflow.genai.simulators.utils.is_databricks_uri", return_value=True):
|
|
model = get_default_simulation_model()
|
|
assert model == "gpt-oss-120b"
|