Add semantic tool filtering Policy to WSO2 API Platform #104
Add semantic tool filtering Policy to WSO2 API Platform #104RakhithaRR merged 12 commits intowso2:mainfrom
Conversation
WalkthroughAdds a new "Semantic Tool Filtering" policy with docs and metadata. The policy extracts a user query and tool definitions (JSONPath or tagged text), generates embeddings, scores tools by cosine similarity, filters by rank or threshold, updates the request body, and provides an in-memory LRU embedding cache. Changes
Sequence DiagramsequenceDiagram
participant Client
participant Policy as SemanticToolFiltering
participant Cache as EmbeddingCacheStore
participant Embedder as EmbeddingProvider
participant Upstream
Client->>Policy: HTTP request (user query + tools)
Policy->>Policy: Extract query & tools (JSONPath or tags)
Policy->>Embedder: Request query embedding
Embedder-->>Policy: Query embedding
loop for each tool
Policy->>Cache: Lookup embedding (apiId + desc hash)
alt cached
Cache-->>Policy: Return embedding
else not cached
Policy->>Embedder: Request tool embedding
Embedder-->>Policy: Tool embedding
Policy->>Cache: Add embedding (LRU-managed)
end
Policy->>Policy: Compute cosine similarity (query, tool)
end
Policy->>Policy: Select tools (By Rank or By Threshold)
Policy->>Policy: Update request body (JSONPath or tags)
Policy-->>Upstream: Forward modified request
Upstream-->>Client: Response
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
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.
Actionable comments posted: 6
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@docs/semantic-tool-filtering/v0.1/docs/semantic-tool-filtering.md`:
- Around line 38-55: The docs currently claim embeddingModel is required and
advertise embedding_provider_dimension, but parseEmbeddingConfig and
policies/semantic-tool-filtering/policy-definition.yaml treat embeddingModel as
optional for AZURE_OPENAI and do not use embedding_provider_dimension; update
the parameter table to mark embeddingModel as optional when embeddingProvider is
AZURE_OPENAI (or indicate conditional requirement), and remove or comment out
embedding_provider_dimension from the sample config (or add a note that it is
unused) so the table and sample match the actual inputs parsed by
parseEmbeddingConfig and the validations in policy-definition.yaml.
In `@policies/semantic-tool-filtering/policy-definition.yaml`:
- Around line 44-49: The policy schema is too permissive compared to
semantictoolfiltering.go: tighten validation so only simple dotted JSON paths
are allowed for toolsJSONPath and require non-empty embeddingEndpoint and
apiKey; update the policy-definition.yaml entries (including the other fields
mentioned around lines 84-95) to add a pattern for toolsJSONPath that matches
simple dot-separated identifiers (no brackets, wildcards or JSONPath operators)
and add minLength: 1 (or a non-empty string pattern) for embeddingEndpoint and
apiKey so invalid configs are rejected at schema validation rather than during
policy initialization.
In `@policies/semantic-tool-filtering/semantictoolfiltering.go`:
- Around line 302-323: The validation error messages use capitalized names but
the parameters were renamed to lowercase; update the error strings in the block
that handles params["limit"] and params["threshold"] (the code that calls
extractInt -> sets p.topK and extractFloat64 -> checks threshold) so they refer
to "limit" and "threshold" (e.g., "'limit' must be a number", "'limit' must be
between 0 and 20", "'threshold' must be a number", "'threshold' must be between
0.0 and 1.0") to match the normalized parameter keys and policy-definition.yaml.
- Around line 1356-1383: The helper buildErrorResponse currently logs an error
and returns an ImmediateResponse 400 which blocks requests; change it so runtime
failures during OnRequest log a warning (use slog.Warn with the error details)
and return a pass-through action instead of an ImmediateResponse — i.e., return
the policy.RequestAction that continues the request unmodified (e.g.,
policy.ContinueRequest or the project’s equivalent) from buildErrorResponse in
SemanticToolFilteringPolicy, ensuring no internal error details are sent to
clients and the original request body/headers remain unchanged.
- Around line 141-145: The current call to embeddingCache.GetAPICache(apiId) in
semantictoolfiltering.go forces a deep copy of the entire cache on each request;
add a lightweight accessor on EmbeddingCacheStore (e.g., GetEntryCount(apiId) or
GetCapacity) that acquires an RLock, returns len(apiCache.Tools) if present, and
RUnlocks, then replace the two-line block that calls GetAPICache and len(...)
with a single call to embeddingCache.GetEntryCount(apiId) (keep the variable
name currentCachedCount) so you avoid cloning the cache in the hot path;
implement the method on the same type that defines GetAPICache
(EmbeddingCacheStore) using the store's mu for concurrency.
- Around line 217-230: The error handling drops the original error (using
fmt.Errorf("...: %w") without err) and on failures in parseEmbeddingConfig,
createEmbeddingProvider, and parseParams the function returns nil instead of a
pass-through policy; update each fmt.Errorf call to include the original err
(e.g., fmt.Errorf("invalid embedding config: %w", err)) and change the failure
path to log the initialization error and return a properly configured
pass-through policy instance (so requests are forwarded unchanged) rather than
nil; apply the same fixes in createEmbeddingProvider (referencing
createEmbeddingProvider, parseEmbeddingConfig, parseParams, and the
GetPolicy/policy constructor path) so initialization errors are logged and
result in a pass-through policy object.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 9a6bd500-7092-45a2-bf50-746926566beb
⛔ Files ignored due to path filters (1)
policies/semantic-tool-filtering/go.sumis excluded by!**/*.sum
📒 Files selected for processing (6)
docs/semantic-tool-filtering/v0.1/docs/semantic-tool-filtering.mddocs/semantic-tool-filtering/v0.1/metadata.jsonpolicies/semantic-tool-filtering/embeddingcache.gopolicies/semantic-tool-filtering/go.modpolicies/semantic-tool-filtering/policy-definition.yamlpolicies/semantic-tool-filtering/semantictoolfiltering.go
…nt based on provider
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
docs/semantic-tool-filtering/v0.1/docs/semantic-tool-filtering.md (1)
126-126: Clarify the filtering behavior explanation.The explanation states that the policy will "filter the
toolsarray to include only the top 3 matches (e.g.,get_weather,book_venue,send_email)," but this is misleading. The example request contains exactly 3 tools and the limit is 3, so all tools happen to be returned. However, the policy ranks tools by semantic relevance to the query, and the actual ordering matters.Consider revising to explain that the policy ranks the tools by semantic similarity and returns the top 3 most relevant ones, noting that
get_weatherwould likely rank highest for the query "Get weather forecast."📝 Suggested clarification
-The policy will interpret the request, calculate embeddings, and filter the `tools` array to include only the top 3 matches (e.g., `get_weather`, `book_venue`, `send_email`). +The policy will calculate embeddings for the query and each tool, rank the tools by semantic similarity, and include only the top 3 most relevant tools in the filtered request. For the query "Get weather forecast," `get_weather` would likely rank highest among the three tools.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@docs/semantic-tool-filtering/v0.1/docs/semantic-tool-filtering.md` at line 126, The sentence about filtering is misleading because the example had exactly three tools; update the wording to state that the policy ranks entries in the tools array by semantic similarity to the request and returns the top 3 most relevant tools (not just "filters to include only the top 3"), and mention that ordering reflects relevance (e.g., for "Get weather forecast" get_weather would likely rank highest); refer to the "tools" array and the policy ranking behavior when making this change.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@docs/semantic-tool-filtering/v0.1/docs/semantic-tool-filtering.md`:
- Line 26: Update the docs table entry for the selectionMode parameter to mark
it as optional (change "Yes" to "No") and note that it defaults to `By Rank`
(SelectionModeTopK); ensure the description still explains the two modes (`By
Rank` / `By Threshold`) and that selectionMode is optional with default
SelectionModeTopK.
- Line 53: Remove the unused configuration parameter
embedding_provider_dimension from the sample configuration so only the consumed
keys remain; locate the sample config where embedding_provider_dimension = 1024
appears and delete that line, ensuring the remaining keys embedding_provider,
embedding_provider_endpoint, embedding_provider_model, and
embedding_provider_api_key are left intact and documented.
---
Nitpick comments:
In `@docs/semantic-tool-filtering/v0.1/docs/semantic-tool-filtering.md`:
- Line 126: The sentence about filtering is misleading because the example had
exactly three tools; update the wording to state that the policy ranks entries
in the tools array by semantic similarity to the request and returns the top 3
most relevant tools (not just "filters to include only the top 3"), and mention
that ordering reflects relevance (e.g., for "Get weather forecast" get_weather
would likely rank highest); refer to the "tools" array and the policy ranking
behavior when making this change.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: e9c5fd37-ab72-4018-b453-f50b449113c6
📒 Files selected for processing (1)
docs/semantic-tool-filtering/v0.1/docs/semantic-tool-filtering.md
… endpoint parameters
There was a problem hiding this comment.
♻️ Duplicate comments (1)
policies/semantic-tool-filtering/policy-definition.yaml (1)
44-50:⚠️ Potential issue | 🟠 Major
toolsJSONPathregex currently rejects indexed paths and accepts non-JSONPath strings.On Line 50, the pattern allows
tools(missing$.) but rejects indexed paths like$.tools[0].function_declarations. That can block valid configurations and admit malformed ones.🔧 Suggested schema fix
toolsJSONPath: type: string description: | JSONPath expression to extract the tools array from the request body (used when toolsIsJson is true). Example: "$.tools" default: "$.tools" - pattern: '^(?:\$\.)?[A-Za-z_][A-Za-z0-9_]*(?:\.[A-Za-z_][A-Za-z0-9_]*)*$' + pattern: '^\$\.([A-Za-z_][A-Za-z0-9_]*(\[\d+\])?\.)*[A-Za-z_][A-Za-z0-9_]*(\[\d+\])?$'🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@policies/semantic-tool-filtering/policy-definition.yaml` around lines 44 - 50, The toolsJSONPath pattern currently accepts names without a leading '$.' and rejects indexed segments like $.tools[0].function_declarations; update the toolsJSONPath schema pattern so it allows an optional leading '$' or '$.' and permits each path segment to be an identifier (letters/numbers/underscores) optionally followed by a numeric index in brackets (e.g., [0]), with dot-separated repeats; apply this change to the pattern field for toolsJSONPath to allow both indexed JSONPath segments and the existing identifier form.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@policies/semantic-tool-filtering/policy-definition.yaml`:
- Around line 44-50: The toolsJSONPath pattern currently accepts names without a
leading '$.' and rejects indexed segments like $.tools[0].function_declarations;
update the toolsJSONPath schema pattern so it allows an optional leading '$' or
'$.' and permits each path segment to be an identifier
(letters/numbers/underscores) optionally followed by a numeric index in brackets
(e.g., [0]), with dot-separated repeats; apply this change to the pattern field
for toolsJSONPath to allow both indexed JSONPath segments and the existing
identifier form.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: f7fd72a7-4905-4316-bace-f8f9ca62579a
📒 Files selected for processing (1)
policies/semantic-tool-filtering/policy-definition.yaml
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@policies/semantic-tool-filtering/methodsforv1alpha.go`:
- Around line 129-145: The current logic uses descHash (from
p.getCacheKey(toolDesc)) as the sole identity and stores a single toolMap in
toolMapByHash which causes collisions to overwrite earlier tools; modify the
flow in the block around extractToolDescription/toolEmbeddingRequest so that
you: 1) extract a canonical toolName using a shared extractor that checks the
top-level "name" and falls back to function.name (e.g., reuse or add an
extractor function used elsewhere), 2) build the embedding input by
concatenating name + description (not description alone) before hashing via
p.getCacheKey, and 3) change toolMapByHash to store a slice/array of toolMap
entries per hash (append the current toolMap) and likewise preserve multiple
toolDescMap entries if needed so collisions don’t drop tools; update
construction of toolEmbeddingRequest (Name, Description, HashKey) accordingly so
scoring/filtering iterates all maps for a given hash.
In `@policies/semantic-tool-filtering/methodsforv1alpha2.go`:
- Around line 133-149: The current code uses descHash (from
p.getCacheKey(toolDesc)) as the sole identity and stores a single toolMap per
hash, causing distinct tools with identical descriptions to collide; change
extractToolDescription usage to a shared extractor that also extracts a fallback
name (use top-level "name" or fall back to "function.name"), build the embedding
text as "<name>: <description>" when creating toolEmbeddingRequest (use Name and
Description together for HashKey), and change toolMapByHash to
map[string][]map[string]interface{} (store/append a slice of toolMap for each
descHash) so collisions keep all candidates; apply the same fix pattern to the
other occurrences mentioned (around the blocks using getCacheKey,
toolEmbeddingRequest, toolDescMap, and toolMapByHash at the other ranges).
In `@policies/semantic-tool-filtering/semantictoolfiltering.go`:
- Around line 139-198: The code currently truncates toolEntriesToCache using
availableSlots before calling embeddingCache.BulkAddTools, which prevents
cache-refresh of updated entries and ignores concurrent changes; instead stop
limiting or slicing toolEntriesToCache in the loop—collect every uncached
ToolEntry and always pass the full slice into embeddingCache.BulkAddTools (do
not pre-truncate by availableSlots), then use the BulkAddResult returned by
BulkAddTools to compute tools actually added vs skipped and update the
toolsCached/willCache/notCached logging accordingly so capacity enforcement and
duplicate-removal happen under the cache's lock.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 5d82aadd-9aa9-4cc4-b5fb-ba11ff4e17ae
⛔ Files ignored due to path filters (1)
policies/semantic-tool-filtering/go.sumis excluded by!**/*.sum
📒 Files selected for processing (5)
policies/semantic-tool-filtering/go.modpolicies/semantic-tool-filtering/methodsforv1alpha.gopolicies/semantic-tool-filtering/methodsforv1alpha2.gopolicies/semantic-tool-filtering/policy-definition.yamlpolicies/semantic-tool-filtering/semantictoolfiltering.go
✅ Files skipped from review due to trivial changes (1)
- policies/semantic-tool-filtering/policy-definition.yaml
🚧 Files skipped from review as they are similar to previous changes (1)
- policies/semantic-tool-filtering/go.mod
| toolDesc := extractToolDescription(toolMap) | ||
| if toolDesc == "" { | ||
| slog.Warn("SemanticToolFiltering: No description found for tool, skipping", | ||
| "toolName", toolMap["name"]) | ||
| continue | ||
| } | ||
|
|
||
| toolName, _ := toolMap["name"].(string) | ||
| descHash := p.getCacheKey(toolDesc) | ||
|
|
||
| embeddingRequests = append(embeddingRequests, toolEmbeddingRequest{ | ||
| Name: toolName, | ||
| Description: toolDesc, | ||
| HashKey: descHash, | ||
| }) | ||
| toolDescMap[descHash] = toolDesc | ||
| toolMapByHash[descHash] = toolMap |
There was a problem hiding this comment.
Don't use a description hash as the only JSON tool identity.
These branches compute descHash from toolDesc and keep a single toolMap per hash. If two JSON tools share the same description, the later one overwrites the earlier one and only one survives scoring/filtering. toolName also comes only from the top-level name, so function-style tools reach the cache with an empty name and same-name replacement stops being reliable. Use a shared extractor that falls back to function.name, builds the embedding text from name + description, and preserves []toolMap per hash if collisions are still possible.
Also applies to: 153-170, 418-432, 438-453
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@policies/semantic-tool-filtering/methodsforv1alpha.go` around lines 129 -
145, The current logic uses descHash (from p.getCacheKey(toolDesc)) as the sole
identity and stores a single toolMap in toolMapByHash which causes collisions to
overwrite earlier tools; modify the flow in the block around
extractToolDescription/toolEmbeddingRequest so that you: 1) extract a canonical
toolName using a shared extractor that checks the top-level "name" and falls
back to function.name (e.g., reuse or add an extractor function used elsewhere),
2) build the embedding input by concatenating name + description (not
description alone) before hashing via p.getCacheKey, and 3) change toolMapByHash
to store a slice/array of toolMap entries per hash (append the current toolMap)
and likewise preserve multiple toolDescMap entries if needed so collisions don’t
drop tools; update construction of toolEmbeddingRequest (Name, Description,
HashKey) accordingly so scoring/filtering iterates all maps for a given hash.
| toolDesc := extractToolDescription(toolMap) | ||
| if toolDesc == "" { | ||
| slog.Warn("SemanticToolFiltering: No description found for tool, skipping", | ||
| "toolName", toolMap["name"]) | ||
| continue | ||
| } | ||
|
|
||
| toolName, _ := toolMap["name"].(string) | ||
| descHash := p.getCacheKey(toolDesc) | ||
|
|
||
| embeddingRequests = append(embeddingRequests, toolEmbeddingRequest{ | ||
| Name: toolName, | ||
| Description: toolDesc, | ||
| HashKey: descHash, | ||
| }) | ||
| toolDescMap[descHash] = toolDesc | ||
| toolMapByHash[descHash] = toolMap |
There was a problem hiding this comment.
Don't use a description hash as the only JSON tool identity.
The v1alpha2 JSON paths have the same lossiness as the v1alpha implementation: descHash is derived from toolDesc, but only one toolMap is kept for each hash. Different tools with the same description therefore collapse into one candidate before scoring. toolName is also read only from the top-level name, so function-style tools enter the cache with an empty name. Reuse a shared extractor that falls back to function.name, builds name + description for the embedding text, and keeps []toolMap per hash if collisions can still occur.
Also applies to: 157-173, 422-436, 442-457
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@policies/semantic-tool-filtering/methodsforv1alpha2.go` around lines 133 -
149, The current code uses descHash (from p.getCacheKey(toolDesc)) as the sole
identity and stores a single toolMap per hash, causing distinct tools with
identical descriptions to collide; change extractToolDescription usage to a
shared extractor that also extracts a fallback name (use top-level "name" or
fall back to "function.name"), build the embedding text as "<name>:
<description>" when creating toolEmbeddingRequest (use Name and Description
together for HashKey), and change toolMapByHash to
map[string][]map[string]interface{} (store/append a slice of toolMap for each
descHash) so collisions keep all candidates; apply the same fix pattern to the
other occurrences mentioned (around the blocks using getCacheKey,
toolEmbeddingRequest, toolDescMap, and toolMapByHash at the other ranges).
| // Calculate available slots for caching new tools | ||
| apiCache := embeddingCache.GetAPICache(apiId) | ||
| currentCachedCount := 0 | ||
| if apiCache != nil { | ||
| currentCachedCount = len(apiCache) | ||
| } | ||
| availableSlots := maxToolsPerAPI - currentCachedCount | ||
| if availableSlots < 0 { | ||
| availableSlots = 0 | ||
| } | ||
|
|
||
| slog.Debug("SemanticToolFiltering: Available cache slots", | ||
| "currentCached", currentCachedCount, | ||
| "maxToolsPerAPI", maxToolsPerAPI, | ||
| "availableSlots", availableSlots) | ||
|
|
||
| // Generate embeddings for ALL uncached tools (for similarity calculation) | ||
| // but only cache the ones that fit | ||
| var toolEntriesToCache []ToolEntry | ||
| toolsCached := 0 | ||
|
|
||
| for _, req := range uncachedRequests { | ||
| embedding, err := p.embeddingProvider.GetEmbedding(req.Description) | ||
| if err != nil { | ||
| slog.Warn("SemanticToolFiltering: Error generating tool embedding, skipping", | ||
| "error", err, "toolName", req.Name) | ||
| continue | ||
| } | ||
|
|
||
| // Add to results so this tool can be used in similarity calculations | ||
| results[req.HashKey] = toolEmbeddingResult{ | ||
| Name: req.Name, | ||
| HashKey: req.HashKey, | ||
| Embedding: embedding, | ||
| FromCache: false, | ||
| } | ||
|
|
||
| // Only cache if we have available slots | ||
| if toolsCached < availableSlots { | ||
| toolEntriesToCache = append(toolEntriesToCache, ToolEntry{ | ||
| HashKey: req.HashKey, | ||
| Name: req.Name, | ||
| Embedding: embedding, | ||
| }) | ||
| toolsCached++ | ||
| } else { | ||
| slog.Debug("SemanticToolFiltering: Tool processed but not cached (limit reached)", | ||
| "toolName", req.Name) | ||
| } | ||
| } | ||
|
|
||
| slog.Debug("SemanticToolFiltering: Embedding generation complete", | ||
| "totalProcessed", len(results), | ||
| "newlyGenerated", len(uncachedRequests), | ||
| "willCache", toolsCached, | ||
| "notCached", len(uncachedRequests)-toolsCached) | ||
|
|
||
| // Bulk add embeddings that fit in cache | ||
| if len(toolEntriesToCache) > 0 { | ||
| bulkResult := embeddingCache.BulkAddTools(apiId, toolEntriesToCache) |
There was a problem hiding this comment.
Let BulkAddTools enforce capacity under the lock.
BulkAddTools already removes same-name entries and recomputes available slots from the live cache state. Truncating toolEntriesToCache before that call means a full cache cannot refresh a tool whose description hash changed, and any concurrent cache changes between the count and the bulk add are ignored. Pass every uncached entry down and use BulkAddResult for the final added/skipped accounting.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@policies/semantic-tool-filtering/semantictoolfiltering.go` around lines 139 -
198, The code currently truncates toolEntriesToCache using availableSlots before
calling embeddingCache.BulkAddTools, which prevents cache-refresh of updated
entries and ignores concurrent changes; instead stop limiting or slicing
toolEntriesToCache in the loop—collect every uncached ToolEntry and always pass
the full slice into embeddingCache.BulkAddTools (do not pre-truncate by
availableSlots), then use the BulkAddResult returned by BulkAddTools to compute
tools actually added vs skipped and update the toolsCached/willCache/notCached
logging accordingly so capacity enforcement and duplicate-removal happen under
the cache's lock.
Purpose
Goals
Approach
User stories
Samples
Summary by CodeRabbit
New Features
Documentation