agent-starter-pack 0.1.7__py3-none-any.whl → 0.2.1__py3-none-any.whl

src/cli/utils/version.py

@@ -0,0 +1,87 @@
+# Copyright 2025 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
+#
+#     http://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.
+
+"""Version checking utilities for the CLI."""
+
+import logging
+from importlib.metadata import PackageNotFoundError, version
+
+import requests
+from packaging import version as pkg_version
+from rich.console import Console
+
+console = Console()
+
+PACKAGE_NAME = "agent-starter-pack"
+
+
+def get_current_version() -> str:
+    """Get the current installed version of the package."""
+    try:
+        return version(PACKAGE_NAME)
+    except PackageNotFoundError:
+        # For development environments where package isn't installed
+        return "0.0.0"  # Default if version can't be determined
+
+
+def get_latest_version() -> str:
+    """Get the latest version available on PyPI."""
+    try:
+        response = requests.get(f"https://pypi.org/pypi/{PACKAGE_NAME}/json", timeout=2)
+        if response.status_code == 200:
+            return response.json()["info"]["version"]
+        return "0.0.0"
+    except Exception:
+        return "0.0.0"  # Default if PyPI can't be reached
+
+
+def check_for_updates() -> tuple[bool, str, str]:
+    """Check if a newer version of the package is available.
+
+    Returns:
+        Tuple of (needs_update, current_version, latest_version)
+    """
+    current = get_current_version()
+    latest = get_latest_version()
+
+    needs_update = pkg_version.parse(latest) > pkg_version.parse(current)
+
+    return needs_update, current, latest
+
+
+def display_update_message() -> None:
+    """Check for updates and display a message if an update is available."""
+    try:
+        needs_update, current, latest = check_for_updates()
+
+        if needs_update:
+            console.print(
+                f"\n[yellow]⚠️  Update available: {current} → {latest}[/]",
+                highlight=False,
+            )
+            console.print(
+                f"[yellow]Run `pip install --upgrade {PACKAGE_NAME}` to update.",
+                highlight=False,
+            )
+            console.print(
+                f"[yellow]Or, if you used pipx: `pipx upgrade {PACKAGE_NAME}`",
+                highlight=False,
+            )
+            console.print(
+                f"[yellow]Or, if you used uv: `uv pip install --upgrade {PACKAGE_NAME}`",
+                highlight=False,
+            )
+    except Exception as e:
+        # Don't let version checking errors affect the CLI
+        logging.debug(f"Error checking for updates: {e}")

src/data_ingestion/README.md

@@ -1,8 +1,8 @@
 # Data Ingestion Pipeline
 
-This pipeline automates the ingestion of data into Vertex AI Search, streamlining the process of building Retrieval Augmented Generation (RAG) applications. 
+This pipeline automates the ingestion of data into Vertex AI{%- if cookiecutter.datastore_type == "vertex_ai_vector_search" %} Vector{%- endif %} Search, streamlining the process of building Retrieval Augmented Generation (RAG) applications. 
 
-It orchestrates the complete workflow: loading data, chunking it into manageable segments, generating embeddings using Vertex AI Embeddings, and importing the processed data into your Vertex AI Search datastore.
+It orchestrates the complete workflow: loading data, chunking it into manageable segments, generating embeddings using Vertex AI Embeddings, and importing the processed data into your Vertex AI{%- if cookiecutter.datastore_type == "vertex_ai_vector_search" %} Vector{%- endif %} Search datastore.
 
 You can trigger the pipeline for an initial data load or schedule it to run periodically, ensuring your search index remains current. Vertex AI Pipelines provides the orchestration and monitoring capabilities for this process.
 
@@ -37,7 +37,7 @@ uv sync --frozen
 **c. Execute the Pipeline:**
 
 Run the following command to execute the data ingestion pipeline. Replace the placeholder values with your actual project details.
