Init OpensearchVectorClient with os_async_client (#16767)

run-llama · Oct 31, 2024 · 4164d64 · 4164d64
1 parent 4216c88
commit 4164d64
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 2 deletions.
diff --git a/...ex-integrations/vector_stores/llama-index-vector-stores-opensearch/CHANGELOG.md b/...ex-integrations/vector_stores/llama-index-vector-stores-opensearch/CHANGELOG.md
@@ -1,5 +1,9 @@
 # CHANGELOG — llama-index-vector-stores-opensearch
 
+## [0.4.1]
+
+- Added ability to create OpensearchVectorClient with custom os_async_client (like os_client)
+
 ## [0.2.2]
 
 - Fixed issue where Opensearch Serverless does not support painless scripting so handling the case where is_aoss is set and using knn_score script instead.

diff --git a/..._stores/llama-index-vector-stores-opensearch/llama_index/vector_stores/opensearch/base.py b/..._stores/llama-index-vector-stores-opensearch/llama_index/vector_stores/opensearch/base.py
@@ -56,6 +56,8 @@ class OpensearchVectorClient:
         settings: Optional[dict]: Settings for the Opensearch index creation. Defaults to:
             {"index": {"knn": True, "knn.algo_param.ef_search": 100}}
         space_type (Optional[str]): space type for distance metric calculation. Defaults to: l2
+        os_client (Optional[OSClient]): Custom synchronous client (see OpenSearch from opensearch-py)
+        os_async_client (Optional[OSClient]): Custom asynchronous client (see AsyncOpenSearch from opensearch-py)
         **kwargs: Optional arguments passed to the OpenSearch client from opensearch-py.
 
     """
@@ -74,6 +76,7 @@ def __init__(
         max_chunk_bytes: int = 1 * 1024 * 1024,
         search_pipeline: Optional[str] = None,
         os_client: Optional[OSClient] = None,
+        os_async_client: Optional[OSClient] = None,
         **kwargs: Any,
     ):
         """Init params."""
@@ -117,7 +120,7 @@ def __init__(
         self._os_client = os_client or self._get_opensearch_client(
             self._endpoint, **kwargs
         )
-        self._os_async_client = self._get_async_opensearch_client(
+        self._os_async_client = os_async_client or self._get_async_opensearch_client(
             self._endpoint, **kwargs
         )
         self._os_version = self._get_opensearch_version()

diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-opensearch/pyproject.toml b/llama-index-integrations/vector_stores/llama-index-vector-stores-opensearch/pyproject.toml
@@ -27,7 +27,7 @@ exclude = ["**/BUILD"]
 license = "MIT"
 name = "llama-index-vector-stores-opensearch"
 readme = "README.md"
-version = "0.4.0"
+version = "0.4.1"
 
 [tool.poetry.dependencies]
 python = ">=3.8.1,<4.0"