131 lines
5.5 KiB
Python
131 lines
5.5 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
orchestrator.py — OSVauco multi-agent orchestrator using ADK 2.0 collaborative workflow.
|
|
|
|
Architecture:
|
|
root_agent (coordinator / LLM-driven delegation)
|
|
├── rag_agent — retrieves answers from private knowledge base
|
|
├── gcp_ops_agent — handles GCP operations questions and script generation
|
|
├── memory_agent — loads/stores long-term memory via Memory Bank
|
|
└── farewell_agent — session closings
|
|
|
|
Pattern: Coordinator with sub_agents list.
|
|
Requires: google-adk >= 2.0.0
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import logging
|
|
from typing import Optional
|
|
|
|
from google.adk.agents import Agent
|
|
from google.adk.agents.callback_context import CallbackContext
|
|
from google.adk.tools.load_memory_tool import load_memory_tool
|
|
from google.adk.tools.preload_memory_tool import preload_memory_tool
|
|
from google.genai.types import Content, Part
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
PROJECT_ID = os.environ.get("PROJECT_ID", "propane-will-491900-m5")
|
|
REGION = os.environ.get("REGION", "us-central1")
|
|
RAG_CORPUS_NAME = os.environ.get("RAG_CORPUS_NAME", "")
|
|
ORCHESTRATOR_MODEL = os.environ.get("ORCHESTRATOR_MODEL", "gemini-2.5-flash")
|
|
SUBAGENT_MODEL = os.environ.get("SUBAGENT_MODEL", "gemini-2.5-flash")
|
|
|
|
# ── Safety callbacks ──────────────────────────────────────────────────────────
|
|
BLOCKED_PATTERNS = [
|
|
"ignore previous instructions",
|
|
"ignore all instructions",
|
|
"drop table",
|
|
"system prompt",
|
|
"jailbreak",
|
|
"disregard safety",
|
|
]
|
|
|
|
def before_model_callback(callback_context: CallbackContext, llm_request) -> Optional[Content]:
|
|
try:
|
|
user_text = ""
|
|
if llm_request.contents:
|
|
last = llm_request.contents[-1]
|
|
if last.parts:
|
|
user_text = last.parts[0].text.lower()
|
|
for pattern in BLOCKED_PATTERNS:
|
|
if pattern in user_text:
|
|
logger.warning("Blocked pattern detected: '%s'", pattern)
|
|
return Content(parts=[Part(text="I cannot process that request. Please rephrase.")])
|
|
except Exception as exc:
|
|
logger.error("before_model_callback error: %s", exc)
|
|
return None
|
|
|
|
def before_tool_callback(tool, args: dict, tool_context) -> Optional[dict]:
|
|
dangerous = ["DROP", "DELETE FROM", "TRUNCATE", "--", ";"]
|
|
for val in args.values():
|
|
if isinstance(val, str):
|
|
for d in dangerous:
|
|
if d.upper() in val.upper():
|
|
raise ValueError(f"Tool argument rejected by safety guardrail: '{val}'")
|
|
return None
|
|
|
|
# ── Sub-agents ────────────────────────────────────────────────────────────────
|
|
rag_agent = Agent(
|
|
model=SUBAGENT_MODEL,
|
|
name="rag_agent",
|
|
description="Retrieves answers from the private OSVauco knowledge base (RAG corpus).",
|
|
instruction=(
|
|
"You are a knowledge retrieval specialist. "
|
|
"Use search_knowledge_base to find relevant information and return well-cited answers."
|
|
),
|
|
before_model_callback=before_model_callback,
|
|
)
|
|
|
|
gcp_ops_agent = Agent(
|
|
model=SUBAGENT_MODEL,
|
|
name="gcp_ops_agent",
|
|
description="Answers GCP operations questions: scripts, IAM, billing, Cloud Run, ADK deployments.",
|
|
instruction=(
|
|
"You are a GCP operations expert for project propane-will-491900-m5 in us-central1. "
|
|
"Provide accurate gcloud CLI commands, IAM patterns, and ADK deployment guidance. "
|
|
"Always include cost-safety reminders (teardown, billing budgets)."
|
|
),
|
|
before_model_callback=before_model_callback,
|
|
)
|
|
|
|
memory_agent = Agent(
|
|
model=SUBAGENT_MODEL,
|
|
name="memory_agent",
|
|
description="Manages long-term memory: loads past context and stores new facts for future sessions.",
|
|
instruction=(
|
|
"You manage the agent's long-term memory. "
|
|
"Use load_memory_tool to retrieve past facts and preload_memory_tool to store important new facts. "
|
|
"Only persist high-value, factual information — not transient conversation."
|
|
),
|
|
tools=[load_memory_tool, preload_memory_tool],
|
|
before_model_callback=before_model_callback,
|
|
)
|
|
|
|
farewell_agent = Agent(
|
|
model=SUBAGENT_MODEL,
|
|
name="farewell_agent",
|
|
description="Handles session closings, summaries, and goodbye messages.",
|
|
instruction="Generate a concise, friendly session summary and closing message.",
|
|
)
|
|
|
|
# ── Root orchestrator ─────────────────────────────────────────────────────────
|
|
root_agent = Agent(
|
|
model=ORCHESTRATOR_MODEL,
|
|
name="oavauco_orchestrator",
|
|
description="OSVauco root orchestrator — delegates to specialist sub-agents.",
|
|
instruction=(
|
|
"You are the OSVauco orchestrator for project propane-will-491900-m5. "
|
|
"Delegate to sub-agents based on the user's intent:\n"
|
|
"- Knowledge base questions → rag_agent\n"
|
|
"- GCP operations, scripts, IAM, billing → gcp_ops_agent\n"
|
|
"- Memory recall or storage → memory_agent\n"
|
|
"- Session endings → farewell_agent\n"
|
|
"Always synthesize sub-agent responses into a clear, concise final answer."
|
|
),
|
|
sub_agents=[rag_agent, gcp_ops_agent, memory_agent, farewell_agent],
|
|
before_model_callback=before_model_callback,
|
|
before_tool_callback=before_tool_callback,
|
|
)
|