refactor(alloydb): migrate batch processing notebook to genai SDK (#14537)

* refactor(alloydb): migrate batch processing notebook to genai SDK

 Updates snippets in the `embeddings_batch_processing.ipynb` notebook to use the newer `google-genai` SDK.

* Addressed a comment.

Author: XrossFox

alloydb/notebooks/embeddings_batch_processing.ipynb

@@ -77,7 +77,8 @@
     "    google-cloud-alloydb-connector[asyncpg]==1.4.0 \\\n",
     "    sqlalchemy==2.0.36 \\\n",
     "    pandas==2.2.3 \\\n",
-    "    vertexai==1.70.0 \\\n",
+    "    google-cloud-aiplatform==1.165.1 \\\n",
+    "    google-genai==2.19.0 \\\n",
     "    asyncio==3.4.3 \\\n",
     "    greenlet==3.1.1 \\\n",
     "    --quiet"
@@ -792,119 +793,133 @@
    },
    "outputs": [],
    "source": [
-    "from google.api_core.exceptions import ResourceExhausted\n",
-    "from typing import Union\n",
-    "from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel\n",
+    "from typing import Any, AsyncIterator, List, Optional, Union\n",
     "\n",
+    "from google import genai\n",
+    "from google.genai import types\n",
     "\n",
     "async def embed_text(\n",
     "    batch_data: List[dict[str, Any]],\n",
-    "    model: TextEmbeddingModel,\n",
     "    cols_to_embed: List[str],\n",
+    "    client: genai.Client,\n",
+    "    model_name: str = \"text-embedding-004\",\n",
     "    task_type: str = \"SEMANTIC_SIMILARITY\",\n",
     "    retries: int = 100,\n",
     "    delay: int = 30,\n",
     ") -> List[dict[str, Union[List[float], str]]]:\n",
-    "    \"\"\"Embeds text data from a batch of records using a Vertex AI embedding model.\n",
+    "    \"\"\"Embeds text data from a batch of records using the google-genai SDK.\n",
     "\n",
     "    Args:\n",
-    "      batch_data: A data batch containing records with text data to embed.\n",
-    "      model: The Vertex AI `TextEmbeddingModel` to use for generating embeddings.\n",
-    "      cols_to_embed: A list of column names containing the data to be embedded.\n",
-    "      task_type: The task type for the embedding model. Defaults to\n",
-    "        \"SEMANTIC_SIMILARITY\".\n",
-    "        Supported task types: https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/task-types\n",
-    "      retries: The maximum number of times to retry embedding generation in case\n",
-    "        of errors. Defaults to 100.\n",
-    "      delay: The delay in seconds between retries. Defaults to 30.\n",
+    "        batch_data: A data batch containing records with text data to embed.\n",
+    "        cols_to_embed: A list of column names containing the data to be embedded.\n",
+    "        model_name: The embedding model ID.\n",
+    "        task_type: The task type for the embedding model.\n",
+    "        retries: The maximum number of retry attempts in case of errors.\n",
+    "        delay: The delay in seconds between retries.\n",
+    "        client: An optional genai.Client instance for connection reuse.\n",
     "\n",
     "    Returns:\n",
-    "      A list of records containing ids and embeddings.\n",
-    "      Example:\n",
-    "        [\n",
-    "          {\n",
-    "            'id': 'id1',\n",
-    "            'col1_embedding': [1.0, 1.1, ...],\n",
-    "            'col2_embedding': [2.0, 2.1, ...],\n",
-    "            ...\n",
-    "          },\n",
-    "          ...\n",
-    "        ]\n",
-    "      where col1 and col2 are columns containing data to be embedded.\n",
+    "        A list of records containing IDs and mapped column embeddings.\n",
     "    Raises:\n",
-    "      Exception: Raises the encountered exception if all retries fail.\n",
+    "        Exception: Raises the encountered exception if all retries fail.\n",
     "    \"\"\"\n",
     "    logger = logging.getLogger(\"embed_objects\")\n",
     "    global total_char_count\n",
     "\n",
-    "    # Place all of the embeddings into a single list\n",
+    "    # Extract non-empty text strings to embed\n",
     "    inputs = []\n",
     "    for row in batch_data:\n",
     "        for col in cols_to_embed:\n",
     "            if col in row and row[col]:\n",
-    "                inputs.append(TextEmbeddingInput(row[col], task_type))\n",
+    "                inputs.append(str(row[col]))\n",
     "\n",
-    "    # Retry loop\n",
     "    for attempt in range(retries):\n",
     "        try:\n",
-    "            # Get embeddings for the text data\n",
-    "            embeddings = await model.get_embeddings_async(inputs)\n",
+    "            # Asynchronous API call using client_instance.aio\n",
+    "            response = await client.aio.models.embed_content(\n",
+    "                model=model_name,\n",
+    "                contents=inputs,\n",
+    "                config=types.EmbedContentConfig(\n",
+    "                    task_type=task_type,\n",
+    "                ),\n",
+    "            )\n",
+    "\n",
+    "            # Track character metrics\n",
+    "            total_char_count += sum(len(text) for text in inputs)\n",
     "\n",
-    "            # Increase total char count\n",
-    "            total_char_count += sum([len(input.text) for input in inputs])\n",
+    "            # Map response embeddings back to dataset record structure\n",
+    "            embeddings_list = (\n",
+    "                response.embeddings\n",
+    "                if response.embeddings\n",
+    "                else [response.embedding]\n",
+    "            )\n",
+    "            embedding_iter = iter(embeddings_list)\n",
     "\n",
-    "            # group the results together by id\n",
-    "            embedding_iter = iter(embeddings)\n",
     "            results = []\n",
     "            for row in batch_data:\n",
-    "                r = {\"id\": row[\"id\"]}\n",
+    "                r = {\"id\": row.get(\"id\")}\n",
     "                for col in cols_to_embed:\n",
     "                    if col in row and row[col]:\n",
-    "                        r[f\"{col}_embedding\"] = str(next(embedding_iter).values)\n",
+    "                        embedding_obj = next(embedding_iter)\n",
+    "                        r[f\"{col}_embedding\"] = str(embedding_obj.values)\n",
     "                    else:\n",
     "                        r[f\"{col}_embedding\"] = None\n",
     "                results.append(r)\n",
+    "\n",
     "            return results\n",
     "\n",
     "        except Exception as e:\n",
-    "            if attempt < retries - 1:  # Retry only if attempts are left\n",
-    "                logger.warning(f\"Error: {e}. Retrying in {delay} seconds...\")\n",
-    "                await asyncio.sleep(delay)  # Wait before retrying\n",
+    "            if attempt < retries - 1:\n",
+    "                logger.warning(\n",
+    "                    f\"Error: {e}. Retrying in {delay} seconds (attempt {attempt + 1}/{retries})...\"\n",
+    "                )\n",
+    "                await asyncio.sleep(delay)\n",
     "            else:\n",
-    "                logger.error(f\"Failed to get embeddings for data: {batch_data} after {retries} attempts.\")\n",
+    "                logger.error(\n",
+    "                    f\"Failed to get embeddings after {retries} attempts: {e}\"\n",
+    "                )\n",
+    "                raise e\n",
+    "\n",
     "    return []\n",
     "\n",
     "\n",
     "async def embed_objects_concurrently(\n",
     "    cols_to_embed: List[str],\n",
     "    batch_data: AsyncIterator[List[dict[str, Any]]],\n",
-    "    model: TextEmbeddingModel,\n",
-    "    task_type: str,\n",
+    "    client: genai.Client,\n",
+    "    model_name: str = \"text-embedding-004\",\n",
+    "    task_type: str = \"SEMANTIC_SIMILARITY\",\n",
     "    max_concurrency: int = 5,\n",
     ") -> AsyncIterator[List[dict[str, Union[str, List[float]]]]]:\n",
     "    \"\"\"Embeds text data concurrently from an asynchronous batch data generator.\n",
     "\n",
     "    Args:\n",
-    "      cols_to_embed: A list of column names containing the data to be embedded.\n",
-    "      batch_data: A data batch containing records with text data to embed.\n",
-    "      model: The Vertex AI `TextEmbeddingModel` to use for generating embeddings.\n",
-    "      task_type: The task type for the embedding model.\n",
-    "        Supported task types: https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/task-types\n",
-    "      max_concurrency: The maximum number of embedding tasks to run concurrently.\n",
-    "        Defaults to 5.\n",
+    "        cols_to_embed: A list of column names containing the data to be embedded.\n",
+    "        batch_data: An async generator yielding data batches with records.\n",
+    "        model_name: The embedding model ID.\n",
+    "        task_type: The task type for the embedding model.\n",
+    "        max_concurrency: The maximum number of concurrent tasks to execute.\n",
+    "        client: An optional genai.Client instance for connection reuse.\n",
+    "\n",
     "    Yields:\n",
-    "      A list of records containing ids and embeddings.\n",
+    "        A list of records containing IDs and mapped column embeddings.\n",
     "    \"\"\"\n",
     "    logger = logging.getLogger(\"embed_objects\")\n",
     "\n",
-    "    # Keep track of pending tasks\n",
     "    pending: set[asyncio.Task] = set()\n",
     "    has_next = True\n",
+    "\n",
     "    while pending or has_next:\n",
     "        while len(pending) < max_concurrency and has_next:\n",
     "            try:\n",
     "                data = await batch_data.__anext__()\n",
-    "                coro = embed_text(data, model, cols_to_embed, task_type)\n",
+    "                coro = embed_text(\n",
+    "                    batch_data=data,\n",
+    "                    cols_to_embed=cols_to_embed,\n",
+    "                    model_name=model_name,\n",
+    "                    task_type=task_type,\n",
+    "                    client=client,\n",
+    "                )\n",
     "                pending.add(asyncio.ensure_future(coro))\n",
     "            except StopAsyncIteration:\n",
     "                has_next = False\n",
@@ -915,7 +930,9 @@
     "            )\n",
     "            for task in done:\n",
     "                result = task.result()\n",
-    "                logger.info(f\"Embedding task completed: Processed {len(result)} rows.\")\n",
+    "                logger.info(\n",
+    "                    f\"Embedding task completed: Processed {len(result)} rows.\"\n",
+    "                )\n",
     "                yield result"
    ]
   },
