add get_top_k_cosine_similarity method to get max top k score and index (#5059)

# Row-wise cosine similarity between two equal-width matrices and return
the max top_k score and index, the score all greater than
threshold_score.

Co-authored-by: Dev 2049 <dev.dev2049@gmail.com>

tests/unit_tests/test_math_utils.py

@@ -2,8 +2,19 @@
 from typing import List
 
 import numpy as np
+import pytest
 
-from langchain.math_utils import cosine_similarity
+from langchain.math_utils import cosine_similarity, cosine_similarity_top_k
+
+
+@pytest.fixture
+def X() -> List[List[float]]:
+    return [[1.0, 2.0, 3.0], [0.0, 1.0, 0.0], [1.0, 2.0, 0.0]]
+
+
+@pytest.fixture
+def Y() -> List[List[float]]:
+    return [[0.5, 1.0, 1.5], [1.0, 0.0, 0.0], [2.0, 5.0, 2.0], [0.0, 0.0, 0.0]]
 
 
 def test_cosine_similarity_zero() -> None:
@@ -27,13 +38,41 @@ def test_cosine_similarity_empty() -> None:
     assert len(cosine_similarity(empty_list, np.random.random((3, 3)))) == 0
 
 
-def test_cosine_similarity() -> None:
-    X = [[1.0, 2.0, 3.0], [0.0, 1.0, 0.0], [1.0, 2.0, 0.0]]
-    Y = [[0.5, 1.0, 1.5], [1.0, 0.0, 0.0], [2.0, 5.0, 2.0]]
+def test_cosine_similarity(X: List[List[float]], Y: List[List[float]]) -> None:
     expected = [
-        [1.0, 0.26726124, 0.83743579],
-        [0.53452248, 0.0, 0.87038828],
-        [0.5976143, 0.4472136, 0.93419873],
+        [1.0, 0.26726124, 0.83743579, 0.0],
+        [0.53452248, 0.0, 0.87038828, 0.0],
+        [0.5976143, 0.4472136, 0.93419873, 0.0],
     ]
     actual = cosine_similarity(X, Y)
     assert np.allclose(expected, actual)
+
+
+def test_cosine_similarity_top_k(X: List[List[float]], Y: List[List[float]]) -> None:
+    expected_idxs = [(0, 0), (2, 2), (1, 2), (0, 2), (2, 0)]
+    expected_scores = [1.0, 0.93419873, 0.87038828, 0.83743579, 0.5976143]
+    actual_idxs, actual_scores = cosine_similarity_top_k(X, Y)
+    assert actual_idxs == expected_idxs
+    assert np.allclose(expected_scores, actual_scores)
+
+
+def test_cosine_similarity_score_threshold(
+    X: List[List[float]], Y: List[List[float]]
+) -> None:
+    expected_idxs = [(0, 0), (2, 2)]
+    expected_scores = [1.0, 0.93419873]
+    actual_idxs, actual_scores = cosine_similarity_top_k(
+        X, Y, top_k=None, score_threshold=0.9
+    )
+    assert actual_idxs == expected_idxs
+    assert np.allclose(expected_scores, actual_scores)
+
+
+def test_cosine_similarity_top_k_and_score_threshold(
+    X: List[List[float]], Y: List[List[float]]
+) -> None:
+    expected_idxs = [(0, 0), (2, 2), (1, 2), (0, 2)]
+    expected_scores = [1.0, 0.93419873, 0.87038828, 0.83743579]
+    actual_idxs, actual_scores = cosine_similarity_top_k(X, Y, score_threshold=0.8)
+    assert actual_idxs == expected_idxs
+    assert np.allclose(expected_scores, actual_scores)