fix(llm-transformer): remove unwanted template params from guardrail policies#1689
fix(llm-transformer): remove unwanted template params from guardrail policies#1689
Conversation
…policies When policies like word-count-guardrail are attached to LLM providers, LLM template extraction parameters (promptTokens, completionTokens, totalTokens, remainingTokens, requestModel, responseModel) were being unconditionally merged into every policy's parameters. These template-level extraction fields are not needed by user-defined policies and were appearing as unwanted parameters in the policy engine config dump. This change ensures that only user-provided parameters are assigned to policies, preventing leakage of LLM template extraction fields. Fixes wso2#1592
SummaryThis pull request refines parameter handling in LLM provider transformation to prevent template-level extraction fields from being merged into user-defined policy parameters. Changesgateway/gateway-controller/pkg/utils/llm_transformer.go (+16/-15 lines)
gateway/gateway-controller/pkg/utils/llm_transformer_test.go (+6/-6 lines)
OutcomeWhen LLM providers are deployed with attached guardrail policies, the configuration now contains only the parameters explicitly defined by users. Template-level extraction parameters that were previously appearing in policy engine configuration dumps (such as promptTokens, completionTokens, totalTokens, remainingTokens, requestModel, responseModel) are no longer included in the policy parameter sets, resulting in cleaner and more predictable policy configurations. WalkthroughTemplate-derived extraction parameters are no longer merged into policy parameters. The transformer now constructs policy parameters solely from user-specified Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 golangci-lint (2.11.4)level=error msg="[linters_context] typechecking error: pattern ./...: directory prefix . does not contain modules listed in go.work or their selected dependencies" Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Pull request overview
Fixes LLM provider/proxy policy transformation so that LLM template extraction fields (e.g., token/model identifiers) are no longer merged into user-defined policy parameters (notably guardrail policies), reducing config noise and aligning parameters with the policy schema.
Changes:
- Replaced
mergeParams(pathEntry.Params, templateParams)withcopyParams(pathEntry.Params)when attaching policies to operations (proxy + provider, AllowAll + DenyAll flows). - Added
copyParams()helper to clone user-provided params without injecting template-derived fields.
| // copyParams creates a shallow copy of the parameters map and returns a pointer to it. | ||
| // Returns nil if the input map is nil or empty. | ||
| func copyParams(params map[string]interface{}) *map[string]interface{} { | ||
| if len(params) == 0 { | ||
| return nil | ||
| } | ||
| copied := make(map[string]interface{}, len(params)) | ||
| for k, v := range params { | ||
| copied[k] = v | ||
| } | ||
| return &copied | ||
| } |
There was a problem hiding this comment.
New helper copyParams and the behavioral change (policies should no longer receive template extraction keys like requestModel/promptTokens/etc.) aren’t covered by unit tests. Since this package already has tests for related helpers (buildTemplateParams/mergeParams), add tests for copyParams and a transformer-level test asserting that attached policies only include user-provided params (and not template extraction fields).
There was a problem hiding this comment.
🧹 Nitpick comments (1)
gateway/gateway-controller/pkg/utils/llm_transformer_test.go (1)
984-1006: Optional: strengthen assertions to catch any future template-param leak.The current checks only verify
requestModelis absent. Since the fix is meant to exclude all template extraction fields (responseModel,promptTokens,completionTokens,totalTokens,remainingTokens,requestModel), consider asserting the params map length equals 1 (or explicitly checking each excluded key). That way, if a future regression re-introduces any template field — not justrequestModel— this test will catch it.♻️ Suggested tightening
- // User-defined policies should only contain user params, not template params - _, hasRequestModel := (*responsesPolicy.Params)["requestModel"] - assert.False(t, hasRequestModel, "template params should not be merged into user-defined policies") - assert.Equal(t, "value", (*responsesPolicy.Params)["userParam"]) + // User-defined policies should only contain user params, not template params + assert.Len(t, *responsesPolicy.Params, 1, "only user-provided params should be present") + assert.Equal(t, "value", (*responsesPolicy.Params)["userParam"]) + for _, k := range []string{"requestModel", "responseModel", "promptTokens", "completionTokens", "totalTokens", "remainingTokens"} { + _, present := (*responsesPolicy.Params)[k] + assert.Falsef(t, present, "template param %q should not be merged into user-defined policies", k) + }Apply the same to the
wildcardPolicyblock.Also note: there is no dedicated unit test for the new
copyParamshelper introduced inllm_transformer.go. A small test covering nil input, empty map, and shallow-copy semantics (mutating the copy does not affect the original) would be worth adding.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@gateway/gateway-controller/pkg/utils/llm_transformer_test.go` around lines 984 - 1006, Strengthen the test assertions to ensure no template extraction fields leak by asserting the params map contains only the expected user key: replace the current single-key absence check for "requestModel" with an explicit size check (assert.Len == 1) and/or explicit absence checks for all template keys (responseModel, promptTokens, completionTokens, totalTokens, remainingTokens, requestModel) against responsesPolicy.Params and wildcardPolicy.Params (keep the existing assert.Equal on "userParam"); additionally add a small unit test for the new copyParams helper in llm_transformer.go that verifies behavior for nil input, empty map, and shallow-copy semantics (mutating the returned map does not change the original).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@gateway/gateway-controller/pkg/utils/llm_transformer_test.go`:
- Around line 984-1006: Strengthen the test assertions to ensure no template
extraction fields leak by asserting the params map contains only the expected
user key: replace the current single-key absence check for "requestModel" with
an explicit size check (assert.Len == 1) and/or explicit absence checks for all
template keys (responseModel, promptTokens, completionTokens, totalTokens,
remainingTokens, requestModel) against responsesPolicy.Params and
wildcardPolicy.Params (keep the existing assert.Equal on "userParam");
additionally add a small unit test for the new copyParams helper in
llm_transformer.go that verifies behavior for nil input, empty map, and
shallow-copy semantics (mutating the returned map does not change the original).
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 5578f94f-ae2a-4f1b-bf42-fd93020f64e0
📒 Files selected for processing (1)
gateway/gateway-controller/pkg/utils/llm_transformer_test.go
Purpose
Fixes unwanted LLM template extraction parameters being attached to guardrail policies when deployed with LLM providers. When policies like
word-count-guardrailare attached to an LLM provider, the policyengine config dump shows extraneous parameters (
promptTokens,completionTokens,totalTokens,remainingTokens,requestModel,responseModel) that are not defined in the policy schemaand should not be present.
Approach
mergeParams(pathEntry.Params, templateParams)withcopyParams(pathEntry.Params)in three locations withinllm_transformer.go(transformProxy and transformProvider methods for both AllowAlland DenyAll modes)
copyParams()helper function that creates a shallow copy of user-provided parameters without merging LLM template extraction fieldsbuildTemplateParams()) are preserved for other uses but no longer injected into user policy parametersTesting
word-count-guardrailattached/config_dumpendpoint that only user-provided params (request,response) appearFixes #1592