Add Cohere Reranker #269
Add Cohere Reranker #269
Conversation
| reranked_query_results: List[KBQueryResult] = [] | ||
| for result in results: | ||
| texts = [doc.text for doc in result.documents] | ||
| response = self._client.rerank(query=result.query, |
There was a problem hiding this comment.
Also needs to be wrapped in try clause.
Transient errors like rate limits etc should be retried (if the Cohere client itself doesn't do that for us already).
Errors that are caused by wrong configuration (like wrong model name or bad API key) need to be re-raised with an actionable error message
There was a problem hiding this comment.
Cohere retries internally. Since Cohere does not return different error types it is hard to understand what the message is. For now I am raising a RuntimeError from the actual error.
|
|
||
| def test_bad_api_key(should_run_test, query_result): | ||
| from cohere import CohereAPIError | ||
| with pytest.raises(CohereAPIError, match="invalid api token"): |
There was a problem hiding this comment.
We try to eliminate underlying service's errors like CohereAPIError or OpenAIError, and replace them with actionable error message (like something the user needs to change in the Canopy config, or the explicit env var to set).
In the future we will have our own error types like EncoderError, AuthenticationError etc. In the meantime simply re-raise RuntimeError for all of these cases (the CLI catches RuntimeError and prints them nicely)
There was a problem hiding this comment.
I checked the client, client does not return a specific error for different errors we always get a CohereAPIError. For now I am raising RuntimeError from that error, if they improve the client we can write actionable messages.
| from cohere import CohereAPIError | ||
| with pytest.raises(CohereAPIError, match="invalid api token"): | ||
| CohereReranker(api_key="bad key").rerank([query_result]) | ||
|
|
There was a problem hiding this comment.
Missing more negative tests - wrong model name, bad input (e.g. not strings) etc.
There was a problem hiding this comment.
Added wrong model name, bad input is not possible since we validate our data with pydantic.
github-merge-queue
bot
removed this pull request from the merge queue due to failed status checks
| ) for r in results | ||
| for d in rr.documents | ||
| ], | ||
| debug_info={"db_result": QueryResult( |
There was a problem hiding this comment.
I just relised that we want debug info to be only dicts with literals like str or int. This allows easier serialisation of this object
Problem
We currently do not rerank results. Reranking the results can result in better quality responses.
Solution
Added Cohere reranker
Type of Change
Test Plan
Added the relevant tests.