mirror of
https://github.com/alibaba/higress.git
synced 2026-03-01 15:10:56 +08:00
add model-mapper plugin & optimize model-router plugin (#1538)
This commit is contained in:
澄潭
GitHub
301
plugins/wasm-cpp/extensions/model_mapper/plugin_test.cc
Normal file
301
plugins/wasm-cpp/extensions/model_mapper/plugin_test.cc
Normal file
@@ -0,0 +1,301 @@
|
||||
// Copyright (c) 2022 Alibaba Group Holding Ltd.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include "extensions/model_mapper/plugin.h"
|
||||
|
||||
#include <cstddef>
|
||||
|
||||
#include "gmock/gmock.h"
|
||||
#include "gtest/gtest.h"
|
||||
#include "include/proxy-wasm/context.h"
|
||||
#include "include/proxy-wasm/null.h"
|
||||
|
||||
namespace proxy_wasm {
|
||||
namespace null_plugin {
|
||||
namespace model_mapper {
|
||||
|
||||
NullPluginRegistry* context_registry_;
|
||||
RegisterNullVmPluginFactory register_model_mapper_plugin("model_mapper", []() {
|
||||
return std::make_unique<NullPlugin>(model_mapper::context_registry_);
|
||||
});
|
||||
|
||||
class MockContext : public proxy_wasm::ContextBase {
|
||||
public:
|
||||
MockContext(WasmBase* wasm) : ContextBase(wasm) {}
|
||||
MOCK_METHOD(BufferInterface*, getBuffer, (WasmBufferType));
|
||||
MOCK_METHOD(WasmResult, log, (uint32_t, std::string_view));
|
||||
MOCK_METHOD(WasmResult, setBuffer,
|
||||
(WasmBufferType, size_t, size_t, std::string_view));
|
||||
MOCK_METHOD(WasmResult, getHeaderMapValue,
|
||||
(WasmHeaderMapType /* type */, std::string_view /* key */,
|
||||
std::string_view* /*result */));
|
||||
MOCK_METHOD(WasmResult, replaceHeaderMapValue,
|
||||
(WasmHeaderMapType /* type */, std::string_view /* key */,
|
||||
std::string_view /* value */));
|
||||
MOCK_METHOD(WasmResult, removeHeaderMapValue,
|
||||
(WasmHeaderMapType /* type */, std::string_view /* key */));
|
||||
MOCK_METHOD(WasmResult, addHeaderMapValue,
|
||||
(WasmHeaderMapType, std::string_view, std::string_view));
|
||||
MOCK_METHOD(WasmResult, getProperty, (std::string_view, std::string*));
|
||||
MOCK_METHOD(WasmResult, setProperty, (std::string_view, std::string_view));
|
||||
};
|
||||
class ModelMapperTest : public ::testing::Test {
|
||||
protected:
|
||||
ModelMapperTest() {
|
||||
// Initialize test VM
|
||||
test_vm_ = createNullVm();
|
||||
wasm_base_ = std::make_unique<WasmBase>(
|
||||
std::move(test_vm_), "test-vm", "", "",
|
||||
std::unordered_map<std::string, std::string>{},
|
||||
AllowedCapabilitiesMap{});
|
||||
wasm_base_->load("model_mapper");
|
||||
wasm_base_->initialize();
|
||||
// Initialize host side context
|
||||
mock_context_ = std::make_unique<MockContext>(wasm_base_.get());
|
||||
current_context_ = mock_context_.get();
|
||||
// Initialize Wasm sandbox context
|
||||
root_context_ = std::make_unique<PluginRootContext>(0, "");
|
||||
context_ = std::make_unique<PluginContext>(1, root_context_.get());
|
||||
|
||||
ON_CALL(*mock_context_, log(testing::_, testing::_))
|
||||
.WillByDefault([](uint32_t, std::string_view m) {
|
||||
std::cerr << m << "\n";
|
||||
return WasmResult::Ok;
|
||||
});
|
||||
|
||||
ON_CALL(*mock_context_, getBuffer(testing::_))
|
||||
.WillByDefault([&](WasmBufferType type) {
|
||||
if (type == WasmBufferType::HttpRequestBody) {
|
||||
return &body_;
|
||||
}
|
||||
return &config_;
|
||||
});
|
||||
ON_CALL(*mock_context_, getHeaderMapValue(WasmHeaderMapType::RequestHeaders,
|
||||
testing::_, testing::_))
|
||||
.WillByDefault([&](WasmHeaderMapType, std::string_view header,
|
||||
std::string_view* result) {
|
||||
if (header == "content-type") {
|
||||
*result = "application/json";
|
||||
} else if (header == "content-length") {
|
||||
*result = "1024";
|
||||
} else if (header == ":path") {
|
||||
*result = path_;
|
||||
}
|
||||
return WasmResult::Ok;
|
||||
});
|
||||
ON_CALL(*mock_context_,
|
||||
replaceHeaderMapValue(WasmHeaderMapType::RequestHeaders, testing::_,
|
||||
testing::_))
|
||||
.WillByDefault([&](WasmHeaderMapType, std::string_view key,
|
||||
std::string_view value) { return WasmResult::Ok; });
|
||||
ON_CALL(*mock_context_,
|
||||
removeHeaderMapValue(WasmHeaderMapType::RequestHeaders, testing::_))
|
||||
.WillByDefault([&](WasmHeaderMapType, std::string_view key) {
|
||||
return WasmResult::Ok;
|
||||
});
|
||||
ON_CALL(*mock_context_, addHeaderMapValue(WasmHeaderMapType::RequestHeaders,
|
||||
testing::_, testing::_))
|
||||
.WillByDefault([&](WasmHeaderMapType, std::string_view header,
|
||||
std::string_view value) { return WasmResult::Ok; });
|
||||
ON_CALL(*mock_context_, getProperty(testing::_, testing::_))
|
||||
.WillByDefault([&](std::string_view path, std::string* result) {
|
||||
if (absl::StartsWith(path, "route_name")) {
|
||||
*result = route_name_;
|
||||
} else if (absl::StartsWith(path, "cluster_name")) {
|
||||
*result = service_name_;
|
||||
}
|
||||
return WasmResult::Ok;
|
||||
});
|
||||
ON_CALL(*mock_context_, setProperty(testing::_, testing::_))
|
||||
.WillByDefault(
|
||||
[&](std::string_view, std::string_view) { return WasmResult::Ok; });
|
||||
}
|
||||
~ModelMapperTest() override {}
|
||||
std::unique_ptr<WasmBase> wasm_base_;
|
||||
std::unique_ptr<WasmVm> test_vm_;
|
||||
std::unique_ptr<MockContext> mock_context_;
|
||||
std::unique_ptr<PluginRootContext> root_context_;
|
||||
std::unique_ptr<PluginContext> context_;
|
||||
std::string route_name_;
|
||||
std::string service_name_;
|
||||
std::string path_;
|
||||
BufferBase body_;
|
||||
BufferBase config_;
|
||||
};
|
||||
|
||||
TEST_F(ModelMapperTest, ModelMappingTest) {
|
||||
std::string configuration = R"(
|
||||
{
|
||||
"modelMapping": {
|
||||
"*": "qwen-long",
|
||||
"gpt-4*": "qwen-max",
|
||||
"gpt-4o": "qwen-turbo",
|
||||
"gpt-4o-mini": "qwen-plus",
|
||||
"text-embedding-v1": ""
|
||||
}
|
||||
})";
|
||||
|
||||
config_.set(configuration);
|
||||
EXPECT_TRUE(root_context_->configure(configuration.size()));
|
||||
|
||||
path_ = "/v1/chat/completions";
|
||||
std::string request_json = R"({"model": "gpt-3.5"})";
|
||||
EXPECT_CALL(*mock_context_,
|
||||
setBuffer(testing::_, testing::_, testing::_, testing::_))
|
||||
.WillOnce([&](WasmBufferType, size_t, size_t, std::string_view body) {
|
||||
EXPECT_EQ(body, R"({"model":"qwen-long"})");
|
||||
return WasmResult::Ok;
|
||||
});
|
||||
|
||||
body_.set(request_json);
|
||||
EXPECT_EQ(context_->onRequestHeaders(0, false),
|
||||
FilterHeadersStatus::StopIteration);
|
||||
EXPECT_EQ(context_->onRequestBody(20, true), FilterDataStatus::Continue);
|
||||
|
||||
request_json = R"({"model": "gpt-4"})";
|
||||
EXPECT_CALL(*mock_context_,
|
||||
setBuffer(testing::_, testing::_, testing::_, testing::_))
|
||||
.WillOnce([&](WasmBufferType, size_t, size_t, std::string_view body) {
|
||||
EXPECT_EQ(body, R"({"model":"qwen-max"})");
|
||||
return WasmResult::Ok;
|
||||
});
|
||||
|
||||
body_.set(request_json);
|
||||
EXPECT_EQ(context_->onRequestHeaders(0, false),
|
||||
FilterHeadersStatus::StopIteration);
|
||||
EXPECT_EQ(context_->onRequestBody(20, true), FilterDataStatus::Continue);
|
||||
|
||||
request_json = R"({"model": "gpt-4o"})";
|
||||
EXPECT_CALL(*mock_context_,
|
||||
setBuffer(testing::_, testing::_, testing::_, testing::_))
|
||||
.WillOnce([&](WasmBufferType, size_t, size_t, std::string_view body) {
|
||||
EXPECT_EQ(body, R"({"model":"qwen-turbo"})");
|
||||
return WasmResult::Ok;
|
||||
});
|
||||
|
||||
body_.set(request_json);
|
||||
EXPECT_EQ(context_->onRequestHeaders(0, false),
|
||||
FilterHeadersStatus::StopIteration);
|
||||
EXPECT_EQ(context_->onRequestBody(20, true), FilterDataStatus::Continue);
|
||||
|
||||
request_json = R"({"model": "gpt-4o-mini"})";
|
||||
EXPECT_CALL(*mock_context_,
|
||||
setBuffer(testing::_, testing::_, testing::_, testing::_))
|
||||
.WillOnce([&](WasmBufferType, size_t, size_t, std::string_view body) {
|
||||
EXPECT_EQ(body, R"({"model":"qwen-plus"})");
|
||||
return WasmResult::Ok;
|
||||
});
|
||||
|
||||
body_.set(request_json);
|
||||
EXPECT_EQ(context_->onRequestHeaders(0, false),
|
||||
FilterHeadersStatus::StopIteration);
|
||||
EXPECT_EQ(context_->onRequestBody(20, true), FilterDataStatus::Continue);
|
||||
|
||||
request_json = R"({"model": "text-embedding-v1"})";
|
||||
EXPECT_CALL(*mock_context_,
|
||||
setBuffer(testing::_, testing::_, testing::_, testing::_))
|
||||
.Times(0);
|
||||
|
||||
body_.set(request_json);
|
||||
EXPECT_EQ(context_->onRequestHeaders(0, false),
|
||||
FilterHeadersStatus::StopIteration);
|
||||
EXPECT_EQ(context_->onRequestBody(20, true), FilterDataStatus::Continue);
|
||||
|
||||
request_json = R"({})";
|
||||
EXPECT_CALL(*mock_context_,
|
||||
setBuffer(testing::_, testing::_, testing::_, testing::_))
|
||||
.WillOnce([&](WasmBufferType, size_t, size_t, std::string_view body) {
|
||||
EXPECT_EQ(body, R"({"model":"qwen-long"})");
|
||||
return WasmResult::Ok;
|
||||
});
|
||||
|
||||
body_.set(request_json);
|
||||
EXPECT_EQ(context_->onRequestHeaders(0, false),
|
||||
FilterHeadersStatus::StopIteration);
|
||||
EXPECT_EQ(context_->onRequestBody(20, true), FilterDataStatus::Continue);
|
||||
}
|
||||
|
||||
TEST_F(ModelMapperTest, RouteLevelModelMappingTest) {
|
||||
std::string configuration = R"(
|
||||
{
|
||||
"_rules_": [
|
||||
{
|
||||
"_match_route_": ["route-a"],
|
||||
"_match_service_": ["service-1"],
|
||||
"modelMapping": {
|
||||
"*": "qwen-long"
|
||||
}
|
||||
},
|
||||
{
|
||||
"_match_route_": ["route-b"],
|
||||
"_match_service_": ["service-2"],
|
||||
"modelMapping": {
|
||||
"*": "qwen-max"
|
||||
}
|
||||
},
|
||||
{
|
||||
"_match_route_": ["route-b"],
|
||||
"_match_service_": ["service-3"],
|
||||
"modelMapping": {
|
||||
"*": "qwen-turbo"
|
||||
}
|
||||
}
|
||||
]})";
|
||||
|
||||
config_.set(configuration);
|
||||
EXPECT_TRUE(root_context_->configure(configuration.size()));
|
||||
|
||||
path_ = "/api/v1/chat/completions";
|
||||
std::string request_json = R"({"model": "gpt-4"})";
|
||||
body_.set(request_json);
|
||||
route_name_ = "route-a";
|
||||
service_name_ = "outbound|80||service-1";
|
||||
EXPECT_CALL(*mock_context_,
|
||||
setBuffer(testing::_, testing::_, testing::_, testing::_))
|
||||
.WillOnce([&](WasmBufferType, size_t, size_t, std::string_view body) {
|
||||
EXPECT_EQ(body, R"({"model":"qwen-long"})");
|
||||
return WasmResult::Ok;
|
||||
});
|
||||
EXPECT_EQ(context_->onRequestHeaders(0, false),
|
||||
FilterHeadersStatus::StopIteration);
|
||||
EXPECT_EQ(context_->onRequestBody(20, true), FilterDataStatus::Continue);
|
||||
|
||||
route_name_ = "route-b";
|
||||
service_name_ = "outbound|80||service-2";
|
||||
EXPECT_CALL(*mock_context_,
|
||||
setBuffer(testing::_, testing::_, testing::_, testing::_))
|
||||
.WillOnce([&](WasmBufferType, size_t, size_t, std::string_view body) {
|
||||
EXPECT_EQ(body, R"({"model":"qwen-max"})");
|
||||
return WasmResult::Ok;
|
||||
});
|
||||
EXPECT_EQ(context_->onRequestHeaders(0, false),
|
||||
FilterHeadersStatus::StopIteration);
|
||||
EXPECT_EQ(context_->onRequestBody(20, true), FilterDataStatus::Continue);
|
||||
|
||||
route_name_ = "route-b";
|
||||
service_name_ = "outbound|80||service-3";
|
||||
EXPECT_CALL(*mock_context_,
|
||||
setBuffer(testing::_, testing::_, testing::_, testing::_))
|
||||
.WillOnce([&](WasmBufferType, size_t, size_t, std::string_view body) {
|
||||
EXPECT_EQ(body, R"({"model":"qwen-turbo"})");
|
||||
return WasmResult::Ok;
|
||||
});
|
||||
EXPECT_EQ(context_->onRequestHeaders(0, false),
|
||||
FilterHeadersStatus::StopIteration);
|
||||
EXPECT_EQ(context_->onRequestBody(20, true), FilterDataStatus::Continue);
|
||||
}
|
||||
|
||||
} // namespace model_mapper
|
||||
} // namespace null_plugin
|
||||
} // namespace proxy_wasm
|
||||
Reference in New Issue
Block a user