Generative AI Google Brain

embeddings-anomaly-sampling-engine

4 min read

🛡️ AI Brand Safety: Three-Tier Agent Anomaly Detection

Source notebook

Repo path: embeddings/anomaly_sampling_engine.ipynb · Open on GitHub · advanced

Builds ADK agent anomaly detection with Gemini baselines, Vector Search scoring, and tiered audits.

Summary

It teaches how to deploy an ADK e-commerce support agent to Agent Engine, synthesize safe and unsafe datasets with Gemini, embed baseline prompts, tool calls, and responses, and provision separate Vector Search indices. It then scores full execution traces against safe neighborhoods, assigns audit tiers with cosine-distance thresholds, logs or reuses BigQuery traces, and visualizes distributions for threshold tuning.

Key code patterns

Vertex AI setup

vertexai.init(project=PROJECT_ID, location=LOCATION)
client = Client()
genai_client = genai.Client(vertexai=True, project=project_id, location=location)
embedding_model = TextEmbeddingModel.from_pretrained("text-embedding-004")

Initializes Vertex AI, Agent Platform client access, Gemini calls, and the embedding model used for anomaly scoring.

ADK agent deployment

app = vertexai.agent_engines.AdkApp(agent=root_agent)
remote_app = deploy_client.agent_engines.create(
    agent=app,
    config={
        "staging_bucket": f"gs://{BUCKET_NAME}",
        "requirements": ["google-cloud-aiplatform[adk,agent_engines]"],
    },
)

Packages the ADK agent and deploys it to Agent Engine with a Cloud Storage staging bucket.

Golden dataset generation

generation_config = GenerateContentConfig(
    temperature=0.9,
    response_mime_type="application/json",
)
response = genai_client.models.generate_content(
    model="gemini-2.5-flash",
    contents=eval_generation_prompt,
    config=generation_config,
)
eval_data_dict = json.loads(response.text)

Uses Gemini to generate structured JSON evaluation data for safe and unsafe prompts.

Three trace indices

prompt_engine = BrandSafetySamplingEngine(PROJECT_ID, LOCATION, BUCKET_NAME)
tool_engine = BrandSafetySamplingEngine(PROJECT_ID, LOCATION, BUCKET_NAME)
response_engine = BrandSafetySamplingEngine(PROJECT_ID, LOCATION, BUCKET_NAME)
future_prompts = executor.submit(prompt_engine.generate_targeted_golden_dataset, INDUSTRY_VERTICAL, "user queries", 50, AGENT_INSTRUCTION)
future_tools = executor.submit(tool_engine.generate_targeted_golden_dataset, INDUSTRY_VERTICAL, "tool calls", 50, AGENT_INSTRUCTION)
future_responses = executor.submit(response_engine.generate_targeted_golden_dataset, INDUSTRY_VERTICAL, "agent responses", 50, AGENT_INSTRUCTION)

Builds separate safe neighborhoods for prompts, tool-call JSON, and final responses.

Risk-tier scoring

is_novelty_anomaly = (
    prompt_eval["distance"] > TRACE_THRESHOLDS["prompt"]
    or (tool_eval and tool_eval["distance"] > TRACE_THRESHOLDS["tool"])
    or response_eval["distance"] > TRACE_THRESHOLDS["response"]
)
if safety_metadata.get("finish_reason") == "SAFETY" or is_novelty_anomaly:
    risk_tier = "TIER 1 (NOVELTY/CRITICAL)"
    audit_required = True

Promotes statistically novel or safety-blocked traces to exhaustive audit.

Models & APIs used

  • Models: gemini-2.5-flash, text-embedding-004
  • APIs / services: Vertex AI, Agent Platform API, Agent Engine, Vector Search, BigQuery, Cloud Storage
  • SDKs / libraries: google-genai, google-cloud-aiplatform, google-adk, google-auth, google-cloud-bigquery, google-cloud-storage, scikit-learn, numpy, pandas, ipywidgets, matplotlib, seaborn, pydantic, db-dtypes

When to use this

Use this pattern when monitoring an ADK or Agent Engine agent for brand safety by comparing prompts, tool calls, and responses to known-safe embedding neighborhoods.

Gotchas & caveats

  • Requires an existing Google Cloud project and Agent Platform API enabled.
  • Colab runs require google.colab auth.authenticate_user().
  • PROJECT_ID, LOCATION, BUCKET_NAME, BigQuery IDs, and existing Vector Search endpoint names must be configured correctly.
  • Agent Engine deployment can take around 10 minutes.
  • The notebook staggers Vector Search index deployment requests to avoid Google Cloud limits.
  • Notebook-scale synthetic trace counts may be too small to reach the stated 95% statistical confidence.
  • Thresholds are tunable defaults and should be calibrated with plotted cosine-distance distributions.
  • RUN_NEW_INFERENCE=False depends on previously logged BigQuery traces.

Best practices

  • Compare prompts, tool calls, and responses against separate Vector Search indices to reduce structural false positives.
  • Generate industry-specific golden baselines using the agent instruction and INDUSTRY_VERTICAL.
  • Request JSON from Gemini with response_mime_type=“application/json” and parse it with json.loads.
  • Log full execution traces to BigQuery as a flight recorder for later threshold tuning without new LLM calls.
  • Route safety finish reasons and novelty anomalies to Tier 1 for 100% audit.
  • Use Tier 2 sampling for negative sentiment and Tier 3 pulse sampling for baseline drift detection.
  • Visualize prompt, tool, and response distance distributions before finalizing thresholds.
  • Stagger parallel index deployments when provisioning multiple Vector Search indices.

Graph View

Backlinks