diff --git a/rag/utils/opensearch_conn.py b/rag/utils/opensearch_conn.py index 2239102ef3..e9909d978c 100644 --- a/rag/utils/opensearch_conn.py +++ b/rag/utils/opensearch_conn.py @@ -666,6 +666,23 @@ class OSConnection(DocStoreConnection): def get_doc_ids(self, res): return [d["_id"] for d in res["hits"]["hits"]] + def get_scores(self, res) -> dict[str, float]: + """ + Map hit `_id` to its raw `_score`. Used by rag/nlp/search.py:_knn_scores() + to recover the cosine similarity returned by a KNN-only second-pass search + without pulling the chunk vectors out of the index. OpenSearch hit headers + carry `_score` exactly like Elasticsearch, so this mirrors + ESConnectionBase.get_scores. + """ + out = {} + for d in res.get("hits", {}).get("hits", []): + doc_id = d.get("_id") + if doc_id is None: + continue + score = d.get("_score") + out[doc_id] = float(score) if score is not None else 0.0 + return out + def __getSource(self, res): rr = [] for d in res["hits"]["hits"]: