tests: upgrade embedding model in tests (#419)

Author: vishwarajanand

tests/test_embeddings.py

@@ -19,7 +19,11 @@
 import pytest_asyncio
 from langchain_core.documents import Document
 
-from langchain_google_alloydb_pg import AlloyDBEmbeddings, AlloyDBEngine
+from langchain_google_alloydb_pg import (
+    AlloyDBEmbeddings,
+    AlloyDBEngine,
+    AlloyDBModelManager,
+)
 
 project_id = os.environ["PROJECT_ID"]
 region = os.environ["REGION"]
@@ -62,10 +66,20 @@ async def sync_engine(self):
 
     @pytest.fixture(scope="module")
     def model_id(self) -> str:
-        return "textembedding-gecko@001"
+        return "text-embedding-005"
 
     @pytest_asyncio.fixture
-    def embeddings(self, engine, model_id):
+    async def embeddings(self, engine, model_id):
+        model_manager = await AlloyDBModelManager.create(engine=engine)
+        model = await model_manager.aget_model(model_id=model_id)
+        if not model:
+            # create model if not exists
+            await model_manager.acreate_model(
+                model_id=model_id,
+                model_provider="google",
+                model_qualified_name=model_id,  # assuming model is built-in
+                model_type="text_embedding",
+            )
         return AlloyDBEmbeddings.create_sync(engine=engine, model_id=model_id)
 
     async def test_model_exists(self, sync_engine):

tests/test_model_manager.py

@@ -28,7 +28,7 @@ def get_env_var(key: str, desc: str) -> str:
     return v
 
 
-EMBEDDING_MODEL_NAME = "textembedding-gecko@003" + str(uuid.uuid4()).replace("-", "_")
+EMBEDDING_MODEL_NAME = "text-embedding-005" + str(uuid.uuid4()).replace("-", "_")
 
 
 @pytest.mark.asyncio
@@ -79,7 +79,7 @@ async def test_acreate_model(self, model_manager):
         await model_manager.acreate_model(
             model_id=EMBEDDING_MODEL_NAME,
             model_provider="google",
-            model_qualified_name="textembedding-gecko@003",
+            model_qualified_name="text-embedding-005",
             model_type="text_embedding",
         )
 

tests/test_vectorstore_embeddings.py

@@ -23,6 +23,7 @@
 from langchain_google_alloydb_pg import (
     AlloyDBEmbeddings,
     AlloyDBEngine,
+    AlloyDBModelManager,
     AlloyDBVectorStore,
     Column,
 )
@@ -31,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 = "textembedding-gecko@001"
+DEFAULT_EMBEDDING_MODEL = "text-embedding-005"
 VECTOR_SIZE = 768
 
 
@@ -106,6 +107,16 @@ async def engine(
 
     @pytest_asyncio.fixture(scope="class")
     async def embeddings_service(self, engine):
+        model_manager = await AlloyDBModelManager.create(engine=engine)
+        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=DEFAULT_EMBEDDING_MODEL,  # assuming model is built-in
+                model_type="text_embedding",
+            )
         return await AlloyDBEmbeddings.create(engine, DEFAULT_EMBEDDING_MODEL)
 
     @pytest_asyncio.fixture(scope="class")
@@ -197,7 +208,7 @@ async def test_similarity_search_with_relevance_scores_threshold_cosine(self, vs
         )
         assert len(results) == 4
 
-        score_threshold = {"score_threshold": 0.73}
+        score_threshold = {"score_threshold": 0.65}
         results = await vs.asimilarity_search_with_relevance_scores(
             "foo", **score_threshold
         )