feat(nodes): add Cohere Rerank pipeline node for RAG quality improvement#499
feat(nodes): add Cohere Rerank pipeline node for RAG quality improvement#499nihalnihalani wants to merge 10 commits intorocketride-org:developfrom
Conversation
nihalnihalani
requested review from
Rod-Christensen,
jmaionchi and
stepmikhaylov
as code owners
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds a Cohere-based rerank node: a global lifecycle managing a shared reranker, a Cohere Rerank client wrapper with error mapping and thresholding, an instance method to rerank and emit documents/answers, service metadata and requirements, and comprehensive unit tests. Changes
Sequence DiagramsequenceDiagram
participant Client as Question Input
participant IGlobal
participant IInstance
participant RerankClient
participant CohereAPI as Cohere API
participant Output as Output Lanes
Note over IGlobal: Initialization
IGlobal->>IGlobal: validateConfig (install deps, load config, check apikey)
IGlobal->>IGlobal: beginGlobal (create _reranker = RerankClient(...))
Note over Client,Output: Processing
Client->>IInstance: writeQuestions(question)
IInstance->>IInstance: extract query & documents
IInstance->>RerankClient: rerank_with_threshold(query, documents)
RerankClient->>CohereAPI: rerank(...)
CohereAPI-->>RerankClient: ranked results
RerankClient-->>IInstance: results (index, relevance_score, text)
IInstance->>Output: write `documents` (reranked, metadata preserved)
IInstance->>Output: write `answers`
Note over IGlobal: Teardown
IGlobal->>IGlobal: endGlobal (_reranker = None)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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 `@nodes/src/nodes/rerank_cohere/IInstance.py`:
- Around line 101-108: The code currently only calls
self.instance.writeAnswers(question) inside the if reranked_docs block, which
drops the question when min_score filters out all documents; change the flow so
question.documents is always set (set to reranked_docs if truthy or an empty
list otherwise) and call self.instance.writeAnswers(question) unconditionally,
while keeping self.instance.writeDocuments(reranked_docs) only inside the if
reranked_docs branch; update references to reranked_docs, question.documents,
self.instance.writeDocuments and self.instance.writeAnswers accordingly.
- Around line 64-99: The reranking uses doc_texts (built from a filtered subset
of question.documents) but later applies result['index'] against the original
unfiltered list, causing misaligned metadata; to fix, while building doc_texts
in the loop over question.documents also build a parallel source_docs list (or
indices list) that records the original document objects (or their indices),
call reranker.rerank_with_threshold with doc_texts as before, and when iterating
reranked results use original_doc = source_docs[result['index']] (or fetch by
the recorded original index) to construct the new Doc (preserving metadata)
before appending to reranked_docs so mapping remains correct.
In `@nodes/src/nodes/rerank_cohere/rerank_client.py`:
- Around line 51-57: The constructor assigns self._top_n and self._min_score
from config without validation; add checks in the constructor to raise
ValueError if top_n < 1 or if min_score is outside [0.0, 1.0] (use the names
self._top_n and self._min_score), and also add the same validation immediately
before any Cohere API call (the method that invokes cohere.rerank at the sites
around the existing calls on lines ~84-85 and ~129-135) to defensively reject
bad inputs rather than passing them to the remote API or using them for
filtering; ensure error messages clearly state the invalid parameter and
expected range.
In `@nodes/src/nodes/rerank_cohere/services.json`:
- Around line 151-170: The preset profiles "rerank_cohere.rerank-v3.5" and
"rerank_cohere.rerank-v3.0" only expose "llm.cloud.apikey" so UI users cannot
adjust "top_n" and "min_score"; update the service definitions for
rerank_cohere.rerank-v3.5 and rerank_cohere.rerank-v3.0 to include "top_n" and
"min_score" in their "properties" arrays (matching how rerank_cohere.custom
lists "model", "top_n", "min_score", "llm.cloud.apikey") so the knobs are
surfaced for the preset profiles.
- Around line 119-123: The preset using the invalid model ID "rerank-v3.0"
should be updated to a valid Cohere model identifier or removed; replace the
"model" value in the "rerank-v3.0" preset with "rerank-english-v3.0" for English
data or "rerank-multilingual-v3.0" for non-English data (or delete the entire
preset entry if obsolete) so requests to Cohere use a supported rerank model.
In `@test/nodes/test_rerank_cohere.py`:
- Around line 479-504: The shared _mock_rocketlib.warning mock is not cleared
between tests, so update test_validate_config_invalid_apikey to reset or
re-create the mock before asserting: call _mock_rocketlib.warning.reset_mock()
(or reassign a fresh Mock) at the start of test_validate_config_invalid_apikey,
or alternatively patch rerank_cohere.IGlobal.warning inside that test; ensure
the assertion after iglobal.validateConfig() uses the freshly reset/mock
instance so it only passes if IGlobal.validateConfig actually issued a new
warning.
🪄 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: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: d6e93981-4a71-4dc1-ad1b-06f730919684
📒 Files selected for processing (7)
nodes/src/nodes/rerank_cohere/IGlobal.pynodes/src/nodes/rerank_cohere/IInstance.pynodes/src/nodes/rerank_cohere/__init__.pynodes/src/nodes/rerank_cohere/requirements.txtnodes/src/nodes/rerank_cohere/rerank_client.pynodes/src/nodes/rerank_cohere/services.jsontest/nodes/test_rerank_cohere.py
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (4)
test/nodes/test_rerank_cohere.py (1)
479-504:⚠️ Potential issue | 🟡 MinorReset
_mock_rocketlib.warningbetween tests to avoid false positives.
_mock_rocketlib.warningis module-scoped, so aftertest_validate_config_missing_apikey()calls it,assert_called()intest_validate_config_invalid_apikey()can pass even ifvalidateConfig()doesn't actually emit a warning. Reset the mock before each test.🛠️ Proposed fix
def test_validate_config_missing_apikey(self): """ValidateConfig warns when API key is missing.""" iglobal = self._make_iglobal() + _mock_rocketlib.warning.reset_mock() with patch('rerank_cohere.IGlobal.Config') as mock_config_cls: mock_config_cls.getNodeConfig.return_value = {'apikey': '', 'model': 'rerank-v3.5'} iglobal.validateConfig() - _mock_rocketlib.warning.assert_called() + _mock_rocketlib.warning.assert_called_once() def test_validate_config_invalid_apikey(self): """ValidateConfig warns when API key is rejected by Cohere.""" iglobal = self._make_iglobal() + _mock_rocketlib.warning.reset_mock() mock_client_instance = Mock() mock_client_instance.rerank.side_effect = _MockUnauthorizedError('invalid key') ... - _mock_rocketlib.warning.assert_called() + _mock_rocketlib.warning.assert_called_once()🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@test/nodes/test_rerank_cohere.py` around lines 479 - 504, The module-scoped mock _mock_rocketlib.warning is not being reset between tests causing false positives; update the tests (e.g., test_validate_config_missing_apikey and test_validate_config_invalid_apikey) to call _mock_rocketlib.warning.reset_mock() before invoking iglobal.validateConfig() (or reset it in a shared setUp/fixture) so each test asserts against a fresh mock; ensure you add the reset call at the start of test_validate_config_invalid_apikey (or in the test class setUp) so _mock_rocketlib.warning.assert_called() only reflects the current test's behavior.nodes/src/nodes/rerank_cohere/IInstance.py (2)
106-113:⚠️ Potential issue | 🟠 MajorAlways forward the question to the
answerslane, even when all documents are filtered.When
min_scorefilters out every document, the question is silently dropped becausewriteAnswersis only called insideif reranked_docs:. Downstream nodes expecting a question will never receive it, breaking pipeline continuity.🛠️ Proposed fix
- # Write reranked documents to the documents output lane - if reranked_docs: - self.instance.writeDocuments(reranked_docs) - - # Also write an answer with the reranked documents attached - # to the question for downstream LLM consumption - question.documents = reranked_docs - self.instance.writeAnswers(question) + # Always update question.documents with reranked results (may be empty) + question.documents = reranked_docs + + # Write reranked documents to the documents output lane if any + if reranked_docs: + self.instance.writeDocuments(reranked_docs) + + # Always forward the question downstream for pipeline continuity + self.instance.writeAnswers(question)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@nodes/src/nodes/rerank_cohere/IInstance.py` around lines 106 - 113, The code currently only calls self.instance.writeAnswers(question) when reranked_docs is truthy, causing the question to be dropped if min_score filters out all documents; modify the logic in IInstance (around writeDocuments / writeAnswers) so that writeAnswers is always called regardless of reranked_docs, setting question.documents = reranked_docs (which may be empty) before calling writeAnswers, and only conditionally call self.instance.writeDocuments(reranked_docs) when reranked_docs is non-empty; keep the existing symbols self.instance.writeDocuments, self.instance.writeAnswers, question.documents, reranked_docs, and min_score to locate and update the block.
69-89:⚠️ Potential issue | 🔴 CriticalIndex mapping is broken when documents with empty content are skipped.
doc_textsis built by filtering out documents with emptypage_content, butresult['index']from Cohere refers to indices indoc_texts, notquestion.documents. When a blank document is skipped, the indices become misaligned, causing wrong document/metadata pairing.🐛 Proposed fix to track source documents
doc_texts = [] + source_docs = [] for doc in question.documents: if hasattr(doc, 'page_content') and doc.page_content: doc_texts.append(doc.page_content) + source_docs.append(doc) elif isinstance(doc, dict) and doc.get('page_content'): doc_texts.append(doc['page_content']) + source_docs.append(doc) if not doc_texts: raise ValueError('No document content found to rerank') # Rerank the documents with threshold filtering reranked = reranker.rerank_with_threshold( query=query_text, documents=doc_texts, ) # Build reranked Doc objects preserving original metadata reranked_docs = [] for result in reranked: original_idx = result['index'] - original_doc = question.documents[original_idx] + original_doc = source_docs[original_idx]🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@nodes/src/nodes/rerank_cohere/IInstance.py` around lines 69 - 89, The reranking index mapping is broken because doc_texts filters out empty page_content so result['index'] refers to positions in doc_texts, not question.documents; fix by building and using an index map when constructing doc_texts (e.g., track original indices for each appended page_content), pass only the filtered texts to reranker.rerank_with_threshold, and when building reranked_docs use original_idx = index_map[result['index']] to retrieve the correct original_doc and its metadata for preservation (refer to variables/functions: doc_texts, question.documents, reranker.rerank_with_threshold, reranked, reranked_docs, result['index']).nodes/src/nodes/rerank_cohere/rerank_client.py (1)
89-91:⚠️ Potential issue | 🟠 MajorAdd validation for
top_nandmin_scoreconfiguration values.The constructor accepts
top_nandmin_scorefrom config without validation. Invalid values (e.g.,top_n < 1ormin_scoreoutside[0.0, 1.0]) will either cause Cohere API errors or unexpected filtering behavior.🛠️ Proposed validation
self._model = config.get('model', 'rerank-v3.5') self._top_n = config.get('top_n', 5) self._min_score = config.get('min_score', 0.0) + + if not isinstance(self._top_n, int) or self._top_n < 1: + raise ValueError('top_n must be an integer >= 1') + if not isinstance(self._min_score, (int, float)) or not (0.0 <= self._min_score <= 1.0): + raise ValueError('min_score must be a number between 0.0 and 1.0') apikey = config.get('apikey', '')🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@nodes/src/nodes/rerank_cohere/rerank_client.py` around lines 89 - 91, Validate the configuration values for top_n and min_score in the Rerank client constructor: ensure self._top_n (from config.get('top_n')) is an integer >= 1 (coerce to int if appropriate or raise ValueError) and ensure self._min_score (from config.get('min_score')) is a float within [0.0, 1.0] (raise ValueError if outside this range); implement these checks in the __init__ where self._top_n and self._min_score are set (e.g., in RerankClient.__init__ or the initializer that sets self._model, self._top_n, self._min_score) and provide clear error messages naming the invalid parameter.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@nodes/src/nodes/rerank_cohere/IGlobal.py`:
- Around line 91-106: Wrap the RerankClient construction in a try/except that
catches ValueError raised when the API key is missing (around the
RerankClient(self.glb.logicalType, config, bag) call), log a clear warning/error
(including the config or logicalType for context) and set self._reranker to None
(or re-raise a more descriptive exception) so beginGlobal() doesn't let the
ValueError propagate uncaught; reference Config.getNodeConfig, RerankClient, and
self._reranker to locate the code to change.
- Around line 68-79: Extract the repeated whitespace-collapse and truncation
logic into a small helper (e.g., normalize_and_truncate_message(message: str,
max_len: int = 500) -> str) and call it from both except handlers for
UnauthorizedError and BadRequestError instead of duplicating the code; the
helper should use the same regex collapse (re.sub(r'\s+', ' ', ...).strip()),
truncate to max_len and append '...' when trimmed, and then pass the returned
string to warning(...) in both except blocks (define the helper near the top of
the module or inside the same class/namespace so
UnauthorizedError/BadRequestError handlers can call it).
In `@nodes/src/nodes/rerank_cohere/rerank_client.py`:
- Around line 141-149: The loop that builds results in rerank_client.py should
be replaced with a list comprehension to be more concise: construct results = [{
'index': r.index, 'relevance_score': r.relevance_score, 'document':
documents[r.index] } for r in response.results], using the existing
response.results, documents, and result attributes (result.index and
result.relevance_score) so behavior and ordering remain identical.
---
Duplicate comments:
In `@nodes/src/nodes/rerank_cohere/IInstance.py`:
- Around line 106-113: The code currently only calls
self.instance.writeAnswers(question) when reranked_docs is truthy, causing the
question to be dropped if min_score filters out all documents; modify the logic
in IInstance (around writeDocuments / writeAnswers) so that writeAnswers is
always called regardless of reranked_docs, setting question.documents =
reranked_docs (which may be empty) before calling writeAnswers, and only
conditionally call self.instance.writeDocuments(reranked_docs) when
reranked_docs is non-empty; keep the existing symbols
self.instance.writeDocuments, self.instance.writeAnswers, question.documents,
reranked_docs, and min_score to locate and update the block.
- Around line 69-89: The reranking index mapping is broken because doc_texts
filters out empty page_content so result['index'] refers to positions in
doc_texts, not question.documents; fix by building and using an index map when
constructing doc_texts (e.g., track original indices for each appended
page_content), pass only the filtered texts to reranker.rerank_with_threshold,
and when building reranked_docs use original_idx = index_map[result['index']] to
retrieve the correct original_doc and its metadata for preservation (refer to
variables/functions: doc_texts, question.documents,
reranker.rerank_with_threshold, reranked, reranked_docs, result['index']).
In `@nodes/src/nodes/rerank_cohere/rerank_client.py`:
- Around line 89-91: Validate the configuration values for top_n and min_score
in the Rerank client constructor: ensure self._top_n (from config.get('top_n'))
is an integer >= 1 (coerce to int if appropriate or raise ValueError) and ensure
self._min_score (from config.get('min_score')) is a float within [0.0, 1.0]
(raise ValueError if outside this range); implement these checks in the __init__
where self._top_n and self._min_score are set (e.g., in RerankClient.__init__ or
the initializer that sets self._model, self._top_n, self._min_score) and provide
clear error messages naming the invalid parameter.
In `@test/nodes/test_rerank_cohere.py`:
- Around line 479-504: The module-scoped mock _mock_rocketlib.warning is not
being reset between tests causing false positives; update the tests (e.g.,
test_validate_config_missing_apikey and test_validate_config_invalid_apikey) to
call _mock_rocketlib.warning.reset_mock() before invoking
iglobal.validateConfig() (or reset it in a shared setUp/fixture) so each test
asserts against a fresh mock; ensure you add the reset call at the start of
test_validate_config_invalid_apikey (or in the test class setUp) so
_mock_rocketlib.warning.assert_called() only reflects the current test's
behavior.
🪄 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: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: e0ffc43a-5c99-4445-9a1a-fec5f7bd70e7
📒 Files selected for processing (5)
nodes/src/nodes/rerank_cohere/IGlobal.pynodes/src/nodes/rerank_cohere/IInstance.pynodes/src/nodes/rerank_cohere/requirements.txtnodes/src/nodes/rerank_cohere/rerank_client.pytest/nodes/test_rerank_cohere.py
| except UnauthorizedError as e: | ||
| message = re.sub(r'\s+', ' ', str(e)).strip() | ||
| if len(message) > 500: | ||
| message = message[:500].rstrip() + '...' | ||
| warning(message) | ||
| return | ||
| except BadRequestError as e: | ||
| message = re.sub(r'\s+', ' ', str(e)).strip() | ||
| if len(message) > 500: | ||
| message = message[:500].rstrip() + '...' | ||
| warning(message) | ||
| return |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Consider extracting duplicate message truncation logic.
The same message normalization pattern (regex whitespace collapse + 500-char truncation) is repeated for both UnauthorizedError and BadRequestError. This could be a small helper function.
♻️ Suggested helper extraction
+def _truncate_error_message(e: Exception, max_len: int = 500) -> str:
+ """Normalize and truncate error message for warning output."""
+ message = re.sub(r'\s+', ' ', str(e)).strip()
+ if len(message) > max_len:
+ message = message[:max_len].rstrip() + '...'
+ return message
+
+
class IGlobal(IGlobalBase):
...
except UnauthorizedError as e:
- message = re.sub(r'\s+', ' ', str(e)).strip()
- if len(message) > 500:
- message = message[:500].rstrip() + '...'
+ message = _truncate_error_message(e)
warning(message)
return
except BadRequestError as e:
- message = re.sub(r'\s+', ' ', str(e)).strip()
- if len(message) > 500:
- message = message[:500].rstrip() + '...'
+ message = _truncate_error_message(e)
warning(message)
return🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@nodes/src/nodes/rerank_cohere/IGlobal.py` around lines 68 - 79, Extract the
repeated whitespace-collapse and truncation logic into a small helper (e.g.,
normalize_and_truncate_message(message: str, max_len: int = 500) -> str) and
call it from both except handlers for UnauthorizedError and BadRequestError
instead of duplicating the code; the helper should use the same regex collapse
(re.sub(r'\s+', ' ', ...).strip()), truncate to max_len and append '...' when
trimmed, and then pass the returned string to warning(...) in both except blocks
(define the helper near the top of the module or inside the same class/namespace
so UnauthorizedError/BadRequestError handlers can call it).
There was a problem hiding this comment.
@coderabbitai Agreed — the duplicate message normalization pattern (regex whitespace collapse + 500-char truncation) for both UnauthorizedError and BadRequestError is a good candidate for extraction. We will introduce a small _truncate_error_message helper as suggested to eliminate the duplication. Thank you for the clean suggestion.
There was a problem hiding this comment.
Rate Limit Exceeded
@nihalnihalani have exceeded the limit for the number of chat messages per hour. Please wait 0 minutes and 23 seconds before sending another message.
| else: | ||
| # Ensure cohere dependency is installed before importing rerank_client | ||
| from depends import depends | ||
|
|
||
| requirements = os.path.dirname(os.path.realpath(__file__)) + '/requirements.txt' | ||
| depends(requirements) | ||
|
|
||
| # Import the rerank client | ||
| from .rerank_client import RerankClient | ||
|
|
||
| # Get our bag | ||
| bag = self.IEndpoint.endpoint.bag | ||
|
|
||
| # Get the passed configuration | ||
| config = Config.getNodeConfig(self.glb.logicalType, self.glb.connConfig) | ||
| self._reranker = RerankClient(self.glb.logicalType, config, bag) |
There was a problem hiding this comment.
Wrap RerankClient instantiation in try-except to handle missing API key gracefully.
If Config.getNodeConfig returns a config without an apikey (or with an empty value), RerankClient.__init__ will raise ValueError('Cohere API key is required'). This exception propagates uncaught from beginGlobal(), potentially leaving the node in a failed state without a clear error message to the user.
Consider catching the ValueError and either re-raising a more descriptive exception or logging an appropriate warning.
🛠️ Proposed fix
# Get the passed configuration
config = Config.getNodeConfig(self.glb.logicalType, self.glb.connConfig)
- self._reranker = RerankClient(self.glb.logicalType, config, bag)
+ try:
+ self._reranker = RerankClient(self.glb.logicalType, config, bag)
+ except ValueError as e:
+ raise Exception(f'Failed to initialize Cohere reranker: {e}') from eBased on context snippet 2 from rerank_client.py (lines 93-95), which shows RerankClient raises ValueError when apikey is empty.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@nodes/src/nodes/rerank_cohere/IGlobal.py` around lines 91 - 106, Wrap the
RerankClient construction in a try/except that catches ValueError raised when
the API key is missing (around the RerankClient(self.glb.logicalType, config,
bag) call), log a clear warning/error (including the config or logicalType for
context) and set self._reranker to None (or re-raise a more descriptive
exception) so beginGlobal() doesn't let the ValueError propagate uncaught;
reference Config.getNodeConfig, RerankClient, and self._reranker to locate the
code to change.
There was a problem hiding this comment.
@coderabbitai Thank you — this is a valid concern. If Config.getNodeConfig returns a config with an empty or missing apikey, the RerankClient.init will raise a ValueError that propagates uncaught from beginGlobal(), potentially leaving the node in a failed state without a user-friendly error message. We will wrap the RerankClient instantiation in a try-except block to catch ValueError and surface a clear warning through the node logging system. This will be prioritized.
There was a problem hiding this comment.
Rate Limit Exceeded
@nihalnihalani have exceeded the limit for the number of chat messages per hour. Please wait 1 minutes and 42 seconds before sending another message.
| results = [] | ||
| for result in response.results: | ||
| results.append( | ||
| { | ||
| 'index': result.index, | ||
| 'relevance_score': result.relevance_score, | ||
| 'document': documents[result.index], | ||
| } | ||
| ) |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Consider using a list comprehension for building results.
The loop can be expressed more concisely as a list comprehension.
♻️ Suggested refactor
- results = []
- for result in response.results:
- results.append(
- {
- 'index': result.index,
- 'relevance_score': result.relevance_score,
- 'document': documents[result.index],
- }
- )
-
- return results
+ return [
+ {
+ 'index': result.index,
+ 'relevance_score': result.relevance_score,
+ 'document': documents[result.index],
+ }
+ for result in response.results
+ ]📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| results = [] | |
| for result in response.results: | |
| results.append( | |
| { | |
| 'index': result.index, | |
| 'relevance_score': result.relevance_score, | |
| 'document': documents[result.index], | |
| } | |
| ) | |
| return [ | |
| { | |
| 'index': result.index, | |
| 'relevance_score': result.relevance_score, | |
| 'document': documents[result.index], | |
| } | |
| for result in response.results | |
| ] |
🧰 Tools
🪛 Ruff (0.15.7)
[warning] 143-149: Use a list comprehension to create a transformed list
(PERF401)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@nodes/src/nodes/rerank_cohere/rerank_client.py` around lines 141 - 149, The
loop that builds results in rerank_client.py should be replaced with a list
comprehension to be more concise: construct results = [{ 'index': r.index,
'relevance_score': r.relevance_score, 'document': documents[r.index] } for r in
response.results], using the existing response.results, documents, and result
attributes (result.index and result.relevance_score) so behavior and ordering
remain identical.
There was a problem hiding this comment.
@coderabbitai Fair point — the explicit loop can be expressed more concisely as a list comprehension without sacrificing readability. We will refactor this to the suggested list comprehension form. Thank you for the nitpick.
There was a problem hiding this comment.
Rate Limit Exceeded
@nihalnihalani have exceeded the limit for the number of chat messages per hour. Please wait 1 minutes and 31 seconds before sending another message.
asclearuc
left a comment
asclearuc
left a comment
There was a problem hiding this comment.
Thank you for this contribution — the node structure, exception hierarchy, and test coverage are solid. A few inline comments have been left on specific lines.
Also, please pay attention to COdeRabbit comments
| "title": "Rerank v3.5", | ||
| "model": "rerank-v3.5", | ||
| "top_n": 5, | ||
| "min_score": 0.0, |
There was a problem hiding this comment.
The constructor already validates these at runtime, but there is no UI-level constraint to prevent invalid values from being saved in the first place. Add range constraints to the field definitions (and to avoid CodeRabbit comment):
"top_n": {
"type": "number",
"title": "Top N",
"description": "Number of top results to return",
"minimum": 1
},
"min_score": {
"type": "number",
"title": "Min Score",
"description": "Minimum relevance score threshold (0.0-1.0)",
"minimum": 0.0,
"maximum": 1.0
}There was a problem hiding this comment.
@asclearuc Agreed. Adding minimum/maximum constraints at the UI level (services.json field definitions) provides a much better user experience than only catching invalid values at runtime. We will add the range constraints you suggested — minimum: 1 for top_n, and minimum: 0.0 / maximum: 1.0 for min_score — to prevent invalid values from being saved in the first place. This also addresses the related CodeRabbit comment. Thank you for the concrete suggestion.
There was a problem hiding this comment.
Fixed in a926179. Added "minimum": 1 constraint to top_n and "minimum": 0.0, "maximum": 1.0 constraints to min_score in the field definitions. Also exposed both fields in the preset profiles (rerank-v3.5 and rerank-v3.0) so users can tune them from the UI.
| raise RerankBadRequestError(f'Invalid rerank request: {e}') from e | ||
| except InternalServerError as e: | ||
| raise RerankServerError(f'Cohere server error: {e}') from e | ||
|
|
There was a problem hiding this comment.
If the Cohere SDK raises anything outside the four known error types, raw SDK exceptions propagate out of the client. A fallback keeps the exception hierarchy intact:
except Exception as e:
raise RerankServerError(f'Unexpected Cohere error: {e}') from eThere was a problem hiding this comment.
@asclearuc Good point. If the Cohere SDK raises an exception type outside the four known error types (UnauthorizedError, BadRequestError, TooManyRequestsError, InternalServerError), raw SDK exceptions will propagate unhandled. We will add a fallback except Exception clause that wraps unknown errors in RerankServerError to keep the exception hierarchy intact and provide consistent error handling for downstream consumers. Thank you for this.
There was a problem hiding this comment.
Fixed in prior commit (already addressed). A catch-all except Exception as e: raise RerankServerError(...) was added as the last handler, so any unknown SDK exception is wrapped in our error hierarchy instead of propagating raw.
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (2)
nodes/src/nodes/rerank_cohere/services.json (1)
154-164:⚠️ Potential issue | 🟠 MajorPreset profiles still hide the advertised tuning knobs.
Only
customexposestop_nandmin_score. Users who choosererank-v3.5orrerank-english-v3.0still can’t change rerank breadth or threshold in the UI, even though the runtime and schema support both.🛠️ Suggested fix
"rerank_cohere.rerank-v3.5": { "object": "rerank-v3.5", "properties": [ + "top_n", + "min_score", "llm.cloud.apikey" ] }, "rerank_cohere.rerank-v3.0": { "object": "rerank-english-v3.0", "properties": [ + "top_n", + "min_score", "llm.cloud.apikey" ] },🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@nodes/src/nodes/rerank_cohere/services.json` around lines 154 - 164, The preset profiles "rerank_cohere.rerank-v3.5" and "rerank_cohere.rerank-v3.0" currently only list "llm.cloud.apikey" so the UI cannot expose rerank tuning knobs; update the service definitions for the symbols "rerank_cohere.rerank-v3.5" and "rerank_cohere.rerank-v3.0" to include the "top_n" and "min_score" properties (same keys used by the existing "custom" profile) so the UI can surface breadth and threshold controls (ensure the property names match the runtime/schema exactly).nodes/src/nodes/rerank_cohere/rerank_client.py (1)
89-96:⚠️ Potential issue | 🟠 MajorRe-validate override values before each Cohere call.
Only the constructor defaults are checked today.
rerank(top_n=0)orrerank(top_n=2.5)still reaches Cohere, andrerank_with_threshold(min_score=1.5)bypasses the local range guard and silently filters everything. The Cohere Rerank API definestop_nas an integer>= 1, so these public override paths need the same validation as config-time values. (docs.cohere.com)Also applies to: 127-135, 174-179
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@nodes/src/nodes/rerank_cohere/rerank_client.py` around lines 89 - 96, The constructor-level checks for model config are fine but you must re-validate public override parameters before each Cohere API call: in the public methods rerank and rerank_with_threshold (and any wrapper that forwards overrides to the Cohere call), assert that top_n is an integer and >= 1 (raise ValueError otherwise) and that min_score is numeric and within [0.0, 1.0] (raise ValueError otherwise); perform these checks at the start of rerank and rerank_with_threshold (and before delegating to any internal helper that calls Cohere) so invalid overrides (e.g., 0, 2.5, 1.5) never reach the API. Ensure error messages include the bad value and expected range/type and keep the same exception type (ValueError) as in the constructor.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@nodes/src/nodes/rerank_cohere/IInstance.py`:
- Around line 111-118: The node currently writes reranked docs with
self.instance.writeDocuments(reranked_docs) but then forwards the Question via
self.instance.writeQuestions(question), which breaks the advertised outputs
(documents and answers); change the emission to populate the advertised answers
lane by constructing the expected Answer payload (e.g., reuse question metadata,
include reranked_docs as the answer content or create an Answer object) and call
self.instance.writeAnswers(answer) instead of writeQuestions, or alternatively
add a "questions" output to services.json and update docs/tests if you intend to
keep writeQuestions; adjust the code paths around reranked_docs, question,
writeDocuments, writeAnswers, and services.json so the emitted lane matches the
node contract.
In `@nodes/src/nodes/rerank_cohere/rerank_client.py`:
- Around line 89-100: The constructor currently accepts an empty string for
model (set via self._model = config.get('model', 'rerank-v3.5')) which defers
failure to runtime; add the same guard used for the API key by validating
self._model after assignment and raise ValueError if it is empty or only
whitespace (e.g., "Cohere model is required" or similar). Update the
RerankClient __init__ validation near the existing apikey check so invalid model
values fail fast locally.
---
Duplicate comments:
In `@nodes/src/nodes/rerank_cohere/rerank_client.py`:
- Around line 89-96: The constructor-level checks for model config are fine but
you must re-validate public override parameters before each Cohere API call: in
the public methods rerank and rerank_with_threshold (and any wrapper that
forwards overrides to the Cohere call), assert that top_n is an integer and >= 1
(raise ValueError otherwise) and that min_score is numeric and within [0.0, 1.0]
(raise ValueError otherwise); perform these checks at the start of rerank and
rerank_with_threshold (and before delegating to any internal helper that calls
Cohere) so invalid overrides (e.g., 0, 2.5, 1.5) never reach the API. Ensure
error messages include the bad value and expected range/type and keep the same
exception type (ValueError) as in the constructor.
In `@nodes/src/nodes/rerank_cohere/services.json`:
- Around line 154-164: The preset profiles "rerank_cohere.rerank-v3.5" and
"rerank_cohere.rerank-v3.0" currently only list "llm.cloud.apikey" so the UI
cannot expose rerank tuning knobs; update the service definitions for the
symbols "rerank_cohere.rerank-v3.5" and "rerank_cohere.rerank-v3.0" to include
the "top_n" and "min_score" properties (same keys used by the existing "custom"
profile) so the UI can surface breadth and threshold controls (ensure the
property names match the runtime/schema exactly).
🪄 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: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: dd1e48c6-940f-4781-8020-f2a528bb0277
📒 Files selected for processing (4)
nodes/src/nodes/rerank_cohere/IGlobal.pynodes/src/nodes/rerank_cohere/IInstance.pynodes/src/nodes/rerank_cohere/rerank_client.pynodes/src/nodes/rerank_cohere/services.json
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@nodes/src/nodes/rerank_cohere/IInstance.py`:
- Around line 49-51: Replace the bare Exception raised when the reranker is
missing with a more specific RuntimeError: in the method accessing
self.IGlobal._reranker (the block that sets reranker = self.IGlobal._reranker
and currently raises Exception('Reranker not initialized')), change the raise to
raise RuntimeError('Reranker not initialized') so callers can catch a more
appropriate built-in exception type.
🪄 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: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: e6a0a527-b8fe-4a04-8f5f-768a25284c72
📒 Files selected for processing (1)
nodes/src/nodes/rerank_cohere/IInstance.py
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
nodes/src/nodes/rerank_cohere/rerank_client.py (1)
89-100:⚠️ Potential issue | 🟠 MajorFail fast on blank
modeland whitespace-onlyapikey.Line 89 currently allows empty/whitespace model values, and Line 99 accepts whitespace-only API keys. Both defer config failures to runtime API calls instead of failing locally at init.
🛠️ Proposed fix
self._model = config.get('model', 'rerank-v3.5') self._top_n = config.get('top_n', 5) self._min_score = config.get('min_score', 0.0) + if not isinstance(self._model, str) or not self._model.strip(): + raise ValueError('Cohere model name is required') if self._top_n < 1: raise ValueError(f'top_n must be >= 1, got {self._top_n}') if not (0.0 <= self._min_score <= 1.0): raise ValueError(f'min_score must be between 0.0 and 1.0, got {self._min_score}') - apikey = config.get('apikey', '') - if not apikey: + apikey = config.get('apikey', '') + if not isinstance(apikey, str) or not apikey.strip(): raise ValueError('Cohere API key is required') - self._client = CohereClient(api_key=apikey) + self._client = CohereClient(api_key=apikey.strip())🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@nodes/src/nodes/rerank_cohere/rerank_client.py` around lines 89 - 100, Validate and fail fast in the constructor: ensure the configured model string is not empty or whitespace-only by stripping config.get('model', 'rerank-v3.5') and raising ValueError if blank (referencing self._model), and normalize/check the API key by stripping apikey = config.get('apikey', '') and raising ValueError if the stripped apikey is empty (referencing the apikey variable); keep the existing checks for self._top_n and self._min_score unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@nodes/src/nodes/rerank_cohere/services.json`:
- Line 55: The services.json entry referencing "cohere.svg" is unresolved and
causes a broken node icon; either add the cohere.svg asset to the shared node
icon directory (ensuring the filename and case match exactly) or update the
"icon" value in nodes/rerank_cohere/services.json to an existing asset name (or
a safe default like the generic node icon) so the node picker can load a valid
image; after changing, verify the asset is included in the build/static assets
so the icon resolves at runtime.
---
Duplicate comments:
In `@nodes/src/nodes/rerank_cohere/rerank_client.py`:
- Around line 89-100: Validate and fail fast in the constructor: ensure the
configured model string is not empty or whitespace-only by stripping
config.get('model', 'rerank-v3.5') and raising ValueError if blank (referencing
self._model), and normalize/check the API key by stripping apikey =
config.get('apikey', '') and raising ValueError if the stripped apikey is empty
(referencing the apikey variable); keep the existing checks for self._top_n and
self._min_score unchanged.
🪄 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: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: b801e566-f76a-4ea5-9c46-3d086a4a7df6
📒 Files selected for processing (4)
nodes/src/nodes/rerank_cohere/IInstance.pynodes/src/nodes/rerank_cohere/rerank_client.pynodes/src/nodes/rerank_cohere/services.jsontest/nodes/test_rerank_cohere.py
| // Optional: | ||
| // The icon is the icon to display in the UI for this node | ||
| // | ||
| "icon": "cohere.svg", |
There was a problem hiding this comment.
cohere.svg icon reference appears unresolved.
Line 55 points to cohere.svg, but the shared node icon directory currently does not include that asset. This can render as a broken/missing icon in the node picker.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@nodes/src/nodes/rerank_cohere/services.json` at line 55, The services.json
entry referencing "cohere.svg" is unresolved and causes a broken node icon;
either add the cohere.svg asset to the shared node icon directory (ensuring the
filename and case match exactly) or update the "icon" value in
nodes/rerank_cohere/services.json to an existing asset name (or a safe default
like the generic node icon) so the node picker can load a valid image; after
changing, verify the asset is included in the build/static assets so the icon
resolves at runtime.
There was a problem hiding this comment.
Fixed in 88a08ba. Added cohere.svg to packages/shared-ui/src/assets/nodes/ and registered it in the iconMap in get-icon-path.ts so the node picker renders it correctly.
There was a problem hiding this comment.
Rate Limit Exceeded
@nihalnihalani have exceeded the limit for the number of chat messages per hour. Please wait 2 minutes and 17 seconds before sending another message.
| # --------------------------------------------------------------------------- | ||
| # Custom exception hierarchy for circuit breaker compatibility. | ||
| # The class names contain 'RateLimit', 'Authentication', etc. so that | ||
| # PR #4's _is_retryable heuristic (which inspects type(exc).__name__) | ||
| # can correctly classify them. | ||
| # --------------------------------------------------------------------------- | ||
|
|
||
|
|
||
| class RerankError(Exception): | ||
| """Base error for rerank operations.""" | ||
|
|
||
| pass | ||
|
|
||
|
|
||
| class RerankAuthenticationError(RerankError): | ||
| """Non-retryable: Invalid API key.""" | ||
|
|
||
| pass | ||
|
|
||
|
|
||
| class RerankRateLimitError(RerankError): | ||
| """Retryable: Rate limit exceeded.""" | ||
|
|
||
| pass | ||
|
|
||
|
|
||
| class RerankBadRequestError(RerankError): | ||
| """Non-retryable: Invalid request parameters.""" | ||
|
|
||
| pass | ||
|
|
||
|
|
||
| class RerankServerError(RerankError): | ||
| """Retryable: Server-side error.""" | ||
|
|
||
| pass |
nihalnihalani
force-pushed
the
feature/cohere-rerank-node
branch
from
863a03b to
4e26890
Compare
github-actions
bot
removed
module:nodes
nihalnihalani
force-pushed
the
feature/cohere-rerank-node
branch
from
4e26890 to
863a03b
Compare
github-actions
bot
added
module:nodes
- Create rerank_cohere node following standard IGlobal/IInstance pattern - Support rerank-v3.5 and rerank-v3.0 models with custom model option - Configurable top_n and min_score threshold for result filtering - Full services.json with profiles, fields, shape, and test cases - Comprehensive unit tests with mocked Cohere client (27 tests) - Improves RAG pipeline quality by reordering retrieved documents by relevance Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- Deep copy question in writeQuestions to prevent mutation in fan-out pipelines - Preserve Cohere exception types with custom error classes for circuit breaker compatibility - Add depends() call in beginGlobal before importing rerank_client - Pin cohere dependency to >=5.0.0,<6.0.0 - Update tests to verify no input mutation and correct exception types Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- Fix index alignment: track original indices when filtering doc_texts so result['index'] maps back to the correct question.documents entry - Always forward the question downstream even when min_score filters out all documents (set question.documents = []) - Validate top_n >= 1 and 0.0 <= min_score <= 1.0 in RerankClient constructor - Change "rerank-v3.0" model name to "rerank-english-v3.0" to match the actual Cohere model identifier Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
…wer for answers lane The previous fix changed writeAnswers to writeQuestions, but "questions" is an input-only lane per services.json. The lanes declaration maps "questions" -> ["documents", "answers"], so writing to "questions" silently dropped data. Construct an Answer object from the reranked document texts and call writeAnswers so downstream nodes on the "answers" lane receive the reranked results. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- Replace bare Exception with RuntimeError for uninitialized reranker
- Always forward answer to answers lane (even when all docs filtered)
- Set question.documents to reranked results for downstream nodes
- Fix tile expression: remove nested ${...} that won't resolve
- Expose top_n and min_score tuning knobs in preset profiles
- Add per-call validation for top_n and min_score overrides
- Reset warning mock between tests to avoid false positives
- Replace invalid-apikey test with empty-model validation test
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
The services.json for rerank_cohere references cohere.svg but the asset did not exist in shared-ui, causing a broken icon in the node picker. Add the SVG and register it in the icon map. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
…ces.json The build was failing with 'Field name model was not found' because the services.json schema didn't expose the 'model' field that the rerank engine expects to read from the configuration. Added the field definition with sensible defaults and description. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
nihalnihalani
force-pushed
the
feature/cohere-rerank-node
branch
from
863a03b to
1eb40e6
Compare
…-org#499) - Add model name validation in RerankClient constructor - Fix test mock target for IGlobal.Config (module vs class shadowing) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Review Feedback AddressedAll critical and major review items have been implemented: Critical Fix
Major Fixes
Additional Fixes
Ready for re-review. |
|
No description provided. |
🚀 Merge RequestAll critical and major review feedback addressed:
This PR is MERGEABLE with no conflicts. Ready for re-review and merge. @asclearuc |
|
@asclearuc — All 10 review items addressed (doc index fix, output lane fix, model ID fix, validation, preset exposure). 35/35 tests passing, MERGEABLE with no conflicts. Ready for re-review. Thanks! 🙏 |
2 participants
#Hack-with-bay-2
Contribution Type
Feature — Add Cohere Rerank pipeline node for RAG quality improvement
Problem / Use Case
RocketRide's RAG pipeline currently retrieves documents via vector similarity search (Pinecone, Chroma, Qdrant, etc.), but vector search alone has a known limitation: it retrieves semantically similar documents but doesn't optimize for relevance to the specific query. This leads to:
Industry research shows that adding a reranking step between retrieval and generation improves answer quality by 15-30%.
Proposed Solution
Create a new
rerank_coherepipeline node that slots into the RAG pipeline between vector search and LLM:RerankClient(rerank_client.py) — Wraps CohereClientV2API withrerank()andrerank_with_threshold()methodsRerankAuthenticationError,RerankRateLimitError,RerankBadRequestError,RerankServerError— naming compatible with circuit breaker (PR feat(nodes): add circuit breaker and retry with backoff for LLM nodes #497) substring heuristicsIGlobal— Standard pattern withvalidateConfigAPI key validation anddepends()inbeginGlobalIInstance— Deep copies question before mutation (safe for fan-out pipelines), reranks documents, filters bymin_scorethresholdservices.json— Profiles for rerank-v3.5 (default), rerank-v3.0, custom. Configurabletop_nandmin_score.Value Added
top_nfiltering reduces token count sent to LLMmin_scorethreshold drops low-relevance documents entirelyWhy This Feature Fits This Codebase
The node follows the exact IGlobal/IInstance pattern used by all 56+ existing nodes. The
services.jsonstructure mirrors other nodes — for example, the profile-based configuration (preconfig.profiles) with conditional field display matches the pattern innodes/src/nodes/llm_openai/services.json(lines 106-160) andnodes/src/nodes/chroma/services.json(lines 188-210).The
IInstance.writeQuestions()method follows the same input/output contract as vector DB nodes: it receives questions with attached documents (from upstream vector search), processes them, and writes to thedocumentsandanswersoutput lanes — the same lane structure defined innodes/src/nodes/chroma/services.jsonlines 80-90.The
depends()call inbeginGlobal()follows the dependency auto-loading pattern established across all nodes (e.g.,nodes/src/nodes/llm_openai/IGlobal.pylines 40-43).Deep copying the question object (
copy.deepcopy(question)) prevents mutation bugs in fan-out pipelines where the same question object is shared across multiple branches — a real concern given the engine's parallel pipe execution model inpackages/server/engine-lib/engLib/python/.Validation
Tests cover: rerank ordering, top_n filtering, min_score thresholds, error handling (auth, rate limit, bad request, server errors), IGlobal lifecycle, IInstance end-to-end, metadata preservation, no-mutation verification, circuit breaker exception compatibility.
How This Could Be Extended
rank_fieldsparameter for structured document rerankingrerank_baseclass for additional reranking providers (Jina, Voyage)🤖 Generated with Claude Code
Summary by CodeRabbit