Native IRIS Vector datatype

Author: daimor

docker-compose.yml

@@ -1,7 +1,8 @@
 version: "3"
 services:
   iris:
-    image: intersystemsdc/iris-community:2023.3-zpm
+    # image: intersystemsdc/iris-community:2023.3-zpm
+    image: caretdev/iris-community:2024.1-vecdb
     ports:
       - 6172:1972
       - 6173:52773

langchain_iris/vectorstores.py

@@ -40,6 +40,7 @@
     func,
 )
 from sqlalchemy_iris import IRISListBuild
+from sqlalchemy_iris import IRISVector as IRISVectorType
 
 from sqlalchemy.orm import Session
 
@@ -75,10 +76,13 @@ class DistanceStrategy(str, enum.Enum):
 
 class IRISVector(VectorStore):
     _conn = None
+    native_vector = False
+    native_vector_cosine_similarity = False
 
     def __init__(
         self,
         embedding_function: Embeddings,
+        dimension: int,
         connection_string: Optional[str] = None,
         collection_name: str = _LANGCHAIN_DEFAULT_COLLECTION_NAME,
         pre_delete_collection: bool = False,
@@ -91,6 +95,7 @@ def __init__(
     ) -> None:
         self.connection_string = connection_string or "iris+emb:///"
         self.embedding_function = embedding_function
+        self.dimension = dimension
         self.collection_name = collection_name
         self.pre_delete_collection = pre_delete_collection
         self.collection_metadata = collection_metadata
@@ -113,6 +118,8 @@ def __post_init__(
         self.create_vector_functions()
 
     def create_vector_functions(self) -> None:
+        if self.native_vector:
+            return
         try:
             with Session(self._conn) as session:
                 session.execute(
@@ -188,6 +195,19 @@ def create_vector_functions(self) -> None:
 
     @property
     def distance_strategy(self) -> str:
+        if self.native_vector:
+            if self._distance_strategy == DistanceStrategy.COSINE:
+                return self.table.c.embedding.cosine
+            elif self._distance_strategy == DistanceStrategy.MAX_INNER_PRODUCT:
+                return self.table.c.embedding.max_inner_product
+            # elif self._distance_strategy == DistanceStrategy.EUCLIDEAN:
+            #     return "langchain_l2_distance"
+            else:
+                raise ValueError(
+                    f"Got unexpected value for distance: {self._distance_strategy}. "
+                    f"Should be one of {', '.join([ds.value for ds in DistanceStrategy])}."
+                )
+
         if self._distance_strategy == DistanceStrategy.EUCLIDEAN:
             return "langchain_l2_distance"
         elif self._distance_strategy == DistanceStrategy.COSINE:
@@ -203,6 +223,14 @@ def distance_strategy(self) -> str:
     def connect(self) -> Connection:
         engine = create_engine(self.connection_string, **self.engine_args)
         conn = engine.connect()
+        try:
+            if conn.dialect.supports_vectors:
+                self.native_vector = True
+                self.native_vector_cosine_similarity = (
+                    conn.dialect.vector_cosine_similarity
+                )
+        except:  # noqa
+            pass
         return conn
 
     def __del__(self) -> None:
@@ -220,7 +248,14 @@ def table(self) -> Table:
             self.collection_name,
             Base.metadata,
             Column("id", VARCHAR(40), primary_key=True, default=uuid.uuid4),
-            Column("embedding", IRISListBuild(16000, float)),
+            Column(
+                "embedding",
+                (
+                    IRISVectorType(self.dimension)
+                    if self.native_vector
+                    else IRISListBuild(self.dimension, float)
+                ),
+            ),
             Column("document", TEXT, nullable=True),
             Column("metadata", TEXT, nullable=True),
             extend_existing=True,
@@ -278,6 +313,13 @@ def _select_relevance_score_fn(self) -> Callable[[float], float]:
                 "Consider providing relevance_score_fn to IRISVector constructor."
             )
 
+    @staticmethod
+    def _cosine_relevance_score_fn(distance: float) -> float:
+        print('_cosine_relevance_score_fn', distance)
+        """Normalize the distance to a score on a scale [0, 1]."""
+
+        return round(1.0 - distance, 15)
+
     @classmethod
     def from_embeddings(
         cls: Type[IRISVector],
@@ -299,8 +341,11 @@ def from_embeddings(
         texts = [t[0] for t in text_embeddings]
         embeddings = [t[1] for t in text_embeddings]
 
+        dimension = len(embeddings[0])
+
         store = cls(
             collection_name=collection_name,
+            dimension=dimension,
             distance_strategy=distance_strategy,
             embedding_function=embedding,
             pre_delete_collection=pre_delete_collection,
@@ -330,8 +375,12 @@ def from_texts(
         Return VectorStore initialized from texts and embeddings.
         """
 
+        sample_embedding = embedding.embed_query("Hello IRISVector!")
+        dimension = len(sample_embedding)
+
         store = cls(
             collection_name=collection_name,
+            dimension=dimension,
             distance_strategy=distance_strategy,
             embedding_function=embedding,
             pre_delete_collection=pre_delete_collection,
@@ -483,9 +532,13 @@ def similarity_search_with_score_by_vector(
             results: Sequence[Row] = (
                 session.query(
                     self.table,
-                    self.table.c.embedding.func(
-                        self.distance_strategy, embedding
-                    ).label("distance"),
+                    (
+                        self.distance_strategy(embedding).label("distance")
+                        if self.native_vector
+                        else self.table.c.embedding.func(
+                            self.distance_strategy, embedding
+                        ).label("distance")
+                    ),
                 )
                 .filter(filter_by)
                 .order_by(asc("distance"))
@@ -499,7 +552,7 @@ def similarity_search_with_score_by_vector(
                     page_content=result.document,
                     metadata=json.loads(result.metadata),
                 ),
-                float(result.distance) if self.embedding_function is not None else None,
+                round(float(result.distance), 15) if self.embedding_function is not None else None,
             )
             for result in results
         ]

setup.py

@@ -7,7 +7,7 @@
     ),
     install_requires=[
         "langchain==0.0.348",
-        "sqlalchemy-iris>=0.12.0",
+        "sqlalchemy-iris>=0.13.0",
     ],
     python_requires=">3.7,<3.12",
 )

tests/test_vectorstores.py

@@ -8,7 +8,7 @@
 from langchain.docstore.document import Document
 from sqlalchemy.orm import Session
 
-from langchain_iris import IRISVector, DistanceStrategy
+from langchain_iris import IRISVector
 from langchain.embeddings.fake import DeterministicFakeEmbedding
 
 
@@ -138,7 +138,7 @@ def test_irisvector_with_filter_distant_match(
     )
     output = docsearch.similarity_search_with_score("foo", k=1, filter={"page": "2"})
     assert output == [
-        (Document(page_content="baz", metadata={"page": "2"}), 0.0013003906671379406)
+        (Document(page_content="baz", metadata={"page": "2"}), 0.001300390667138)
     ]
 
 
@@ -232,13 +232,10 @@ def test_irisvector_relevance_score(collection_name, connection_string) -> None:
     )
 
     output = docsearch.similarity_search_with_relevance_scores("foo", k=3)
-    print([[doc.page_content, score] for doc, score in output])
-    # [1.0, 0.07647878038412692, -0.10694835587301355]
-    # [1.0, 0.07647878038412692, -0.10694835587301355]
     assert output == [
         (Document(page_content="foo", metadata={"page": "0"}), 1.0),
-        (Document(page_content="bar", metadata={"page": "1"}), 0.9996744261675065),
-        (Document(page_content="baz", metadata={"page": "2"}), 0.9986996093328621),
+        (Document(page_content="bar", metadata={"page": "1"}), 0.999674426167506),
+        (Document(page_content="baz", metadata={"page": "2"}), 0.998699609332862),
     ]