Remove dependency to langchain-community (#138)

* Remove dependency to langchain-community

* increase AstraDB client request timeout for integration testing

---------

Co-authored-by: Stefano Lottini <[email protected]>

Author: cbornet

libs/astradb/langchain_astradb/vectorstores.py

@@ -28,7 +28,6 @@
     CollectionRerankOptions,
     VectorServiceOptions,
 )
-from langchain_community.vectorstores.utils import maximal_marginal_relevance
 from langchain_core.runnables.utils import gather_with_concurrency
 from langchain_core.vectorstores import VectorStore
 from typing_extensions import override
@@ -56,6 +55,14 @@
     _DefaultVSDocumentCodec,
 )
 
+is_simd_available: bool = False
+try:
+    import simsimd as simd
+
+    is_simd_available = True
+except ImportError:
+    pass
+
 if TYPE_CHECKING:
     from collections.abc import AsyncIterable, Awaitable, Iterable, Sequence
 
@@ -310,6 +317,75 @@ def _describe_error(_errd: Exception) -> list[str]:
     return err_msg
 
 
+_Matrix = Union[list[list[float]], list[np.ndarray], np.ndarray]
+
+
+def _cosine_similarity(x: _Matrix, y: _Matrix) -> np.ndarray:
+    """Row-wise cosine similarity between two equal-width matrices."""
+    if len(x) == 0 or len(y) == 0:
+        return np.array([])
+
+    x = np.array(x)
+    y = np.array(y)
+    if x.shape[1] != y.shape[1]:
+        msg = (
+            f"Number of columns in X and Y must be the same. X has shape {x.shape} "
+            f"and Y has shape {y.shape}."
+        )
+        raise ValueError(msg)
+
+    if is_simd_available:
+        x = np.array(x, dtype=np.float32)
+        y = np.array(y, dtype=np.float32)
+        return 1 - np.array(simd.cdist(x, y, metric="cosine"))
+
+    logger.debug(
+        "Unable to use simsimd, defaulting to NumPy implementation. If you want "
+        "to use simsimd please install with `pip install simsimd`."
+    )
+    x_norm = np.linalg.norm(x, axis=1)
+    y_norm = np.linalg.norm(y, axis=1)
+    # Ignore divide by zero errors run time warnings as those are handled below.
+    with np.errstate(divide="ignore", invalid="ignore"):
+        similarity: np.ndarray = np.dot(x, y.T) / np.outer(x_norm, y_norm)
+    similarity[np.isnan(similarity) | np.isinf(similarity)] = 0.0
+    return similarity
+
+
+def _maximal_marginal_relevance(
+    query_embedding: np.ndarray,
+    embedding_list: list[list[float]],
+    lambda_mult: float = 0.5,
+    k: int = 4,
+) -> list[int]:
+    """Calculate maximal marginal relevance."""
+    if min(k, len(embedding_list)) <= 0:
+        return []
+    if query_embedding.ndim == 1:
+        query_embedding = np.expand_dims(query_embedding, axis=0)
+    similarity_to_query = _cosine_similarity(query_embedding, embedding_list)[0]
+    most_similar = int(np.argmax(similarity_to_query))
+    idxs = [most_similar]
+    selected = np.array([embedding_list[most_similar]])
+    while len(idxs) < min(k, len(embedding_list)):
+        best_score = -np.inf
+        idx_to_add = -1
+        similarity_to_selected = _cosine_similarity(embedding_list, selected)
+        for i, query_score in enumerate(similarity_to_query):
+            if i in idxs:
+                continue
+            redundant_score = max(similarity_to_selected[i])
+            equation_score = (
+                lambda_mult * query_score - (1 - lambda_mult) * redundant_score
+            )
+            if equation_score > best_score:
+                best_score = equation_score
+                idx_to_add = i
+        idxs.append(idx_to_add)
+        selected = np.append(selected, [embedding_list[idx_to_add]], axis=0)
+    return idxs
+
+
 class AstraDBVectorStoreError(Exception):
     """An exception during vector-store activities.
 
@@ -3308,9 +3384,9 @@ def _get_mmr_hits(
         lambda_mult: float,
         prefetch_hit_pairs: list[tuple[Document, list[float]]],
     ) -> list[Document]:
-        mmr_chosen_indices = maximal_marginal_relevance(
+        mmr_chosen_indices = _maximal_marginal_relevance(
             np.array(embedding, dtype=np.float32),
-            [hit_pair[1] for hit_pair in prefetch_hit_pairs],
+            [embedding for _, embedding in prefetch_hit_pairs],
             k=k,
             lambda_mult=lambda_mult,
         )