linter fix and add vector index tests

Author: dishaprakash

samples/index_tuning_sample/index_search.py

@@ -32,14 +32,12 @@
     vector_table_name,
 )
 from langchain_google_vertexai import VertexAIEmbeddings
+
+from langchain_google_alloydb_pg import AlloyDBEngine, AlloyDBVectorStore
 from langchain_google_alloydb_pg.indexes import (
     HNSWIndex,
     HNSWQueryOptions,
     IVFFlatIndex,
-)
-
-from langchain_google_alloydb_pg import AlloyDBEngine, AlloyDBVectorStore
-from langchain_google_alloydb_pg.indexes import (
     IVFIndex,
     ScaNNIndex,
 )

src/langchain_google_alloydb_pg/__init__.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from langchain_postgres import Column
+
 from .chat_message_history import AlloyDBChatMessageHistory
 from .checkpoint import AlloyDBSaver
 from .embeddings import AlloyDBEmbeddings
@@ -20,7 +22,6 @@
 from .model_manager import AlloyDBModel, AlloyDBModelManager
 from .vectorstore import AlloyDBVectorStore
 from .version import __version__
-from langchain_postgres import Column
 
 __all__ = [
     "AlloyDBEngine",

src/langchain_google_alloydb_pg/async_vectorstore.py

@@ -28,6 +28,9 @@
 from langchain_core.documents import Document
 from langchain_core.embeddings import Embeddings
 from langchain_core.vectorstores import VectorStore, utils
+from sqlalchemy import RowMapping, text
+from sqlalchemy.ext.asyncio import AsyncEngine
+
 from langchain_google_alloydb_pg.indexes import (
     DEFAULT_DISTANCE_STRATEGY,
     DEFAULT_INDEX_NAME_SUFFIX,
@@ -36,8 +39,6 @@
     ExactNearestNeighbor,
     QueryOptions,
 )
-from sqlalchemy import RowMapping, text
-from sqlalchemy.ext.asyncio import AsyncEngine
 
 from .engine import AlloyDBEngine
 

src/langchain_google_alloydb_pg/indexes.py

@@ -16,17 +16,17 @@
 from dataclasses import dataclass, field
 
 from langchain_postgres.v2.indexes import (
-  BaseIndex,
-  DistanceStrategy,
-  QueryOptions,
-  StrategyMixin,
-  DEFAULT_DISTANCE_STRATEGY,
-  DEFAULT_INDEX_NAME_SUFFIX,
-  ExactNearestNeighbor,
-  HNSWIndex,
-  HNSWQueryOptions,
-  IVFFlatIndex,
-  IVFFlatQueryOptions,
+    DEFAULT_DISTANCE_STRATEGY,
+    DEFAULT_INDEX_NAME_SUFFIX,
+    BaseIndex,
+    DistanceStrategy,
+    ExactNearestNeighbor,
+    HNSWIndex,
+    HNSWQueryOptions,
+    IVFFlatIndex,
+    IVFFlatQueryOptions,
+    QueryOptions,
+    StrategyMixin,
 )
 
 

src/langchain_google_alloydb_pg/vectorstore.py

@@ -20,6 +20,7 @@
 from langchain_core.documents import Document
 from langchain_core.embeddings import Embeddings
 from langchain_core.vectorstores import VectorStore
+
 from langchain_google_alloydb_pg.indexes import (
     DEFAULT_DISTANCE_STRATEGY,
     BaseIndex,

tests/test_async_vectorstore_index.py

@@ -21,16 +21,16 @@
 import pytest_asyncio
 from langchain_core.documents import Document
 from langchain_core.embeddings import DeterministicFakeEmbedding
+from sqlalchemy import text
+
+from langchain_google_alloydb_pg import AlloyDBEngine
+from langchain_google_alloydb_pg.async_vectorstore import AsyncAlloyDBVectorStore
 from langchain_google_alloydb_pg.indexes import (
     DEFAULT_INDEX_NAME_SUFFIX,
     DistanceStrategy,
     HNSWIndex,
     IVFFlatIndex,
 )
-from sqlalchemy import text
-
-from langchain_google_alloydb_pg import AlloyDBEngine
-from langchain_google_alloydb_pg.async_vectorstore import AsyncAlloyDBVectorStore
 
 DEFAULT_TABLE = "test_table" + str(uuid.uuid4()).replace("-", "_")
 DEFAULT_INDEX_NAME = DEFAULT_TABLE + DEFAULT_INDEX_NAME_SUFFIX

tests/test_async_vectorstore_search.py

@@ -19,17 +19,15 @@
 import pytest_asyncio
 from langchain_core.documents import Document
 from langchain_core.embeddings import DeterministicFakeEmbedding
-from langchain_google_alloydb_pg.indexes import (
-    DistanceStrategy,
-    HNSWQueryOptions,
-)
 from metadata_filtering_data import FILTERING_TEST_CASES, METADATAS
 from PIL import Image
 from sqlalchemy import text
 
 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,
+    HNSWQueryOptions,
     ScaNNQueryOptions,
 )
 

tests/test_indexes.py

@@ -14,8 +14,12 @@
 
 import warnings
 
-from langchain_google_alloydb_pg.indexes import (
+from langchain_google_alloydb_pg.indexes import (  # type: ignore
     DistanceStrategy,
+    HNSWIndex,
+    HNSWQueryOptions,
+    IVFFlatIndex,
+    IVFFlatQueryOptions,
     IVFIndex,
     IVFQueryOptions,
     ScaNNIndex,
@@ -44,6 +48,42 @@ def test_distance_strategy(self):
         scann_index = ScaNNIndex(distance_strategy=DistanceStrategy.INNER_PRODUCT)
         assert scann_index.get_index_function() == "dot_prod"
 
+    def test_ivfflat_index(self):
+        index = IVFFlatIndex(name="test_index", lists=200)
+        assert index.index_type == "ivfflat"
+        assert index.lists == 200
+        assert index.index_options() == "(lists = 200)"
+
+    def test_ivfflat_query_options(self):
+        options = IVFFlatQueryOptions(probes=2)
+        assert options.to_parameter() == ["ivfflat.probes = 2"]
+
+        with warnings.catch_warnings(record=True) as w:
+            options.to_string()
+            assert len(w) == 1
+            assert "to_string is deprecated, use to_parameter instead." in str(
+                w[-1].message
+            )
+
+    def test_hnsw_index(self):
+        index = HNSWIndex(name="test_index", m=32, ef_construction=128)
+        assert index.index_type == "hnsw"
+        assert index.m == 32
+        assert index.ef_construction == 128
+        assert index.index_options() == "(m = 32, ef_construction = 128)"
+
+    def test_hnsw_query_options(self):
+        options = HNSWQueryOptions(ef_search=80)
+        assert options.to_parameter() == ["hnsw.ef_search = 80"]
+
+        with warnings.catch_warnings(record=True) as w:
+            options.to_string()
+
+            assert len(w) == 1
+            assert "to_string is deprecated, use to_parameter instead." in str(
+                w[-1].message
+            )
+
     def test_ivf_index(self):
         index = IVFIndex(name="test_index", lists=200)
         assert index.index_type == "ivf"

tests/test_loader.py

@@ -25,7 +25,7 @@
     AlloyDBDocumentSaver,
     AlloyDBEngine,
     AlloyDBLoader,
-    Column
+    Column,
 )
 
 project_id = os.environ["PROJECT_ID"]

tests/test_vectorstore_embeddings.py

@@ -18,16 +18,16 @@
 import pytest
 import pytest_asyncio
 from langchain_core.documents import Document
-from langchain_google_alloydb_pg.indexes import DistanceStrategy, HNSWQueryOptions
 from sqlalchemy import text
 
 from langchain_google_alloydb_pg import (
     AlloyDBEmbeddings,
     AlloyDBEngine,
     AlloyDBModelManager,
     AlloyDBVectorStore,
-    Column
+    Column,
 )
+from langchain_google_alloydb_pg.indexes import DistanceStrategy, HNSWQueryOptions
 
 DEFAULT_TABLE = "test_table" + str(uuid.uuid4()).replace("-", "_")
 DEFAULT_TABLE_SYNC = "test_table" + str(uuid.uuid4()).replace("-", "_")