Don't hardcode PGVector distance strategies (langchain-ai#10265)

- Description: Remove hardcoded/duplicated distance strategies in the PGVector store. - Issue: NA - Dependencies: NA - Tag maintainer: for a quicker response, tag the relevant maintainer (see below), - Twitter handle: @archmonkeymojo --------- Co-authored-by: Bagatur <[email protected]>
getjavelin · Sep 7, 2023 · e220fa9 · e220fa9
1 parent f081b27
commit e220fa9
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/libs/langchain/langchain/vectorstores/pgvector.py b/libs/langchain/langchain/vectorstores/pgvector.py
@@ -349,16 +349,16 @@ def similarity_search_with_score(
 
     @property
     def distance_strategy(self) -> Any:
-        if self._distance_strategy == "l2":
+        if self._distance_strategy == DistanceStrategy.EUCLIDEAN:
             return self.EmbeddingStore.embedding.l2_distance
-        elif self._distance_strategy == "cosine":
+        elif self._distance_strategy == DistanceStrategy.COSINE:
             return self.EmbeddingStore.embedding.cosine_distance
-        elif self._distance_strategy == "inner":
+        elif self._distance_strategy == DistanceStrategy.MAX_INNER_PRODUCT:
             return self.EmbeddingStore.embedding.max_inner_product
         else:
             raise ValueError(
                 f"Got unexpected value for distance: {self._distance_strategy}. "
-                f"Should be one of `l2`, `cosine`, `inner`."
+                f"Should be one of {', '.join([ds.value for ds in DistanceStrategy])}."
             )
 
     def similarity_search_with_score_by_vector(