feat: a flexible topic detector rail based on embeddings models #1497
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Contributor
Documentation preview |
m-misiura
force-pushed
the
feature/generic-embedding-detector
branch
from
1c8c918 to
67d8af4
Compare
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
m-misiura
force-pushed
the
feature/generic-embedding-detector
branch
2 times, most recently
from
854a932 to
51b1f52
Compare
RobGeada
approved these changes
Nov 7, 2025
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Greptile Overview
Greptile Summary
adds lightweight embedding-based topic detection rail using semantic similarity to block off-topic queries without LLM calls
Key changes:
- new
EmbeddingTopicDetectorclass computes cosine similarity between query and pre-defined topic examples - supports both input and output rails with configurable threshold and top_k parameters
- caches detector instances for performance
- configuration enabled via
ConfigDict(extra="allow")inRailsConfigData - includes comprehensive tests and documentation
Issues found:
- critical bug in cosine similarity calculation where
or 1fallback incorrectly applies to division result instead of norm product - cache key missing
top_kparameter, causing stale detector reuse when top_k changes
Confidence Score: 2/5
- PR has critical logic bugs that will cause incorrect topic detection behavior in production
- two critical logic errors in actions.py: (1) cosine similarity fallback produces incorrect results for orthogonal vectors, (2) cache key omission causes wrong detector to be used when top_k changes - both directly impact core functionality
- nemoguardrails/library/embedding_topic_detector/actions.py requires immediate fixes for similarity calculation and cache key
Important Files Changed
File Analysis
| Filename | Score | Overview |
|---|---|---|
| nemoguardrails/library/embedding_topic_detector/actions.py | 2/5 | adds embedding-based topic detection with cosine similarity; has critical bugs in similarity calculation and cache key generation |
| nemoguardrails/rails/llm/config.py | 5/5 | adds model_config = ConfigDict(extra="allow") to allow custom rail configs like embedding_topic_detector |
| tests/test_embedding_topic_detector.py | 4/5 | comprehensive tests covering on/off topic detection, empty queries, and output rails; could benefit from edge case testing |
Sequence Diagram
sequenceDiagram
participant User
participant Rails as NeMo Guardrails
participant Flow as Embedding Topic Flow
participant Action as EmbeddingTopicDetector
participant Cache as Detector Cache
participant Model as Embedding Model
participant LLM
User->>Rails: Send message
Rails->>Flow: Execute input flow
Flow->>Action: embedding_topic_check(context)
alt Detector not in cache
Action->>Cache: Check cache with key
Cache-->>Action: Not found
Action->>Model: init_embedding_model()
Model-->>Action: Embedding model instance
Action->>Model: encode(examples)
Model-->>Action: Pre-computed embeddings
Action->>Cache: Store detector
else Detector in cache
Action->>Cache: Retrieve cached detector
Cache-->>Action: Detector instance
end
Action->>Model: encode_async([query])
Model-->>Action: Query embedding
Action->>Action: Compute cosine similarity
Action->>Action: Sort by similarity, take top_k
Action->>Action: Calculate category scores
Action->>Action: Compare max_score vs threshold
Action-->>Flow: {on_topic, confidence, category}
alt Off-topic detected
Flow->>Flow: Check enable_rails_exceptions
alt Exceptions enabled
Flow->>Rails: Send OffTopicRailException
else Exceptions disabled
Flow->>Rails: bot refuse to respond
end
Flow->>Rails: abort
else On-topic
Flow->>Rails: Continue to LLM
Rails->>LLM: Forward message
LLM-->>Rails: Generate response
Rails->>Flow: Execute output flow (optional)
Flow->>Action: embedding_topic_check_output(bot_message)
Action-->>Flow: {on_topic, confidence, category}
alt Output off-topic
Flow->>Rails: bot refuse to respond / abort
else Output on-topic
Rails-->>User: Return response
end
end
7 files reviewed, 2 comments
Comment on lines
54
to
57
| float( | ||
| np.dot(query_emb, emb) | ||
| / (np.linalg.norm(query_emb) * np.linalg.norm(emb) or 1) | ||
| ), |
Contributor
There was a problem hiding this comment.
logic: or 1 fallback applies to entire expression not just norms - if dot product is 0, the division result becomes 0, then 0 or 1 evaluates to 1, giving false similarity of 1.0 for orthogonal vectors
Suggested change
| float( | |
| np.dot(query_emb, emb) | |
| / (np.linalg.norm(query_emb) * np.linalg.norm(emb) or 1) | |
| ), | |
| np.dot(query_emb, emb) | |
| / ((np.linalg.norm(query_emb) * np.linalg.norm(emb)) or 1e-10) |
|
|
||
| async def _check(context: Optional[dict], llm_task_manager, message_key: str) -> dict: | ||
| config = llm_task_manager.config.rails.config.embedding_topic_detector | ||
| cache_key = f"{config['embedding_model']}_{config['embedding_engine']}_{config.get('threshold', 0.75)}" |
Contributor
There was a problem hiding this comment.
logic: cache key doesn't include top_k - changing top_k won't create new detector instance, will use stale cached detector with old top_k value
m-misiura
force-pushed
the
feature/generic-embedding-detector
branch
from
d6a80ad to
b40ecfb
Compare
Contributor
There was a problem hiding this comment.
Greptile Overview
Greptile Summary
Adds lightweight embedding-based topic detection rail using semantic similarity, providing fast alternative to LLM-based detection.
Key Changes:
- New
EmbeddingTopicDetectorclass computes cosine similarity between query and pre-defined example embeddings - Configurable via
rails.config.embedding_topic_detectorwith support for different embedding models/engines - Supports both input and output rails with Colang 1.x and 2.x flows
- Detector instances are cached based on model, engine, threshold, and top_k parameters
- Comprehensive test suite covering on/off-topic scenarios and edge cases
Issues Found:
- Critical: Cache key on actions.py:80 doesn't include
examplesdictionary, so changing examples won't invalidate cache - Cosine similarity division by zero protection already fixed in previous comment
Confidence Score: 3/5
- This PR has a critical caching bug that could cause incorrect behavior in production
- Score reflects one critical bug (cache key missing examples) that would cause stale detectors to be reused when examples change, plus one already-noted cosine similarity issue. Documentation and tests are solid. The feature itself is well-designed.
- nemoguardrails/library/embedding_topic_detector/actions.py needs the cache key bug fixed before merge
Important Files Changed
File Analysis
| Filename | Score | Overview |
|---|---|---|
| nemoguardrails/library/embedding_topic_detector/actions.py | 3/5 | Core implementation with two critical bugs: cache key missing examples hash (will reuse stale detectors) and previous cosine similarity fix incomplete |
| nemoguardrails/rails/llm/config.py | 5/5 | Added model_config = ConfigDict(extra="allow") to allow dynamic rail configurations like embedding_topic_detector |
| tests/test_embedding_topic_detector.py | 5/5 | Comprehensive test coverage including off-topic blocking, detector logic, empty query handling, and output rail functionality |
Sequence Diagram
sequenceDiagram
participant User
participant API as Chat API
participant Flow as Colang Flow
participant Action as embedding_topic_check
participant Detector as EmbeddingTopicDetector
participant Cache as _detector_cache
participant Model as Embedding Model
User->>API: POST /v1/chat/completions
API->>Flow: Process input message
Flow->>Action: Execute embedding_topic_check
Action->>Cache: Check cache_key exists
alt Cache Miss
Cache-->>Action: Not found
Action->>Detector: Create new EmbeddingTopicDetector
Detector->>Model: init_embedding_model()
Detector->>Model: encode(examples)
Model-->>Detector: Pre-computed embeddings
Action->>Cache: Store detector
else Cache Hit
Cache-->>Action: Return cached detector
end
Action->>Detector: detect(query)
Detector->>Model: encode_async([query])
Model-->>Detector: query_embedding
Detector->>Detector: Calculate cosine similarities
Detector->>Detector: Get top_k most similar
Detector->>Detector: Compute category scores
Detector-->>Action: {on_topic, confidence, top_category, scores}
alt Off-topic (confidence < threshold)
Action-->>Flow: on_topic: false
Flow->>Flow: bot refuse to respond
Flow->>API: Abort
API-->>User: "I'm sorry, I can't respond to that."
else On-topic
Action-->>Flow: on_topic: true
Flow->>API: Continue to LLM
API-->>User: Generated response
end
7 files reviewed, 1 comment
|
|
||
| async def _check(context: Optional[dict], llm_task_manager, message_key: str) -> dict: | ||
| config = llm_task_manager.config.rails.config.embedding_topic_detector | ||
| cache_key = f"{config['embedding_model']}_{config['embedding_engine']}_{config.get('threshold', 0.75)}_{config.get('top_k', 3)}" |
Contributor
There was a problem hiding this comment.
logic: cache key doesn't include examples - changing examples won't create new detector instance, will use stale cached detector with old examples
The cache key should include a hash or identifier of the examples dictionary to ensure different example sets create different detector instances.
Signed-off-by: mmisiura <[email protected]>
m-misiura
force-pushed
the
feature/generic-embedding-detector
branch
from
fa896c8 to
e48416c
Compare
Contributor
There was a problem hiding this comment.
Greptile Overview
Greptile Summary
Adds a lightweight embedding-based topic detection rail that uses semantic similarity to block off-topic queries without LLM calls.
Key Changes:
- New
EmbeddingTopicDetectorclass computes cosine similarity between query embeddings and pre-configured topic examples - Efficient caching mechanism using hash-based keys (includes model, threshold, top_k, and examples hash)
- Supports both input and output rails with proper exception handling for Colang 1.0 and 2.0
- Configuration integrated via
RailsConfigDatawithextra='allow'to support dynamic rail configs - Comprehensive test coverage for on-topic/off-topic detection, edge cases, and empty query handling
Implementation Quality:
- Proper numerical stability with
1e-10fallback for zero-norm vectors in cosine similarity - Top-k averaging provides robust category scoring across multiple examples
- Async embedding encoding for non-blocking performance
Confidence Score: 4/5
- Safe to merge with minor considerations for edge case behavior
- Well-structured implementation with proper caching, numerical stability, comprehensive tests, and clear documentation. The code follows existing patterns and handles edge cases appropriately. One area of consideration is the behavior when top_k exceeds available examples, though the implementation handles this gracefully.
- No files require special attention - implementation is solid across all changes
Important Files Changed
File Analysis
| Filename | Score | Overview |
|---|---|---|
| nemoguardrails/library/embedding_topic_detector/actions.py | 4/5 | Core detector logic with proper caching, cosine similarity calculation handles edge cases correctly with 1e-10 fallback |
| nemoguardrails/rails/llm/config.py | 5/5 | Adds extra='allow' to RailsConfigData to support dynamic rail configurations, consistent with existing pattern |
| tests/test_embedding_topic_detector.py | 4/5 | Comprehensive test coverage including edge cases, integration tests, and output rail verification |
Sequence Diagram
sequenceDiagram
participant User
participant Rails as NeMo Guardrails
participant Flow as Embedding Topic Check Flow
participant Action as embedding_topic_check Action
participant Cache as Detector Cache
participant Detector as EmbeddingTopicDetector
participant EmbedModel as Embedding Model
User->>Rails: Send message
Rails->>Flow: Trigger input flow
Flow->>Action: Execute action
Action->>Cache: Check cache with key
alt Cache miss
Cache->>Detector: Create new detector
Detector->>EmbedModel: Initialize model
Detector->>EmbedModel: Encode example queries
EmbedModel-->>Detector: Return embeddings
Detector->>Cache: Store detector
end
Cache-->>Action: Return detector
Action->>Detector: detect(query)
Detector->>EmbedModel: Encode query async
EmbedModel-->>Detector: Query embedding
Detector->>Detector: Calculate cosine similarities
Detector->>Detector: Average top-k scores
Detector->>Detector: Compare with threshold
Detector-->>Action: Return result {on_topic, confidence, top_category}
Action-->>Flow: Return result
alt Off-topic detected
Flow->>Rails: Abort with refuse message
Rails->>User: "I'm sorry, I can't respond to that."
else On-topic
Flow->>Rails: Continue processing
Rails->>User: Normal response
end
7 files reviewed, no comments
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
This PR adds embedding-based topic detection rail using semantic similarity to block off-topic queries. This additional detector would give users a lightweight, fast and flexible alternative for topic detection without the need to make any LLMs calls.
Observed behaviour
Start up the server
Start up a server with a config such as this one
Send a request which should be far away from the pre-specified theme
Send a request which should be within the pre-specified theme
Checklist