Skip to content

Add semantic tool filtering Policy to WSO2 API Platform #104

Merged
RakhithaRR merged 12 commits intowso2:mainfrom
NaveenSandaruwan:addSemanticToolFiltering
Mar 27, 2026
Merged

Add semantic tool filtering Policy to WSO2 API Platform #104
RakhithaRR merged 12 commits intowso2:mainfrom
NaveenSandaruwan:addSemanticToolFiltering

Conversation

@NaveenSandaruwan
Copy link
Copy Markdown
Contributor

@NaveenSandaruwan NaveenSandaruwan commented Mar 24, 2026

Purpose

This PR introduces a Semantic Tool Filtering Policy for the WSO2 AI Gateway to solve the problem of "tool over-exposure" in Agentic applications. Currently, exposing all available tools (50-100+) to an LLM leads to degraded performance, increased risk of hallucinations, and high token costs.

Goals

  • Reduce LLM Hallucinations: Minimize the "search space" to improve the mathematical probability of correct tool selection.
  • Optimize Performance: Prevent context window overflow and reduce system latency by sending only relevant content to the LLM.
  • Cost Efficiency: Lower token usage by filtering out irrelevant tool descriptions before request reach to LLM .

Approach

The policy implements a semantic retrieval and caching process:

  1. Semantic Ranking: Converts the user query and tool metadata into high-dimensional vectors to calculate Cosine Similarity.
  2. Dynamic Filtering: Injects only the top-ranked tools (e.g., top 5) into the LLM prompt.
  3. Caching Mechanism: Uses a hashing strategy to store tool embeddings in a vector database. This avoids redundant embedding generation for recurring tool descriptions at runtime.

User stories

As a developer building sophisticated Agentic applications, I want the AI Gateway to automatically select the most relevant tools for a user's request so that the application remains fast, accurate, and cost-effective even as my toolset grows.

Samples

Semantic tool filtering process

Summary by CodeRabbit

  • New Features

    • Added Semantic Tool Filtering policy to semantically filter available tools by relevance, supporting top‑K (rank) and threshold selection, JSONPath/tag-based extraction, and mixed extraction modes.
    • Added an in-memory embedding cache to reduce repeated embedding calls and improve performance.
  • Documentation

    • Added comprehensive docs, metadata, configuration examples, and registration instructions with sample payloads for the new policy.

@coderabbitai
Copy link
Copy Markdown
Contributor

coderabbitai bot commented Mar 24, 2026

Walkthrough

Adds 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

Cohort / File(s) Summary
Documentation
docs/semantic-tool-filtering/v0.1/docs/semantic-tool-filtering.md, docs/semantic-tool-filtering/v0.1/metadata.json
New policy documentation and metadata describing behavior, config.toml system parameters, JSONPath/text extraction options, selection modes, registration, and example request/response scenarios.
Policy schema & module
policies/semantic-tool-filtering/policy-definition.yaml, policies/semantic-tool-filtering/go.mod
New policy definition (system & runtime params, validation) and Go module declaration with dependencies.
Implementation — core policy
policies/semantic-tool-filtering/semantictoolfiltering.go
Policy implementation: config parsing/validation, embedding provider integration, cosine similarity, selection modes (TopK/Threshold), JSONPath-based request updates, text tag extraction/rebuild, Mode()/OnResponse, and exported GetPolicy/GetPolicyV2.
Implementation — request handlers
policies/semantic-tool-filtering/methodsforv1alpha.go, policies/semantic-tool-filtering/methodsforv1alpha2.go
v1alpha and v1alpha2 request entrypoints and handlers: JSON, text, and mixed processing flows, error handling that passes through on failure, and body read/modify logic.
Implementation — embedding cache
policies/semantic-tool-filtering/embeddingcache.go
Thread-safe in-memory cache keyed by API ID + description hash with API-level and per-API LRU eviction, bulk-add semantics, cache stats and management APIs, and hashing utility. Key area for review: concurrency, eviction correctness, and deep-copy semantics.

Sequence Diagram

sequenceDiagram
    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
Loading

Estimated code review effort

🎯 5 (Critical) | ⏱️ ~120 minutes

Poem

🐇 I hop through queries and vectors bright,
I cache tool echoes in the soft moonlight.
With cosine twirls I sort and prune,
Pruning the list so models swoon.
Hooray — fewer choices, clearer tune!

🚥 Pre-merge checks | ✅ 3
✅ Passed checks (3 passed)
Check name Status Explanation
Title check ✅ Passed The title accurately summarizes the main change: adding a new Semantic Tool Filtering Policy to the WSO2 API Platform, which aligns with all files in the changeset.
Description check ✅ Passed The PR description includes Purpose, Goals, Approach, User stories, and Samples sections from the template. Most required sections are well-documented with clear explanations of the problem, solution, and implementation strategy.
Docstring Coverage ✅ Passed Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@NaveenSandaruwan NaveenSandaruwan changed the title Add semantic tool filtering Add semantic tool filtering Policy to WSO2 API Platform Mar 24, 2026
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

📥 Commits

Reviewing files that changed from the base of the PR and between 619237c and f12c109.

⛔ Files ignored due to path filters (1)
  • policies/semantic-tool-filtering/go.sum is excluded by !**/*.sum
📒 Files selected for processing (6)
  • docs/semantic-tool-filtering/v0.1/docs/semantic-tool-filtering.md
  • docs/semantic-tool-filtering/v0.1/metadata.json
  • policies/semantic-tool-filtering/embeddingcache.go
  • policies/semantic-tool-filtering/go.mod
  • policies/semantic-tool-filtering/policy-definition.yaml
  • policies/semantic-tool-filtering/semantictoolfiltering.go

Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 tools array 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_weather would 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

📥 Commits

Reviewing files that changed from the base of the PR and between f12c109 and 9025ca9.

📒 Files selected for processing (1)
  • docs/semantic-tool-filtering/v0.1/docs/semantic-tool-filtering.md

Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

♻️ Duplicate comments (1)
policies/semantic-tool-filtering/policy-definition.yaml (1)

44-50: ⚠️ Potential issue | 🟠 Major

toolsJSONPath regex 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

📥 Commits

Reviewing files that changed from the base of the PR and between 9025ca9 and a2719ff.

📒 Files selected for processing (1)
  • policies/semantic-tool-filtering/policy-definition.yaml

Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

📥 Commits

Reviewing files that changed from the base of the PR and between a2719ff and 6e721a7.

⛔ Files ignored due to path filters (1)
  • policies/semantic-tool-filtering/go.sum is excluded by !**/*.sum
📒 Files selected for processing (5)
  • policies/semantic-tool-filtering/go.mod
  • policies/semantic-tool-filtering/methodsforv1alpha.go
  • policies/semantic-tool-filtering/methodsforv1alpha2.go
  • policies/semantic-tool-filtering/policy-definition.yaml
  • policies/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

Comment on lines +129 to +145
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
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

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.

Comment on lines +133 to +149
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
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

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).

Comment on lines +139 to +198
// 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)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

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.

@RakhithaRR RakhithaRR merged commit 7f39242 into wso2:main Mar 27, 2026
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants