mlflow--mlflow
wehub-resource-sync
451 行
14 KiB
Python
451 行
14 KiB
Python
# TODO: ADD MORE TESTS
|
|
import json
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
from fastapi.encoders import jsonable_encoder
|
|
|
|
from mlflow.gateway.config import EndpointConfig
|
|
from mlflow.gateway.providers.openai import OpenAIProvider
|
|
from mlflow.gateway.schemas import chat
|
|
from mlflow.gateway.uc_function_utils import (
|
|
_quote_identifier,
|
|
get_execute_function_sql_stmt,
|
|
uc_type_to_json_schema_type,
|
|
)
|
|
|
|
from tests.gateway.tools import (
|
|
MockAsyncResponse,
|
|
MockHttpClient,
|
|
)
|
|
|
|
|
|
def chat_config():
|
|
return {
|
|
"name": "chat",
|
|
"endpoint_type": "llm/v1/chat",
|
|
"model": {
|
|
"provider": "openai",
|
|
"name": "gpt-4o-mini",
|
|
"config": {
|
|
"openai_api_base": "https://api.openai.com/v1",
|
|
"openai_api_key": "key",
|
|
},
|
|
},
|
|
}
|
|
|
|
|
|
def test_uc_type_to_json_schema_type():
|
|
# TODO: Improve the test coverage
|
|
|
|
# Test with primitive types
|
|
assert uc_type_to_json_schema_type("long") == {"type": "integer"}
|
|
assert uc_type_to_json_schema_type("boolean") == {"type": "boolean"}
|
|
assert uc_type_to_json_schema_type("date") == {"type": "string", "format": "date"}
|
|
|
|
# Test with complex types
|
|
assert uc_type_to_json_schema_type({"type": "array", "elementType": "integer"}) == {
|
|
"type": "array",
|
|
"items": {"type": "integer"},
|
|
}
|
|
assert uc_type_to_json_schema_type({
|
|
"type": "map",
|
|
"keyType": "string",
|
|
"valueType": "integer",
|
|
}) == {"type": "object", "additionalProperties": {"type": "integer"}}
|
|
assert uc_type_to_json_schema_type({
|
|
"type": "struct",
|
|
"fields": [{"name": "field1", "type": "integer"}],
|
|
}) == {"type": "object", "properties": {"field1": {"type": "integer"}}}
|
|
|
|
# Test with unsupported types
|
|
with pytest.raises(TypeError, match=r"Type interval is not supported\."):
|
|
uc_type_to_json_schema_type("interval")
|
|
|
|
with pytest.raises(TypeError, match=r"Unknown type unknown. Try upgrading this package\."):
|
|
uc_type_to_json_schema_type("unknown")
|
|
|
|
with pytest.raises(TypeError, match=r"Only support STRING key type for MAP but got integer\."):
|
|
uc_type_to_json_schema_type({"type": "map", "keyType": "integer", "valueType": "integer"})
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_uc_functions(monkeypatch):
|
|
from databricks.sdk.service.catalog import ColumnTypeName
|
|
from databricks.sdk.service.sql import StatementState
|
|
|
|
monkeypatch.setenv("MLFLOW_ENABLE_UC_FUNCTIONS", "true")
|
|
monkeypatch.setenv("DATABRICKS_WAREHOUSE_ID", "1234")
|
|
|
|
config = chat_config()
|
|
first_resp = {
|
|
"id": "chatcmpl-abc123",
|
|
"object": "chat.completion",
|
|
"created": 1677858242,
|
|
"model": "gpt-4o-mini",
|
|
"usage": {
|
|
"prompt_tokens": 13,
|
|
"completion_tokens": 7,
|
|
"total_tokens": 20,
|
|
},
|
|
"choices": [
|
|
{
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_id",
|
|
"function": {
|
|
"arguments": '{"x": 1, "y": 2}',
|
|
"name": "test__uc__func",
|
|
},
|
|
"type": "function",
|
|
},
|
|
],
|
|
},
|
|
"index": 0,
|
|
}
|
|
],
|
|
"headers": {"Content-Type": "application/json"},
|
|
}
|
|
|
|
second_resp = {
|
|
"id": "chatcmpl-abc123",
|
|
"object": "chat.completion",
|
|
"created": 1677858242,
|
|
"model": "gpt-4o-mini",
|
|
"usage": {
|
|
"prompt_tokens": 13,
|
|
"completion_tokens": 7,
|
|
"total_tokens": 20,
|
|
},
|
|
"choices": [
|
|
{
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": "1 + 2 = 3",
|
|
},
|
|
"finish_reason": "stop",
|
|
"index": 0,
|
|
}
|
|
],
|
|
"headers": {"Content-Type": "application/json"},
|
|
}
|
|
|
|
mock_client_session = MockHttpClient()
|
|
mock_client_session.post.side_effect = [
|
|
MockAsyncResponse(first_resp),
|
|
MockAsyncResponse(second_resp),
|
|
]
|
|
|
|
params = []
|
|
for name, typ in [("x", "integer"), ("y", "integer")]:
|
|
p = mock.MagicMock()
|
|
p.type_json = json.dumps({"type": typ})
|
|
p.name = name
|
|
params.append(p)
|
|
|
|
mock_func_info = mock.MagicMock()
|
|
mock_func_info.input_params.parameters = params
|
|
mock_func_info.data_type = ColumnTypeName.INT
|
|
mock_func_info.name = "func"
|
|
mock_func_info.catalog_name = "test"
|
|
mock_func_info.schema_name = "uc"
|
|
mock_func_info.full_name = "test.uc.func"
|
|
|
|
mock_workspace_client = mock.MagicMock()
|
|
mock_workspace_client.functions.get.return_value = mock_func_info
|
|
|
|
mock_statement_result = mock.MagicMock()
|
|
mock_statement_result.result.data_array = [[3]]
|
|
mock_statement_result.status.state = StatementState.SUCCEEDED
|
|
mock_statement_result.manifest.truncated = "manifest"
|
|
mock_workspace_client.statement_execution.execute_statement.return_value = mock_statement_result
|
|
|
|
with (
|
|
mock.patch("aiohttp.ClientSession", return_value=mock_client_session),
|
|
mock.patch(
|
|
"mlflow.gateway.providers.openai._get_workspace_client",
|
|
return_value=mock_workspace_client,
|
|
) as mock_workspace_client,
|
|
):
|
|
provider = OpenAIProvider(EndpointConfig(**config))
|
|
payload = {
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": "Tell me a joke",
|
|
}
|
|
],
|
|
"temperature": 0.5,
|
|
"tools": [
|
|
{
|
|
"type": "uc_function",
|
|
"uc_function": {
|
|
"name": "test.uc.func",
|
|
},
|
|
}
|
|
],
|
|
}
|
|
response = await provider.chat(chat.RequestPayload(**payload))
|
|
assert mock_client_session.post.call_count == 2
|
|
|
|
assert jsonable_encoder(response) == {
|
|
"id": "chatcmpl-abc123",
|
|
"object": "chat.completion",
|
|
"created": 1677858242,
|
|
"model": "gpt-4o-mini",
|
|
"provider": "openai",
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": r"""
|
|
<uc_function_call>
|
|
{
|
|
"id": "call_id",
|
|
"name": "test.uc.func",
|
|
"arguments": "{\"x\": 1, \"y\": 2}"
|
|
}
|
|
</uc_function_call>
|
|
|
|
<uc_function_result>
|
|
{
|
|
"tool_call_id": "call_id",
|
|
"content": "{\"format\": \"SCALAR\", \"value\": \"3\", \"truncated\": \"manifest\"}"
|
|
}
|
|
</uc_function_result>
|
|
|
|
1 + 2 = 3""".lstrip(),
|
|
"tool_calls": None,
|
|
"refusal": None,
|
|
},
|
|
"finish_reason": "stop",
|
|
}
|
|
],
|
|
# The server makes two requests, the token usage should be doubled
|
|
"usage": {"prompt_tokens": 13 * 2, "completion_tokens": 7 * 2, "total_tokens": 20 * 2},
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_uc_functions_user_defined_functions(monkeypatch):
|
|
from databricks.sdk.service.catalog import ColumnTypeName
|
|
from databricks.sdk.service.sql import StatementState
|
|
|
|
monkeypatch.setenv("MLFLOW_ENABLE_UC_FUNCTIONS", "true")
|
|
monkeypatch.setenv("DATABRICKS_WAREHOUSE_ID", "1234")
|
|
|
|
config = chat_config()
|
|
first_resp = {
|
|
"id": "chatcmpl-abc123",
|
|
"object": "chat.completion",
|
|
"created": 1677858242,
|
|
"model": "gpt-4o-mini",
|
|
"usage": {
|
|
"prompt_tokens": 13,
|
|
"completion_tokens": 7,
|
|
"total_tokens": 20,
|
|
},
|
|
"choices": [
|
|
{
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_id_1",
|
|
"function": {
|
|
"arguments": '{"x": 1, "y": 2}',
|
|
"name": "test__uc__func",
|
|
},
|
|
"type": "function",
|
|
},
|
|
{
|
|
"id": "call_id_2",
|
|
"function": {
|
|
"arguments": '{"x": 3, "y": 4}',
|
|
"name": "multiply",
|
|
},
|
|
"type": "function",
|
|
},
|
|
],
|
|
},
|
|
"index": 0,
|
|
}
|
|
],
|
|
"headers": {"Content-Type": "application/json"},
|
|
}
|
|
|
|
second_resp = {
|
|
"id": "chatcmpl-abc123",
|
|
"object": "chat.completion",
|
|
"created": 1677858242,
|
|
"model": "gpt-4o-mini",
|
|
"usage": {
|
|
"prompt_tokens": 13,
|
|
"completion_tokens": 7,
|
|
"total_tokens": 20,
|
|
},
|
|
"choices": [
|
|
{
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": "1 + 2 = 3",
|
|
},
|
|
"finish_reason": "stop",
|
|
"index": 0,
|
|
}
|
|
],
|
|
"headers": {"Content-Type": "application/json"},
|
|
}
|
|
|
|
mock_client_session = MockHttpClient()
|
|
mock_client_session.post.side_effect = [
|
|
MockAsyncResponse(first_resp),
|
|
MockAsyncResponse(second_resp),
|
|
]
|
|
|
|
params = []
|
|
for name, typ in [("x", "integer"), ("y", "integer")]:
|
|
p = mock.MagicMock()
|
|
p.type_json = json.dumps({"type": typ})
|
|
p.name = name
|
|
params.append(p)
|
|
|
|
mock_func_info = mock.MagicMock()
|
|
mock_func_info.input_params.parameters = params
|
|
mock_func_info.data_type = ColumnTypeName.INT
|
|
mock_func_info.name = "func"
|
|
mock_func_info.catalog_name = "test"
|
|
mock_func_info.schema_name = "uc"
|
|
mock_func_info.full_name = "test.uc.func"
|
|
|
|
mock_workspace_client = mock.MagicMock()
|
|
mock_workspace_client.functions.get.return_value = mock_func_info
|
|
|
|
mock_statement_result = mock.MagicMock()
|
|
mock_statement_result.result.data_array = [[3]]
|
|
mock_statement_result.status.state = StatementState.SUCCEEDED
|
|
mock_statement_result.manifest.truncated = "manifest"
|
|
mock_workspace_client.statement_execution.execute_statement.return_value = mock_statement_result
|
|
|
|
with (
|
|
mock.patch("aiohttp.ClientSession", return_value=mock_client_session),
|
|
mock.patch(
|
|
"mlflow.gateway.providers.openai._get_workspace_client",
|
|
return_value=mock_workspace_client,
|
|
) as mock_workspace_client,
|
|
):
|
|
provider = OpenAIProvider(EndpointConfig(**config))
|
|
payload = {
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": "What is 1 + 2? What is 3 * 4?",
|
|
}
|
|
],
|
|
"temperature": 0.5,
|
|
"tools": [
|
|
{
|
|
"type": "uc_function",
|
|
"uc_function": {
|
|
"name": "test.uc.func",
|
|
},
|
|
},
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"description": "Multiply numbers",
|
|
"name": "multiply",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"x": {
|
|
"type": "integer",
|
|
"description": "First number",
|
|
},
|
|
"y": {
|
|
"type": "integer",
|
|
"description": "Second number",
|
|
},
|
|
},
|
|
"required": ["x", "y"],
|
|
},
|
|
},
|
|
},
|
|
],
|
|
}
|
|
response = await provider.chat(chat.RequestPayload(**payload))
|
|
assert mock_client_session.post.call_count == 1
|
|
|
|
assert jsonable_encoder(response) == {
|
|
"id": "chatcmpl-abc123",
|
|
"object": "chat.completion",
|
|
"created": 1677858242,
|
|
"model": "gpt-4o-mini",
|
|
"provider": "openai",
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": r"""
|
|
<uc_function_call>
|
|
{
|
|
"id": "call_id_1",
|
|
"name": "test.uc.func",
|
|
"arguments": "{\"x\": 1, \"y\": 2}"
|
|
}
|
|
</uc_function_call>
|
|
|
|
<uc_function_result>
|
|
{
|
|
"tool_call_id": "call_id_1",
|
|
"content": "{\"format\": \"SCALAR\", \"value\": \"3\", \"truncated\": \"manifest\"}"
|
|
}
|
|
</uc_function_result>""".lstrip(),
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_id_2",
|
|
"type": "function",
|
|
"function": {
|
|
"arguments": '{"x": 3, "y": 4}',
|
|
"name": "multiply",
|
|
},
|
|
},
|
|
],
|
|
"refusal": None,
|
|
},
|
|
"finish_reason": None,
|
|
}
|
|
],
|
|
"usage": {"prompt_tokens": 13, "completion_tokens": 7, "total_tokens": 20},
|
|
}
|
|
|
|
|
|
def test_quote_identifier():
|
|
assert _quote_identifier("catalog.schema.func") == "`catalog`.`schema`.`func`"
|
|
# Strips existing backticks
|
|
assert _quote_identifier("`catalog`.`schema`.`func`") == "`catalog`.`schema`.`func`"
|
|
# Handles special characters
|
|
assert _quote_identifier("func;DROP TABLE") == "`func;DROP TABLE`"
|
|
# Rejects embedded backticks
|
|
with pytest.raises(ValueError, match="Backticks are not allowed"):
|
|
_quote_identifier("func`tion")
|
|
|
|
|
|
def test_get_execute_function_sql_stmt_quotes_function_name():
|
|
from databricks.sdk.service.catalog import ColumnTypeName, FunctionInfo
|
|
|
|
function = mock.Mock(spec=FunctionInfo)
|
|
function.full_name = "catalog.schema.func); DROP TABLE users; --"
|
|
function.data_type = ColumnTypeName.INT
|
|
function.input_params = None
|
|
|
|
result = get_execute_function_sql_stmt(function, {})
|
|
|
|
assert result.statement == "SELECT `catalog`.`schema`.`func); DROP TABLE users; --`()"
|