@@ -1045,9 +1062,9 @@
    },
    "outputs": [],
    "source": [
-    "import vertexai\n",
     "import time\n",
-    "from vertexai.language_models import TextEmbeddingModel\n",
+    "\n",
+    "from google import genai\n",
     "\n",
     "### Define variables ###\n",
     "\n",
@@ -1080,40 +1097,40 @@
     "):\n",
     "    \"\"\"Orchestrates the end-to-end workflow for generating and storing embeddings.\n",
     "\n",
-    "    The workflow includes the following major steps:\n",
-    "\n",
-    "    1. Data Retrieval: Fetches data from the database that requires embedding.\n",
-    "    2. Batching: Divides the data into batches for optimized processing.\n",
-    "    3. Embedding Generation: Generates embeddings concurrently for the batched\n",
-    "        data using the Vertex AI model.\n",
-    "    4. Database Update: Updates the database concurrently with the generated\n",
-    "        embeddings.\n",
+    "    Workflow Steps:\n",
+    "    1. Connection Initialization: Initializes DB pool.\n",
+    "    2. Client Initialization: Instantiates shared GenAI client for workflow lifetime.\n",
+    "    3. Data Retrieval & Batching: Fetches and chunks source data into async streams.\n",
+    "    4. Embedding Generation: Uses google-genai client asynchronously and concurrently.\n",
+    "    5. Database Update: Stores output embeddings concurrently in the database.\n",
     "\n",
     "    Args:\n",
-    "        pool_size: The size of the database connection pool. Defaults to 10.\n",
-    "        embed_data_concurrency: The maximum number of concurrent tasks for generating embeddings.\n",
-    "            Defaults to 20.\n",
-    "        batch_update_concurrency: The maximum number of concurrent tasks for updating the database.\n",
-    "            Defaults to 10.\n",
+    "        pool_size: The size of the database connection pool.\n",
+    "        embed_data_concurrency: Max concurrent tasks for generating embeddings.\n",
+    "        batch_update_concurrency: Max concurrent tasks for database updates.\n",
     "    \"\"\"\n",
-    "    # Set up connections to the database\n",
+    "    # Set up database connection pool\n",
     "    pool = await init_connection_pool(connector, database_name, pool_size=pool_size)\n",
     "\n",
-    "    # Initialise VertexAI and the model to be used to generate embeddings\n",
-    "    vertexai.init(project=project_id, location=region)\n",
-    "    model = TextEmbeddingModel.from_pretrained(model_name)\n",
+    "    # Initialize single GenAI client instance for the top-level workflow lifetime\n",
+    "    client = genai.Client(vertexai=True, project=project_id, location=region)\n",
     "\n",
     "    start_time = time.monotonic()\n",
     "\n",
-    "    # Fetch source data from the database\n",
+    "    # Fetch source data from database\n",
     "    source_data = get_source_data(pool, cols_to_embed)\n",
     "\n",
-    "    # Divide the source data into batches for efficient processing\n",
+    "    # Divide source data into asynchronous batches\n",
     "    batch_data = batch_source_data(source_data, cols_to_embed)\n",
     "\n",
-    "    # Generate embeddings for the batched data concurrently\n",
+    "    # Generate embeddings concurrently using the GenAI SDK\n",
     "    embeddings_data = embed_objects_concurrently(\n",
-    "        cols_to_embed, batch_data, model, task, max_concurrency=embed_data_concurrency\n",
+    "        cols_to_embed=cols_to_embed,\n",
+    "        batch_data=batch_data,\n",
+    "        model_name=model_name,\n",
+    "        task_type=task,\n",
+    "        max_concurrency=embed_data_concurrency,\n",
+    "        client=client,\n",
     "    )\n",
     "\n",
     "    # Update the database with the generated embeddings concurrently\n",
@@ -1124,7 +1141,7 @@
     "    end_time = time.monotonic()\n",
     "    elapsed_time = end_time - start_time\n",
     "\n",
-    "    # Release database connections and close the connector\n",
+    "    # Release database resources\n",
     "    await pool.dispose()\n",
     "    await connector.close()\n",
     "\n",