fix!: Update AlloyDBModel based on ml extension v1.5.2 (#500)

* chore: Remove generate_headers_fn from model manager

Removed the generate_headers_fn parameter from the model manager.

* Update test_vectorstore_embeddings.py

* Refactor embeddings_service to use synchronous methods

* Corrected model manager instantiation

Fix async call to create_sync for model manager

* Convert embeddings_service to async fixture

* Update test_embeddings.py

* Update model qualified name in tests

* Update test_embeddings.py

* Update model_manager.py

* Add batch transform functions and update version check

* Fix input_batch_transform_fn assignment syntax

* Fix syntax error in model_manager.py

* Update google_ml_integration extension version check

* Update google_ml_integration version requirement to 1.3

* Update google_ml_integration extension version to 1.5.2

Author: dishaprakash

src/langchain_google_alloydb_pg/model_manager.py

@@ -36,6 +36,8 @@ def __init__(
         input_transform_fn: Optional[str],
         output_transform_fn: Optional[str],
         generate_headers_fn: Optional[str] = None,
+        input_batch_transform_fn: Optional[str] = None,
+        output_batch_transform_fn: Optional[str] = None,
         **kwargs: Any,
     ):
         self.model_id = model_id
@@ -49,6 +51,8 @@ def __init__(
         self.output_transform_fn = output_transform_fn
         # List models is returning column name "header_gen_fn"
         self.generate_headers_fn = generate_headers_fn or kwargs.get("header_gen_fn")
+        self.input_batch_transform_fn = input_batch_transform_fn
+        self.output_batch_transform_fn = output_batch_transform_fn
 
 
 class AlloyDBModelManager:
@@ -170,14 +174,14 @@ async def __avalidate(self) -> None:
         """Private async function to validate prerequisites.
 
         Raises:
-            Exception if google_ml_integration EXTENSION is not 1.3.
+            Exception if google_ml_integration EXTENSION is not 1.5.2.
             Exception if google_ml_integration.enable_model_support DB Flag not set.
         """
         extension_version = await self.__fetch_google_ml_extension()
         db_flag = await self.__fetch_db_flag()
-        if extension_version < "1.3":
+        if extension_version < "1.5.2":
             raise Exception(
-                "Please upgrade google_ml_integration EXTENSION to version 1.3 or above."
+                "Please upgrade google_ml_integration EXTENSION to version 1.5.2 or above."
             )
         if db_flag != "on":
             raise Exception(
@@ -214,13 +218,15 @@ async def __aget_model(self, model_id: str) -> Optional[AlloyDBModel]:
                 model_qualified_name VARCHAR,
                 model_auth_type google_ml.auth_type,
                 model_auth_id VARCHAR,
-                generate_headers_fn VARCHAR,
+                header_gen_fn VARCHAR,
                 input_transform_fn VARCHAR,
-                output_transform_fn VARCHAR)"""
+                output_transform_fn VARCHAR,
+                input_batch_transform_fn VARCHAR,
+                output_batch_transform_fn VARCHAR)"""
 
         try:
             result = await self.__query_db(query)
-        except Exception:
+        except Exception as e:
             return None
         data_class = self.__convert_dict_to_dataclass(result)[0]
         return data_class
@@ -285,13 +291,13 @@ async def __adrop_model(self, model_id: str) -> None:
             await conn.commit()
 
     async def __fetch_google_ml_extension(self) -> str:
-        """Creates the Google ML Extension if it does not exist and returns the version number (Default creates version 1.3)."""
+        """Creates the Google ML Extension if it does not exist and returns the version number (Default creates version 1.5.2)."""
         create_extension_query = """
         DO $$
         BEGIN
         IF NOT EXISTS (
           SELECT 1 FROM pg_extension WHERE extname = 'google_ml_integration' )
-          THEN CREATE EXTENSION google_ml_integration VERSION '1.3' CASCADE;
+          THEN CREATE EXTENSION google_ml_integration VERSION '1.5.2' CASCADE;
         END IF;
         END
         $$;

tests/test_embeddings.py

@@ -31,6 +31,7 @@
 instance_id = os.environ["INSTANCE_ID"]
 db_name = os.environ["DATABASE_ID"]
 table_name = "test-table" + str(uuid.uuid4())
+embedding_model = "text-embedding-005" + str(uuid.uuid4()).replace("-", "_")
 
 
 @pytest.mark.asyncio
@@ -66,7 +67,7 @@ async def sync_engine(self):
 
     @pytest.fixture(scope="module")
     def model_id(self) -> str:
-        return "text-embedding-005"
+        return embedding_model
 
     @pytest_asyncio.fixture
     async def embeddings(self, engine, model_id):
@@ -77,7 +78,7 @@ async def embeddings(self, engine, model_id):
             await model_manager.acreate_model(
                 model_id=model_id,
                 model_provider="google",
-                model_qualified_name=model_id,  # assuming model is built-in
+                model_qualified_name="text-embedding-005",  # assuming model is built-in
                 model_type="text_embedding",
             )
         return AlloyDBEmbeddings.create_sync(engine=engine, model_id=model_id)

tests/test_vectorstore_embeddings.py

@@ -32,7 +32,7 @@
 DEFAULT_TABLE = "test_table" + str(uuid.uuid4()).replace("-", "_")
 DEFAULT_TABLE_SYNC = "test_table" + str(uuid.uuid4()).replace("-", "_")
 CUSTOM_TABLE = "test_table_custom" + str(uuid.uuid4()).replace("-", "_")
-DEFAULT_EMBEDDING_MODEL = "text-embedding-005"
+DEFAULT_EMBEDDING_MODEL = "text-embedding-005" + str(uuid.uuid4()).replace("-", "_")
 VECTOR_SIZE = 768
 
 
@@ -114,10 +114,11 @@ async def embeddings_service(self, engine):
             await model_manager.acreate_model(
                 model_id=DEFAULT_EMBEDDING_MODEL,
                 model_provider="google",
-                model_qualified_name=DEFAULT_EMBEDDING_MODEL,  # assuming model is built-in
+                model_qualified_name="text-embedding-005",  # assuming model is built-in
                 model_type="text_embedding",
             )
-        return await AlloyDBEmbeddings.create(engine, DEFAULT_EMBEDDING_MODEL)
+        yield await AlloyDBEmbeddings.create(engine, DEFAULT_EMBEDDING_MODEL)
+        await model_manager.adrop_model(DEFAULT_EMBEDDING_MODEL)
 
     @pytest_asyncio.fixture(scope="class")
     async def vs(self, engine, embeddings_service):
@@ -308,8 +309,19 @@ async def engine_sync(
         await engine.close()
 
     @pytest_asyncio.fixture(scope="class")
-    def embeddings_service(self, engine_sync):
+    async def embeddings_service(self, engine_sync):
+        model_manager = AlloyDBModelManager.create_sync(engine=engine_sync)
+        model = await model_manager.aget_model(model_id=DEFAULT_EMBEDDING_MODEL)
+        if not model:
+            # create model if not exists
+            await model_manager.acreate_model(
+                model_id=DEFAULT_EMBEDDING_MODEL,
+                model_provider="google",
+                model_qualified_name="text-embedding-005",  # assuming model is built-in
+                model_type="text_embedding",
+            )
         return AlloyDBEmbeddings.create_sync(engine_sync, DEFAULT_EMBEDDING_MODEL)
+        await model_manager.adrop_model(DEFAULT_EMBEDDING_MODEL)
 
     @pytest_asyncio.fixture(scope="class")
     async def vs_custom(self, engine_sync, embeddings_service):