jiazhizhong/higress
1
0
Fork 0
mirror of https://github.com/alibaba/higress.git synced 2026-06-10 05:07:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix model-router plugin (#1432)

Browse Source
This commit is contained in:
澄潭
2024-10-25 16:48:37 +08:00
committed by GitHub
parent 496d365a95
commit d309bf2e25
2 changed files with 53 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
plugins/wasm-cpp/extensions/model_router/plugin.cc
View File

@@ -101,7 +101,7 @@ bool PluginRootContext::configure(size_t configuration_size) {
configuration_data->view())); configuration_data->view()));
return false; return false;
} }
if (!parseAuthRuleConfig(result.value())) { if (!parseRuleConfig(result.value())) {
LOG_WARN(absl::StrCat("cannot parse plugin configuration JSON string: ", LOG_WARN(absl::StrCat("cannot parse plugin configuration JSON string: ",
configuration_data->view())); configuration_data->view()));
return false; return false;

56
plugins/wasm-cpp/extensions/model_router/plugin_test.cc
View File

@@ -40,6 +40,11 @@ class MockContext : public proxy_wasm::ContextBase {
MOCK_METHOD(WasmResult, getHeaderMapValue, MOCK_METHOD(WasmResult, getHeaderMapValue,
(WasmHeaderMapType /* type */, std::string_view /* key */, (WasmHeaderMapType /* type */, std::string_view /* key */,
std::string_view* /*result */)); 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, MOCK_METHOD(WasmResult, addHeaderMapValue,
(WasmHeaderMapType, std::string_view, std::string_view)); (WasmHeaderMapType, std::string_view, std::string_view));
MOCK_METHOD(WasmResult, getProperty, (std::string_view, std::string*)); MOCK_METHOD(WasmResult, getProperty, (std::string_view, std::string*));
@@ -87,6 +92,16 @@ class ModelRouterTest : public ::testing::Test {
} }
return WasmResult::Ok; 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, ON_CALL(*mock_context_, addHeaderMapValue(WasmHeaderMapType::RequestHeaders,
testing::_, testing::_)) testing::_, testing::_))
.WillByDefault([&](WasmHeaderMapType, std::string_view header, .WillByDefault([&](WasmHeaderMapType, std::string_view header,
@@ -128,10 +143,10 @@ TEST_F(ModelRouterTest, RewriteModelAndHeader) {
return WasmResult::Ok; return WasmResult::Ok;
}); });
EXPECT_CALL( EXPECT_CALL(*mock_context_,
*mock_context_, replaceHeaderMapValue(testing::_,
addHeaderMapValue(testing::_, std::string_view("x-higress-llm-provider"), std::string_view("x-higress-llm-provider"),
std::string_view("qwen"))); std::string_view("qwen")));
body_.set(request_json); body_.set(request_json);
EXPECT_EQ(context_->onRequestHeaders(0, false), EXPECT_EQ(context_->onRequestHeaders(0, false),
@@ -139,6 +154,39 @@ TEST_F(ModelRouterTest, RewriteModelAndHeader) {
EXPECT_EQ(context_->onRequestBody(28, true), FilterDataStatus::Continue); EXPECT_EQ(context_->onRequestBody(28, true), FilterDataStatus::Continue);
} }
TEST_F(ModelRouterTest, RouteLevelRewriteModelAndHeader) {
std::string configuration = R"(
{
"_rules_": [
{
"_match_route_": ["route-a"],
"enable": true
}
]})";
config_.set(configuration);
EXPECT_TRUE(root_context_->configure(configuration.size()));
std::string request_json = R"({"model": "qwen/qwen-long"})";
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_CALL(*mock_context_,
replaceHeaderMapValue(testing::_,
std::string_view("x-higress-llm-provider"),
std::string_view("qwen")));
body_.set(request_json);
route_name_ = "route-a";
EXPECT_EQ(context_->onRequestHeaders(0, false),
FilterHeadersStatus::StopIteration);
EXPECT_EQ(context_->onRequestBody(28, true), FilterDataStatus::Continue);
}
} // namespace model_router } // namespace model_router
} // namespace null_plugin } // namespace null_plugin
} // namespace proxy_wasm } // namespace proxy_wasm