-
+{%- if cookiecutter.datastore_type == "vertex_ai_search" %}
 ```bash
 PROJECT_ID="YOUR_PROJECT_ID"
 REGION="us-central1"
@@ -47,22 +47,26 @@ uv run data_ingestion_pipeline/submit_pipeline.py \
     --region=$REGION \
     --data-store-region=$DATA_STORE_REGION \
     --data-store-id="sample-datastore" \
-    --service-account="vertexai-pipelines-sa@$PROJECT_ID.iam.gserviceaccount.com" \
-    --pipeline-root="gs://$PROJECT_ID-pipeline-artifacts" \
+    --service-account="{{cookiecutter.project_name}}-rag@$PROJECT_ID.iam.gserviceaccount.com" \
+    --pipeline-root="gs://$PROJECT_ID-{{cookiecutter.project_name}}-rag" \
     --pipeline-name="data-ingestion-pipeline"
 ```
-
-**Parameter Explanation:**
-
-*   `--project-id`: Your Google Cloud project ID.
-*   `--region`: The region where Vertex AI Pipelines will run (e.g., `us-central1`).
-*   `--data-store-region`: The region for Vertex AI Search operations (e.g., `us` or `eu`).
-*   `--data-store-id`: The ID of your Vertex AI Search datastore.
-*   `--service-account`: The service account email used for pipeline execution.  Ensure this service account has the necessary permissions (e.g., Vertex AI User, Storage Object Admin).
-*   `--pipeline-root`: The Google Cloud Storage (GCS) bucket for storing pipeline artifacts.
-*   `--pipeline-name`: A descriptive name for your pipeline.
-*   `--schedule-only` (Optional): If specified, the pipeline will only be scheduled and not executed immediately. Requires `--cron-schedule`.
-*   `--cron-schedule` (Optional): A cron expression defining the pipeline's schedule (e.g., `"0 9 * * 1"` for every Monday at 9:00 AM UTC).
+{%- elif cookiecutter.datastore_type == "vertex_ai_vector_search" %}
+```bash
+PROJECT_ID="YOUR_PROJECT_ID"
+REGION="us-central1"
+VECTOR_SEARCH_INDEX="YOUR_VECTOR_SEARCH_INDEX"
+VECTOR_SEARCH_INDEX_ENDPOINT="YOUR_VECTOR_SEARCH_INDEX_ENDPOINT"
+uv run data_ingestion_pipeline/submit_pipeline.py \
+    --project-id=$PROJECT_ID \
+    --region=$REGION \
+    --vector-search-index=$VECTOR_SEARCH_INDEX \
+    --vector-search-index-endpoint=$VECTOR_SEARCH_INDEX_ENDPOINT \
+    --service-account="{{cookiecutter.project_name}}-rag@$PROJECT_ID.iam.gserviceaccount.com" \
+    --pipeline-root="gs://$PROJECT_ID-{{cookiecutter.project_name}}-rag" \
+    --pipeline-name="data-ingestion-pipeline"
+```
+{%- endif %}
 
 **d. Pipeline Scheduling and Execution:**
 
@@ -74,6 +78,7 @@ The pipeline's configuration and execution status will be printed to the console
 
 ## Testing Your RAG Application
 
-Once the data ingestion pipeline completes successfully, you can test your RAG application with Vertex AI Search.
-
+Once the data ingestion pipeline completes successfully, you can test your RAG application with Vertex AI{%- if cookiecutter.datastore_type == "vertex_ai_vector_search" %} Vector{%- endif %} Search.
+{%- if cookiecutter.datastore_type == "vertex_ai_search" %}
 > **Troubleshooting:** If you encounter the error `"google.api_core.exceptions.InvalidArgument: 400 The embedding field path: embedding not found in schema"` after the initial data ingestion, wait a few minutes and try again. This delay allows Vertex AI Search to fully index the ingested data.
+{%- endif %}

src/data_ingestion/data_ingestion_pipeline/components/ingest_data.py

@@ -11,12 +11,13 @@
 # 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.
+# ruff: noqa
 
 from kfp.dsl import Dataset, Input, component
-
+{% if cookiecutter.datastore_type == "vertex_ai_search" %}
 
 @component(
-    base_image="us-docker.pkg.dev/production-ai-template/starter-pack/data_processing:0.1"
+    base_image="us-docker.pkg.dev/production-ai-template/starter-pack/data_processing:0.2"
 )
 def ingest_data(
     project_id: str,
@@ -37,6 +38,7 @@ def ingest_data(
     """
     import json
     import logging
+    import time
 
     from google.api_core.client_options import ClientOptions
     from google.cloud import discoveryengine
@@ -173,3 +175,134 @@ def ingest_data(
         input_files_uri=input_files.uri,
     )
     logging.info("Data import completed")
+    logging.info(
+        "Sleeping for 3 minutes to allow Vertex AI Search to properly index the data..."
+    )
+    time.sleep(180)  # Sleep for 180 seconds (3 minutes)
+    logging.info("Sleep completed. Data indexing should now be complete.")
+{% elif cookiecutter.datastore_type == "vertex_ai_vector_search" %}
+from google_cloud_pipeline_components.types.artifact_types import BQTable
+
+
+@component(
+    base_image="us-docker.pkg.dev/production-ai-template/starter-pack/data_processing:0.2"
+)
+def ingest_data(
+    project_id: str,
+    location: str,
+    vector_search_index: str,
+    vector_search_index_endpoint: str,
+    vector_search_data_bucket_name: str,
+    schedule_time: str,
+    ingestion_batch_size: int,
+    input_table: Input[BQTable],
+    is_incremental: bool = True,
+    look_back_days: int = 1,
+) -> None:
+    """Process and ingest documents into Vertex AI Vector Search.
+
+    Args:
+        project_id: Google Cloud project ID
+    """
+    import logging
+    from datetime import datetime, timedelta
+
+    import bigframes.pandas as bpd
+    from google.cloud import aiplatform
+    from langchain_google_vertexai import VectorSearchVectorStore
+    from langchain_google_vertexai import VertexAIEmbeddings
+
+    # Initialize logging
+    logging.basicConfig(level=logging.INFO)
+
+    # Initialize clients
+    logging.info("Initializing clients...")
+    bpd.options.bigquery.project = project_id
+    bpd.options.bigquery.location = location
+    logging.info("Clients initialized.")
+
+    # Set date range for data fetch
+    schedule_time_dt: datetime = datetime.fromisoformat(
+        schedule_time.replace("Z", "+00:00")
+    )
+    if schedule_time_dt.year == 1970:
+        logging.warning(
+            "Pipeline schedule not set. Setting schedule_time to current date."
+        )
+        schedule_time_dt = datetime.now()
+
+    # Note: The following line sets the schedule time 5 years back to allow sample data to be present.
+    # For your use case, please comment out the following line to use the actual schedule time.
+    schedule_time_dt = schedule_time_dt - timedelta(days=5 * 365)
+
+    START_DATE: datetime = schedule_time_dt - timedelta(
+        days=look_back_days
+    )  # Start date for data processing window
+    END_DATE: datetime = schedule_time_dt  # End date for data processing window
+
+    logging.info(f"Date range set: START_DATE={START_DATE}, END_DATE={END_DATE}")
+
+    dataset = input_table.metadata["datasetId"]
+    table = input_table.metadata["tableId"]
+
+    query = f"""
+        SELECT
+            question_id
+            , last_edit_date
+            , full_text_md
+            , text_chunk
+            , chunk_id
+            , embedding
+        FROM  {project_id}.{dataset}.{table}
+        WHERE TRUE
+                {f'AND DATETIME(creation_timestamp) BETWEEN DATETIME("{START_DATE}") AND DATETIME("{END_DATE}")' if is_incremental else ""}
+    """
+    df = (
+        bpd.read_gbq(query)
+        .sort_values("last_edit_date", ascending=False)
+        .drop_duplicates("question_id")
+        .reset_index(drop=True)
+    )
+
+    aiplatform.init(
+        project=project_id,
+        location=location,
+        staging_bucket=vector_search_data_bucket_name,
+    )
+
+    embedding_model = VertexAIEmbeddings(model_name="text-embedding-005")
+    my_index = aiplatform.MatchingEngineIndex(vector_search_index)
+    my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint(
+        vector_search_index_endpoint
+    )
+    vector_store = VectorSearchVectorStore.from_components(
+        project_id=project_id,
+        region=location,
+        gcs_bucket_name=vector_search_data_bucket_name.replace("gs://", ""),
+        index_id=my_index.name,
+        endpoint_id=my_index_endpoint.name,
+        embedding=embedding_model,
+        stream_update=True,
+    )
+
+    for batch_num, start in enumerate(range(0, len(df), ingestion_batch_size)):
+        ids = (
+            df.iloc[start : start + ingestion_batch_size]
+            .question_id.astype(str)
+            .tolist()
+        )
+        texts = df.iloc[start : start + ingestion_batch_size].text_chunk.tolist()
+        embeddings = df.iloc[start : start + ingestion_batch_size].embedding.tolist()
+        metadatas = (
+            df.iloc[start : start + ingestion_batch_size]
+            .drop(columns=["embedding", "last_edit_date"])
+            .to_dict(orient="records")
+        )
+        vector_store.add_texts_with_embeddings(
+            ids=ids,
+            texts=texts,
+            embeddings=embeddings,
+            metadatas=metadatas,
+            is_complete_overwrite=True,
+        )
+{% endif %}

src/data_ingestion/data_ingestion_pipeline/components/process_data.py

@@ -11,6 +11,7 @@
 # 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.
+# ruff: noqa
 
 """
 This component is derived from the notebook:
@@ -21,9 +22,9 @@ It leverages BigQuery for data processing. We also suggest looking at remote fun
 
 from kfp.dsl import Dataset, Output, component
 
-
+{% if cookiecutter.datastore_type == "vertex_ai_search" %}
 @component(
-    base_image="us-docker.pkg.dev/production-ai-template/starter-pack/data_processing:0.1"
+    base_image="us-docker.pkg.dev/production-ai-template/starter-pack/data_processing:0.2"
 )
 def process_data(
     project_id: str,
@@ -319,3 +320,276 @@ def process_data(
     )
     extract_job.result()
     logging.info("Exported to JSONL.")
+{% elif cookiecutter.datastore_type == "vertex_ai_vector_search" %}
+from google_cloud_pipeline_components.types.artifact_types import BQTable
+
+
+@component(
+    base_image="us-docker.pkg.dev/production-ai-template/starter-pack/data_processing:0.2",
+)
+def process_data(
+    project_id: str,
+    schedule_time: str,
+    output_table: Output[BQTable],
+    is_incremental: bool = True,
+    look_back_days: int = 1,
+    chunk_size: int = 1500,
+    chunk_overlap: int = 20,
+    destination_dataset: str = "stackoverflow_data",
+    destination_table: str = "incremental_questions_embeddings",
+    deduped_table: str = "questions_embeddings",
+    location: str = "us-central1",
+) -> None:
+    """Process StackOverflow questions and answers by:
+    1. Fetching data from BigQuery
+    2. Converting HTML to markdown
+    3. Splitting text into chunks
+    4. Generating embeddings
+    5. Storing results in BigQuery
+    6. Exporting to JSONL
+
+    Args:
+        output_files: Output dataset path
+        is_incremental: Whether to process only recent data
+        look_back_days: Number of days to look back for incremental processing
+        chunk_size: Size of text chunks
+        chunk_overlap: Overlap between chunks
+        destination_dataset: BigQuery dataset for storing results
+        destination_table: Table for storing incremental results
+        deduped_table: Table for storing deduplicated results
+        location: BigQuery location
+    """
+    import logging
+    from datetime import datetime, timedelta
+
+    import backoff
+    import bigframes.ml.llm as llm
+    import bigframes.pandas as bpd
+    import google.api_core.exceptions
+    import swifter
+    from google.cloud import bigquery
+    from langchain.text_splitter import RecursiveCharacterTextSplitter
+    from markdownify import markdownify
+
+    # Initialize logging
+    logging.basicConfig(level=logging.INFO)
+    logging.info(f"Using {swifter} for apply operations.")
+
+    # Initialize clients
+    logging.info("Initializing clients...")
+    bq_client = bigquery.Client(project=project_id, location=location)
+    bpd.options.bigquery.project = project_id
+    bpd.options.bigquery.location = location
+    logging.info("Clients initialized.")
+
+    # Set date range for data fetch
+    schedule_time_dt: datetime = datetime.fromisoformat(
+        schedule_time.replace("Z", "+00:00")
+    )
+    if schedule_time_dt.year == 1970:
+        logging.warning(
+            "Pipeline schedule not set. Setting schedule_time to current date."
+        )
+        schedule_time_dt = datetime.now()
+
+    # Note: The following line sets the schedule time 5 years back to allow sample data to be present.
+    # For your use case, please comment out the following line to use the actual schedule time.
+    schedule_time_dt = schedule_time_dt - timedelta(days=5 * 365)
+
+    START_DATE: datetime = schedule_time_dt - timedelta(
+        days=look_back_days
+    )  # Start date for data processing window
+    END_DATE: datetime = schedule_time_dt  # End date for data processing window
+
+    logging.info(f"Date range set: START_DATE={START_DATE}, END_DATE={END_DATE}")
+
+    def fetch_stackoverflow_data(
+        dataset_suffix: str, start_date: str, end_date: str
+    ) -> bpd.DataFrame:
+        """Fetch StackOverflow data from BigQuery."""
+        query = f"""
+            SELECT
+                creation_date,
+                last_edit_date,
+                question_id,
+                question_title,
+                question_body AS question_text,
+                answers
+            FROM `production-ai-template.stackoverflow_qa_{dataset_suffix}.stackoverflow_python_questions_and_answers`
+            WHERE TRUE
+                {f'AND TIMESTAMP_TRUNC(creation_date, DAY) BETWEEN TIMESTAMP("{start_date}") AND TIMESTAMP("{end_date}")' if is_incremental else ""}
+        """
+        logging.info("Fetching StackOverflow data from BigQuery...")
+        return bpd.read_gbq(query)
+
+    def convert_html_to_markdown(html: str) -> str:
+        """Convert HTML into Markdown for easier parsing and rendering after LLM response."""
+        return markdownify(html).strip()
+
+    def create_answers_markdown(answers: list) -> str:
+        """Convert each answer's HTML to markdown and concatenate into a single markdown text."""
+        answers_md = ""
+        for index, answer_record in enumerate(answers):
+            answers_md += (
+                f"\n\n## Answer {index + 1}:\n"  # Answer number is H2 heading size
+            )
+            answers_md += convert_html_to_markdown(answer_record["body"])
+        return answers_md
+
+    def create_table_if_not_exist(
+        df: bpd.DataFrame,
+        project_id: str,
+        dataset_id: str,
+        table_id: str,
+        partition_column: str,
+        location: str = location,
+    ) -> None:
+        """Create BigQuery table with time partitioning if it doesn't exist."""
+        table_schema = bq_client.get_table(df.head(0).to_gbq()).schema
+        table = bigquery.Table(
+            f"{project_id}.{dataset_id}.{table_id}", schema=table_schema
+        )
+        table.time_partitioning = bigquery.TimePartitioning(
+            type_=bigquery.TimePartitioningType.DAY, field=partition_column
+        )
+
+        dataset = bigquery.Dataset(f"{project_id}.{dataset_id}")
+        dataset.location = location
+        bq_client.create_dataset(dataset, exists_ok=True)
+        bq_client.create_table(table=table, exists_ok=True)
+
+    # Fetch and preprocess data
+    logging.info("Fetching and preprocessing data...")
+    df = fetch_stackoverflow_data(
+        start_date=START_DATE.strftime("%Y-%m-%d"),
+        end_date=END_DATE.strftime("%Y-%m-%d"),
+        dataset_suffix=location.lower().replace("-", "_"),
+    )
+    df = (
+        df.sort_values("last_edit_date", ascending=False)
+        .drop_duplicates("question_id")
+        .reset_index(drop=True)
+    )
+    logging.info("Data fetched and preprocessed.")
+
+    # Convert content to markdown
+    logging.info("Converting content to markdown...")
+
+    # Create markdown fields efficiently
+    df["question_title_md"] = (
+        "# " + df["question_title"] + "\n"
+    )  # Title is H1 heading size
+    df["question_text_md"] = (
+        df["question_text"].to_pandas().swifter.apply(convert_html_to_markdown) + "\n"
+    )
+    df["answers_md"] = df["answers"].to_pandas().swifter.apply(create_answers_markdown)
+
+    # Create a column containing the whole markdown text
+    df["full_text_md"] = (
+        df["question_title_md"] + df["question_text_md"] + df["answers_md"]
+    )
+    logging.info("Content converted to markdown.")
+
+    # Keep only necessary columns
+    df = df[["last_edit_date", "question_id", "question_text", "full_text_md"]]
+
+    # Split text into chunks
+    logging.info("Splitting text into chunks...")
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=chunk_size,
+        chunk_overlap=chunk_overlap,
+        length_function=len,
+    )
+
+    df["text_chunk"] = (
+        df["full_text_md"]
+        .to_pandas()
+        .astype(object)
+        .swifter.apply(text_splitter.split_text)
+    )
+    logging.info("Text split into chunks.")
+
+    # Create chunk IDs and explode chunks into rows
+    logging.info("Creating chunk IDs and exploding chunks into rows...")
+    chunk_ids = [
+        str(idx) for text_chunk in df["text_chunk"] for idx in range(len(text_chunk))
+    ]
+    df = df.explode("text_chunk").reset_index(drop=True)
+    df["chunk_id"] = df["question_id"].astype("string") + "__" + chunk_ids
+    logging.info("Chunk IDs created and chunks exploded.")
+
+    # Generate embeddings
+    logging.info("Generating embeddings...")
+
+    # The first invocation in a new project might fail due to permission propagation.
+    @backoff.on_exception(
+        backoff.expo, google.api_core.exceptions.InvalidArgument, max_tries=10
+    )
+    def create_embedder() -> llm.TextEmbeddingGenerator:
+        return llm.TextEmbeddingGenerator(model_name="text-embedding-005")
+
+    embedder = create_embedder()
+
+    embeddings_df = embedder.predict(df["text_chunk"])
+    logging.info("Embeddings generated.")
+
+    df = df.assign(
+        embedding=embeddings_df["ml_generate_embedding_result"],
+        embedding_statistics=embeddings_df["ml_generate_embedding_statistics"],
+        embedding_status=embeddings_df["ml_generate_embedding_status"],
+        creation_timestamp=datetime.now(),
+    )
+
+    # Store results in BigQuery
+    PARTITION_DATE_COLUMN = "creation_timestamp"
+
+    # Create and populate incremental table
+    logging.info("Creating and populating incremental table...")
+    create_table_if_not_exist(
+        df=df,
+        project_id=project_id,
+        dataset_id=destination_dataset,
+        table_id=destination_table,
+        partition_column=PARTITION_DATE_COLUMN,
+    )
+
+    if_exists_mode = "append" if is_incremental else "replace"
+    df.to_gbq(
+        destination_table=f"{destination_dataset}.{destination_table}",
+        if_exists=if_exists_mode,
+    )
+    logging.info("Incremental table created and populated.")
+
+    # Create deduplicated table
+    logging.info("Creating deduplicated table...")
+    df_questions = bpd.read_gbq(
+        f"{destination_dataset}.{destination_table}", use_cache=False
+    )
+    max_date_df = (
+        df_questions.groupby("question_id")["creation_timestamp"].max().reset_index()
+    )
+    df_questions_dedup = max_date_df.merge(
+        df_questions, how="inner", on=["question_id", "creation_timestamp"]
+    )
+
+    create_table_if_not_exist(
+        df=df_questions_dedup,
+        project_id=project_id,
+        dataset_id=destination_dataset,
+        table_id=deduped_table,
+        partition_column=PARTITION_DATE_COLUMN,
+    )
+
+    df_questions_dedup.to_gbq(
+        destination_table=f"{destination_dataset}.{deduped_table}",
+        if_exists="replace",
+    )
+    logging.info("Deduplicated table created and populated.")
+    # Set artifact metadata (important!)
+    output_table.uri = (
+        f"bq://{project_id}.{destination_dataset}.{deduped_table}"  # Full BQ URI
+    )
+    output_table.metadata["projectId"] = project_id
+    output_table.metadata["datasetId"] = destination_dataset
+    output_table.metadata["tableId"] = deduped_table
+{% endif %}

src/data_ingestion/data_ingestion_pipeline/pipeline.py

@@ -21,15 +21,22 @@ from kfp import dsl
 def pipeline(
     project_id: str,
     location: str,
-    data_store_region: str,
-    data_store_id: str,
     is_incremental: bool = True,
     look_back_days: int = 1,
     chunk_size: int = 1500,
     chunk_overlap: int = 20,
-    destination_dataset: str = "stackoverflow_data",
     destination_table: str = "incremental_questions_embeddings",
     deduped_table: str = "questions_embeddings",
+    destination_dataset: str = "{{cookiecutter.project_name | replace('-', '_')}}_stackoverflow_data",
+{%- if cookiecutter.datastore_type == "vertex_ai_search" %}
+    data_store_region: str = "",
+    data_store_id: str = "",
+{%- elif cookiecutter.datastore_type == "vertex_ai_vector_search" %}
+    vector_search_index: str = "",
+    vector_search_index_endpoint: str = "",
+    vector_search_data_bucket_name: str = "",
+    ingestion_batch_size: int = 1000,
+{%- endif %}
 ) -> None:
     """Processes data and ingests it into a datastore for RAG Retrieval"""
 
@@ -45,9 +52,10 @@ def pipeline(
         destination_table=destination_table,
         deduped_table=deduped_table,
         location=location,
-        embedding_column="embedding",
+{%- if cookiecutter.datastore_type == "vertex_ai_search" %}
+        embedding_column="embedding",{% endif %}
     ).set_retry(num_retries=2)
-
+{% if cookiecutter.datastore_type == "vertex_ai_search" %}
     # Ingest the processed data into Vertex AI Search datastore
     ingest_data(
         project_id=project_id,
@@ -56,3 +64,18 @@ def pipeline(
         data_store_id=data_store_id,
         embedding_column="embedding",
     ).set_retry(num_retries=2)
+{% elif cookiecutter.datastore_type == "vertex_ai_vector_search" %}
+    # Ingest the processed data into Vertex AI Vector Search
+    ingest_data(
+        project_id=project_id,
+        location=location,
+        vector_search_index=vector_search_index,
+        vector_search_index_endpoint=vector_search_index_endpoint,
+        vector_search_data_bucket_name=vector_search_data_bucket_name,
+        input_table=processed_data.output,
+        schedule_time=dsl.PIPELINE_JOB_SCHEDULE_TIME_UTC_PLACEHOLDER,
+        is_incremental=False,
+        look_back_days=look_back_days,
+        ingestion_batch_size=ingestion_batch_size,
+    ).set_retry(num_retries=2)
+{% endif %}