diff --git a/tests/integration/test_spanner_loader.py b/tests/integration/test_spanner_loader.py index e046bf5..fd8f028 100644 --- a/tests/integration/test_spanner_loader.py +++ b/tests/integration/test_spanner_loader.py @@ -16,7 +16,7 @@ import uuid import pytest -from google.cloud.spanner import Client +from google.cloud.spanner import Client # type: ignore from langchain_core.documents import Document from langchain_google_spanner.loader import Column, SpannerDocumentSaver, SpannerLoader diff --git a/tests/integration/test_spanner_vector_store.py b/tests/integration/test_spanner_vector_store.py index bb71125..81dd1b9 100644 --- a/tests/integration/test_spanner_vector_store.py +++ b/tests/integration/test_spanner_vector_store.py @@ -26,16 +26,13 @@ QueryParameters, SpannerVectorStore, TableColumn, - VectorSearchIndex, ) project_id = os.environ["PROJECT_ID"] instance_id = os.environ["INSTANCE_ID"] google_database = os.environ["GOOGLE_DATABASE"] -pg_database = os.environ.get("PG_DATABASE", None) -zone = os.environ.get("GOOGLE_DATABASE_ZONE", "us-west2") +pg_database = os.environ["PG_DATABASE"] table_name = "test_table" + str(uuid.uuid4()).replace("-", "_") -table_name_ANN = "products" OPERATION_TIMEOUT_SECONDS = 240 @@ -392,221 +389,6 @@ def test_spanner_vector_search_data4(self, setup_database): assert len(docs) == 3 -class TestSpannerVectorStoreGoogleSQL_ANN: - @pytest.fixture(scope="class") - def setup_database(self, client): - distance_strategy = DistanceStrategy.COSINE - SpannerVectorStore.init_vector_store_table( - instance_id=instance_id, - database_id=google_database, - table_name=table_name_ANN, - id_column=TableColumn("productId", type="INT64"), - vector_size=758, - embedding_column=TableColumn( - name="productDescriptionEmbedding", - type="ARRAY", - is_null=True, - ), - metadata_columns=[ - TableColumn(name="categoryId", type="INT64", is_null=False), - TableColumn(name="productName", type="STRING(MAX)", is_null=False), - TableColumn( - name="productDescription", type="STRING(MAX)", is_null=False - ), - TableColumn(name="inventoryCount", type="INT64", is_null=False), - TableColumn(name="priceInCents", type="INT64", is_null=True), - ], - secondary_indexes=[ - VectorSearchIndex( - index_name="ProductDescriptionEmbeddingIndex", - columns=["productDescriptionEmbedding"], - nullable_column=True, - num_branches=1000, - tree_depth=3, - index_type=distance_strategy, - num_leaves=100000, - ), - ], - ) - - raw_data = [ - ( - 1, - 1, - "Cymbal Helios Helmet", - "Safety meets style with the Cymbal children's bike helmet. Its lightweight design, superior ventilation, and adjustable fit ensure comfort and protection on every ride. Stay bright and keep your child safe under the sun with Cymbal Helios!", - 100, - 10999, - ), - ( - 1, - 2, - "Cymbal Sprout", - "Let their cycling journey begin with the Cymbal Sprout, the ideal balance bike for beginning riders ages 2-4 years. Its lightweight frame, low seat height, and puncture-proof tires promote stability and confidence as little ones learn to balance and steer. Watch them sprout into cycling enthusiasts with Cymbal Sprout!", - 10, - 13999, - ), - ( - 1, - 3, - "Cymbal Spark Jr.", - "Light, vibrant, and ready for adventure, the Spark Jr. is the perfect first bike for young riders (ages 5-8). Its sturdy frame, easy-to-use brakes, and puncture-resistant tires inspire confidence and endless playtime. Let the spark of cycling ignite with Cymbal!", - 34, - 13900, - ), - ( - 1, - 4, - "Cymbal Summit", - "Conquering trails is a breeze with the Summit mountain bike. Its lightweight aluminum frame, responsive suspension, and powerful disc brakes provide exceptional control and comfort for experienced bikers navigating rocky climbs or shredding downhill. Reach new heights with Cymbal Summit!", - 0, - 79999, - ), - ( - 1, - 5, - "Cymbal Breeze", - "Cruise in style and embrace effortless pedaling with the Breeze electric bike. Its whisper-quiet motor and long-lasting battery let you conquer hills and distances with ease. Enjoy scenic rides, commutes, or errands with a boost of confidence from Cymbal Breeze!", - 72, - 129999, - ), - ( - 1, - 6, - "Cymbal Trailblazer Backpack", - "Carry all your essentials in style with the Trailblazer backpack. Its water-resistant material, multiple compartments, and comfortable straps keep your gear organized and accessible, allowing you to focus on the adventure. Blaze new trails with Cymbal Trailblazer!", - 24, - 7999, - ), - ( - 1, - 7, - "Cymbal Phoenix Lights", - "See and be seen with the Phoenix bike lights. Powerful LEDs and multiple light modes ensure superior visibility, enhancing your safety and enjoyment during day or night rides. Light up your journey with Cymbal Phoenix!", - 87, - 3999, - ), - ( - 1, - 8, - "Cymbal Windstar Pump", - "Flat tires are no match for the Windstar pump. Its compact design, lightweight construction, and high-pressure capacity make inflating tires quick and effortless. Get back on the road in no time with Cymbal Windstar!", - 36, - 24999, - ), - ( - 1, - 9, - "Cymbal Odyssey Multi-Tool", - "Be prepared for anything with the Odyssey multi-tool. This handy gadget features essential tools like screwdrivers, hex wrenches, and tire levers, keeping you ready for minor repairs and adjustments on the go. Conquer your journey with Cymbal Odyssey!", - 52, - 999, - ), - ( - 1, - 10, - "Cymbal Nomad Water Bottle", - "Stay hydrated on every ride with the Nomad water bottle. Its sleek design, BPA-free construction, and secure lock lid make it the perfect companion for staying refreshed and motivated throughout your adventures. Hydrate and explore with Cymbal Nomad!", - 42, - 1299, - ), - ] - - columns = [ - "categoryId", - "productId", - "productName", - "productDescription", - "createTime", - "inventoryCount", - "priceInCents", - ] - - model_ddl_statements = [ - f""" - CREATE MODEL IF NOT EXISTS EmbeddingsModel INPUT( - content STRING(MAX), - ) OUTPUT( - embeddings STRUCT, values ARRAY>, - ) REMOTE OPTIONS ( - endpoint = '//aiplatform.googleapis.com/projects/{project_id}/locations/{zone}/publishers/google/models/text-embedding-004' - ) - """, - f""" - CREATE MODEL IF NOT EXISTS LLMModel INPUT( - prompt STRING(MAX), - ) OUTPUT( - content STRING(MAX), - ) REMOTE OPTIONS ( - endpoint = '//aiplatform.googleapis.com/projects/{project_id}/locations/{zone}/publishers/google/models/gemini-pro', - default_batch_size = 1 - ) - """, - """ - UPDATE products p1 - SET productDescriptionEmbedding = - ( - SELECT embeddings.values from ML.PREDICT( - MODEL EmbeddingsModel, - (SELECT productDescription as content FROM products p2 where p2.productId=p1.productId) - ) - ) - WHERE categoryId=1 - """, - ] - database = client.instance(instance_id).database(google_database) - - def create_models(): - operation = database.update_ddl(model_ddl_statements) - return operation.result(OPERATION_TIMEOUT_SECONDS) - - def get_embeddings(self): - sql = """SELECT embeddings.values FROM ML.PREDICT( - MODEL EmbeddingsModel, - (SELECT "I'd like to buy a starter bike for my 3 year old child" as content) - )""" - - with database.snapshot() as snapshot: - res = snapshot.execute_sql(sql) - return list(res) - - yield raw_data, columns, create_models, get_embeddings - - print("\nPerforming GSQL cleanup after each ANN test...") - - operation = database.update_ddl( - [ - f"DROP TABLE IF EXISTS {table_name_ANN}", - "DROP MODEL IF EXISTS EmbeddingsModel", - "DROP MODEL IF EXISTS LLMModel", - "DROP Index IF EXISTS ProductDescriptionEmbeddingIndex", - ] - ) - if False: # Creating a vector index takes 30+ minutes, so avoiding this. - operation.result(OPERATION_TIMEOUT_SECONDS) - - # Code to perform teardown after each test goes here - print("\nGSQL Cleanup complete.") - - def test_ann_add_data1(self, setup_database): - raw_data, columns, create_models, get_embeddings = setup_database - - # Retrieve embeddings using ML_PREDICT. - embeddings = get_embeddings() - print("embeddings", embeddings) - - db = SpannerVectorStore( - instance_id=instance_id, - database_id=google_database, - table_name=table_name_ANN, - id_column="categoryId", - ignore_metadata_columns=[], - embedding_service=embeddings, - metadata_json_column="metadata", - ) - _ = db - - class TestSpannerVectorStorePGSQL: @pytest.fixture(scope="class") def setup_database(self, client):