agent-starter-pack 0.1.6__py3-none-any.whl → 0.2.0__py3-none-any.whl

src/data_ingestion/data_ingestion_pipeline/submit_pipeline.py

@@ -40,6 +40,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--region", default=os.getenv("REGION"), help="Vertex AI Pipelines region"
     )
+{%- if cookiecutter.datastore_type == "vertex_ai_search" %}
     parser.add_argument(
         "--data-store-region",
         default=os.getenv("DATA_STORE_REGION"),
@@ -48,6 +49,23 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--data-store-id", default=os.getenv("DATA_STORE_ID"), help="Data store ID"
     )
+{%- elif cookiecutter.datastore_type == "vertex_ai_vector_search" %}
+    parser.add_argument(
+        "--vector-search-index",
+        default=os.getenv("VECTOR_SEARCH_INDEX"),
+        help="Vector Search Index",
+    )
+    parser.add_argument(
+        "--vector-search-index-endpoint",
+        default=os.getenv("VECTOR_SEARCH_INDEX_ENDPOINT"),
+        help="Vector Search Index Endpoint",
+    )
+    parser.add_argument(
+        "--vector-search-data-bucket-name",
+        default=os.getenv("VECTOR_SEARCH_BUCKET"),
+        help="Vector Search Data Bucket Name",
+    )
+{%- endif %}
     parser.add_argument(
         "--service-account",
         default=os.getenv("SERVICE_ACCOUNT"),
@@ -85,12 +103,22 @@ def parse_args() -> argparse.Namespace:
     required_params = {
         "project_id": parsed_args.project_id,
         "region": parsed_args.region,
-        "data_store_region": parsed_args.data_store_region,
-        "data_store_id": parsed_args.data_store_id,
         "service_account": parsed_args.service_account,
         "pipeline_root": parsed_args.pipeline_root,
         "pipeline_name": parsed_args.pipeline_name,
     }
+{%- if cookiecutter.datastore_type == "vertex_ai_search" %}
+    required_params["data_store_region"] = parsed_args.data_store_region
+    required_params["data_store_id"] = parsed_args.data_store_id
+{%- elif cookiecutter.datastore_type == "vertex_ai_vector_search" %}
+    required_params["vector_search_index"] = parsed_args.vector_search_index
+    required_params["vector_search_index_endpoint"] = (
+        parsed_args.vector_search_index_endpoint
+    )
+    required_params["vector_search_data_bucket_name"] = (
+        parsed_args.vector_search_data_bucket_name
+    )
+{%- endif %}
 
     for param_name, param_value in required_params.items():
         if param_value is None:
@@ -118,16 +146,9 @@ if __name__ == "__main__":
     # Print configuration
     logging.info("\nConfiguration:")
     logging.info("--------------")
-    logging.info(f"project_id: {args.project_id}")
-    logging.info(f"region: {args.region}")
-    logging.info(f"region_data_store: {args.data_store_region}")
-    logging.info(f"data_store_id: {args.data_store_id}")
-    logging.info(f"service_account: {args.service_account}")
-    logging.info(f"pipeline_root: {args.pipeline_root}")
-    logging.info(f"cron_schedule: {args.cron_schedule}")
-    logging.info(f"pipeline_name: {args.pipeline_name}")
-    logging.info(f"disable_caching: {args.disable_caching}")
-    logging.info(f"schedule_only: {args.schedule_only}")
+    # Print all arguments dynamically
+    for arg_name, arg_value in vars(args).items():
+        logging.info(f"{arg_name}: {arg_value}")
     logging.info("--------------\n")
 
     compiler.Compiler().compile(pipeline_func=pipeline, package_path=PIPELINE_FILE_NAME)
@@ -142,10 +163,24 @@ if __name__ == "__main__":
         "parameter_values": {
             "project_id": args.project_id,
             "location": args.region,
-            "data_store_region": args.data_store_region,
-            "data_store_id": args.data_store_id,
         },
     }
