Use orchestration lib (#22)

* Utilize VectorStore from Elasticsearch client lib

* remove non-existing dependency group installation

* type ignores for mocking methods

* bring back integration_test dependency group

* bring back ElasticsearchStore.connect_to_elasticsearch

* make custom_query and doc_builder kw-only

* keep original strategies

* deprecate old retrieval classes

* bring back previous integration tests

Author: maxjakob

libs/elasticsearch/langchain_elasticsearch/__init__.py

@@ -1,21 +1,40 @@
+from elasticsearch.helpers.vectorstore import (
+    BM25Strategy,
+    DenseVectorScriptScoreStrategy,
+    DenseVectorStrategy,
+    DistanceMetric,
+    RetrievalStrategy,
+    SparseVectorStrategy,
+)
+
 from langchain_elasticsearch.cache import ElasticsearchCache
 from langchain_elasticsearch.chat_history import ElasticsearchChatMessageHistory
 from langchain_elasticsearch.embeddings import ElasticsearchEmbeddings
 from langchain_elasticsearch.retrievers import ElasticsearchRetriever
 from langchain_elasticsearch.vectorstores import (
     ApproxRetrievalStrategy,
+    BM25RetrievalStrategy,
     ElasticsearchStore,
     ExactRetrievalStrategy,
     SparseRetrievalStrategy,
 )
 
 __all__ = [
-    "ApproxRetrievalStrategy",
     "ElasticsearchCache",
     "ElasticsearchChatMessageHistory",
     "ElasticsearchEmbeddings",
     "ElasticsearchRetriever",
     "ElasticsearchStore",
+    # retrieval strategies
+    "BM25Strategy",
+    "DenseVectorScriptScoreStrategy",
+    "DenseVectorStrategy",
+    "DistanceMetric",
+    "RetrievalStrategy",
+    "SparseVectorStrategy",
+    # deprecated retrieval strategies
+    "ApproxRetrievalStrategy",
+    "BM25RetrievalStrategy",
     "ExactRetrievalStrategy",
     "SparseRetrievalStrategy",
 ]

libs/elasticsearch/langchain_elasticsearch/_utilities.py

@@ -1,12 +1,8 @@
 from enum import Enum
-from typing import List, Union
 
-import numpy as np
 from elasticsearch import BadRequestError, ConflictError, Elasticsearch, NotFoundError
 from langchain_core import __version__ as langchain_version
 
-Matrix = Union[List[List[float]], List[np.ndarray], np.ndarray]
-
 
 class DistanceStrategy(str, Enum):
     """Enumerator of the Distance strategies for calculating distances
@@ -19,77 +15,16 @@ class DistanceStrategy(str, Enum):
     COSINE = "COSINE"
 
 
+def user_agent(prefix: str) -> str:
+    return f"{prefix}/{langchain_version}"
+
+
 def with_user_agent_header(client: Elasticsearch, header_prefix: str) -> Elasticsearch:
     headers = dict(client._headers)
-    headers.update({"user-agent": f"{header_prefix}/{langchain_version}"})
+    headers.update({"user-agent": f"{user_agent(header_prefix)}"})
     return client.options(headers=headers)
 
 
-def maximal_marginal_relevance(
-    query_embedding: np.ndarray,
-    embedding_list: list,
-    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
-
-
-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]:
-        raise ValueError(
-            f"Number of columns in X and Y must be the same. X has shape {X.shape} "
-            f"and Y has shape {Y.shape}."
-        )
-    try:
-        import simsimd as simd  # type: ignore
-
-        X = np.array(X, dtype=np.float32)
-        Y = np.array(Y, dtype=np.float32)
-        Z = 1 - simd.cdist(X, Y, metric="cosine")
-        if isinstance(Z, float):
-            return np.array([Z])
-        return np.array(Z)
-    except ImportError:
-        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.dot(X, Y.T) / np.outer(X_norm, Y_norm)
-        similarity[np.isnan(similarity) | np.isinf(similarity)] = 0.0
-        return similarity
-
-
 def model_must_be_deployed(client: Elasticsearch, model_id: str) -> None:
     try:
         dummy = {"x": "y"}
@@ -106,11 +41,3 @@ def model_must_be_deployed(client: Elasticsearch, model_id: str) -> None:
         # This error is expected because we do not know the expected document
         # shape and just use a dummy doc above.
         pass
-
-
-def model_is_deployed(es_client: Elasticsearch, model_id: str) -> bool:
-    try:
-        model_must_be_deployed(es_client, model_id)
-        return True
-    except NotFoundError:
-        return False

libs/elasticsearch/langchain_elasticsearch/embeddings.py

@@ -3,6 +3,7 @@
 from typing import TYPE_CHECKING, List, Optional
 
 from elasticsearch import Elasticsearch
+from elasticsearch.helpers.vectorstore import EmbeddingService
 from langchain_core.embeddings import Embeddings
 from langchain_core.utils import get_from_env
 
@@ -206,3 +207,45 @@ def embed_query(self, text: str) -> List[float]:
             List[float]: The embedding for the input query text.
         """
         return self._embedding_func([text])[0]
+
+
+class EmbeddingServiceAdapter(EmbeddingService):
+    """
+    Adapter for LangChain Embeddings to support the EmbeddingService interface from
+    elasticsearch.helpers.vectorstore.
+    """
+
+    def __init__(self, langchain_embeddings: Embeddings):
+        self._langchain_embeddings = langchain_embeddings
+
+    def __eq__(self, other):  # type: ignore[no-untyped-def]
+        if isinstance(other, self.__class__):
+            return self.__dict__ == other.__dict__
+        else:
+            return False
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        """
+        Generate embeddings for a list of documents.
+
+        Args:
+            texts (List[str]): A list of document text strings to generate embeddings
+                for.
+
+        Returns:
+            List[List[float]]: A list of embeddings, one for each document in the input
+                list.
+        """
+        return self._langchain_embeddings.embed_documents(texts)
+
+    def embed_query(self, text: str) -> List[float]:
+        """
+        Generate an embedding for a single query text.
+
+        Args:
+            text (str): The query text to generate an embedding for.
+
+        Returns:
+            List[float]: The embedding for the input query text.
+        """
+        return self._langchain_embeddings.embed_query(text)