feat(genai/embeddings): migrate code retrieval and model tuning examples to genai

genai/embeddings/code_retrieval_example.py

@@ -0,0 +1,78 @@
+# Copyright 2026 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
+#
+#    https://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.
+
+# [START aiplatform_genai_embedding_code_retrieval]
+
+import os
+
+from google import genai
+
+# TODO (Developer) set the following environment variables.
+PROJECT_ID = os.getenv("PROJECT_ID")
+LOCATION_ID = os.getenv("LOCATION_ID", "us-central1")
+MODEL_NAME = os.getenv("MODEL_NAME", "gemini-embedding-001")
+
+QUERY_LINES = ["Retrieve a function that adds two numbers"]
+CODE_RETRIEVAL_QUERY = "CODE_RETRIEVAL_QUERY"
+RETRIEVAL_DOCUMENT = "RETRIEVAL_DOCUMENT"
+SOURCE_CODE = [
+    "def func(a, b): return a + b",
+    "def func(a, b): return a - b",
+    "def func(a, b): return (a ** 2 + b ** 2) ** 0.5",
+]
+
+
+def embed_test() -> (
+    tuple[genai.types.EmbedContentConfig, genai.types.EmbedContentResponse]
+):
+    """Generates embeddings for source code indexing and code search queries using the Gemini API.
+
+    Returns:
+        tuple[genai.types.EmbedContentConfig, genai.types.EmbedContentConfig]: A tuple containing
+        the final source code indexing response and search query embedding response.
+    """
+    client = genai.Client(enterprise=True, project=PROJECT_ID, location=LOCATION_ID)
+
+    # Index Source Code
+    for line in SOURCE_CODE:
+        config = genai.types.EmbedContentConfig(task_type=RETRIEVAL_DOCUMENT)
+
+        index_response = client.models.embed_content(
+            model=MODEL_NAME, contents=line, config=config
+        )
+
+        print(
+            f"Task: {RETRIEVAL_DOCUMENT} | "
+            f"Vector length: {len(index_response.embeddings)} | "
+            f"Preview: {index_response.embeddings[:3]}..."
+        )
+
+    # Embed Search Prompts
+    for line in QUERY_LINES:
+        config = genai.types.EmbedContentConfig(task_type=CODE_RETRIEVAL_QUERY)
+
+        query_response = client.models.embed_content(
+            model=MODEL_NAME, contents=line, config=config
+        )
+
+        print(
+            f"Task: {CODE_RETRIEVAL_QUERY} | "
+            f"Vector length: {len(query_response.embeddings)} | "
+            f"Preview: {query_response.embeddings[:3]}..."
+        )
+
+    return index_response, query_response
+
+
+# [END aiplatform_genai_embedding_code_retrieval]

genai/embeddings/model_tuning_example.py

@@ -0,0 +1,78 @@
+# Copyright 2026 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
+#
+#    https://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.
+
+# [START aiplatform_genai_embedding_model_tuning]
+import os
+
+from google.cloud import aiplatform
+
+# TODO (Developer) set the following environment variables.
+PROJECT_ID = os.getenv("PROJECT_ID")
+LOCATION_ID = os.getenv("LOCATION_ID", "us-central1")
+MODEL_NAME = os.getenv("MODEL_NAME", "text-embedding-004")
+# A storage bucket: gs://your-bucket-name/embedding-tuning-output
+OUTPUT_URI = os.getenv("OUTPUT_DIR")
+
+TRAIN_LABEL_PATH = (
+    "gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/train.tsv"
+)
+TEST_LABEL_PATH = (
+    "gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/test.tsv"
+)
+CORPUS_PATH = (
+    "gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/corpus.jsonl"
+)
+QUERIES_PATH = (
+    "gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/queries.jsonl"
+)
+
+ACCELERATOR_TYPE = "NVIDIA_L4"
+
+# Official Google Cloud KFP pipeline template URI for text embedding model tuning
+EMBEDDING_TUNING_PIPELINE_URI = "https://us-kfp.pkg.dev/ml-pipeline/llm-text-embedding/tune-text-embedding-model/v1.1.3"
+
+
+def tune_embedding_model() -> aiplatform.PipelineJob:
+    """Tune an embedding model using the specified parameters."""
+
+    aiplatform.init(project=PROJECT_ID, location=LOCATION_ID)
+
+    # Configure parameters expected by the embedding tuning pipeline template
+    pipeline_parameters = {
+        "base_model_version_id": MODEL_NAME,
+        "corpus_path": CORPUS_PATH,
+        "queries_path": QUERIES_PATH,
+        "train_label_path": TRAIN_LABEL_PATH,
+        "test_label_path": TEST_LABEL_PATH,
+        "accelerator_type": ACCELERATOR_TYPE,
+    }
+
+    # Instantiate the Vertex AI Pipeline job
+    pipeline_job = aiplatform.PipelineJob(
+        display_name="tune-text-embedding-model-job",
+        template_path=EMBEDDING_TUNING_PIPELINE_URI,
+        pipeline_root=OUTPUT_URI,
+        parameter_values=pipeline_parameters,
+        project=PROJECT_ID,
+        location=LOCATION_ID,
+    )
+
+    pipeline_job.submit()
+
+    print(f"Pipeline submitted successfully: {pipeline_job.resource_name}")
+
+    return pipeline_job
+
+
+# [END aiplatform_genai_embedding_model_tuning]

genai/embeddings/requirements-test.txt

@@ -1,2 +1,2 @@
-google-api-core==2.24.0
-pytest==9.0.3; python_version >= "3.10"
+google-api-core==2.33.0
+pytest==9.1.1

genai/embeddings/requirements.txt

@@ -1 +1,2 @@
-google-genai==1.42.0
+google-genai==2.16.0
+google-cloud-aiplatform[pipelines]==1.163.0

genai/embeddings/test_embeddings_examples.py

@@ -18,7 +18,9 @@
 
 import os
 
+import code_retrieval_example
 import embeddings_docretrieval_with_txt
+import model_tuning_example
 
 os.environ["GOOGLE_GENAI_USE_ENTERPRISE"] = "True"
 os.environ["GOOGLE_CLOUD_LOCATION"] = "us-central1"
@@ -29,3 +31,12 @@
 def test_embeddings_docretrieval_with_txt() -> None:
     response = embeddings_docretrieval_with_txt.embed_content()
     assert response
+
+
+def test_code_retrieval_example() -> None:
+    response = code_retrieval_example.embed_test()
+    assert response
+
+def test_model_tuning_example() -> None:
+    response = model_tuning_example.tune_embedding_model()
+    assert response