From 6bd5302d8c1d08eb1fe25f26bc7ac7290c3c59c5 Mon Sep 17 00:00:00 2001 From: Jorge Date: Tue, 12 Mar 2024 17:43:36 +0100 Subject: [PATCH] Added tests --- .../integration_tests/test_vectorstores.py | 233 ++++++++++++++++-- 1 file changed, 212 insertions(+), 21 deletions(-) diff --git a/libs/vertexai/tests/integration_tests/test_vectorstores.py b/libs/vertexai/tests/integration_tests/test_vectorstores.py index fc3df8b8..f7ff2e31 100644 --- a/libs/vertexai/tests/integration_tests/test_vectorstores.py +++ b/libs/vertexai/tests/integration_tests/test_vectorstores.py @@ -8,10 +8,10 @@ - GCS_BUCKET_NAME: Name of a Google Cloud Storage Bucket - INDEX_ID: Id of the Vector Search index. - ENDPOINT_ID: Id of the Vector Search endpoint. -If required to run slow tests, environment variable 'RUN_SLOW_TESTS' must be set """ import os +from typing import List from uuid import uuid4 import pytest @@ -20,6 +20,10 @@ MatchingEngineIndex, MatchingEngineIndexEndpoint, ) +from google.cloud.aiplatform.matching_engine.matching_engine_index_endpoint import ( + Namespace, + NumericNamespace, +) from langchain_core.documents import Document from langchain_google_vertexai.embeddings import VertexAIEmbeddings @@ -57,6 +61,22 @@ def datastore_document_storage( return DataStoreDocumentStorage(datastore_client=ds_client) +@pytest.fixture +def vector_store() -> VectorSearchVectorStore: + embeddings = VertexAIEmbeddings(model_name="textembedding-gecko-default") + + vector_store = VectorSearchVectorStore.from_components( + project_id=os.environ["PROJECT_ID"], + region=os.environ["REGION"], + gcs_bucket_name=os.environ["GCS_BUCKET_NAME"], + index_id=os.environ["INDEX_ID"], + endpoint_id=os.environ["ENDPOINT_ID"], + embedding=embeddings, + ) + + return vector_store + + @pytest.mark.extended def test_vector_search_sdk_manager(sdk_manager: VectorSearchSDKManager): gcs_client = sdk_manager.get_gcs_client() @@ -125,18 +145,7 @@ def test_public_endpoint_vector_searcher(sdk_manager: VectorSearchSDKManager): @pytest.mark.extended -def test_vector_store(): - embeddings = VertexAIEmbeddings(model_name="textembedding-gecko-default") - - vector_store = VectorSearchVectorStore.from_components( - project_id=os.environ["PROJECT_ID"], - region=os.environ["REGION"], - gcs_bucket_name=os.environ["GCS_BUCKET_NAME"], - index_id=os.environ["INDEX_ID"], - endpoint_id=os.environ["ENDPOINT_ID"], - embedding=embeddings, - ) - +def test_vector_store(vector_store: VectorSearchVectorStore): assert isinstance(vector_store, VectorSearchVectorStore) query = "What are your favourite animals?" @@ -153,7 +162,20 @@ def test_vector_store(): @pytest.mark.extended -def test_vector_store_update_index(): +def test_vector_store_filtering(vector_store: VectorSearchVectorStore): + documents = vector_store.similarity_search( + "I want some pants", + filter=[Namespace(name="color", allow_tokens=["blue"])], + numeric_filter=[NumericNamespace(name="price", value_float=20.0, op="LESS")], + ) + + assert len(documents) > 0 + + assert all(document.metadata["price"] < 20.0 for document in documents) + + +@pytest.mark.extended +def test_vector_store_update_index(sample_documents: list[Document]): embeddings = VertexAIEmbeddings(model_name="textembedding-gecko-default") vector_store = VectorSearchVectorStore.from_components( @@ -165,10 +187,179 @@ def test_vector_store_update_index(): embedding=embeddings, ) - vector_store.add_texts( - texts=[ - "Lions are my favourite animals", - "There are two apples on the table", - "Today is raining a lot in Madrid", - ] - ) + vector_store.add_documents(documents=sample_documents, is_complete_overwrite=True) + + +@pytest.fixture +def sample_documents() -> List[Document]: + record_data = [ + { + "description": "A versatile pair of dark-wash denim jeans." + "Made from durable cotton with a classic straight-leg cut, these jeans" + " transition easily from casual days to dressier occasions.", + "price": 65.00, + "color": "blue", + "season": ["fall", "winter", "spring"], + }, + { + "description": "A lightweight linen button-down shirt in a crisp white." + " Perfect for keeping cool with breathable fabric and a relaxed fit.", + "price": 34.99, + "color": "white", + "season": ["summer", "spring"], + }, + { + "description": "A soft, chunky knit sweater in a vibrant forest green. " + "The oversized fit and cozy wool blend make this ideal for staying warm " + "when the temperature drops.", + "price": 89.99, + "color": "green", + "season": ["fall", "winter"], + }, + { + "description": "A classic crewneck t-shirt in a soft, heathered blue. " + "Made from comfortable cotton jersey, this t-shirt is a wardrobe essential " + "that works for every season.", + "price": 19.99, + "color": "blue", + "season": ["fall", "winter", "summer", "spring"], + }, + { + "description": "A flowing midi-skirt in a delicate floral print. " + "Lightweight and airy, this skirt adds a touch of feminine style " + "to warmer days.", + "price": 45.00, + "color": "white", + "season": ["spring", "summer"], + }, + { + "description": "A pair of tailored black trousers in a comfortable stretch " + "fabric. Perfect for work or dressier events, these trousers provide a" + " sleek, polished look.", + "price": 59.99, + "color": "black", + "season": ["fall", "winter", "spring"], + }, + { + "description": "A cozy fleece hoodie in a neutral heather grey. " + "This relaxed sweatshirt is perfect for casual days or layering when the " + "weather turns chilly.", + "price": 39.99, + "color": "grey", + "season": ["fall", "winter", "spring"], + }, + { + "description": "A bright yellow raincoat with a playful polka dot pattern. " + "This waterproof jacket will keep you dry and add a touch of cheer to " + "rainy days.", + "price": 75.00, + "color": "yellow", + "season": ["spring", "fall"], + }, + { + "description": "A pair of comfortable khaki chino shorts. These versatile " + "shorts are a summer staple, perfect for outdoor adventures or relaxed" + " weekends.", + "price": 34.99, + "color": "khaki", + "season": ["summer"], + }, + { + "description": "A bold red cocktail dress with a flattering A-line " + "silhouette. This statement piece is made from a luxurious satin fabric, " + "ensuring a head-turning look.", + "price": 125.00, + "color": "red", + "season": ["fall", "winter", "summer", "spring"], + }, + { + "description": "A pair of classic white sneakers crafted from smooth " + "leather. These timeless shoes offer a clean and polished look, perfect " + "for everyday wear.", + "price": 79.99, + "color": "white", + "season": ["fall", "winter", "summer", "spring"], + }, + { + "description": "A chunky cable-knit scarf in a rich burgundy color. " + "Made from a soft wool blend, this scarf will provide warmth and a touch " + "of classic style to cold-weather looks.", + "price": 45.00, + "color": "burgundy", + "season": ["fall", "winter"], + }, + { + "description": "A lightweight puffer vest in a vibrant teal hue. " + "This versatile piece adds a layer of warmth without bulk, transitioning" + " perfectly between seasons.", + "price": 65.00, + "color": "teal", + "season": ["fall", "spring"], + }, + { + "description": "A pair of high-waisted leggings in a sleek black." + " Crafted from a moisture-wicking fabric with plenty of stretch, " + "these leggings are perfect for workouts or comfortable athleisure style.", + "price": 49.99, + "color": "black", + "season": ["fall", "winter", "summer", "spring"], + }, + { + "description": "A denim jacket with a faded wash and distressed details. " + "This wardrobe staple adds a touch of effortless cool to any outfit.", + "price": 79.99, + "color": "blue", + "season": ["fall", "spring", "summer"], + }, + { + "description": "A woven straw sunhat with a wide brim. This stylish " + "accessory provides protection from the sun while adding a touch of " + "summery elegance.", + "price": 32.00, + "color": "beige", + "season": ["summer"], + }, + { + "description": "A graphic tee featuring a vintage band logo. " + "Made from a soft cotton blend, this casual tee adds a touch of " + "personal style to everyday looks.", + "price": 24.99, + "color": "white", + "season": ["fall", "winter", "summer", "spring"], + }, + { + "description": "A pair of well-tailored dress pants in a neutral grey. " + "Made from a wrinkle-resistant blend, these pants look sharp and " + "professional for workwear or formal occasions.", + "price": 69.99, + "color": "grey", + "season": ["fall", "winter", "summer", "spring"], + }, + { + "description": "A pair of classic leather ankle boots in a rich brown hue." + " Featuring a subtle stacked heel and sleek design, these boots are perfect" + " for elevating outfits in cooler seasons.", + "price": 120.00, + "color": "brown", + "season": ["fall", "winter", "spring"], + }, + { + "description": "A vibrant swimsuit with a bold geometric pattern. This fun " + "and eye-catching piece is perfect for making a splash by the pool or at " + "the beach.", + "price": 55.00, + "color": "multicolor", + "season": ["summer"], + }, + ] + + documents = [] + for record in record_data: + record = record.copy() + page_content = record.pop("description") + if isinstance(page_content, str): + metadata = {**record} + document = Document(page_content=page_content, metadata=metadata) + documents.append(document) + + return documents