agent-starter-pack 0.3.3__py3-none-any.whl → 0.21.0__py3-none-any.whl

agent_starter_pack/deployment_targets/agent_engine/{{cookiecutter.agent_directory}}/agent_engine_app.py

@@ -0,0 +1,263 @@
+# 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.
+
+# mypy: disable-error-code="attr-defined,arg-type"
+{%- if cookiecutter.is_adk %}
+{%- if cookiecutter.is_a2a %}
+import asyncio
+{%- endif %}
+import logging
+import os
+from typing import Any
+
+{% if cookiecutter.is_a2a -%}
+import nest_asyncio
+{% endif -%}
+import vertexai
+{%- if cookiecutter.is_a2a %}
+from a2a.types import AgentCapabilities, AgentCard, TransportProtocol
+from google.adk.a2a.executor.a2a_agent_executor import A2aAgentExecutor
+from google.adk.a2a.utils.agent_card_builder import AgentCardBuilder
+from google.adk.apps.app import App
+{%- endif %}
+from google.adk.artifacts import GcsArtifactService, InMemoryArtifactService
+{%- if cookiecutter.is_a2a %}
+from google.adk.runners import Runner
+from google.adk.sessions import InMemorySessionService
+{%- endif %}
+from google.cloud import logging as google_cloud_logging
+{%- if cookiecutter.is_adk_live %}
+from vertexai.agent_engines.templates.adk import AdkApp
+from vertexai.preview.reasoning_engines import AdkApp as PreviewAdkApp
+{%- elif cookiecutter.is_a2a %}
+from vertexai.preview.reasoning_engines import A2aAgent
+{%- else %}
+from vertexai.agent_engines.templates.adk import AdkApp
+{%- endif %}
+{%- if cookiecutter.is_adk or cookiecutter.is_adk_live %}
+
+from {{cookiecutter.agent_directory}}.agent import app as adk_app
+{%- else %}
+
+{%- endif %}
+from {{cookiecutter.agent_directory}}.app_utils.telemetry import setup_telemetry
+from {{cookiecutter.agent_directory}}.app_utils.typing import Feedback
+{%- if cookiecutter.is_a2a %}
+
+
+class AgentEngineApp(A2aAgent):
+    @staticmethod
+    def create(
+        app: App | None = None,
+        artifact_service: Any = None,
+        session_service: Any = None,
+    ) -> Any:
+        """Create an AgentEngineApp instance.
+
+        This method detects whether it's being called in an async context (like notebooks
+        or Agent Engine) and handles agent card creation appropriately.
+        """
+        if app is None:
+            app = adk_app
+
+        def create_runner() -> Runner:
+            """Create a Runner for the AgentEngineApp."""
+            return Runner(
+                app=app,
+                session_service=session_service,
+                artifact_service=artifact_service,
+            )
+
+        # Build agent card in an async context if needed
+        try:
+            asyncio.get_running_loop()
+            # Running event loop detected - enable nested asyncio.run()
+            nest_asyncio.apply()
+        except RuntimeError:
+            pass
+
+        agent_card = asyncio.run(AgentEngineApp.build_agent_card(app=app))
+
+        return AgentEngineApp(
+            agent_executor_builder=lambda: A2aAgentExecutor(runner=create_runner()),
+            agent_card=agent_card,
+        )
+
+    @staticmethod
+    async def build_agent_card(app: App) -> AgentCard:
+        """Builds the Agent Card dynamically from the app."""
+        agent_card_builder = AgentCardBuilder(
+            agent=app.root_agent,
+            # Agent Engine does not support streaming yet
+            capabilities=AgentCapabilities(streaming=False),
+            rpc_url="http://localhost:9999/",
+            agent_version=os.getenv("AGENT_VERSION", "0.1.0"),
+        )
+        agent_card = await agent_card_builder.build()
+        agent_card.preferred_transport = TransportProtocol.http_json  # Http Only.
+        agent_card.supports_authenticated_extended_card = True
+        return agent_card
+{% else %}
+
+
+class AgentEngineApp(AdkApp):
+{%- endif %}
+    def set_up(self) -> None:
+        """Initialize the agent engine app with logging and telemetry."""
+        vertexai.init()
+        setup_telemetry()
+        super().set_up()
+        logging.basicConfig(level=logging.INFO)
+        logging_client = google_cloud_logging.Client()
+        self.logger = logging_client.logger(__name__)
+        if gemini_location:
+            os.environ["GOOGLE_CLOUD_LOCATION"] = gemini_location
+
+    def register_feedback(self, feedback: dict[str, Any]) -> None:
+        """Collect and log feedback."""
+        feedback_obj = Feedback.model_validate(feedback)
+        self.logger.log_struct(feedback_obj.model_dump(), severity="INFO")
+
+    def register_operations(self) -> dict[str, list[str]]:
+        """Registers the operations of the Agent."""
+        operations = super().register_operations()
+        operations[""] = operations.get("", []) + ["register_feedback"]
+{%- if cookiecutter.is_adk_live %}
+        # Add bidi_stream_query for adk_live
+        operations["bidi_stream"] = ["bidi_stream_query"]
+{%- endif %}
+        return operations
+{%- if cookiecutter.is_a2a %}
+
+    def clone(self) -> "AgentEngineApp":
+        """Returns a clone of the Agent Engine application."""
+        return self
+{%- endif %}
+{%- if cookiecutter.is_adk_live %}
+
+
+# Add bidi_stream_query support from preview AdkApp for adk_live
+AgentEngineApp.bidi_stream_query = PreviewAdkApp.bidi_stream_query
+{%- endif %}
+
+
+gemini_location = os.environ.get("GOOGLE_CLOUD_LOCATION")
+logs_bucket_name = os.environ.get("LOGS_BUCKET_NAME")
+{%- if cookiecutter.is_a2a %}
+agent_engine = AgentEngineApp.create(
+    app=adk_app,
+    artifact_service=(
+        GcsArtifactService(bucket_name=logs_bucket_name)
+        if logs_bucket_name
+        else InMemoryArtifactService()
+    ),
+    session_service=InMemorySessionService(),
+)
+{%- else %}
+agent_engine = AgentEngineApp(
+    app=adk_app,
+    artifact_service_builder=lambda: GcsArtifactService(bucket_name=logs_bucket_name)
+    if logs_bucket_name
+    else InMemoryArtifactService(),
+)
+{%- endif -%}
+{% else %}
+
+import asyncio
+import os
+from typing import Any
+
+import nest_asyncio
+from a2a.types import AgentCapabilities, AgentCard, AgentSkill, TransportProtocol
+from google.cloud import logging as google_cloud_logging
+from vertexai.preview.reasoning_engines import A2aAgent
+
+from {{cookiecutter.agent_directory}}.agent import root_agent
+from {{cookiecutter.agent_directory}}.app_utils.executor.a2a_agent_executor import (
+    LangGraphAgentExecutor,
+)
+from {{cookiecutter.agent_directory}}.app_utils.typing import Feedback
+
+
+class AgentEngineApp(A2aAgent):
+    """Agent Engine App with A2A Protocol support for LangGraph agents."""
+
+    @staticmethod
+    def create() -> "AgentEngineApp":
+        """Create an AgentEngineApp instance with A2A support.
+
+        This method handles agent card creation in async context.
+        """
+        # Handle nested asyncio contexts (like notebooks or Agent Engine)
+        try:
+            asyncio.get_running_loop()
+            nest_asyncio.apply()
+        except RuntimeError:
+            pass
+
+        agent_card = asyncio.run(AgentEngineApp.build_agent_card())
+
+        return AgentEngineApp(
+            agent_executor_builder=lambda: LangGraphAgentExecutor(graph=root_agent),
+            agent_card=agent_card,
+        )
+
+    @staticmethod
+    async def build_agent_card() -> AgentCard:
+        """Build the Agent Card for the LangGraph agent."""
+        skill = AgentSkill(
+            id="root_agent-get_weather",
+            name="get_weather",
+            description="Simulates a web search. Use it get information on weather.",
+            tags=["llm", "tools"],
+            examples=["What's the weather in San Francisco?"],
+        )
+        agent_card = AgentCard(
+            name="root_agent",
+            description="A base ReAct agent using LangGraph with Agent2Agent (A2A) Protocol support",
+            url="http://localhost:9999/",  # RPC URL for Agent Engine
+            version=os.getenv("AGENT_VERSION", "0.1.0"),
+            default_input_modes=["text/plain"],
+            default_output_modes=["text/plain"],
+            capabilities=AgentCapabilities(
+                streaming=False
+            ),  # Agent Engine does not support streaming yet
+            skills=[skill],
+        )
+
+        agent_card.preferred_transport = TransportProtocol.http_json  # Http Only.
+        agent_card.supports_authenticated_extended_card = True
+
+        return agent_card
+
+    def register_feedback(self, feedback: dict[str, Any]) -> None:
+        """Collect and log feedback."""
+        feedback_obj = Feedback.model_validate(feedback)
+        logging_client = google_cloud_logging.Client()
+        logger = logging_client.logger(__name__)
+        logger.log_struct(feedback_obj.model_dump(), severity="INFO")
+
+    def register_operations(self) -> dict[str, list[str]]:
+        """Registers the operations of the Agent."""
+        operations = super().register_operations()
+        operations[""] = operations.get("", []) + ["register_feedback"]
+        return operations
+
+    def clone(self) -> "AgentEngineApp":
+        """Returns a clone of the Agent Engine application."""
+        return self
+
+
+agent_engine = AgentEngineApp.create()
+{%- endif %}

