llama-index-vector-stores-opensearch 0.2.0__tar.gz → 0.2.2__tar.gz

{llama_index_vector_stores_opensearch-0.2.0 → llama_index_vector_stores_opensearch-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: llama-index-vector-stores-opensearch
-Version: 0.2.0
+Version: 0.2.2
 Summary: llama-index vector_stores opensearch integration
 License: MIT
 Author: Your Name

{llama_index_vector_stores_opensearch-0.2.0 → llama_index_vector_stores_opensearch-0.2.2}/llama_index/vector_stores/opensearch/base.py

@@ -56,6 +56,7 @@ class OpensearchVectorClient:
             This includes engine, metric, and other config params. Defaults to:
             {"name": "hnsw", "space_type": "l2", "engine": "faiss",
             "parameters": {"ef_construction": 256, "m": 48}}
+        space_type (Optional[str]): space type for distance metric calculation. Defaults to: l2
         **kwargs: Optional arguments passed to the OpenSearch client from opensearch-py.
 
     """
@@ -69,6 +70,7 @@ class OpensearchVectorClient:
         text_field: str = "content",
         method: Optional[dict] = None,
         engine: Optional[str] = "nmslib",
+        space_type: Optional[str] = "l2",
         max_chunk_bytes: int = 1 * 1024 * 1024,
         search_pipeline: Optional[str] = None,
         os_client: Optional[OSClient] = None,
@@ -94,6 +96,7 @@ class OpensearchVectorClient:
 
         self._search_pipeline = search_pipeline
         http_auth = kwargs.get("http_auth")
+        self.space_type = space_type
         self.is_aoss = self._is_aoss_enabled(http_auth=http_auth)
         # initialize mapping
         idx_conf = {
@@ -309,9 +312,11 @@ class OpensearchVectorClient:
 
         If there are no filters do approx-knn search.
         If there are (pre)-filters, do an exhaustive exact knn search using 'painless
-            scripting'.
+            scripting' if the version of Opensearch supports it, otherwise uses knn_score scripting score.
 
-        Note that approximate knn search does not support pre-filtering.
+        Note:
+            -AWS Opensearch Serverless does not support the painless scripting functionality at this time according to AWS.
+            -Also note that approximate knn search does not support pre-filtering.
 
         Args:
             query_embedding: Vector embedding to query.
@@ -328,16 +333,25 @@ class OpensearchVectorClient:
             search_query = self._default_approximate_search_query(
                 query_embedding, k, vector_field=embedding_field
             )
+        elif self.is_aoss:
+            # if is_aoss is set we are using Opensearch Serverless AWS offering which cannot use
+            # painless scripting so default scoring script returned will be just normal knn_score script
+            search_query = self._default_scoring_script_query(
+                query_embedding,
+                k,
+                space_type=self.space_type,
+                pre_filter={"bool": {"filter": pre_filter}},
+                vector_field=embedding_field,
+            )
         else:
             # https://opensearch.org/docs/latest/search-plugins/knn/painless-functions/
-            search_query = self._default_painless_scripting_query(
+            search_query = self._default_scoring_script_query(
                 query_embedding,
                 k,
                 space_type="l2Squared",
                 pre_filter={"bool": {"filter": pre_filter}},
                 vector_field=embedding_field,
             )
-
         return search_query
 
     def _hybrid_search_query(
@@ -382,7 +396,9 @@ class OpensearchVectorClient:
     def __get_painless_scripting_source(
         self, space_type: str, vector_field: str = "embedding"
     ) -> str:
-        """For Painless Scripting, it returns the script source based on space type."""
+        """For Painless Scripting, it returns the script source based on space type.
+        This does not work with Opensearch Serverless currently.
+        """
         source_value = (
             f"(1.0 + {space_type}(params.query_value, doc['{vector_field}']))"
         )
@@ -391,7 +407,29 @@ class OpensearchVectorClient:
         else:
             return f"1/{source_value}"
 
-    def _default_painless_scripting_query(
+    def _get_knn_scoring_script(self, space_type, vector_field, query_vector):
+        """Default scoring script that will work with AWS Opensearch Serverless."""
+        return {
+            "source": "knn_score",
+            "lang": "knn",
+            "params": {
+                "field": vector_field,
+                "query_value": query_vector,
+                "space_type": space_type,
+            },
+        }
+
+    def _get_painless_scoring_script(self, space_type, vector_field, query_vector):
+        source = self.__get_painless_scripting_source(space_type, vector_field)
+        return {
+            "source": source,
+            "params": {
+                "field": vector_field,
+                "query_value": query_vector,
+            },
+        }
+
+    def _default_scoring_script_query(
         self,
         query_vector: List[float],
         k: int = 4,
@@ -399,23 +437,31 @@ class OpensearchVectorClient:
         pre_filter: Optional[Union[Dict, List]] = None,
         vector_field: str = "embedding",
     ) -> Dict:
-        """For Painless Scripting Search, this is the default query."""
+        """For Scoring Script Search, this is the default query. Has to account for Opensearch Service
+        Serverless which does not support painless scripting functions so defaults to knn_score.
+        """
         if not pre_filter:
             pre_filter = MATCH_ALL_QUERY
 
-        source = self.__get_painless_scripting_source(space_type, vector_field)
+        # check if we can use painless scripting or have to use default knn_score script
+        if self.is_aoss:
+            if space_type == "l2Squared":
+                raise ValueError(
+                    "Unsupported space type for aoss. Can only use l1, l2, cosinesimil."
+                )
+            script = self._get_knn_scoring_script(
+                space_type, vector_field, query_vector
+            )
+        else:
+            script = self._get_painless_scoring_script(
+                space_type, vector_field, query_vector
+            )
         return {
             "size": k,
             "query": {
                 "script_score": {
                     "query": pre_filter,
-                    "script": {
-                        "source": source,
-                        "params": {
-                            "field": vector_field,
-                            "query_value": query_vector,
-                        },
-                    },
+                    "script": script,
                 }
             },
         }
@@ -466,7 +512,9 @@ class OpensearchVectorClient:
         search_query = {
             "query": {"term": {"metadata.doc_id.keyword": {"value": doc_id}}}
         }
-        await self._os_client.delete_by_query(index=self._index, body=search_query)
+        await self._os_client.delete_by_query(
+            index=self._index, body=search_query, refresh=True
+        )
 
     async def delete_nodes(
         self,
@@ -490,12 +538,16 @@ class OpensearchVectorClient:
         if filters:
             query["query"]["bool"]["filter"].extend(self._parse_filters(filters))
 
-        await self._os_client.delete_by_query(index=self._index, body=query)
+        await self._os_client.delete_by_query(
+            index=self._index, body=query, refresh=True
+        )
 
     async def clear(self) -> None:
         """Clears index."""
         query = {"query": {"bool": {"filter": []}}}
-        await self._os_client.delete_by_query(index=self._index, body=query)
+        await self._os_client.delete_by_query(
+            index=self._index, body=query, refresh=True
+        )
 
     async def aquery(
         self,

{llama_index_vector_stores_opensearch-0.2.0 → llama_index_vector_stores_opensearch-0.2.2}/pyproject.toml

@@ -27,7 +27,7 @@ exclude = ["**/BUILD"]
 license = "MIT"
 name = "llama-index-vector-stores-opensearch"
 readme = "README.md"
-version = "0.2.0"
+version = "0.2.2"
 
 [tool.poetry.dependencies]
 python = ">=3.8.1,<4.0"