Update opensearch vectorstore to PydanticVectorStore class (#11400)

run-llama · Feb 26, 2024 · 52383c7 · 52383c7
1 parent 4077fee
commit 52383c7
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
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
@@ -2,11 +2,12 @@
 import json
 import uuid
 from typing import Any, Dict, Iterable, List, Optional, Union, cast
+from llama_index.core.bridge.pydantic import PrivateAttr
 
 from llama_index.core.schema import BaseNode, MetadataMode, TextNode
 from llama_index.core.vector_stores.types import (
     MetadataFilters,
-    VectorStore,
+    BasePydanticVectorStore,
     VectorStoreQuery,
     VectorStoreQueryMode,
     VectorStoreQueryResult,
@@ -419,7 +420,7 @@ def query(
         return VectorStoreQueryResult(nodes=nodes, ids=ids, similarities=scores)
 
 
-class OpensearchVectorStore(VectorStore):
+class OpensearchVectorStore(BasePydanticVectorStore):
     """Elasticsearch/Opensearch vector store.
 
     Args:
@@ -428,12 +429,14 @@ class OpensearchVectorStore(VectorStore):
     """
 
     stores_text: bool = True
+    _client: OpensearchVectorClient = PrivateAttr(default=None)
 
     def __init__(
         self,
         client: OpensearchVectorClient,
     ) -> None:
         """Initialize params."""
+        super().__init__()
         self._client = client
 
     @property

diff --git a/...vector_stores/llama-index-vector-stores-opensearch/tests/test_vector_stores_opensearch.py b/...vector_stores/llama-index-vector-stores-opensearch/tests/test_vector_stores_opensearch.py
@@ -1,7 +1,7 @@
-from llama_index.core.vector_stores.types import VectorStore
+from llama_index.core.vector_stores.types import BasePydanticVectorStore
 from llama_index.vector_stores.opensearch import OpensearchVectorStore
 
 
 def test_class():
     names_of_base_classes = [b.__name__ for b in OpensearchVectorStore.__mro__]
-    assert VectorStore.__name__ in names_of_base_classes
+    assert BasePydanticVectorStore.__name__ in names_of_base_classes