diff --git a/docs/index.rst b/docs/index.rst index a9683863..1f53b5b2 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -7,6 +7,7 @@ API Reference langchain_google_alloydb_pg/engine langchain_google_alloydb_pg/vectorstore + langchain_google_alloydb_pg/indexes langchain_google_alloydb_pg/loader langchain_google_alloydb_pg/history diff --git a/docs/langchain_google_alloydb_pg/indexes.rst b/docs/langchain_google_alloydb_pg/indexes.rst new file mode 100644 index 00000000..fb8156ca --- /dev/null +++ b/docs/langchain_google_alloydb_pg/indexes.rst @@ -0,0 +1,7 @@ +Indexes +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. automodule:: langchain_google_alloydb_pg.indexes + :members: + :private-members: + :noindex: diff --git a/src/langchain_google_alloydb_pg/async_vectorstore.py b/src/langchain_google_alloydb_pg/async_vectorstore.py index 7437f5a3..83e7fbfe 100644 --- a/src/langchain_google_alloydb_pg/async_vectorstore.py +++ b/src/langchain_google_alloydb_pg/async_vectorstore.py @@ -16,6 +16,7 @@ from __future__ import annotations import base64 +import logging import re from typing import Any, Optional @@ -26,6 +27,8 @@ from langchain_postgres.v2.async_vectorstore import AsyncPGVectorStore from sqlalchemy import text +logger = logging.getLogger(__name__) + class AsyncAlloyDBVectorStore(AsyncPGVectorStore): """Google AlloyDB Vector Store class""" @@ -134,8 +137,12 @@ async def asimilarity_search_image( embedding=embedding, k=k, filter=filter, **kwargs ) - async def set_maintenance_work_mem(self, num_leaves: int, vector_size: int) -> None: + async def aset_maintenance_work_mem( + self, num_leaves: Optional[int], vector_size: int + ) -> None: """Set database maintenance work memory (for ScaNN index creation).""" + if not num_leaves: + return # Required index memory in MB buffer = 1 index_memory_required = ( @@ -146,6 +153,125 @@ async def set_maintenance_work_mem(self, num_leaves: int, vector_size: int) -> N await conn.execute(text(query)) await conn.commit() + set_maintenance_work_mem = aset_maintenance_work_mem + + async def ainitialize_auto_vector_embeddings( + self, + model_id: str, + content_column: Optional[str] = None, + embedding_column: Optional[str] = None, + schema_name: Optional[str] = None, + ) -> None: + """Asynchronously initialize auto vector embeddings. + + Args: + model_id: The ID of the model to use for embeddings. + content_column: Optional name of the content column. Defaults to self.content_column. + embedding_column: Optional name of the embedding column. Defaults to self.embedding_column. + schema_name: Optional name of the database schema. Defaults to self.schema_name. + """ + content_col = content_column or self.content_column + embedding_col = embedding_column or self.embedding_column + schema = schema_name or getattr(self, "schema_name", "public") + + if not content_col: + raise ValueError( + "content_column must be provided or configured on the vector store." + ) + if not embedding_col: + raise ValueError( + "embedding_column must be provided or configured on the vector store." + ) + + table_identifier = ( + f'"{schema}"."{self.table_name}"' if schema else f'"{self.table_name}"' + ) + query = "CALL ai.initialize_embeddings(:model_id, :table_name, :content_column, :embedding_column)" + async with self.engine.connect() as conn: + await conn.execute( + text(query), + { + "model_id": model_id, + "table_name": table_identifier, + "content_column": content_col, + "embedding_column": embedding_col, + }, + ) + await conn.commit() + + async def aenable_columnar_engine( + self, + columns: Optional[list[str]] = None, + ) -> None: + """Asynchronously add the table and its columns to the columnar engine. + + Args: + columns: Optional list of column names to add to the columnar engine. + """ + if columns: + columns_str = ",".join(columns) + query = "SELECT google_columnar_engine_add(relation => :table_name, columns => :columns)" + params = {"table_name": self.table_name, "columns": columns_str} + else: + query = "SELECT google_columnar_engine_add(:table_name)" + params = {"table_name": self.table_name} + + async with self.engine.connect() as conn: + await conn.execute(text(query), params) + await conn.commit() + + async def aenable_auto_columnarization(self) -> None: + """Asynchronously trigger auto-columnarization recommendations.""" + query = "SELECT google_columnar_engine_recommend('AUTO_COLUMNARIZATION')" + async with self.engine.connect() as conn: + await conn.execute(text(query)) + await conn.commit() + + async def adefine_vector_assist_spec(self) -> list[dict]: + """Asynchronously define a Vector Assist spec for the current table.""" + query = "SELECT * FROM vector_assist.define_spec(table_name => :table_name, vector_column_name => :embedding_column)" + params = { + "table_name": self.table_name, + "embedding_column": self.embedding_column, + } + async with self.engine.connect() as conn: + result = await conn.execute(text(query), params) + return [dict(row) for row in result.mappings()] + + async def aapply_vector_assist_spec(self) -> list[dict]: + """Asynchronously apply the Vector Assist spec for the current table.""" + query = "SELECT * FROM vector_assist.apply_spec(table_name => :table_name, column_name => :embedding_column)" + params = { + "table_name": self.table_name, + "embedding_column": self.embedding_column, + } + async with self.engine.connect() as conn: + result = await conn.execute(text(query), params) + return [dict(row) for row in result.mappings()] + + async def aget_vector_assist_recommendations(self) -> list[dict]: + """Asynchronously get Vector Assist recommendations for the current table.""" + # First we need to get the spec ID for the current table + specs = await self.adefine_vector_assist_spec() + if not specs: + logger.warning( + "No vector assist spec found for table '%s'.", self.table_name + ) + return [] + + spec_id = specs[0].get("vector_spec_id") + if spec_id is None: + logger.warning( + "Vector assist spec for table '%s' does not contain 'vector_spec_id'.", + self.table_name, + ) + return [] + + query = "SELECT * FROM vector_assist.get_recommendations(:spec_id)" + async with self.engine.connect() as conn: + result = await conn.execute(text(query), {"spec_id": spec_id}) + return [dict(row) for row in result.mappings()] + def add_images( self, uris: list[str], diff --git a/src/langchain_google_alloydb_pg/engine.py b/src/langchain_google_alloydb_pg/engine.py index b712abd6..aa698889 100644 --- a/src/langchain_google_alloydb_pg/engine.py +++ b/src/langchain_google_alloydb_pg/engine.py @@ -621,6 +621,127 @@ def init_checkpoint_table( """ self._run_as_sync(self._ainit_checkpoint_table(table_name, schema_name)) + async def _aforecast( + self, + model_id: str, + source_table: str, + timestamp_col: str, + data_col: str, + horizon: int, + source_query: Optional[str] = None, + conf_level: Optional[float] = None, + ) -> list[dict]: + if not model_id: + raise ValueError("model_id must be provided.") + if not source_table: + raise ValueError("source_table must be provided.") + if not timestamp_col: + raise ValueError("timestamp_col must be provided.") + if not data_col: + raise ValueError("data_col must be provided.") + if horizon <= 0: + raise ValueError("horizon must be a positive integer.") + if conf_level is not None and not (0 < conf_level < 1): + raise ValueError("conf_level must be between 0 and 1.") + + args = [ + "model_id => :model_id", + "source_table => :source_table", + "timestamp_col => :timestamp_col", + "data_col => :data_col", + "horizon => :horizon", + ] + params: dict[str, Any] = { + "model_id": model_id, + "source_table": source_table, + "timestamp_col": timestamp_col, + "data_col": data_col, + "horizon": horizon, + } + if source_query is not None: + args.append("source_query => :source_query") + params["source_query"] = source_query + if conf_level is not None: + args.append("conf_level => :conf_level") + params["conf_level"] = conf_level + + query = f"SELECT * FROM google_ml.forecast({', '.join(args)})" + async with self._pool.connect() as conn: + result = await conn.execute(text(query), params) + return [dict(row) for row in result.mappings()] + + async def aforecast( + self, + model_id: str, + source_table: str, + timestamp_col: str, + data_col: str, + horizon: int, + source_query: Optional[str] = None, + conf_level: Optional[float] = None, + ) -> list[dict]: + """Asynchronously get forecasting from AlloyDB AI. + + Args: + model_id: The ID of the time series forecasting model. + source_table: The table to read historical time series data from. + timestamp_col: The column containing the timestamp. + data_col: The column containing the data to forecast. + horizon: Number of future time steps to forecast. + source_query: Optional query to filter historical data. + conf_level: Optional confidence level for prediction intervals. + + Returns: + A list of dictionaries with forecast_timestamp, forecast_value, and intervals. + """ + return await self._run_as_async( + self._aforecast( + model_id, + source_table, + timestamp_col, + data_col, + horizon, + source_query, + conf_level, + ) + ) + + def forecast( + self, + model_id: str, + source_table: str, + timestamp_col: str, + data_col: str, + horizon: int, + source_query: Optional[str] = None, + conf_level: Optional[float] = None, + ) -> list[dict]: + """Synchronously get forecasting from AlloyDB AI. + + Args: + model_id: The ID of the time series forecasting model. + source_table: The table to read historical time series data from. + timestamp_col: The column containing the timestamp. + data_col: The column containing the data to forecast. + horizon: Number of future time steps to forecast. + source_query: Optional query to filter historical data. + conf_level: Optional confidence level for prediction intervals. + + Returns: + A list of dictionaries with forecast_timestamp, forecast_value, and intervals. + """ + return self._run_as_sync( + self._aforecast( + model_id, + source_table, + timestamp_col, + data_col, + horizon, + source_query, + conf_level, + ) + ) + async def _aload_table_schema( self, table_name: str, schema_name: str = "public" ) -> Table: diff --git a/src/langchain_google_alloydb_pg/indexes.py b/src/langchain_google_alloydb_pg/indexes.py index 48f5974f..164320f2 100644 --- a/src/langchain_google_alloydb_pg/indexes.py +++ b/src/langchain_google_alloydb_pg/indexes.py @@ -14,6 +14,7 @@ import warnings from dataclasses import dataclass, field +from typing import Optional from langchain_postgres.v2.indexes import ( DEFAULT_DISTANCE_STRATEGY, @@ -62,8 +63,18 @@ def to_string(self) -> str: @dataclass class ScaNNIndex(BaseIndex): + """ScaNN index configuration for AlloyDB. + + Args: + mode (Optional[str]): Index mode (e.g. 'AUTO' for auto-tuned indexing). Defaults to None. + num_leaves (Optional[int]): Number of leaves in index clusters. Defaults to 5. + quantizer (str): Quantizer type. Defaults to 'sq8'. + extension_name (str): Extension name. Defaults to 'alloydb_scann'. + """ + index_type: str = "ScaNN" - num_leaves: int = 5 + mode: Optional[str] = None + num_leaves: Optional[int] = 5 quantizer: str = field( default="sq8", init=False ) # Disable `quantizer` initialization currently only supports the value "sq8" @@ -71,6 +82,12 @@ class ScaNNIndex(BaseIndex): def index_options(self) -> str: """Set index query options for vector store initialization.""" + if self.mode is not None: + if self.mode.upper() != "AUTO": + raise ValueError( + f"Invalid mode '{self.mode}'. Only mode='AUTO' is currently supported." + ) + return "(mode = 'AUTO')" return f"(num_leaves = {self.num_leaves}, quantizer = {self.quantizer})" def get_index_function(self) -> str: @@ -84,15 +101,36 @@ def get_index_function(self) -> str: @dataclass class ScaNNQueryOptions(QueryOptions): - num_leaves_to_search: int = 1 + """Query options for ScaNN index. + + Args: + num_leaves_to_search (Optional[int]): Absolute number of leaves to search. Defaults to 1. + pre_reordering_num_neighbors (int): Number of neighbors to consider before reordering. Defaults to -1. + pct_leaves_to_search (Optional[float]): Percentage of leaves to search (0.0 to 1.0 or proportion). + When specified, this takes precedence over `num_leaves_to_search`. + """ + + num_leaves_to_search: Optional[int] = 1 pre_reordering_num_neighbors: int = -1 + pct_leaves_to_search: Optional[float] = None def to_parameter(self) -> list[str]: """Convert index attributes to list of configurations.""" - return [ - f"scann.num_leaves_to_search = {self.num_leaves_to_search}", - f"scann.pre_reordering_num_neighbors = {self.pre_reordering_num_neighbors}", - ] + params = [] + if self.pct_leaves_to_search is not None: + if self.num_leaves_to_search is not None and self.num_leaves_to_search != 1: + warnings.warn( + "Both 'pct_leaves_to_search' and 'num_leaves_to_search' were provided. " + "'pct_leaves_to_search' takes precedence.", + UserWarning, + ) + params.append(f"scann.pct_leaves_to_search = {self.pct_leaves_to_search}") + elif self.num_leaves_to_search is not None: + params.append(f"scann.num_leaves_to_search = {self.num_leaves_to_search}") + params.append( + f"scann.pre_reordering_num_neighbors = {self.pre_reordering_num_neighbors}" + ) + return params def to_string(self) -> str: """Convert index attributes to string.""" @@ -100,4 +138,4 @@ def to_string(self) -> str: "to_string is deprecated, use to_parameter instead.", DeprecationWarning, ) - return f"scann.num_leaves_to_search = {self.num_leaves_to_search}, scann.pre_reordering_num_neighbors = {self.pre_reordering_num_neighbors}" + return ", ".join(self.to_parameter()) diff --git a/src/langchain_google_alloydb_pg/vectorstore.py b/src/langchain_google_alloydb_pg/vectorstore.py index 09fba583..97064f0b 100644 --- a/src/langchain_google_alloydb_pg/vectorstore.py +++ b/src/langchain_google_alloydb_pg/vectorstore.py @@ -167,6 +167,48 @@ def create_sync( vs = engine._run_as_sync(coro) return cls(cls._PGVectorStore__create_key, engine, vs) # type: ignore + async def ainitialize_auto_vector_embeddings( + self, + model_id: str, + content_column: Optional[str] = None, + embedding_column: Optional[str] = None, + schema_name: Optional[str] = None, + ) -> None: + """Generate and manage auto vector embeddings for large tables. + + Args: + model_id (str): The model id used for generating embeddings. + content_column (Optional[str]): Name of the content column. + embedding_column (Optional[str]): Name of the embedding column. + schema_name (Optional[str]): Name of the database schema. + """ + await self._engine._run_as_async( + self._PGVectorStore__vs.ainitialize_auto_vector_embeddings( # type: ignore + model_id, content_column, embedding_column, schema_name + ) + ) + + def initialize_auto_vector_embeddings( + self, + model_id: str, + content_column: Optional[str] = None, + embedding_column: Optional[str] = None, + schema_name: Optional[str] = None, + ) -> None: + """Generate and manage auto vector embeddings for large tables. + + Args: + model_id (str): The model id used for generating embeddings. + content_column (Optional[str]): Name of the content column. + embedding_column (Optional[str]): Name of the embedding column. + schema_name (Optional[str]): Name of the database schema. + """ + self._engine._run_as_sync( + self._PGVectorStore__vs.ainitialize_auto_vector_embeddings( # type: ignore + model_id, content_column, embedding_column, schema_name + ) + ) + async def aadd_images( self, uris: list[str], @@ -222,15 +264,91 @@ async def asimilarity_search_image( ) async def aset_maintenance_work_mem( - self, num_leaves: int, vector_size: int + self, num_leaves: Optional[int], vector_size: int ) -> None: """Set database maintenance work memory (for ScaNN index creation).""" await self._engine._run_as_async( - self._PGVectorStore__vs.set_maintenance_work_mem(num_leaves, vector_size) # type: ignore + self._PGVectorStore__vs.aset_maintenance_work_mem(num_leaves, vector_size) # type: ignore ) - def set_maintenance_work_mem(self, num_leaves: int, vector_size: int) -> None: + def set_maintenance_work_mem( + self, num_leaves: Optional[int], vector_size: int + ) -> None: """Set database maintenance work memory (for ScaNN index creation).""" self._engine._run_as_sync( - self._PGVectorStore__vs.set_maintenance_work_mem(num_leaves, vector_size) # type: ignore + self._PGVectorStore__vs.aset_maintenance_work_mem(num_leaves, vector_size) # type: ignore + ) + + async def aenable_columnar_engine( + self, + columns: Optional[list[str]] = None, + ) -> None: + """Asynchronously add the table and its columns to the columnar engine. + + Args: + columns: Optional list of column names to add to the columnar engine. + """ + await self._engine._run_as_async( + self._PGVectorStore__vs.aenable_columnar_engine(columns) # type: ignore + ) + + def enable_columnar_engine( + self, + columns: Optional[list[str]] = None, + ) -> None: + """Add the table and its columns to the columnar engine. + + Args: + columns: Optional list of column names to add to the columnar engine. + """ + self._engine._run_as_sync( + self._PGVectorStore__vs.aenable_columnar_engine(columns) # type: ignore + ) + + async def aenable_auto_columnarization(self) -> None: + """Asynchronously trigger auto-columnarization recommendations.""" + await self._engine._run_as_async( + self._PGVectorStore__vs.aenable_auto_columnarization() # type: ignore + ) + + def enable_auto_columnarization(self) -> None: + """Trigger auto-columnarization recommendations.""" + self._engine._run_as_sync( + self._PGVectorStore__vs.aenable_auto_columnarization() # type: ignore + ) + + async def adefine_vector_assist_spec(self) -> list[dict]: + """Asynchronously define a Vector Assist spec for the current table.""" + return await self._engine._run_as_async( + self._PGVectorStore__vs.adefine_vector_assist_spec() # type: ignore + ) + + def define_vector_assist_spec(self) -> list[dict]: + """Define a Vector Assist spec for the current table.""" + return self._engine._run_as_sync( + self._PGVectorStore__vs.adefine_vector_assist_spec() # type: ignore + ) + + async def aapply_vector_assist_spec(self) -> list[dict]: + """Asynchronously apply the Vector Assist spec for the current table.""" + return await self._engine._run_as_async( + self._PGVectorStore__vs.aapply_vector_assist_spec() # type: ignore + ) + + def apply_vector_assist_spec(self) -> list[dict]: + """Apply the Vector Assist spec for the current table.""" + return self._engine._run_as_sync( + self._PGVectorStore__vs.aapply_vector_assist_spec() # type: ignore + ) + + async def aget_vector_assist_recommendations(self) -> list[dict]: + """Asynchronously get Vector Assist recommendations for the current table.""" + return await self._engine._run_as_async( + self._PGVectorStore__vs.aget_vector_assist_recommendations() # type: ignore + ) + + def get_vector_assist_recommendations(self) -> list[dict]: + """Get Vector Assist recommendations for the current table.""" + return self._engine._run_as_sync( + self._PGVectorStore__vs.aget_vector_assist_recommendations() # type: ignore ) diff --git a/tests/test_async_vectorstore.py b/tests/test_async_vectorstore.py index 8dc93762..46dd0693 100644 --- a/tests/test_async_vectorstore.py +++ b/tests/test_async_vectorstore.py @@ -16,6 +16,7 @@ import os import uuid from typing import Sequence +from unittest.mock import AsyncMock, MagicMock, patch import pytest import pytest_asyncio @@ -27,6 +28,10 @@ from langchain_google_alloydb_pg import AlloyDBEngine, Column from langchain_google_alloydb_pg.async_vectorstore import AsyncAlloyDBVectorStore +from langchain_google_alloydb_pg.indexes import ( + DistanceStrategy, + ScaNNIndex, +) DEFAULT_TABLE = "test_table" + str(uuid.uuid4()) DEFAULT_TABLE_SYNC = "test_table_sync" + str(uuid.uuid4()) @@ -473,3 +478,236 @@ async def test_create_vectorstore_with_init(self, engine): embedding_column="myembedding", metadata_columns=["random_column"], # invalid metadata column ) + + async def test_live_columnar_engine(self, vs): + """Test enabling columnar engine against live AlloyDB instance.""" + try: + await vs.aenable_columnar_engine(["content"]) + await vs.aenable_columnar_engine() + except Exception as e: + pytest.skip(f"Columnar engine not supported/enabled on instance: {e}") + + async def test_live_auto_columnarization(self, vs): + """Test triggering auto columnarization recommendations against live AlloyDB instance.""" + try: + await vs.aenable_auto_columnarization() + except Exception as e: + pytest.skip(f"Auto columnarization not supported/enabled on instance: {e}") + + async def test_live_vector_assist(self, vs): + """Test vector assist spec definition, application, and recommendations against live AlloyDB instance.""" + try: + specs = await vs.adefine_vector_assist_spec() + assert isinstance(specs, list) + apply_res = await vs.aapply_vector_assist_spec() + assert isinstance(apply_res, list) + recs = await vs.aget_vector_assist_recommendations() + assert isinstance(recs, list) + except Exception as e: + pytest.skip(f"Vector assist not supported/enabled on instance: {e}") + + +@pytest.mark.asyncio +class TestAsyncVectorStoreUnit: + @pytest.fixture + def vs(self): + vs = AsyncAlloyDBVectorStore.__new__(AsyncAlloyDBVectorStore) + vs.engine = MagicMock() + vs.schema_name = "public" + vs.table_name = "test_table" + vs.content_column = "content" + vs.embedding_column = "embedding" + return vs + + async def test_aenable_columnar_engine(self, vs): + """Test enabling the columnar engine executes queries on engine with columns.""" + with patch.object(vs.engine, "connect") as mock_connect: + mock_conn = AsyncMock() + mock_connect.return_value.__aenter__.return_value = mock_conn + await vs.aenable_columnar_engine(["content"]) + call_args = mock_conn.execute.call_args + assert str(call_args[0][0]) == "SELECT google_columnar_engine_add(relation => :table_name, columns => :columns)" + assert call_args[0][1] == {"table_name": "test_table", "columns": "content"} + + async def test_aenable_columnar_engine_without_columns(self, vs): + """Test enabling columnar engine without specifying columns.""" + with patch.object(vs.engine, "connect") as mock_connect: + mock_conn = AsyncMock() + mock_connect.return_value.__aenter__.return_value = mock_conn + await vs.aenable_columnar_engine() + call_args = mock_conn.execute.call_args + assert str(call_args[0][0]) == "SELECT google_columnar_engine_add(:table_name)" + assert call_args[0][1] == {"table_name": "test_table"} + + async def test_aenable_auto_columnarization(self, vs): + """Test enabling auto columnarization executes queries on engine.""" + with patch.object(vs.engine, "connect") as mock_connect: + mock_conn = AsyncMock() + mock_connect.return_value.__aenter__.return_value = mock_conn + await vs.aenable_auto_columnarization() + call_args = mock_conn.execute.call_args + assert str(call_args[0][0]) == "SELECT google_columnar_engine_recommend('AUTO_COLUMNARIZATION')" + + async def test_adefine_vector_assist_spec(self, vs): + """Test definition of vector assist specification.""" + with patch.object(vs.engine, "connect") as mock_connect: + mock_conn = AsyncMock() + mock_result = MagicMock() + mock_result.mappings.return_value = [{"spec": "ok"}] + mock_conn.execute.return_value = mock_result + mock_connect.return_value.__aenter__.return_value = mock_conn + res = await vs.adefine_vector_assist_spec() + assert res == [{"spec": "ok"}] + call_args = mock_conn.execute.call_args + assert str(call_args[0][0]) == "SELECT * FROM vector_assist.define_spec(table_name => :table_name, vector_column_name => :embedding_column)" + assert call_args[0][1] == { + "table_name": "test_table", + "embedding_column": "embedding", + } + + async def test_aapply_vector_assist_spec(self, vs): + """Test applying vector assist specifications.""" + with patch.object(vs.engine, "connect") as mock_connect: + mock_conn = AsyncMock() + mock_result = MagicMock() + mock_result.mappings.return_value = [{"apply": "ok"}] + mock_conn.execute.return_value = mock_result + mock_connect.return_value.__aenter__.return_value = mock_conn + res = await vs.aapply_vector_assist_spec() + assert res == [{"apply": "ok"}] + call_args = mock_conn.execute.call_args + assert str(call_args[0][0]) == "SELECT * FROM vector_assist.apply_spec(table_name => :table_name, column_name => :embedding_column)" + assert call_args[0][1] == { + "table_name": "test_table", + "embedding_column": "embedding", + } + + async def test_aget_vector_assist_recommendations(self, vs): + """Test retrieving vector assist recommendations.""" + with patch.object( + vs, + "adefine_vector_assist_spec", + return_value=[{"vector_spec_id": "spec123"}], + ): + with patch.object(vs.engine, "connect") as mock_connect: + mock_conn = AsyncMock() + mock_result = MagicMock() + mock_result.mappings.return_value = [{"rec": "ok"}] + mock_conn.execute.return_value = mock_result + mock_connect.return_value.__aenter__.return_value = mock_conn + res = await vs.aget_vector_assist_recommendations() + assert res == [{"rec": "ok"}] + call_args = mock_conn.execute.call_args + assert str(call_args[0][0]) == "SELECT * FROM vector_assist.get_recommendations(:spec_id)" + assert call_args[0][1] == {"spec_id": "spec123"} + + async def test_aget_vector_assist_recommendations_spec_id_zero(self, vs): + """Test retrieving vector assist recommendations when spec_id is 0.""" + with patch.object( + vs, + "adefine_vector_assist_spec", + return_value=[{"vector_spec_id": 0}], + ): + with patch.object(vs.engine, "connect") as mock_connect: + mock_conn = AsyncMock() + mock_result = MagicMock() + mock_result.mappings.return_value = [{"rec": "ok_zero"}] + mock_conn.execute.return_value = mock_result + mock_connect.return_value.__aenter__.return_value = mock_conn + res = await vs.aget_vector_assist_recommendations() + assert res == [{"rec": "ok_zero"}] + call_args = mock_conn.execute.call_args + assert str(call_args[0][0]) == "SELECT * FROM vector_assist.get_recommendations(:spec_id)" + assert call_args[0][1] == {"spec_id": 0} + + async def test_aget_vector_assist_recommendations_empty_specs(self, vs): + """Test retrieving vector assist recommendations when no specs exist.""" + with patch.object(vs, "adefine_vector_assist_spec", return_value=[]): + res = await vs.aget_vector_assist_recommendations() + assert res == [] + + async def test_aget_vector_assist_recommendations_no_spec_id(self, vs): + """Test retrieving vector assist recommendations when spec has no ID.""" + with patch.object( + vs, "adefine_vector_assist_spec", return_value=[{"other_key": "val"}] + ): + res = await vs.aget_vector_assist_recommendations() + assert res == [] + + async def test_ainitialize_auto_vector_embeddings(self, vs): + """Test initializing auto vector embeddings asynchronously.""" + with patch.object(vs.engine, "connect") as mock_connect: + mock_conn = AsyncMock() + mock_connect.return_value.__aenter__.return_value = mock_conn + await vs.ainitialize_auto_vector_embeddings( + model_id="test-model", + ) + call_args = mock_conn.execute.call_args + assert str(call_args[0][0]) == "CALL ai.initialize_embeddings(:model_id, :table_name, :content_column, :embedding_column)" + assert call_args[0][1] == { + "model_id": "test-model", + "table_name": '"public"."test_table"', + "content_column": "content", + "embedding_column": "embedding", + } + + async def test_ainitialize_auto_vector_embeddings_custom_columns(self, vs): + """Test initializing auto vector embeddings with custom columns and schema.""" + with patch.object(vs.engine, "connect") as mock_connect: + mock_conn = AsyncMock() + mock_connect.return_value.__aenter__.return_value = mock_conn + await vs.ainitialize_auto_vector_embeddings( + model_id="test-model", + content_column="custom_content", + embedding_column="custom_embedding", + schema_name="myschema", + ) + call_args = mock_conn.execute.call_args + assert str(call_args[0][0]) == "CALL ai.initialize_embeddings(:model_id, :table_name, :content_column, :embedding_column)" + assert call_args[0][1] == { + "model_id": "test-model", + "table_name": '"myschema"."test_table"', + "content_column": "custom_content", + "embedding_column": "custom_embedding", + } + + async def test_ainitialize_auto_vector_embeddings_missing_columns(self, vs): + """Test error raised when required column names are missing.""" + vs.content_column = None + with pytest.raises( + ValueError, match="content_column must be provided or configured" + ): + await vs.ainitialize_auto_vector_embeddings(model_id="test-model") + + async def test_aset_maintenance_work_mem_none(self, vs): + """Test setting maintenance work mem with None returns without executing SQL.""" + with patch.object(vs.engine, "connect") as mock_connect: + await vs.aset_maintenance_work_mem(None, 768) + assert not mock_connect.called + + async def test_aset_maintenance_work_mem_valid(self, vs): + """Test setting maintenance work mem with valid num_leaves executes SQL.""" + with patch.object(vs.engine, "connect") as mock_connect: + mock_conn = AsyncMock() + mock_connect.return_value.__aenter__.return_value = mock_conn + await vs.aset_maintenance_work_mem(10, 768) + call_args = mock_conn.execute.call_args + assert str(call_args[0][0]) == "SET maintenance_work_mem TO '2 MB';" + + async def test_aapply_vector_index_scann_auto(self, vs): + """Test applying ScaNN index in AUTO mode without live DB.""" + index = ScaNNIndex( + name="scann_auto", + mode="AUTO", + distance_strategy=DistanceStrategy.COSINE_DISTANCE, + ) + with patch.object(vs.engine, "connect") as mock_connect: + mock_conn = AsyncMock() + mock_connect.return_value.__aenter__.return_value = mock_conn + await vs.aapply_vector_index(index) + executed_sqls = [str(call[0][0]) for call in mock_conn.execute.call_args_list] + assert any("CREATE EXTENSION IF NOT EXISTS alloydb_scann" in s for s in executed_sqls) + assert any( + 'CREATE INDEX "scann_auto" ON "public"."test_table" USING ScaNN (embedding cosine) WITH (mode = \'AUTO\')' in s + for s in executed_sqls + ) diff --git a/tests/test_async_vectorstore_index.py b/tests/test_async_vectorstore_index.py index d1befd05..b5d2b29d 100644 --- a/tests/test_async_vectorstore_index.py +++ b/tests/test_async_vectorstore_index.py @@ -31,6 +31,7 @@ HNSWIndex, IVFFlatIndex, IVFIndex, + ScaNNIndex, ) UUID_STR = str(uuid.uuid4()).replace("-", "_") @@ -225,3 +226,16 @@ async def test_aapply_hybrid_search_index_table_with_tsv_column(self, engine): await vs.adrop_vector_index(tsv_index_name) is_valid_index = await vs.is_valid_index(tsv_index_name) assert is_valid_index == False + + async def test_aapply_alloydb_scann_index_auto_mode(self, vs): + index = ScaNNIndex( + name="auto_scann_index", + mode="AUTO", + distance_strategy=DistanceStrategy.COSINE_DISTANCE, + ) + try: + await vs.aapply_vector_index(index) + assert await vs.is_valid_index("auto_scann_index") + await vs.adrop_vector_index("auto_scann_index") + except Exception as e: + pytest.skip(f"alloydb_scann index not supported on instance: {e}") diff --git a/tests/test_engine.py b/tests/test_engine.py index 8d22a7ef..d1b01b90 100644 --- a/tests/test_engine.py +++ b/tests/test_engine.py @@ -15,6 +15,7 @@ import os import uuid from typing import Sequence +from unittest.mock import AsyncMock, MagicMock, patch import asyncpg # type: ignore import pytest @@ -42,7 +43,7 @@ VECTOR_SIZE = 768 embeddings_service = DeterministicFakeEmbedding(size=VECTOR_SIZE) -host = os.environ["IP_ADDRESS"] +host = os.environ.get("IP_ADDRESS", "127.0.0.1") def get_env_var(key: str, desc: str) -> str: @@ -617,3 +618,207 @@ async def test_init_table_hybrid_search(self, engine): ] for row in results: assert row in expected + + +class TestEngineUnit: + @pytest.fixture + def engine(self): + eng = AlloyDBEngine.__new__(AlloyDBEngine) + eng._pool = MagicMock() + + def mock_run_sync(coro): + coro.close() + ret = eng._run_as_sync.return_value + if isinstance(ret, MagicMock): + return [{"prediction": 1.0}] + return ret + + eng._run_as_sync = MagicMock(side_effect=mock_run_sync) + + async def mock_run_async(coro): + return await coro + + eng._run_as_async = mock_run_async + return eng + + @pytest.mark.asyncio + async def test_aforecast(self, engine): + """Test that aforecast calls the underlying google_ml.forecast table function asynchronously.""" + with patch.object(engine._pool, "connect") as mock_connect: + mock_conn = AsyncMock() + mock_result = MagicMock() + mock_result.mappings.return_value = [ + {"prediction": 1.0}, + {"prediction": 2.0}, + ] + mock_conn.execute.return_value = mock_result + mock_connect.return_value.__aenter__.return_value = mock_conn + + results = await engine.aforecast( + model_id="test_model", + source_table="test_table", + source_query=None, + data_col="data", + timestamp_col="ts", + horizon=5, + ) + assert len(results) == 2 + assert results[0]["prediction"] == 1.0 + call_args = mock_conn.execute.call_args + assert "SELECT * FROM google_ml.forecast" in str(call_args[0][0]) + assert "source_query" not in str(call_args[0][0]) + assert "conf_level" not in str(call_args[0][0]) + assert call_args[0][1] == { + "model_id": "test_model", + "source_table": "test_table", + "timestamp_col": "ts", + "data_col": "data", + "horizon": 5, + } + + @pytest.mark.asyncio + async def test_aforecast_with_optional_params(self, engine): + """Test aforecast with source_query and conf_level.""" + with patch.object( + engine, "_aforecast", new_callable=AsyncMock + ) as mock_aforecast: + mock_aforecast.return_value = [{"prediction": 42.0}] + results = await engine.aforecast( + model_id="test_model", + source_table="test_table", + source_query="SELECT * FROM data", + data_col="data", + timestamp_col="ts", + horizon=10, + conf_level=0.95, + ) + assert len(results) == 1 + assert results[0]["prediction"] == 42.0 + mock_aforecast.assert_called_once_with( + "test_model", + "test_table", + "ts", + "data", + 10, + "SELECT * FROM data", + 0.95, + ) + + def test_forecast(self, engine): + """Test that forecast evaluates via _run_as_sync to proxy the google_ml.forecast.""" + engine._run_as_sync.return_value = [{"prediction": 1.0}] + results = engine.forecast( + model_id="test_model", + source_table="test_table", + source_query=None, + data_col="data", + timestamp_col="ts", + horizon=5, + ) + assert results == [{"prediction": 1.0}] + engine._run_as_sync.assert_called_once() + + def test_forecast_with_optional_params(self, engine): + """Test forecast with source_query and conf_level.""" + engine._run_as_sync.return_value = [{"prediction": 42.0}] + results = engine.forecast( + model_id="test_model", + source_table="test_table", + source_query="SELECT * FROM data", + data_col="data", + timestamp_col="ts", + horizon=10, + conf_level=0.95, + ) + assert results == [{"prediction": 42.0}] + engine._run_as_sync.assert_called_once() + + @pytest.mark.asyncio + async def test_private_aforecast(self, engine): + """Test direct _aforecast execution and mapping parsing.""" + with patch.object(engine._pool, "connect") as mock_connect: + mock_conn = AsyncMock() + mock_result = MagicMock() + mock_result.mappings.return_value = [ + {"forecast_timestamp": "2026-08-07", "forecast_value": 100.0} + ] + mock_conn.execute.return_value = mock_result + mock_connect.return_value.__aenter__.return_value = mock_conn + + results = await engine._aforecast( + model_id="model_1", + source_table="sales", + timestamp_col="date", + data_col="revenue", + horizon=3, + source_query="SELECT * FROM sales WHERE active = true", + conf_level=0.9, + ) + assert len(results) == 1 + assert results[0]["forecast_value"] == 100.0 + call_args = mock_conn.execute.call_args + assert "SELECT * FROM google_ml.forecast" in str(call_args[0][0]) + assert "source_query => :source_query" in str(call_args[0][0]) + assert "conf_level => :conf_level" in str(call_args[0][0]) + assert call_args[0][1] == { + "model_id": "model_1", + "source_table": "sales", + "timestamp_col": "date", + "data_col": "revenue", + "horizon": 3, + "source_query": "SELECT * FROM sales WHERE active = true", + "conf_level": 0.9, + } + + @pytest.mark.asyncio + async def test_aforecast_validation_errors(self, engine): + """Test validation errors for invalid input parameters in _aforecast.""" + with pytest.raises(ValueError, match="model_id must be provided"): + await engine._aforecast( + model_id="", + source_table="sales", + timestamp_col="date", + data_col="revenue", + horizon=3, + ) + with pytest.raises(ValueError, match="source_table must be provided"): + await engine._aforecast( + model_id="model_1", + source_table="", + timestamp_col="date", + data_col="revenue", + horizon=3, + ) + with pytest.raises(ValueError, match="timestamp_col must be provided"): + await engine._aforecast( + model_id="model_1", + source_table="sales", + timestamp_col="", + data_col="revenue", + horizon=3, + ) + with pytest.raises(ValueError, match="data_col must be provided"): + await engine._aforecast( + model_id="model_1", + source_table="sales", + timestamp_col="date", + data_col="", + horizon=3, + ) + with pytest.raises(ValueError, match="horizon must be a positive integer"): + await engine._aforecast( + model_id="model_1", + source_table="sales", + timestamp_col="date", + data_col="revenue", + horizon=0, + ) + with pytest.raises(ValueError, match="conf_level must be between 0 and 1"): + await engine._aforecast( + model_id="model_1", + source_table="sales", + timestamp_col="date", + data_col="revenue", + horizon=3, + conf_level=1.5, + ) diff --git a/tests/test_indexes.py b/tests/test_indexes.py index bc1d04de..63dbc65a 100644 --- a/tests/test_indexes.py +++ b/tests/test_indexes.py @@ -109,6 +109,26 @@ def test_scann_index(self): assert index.quantizer == "sq8" # Check default value assert index.index_options() == "(num_leaves = 10, quantizer = sq8)" + def test_scann_index_auto_mode(self): + index = ScaNNIndex(name="test_index", mode="AUTO") + assert index.index_type == "ScaNN" + assert index.mode == "AUTO" + assert index.index_options() == "(mode = 'AUTO')" + + def test_scann_index_invalid_mode(self): + index = ScaNNIndex(name="test_index", mode="INVALID") + import pytest + + with pytest.raises(ValueError, match="Invalid mode 'INVALID'"): + index.index_options() + + def test_scann_query_options_default(self): + options = ScaNNQueryOptions() + assert options.to_parameter() == [ + "scann.num_leaves_to_search = 1", + "scann.pre_reordering_num_neighbors = -1", + ] + def test_scann_query_options(self): options = ScaNNQueryOptions( num_leaves_to_search=2, pre_reordering_num_neighbors=10 @@ -124,3 +144,37 @@ def test_scann_query_options(self): assert "to_string is deprecated, use to_parameter instead." in str( w[-1].message ) + + def test_scann_query_options_pct_leaves(self): + options = ScaNNQueryOptions( + pre_reordering_num_neighbors=10, + pct_leaves_to_search=0.2, + ) + assert options.to_parameter() == [ + "scann.pct_leaves_to_search = 0.2", + "scann.pre_reordering_num_neighbors = 10", + ] + with warnings.catch_warnings(record=True) as w: + to_str = options.to_string() + assert ( + to_str + == "scann.pct_leaves_to_search = 0.2, scann.pre_reordering_num_neighbors = 10" + ) + + def test_scann_query_options_both_params_warns(self): + options = ScaNNQueryOptions( + num_leaves_to_search=5, + pre_reordering_num_neighbors=10, + pct_leaves_to_search=0.5, + ) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + params = options.to_parameter() + assert len(w) == 1 + assert "Both 'pct_leaves_to_search' and 'num_leaves_to_search' were provided" in str( + w[-1].message + ) + assert params == [ + "scann.pct_leaves_to_search = 0.5", + "scann.pre_reordering_num_neighbors = 10", + ] diff --git a/tests/test_vectorstore.py b/tests/test_vectorstore.py index 0ec411ff..41b8188b 100644 --- a/tests/test_vectorstore.py +++ b/tests/test_vectorstore.py @@ -18,6 +18,7 @@ import uuid from threading import Thread from typing import Sequence +from unittest.mock import AsyncMock, MagicMock, patch import pytest import pytest_asyncio @@ -30,6 +31,10 @@ from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine from langchain_google_alloydb_pg import AlloyDBEngine, AlloyDBVectorStore, Column +from langchain_google_alloydb_pg.indexes import ( + DistanceStrategy, + ScaNNIndex, +) DEFAULT_TABLE = "test_table" + str(uuid.uuid4()) DEFAULT_TABLE_SYNC = "test_table_sync" + str(uuid.uuid4()) @@ -39,7 +44,7 @@ VECTOR_SIZE = 768 embeddings_service = DeterministicFakeEmbedding(size=VECTOR_SIZE) -host = os.environ["IP_ADDRESS"] +host = os.environ.get("IP_ADDRESS", "127.0.0.1") texts = ["foo", "bar", "baz"] metadatas = [{"page": str(i), "source": "google.com"} for i in range(len(texts))] @@ -745,3 +750,210 @@ async def test_from_engine_loop( def test_get_table_name(self, vs): assert vs.get_table_name() == DEFAULT_TABLE + + def test_live_columnar_engine(self, vs): + """Test enabling columnar engine against live AlloyDB instance.""" + try: + vs.enable_columnar_engine(["content"]) + vs.enable_columnar_engine() + except Exception as e: + pytest.skip(f"Columnar engine not supported/enabled on instance: {e}") + + def test_live_auto_columnarization(self, vs): + """Test triggering auto columnarization recommendations against live AlloyDB instance.""" + try: + vs.enable_auto_columnarization() + except Exception as e: + pytest.skip(f"Auto columnarization not supported/enabled on instance: {e}") + + def test_live_vector_assist(self, vs): + """Test vector assist spec definition, application, and recommendations against live AlloyDB instance.""" + try: + specs = vs.define_vector_assist_spec() + assert isinstance(specs, list) + apply_res = vs.apply_vector_assist_spec() + assert isinstance(apply_res, list) + recs = vs.get_vector_assist_recommendations() + assert isinstance(recs, list) + except Exception as e: + pytest.skip(f"Vector assist not supported/enabled on instance: {e}") + + +class TestVectorStoreUnit: + @pytest.fixture + def vs(self): + vs = AlloyDBVectorStore.__new__(AlloyDBVectorStore) + vs._engine = MagicMock() + mock_vs = MagicMock() + vs._PGVectorStore__vs = mock_vs + vs._AlloyDBVectorStore__vs = mock_vs + + def mock_sync(coro): + if hasattr(coro, "close"): + coro.close() + return getattr(vs._engine._run_as_sync, "return_value", None) + + async def mock_async(coro): + if hasattr(coro, "close"): + coro.close() + return getattr(vs._engine._run_as_async, "return_value", None) + + vs._engine._run_as_sync = MagicMock(side_effect=mock_sync) + vs._engine._run_as_async = AsyncMock(side_effect=mock_async) + return vs + + def test_enable_columnar_engine(self, vs): + """Test enabling the columnar engine triggers the appropriate sync method on the underlying store.""" + vs.enable_columnar_engine(["content"]) + vs._PGVectorStore__vs.aenable_columnar_engine.assert_called_once_with(["content"]) + + def test_enable_columnar_engine_without_columns(self, vs): + """Test enabling columnar engine without columns.""" + vs.enable_columnar_engine() + vs._PGVectorStore__vs.aenable_columnar_engine.assert_called_once_with(None) + + @pytest.mark.asyncio + async def test_aenable_columnar_engine(self, vs): + """Test enabling the columnar engine triggers the appropriate async method on the underlying store.""" + await vs.aenable_columnar_engine(["content"]) + vs._PGVectorStore__vs.aenable_columnar_engine.assert_called_once_with(["content"]) + + @pytest.mark.asyncio + async def test_aenable_columnar_engine_without_columns(self, vs): + """Test enabling columnar engine without columns asynchronously.""" + await vs.aenable_columnar_engine() + vs._PGVectorStore__vs.aenable_columnar_engine.assert_called_once_with(None) + + def test_enable_auto_columnarization(self, vs): + """Test enabling auto columnarization triggers the sync engine wrapper.""" + vs.enable_auto_columnarization() + vs._PGVectorStore__vs.aenable_auto_columnarization.assert_called_once_with() + + @pytest.mark.asyncio + async def test_aenable_auto_columnarization(self, vs): + """Test enabling auto columnarization triggers the async engine wrapper.""" + await vs.aenable_auto_columnarization() + vs._PGVectorStore__vs.aenable_auto_columnarization.assert_called_once_with() + + def test_define_vector_assist_spec(self, vs): + """Test definition of vector assist specification.""" + expected = [{"spec": "ok"}] + vs._engine._run_as_sync.return_value = expected + res = vs.define_vector_assist_spec() + assert res == expected + vs._PGVectorStore__vs.adefine_vector_assist_spec.assert_called_once_with() + + @pytest.mark.asyncio + async def test_adefine_vector_assist_spec(self, vs): + """Test definition of vector assist specification asynchronously.""" + expected = [{"spec": "ok"}] + vs._engine._run_as_async.return_value = expected + res = await vs.adefine_vector_assist_spec() + assert res == expected + vs._PGVectorStore__vs.adefine_vector_assist_spec.assert_called_once_with() + + def test_apply_vector_assist_spec(self, vs): + """Test applying vector assist specifications.""" + expected = [{"apply": "ok"}] + vs._engine._run_as_sync.return_value = expected + res = vs.apply_vector_assist_spec() + assert res == expected + vs._PGVectorStore__vs.aapply_vector_assist_spec.assert_called_once_with() + + @pytest.mark.asyncio + async def test_aapply_vector_assist_spec(self, vs): + """Test applying vector assist specifications asynchronously.""" + expected = [{"apply": "ok"}] + vs._engine._run_as_async.return_value = expected + res = await vs.aapply_vector_assist_spec() + assert res == expected + vs._PGVectorStore__vs.aapply_vector_assist_spec.assert_called_once_with() + + def test_get_vector_assist_recommendations(self, vs): + """Test retrieving vector assist recommendations.""" + expected = [{"rec": "ok"}] + vs._engine._run_as_sync.return_value = expected + res = vs.get_vector_assist_recommendations() + assert res == expected + vs._PGVectorStore__vs.aget_vector_assist_recommendations.assert_called_once_with() + + @pytest.mark.asyncio + async def test_aget_vector_assist_recommendations(self, vs): + """Test retrieving vector assist recommendations asynchronously.""" + expected = [{"rec": "ok"}] + vs._engine._run_as_async.return_value = expected + res = await vs.aget_vector_assist_recommendations() + assert res == expected + vs._PGVectorStore__vs.aget_vector_assist_recommendations.assert_called_once_with() + + def test_initialize_auto_vector_embeddings(self, vs): + """Test initializing auto vector embeddings.""" + vs.initialize_auto_vector_embeddings( + model_id="test-model", + ) + vs._PGVectorStore__vs.ainitialize_auto_vector_embeddings.assert_called_once_with( + "test-model", None, None, None + ) + + def test_initialize_auto_vector_embeddings_with_columns(self, vs): + """Test initializing auto vector embeddings with custom columns.""" + vs.initialize_auto_vector_embeddings( + model_id="test-model", + content_column="custom_content", + embedding_column="custom_embedding", + schema_name="myschema", + ) + vs._PGVectorStore__vs.ainitialize_auto_vector_embeddings.assert_called_once_with( + "test-model", "custom_content", "custom_embedding", "myschema" + ) + + @pytest.mark.asyncio + async def test_ainitialize_auto_vector_embeddings(self, vs): + """Test initializing auto vector embeddings asynchronously.""" + await vs.ainitialize_auto_vector_embeddings( + model_id="test-model", + ) + vs._PGVectorStore__vs.ainitialize_auto_vector_embeddings.assert_called_once_with( + "test-model", None, None, None + ) + + @pytest.mark.asyncio + async def test_ainitialize_auto_vector_embeddings_with_columns(self, vs): + """Test initializing auto vector embeddings with custom columns asynchronously.""" + await vs.ainitialize_auto_vector_embeddings( + model_id="test-model", + content_column="custom_content", + embedding_column="custom_embedding", + schema_name="myschema", + ) + vs._PGVectorStore__vs.ainitialize_auto_vector_embeddings.assert_called_once_with( + "test-model", "custom_content", "custom_embedding", "myschema" + ) + + def test_set_maintenance_work_mem_none(self, vs): + """Test setting maintenance work mem with None.""" + vs.set_maintenance_work_mem(None, 768) + vs._PGVectorStore__vs.aset_maintenance_work_mem.assert_called_once_with(None, 768) + + @pytest.mark.asyncio + async def test_aset_maintenance_work_mem_none(self, vs): + """Test setting maintenance work mem with None asynchronously.""" + await vs.aset_maintenance_work_mem(None, 768) + vs._PGVectorStore__vs.aset_maintenance_work_mem.assert_called_once_with(None, 768) + + def test_apply_vector_index_scann_auto(self, vs): + """Test applying ScaNN index in AUTO mode synchronously without live DB.""" + index = ScaNNIndex(name="scann_auto", mode="AUTO") + vs.apply_vector_index(index) + vs._PGVectorStore__vs.aapply_vector_index.assert_called_once_with( + index, None, concurrently=False + ) + + @pytest.mark.asyncio + async def test_aapply_vector_index_scann_auto(self, vs): + """Test applying ScaNN index in AUTO mode asynchronously without live DB.""" + index = ScaNNIndex(name="scann_auto", mode="AUTO") + await vs.aapply_vector_index(index) + vs._PGVectorStore__vs.aapply_vector_index.assert_called_once_with( + index, None, concurrently=False + ) diff --git a/tests/test_vectorstore_index.py b/tests/test_vectorstore_index.py index 310d3d21..4ff37b3c 100644 --- a/tests/test_vectorstore_index.py +++ b/tests/test_vectorstore_index.py @@ -327,3 +327,16 @@ async def test_aapply_alloydb_scann_index_ScaNN(self, omni_vs): assert await omni_vs.ais_valid_index("secondindex") await omni_vs.adrop_vector_index("secondindex") await omni_vs.adrop_vector_index(DEFAULT_INDEX_NAME_OMNI) + + async def test_aapply_alloydb_scann_index_auto_mode(self, omni_vs): + index = ScaNNIndex( + name="auto_scann_index", + mode="AUTO", + distance_strategy=DistanceStrategy.COSINE_DISTANCE, + ) + try: + await omni_vs.aapply_vector_index(index) + assert await omni_vs.ais_valid_index("auto_scann_index") + await omni_vs.adrop_vector_index("auto_scann_index") + except Exception as e: + pytest.skip(f"alloydb_scann index not supported on instance: {e}")