diff --git a/libs/vertexai/langchain_google_vertexai/vectorstores/_searcher.py b/libs/vertexai/langchain_google_vertexai/vectorstores/_searcher.py index 783de5ac..96b2e041 100644 --- a/libs/vertexai/langchain_google_vertexai/vectorstores/_searcher.py +++ b/libs/vertexai/langchain_google_vertexai/vectorstores/_searcher.py @@ -35,7 +35,9 @@ def find_neighbors( filter_: Union[List[Namespace], None] = None, numeric_filter: Union[List[NumericNamespace], None] = None, *, - sparse_embeddings: Optional[List[Dict[str, List[int] | List[float]]]] = None, + sparse_embeddings: Optional[ + List[Dict[str, Union[List[int], List[float]]]] + ] = None, rrf_ranking_alpha: float = 1, ) -> List[List[Dict[str, Any]]]: """Finds the k closes neighbors of each instance of embeddings. @@ -70,7 +72,9 @@ def add_to_index( metadatas: Union[List[dict], None] = None, is_complete_overwrite: bool = False, *, - sparse_embeddings: Optional[List[Dict[str, List[int] | List[float]]]] = None, + sparse_embeddings: Optional[ + List[Dict[str, Union[List[int], List[float]]]] + ] = None, **kwargs: Any, ) -> None: """Adds documents to the index. @@ -191,7 +195,9 @@ def add_to_index( metadatas: Union[List[dict], None] = None, is_complete_overwrite: bool = False, *, - sparse_embeddings: Optional[List[Dict[str, List[int] | List[float]]]] = None, + sparse_embeddings: Optional[ + List[Dict[str, Union[List[int], List[float]]]] + ] = None, **kwargs: Any, ) -> None: """Adds documents to the index. @@ -233,7 +239,9 @@ def find_neighbors( filter_: Union[List[Namespace], None] = None, numeric_filter: Union[List[NumericNamespace], None] = None, *, - sparse_embeddings: Optional[List[Dict[str, List[int] | List[float]]]] = None, + sparse_embeddings: Optional[ + List[Dict[str, Union[List[int], List[float]]]] + ] = None, rrf_ranking_alpha: float = 1, ) -> List[List[Dict[str, Any]]]: """Finds the k closes neighbors of each instance of embeddings. diff --git a/libs/vertexai/langchain_google_vertexai/vectorstores/_utils.py b/libs/vertexai/langchain_google_vertexai/vectorstores/_utils.py index 541c931c..dbd76970 100644 --- a/libs/vertexai/langchain_google_vertexai/vectorstores/_utils.py +++ b/libs/vertexai/langchain_google_vertexai/vectorstores/_utils.py @@ -65,7 +65,7 @@ def batch_update_index( def to_data_points( ids: List[str], embeddings: List[List[float]], - sparse_embeddings: Optional[List[Dict[str, List[int] | List[float]]]] = None, + sparse_embeddings: Optional[List[Dict[str, Union[List[int], List[float]]]]] = None, metadatas: Union[List[Dict[str, Any]], None] = None, ) -> List["meidx_types.IndexDataPoint"]: """Converts triplets id, embedding, metadata into IndexDataPoints instances. diff --git a/libs/vertexai/langchain_google_vertexai/vectorstores/vectorstores.py b/libs/vertexai/langchain_google_vertexai/vectorstores/vectorstores.py index 70b4dc30..9d2b6160 100644 --- a/libs/vertexai/langchain_google_vertexai/vectorstores/vectorstores.py +++ b/libs/vertexai/langchain_google_vertexai/vectorstores/vectorstores.py @@ -278,7 +278,9 @@ def add_texts_with_embeddings( embeddings: List[List[float]], metadatas: Union[List[dict], None] = None, *, - sparse_embeddings: Optional[List[Dict[str, List[int] | List[float]]]] = None, + sparse_embeddings: Optional[ + List[Dict[str, Union[List[int], List[float]]]] + ] = None, ids: Optional[List[str]] = None, is_complete_overwrite: bool = False, **kwargs: Any, diff --git a/libs/vertexai/tests/integration_tests/test_vectorstores.py b/libs/vertexai/tests/integration_tests/test_vectorstores.py index d2fd5aac..6c3c4bbb 100644 --- a/libs/vertexai/tests/integration_tests/test_vectorstores.py +++ b/libs/vertexai/tests/integration_tests/test_vectorstores.py @@ -11,7 +11,7 @@ """ import os -from typing import Dict, List +from typing import Dict, List, Union from uuid import uuid4 import pytest @@ -280,7 +280,7 @@ def test_vector_store_hybrid_search( query = "What are your favourite animals?" embedding = embeddings.embed_query(query) - sparse_embedding: Dict[str, List[int] | List[float]] = { + sparse_embedding: Dict[str, Union[List[int], List[float]]] = { "values": [0.5, 0.7], "dimensions": [2, 4], } @@ -317,7 +317,7 @@ def test_add_texts_with_embeddings( ) assert len(ids1) == 2 - sparse_embeddings: List[Dict[str, List[int] | List[float]]] = [ + sparse_embeddings: List[Dict[str, Union[List[int], List[float]]]] = [ {"values": [0.5, 0.7], "dimensions": [2, 4]} ] * 2 ids2 = vector_store.add_texts_with_embeddings( diff --git a/libs/vertexai/tests/unit_tests/test_vectorstores.py b/libs/vertexai/tests/unit_tests/test_vectorstores.py index add2c6b6..24653013 100644 --- a/libs/vertexai/tests/unit_tests/test_vectorstores.py +++ b/libs/vertexai/tests/unit_tests/test_vectorstores.py @@ -1,4 +1,4 @@ -from typing import Dict, List +from typing import Dict, List, Union from unittest.mock import MagicMock import pytest @@ -11,7 +11,7 @@ def test_to_data_points(): ids = ["Id1"] embeddings = [[0.0, 0.0]] - sparse_embeddings: List[Dict[str, List[int] | List[float]]] = [ + sparse_embeddings: List[Dict[str, Union[List[int], List[float]]]] = [ {"values": [0.9, 0.3], "dimensions": [3, 20]} ] metadatas = [ @@ -112,7 +112,7 @@ def test_add_texts_with_embeddings(): def test_similarity_search_by_vector_with_score_output_shape(): embedding = [0.0, 0.5, 0.8] - sparse_embedding: Dict[str, List[int] | List[float]] = { + sparse_embedding: Dict[str, Union[List[int], List[float]]] = { "values": [0.9, 0.3], "dimensions": [3, 20], }