+{%- if cookiecutter.datastore_type == "vertex_ai_search" %}
+    pipeline_job_params["parameter_values"]["data_store_region"] = (
+        args.data_store_region
+    )
+    pipeline_job_params["parameter_values"]["data_store_id"] = args.data_store_id
+{%- elif cookiecutter.datastore_type == "vertex_ai_vector_search" %}
+    pipeline_job_params["parameter_values"]["vector_search_index"] = (
+        args.vector_search_index
+    )
+    pipeline_job_params["parameter_values"]["vector_search_index_endpoint"] = (
+        args.vector_search_index_endpoint
+    )
+    pipeline_job_params["parameter_values"]["vector_search_data_bucket_name"] = (
+        args.vector_search_data_bucket_name
+    )
+{%- endif %}
 
     # Create pipeline job instance
     job = aiplatform.PipelineJob(**pipeline_job_params)

src/data_ingestion/pyproject.toml

@@ -6,6 +6,7 @@ readme = "README.md"
 requires-python = ">=3.9, <=3.13"
 dependencies = [
     "google-cloud-aiplatform>=1.80.0",
+    "google-cloud-pipeline-components>=2.19.0",
     "kfp>=1.4.0",
 ]
 

src/deployment_targets/agent_engine/app/agent_engine_app.py

@@ -215,7 +215,9 @@ if __name__ == "__main__":
         help="GCP project ID (defaults to application default credentials)",
     )
     parser.add_argument(
-        "--location", default="us-central1", help="GCP region (defaults to us-central1)"
+        "--location",
+        default="us-central1",
+        help="GCP region (defaults to us-central1)",
     )
     parser.add_argument(
         "--agent-name",

src/deployment_targets/cloud_run/tests/unit/test_server.py

@@ -12,16 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import importlib.util
 import json
 import logging
 import os
+import sys
 from collections.abc import Generator
 from unittest.mock import MagicMock, patch
 
 import pytest
 from fastapi.testclient import TestClient
-from google.auth import exceptions as google_auth_exceptions
 from google.auth.credentials import Credentials
 from langchain_core.messages import HumanMessage
 
@@ -65,36 +64,14 @@ def sample_input_chat() -> InputChat:
     )
 
 
-@pytest.fixture(autouse=True)
-def mock_dependencies() -> Generator[None, None, None]:
-    """
-    Mock Vertex AI dependencies for testing.
-    Patches VertexAIEmbeddings (if defined) and ChatVertexAI.
-    """
-    patches = []
-    try:
-        try:
-            importlib.util.find_spec("app.agent.VertexAIEmbeddings")
-        except (ModuleNotFoundError, google_auth_exceptions.DefaultCredentialsError):
-            pass
-        else:
-            patches.append(patch("app.agent.VertexAIEmbeddings"))
-        patches.append(patch("app.agent.ChatVertexAI"))
-
-        for patch_item in patches:
-            mock = patch_item.start()
-            mock.return_value = MagicMock()
-
-        yield
-    except google_auth_exceptions.GoogleAuthError:
-        yield
-
-
 def test_redirect_root_to_docs() -> None:
     """
     Test that the root endpoint (/) redirects to the Swagger UI documentation.
     """
-    with patch("app.server.agent") as _:
+    # Mock the agent module before importing server
+    mock_agent = MagicMock()
+    with patch.dict(sys.modules, {"app.agent": mock_agent}):
+        # Now import server after the mock is in place
         from app.server import app
 
         client = TestClient(app)
@@ -109,8 +86,6 @@ async def test_stream_chat_events() -> None:
     Test the stream endpoint to ensure it correctly handles
     streaming responses and generates the expected events.
     """
-    from app.server import app
-
     input_data = {
         "input": {
             "messages": [
@@ -124,11 +99,18 @@ async def test_stream_chat_events() -> None:
 
     mock_events = [{"content": "Mocked response"}, {"content": "Additional response"}]
 
-    with patch("app.server.agent") as mock_agent:
-        mock_agent.stream.return_value = mock_events
+    # Create a mock agent module
+    mock_agent_module = MagicMock()
+    mock_agent_module.agent = MagicMock()
+    mock_agent_module.agent.stream.return_value = mock_events
+
+    # Patch the module import
+    with patch.dict(sys.modules, {"app.agent": mock_agent_module}):
+        # Import server after the mock is in place
+        from app.server import app
 
         client = TestClient(app)
-        response = client.post("stream_messages", json=input_data)
+        response = client.post("/stream_messages", json=input_data)
 
         assert response.status_code == 200