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

src/data_ingestion/README.md

@@ -1,84 +1,71 @@
+{%- if cookiecutter.datastore_type == "vertex_ai_search" -%}
+{%- set datastore_service_name = "Vertex AI Search" -%}
+{%- elif cookiecutter.datastore_type == "vertex_ai_vector_search" -%}
+{%- set datastore_service_name = "Vertex AI Vector Search" -%}
+{%- else -%}
+{%- set datastore_service_name = "Your Configured Datastore" -%}
+{%- endif -%}
+
 # Data Ingestion Pipeline
 
-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. 
+This pipeline automates the ingestion of data into {{ datastore_service_name }}, 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{%- if cookiecutter.datastore_type == "vertex_ai_vector_search" %} Vector{%- endif %} 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 {{ datastore_service_name }} 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.
 
 ## Prerequisites
 
-Before running the data ingestion pipeline, ensure you have completed the following:
+Before running any commands, ensure you have set your Google Cloud Project ID as an environment variable. This variable will be used by the subsequent `make` commands.
 
-1. **Set up Dev Terraform:** Follow the instructions in the parent [deployment/README.md - Dev Deployment section](../deployment/README.md#dev-deployment) to provision the necessary resources in your development environment using Terraform. This includes deploying a datastore and configuring the required permissions.
+```bash
+export PROJECT_ID="YOUR_PROJECT_ID"
+```
+Replace `"YOUR_PROJECT_ID"` with your actual Google Cloud Project ID.
 
-## Running the Data Ingestion Pipeline
+Now, you can set up the development environment:
 
-After setting up the Terraform infrastructure, you can test the data ingestion pipeline.
+1.  **Set up Dev Environment:** Use the following command from the root of the repository to provision the necessary resources in your development environment using Terraform. This includes deploying a datastore and configuring the required permissions.
 
-> **Note:** The initial pipeline execution might take longer as your project is configured for Vertex AI Pipelines.
+    ```bash
+    make setup-dev-env
+    ```
+    This command requires `terraform` to be installed and configured.
 
-**Steps:**
+## Running the Data Ingestion Pipeline
 
-**a. Navigate to the `data_ingestion` directory:**
+After setting up the infrastructure using `make setup-dev-env`, you can run the data ingestion pipeline.
 
-```bash
-cd data_ingestion
-```
+> **Note:** The initial pipeline execution might take longer as your project is configured for Vertex AI Pipelines.
 
-**b. Install Dependencies:**
+**Steps:**
 
-Install the required Python dependencies using uv:
+**a. Execute the Pipeline:**
+Run the following command from the root of the repository. Ensure the `PROJECT_ID` environment variable is still set in your current shell session (as configured in Prerequisites).
 
 ```bash
-uv sync --frozen
+make data-ingestion
 ```
 
-**c. Execute the Pipeline:**
-
-Run the following command to execute the data ingestion pipeline. Replace the placeholder values with your actual project details.
+This command handles installing dependencies (if needed via `make install`) and submits the pipeline job using the configuration derived from your project setup. The specific parameters passed to the underlying script depend on the `datastore_type` selected during project generation:
 {%- if cookiecutter.datastore_type == "vertex_ai_search" %}
-```bash
-PROJECT_ID="YOUR_PROJECT_ID"
-REGION="us-central1"
-DATA_STORE_REGION="us"
-uv run data_ingestion_pipeline/submit_pipeline.py \
-    --project-id=$PROJECT_ID \
-    --region=$REGION \
-    --data-store-region=$DATA_STORE_REGION \
-    --data-store-id="sample-datastore" \
-    --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"
-```
+*   It will use parameters like `--data-store-id`, `--data-store-region`.
 {%- 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"
-```
+*   It will use parameters like `--vector-search-index`, `--vector-search-index-endpoint`, `--vector-search-data-bucket-name`.
 {%- endif %}
+*   Common parameters include `--project-id`, `--region`, `--service-account`, `--pipeline-root`, and `--pipeline-name`.
 
-**d. Pipeline Scheduling and Execution:**
+**b. Pipeline Scheduling:**
 
-The pipeline, by default, executes immediately. To schedule the pipeline for periodic execution without immediate initiation, use the `--schedule-only` flag in conjunction with `--cron-schedule`. If a schedule doesn't exist, it will be created. If a schedule already exists, its cron expression will be updated to the provided value.
+The `make data-ingestion` command triggers an immediate pipeline run. For production environments, the underlying `submit_pipeline.py` script also supports scheduling options with flags like `--schedule-only` and `--cron-schedule` for periodic execution.
 
-**e. Monitoring Pipeline Progress:**
+**c. Monitoring Pipeline Progress:**
 
-The pipeline's configuration and execution status will be printed to the console. For detailed monitoring, use the Vertex AI Pipelines dashboard in the Google Cloud Console. This dashboard provides real-time insights into the pipeline's progress, logs, and any potential issues.
+The pipeline's configuration and execution status link will be printed to the console upon submission. For detailed monitoring, use the Vertex AI Pipelines dashboard in the Google Cloud Console.
 
 ## Testing Your RAG Application
 
-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.
+Once the data ingestion pipeline completes successfully, you can test your RAG application with {{ datastore_service_name }}.
 {%- 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/deployment_targets/agent_engine/app/agent_engine_app.py

@@ -26,16 +26,12 @@ import vertexai
 from google.cloud import logging as google_cloud_logging
 from langchain_core.runnables import RunnableConfig
 from traceloop.sdk import Instruments, Traceloop
-from vertexai.preview import reasoning_engines
+from vertexai import agent_engines
 
 from app.utils.gcs import create_bucket_if_not_exists
 from app.utils.tracing import CloudTraceLoggingSpanExporter
 from app.utils.typing import Feedback, InputChat, dumpd, ensure_valid_config
 
-logging.basicConfig(
-    level=logging.INFO,
-)
-
 
 class AgentEngineApp:
     """Class for managing agent engine functionality."""
@@ -124,6 +120,8 @@ class AgentEngineApp:
         **kwargs: Any,
     ) -> Any:
         """Process a single input and return the agent's response."""
+        config = ensure_valid_config(config)
+        self.set_tracing_properties(config=config)
         return dumpd(self.runnable.invoke(input=input, config=config, **kwargs))
 
     def register_operations(self) -> Mapping[str, Sequence]:
@@ -151,7 +149,7 @@ def deploy_agent_engine_app(
     requirements_file: str = ".requirements.txt",
     extra_packages: list[str] = ["./app"],
     env_vars: dict[str, str] | None = None,
-) -> reasoning_engines.ReasoningEngine:
+) -> agent_engines.AgentEngine:
     """Deploy the agent engine app to Vertex AI."""
 
     staging_bucket = f"gs://{project}-agent-engine"
@@ -169,7 +167,7 @@ def deploy_agent_engine_app(
 
     # Common configuration for both create and update operations
     agent_config = {
-        "reasoning_engine": agent,
+        "agent_engine": agent,
         "display_name": agent_name,
         "description": "This is a sample custom application in Agent Engine that uses LangGraph",
         "extra_packages": extra_packages,
@@ -178,10 +176,7 @@ def deploy_agent_engine_app(
     agent_config["requirements"] = requirements
 
     # Check if an agent with this name already exists
-    existing_agents = reasoning_engines.ReasoningEngine.list(
-        filter=f"display_name={agent_name}"
-    )
-
+    existing_agents = list(agent_engines.list(filter=f"display_name={agent_name}"))
     if existing_agents:
         # Update the existing agent with new configuration
         logging.info(f"Updating existing agent: {agent_name}")
@@ -189,7 +184,7 @@ def deploy_agent_engine_app(
     else:
         # Create a new agent if none exists
         logging.info(f"Creating new agent: {agent_name}")
-        remote_agent = reasoning_engines.ReasoningEngine.create(**agent_config)
+        remote_agent = agent_engines.create(**agent_config)
 
     config = {
         "remote_agent_engine_id": remote_agent.resource_name,

src/deployment_targets/cloud_run/Dockerfile

@@ -14,7 +14,7 @@
 
 FROM python:3.11-slim
 
-RUN pip install --no-cache-dir uv
+RUN pip install --no-cache-dir uv==0.6.12
 
 WORKDIR /code
 

src/frontends/streamlit/frontend/utils/stream_handler.py

@@ -28,7 +28,7 @@ import streamlit as st
 import vertexai
 from google.auth.exceptions import DefaultCredentialsError
 from langchain_core.messages import AIMessage, ToolMessage
-from vertexai.preview import reasoning_engines
+from vertexai import agent_engines
 
 from frontend.utils.multimodal_utils import format_content
 
@@ -43,7 +43,7 @@ def get_remote_agent(remote_agent_engine_id: str) -> Any:
     project_id = parts[1]
     location = parts[3]
     vertexai.init(project=project_id, location=location)
-    return reasoning_engines.ReasoningEngine(remote_agent_engine_id)
+    return agent_engines.AgentEngine(remote_agent_engine_id)
 
 
 @st.cache_resource
@@ -159,7 +159,7 @@ class Client:
             if self.authenticate_request:
                 headers["Authorization"] = f"Bearer {self.id_token}"
             with requests.post(
-                self.url, json=data, headers=headers, stream=True, timeout=10
+                self.url, json=data, headers=headers, stream=True, timeout=60
             ) as response:
                 for line in response.iter_lines():
                     if line: