diff --git a/api/core/model_runtime/model_providers/jina/rerank/rerank.py b/api/core/model_runtime/model_providers/jina/rerank/rerank.py index 79ca68914f..0350207651 100644 --- a/api/core/model_runtime/model_providers/jina/rerank/rerank.py +++ b/api/core/model_runtime/model_providers/jina/rerank/rerank.py @@ -61,11 +61,19 @@ class JinaRerankModel(RerankModel): rerank_documents = [] for result in results["results"]: + index = result["index"] + if "document" in result: + text = result["document"]["text"] + else: + # llama.cpp rerank maynot return original documents + text = docs[index] + rerank_document = RerankDocument( - index=result["index"], - text=result["document"]["text"], + index=index, + text=text, score=result["relevance_score"], ) + if score_threshold is None or result["relevance_score"] >= score_threshold: rerank_documents.append(rerank_document) diff --git a/api/core/model_runtime/model_providers/localai/rerank/rerank.py b/api/core/model_runtime/model_providers/localai/rerank/rerank.py index 2b0f53bc19..075b44658d 100644 --- a/api/core/model_runtime/model_providers/localai/rerank/rerank.py +++ b/api/core/model_runtime/model_providers/localai/rerank/rerank.py @@ -70,11 +70,19 @@ class LocalaiRerankModel(RerankModel): rerank_documents = [] for result in results["results"]: + index = result["index"] + if "document" in result: + text = result["document"]["text"] + else: + # llama.cpp rerank maynot return original documents + text = docs[index] + rerank_document = RerankDocument( - index=result["index"], - text=result["document"]["text"], + index=index, + text=text, score=result["relevance_score"], ) + if score_threshold is None or result["relevance_score"] >= score_threshold: rerank_documents.append(rerank_document)