|
| 1 | +# Copyright 2025 Google LLC |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +"""Deployment script for LLM Auditor.""" |
| 16 | + |
| 17 | +import os |
| 18 | + |
| 19 | +from absl import app |
| 20 | +from absl import flags |
| 21 | +from dotenv import load_dotenv |
| 22 | +from llm_auditor.agent import root_agent |
| 23 | +import vertexai |
| 24 | +from vertexai import agent_engines |
| 25 | + |
| 26 | +# A2A wrapping |
| 27 | +from a2a.types import AgentSkill |
| 28 | +from google.adk.a2a.executor.a2a_agent_executor import A2aAgentExecutor |
| 29 | +from google.adk.runners import InMemoryRunner |
| 30 | +from vertexai.preview.reasoning_engines.templates.a2a import create_agent_card |
| 31 | +from vertexai.preview.reasoning_engines import A2aAgent |
| 32 | + |
| 33 | +FLAGS = flags.FLAGS |
| 34 | +flags.DEFINE_string("project_id", None, "GCP project ID.") |
| 35 | +flags.DEFINE_string("location", None, "GCP location.") |
| 36 | +flags.DEFINE_string("bucket", None, "GCP bucket.") |
| 37 | +flags.DEFINE_string("resource_id", None, "ReasoningEngine resource ID.") |
| 38 | + |
| 39 | +flags.DEFINE_bool("list", False, "List all agents.") |
| 40 | +flags.DEFINE_bool("create", False, "Creates a new agent.") |
| 41 | +flags.DEFINE_bool("delete", False, "Deletes an existing agent.") |
| 42 | +flags.mark_bool_flags_as_mutual_exclusive(["create", "delete"]) |
| 43 | + |
| 44 | + |
| 45 | +def create() -> None: |
| 46 | + """Creates an agent engine for LLM Auditor.""" |
| 47 | + agent_card = create_agent_card( |
| 48 | + agent_name=root_agent.name, |
| 49 | + description=root_agent.description, |
| 50 | + skills=[AgentSkill( |
| 51 | + id='audit_llm_output', |
| 52 | + name='Audit LLM Output', |
| 53 | + description='Critiques and revises outputs from large language models.', |
| 54 | + tags=['LLM', 'Audit', 'Revision'], |
| 55 | + examples=[ |
| 56 | + 'The earth is flat.', |
| 57 | + 'The capital of France is Berlin.', |
| 58 | + 'The last winner of the Super Bowl was the New England Patriots in 2020.', |
| 59 | + ], |
| 60 | + )] |
| 61 | + ) |
| 62 | + a2a_agent = A2aAgent( |
| 63 | + agent_card=agent_card, |
| 64 | + agent_executor_builder=lambda: A2aAgentExecutor( |
| 65 | + runner=InMemoryRunner( |
| 66 | + app_name=root_agent.name, |
| 67 | + agent=root_agent, |
| 68 | + ) |
| 69 | + ) |
| 70 | + ) |
| 71 | + a2a_agent.set_up() |
| 72 | + |
| 73 | + remote_agent = agent_engines.create( |
| 74 | + a2a_agent, |
| 75 | + display_name=root_agent.name, |
| 76 | + requirements=[ |
| 77 | + "google-adk (>=0.0.2)", |
| 78 | + "google-cloud-aiplatform[agent_engines] (>=1.88.0,<2.0.0)", |
| 79 | + "google-genai (>=1.5.0,<2.0.0)", |
| 80 | + "pydantic (>=2.10.6,<3.0.0)", |
| 81 | + "absl-py (>=2.2.1,<3.0.0)", |
| 82 | + "a2a-sdk>=0.3.22", |
| 83 | + "uvicorn", |
| 84 | + ], |
| 85 | + # In-memory runner |
| 86 | + max_instances=1, |
| 87 | + env_vars ={ |
| 88 | + "NUM_WORKERS": "1" |
| 89 | + }, |
| 90 | + extra_packages=["./llm_auditor"], |
| 91 | + ) |
| 92 | + print(f"Created remote agent: {remote_agent.resource_name}") |
| 93 | + |
| 94 | + |
| 95 | +def delete(resource_id: str) -> None: |
| 96 | + remote_agent = agent_engines.get(resource_id) |
| 97 | + remote_agent.delete(force=True) |
| 98 | + print(f"Deleted remote agent: {resource_id}") |
| 99 | + |
| 100 | + |
| 101 | +def list_agents() -> None: |
| 102 | + remote_agents = agent_engines.list() |
| 103 | + TEMPLATE = ''' |
| 104 | +{agent.name} ("{agent.display_name}") |
| 105 | +- Create time: {agent.create_time} |
| 106 | +- Update time: {agent.update_time} |
| 107 | +''' |
| 108 | + remote_agents_string = '\n'.join(TEMPLATE.format(agent=agent) for agent in remote_agents) |
| 109 | + print(f"All remote agents:\n{remote_agents_string}") |
| 110 | + |
| 111 | +def main(argv: list[str]) -> None: |
| 112 | + del argv # unused |
| 113 | + load_dotenv() |
| 114 | + |
| 115 | + project_id = ( |
| 116 | + FLAGS.project_id |
| 117 | + if FLAGS.project_id |
| 118 | + else os.getenv("GOOGLE_CLOUD_PROJECT") |
| 119 | + ) |
| 120 | + location = ( |
| 121 | + FLAGS.location if FLAGS.location else os.getenv("GOOGLE_CLOUD_LOCATION") |
| 122 | + ) |
| 123 | + bucket = ( |
| 124 | + FLAGS.bucket if FLAGS.bucket |
| 125 | + else os.getenv("GOOGLE_CLOUD_STORAGE_BUCKET") |
| 126 | + ) |
| 127 | + |
| 128 | + print(f"PROJECT: {project_id}") |
| 129 | + print(f"LOCATION: {location}") |
| 130 | + print(f"BUCKET: {bucket}") |
| 131 | + |
| 132 | + if not project_id: |
| 133 | + print("Missing required environment variable: GOOGLE_CLOUD_PROJECT") |
| 134 | + return |
| 135 | + elif not location: |
| 136 | + print("Missing required environment variable: GOOGLE_CLOUD_LOCATION") |
| 137 | + return |
| 138 | + elif not bucket: |
| 139 | + print( |
| 140 | + "Missing required environment variable: GOOGLE_CLOUD_STORAGE_BUCKET" |
| 141 | + ) |
| 142 | + return |
| 143 | + |
| 144 | + vertexai.init( |
| 145 | + project=project_id, |
| 146 | + location=location, |
| 147 | + staging_bucket=f"gs://{bucket}", |
| 148 | + ) |
| 149 | + |
| 150 | + if FLAGS.list: |
| 151 | + list_agents() |
| 152 | + elif FLAGS.create: |
| 153 | + create() |
| 154 | + elif FLAGS.delete: |
| 155 | + if not FLAGS.resource_id: |
| 156 | + print("resource_id is required for delete") |
| 157 | + return |
| 158 | + delete(FLAGS.resource_id) |
| 159 | + else: |
| 160 | + print("Unknown command") |
| 161 | + |
| 162 | + |
| 163 | +if __name__ == "__main__": |
| 164 | + app.run(main) |
0 commit comments