From 14e47f2176047bf24acee0c6e4ff57714761dfd3 Mon Sep 17 00:00:00 2001 From: Averi Kitsch Date: Fri, 7 Aug 2026 14:30:25 -0700 Subject: [PATCH 1/2] fix: Update to GA syntax --- .../vectorstore.py | 94 +++++++---- tests/unit/test_vectorstore_sql.py | 152 ++++++++++++++++++ 2 files changed, 212 insertions(+), 34 deletions(-) create mode 100644 tests/unit/test_vectorstore_sql.py diff --git a/src/langchain_google_cloud_sql_mysql/vectorstore.py b/src/langchain_google_cloud_sql_mysql/vectorstore.py index 64ecc17..73cd26b 100644 --- a/src/langchain_google_cloud_sql_mysql/vectorstore.py +++ b/src/langchain_google_cloud_sql_mysql/vectorstore.py @@ -16,6 +16,7 @@ from __future__ import annotations import json +import warnings from typing import Any, Dict, Iterable, List, Optional, Tuple, Type, Union import numpy as np @@ -27,6 +28,7 @@ from .indexes import ( DEFAULT_QUERY_OPTIONS, DistanceMeasure, + IndexType, QueryOptions, SearchType, VectorIndex, @@ -244,42 +246,51 @@ def apply_vector_index(self, vector_index: VectorIndex): # Construct the default index name if not vector_index.name: vector_index.name = f"{self.table_name}_{DEFAULT_INDEX_NAME_SUFFIX}" - query_template = f"CALL mysql.create_vector_index('{vector_index.name}', '{self.db_name}.{self.table_name}', '{self.embedding_column}', '{{}}');" - self.__exec_apply_vector_index(query_template, vector_index) + index_name = ( + vector_index.name.split(".")[1] + if "." in vector_index.name + else vector_index.name + ) + + distance_measure = ( + vector_index.distance_measure + or self.query_options.distance_measure + or DistanceMeasure.L2_SQUARED + ) + + options = ["QUANTIZER=SQ8", f"DISTANCE_MEASURE={distance_measure.value}"] + if vector_index.num_partitions: + options.append(f"NUM_LEAVES={vector_index.num_partitions}") + + options_str = " ".join(options) + stmt = f"CREATE VECTOR INDEX `{index_name}` ON `{self.table_name}` (`{self.embedding_column}`) USING SCANN {options_str};" + self.engine._execute_outside_tx(stmt) # After applying an index to the table, set the query option search type to be ANN self.query_options.search_type = SearchType.ANN + self.query_options.distance_measure = distance_measure def alter_vector_index(self, vector_index: VectorIndex): existing_index_name = self._get_vector_index_name() if not existing_index_name: raise ValueError("No existing vector index found.") + existing_short_name = ( + existing_index_name.split(".")[1] + if "." in existing_index_name + else existing_index_name + ) if not vector_index.name: - vector_index.name = existing_index_name.split(".")[1] - if existing_index_name.split(".")[1] != vector_index.name: + vector_index.name = existing_short_name + req_short_name = ( + vector_index.name.split(".")[1] + if "." in vector_index.name + else vector_index.name + ) + if existing_short_name != req_short_name: raise ValueError( f"Existing index name {existing_index_name} does not match the new index name {vector_index.name}." ) - query_template = ( - f"CALL mysql.alter_vector_index('{existing_index_name}', '{{}}');" - ) - self.__exec_apply_vector_index(query_template, vector_index) - - def __exec_apply_vector_index(self, query_template: str, vector_index: VectorIndex): - index_options = [] - if vector_index.index_type: - index_options.append(f"index_type={vector_index.index_type.value}") - if vector_index.distance_measure: - index_options.append( - f"distance_measure={vector_index.distance_measure.value}" - ) - if vector_index.num_partitions: - index_options.append(f"num_partitions={vector_index.num_partitions}") - if vector_index.num_neighbors: - index_options.append(f"num_neighbors={vector_index.num_neighbors}") - index_options_query = ",".join(index_options) - - stmt = query_template.format(index_options_query) - self.engine._execute_outside_tx(stmt) + self.drop_vector_index() + self.apply_vector_index(vector_index) def _get_vector_index_name(self): query = "SELECT index_name FROM mysql.vector_indexes WHERE table_name=:table_name;" @@ -295,8 +306,13 @@ def _get_vector_index_name(self): def drop_vector_index(self): existing_index_name = self._get_vector_index_name() if existing_index_name: + index_name = ( + existing_index_name.split(".")[1] + if "." in existing_index_name + else existing_index_name + ) self.engine._execute_outside_tx( - f"CALL mysql.drop_vector_index('{existing_index_name}');" + f"DROP INDEX `{index_name}` ON `{self.table_name}`;" ) self.query_options.search_type = SearchType.KNN return existing_index_name @@ -662,24 +678,34 @@ def _query_collection( if query_options.num_partitions and query_options.search_type == SearchType.KNN: raise ValueError("num_partitions is not supported for the search type KNN") - k = k if k else query_options.num_neighbors + k = k if k is not None else query_options.num_neighbors distance_function = ( f"{query_options.distance_measure.value}_distance" if query_options.distance_measure != DistanceMeasure.DOT_PRODUCT else query_options.distance_measure.value ) bind_params: Dict[str, Any] = {"embedding": str(embedding)} - if query_options.search_type == SearchType.KNN: - filter = f"WHERE {filter}" if filter else "" - stmt = f"SELECT {column_query}, {distance_function}({self.embedding_column}, string_to_vector(:embedding)) AS distance FROM `{self.table_name}` {filter} ORDER BY distance LIMIT {k};" + + search_type = query_options.search_type + if search_type == SearchType.ANN: + if not k or k <= 0: + warnings.warn("A LIMIT is required for ANN search. Falling back to KNN search.") + search_type = SearchType.KNN + elif not query_options.distance_measure: + raise ValueError("Specifying the distance measure is required for ANN search.") + + if search_type == SearchType.KNN: + where_clause = f"WHERE {filter}" if filter else "" + stmt = f"SELECT {column_query}, {distance_function}({self.embedding_column}, string_to_vector(:embedding)) AS distance FROM `{self.table_name}` {where_clause} ORDER BY distance LIMIT {k};" else: - filter = f"AND {filter}" if filter else "" - num_partitions = ( - f",num_partitions={query_options.num_partitions}" + where_clause = f"WHERE {filter}" if filter else "" + num_leaves = ( + f",num_leaves_to_search={query_options.num_partitions}" if query_options.num_partitions else "" ) - stmt = f"SELECT {column_query}, {distance_function}({self.embedding_column}, string_to_vector(:embedding)) AS distance FROM `{self.table_name}` WHERE NEAREST({self.embedding_column}) TO (string_to_vector(:embedding), 'num_neighbors={k}{num_partitions}') {filter} ORDER BY distance;" + ann_options = f"distance_measure={query_options.distance_measure.value}{num_leaves}" + stmt = f"SELECT {column_query}, APPROX_DISTANCE({self.embedding_column}, string_to_vector(:embedding), '{ann_options}') AS distance FROM `{self.table_name}` {where_clause} ORDER BY distance LIMIT {k};" # return self.engine._fetch(stmt) if map_results: diff --git a/tests/unit/test_vectorstore_sql.py b/tests/unit/test_vectorstore_sql.py new file mode 100644 index 0000000..0a4d730 --- /dev/null +++ b/tests/unit/test_vectorstore_sql.py @@ -0,0 +1,152 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest.mock import MagicMock +import pytest +from langchain_community.embeddings import DeterministicFakeEmbedding + +from langchain_google_cloud_sql_mysql import ( + DistanceMeasure, + IndexType, + MySQLVectorStore, + QueryOptions, + SearchType, + VectorIndex, +) + + +@pytest.fixture +def mock_engine(): + engine = MagicMock() + engine._fetch.side_effect = lambda stmt, *args, **kwargs: ( + [{"DATABASE()": "test_db"}] + if "DATABASE()" in stmt + else [ + {"COLUMN_NAME": "langchain_id", "DATA_TYPE": "varchar"}, + {"COLUMN_NAME": "content", "DATA_TYPE": "text"}, + {"COLUMN_NAME": "embedding", "DATA_TYPE": "varbinary"}, + ] + if "information_schema" in stmt.lower() + else [{"COLUMN_NAME": "langchain_id"}, {"COLUMN_NAME": "content"}, {"COLUMN_NAME": "embedding"}] + if "COLUMNS" in stmt + else [] + ) + return engine + + +@pytest.fixture +def vectorstore(mock_engine): + embedding_service = DeterministicFakeEmbedding(size=8) + vs = MySQLVectorStore( + engine=mock_engine, + embedding_service=embedding_service, + table_name="my_table", + ) + mock_engine._execute_outside_tx.reset_mock() + mock_engine._fetch.reset_mock() + return vs + + +def test_apply_vector_index_ga_syntax(vectorstore, mock_engine): + index = VectorIndex( + name="my_idx", + index_type=IndexType.TREE_SQ, + distance_measure=DistanceMeasure.L2_SQUARED, + num_partitions=10, + ) + vectorstore.apply_vector_index(index) + + expected_sql = ( + "CREATE VECTOR INDEX `my_idx` ON `my_table` (`embedding`) " + "USING SCANN QUANTIZER=SQ8 DISTANCE_MEASURE=l2_squared NUM_LEAVES=10;" + ) + mock_engine._execute_outside_tx.assert_called_once_with(expected_sql) + assert vectorstore.query_options.search_type == SearchType.ANN + + +def test_drop_vector_index_ga_syntax(vectorstore, mock_engine): + vectorstore._get_vector_index_name = MagicMock(return_value="test_db.my_idx") + vectorstore.drop_vector_index() + + expected_sql = "DROP INDEX `my_idx` ON `my_table`;" + mock_engine._execute_outside_tx.assert_called_once_with(expected_sql) + assert vectorstore.query_options.search_type == SearchType.KNN + + +def test_alter_vector_index_ga_syntax(vectorstore, mock_engine): + vectorstore._get_vector_index_name = MagicMock(return_value="test_db.my_idx") + new_index = VectorIndex( + name="my_idx", + index_type=IndexType.TREE_SQ, + distance_measure=DistanceMeasure.COSINE, + num_partitions=20, + ) + vectorstore.alter_vector_index(new_index) + + drop_sql = "DROP INDEX `my_idx` ON `my_table`;" + create_sql = ( + "CREATE VECTOR INDEX `my_idx` ON `my_table` (`embedding`) " + "USING SCANN QUANTIZER=SQ8 DISTANCE_MEASURE=cosine NUM_LEAVES=20;" + ) + mock_engine._execute_outside_tx.assert_any_call(drop_sql) + mock_engine._execute_outside_tx.assert_any_call(create_sql) + + +def test_query_collection_knn(vectorstore, mock_engine): + vectorstore.__get_column_names = MagicMock( + return_value=["langchain_id", "content", "vector_to_string(embedding) as embedding"] + ) + query_options = QueryOptions( + distance_measure=DistanceMeasure.L2_SQUARED, + search_type=SearchType.KNN, + ) + vectorstore._query_collection( + embedding=[0.1, 0.2, 0.3], + k=5, + filter="content != 'test'", + query_options=query_options, + ) + + called_stmt = mock_engine._fetch.call_args[0][0] + expected_stmt = ( + "SELECT langchain_id, content, vector_to_string(embedding) as embedding, " + "l2_squared_distance(embedding, string_to_vector(:embedding)) AS distance " + "FROM `my_table` WHERE content != 'test' ORDER BY distance LIMIT 5;" + ) + assert called_stmt == expected_stmt + + +def test_query_collection_ann_ga_syntax(vectorstore, mock_engine): + vectorstore.__get_column_names = MagicMock( + return_value=["langchain_id", "content", "vector_to_string(embedding) as embedding"] + ) + query_options = QueryOptions( + distance_measure=DistanceMeasure.COSINE, + search_type=SearchType.ANN, + num_partitions=10, + ) + vectorstore._query_collection( + embedding=[0.1, 0.2, 0.3], + k=4, + filter="content != 'test'", + query_options=query_options, + ) + + called_stmt = mock_engine._fetch.call_args[0][0] + expected_stmt = ( + "SELECT langchain_id, content, vector_to_string(embedding) as embedding, " + "APPROX_DISTANCE(embedding, string_to_vector(:embedding), 'distance_measure=cosine,num_leaves_to_search=10') AS distance " + "FROM `my_table` WHERE content != 'test' ORDER BY distance LIMIT 4;" + ) + assert called_stmt == expected_stmt From 2b63a96dda717c9bfada6e98e06f3df3ce27ddd3 Mon Sep 17 00:00:00 2001 From: Averi Kitsch Date: Fri, 7 Aug 2026 14:32:45 -0700 Subject: [PATCH 2/2] Delete tests/unit/test_vectorstore_sql.py --- tests/unit/test_vectorstore_sql.py | 152 ----------------------------- 1 file changed, 152 deletions(-) delete mode 100644 tests/unit/test_vectorstore_sql.py diff --git a/tests/unit/test_vectorstore_sql.py b/tests/unit/test_vectorstore_sql.py deleted file mode 100644 index 0a4d730..0000000 --- a/tests/unit/test_vectorstore_sql.py +++ /dev/null @@ -1,152 +0,0 @@ -# Copyright 2024 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from unittest.mock import MagicMock -import pytest -from langchain_community.embeddings import DeterministicFakeEmbedding - -from langchain_google_cloud_sql_mysql import ( - DistanceMeasure, - IndexType, - MySQLVectorStore, - QueryOptions, - SearchType, - VectorIndex, -) - - -@pytest.fixture -def mock_engine(): - engine = MagicMock() - engine._fetch.side_effect = lambda stmt, *args, **kwargs: ( - [{"DATABASE()": "test_db"}] - if "DATABASE()" in stmt - else [ - {"COLUMN_NAME": "langchain_id", "DATA_TYPE": "varchar"}, - {"COLUMN_NAME": "content", "DATA_TYPE": "text"}, - {"COLUMN_NAME": "embedding", "DATA_TYPE": "varbinary"}, - ] - if "information_schema" in stmt.lower() - else [{"COLUMN_NAME": "langchain_id"}, {"COLUMN_NAME": "content"}, {"COLUMN_NAME": "embedding"}] - if "COLUMNS" in stmt - else [] - ) - return engine - - -@pytest.fixture -def vectorstore(mock_engine): - embedding_service = DeterministicFakeEmbedding(size=8) - vs = MySQLVectorStore( - engine=mock_engine, - embedding_service=embedding_service, - table_name="my_table", - ) - mock_engine._execute_outside_tx.reset_mock() - mock_engine._fetch.reset_mock() - return vs - - -def test_apply_vector_index_ga_syntax(vectorstore, mock_engine): - index = VectorIndex( - name="my_idx", - index_type=IndexType.TREE_SQ, - distance_measure=DistanceMeasure.L2_SQUARED, - num_partitions=10, - ) - vectorstore.apply_vector_index(index) - - expected_sql = ( - "CREATE VECTOR INDEX `my_idx` ON `my_table` (`embedding`) " - "USING SCANN QUANTIZER=SQ8 DISTANCE_MEASURE=l2_squared NUM_LEAVES=10;" - ) - mock_engine._execute_outside_tx.assert_called_once_with(expected_sql) - assert vectorstore.query_options.search_type == SearchType.ANN - - -def test_drop_vector_index_ga_syntax(vectorstore, mock_engine): - vectorstore._get_vector_index_name = MagicMock(return_value="test_db.my_idx") - vectorstore.drop_vector_index() - - expected_sql = "DROP INDEX `my_idx` ON `my_table`;" - mock_engine._execute_outside_tx.assert_called_once_with(expected_sql) - assert vectorstore.query_options.search_type == SearchType.KNN - - -def test_alter_vector_index_ga_syntax(vectorstore, mock_engine): - vectorstore._get_vector_index_name = MagicMock(return_value="test_db.my_idx") - new_index = VectorIndex( - name="my_idx", - index_type=IndexType.TREE_SQ, - distance_measure=DistanceMeasure.COSINE, - num_partitions=20, - ) - vectorstore.alter_vector_index(new_index) - - drop_sql = "DROP INDEX `my_idx` ON `my_table`;" - create_sql = ( - "CREATE VECTOR INDEX `my_idx` ON `my_table` (`embedding`) " - "USING SCANN QUANTIZER=SQ8 DISTANCE_MEASURE=cosine NUM_LEAVES=20;" - ) - mock_engine._execute_outside_tx.assert_any_call(drop_sql) - mock_engine._execute_outside_tx.assert_any_call(create_sql) - - -def test_query_collection_knn(vectorstore, mock_engine): - vectorstore.__get_column_names = MagicMock( - return_value=["langchain_id", "content", "vector_to_string(embedding) as embedding"] - ) - query_options = QueryOptions( - distance_measure=DistanceMeasure.L2_SQUARED, - search_type=SearchType.KNN, - ) - vectorstore._query_collection( - embedding=[0.1, 0.2, 0.3], - k=5, - filter="content != 'test'", - query_options=query_options, - ) - - called_stmt = mock_engine._fetch.call_args[0][0] - expected_stmt = ( - "SELECT langchain_id, content, vector_to_string(embedding) as embedding, " - "l2_squared_distance(embedding, string_to_vector(:embedding)) AS distance " - "FROM `my_table` WHERE content != 'test' ORDER BY distance LIMIT 5;" - ) - assert called_stmt == expected_stmt - - -def test_query_collection_ann_ga_syntax(vectorstore, mock_engine): - vectorstore.__get_column_names = MagicMock( - return_value=["langchain_id", "content", "vector_to_string(embedding) as embedding"] - ) - query_options = QueryOptions( - distance_measure=DistanceMeasure.COSINE, - search_type=SearchType.ANN, - num_partitions=10, - ) - vectorstore._query_collection( - embedding=[0.1, 0.2, 0.3], - k=4, - filter="content != 'test'", - query_options=query_options, - ) - - called_stmt = mock_engine._fetch.call_args[0][0] - expected_stmt = ( - "SELECT langchain_id, content, vector_to_string(embedding) as embedding, " - "APPROX_DISTANCE(embedding, string_to_vector(:embedding), 'distance_measure=cosine,num_leaves_to_search=10') AS distance " - "FROM `my_table` WHERE content != 'test' ORDER BY distance LIMIT 4;" - ) - assert called_stmt == expected_stmt