agent_starter_pack/deployment_targets/agent_engine/{{cookiecutter.agent_directory}}/app_utils/deploy.py

@@ -0,0 +1,414 @@
+# 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.
+
+import asyncio
+import datetime
+import importlib
+import inspect
+import json
+import logging
+import warnings
+from typing import Any
+
+import click
+import google.auth
+import vertexai
+from vertexai._genai import _agent_engines_utils
+from vertexai._genai.types import AgentEngine, AgentEngineConfig{%- if cookiecutter.is_adk_live %}, AgentServerMode{%- endif %}
+{%- if cookiecutter.is_adk_live %}
+
+from {{cookiecutter.agent_directory}}.app_utils.gcs import create_bucket_if_not_exists
+{%- endif %}
+
+# Suppress google-cloud-storage version compatibility warning
+warnings.filterwarnings(
+    "ignore", category=FutureWarning, module="google.cloud.aiplatform"
+)
+
+
+def generate_class_methods_from_agent(agent_instance: Any) -> list[dict[str, Any]]:
+    """Generate method specifications with schemas from agent's register_operations().
+
+    See: https://docs.cloud.google.com/agent-builder/agent-engine/use/custom#supported-operations
+    """
+    registered_operations = _agent_engines_utils._get_registered_operations(
+        agent=agent_instance
+    )
+    class_methods_spec = _agent_engines_utils._generate_class_methods_spec_or_raise(
+        agent=agent_instance,
+        operations=registered_operations,
+    )
+    class_methods_list = [
+        _agent_engines_utils._to_dict(method_spec) for method_spec in class_methods_spec
+    ]
+    return class_methods_list
+
+
+def parse_key_value_pairs(kv_string: str | None) -> dict[str, str]:
+    """Parse key-value pairs from a comma-separated KEY=VALUE string."""
+    result = {}
+    if kv_string:
+        for pair in kv_string.split(","):
+            if "=" in pair:
+                key, value = pair.split("=", 1)
+                result[key.strip()] = value.strip()
+            else:
+                logging.warning(f"Skipping malformed key-value pair: {pair}")
+    return result
+
+
+def write_deployment_metadata(
+    remote_agent: Any,
+    metadata_file: str = "deployment_metadata.json",
+) -> None:
+    """Write deployment metadata to file."""
+    metadata = {
+        "remote_agent_engine_id": remote_agent.api_resource.name,
+        "deployment_target": "agent_engine",
+{%- if cookiecutter.is_a2a %}
+        "is_a2a": True,
+{%- else %}
+        "is_a2a": False,
+{%- endif %}
+        "deployment_timestamp": datetime.datetime.now().isoformat(),
+    }
+
+    with open(metadata_file, "w") as f:
+        json.dump(metadata, f, indent=2)
+
+    logging.info(f"Agent Engine ID written to {metadata_file}")
+
+
+def print_deployment_success(
+    remote_agent: Any,
+    location: str,
+    project: str,
+) -> None:
+    """Print deployment success message with console URL."""
+    # Extract agent engine ID and project number for console URL
+    resource_name_parts = remote_agent.api_resource.name.split("/")
+    agent_engine_id = resource_name_parts[-1]
+    project_number = resource_name_parts[1]
+
+{%- if cookiecutter.is_adk_live %}
+    print("\n✅ Deployment successful! Run your agent with: `make playground-remote`")
+{%- elif cookiecutter.is_a2a %}
+    print(
+        "\n✅ Deployment successful! Test your agent: notebooks/adk_a2a_app_testing.ipynb"
+    )
+{%- else %}
+    print("\n✅ Deployment successful!")
+{%- endif %}
+    service_account = remote_agent.api_resource.spec.service_account
+    if service_account:
+        print(f"Service Account: {service_account}")
+    else:
+        default_sa = (
+            f"service-{project_number}@gcp-sa-aiplatform-re.iam.gserviceaccount.com"
+        )
+        print(f"Service Account: {default_sa}")
+{%- if cookiecutter.is_adk and not cookiecutter.is_adk_live and not cookiecutter.is_a2a %}
+    playground_url = f"https://console.cloud.google.com/vertex-ai/agents/locations/{location}/agent-engines/{agent_engine_id}/playground?project={project}"
+    print(f"\n📊 Open Console Playground: {playground_url}\n")
+{%- else %}
+    console_url = f"https://console.cloud.google.com/vertex-ai/agents/locations/{location}/agent-engines/{agent_engine_id}?project={project}"
+    print(f"\n📊 View in Console: {console_url}\n")
+{%- endif %}
+
+
+@click.command()
+@click.option(
+    "--project",
+    default=None,
+    help="GCP project ID (defaults to application default credentials)",
+)
+@click.option(
+    "--location",
+    default="us-central1",
+    help="GCP region (defaults to us-central1)",
+)
+@click.option(
+    "--display-name",
+    default="{{cookiecutter.project_name}}",
+    help="Display name for the agent engine",
+)
+@click.option(
+    "--description",
+    default="{{cookiecutter.agent_description}}",
+    help="Description of the agent",
+)
+@click.option(
+    "--source-packages",
+    multiple=True,
+    default=["./{{cookiecutter.agent_directory}}"],
+    help="Source packages to deploy. Can be specified multiple times (e.g., --source-packages=./app --source-packages=./lib)",
+)
+@click.option(
+    "--entrypoint-module",
+    default="{{cookiecutter.agent_directory}}.agent_engine_app",
+    help="Python module path for the agent entrypoint (required)",
+)
+@click.option(
+    "--entrypoint-object",
+    default="agent_engine",
+    help="Name of the agent instance at module level (required)",
+)
+@click.option(
+    "--requirements-file",
+    default="{{cookiecutter.agent_directory}}/app_utils/.requirements.txt",
+    help="Path to requirements.txt file",
+)
+@click.option(
+    "--set-env-vars",
+    default=None,
+    help="Comma-separated list of environment variables in KEY=VALUE format",
+)
+@click.option(
+    "--labels",
+    default=None,
+    help="Comma-separated list of labels in KEY=VALUE format",
+)
+@click.option(
+    "--service-account",
+    default=None,
+    help="Service account email to use for the agent engine",
+)
+@click.option(
+    "--min-instances",
+    type=int,
+    default=1,
+    help="Minimum number of instances (default: 1)",
+)
+@click.option(
+    "--max-instances",
+    type=int,
+    default=10,
+    help="Maximum number of instances (default: 10)",
+)
+@click.option(
+    "--cpu",
+    default="4",
+    help="CPU limit (default: 4)",
+)
+@click.option(
+    "--memory",
+    default="8Gi",
+    help="Memory limit (default: 8Gi)",
+)
+@click.option(
+    "--container-concurrency",
+    type=int,
+    default=9,
+    help="Container concurrency (default: 9)",
+)
+@click.option(
+    "--num-workers",
+    type=int,
+    default=1,
+    help="Number of worker processes (default: 1)",
+)
+def deploy_agent_engine_app(
+    project: str | None,
+    location: str,
+    display_name: str,
+    description: str,
+    source_packages: tuple[str, ...],
+    entrypoint_module: str,
+    entrypoint_object: str,
+    requirements_file: str,
+    set_env_vars: str | None,
+    labels: str | None,
+    service_account: str | None,
+    min_instances: int,
+    max_instances: int,
+    cpu: str,
+    memory: str,
+    container_concurrency: int,
+    num_workers: int,
+) -> AgentEngine:
+    """Deploy the agent engine app to Vertex AI."""
+
+    logging.basicConfig(level=logging.INFO)
+    logging.getLogger("httpx").setLevel(logging.WARNING)
+
+    # Parse environment variables and labels if provided
+    env_vars = parse_key_value_pairs(set_env_vars)
+    labels_dict = parse_key_value_pairs(labels)
+
+    # Set GOOGLE_CLOUD_REGION to match deployment location
+    env_vars["GOOGLE_CLOUD_REGION"] = location
+
+    # Add NUM_WORKERS from CLI argument (can be overridden via --set-env-vars)
+    if "NUM_WORKERS" not in env_vars:
+        env_vars["NUM_WORKERS"] = str(num_workers)
+
+    # Enable telemetry by default for Agent Engine
+    if "GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY" not in env_vars:
+        env_vars["GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY"] = "true"
+    if "OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT" not in env_vars:
+        env_vars["OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT"] = "true"
+
+    if not project:
+        _, project = google.auth.default()
+
+    print("""
+    ╔═══════════════════════════════════════════════════════════╗
+    ║                                                           ║
+    ║   🤖 DEPLOYING AGENT TO VERTEX AI AGENT ENGINE 🤖         ║
+    ║                                                           ║
+    ╚═══════════════════════════════════════════════════════════╝
+    """)
+
+    # Log deployment parameters
+    click.echo("\n📋 Deployment Parameters:")
+    click.echo(f"  Project: {project}")
+    click.echo(f"  Location: {location}")
+    click.echo(f"  Display Name: {display_name}")
+    click.echo(f"  Min Instances: {min_instances}")
+    click.echo(f"  Max Instances: {max_instances}")
+    click.echo(f"  CPU: {cpu}")
+    click.echo(f"  Memory: {memory}")
+    click.echo(f"  Container Concurrency: {container_concurrency}")
+    if service_account:
+        click.echo(f"  Service Account: {service_account}")
+    if env_vars:
+        click.echo("\n🌍 Environment Variables:")
+        for key, value in sorted(env_vars.items()):
+            click.echo(f"  {key}: {value}")
+
+    source_packages_list = list(source_packages)
+
+    # Initialize vertexai client
+    client = vertexai.Client(
+        project=project,
+        location=location,
+    )
+    vertexai.init(project=project, location=location)
+
+    # Add agent garden labels if configured
+{%- if cookiecutter.agent_garden %}
+{%- if cookiecutter.agent_sample_id and cookiecutter.agent_sample_publisher %}
+    labels_dict["vertex-agent-sample-id"] = "{{cookiecutter.agent_sample_id}}"
+    labels_dict["vertex-agent-sample-publisher"] = "{{cookiecutter.agent_sample_publisher}}"
+    labels_dict["deployed-with"] = "agent-garden"
+{%- endif %}
+{%- endif %}
+
+    # Dynamically import the agent instance to generate class_methods
+    logging.info(f"Importing {entrypoint_module}.{entrypoint_object}")
+    module = importlib.import_module(entrypoint_module)
+    agent_instance = getattr(module, entrypoint_object)
+
+    # If the agent_instance is a coroutine, await it to get the actual instance
+    if inspect.iscoroutine(agent_instance):
+        logging.info(f"Detected coroutine, awaiting {entrypoint_object}...")
+        agent_instance = asyncio.run(agent_instance)
+
+{%- if cookiecutter.is_adk_live %}
+    # For adk_live, use pickle-based deployment (source-based not yet supported with EXPERIMENTAL mode)
+    # Ensure staging bucket exists for pickle serialization
+    staging_bucket_uri = f"gs://{project}-agent-engine"
+
+    create_bucket_if_not_exists(
+        bucket_name=staging_bucket_uri, project=project, location=location
+    )
+
+    config = AgentEngineConfig(
+        display_name=display_name,
+        description=description,
+        env_vars=env_vars,
+        extra_packages=source_packages_list,
+        service_account=service_account,
+        requirements=requirements_file,
+        staging_bucket=staging_bucket_uri,
+        labels=labels_dict,
+        gcs_dir_name=display_name,
+        agent_server_mode=AgentServerMode.EXPERIMENTAL,  # Enable bidi streaming
+        resource_limits={"cpu": cpu, "memory": memory},
+    )
+
+    agent_config = {
+        "agent": agent_instance,
+        "config": config,
+    }
+{%- else %}
+    # Generate class methods spec from register_operations
+    class_methods_list = generate_class_methods_from_agent(agent_instance)
+
+    config = AgentEngineConfig(
+        display_name=display_name,
+        description=description,
+        source_packages=source_packages_list,
+        entrypoint_module=entrypoint_module,
+        entrypoint_object=entrypoint_object,
+        class_methods=class_methods_list,
+        env_vars=env_vars,
+        service_account=service_account,
+        requirements_file=requirements_file,
+        labels=labels_dict,
+        min_instances=min_instances,
+        max_instances=max_instances,
+        resource_limits={"cpu": cpu, "memory": memory},
+        container_concurrency=container_concurrency,
+{%- if cookiecutter.is_adk %}
+        agent_framework="google-adk",
+{%- endif %}
+    )
+{%- endif %}
+
+    # Check if an agent with this name already exists
+    existing_agents = list(client.agent_engines.list())
+    matching_agents = [
+        agent
+        for agent in existing_agents
+        if agent.api_resource.display_name == display_name
+    ]
+
+    # Deploy the agent (create or update)
+    if matching_agents:
+        click.echo(f"\n📝 Updating existing agent: {display_name}")
+    else:
+        click.echo(f"\n🚀 Creating new agent: {display_name}")
+
+    click.echo("🚀 Deploying to Vertex AI Agent Engine (this can take 3-5 minutes)...")
+
+{%- if cookiecutter.is_adk_live %}
+    if matching_agents:
+        remote_agent = client.agent_engines.update(
+            name=matching_agents[0].api_resource.name,
+            agent=agent_instance,
+            config=config,
+        )
+    else:
+        remote_agent = client.agent_engines.create(
+            agent=agent_instance,
+            config=config,
+        )
+{%- else %}
+    if matching_agents:
+        remote_agent = client.agent_engines.update(
+            name=matching_agents[0].api_resource.name, config=config
+        )
+    else:
+        remote_agent = client.agent_engines.create(config=config)
+{%- endif %}
+
+    write_deployment_metadata(remote_agent)
+    print_deployment_success(remote_agent, location, project)
+
+    return remote_agent
+
+
+if __name__ == "__main__":
+    deploy_agent_engine_app()